Refactor dependencies for gateway (#251)

skyplane-project · Apr 6, 2022 · 7759318 · 7759318
1 parent 6bb7806
commit 7759318
Show file tree

Hide file tree

Showing 16 changed files with 85 additions and 65 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -12,16 +12,16 @@ RUN (echo 'net.ipv4.ip_local_port_range = 12000 65535' >> /etc/sysctl.conf) \
 
 # install apt packages
 RUN --mount=type=cache,target=/var/cache/apt apt update \
-    && apt-get install --no-install-recommends -y git wget ca-certificates build-essential graphviz \
+    && apt-get install --no-install-recommends -y ca-certificates build-essential \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
 # install gateway
 COPY scripts/requirements-gateway.txt /tmp/requirements-gateway.txt
-RUN --mount=type=cache,target=/root/.cache/pip pip3 install --no-cache-dir --compile -r /tmp/requirements-gateway.txt && rm -r /tmp/requirements-gateway.txt
+RUN --mount=type=cache,target=/root/.cache/pip pip3 install --no-cache-dir -r /tmp/requirements-gateway.txt && rm -r /tmp/requirements-gateway.txt
 
 WORKDIR /pkg
 COPY . .
-RUN pip3 install -e .
+RUN pip3 install --no-dependencies -e ".[gateway]"
 
 CMD ["python3", "skylark/gateway/gateway_daemon.py"]
diff --git a/scripts/experiment_paras.sh b/scripts/experiment_paras.sh
@@ -49,6 +49,7 @@ cp ${filename} data/results/${experiment}
 skylark replicate-json ${filename} \
     --source-bucket $src_bucket \
     --dest-bucket $dest_bucket \
-    --key-prefix ${key_prefix} |& tee data/results/${experiment}/obj-store-logs.txt
+    --src-key-prefix ${key_prefix} \
+    --dest-key-prefix ${key_prefix} |& tee data/results/${experiment}/obj-store-logs.txt
 tail -1 data/results/${experiment}/obj-store-logs.txt
 echo ${experiment}
diff --git a/scripts/requirements-gateway.txt b/scripts/requirements-gateway.txt
@@ -1,24 +1,26 @@
+# cloud dependencies
 awscrt
 azure-identity
-azure-mgmt-compute
-azure-mgmt-network
-azure-mgmt-resource
+# removed from gateway due to large size
+# azure-mgmt-compute
+# azure-mgmt-network
+# azure-mgmt-resource
 azure-mgmt-storage
 azure-mgmt-authorization
 azure-storage-blob>=12.0.0
 boto3
-click>=7.1.2
-flask
 google-api-python-client
 google-auth
 google-cloud-compute
 google-cloud-storage
 grpcio-status>=1.33.2
-numpy
+# shared dependencies
+cachetools
 oslo.concurrency
 paramiko
-setproctitle
 termcolor
-testresources
 tqdm
+# gateway dependencies
+flask
+pyopenssl
 werkzeug
diff --git a/setup.py b/setup.py
@@ -7,6 +7,7 @@
     packages=["skylark"],
     python_requires=">=3.7",
     install_requires=[
+        # cloud integrations
         "awscrt",
         "azure-identity",
         "azure-mgmt-compute",
@@ -16,36 +17,41 @@
         "azure-mgmt-authorization",
         "azure-storage-blob>=12.0.0",
         "boto3",
-        "click>=7.1.2",
-        "flask",
         "google-api-python-client",
         "google-auth",
         "google-cloud-compute",
         "google-cloud-storage",
         "grpcio-status>=1.33.2",
-        "numpy",
+        # client dependencies
+        "click",
+        "pandas",
+        "questionary",
+        "typer",
+        # shared dependencies
+        "cachetools",
         "oslo.concurrency",
         "paramiko",
-        "pyopenssl",
-        "setproctitle",
         "termcolor",
-        "testresources",
         "tqdm",
-        "werkzeug",
     ],
     extras_require={
-        "all": [
-            "cachetools",
-            "click",
+        "solver": [
             "cvxopt",
             "cvxpy",
             "graphviz",
             "matplotlib",
-            "pandas",
-            "questionary",
-            "typer",
+            "numpy",
+        ],
+        "gateway": [
+            "flask",
+            "pyopenssl",
+            "werkzeug",
+        ],
+        "experiments": [
+            "matplotlib",
+            "numpy",
+            "ray",
         ],
-        "experiments": ["ray"],
         "test": ["black", "ipython", "jupyter_console", "pytest", "pytype"],
     },
     entry_points={"console_scripts": ["skylark=skylark.cli.cli:app"]},

diff --git a/skylark/cli/cli_helper.py b/skylark/cli/cli_helper.py
@@ -355,9 +355,11 @@ def load_aws_config(config: SkylarkConfig) -> SkylarkConfig:
     config.aws_enabled = True
     return config
 
+
 def create_aws_region_config(config):
     AWSAuthentication.save_region_config(config)
 
+
 def load_azure_config(config: SkylarkConfig, force_init: bool = False) -> SkylarkConfig:
     if force_init:
         typer.secho("    Azure credentials will be re-initialized", fg="red")

diff --git a/skylark/compute/aws/aws_auth.py b/skylark/compute/aws/aws_auth.py
@@ -1,13 +1,11 @@
-from pydoc import describe
 import threading
 from typing import Optional
-import typer
 
 import boto3
 
 from skylark.config import SkylarkConfig
 from skylark import config_path
-from skylark import aws_config_path 
+from skylark import aws_config_path
 
 
 class AWSAuthentication:
@@ -19,7 +17,7 @@ def __init__(self, config: Optional[SkylarkConfig] = None, access_key: Optional[
             self.config = config
         else:
             self.config = SkylarkConfig.load_config(config_path)
-        
+
         if access_key and secret_key:
             self.config_mode = "manual"
             self._access_key = access_key
@@ -28,29 +26,29 @@ def __init__(self, config: Optional[SkylarkConfig] = None, access_key: Optional[
             self.config_mode = "iam_inferred"
             self._access_key = None
             self._secret_key = None
-    
+
     @staticmethod
     def save_region_config(config):
         with open(aws_config_path, "w") as f:
             if config.aws_enabled == False:
                 f.write("")
                 return
             region_list = []
-            describe_regions = boto3.client('ec2', region_name="us-east-1").describe_regions()
-            for region in describe_regions['Regions']:
-                if region['OptInStatus'] == 'opt-in-not-required' or region['OptInStatus'] == 'opted-in':
-                    region_text = region['Endpoint']
-                    region_name = region_text[region_text.find('.') + 1 :region_text.find(".amazon")]
+            describe_regions = boto3.client("ec2", region_name="us-east-1").describe_regions()
+            for region in describe_regions["Regions"]:
+                if region["OptInStatus"] == "opt-in-not-required" or region["OptInStatus"] == "opted-in":
+                    region_text = region["Endpoint"]
+                    region_name = region_text[region_text.find(".") + 1 : region_text.find(".amazon")]
                     region_list.append(region_name)
             f.write("\n".join(region_list))
-            typer.secho(f"    AWS region config file saved to {aws_config_path}", fg="green")
+            print(f"    AWS region config file saved to {aws_config_path}")
 
     @staticmethod
     def get_region_config():
         try:
             f = open(aws_config_path, "r")
         except FileNotFoundError:
-            typer.secho("    No AWS config detected! Consquently, the AWS region list is empty. Run 'skylark init' to remedy this.", fg="red")
+            print("    No AWS config detected! Consquently, the AWS region list is empty. Run 'skylark init' to remedy this.")
             return []
         region_list = []
         for region in f.read().split("\n"):

diff --git a/skylark/compute/aws/aws_cloud_provider.py b/skylark/compute/aws/aws_cloud_provider.py
@@ -5,10 +5,8 @@
 from pathlib import Path
 
 import botocore
-import pandas as pd
 from skylark.compute.aws.aws_auth import AWSAuthentication
 from skylark.utils import logger
-import typer
 from skylark import key_root
 
 from oslo_concurrency import lockutils
@@ -17,6 +15,12 @@
 from skylark.compute.cloud_providers import CloudProvider
 from skylark.utils.utils import retry_backoff, wait_for
 
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+    logger.warning("pandas not installed, will not be able to load transfer costs")
+
 
 class AWSCloudProvider(CloudProvider):
     def __init__(self):

diff --git a/skylark/compute/azure/azure_auth.py b/skylark/compute/azure/azure_auth.py
@@ -3,9 +3,6 @@
 import threading
 from typing import Optional
 from azure.identity import DefaultAzureCredential
-from azure.mgmt.compute import ComputeManagementClient
-from azure.mgmt.network import NetworkManagementClient
-from azure.mgmt.resource import ResourceManagementClient
 from azure.mgmt.authorization import AuthorizationManagementClient
 from azure.mgmt.storage import StorageManagementClient
 from azure.storage.blob import BlobServiceClient, ContainerClient
@@ -15,6 +12,22 @@
 from skylark.config import SkylarkConfig
 from skylark import config_path
 
+# optional imports due to large package size
+try:
+    from azure.mgmt.network import NetworkManagementClient
+except ImportError:
+    NetworkManagementClient = None
+
+try:
+    from azure.mgmt.compute import ComputeManagementClient
+except ImportError:
+    ComputeManagementClient = None
+
+try:
+    from azure.mgmt.resource import ResourceManagementClient
+except ImportError:
+    ResourceManagementClient = None
+
 
 class AzureAuthentication:
     __cached_credentials = threading.local()
@@ -55,12 +68,15 @@ def get_token(self, resource: str):
         return self.credential.get_token(resource)
 
     def get_compute_client(self):
+        assert ComputeManagementClient is not None, "ComputeManagementClient is not installed"
         return ComputeManagementClient(self.credential, self.subscription_id)
 
     def get_resource_client(self):
+        assert ResourceManagementClient is not None, "ResourceManagementClient is not installed"
         return ResourceManagementClient(self.credential, self.subscription_id)
 
     def get_network_client(self):
+        assert NetworkManagementClient is not None, "NetworkManagementClient is not installed"
         return NetworkManagementClient(self.credential, self.subscription_id)
 
     def get_authorization_client(self):

diff --git a/skylark/compute/azure/azure_cloud_provider.py b/skylark/compute/azure/azure_cloud_provider.py
@@ -1,21 +1,20 @@
-from multiprocessing import BoundedSemaphore
 import os
 import uuid
+from multiprocessing import BoundedSemaphore
 from pathlib import Path
 from typing import List, Optional
 
 import paramiko
-from skylark.compute.azure.azure_auth import AzureAuthentication
-from skylark.utils import logger
 from skylark import key_root
+from skylark.compute.azure.azure_auth import AzureAuthentication
 from skylark.compute.azure.azure_server import AzureServer
 from skylark.compute.cloud_providers import CloudProvider
-from azure.mgmt.authorization.models import RoleAssignmentCreateParameters
+from skylark.utils import logger
+from skylark.utils.utils import Timer, do_parallel
 
+from azure.mgmt.authorization.models import RoleAssignmentCreateParameters
 from azure.mgmt.compute.models import ResourceIdentityType
 
-from skylark.utils.utils import Timer, do_parallel
-
 
 class AzureCloudProvider(CloudProvider):
     def __init__(self, key_root=key_root / "azure"):

diff --git a/skylark/config.py b/skylark/config.py
@@ -18,12 +18,12 @@ class SkylarkConfig:
     @staticmethod
     def default_config() -> "SkylarkConfig":
         return SkylarkConfig(
-                aws_enabled=False,
-                azure_enabled=False,
-                gcp_enabled=False,
-                azure_subscription_id=None,
-                gcp_project_id=None,
-            )
+            aws_enabled=False,
+            azure_enabled=False,
+            gcp_enabled=False,
+            azure_subscription_id=None,
+            gcp_project_id=None,
+        )
 
     @staticmethod
     def load_config(path) -> "SkylarkConfig":

diff --git a/skylark/utils/cert.py → skylark/gateway/cert.py b/skylark/utils/cert.py → skylark/gateway/cert.py
diff --git a/skylark/gateway/gateway_daemon.py b/skylark/gateway/gateway_daemon.py
@@ -13,7 +13,6 @@
 from typing import Dict
 
 
-import setproctitle
 from skylark.utils import logger
 from skylark import MB, print_header
 from skylark.chunk import ChunkState
@@ -47,7 +46,6 @@ def __init__(self, region: str, outgoing_ports: Dict[str, int], chunk_dir: PathL
         logger.info(f"[gateway_daemon] API started at {self.api_server.url}")
 
     def run(self):
-        setproctitle.setproctitle(f"skylark-gateway-daemon")
         exit_flag = Event()
 
         def exit_handler(signum, frame):

diff --git a/skylark/gateway/gateway_obj_store.py b/skylark/gateway/gateway_obj_store.py
@@ -4,7 +4,6 @@
 from multiprocessing import Event, Manager, Process, Value
 from typing import Dict, Optional
 
-import setproctitle
 from skylark.gateway.chunk_store import ChunkStore
 from skylark.utils import logger
 from skylark.chunk import ChunkRequest
@@ -61,7 +60,6 @@ def stop_workers(self):
 
     def worker_loop(self, worker_id: int):
         # todo should this use processes instead of threads?
-        setproctitle.setproctitle(f"skylark-gateway-obj-store:{worker_id}")
         self.worker_id = worker_id
 
         while not self.exit_flags[worker_id].is_set():

diff --git a/skylark/gateway/gateway_receiver.py b/skylark/gateway/gateway_receiver.py
@@ -7,12 +7,11 @@
 from multiprocessing import Event, Process, Value
 from typing import Tuple
 
-import setproctitle
 from skylark.utils import logger
 from skylark import GB, MB
 from skylark.chunk import WireProtocolHeader
 from skylark.gateway.chunk_store import ChunkStore
-from skylark.utils.cert import generate_self_signed_certificate
+from skylark.gateway.cert import generate_self_signed_certificate
 from skylark.utils.utils import Timer
 
 
@@ -56,7 +55,6 @@ def signal_handler(signal, frame):
                     exit_flag.value = 1
 
                 signal.signal(signal.SIGINT, signal_handler)
-                setproctitle.setproctitle(f"skylark-gateway-receiver:{socket_port}")
 
                 sock.listen()
                 if self.ssl_context is not None:

diff --git a/skylark/gateway/gateway_sender.py b/skylark/gateway/gateway_sender.py
@@ -6,7 +6,6 @@
 from typing import Dict, List, Optional
 
 import requests
-import setproctitle
 from skylark.utils import logger
 from skylark import MB
 from skylark.chunk import ChunkRequest
@@ -57,7 +56,6 @@ def stop_workers(self):
         self.processes = []
 
     def worker_loop(self, worker_id: int, dest_ip: str):
-        setproctitle.setproctitle(f"skylark-gateway-sender:{worker_id}")
         self.worker_id = worker_id
 
         while not self.exit_flags[worker_id].is_set():