Direct replication with randomly generated objects on source node

environment.yml

@@ -23,4 +23,6 @@ dependencies:
     - paramiko
     - questionary
     - ray
-    - tqdm
+    - setproctitle
+    - tqdm
+    - werkzeug

setup.py

@@ -18,12 +18,14 @@
         "graphviz",
         "loguru",
         "matplotlib",
-        "ray",
         "numpy",
         "pandas",
         "paramiko",
         "questionary",
+        "ray",
+        "setproctitle",
         "tqdm",
+        "werkzeug",
     ],
     extras_require={"test": ["black", "pytest", "ipython", "jupyter_console"]}
 )

skylark/__init__.py

@@ -4,3 +4,16 @@
 data_root = skylark_root / "data"
 instance_log_root = skylark_root / "data" / "logs"
 key_root = skylark_root / "data" / "keys"
+
+
+def print_header():
+    header = """=================================================
+  ______  _             _                 _     
+ / _____)| |           | |               | |    
+( (____  | |  _  _   _ | |  _____   ____ | |  _ 
+ \____ \ | |_/ )| | | || | (____ | / ___)| |_/ )
+ _____) )|  _ ( | |_| || | / ___ || |    |  _ ( 
+(______/ |_| \_) \__  | \_)\_____||_|    |_| \_)
+                (____/                          
+================================================="""
+    print(header, flush=True)

skylark/compute/aws/aws_cloud_provider.py

@@ -112,7 +112,7 @@ def get_ubuntu_ami_id(region: str) -> str:
             return image_list[0]["ImageId"]
 
     def provision_instance(
-        self, region: str, instance_class: str, name: Optional[str] = None, ami_id: Optional[str] = None, tags={"skylark": "true"}
+        self, region: str, instance_class: str, name: Optional[str] = None, ami_id: Optional[str] = None, tags={"skylark": "true"}, ebs_volume_size: int = 128
     ) -> AWSServer:
         assert not region.startswith("aws:"), "Region should be AWS region"
         if name is None:
@@ -121,6 +121,7 @@ def provision_instance(
             ami_id = self.get_ubuntu_ami_id(region)
         ec2 = AWSServer.get_boto3_resource("ec2", region)
         AWSServer.make_keyfile(region)
+        # set instance storage to 128GB EBS
         instance = ec2.create_instances(
             ImageId=ami_id,
             InstanceType=instance_class,
@@ -133,6 +134,16 @@ def provision_instance(
                     "Tags": [{"Key": "Name", "Value": name}] + [{"Key": k, "Value": v} for k, v in tags.items()],
                 }
             ],
+            BlockDeviceMappings=[
+                {
+                    "DeviceName": "/dev/sda1",
+                    "Ebs": {
+                        "DeleteOnTermination": True,
+                        "VolumeSize": ebs_volume_size,
+                        "VolumeType": "gp2",
+                    },
+                }
+            ],
         )
         server = AWSServer(f"aws:{region}", instance[0].id)
         server.wait_for_ready()

skylark/gateway/chunk_store.py

@@ -1,8 +1,9 @@
 from dataclasses import asdict, dataclass
-from enum import Enum
+from enum import Enum, auto
 from multiprocessing import Manager
+from os import PathLike
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
 from loguru import logger
 
@@ -70,16 +71,20 @@ def from_dict(in_dict: Dict):
 
 
 class ChunkState(Enum):
-    REGISTERED = "registered"
-    DOWNLOAD_IN_PROGRESS = "download_in_progress"
-    READY_TO_UPLOAD = "ready_to_upload"
-    UPLOAD_IN_PROGRESS = "upload_in_progress"
-    UPLOAD_COMPLETE = "upload_complete"
-    FAILED = "failed"
+    registered = auto()
+    download_in_progress = auto()
+    downloaded = auto()
+    upload_in_progress = auto()
+    upload_complete = auto()
+    failed = auto()
+
+    @staticmethod
+    def from_str(s: str):
+        return ChunkState[s.lower()]
 
 
 class ChunkStore:
-    def __init__(self, chunk_dir: str = "/dev/shm/skylark/chunks"):
+    def __init__(self, chunk_dir: PathLike):
         self.chunk_dir = Path(chunk_dir)
         self.chunk_dir.mkdir(parents=True, exist_ok=True)
 
@@ -90,13 +95,9 @@ def __init__(self, chunk_dir: str = "/dev/shm/skylark/chunks"):
 
         # multiprocess-safe concurrent structures
         self.manager = Manager()
-        self.chunks: Dict[int, Chunk] = self.manager.dict()
+        self.chunk_requests: Dict[int, Chunk] = self.manager.dict()
         self.chunk_status: Dict[int, ChunkState] = self.manager.dict()
 
-        self.pending_chunk_requests: List[ChunkRequest] = self.manager.list()
-        self.downloaded_chunk_requests: List[ChunkRequest] = self.manager.list()
-        self.uploaded_chunk_requests: List[ChunkRequest] = self.manager.list()
-
     def get_chunk_file_path(self, chunk_id: int) -> Path:
         return self.chunk_dir / f"{chunk_id}.chunk"
 
@@ -110,80 +111,63 @@ def set_chunk_status(self, chunk_id: int, new_status: ChunkState):
         self.chunk_status[chunk_id] = new_status
 
     def start_download(self, chunk_id: int):
-        if self.get_chunk_status(chunk_id) == ChunkState.REGISTERED:
-            self.set_chunk_status(chunk_id, ChunkState.DOWNLOAD_IN_PROGRESS)
+        state = self.get_chunk_status(chunk_id)
+        if state in [ChunkState.registered, ChunkState.download_in_progress]:
+            self.set_chunk_status(chunk_id, ChunkState.download_in_progress)
         else:
             raise ValueError(f"Invalid transition start_download from {self.get_chunk_status(chunk_id)}")
 
     def finish_download(self, chunk_id: int, runtime_s: Optional[float] = None):
         # todo log runtime to statistics store
-        if self.get_chunk_status(chunk_id) == ChunkState.DOWNLOAD_IN_PROGRESS:
-            self.set_chunk_status(chunk_id, ChunkState.READY_TO_UPLOAD)
+        state = self.get_chunk_status(chunk_id)
+        if state in [ChunkState.download_in_progress, ChunkState.downloaded]:
+            self.set_chunk_status(chunk_id, ChunkState.downloaded)
         else:
             raise ValueError(f"Invalid transition finish_download from {self.get_chunk_status(chunk_id)}")
 
     def start_upload(self, chunk_id: int):
-        if self.get_chunk_status(chunk_id) == ChunkState.READY_TO_UPLOAD:
-            self.set_chunk_status(chunk_id, ChunkState.UPLOAD_IN_PROGRESS)
+        state = self.get_chunk_status(chunk_id)
+        if state in [ChunkState.downloaded, ChunkState.upload_in_progress]:
+            self.set_chunk_status(chunk_id, ChunkState.upload_in_progress)
         else:
             raise ValueError(f"Invalid transition start_upload from {self.get_chunk_status(chunk_id)}")
 
     def finish_upload(self, chunk_id: int, runtime_s: Optional[float] = None):
         # todo log runtime to statistics store
-        if self.get_chunk_status(chunk_id) == ChunkState.UPLOAD_IN_PROGRESS:
-            self.set_chunk_status(chunk_id, ChunkState.UPLOAD_COMPLETE)
+        state = self.get_chunk_status(chunk_id)
+        if state in [ChunkState.upload_in_progress, ChunkState.upload_complete]:
+            self.set_chunk_status(chunk_id, ChunkState.upload_complete)
         else:
             raise ValueError(f"Invalid transition finish_upload from {self.get_chunk_status(chunk_id)}")
 
     def fail(self, chunk_id: int):
-        if self.get_chunk_status(chunk_id) != ChunkState.UPLOAD_COMPLETE:
-            self.set_chunk_status(chunk_id, ChunkState.FAILED)
+        if self.get_chunk_status(chunk_id) != ChunkState.upload_complete:
+            self.set_chunk_status(chunk_id, ChunkState.failed)
         else:
             raise ValueError(f"Invalid transition fail from {self.get_chunk_status(chunk_id)}")
 
     ###
     # Chunk management
     ###
-    def get_chunks(self):
-        return self.chunks.values()
-
-    def get_chunk(self, chunk_id: int) -> Optional[Chunk]:
-        return self.chunks[chunk_id] if chunk_id in self.chunks else None
-
-    def add_chunk(self, chunk: Chunk):
-        self.chunks[chunk.chunk_id] = chunk
-        self.set_chunk_status(chunk.chunk_id, "registered")
-
-    def get_chunk_requests(self) -> Tuple[List[ChunkRequest], List[ChunkRequest], List[ChunkRequest]]:
-        return list(self.pending_chunk_requests), list(self.downloaded_chunk_requests), list(self.uploaded_chunk_requests)
+    def get_chunk_requests(self, status: Optional[ChunkState] = None) -> List[ChunkRequest]:
+        if status is None:
+            return list(self.chunk_requests.values())
+        else:
+            return [req for i, req in self.chunk_requests.items() if self.get_chunk_status(i) == status]
 
     def get_chunk_request(self, chunk_id: int) -> Optional[ChunkRequest]:
-        for chunk_request in self.pending_chunk_requests:
-            if chunk_request.chunk.chunk_id == chunk_id:
-                return chunk_request
-        for chunk_request in self.downloaded_chunk_requests:
-            if chunk_request.chunk.chunk_id == chunk_id:
-                return chunk_request
-        for chunk_request in self.uploaded_chunk_requests:
-            if chunk_request.chunk.chunk_id == chunk_id:
-                return chunk_request
+        return self.chunk_requests[chunk_id] if chunk_id in self.chunk_requests else None
+
+    def add_chunk_request(self, chunk_request: ChunkRequest, state=ChunkState.registered):
+        logger.debug(f"Adding chunk request {chunk_request.chunk.chunk_id}")
+        self.set_chunk_status(chunk_request.chunk.chunk_id, state)
+        self.chunk_requests[chunk_request.chunk.chunk_id] = chunk_request
+
+    def pop_chunk_request_path(self, chunk_id: int) -> Optional[ChunkRequestHop]:
+        if chunk_id in self.chunk_requests:
+            chunk_request = self.chunk_requests[chunk_id]
+            if len(chunk_request.path) > 0:
+                result = chunk_request.path.pop(0)
+                self.chunk_requests[chunk_id] = chunk_request
+                return result
         return None
-
-    def add_chunk_request_pending(self, chunk_request: ChunkRequest):
-        self.pending_chunk_requests.append(chunk_request)
-
-    def add_chunk_request_downloaded(self, chunk_request: ChunkRequest):
-        self.downloaded_chunk_requests.append(chunk_request)
-
-    def add_chunk_request_uploaded(self, chunk_request: ChunkRequest):
-        self.uploaded_chunk_requests.append(chunk_request)
-
-    def mark_chunk_request_downloaded(self, chunk_request: ChunkRequest):
-        assert chunk_request in self.pending_chunk_requests
-        self.pending_chunk_requests.remove(chunk_request)
-        self.downloaded_chunk_requests.append(chunk_request)
-
-    def mark_chunk_request_uploaded(self, chunk_request: ChunkRequest):
-        assert chunk_request in self.downloaded_chunk_requests
-        self.downloaded_chunk_requests.remove(chunk_request)
-        self.uploaded_chunk_requests.append(chunk_request)