Authentication from config

scripts/experiment.sh

@@ -16,7 +16,7 @@ dest=$2
 
 #key_prefix="synthetic-fake-imagenet/4_16384"
 key_prefix="fake_imagenet"
-bucket_prefix="exps222"
+bucket_prefix="exps"
 src_bucket=(${src//:/ })
 src_bucket=${bucket_prefix}-skylark-${src_bucket[1]}
 dest_bucket=(${dest//:/ })
@@ -32,7 +32,7 @@ echo $filename
 export GOOGLE_APPLICATION_CREDENTIALS="/home/ubuntu/.skylark-shishir-42be5f375b7a.json"
 
 # creats buckets + bucket data and sets env variables
-#python scripts/setup_bucket.py --key-prefix ${key_prefix} --bucket-prefix ${bucket_prefix} --gcp-project skylark-shishir --src-data-path ../${key_prefix}/ --src-region ${src} --dest-region ${dest}
+python scripts/setup_bucket.py --key-prefix ${key_prefix} --bucket-prefix ${bucket_prefix} --gcp-project skylark-shishir --src-data-path ../${key_prefix}/ --src-region ${src} --dest-region ${dest}
 
 
 # TODO:artificially increase the number of chunks 

skylark/compute/server.py

@@ -5,8 +5,8 @@
 from enum import Enum, auto
 from pathlib import Path
 from typing import Dict
-
 import requests
+from skylark import config_file
 from skylark.utils import logger
 from skylark.compute.utils import make_dozzle_command, make_sysctl_tcp_tuning_command
 from skylark.utils.utils import PathLike, Timer, retry_backoff, wait_for
@@ -213,16 +213,12 @@ def check_stderr(tup):
 
         # read AWS config file to get credentials
         # TODO: Integrate this with updated skylark config file
-        docker_envs = ""
-        try:
-            config = configparser.RawConfigParser()
-            config.read(os.path.expanduser("~/.aws/credentials"))
-            aws_access_key_id = config.get("default", "aws_access_key_id")
-            aws_secret_access_key = config.get("default", "aws_secret_access_key")
-            docker_envs += f" -e AWS_ACCESS_KEY_ID='{aws_access_key_id}'"
-            docker_envs += f" -e AWS_SECRET_ACCESS_KEY='{aws_secret_access_key}'"
-        except Exception as e:
-            logger.error(f"Failed to read AWS credentials locally {e}")
+        # copy config file
+        config = config_file.read_text()[:-2] + "}"
+        config = json.dumps(config)  # Convert to JSON string and remove trailing comma/new-line
+        self.run_command(f'mkdir -p /opt; echo "{config}" | sudo tee /opt/{config_file.name} > /dev/null')
+
+        docker_envs = ""  # If needed, add environment variables to docker command
 
         with Timer(f"{desc_prefix}: Docker pull"):
             docker_out, docker_err = self.run_command(f"sudo docker pull {gateway_docker_image}")
@@ -232,6 +228,7 @@ def check_stderr(tup):
             f"-d --rm --log-driver=local --log-opt max-file=16 --ipc=host --network=host --ulimit nofile={1024 * 1024} {docker_envs}"
         )
         docker_run_flags += " --mount type=tmpfs,dst=/skylark,tmpfs-size=$(($(free -b  | head -n2 | tail -n1 | awk '{print $2}')/2))"
+        docker_run_flags += f" -v /opt/{config_file.name}:/pkg/data/{config_file.name}"
         gateway_daemon_cmd = f"python -u /pkg/skylark/gateway/gateway_daemon.py --chunk-dir /skylark/chunks --outgoing-ports '{json.dumps(outgoing_ports)}' --region {self.region_tag}"
         docker_launch_cmd = f"sudo docker run {docker_run_flags} --name skylark_gateway {gateway_docker_image} {gateway_daemon_cmd}"
         start_out, start_err = self.run_command(docker_launch_cmd)

skylark/obj_store/azure_interface.py

@@ -1,13 +1,11 @@
 import os
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Iterator, List
-
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
-
-# from azure.identity import DefaultAzureCredential
+from azure.identity import DefaultAzureCredential, ClientSecretCredential
 from azure.storage.blob import BlobServiceClient
+from skylark.config import load_config
 from skylark.utils import logger
-from skylark.obj_store.azure_keys import azure_storage_credentials
 from skylark.obj_store.object_store_interface import NoSuchObjectException, ObjectStoreInterface, ObjectStoreObject
 
 
@@ -20,20 +18,23 @@ class AzureInterface(ObjectStoreInterface):
     def __init__(self, azure_region, container_name):
         # TODO: the azure region should get corresponding os.getenv()
         self.azure_region = azure_region
-        assert self.azure_region in azure_storage_credentials
-
         self.container_name = container_name
         self.bucket_name = self.container_name  # For compatibility
         self.pending_downloads, self.completed_downloads = 0, 0
         self.pending_uploads, self.completed_uploads = 0, 0
-
-        # Connection strings are stored in azure_keys.py
-        self._connect_str = azure_storage_credentials[self.azure_region]["connection_string"]
-        self.blob_service_client = BlobServiceClient.from_connection_string(self._connect_str)
-        # self.azure_default_credential = DefaultAzureCredential()
-        # self.blob_service_client = BlobServiceClient(account_url=account_url, credential=self.azure_default_credential)
-
-        self.pool = ThreadPoolExecutor(max_workers=256)  # TODO: Figure this out, since azure by default has 15 workers
+        # Authenticate
+        config = load_config()
+        self.subscription_id = config["azure_subscription_id"]
+        self.credential = ClientSecretCredential(
+            tenant_id=config["azure_tenant_id"],
+            client_id=config["azure_client_id"],
+            client_secret=config["azure_client_secret"],
+        )
+        # Create a blob service client
+        self.account_url = "https://{}.blob.core.windows.net".format("skylark" + self.azure_region)
+        self.blob_service_client = BlobServiceClient(account_url=self.account_url, credential=self.credential)
+
+        self.pool = ThreadPoolExecutor(max_workers=256)  # TODO: This might need some tuning
         self.max_concurrency = 1
         self.container_client = None
         if not self.container_exists():

skylark/obj_store/s3_interface.py

@@ -10,6 +10,7 @@
 from awscrt.s3 import S3Client, S3RequestTlsMode, S3RequestType
 
 from skylark.compute.aws.aws_server import AWSServer
+from skylark.config import load_config
 from skylark.obj_store.object_store_interface import NoSuchObjectException, ObjectStoreInterface, ObjectStoreObject
 
 
@@ -33,10 +34,10 @@ def __init__(self, aws_region, bucket_name, use_tls=True, part_size=None, throug
         event_loop_group = EventLoopGroup(num_threads=num_threads, cpu_group=None)
         host_resolver = DefaultHostResolver(event_loop_group)
         bootstrap = ClientBootstrap(event_loop_group, host_resolver)
-
-        # get aws auth info for docker envs
-        aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", None)
-        aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY", None)
+        # Authenticate
+        config = load_config()
+        aws_access_key_id = config["aws_access_key_id"]
+        aws_secret_access_key = config["aws_secret_access_key"]
         if aws_access_key_id and aws_secret_access_key:
             credential_provider = AwsCredentialsProvider.new_static(aws_access_key_id, aws_secret_access_key)
         else:  # use default

skylark/replicate/replicator_client.py

@@ -392,7 +392,7 @@ def fn(s: Server):
                                 tqdm.write(log_line)
                             else:
                                 logger.debug(log_line)
-                        elif t.elapsed > 180 and completed_bytes == 0:
+                        elif t.elapsed > 600 and completed_bytes == 0:
                             logger.error(f"No chunks completed after {int(t.elapsed)}s! There is probably a bug, check logs. Exiting...")
                             return dict(
                                 completed_chunk_ids=completed_chunk_ids,