Support private versus public IPs for gateways

skyplane/api/dataplane.py

@@ -218,9 +218,12 @@ def provision(
     def copy_gateway_logs(self):
         # copy logs from all gateways in parallel
         def copy_log(instance):
+            out_file = self.transfer_dir / f"gateway_{instance.uuid()}.stdout"
+            err_file = self.transfer_dir / f"gateway_{instance.uuid()}.stdout"
+            logger.fs.info(f"[Dataplane.copy_gateway_logs] Copying logs from {instance.uuid()}: {out_file}")
             instance.run_command("sudo docker logs -t skyplane_gateway 2> /tmp/gateway.stderr > /tmp/gateway.stdout")
-            instance.download_file("/tmp/gateway.stdout", self.transfer_dir / f"gateway_{instance.uuid()}.stdout")
-            instance.download_file("/tmp/gateway.stderr", self.transfer_dir / f"gateway_{instance.uuid()}.stderr")
+            instance.download_file("/tmp/gateway.stdout", out_file)
+            instance.download_file("/tmp/gateway.stderr", err_file)
 
         do_parallel(copy_log, self.bound_nodes.values(), n=-1)
 

skyplane/api/provisioner.py

@@ -164,7 +164,6 @@ def _provision_task(self, task: ProvisionerTask):
                 server = self.azure.provision_instance(task.region, task.vm_type, use_spot_instances=task.spot, tags=task.tags)
             elif task.cloud_provider == "gcp":
                 assert self.gcp.auth.enabled(), "GCP credentials not configured"
-                # TODO: specify network tier in
                 server = self.gcp.provision_instance(
                     task.region,
                     task.vm_type,
@@ -248,10 +247,12 @@ def provision(self, authorize_firewall: bool = True, max_jobs: int = 16, spinner
             # NOTE: the following setup is for broadcast only
             if aws_provisioned:
                 authorize_ip_jobs.extend([partial(self.aws.add_ips_to_security_group, r, None) for r in set(aws_regions)])
-            # if gcp_provisioned:
-            #     def authorize_gcp_gateways():
-            #         self.gcp_firewall_rules.add(self.gcp.authorize_gateways(public_ips + private_ips))
-            #     authorize_ip_jobs.append(authorize_gcp_gateways)
+            if gcp_provisioned:
+
+                def authorize_gcp_gateways():
+                    self.gcp_firewall_rules.add(self.gcp.authorize_gateways(public_ips + private_ips))
+
+                authorize_ip_jobs.append(authorize_gcp_gateways)
 
             do_parallel(
                 lambda fn: fn(),

skyplane/api/tracker.py

@@ -1,5 +1,6 @@
 import functools
 from pprint import pprint
+from pprint import pprint
 import json
 import time
 from abc import ABC
@@ -221,12 +222,7 @@ def monitor_single_dst_helper(dst_region):
                 job.finalize()
         except Exception as e:
             UsageClient.log_exception(
-                "finalize job",
-                e,
-                args,
-                self.dataplane.topology.src_region_tag,
-                self.dataplane.topology.dest_region_tags,
-                session_start_timestamp_ms,
+                "finalize job", e, args, self.dataplane.src_region_tag, self.dataplane.dst_region_tag, session_start_timestamp_ms
             )
             raise e
         end_time = int(time.time())
@@ -374,11 +370,5 @@ def query_bytes_dispatched(self):
             return 0
         bytes_total_per_job = {}
         for job_uuid in self.job_complete_chunk_ids.keys():
-            bytes_total_per_job[job_uuid] = sum(
-                [
-                    cr.chunk_length_bytes
-                    for cr in self.job_chunk_requests[job_uuid].values()
-                    # if cr.chunk_id in self.job_complete_chunk_ids[job_uuid]
-                ]
-            )
+            bytes_total_per_job[job_uuid] = sum([cr.chunk_length_bytes for cr in self.job_chunk_requests[job_uuid].values()])
         return sum(bytes_total_per_job.values())

skyplane/api/transfer_job.py

@@ -10,6 +10,8 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
 from queue import Queue
+
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Callable, Generator, List, Optional, Tuple, TypeVar, Dict
 
 from abc import ABC, abstractmethod
@@ -341,7 +343,12 @@ def chunk(self, transfer_pair_generator: Generator[TransferPair, None, None]) ->
                     yield multipart_chunk_queue.get()
 
         if self.transfer_config.multipart_enabled:
-            while not multipart_send_queue.empty() or not multipart_chunk_queue.empty():
+            # wait for processing multipart requests to finish
+            logger.fs.debug("Waiting for multipart threads to finish")
+            # while not multipart_send_queue.empty():
+            # TODO: may be an issue waiting for this in case of force-quit
+            while not multipart_send_queue.empty():
+                logger.fs.debug(f"Remaining in multipart queue: sent {multipart_send_queue.qsize()}")
                 time.sleep(0.1)
             # send sentinel to all threads
             multipart_exit_event.set()
@@ -587,6 +594,7 @@ def dispatch(
         bytes_dispatched = [0] * len(src_gateways)
         n_multiparts = 0
         start = time.time()
+
         for batch in batches:
             # send upload_id mappings to sink gateways
             upload_id_batch = [cr for cr in batch if cr.upload_id_mapping is not None]
@@ -642,7 +650,6 @@ def dispatch(
 
     def finalize(self):
         """Complete the multipart upload requests"""
-        print("Finalizing multipart uploads...")
         typer.secho(f"Finalizing multipart uploads...", fg="bright_black")
         groups = defaultdict(list)
         for req in self.multipart_transfer_list:
@@ -657,9 +664,10 @@ def finalize(self):
 
             def complete_fn(batch):
                 for req in batch:
+                    logger.fs.debug(f"Finalize upload id {req['upload_id']} for key {req['key']}")
                     obj_store_interface.complete_multipart_upload(req["key"], req["upload_id"])
 
-            do_parallel(complete_fn, batches, n=-1)
+            do_parallel(complete_fn, batches, n=8)
 
     def verify(self):
         """Verify the integrity of the transfered destination objects"""

skyplane/api/usage.py

@@ -93,7 +93,7 @@ class UsageStatsToReport:
     #: The source region of the transfer session.
     source_region: str
     #: The destination region of the transfer session.
-    destination_region: str
+    destination_region: str  # TODO: make into dest_regions
     #: The source cloud provider of the transfer session.
     source_cloud_provider: str
     #: The destination cloud provider of the transfer session.

skyplane/compute/ibmcloud/ibmcloud_auth.py

@@ -31,9 +31,7 @@ def __init__(self, config: Optional[SkyplaneConfig] = None):
 
     @imports.inject("ibm_cloud_sdk_core", "ibm_cloud_sdk_core.authenticators", pip_extra="ibmcloud")
     def get_iam_authenticator(ibm_cloud_sdk_core, self):
-        from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
-
-        return IAMAuthenticator(self.config.ibmcloud_iam_key, url=self.config.ibmcloud_iam_endpoint)
+        return ibm_cloud_sdk_core.authenticators.IAMAuthenticator(self.config.ibmcloud_iam_key, url=self.config.ibmcloud_iam_endpoint)
 
     def get_ibmcloud_endpoint(self, region, compute_backend="public"):
         if region is not None:

skyplane/gateway/chunk_store.py

@@ -41,7 +41,7 @@ def add_partition(self, partition_id: str):
             raise ValueError(f"Partition {partition_id} already exists")
         self.chunk_requests[partition_id] = GatewayQueue()
 
-    def add_partition(self, partition_id: str, queue: Optional[GatewayQueue]):
+    def add_partition(self, partition_id: str, queue: GatewayQueue):
         """Create a queue for this partition."""
         print("Adding partition", partition_id, queue)
         if partition_id in self.chunk_requests:

skyplane/gateway/gateway_daemon.py

@@ -218,11 +218,12 @@ def create_gateway_operators_helper(input_queue, program: List[Dict], partition_
                 elif op["op_type"] == "send":
                     # TODO: handle private ips for GCP->GCP
                     target_gateway_info = self.gateway_info[op["target_gateway_id"]]
-                    print("Gateway sender sending to ", target_gateway_info["private_ip_address"])
+                    ip_addr = target_gateway_info["private_ip_address"] if op["private_ip"] else target_gateway_info["public_ip_address"]
+                    print("Gateway sender sending to ", ip_addr, "private", op["private_ip"])
                     operators[handle] = GatewaySender(
                         handle,
                         region=self.region,
-                        ip_addr=target_gateway_info["private_ip_address"],
+                        ip_addr=ip_addr,
                         input_queue=input_queue,
                         output_queue=output_queue,
                         error_event=self.error_event,

skyplane/gateway/gateway_program.py

@@ -32,13 +32,22 @@ def __repr__(self):
 
 
 class GatewaySend(GatewayOperator):
-    def __init__(self, target_gateway_id: str, region: str, num_connections: int = 32, compress: bool = False, encrypt: bool = False):
+    def __init__(
+        self,
+        target_gateway_id: str,
+        region: str,
+        num_connections: int = 32,
+        compress: bool = False,
+        encrypt: bool = False,
+        private_ip: bool = False,
+    ):
         super().__init__("send")
         self.target_gateway_id = target_gateway_id  # gateway to send to
         self.region = region  # region to send to
         self.num_connections = num_connections  # default this for now
         self.compress = compress
         self.encrypt = encrypt
+        self.private_ip = private_ip  # whether to send to private or public IP (private for GCP->GCP)
 
 
 class GatewayReceive(GatewayOperator):

skyplane/gateway/gateway_queue.py

@@ -13,7 +13,7 @@ def put(self, chunk_req):
         self.q.put(chunk_req)
 
     def pop(self, requester_handle=None):
-        self.q.pop()
+        self.q.get()
 
     def get_nowait(self, requester_handle=None):
         return self.q.get_nowait()
@@ -48,7 +48,7 @@ def put(self, chunk_req):
             self.q[handle].put(chunk_req)
 
     def pop(self, requester_handle):
-        self.q[requester_handle].pop()
+        self.q[requester_handle].get()
 
     def get_nowait(self, requester_handle):
         return self.q[requester_handle].get_nowait()

skyplane/gateway/operators/gateway_operator.py

@@ -7,7 +7,7 @@
 import time
 import traceback
 from functools import partial
-from multiprocessing import Event, Process
+from multiprocessing import Event, Process, Queue
 from multiprocessing.managers import DictProxy
 from typing import Dict, List, Optional
 
@@ -35,7 +35,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         chunk_store: ChunkStore,
         n_processes: Optional[int] = 1,
     ):
@@ -163,7 +163,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         chunk_store: ChunkStore,
         ip_addr: str,
         use_tls: Optional[bool] = True,
@@ -358,7 +358,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         chunk_store: ChunkStore,
         size_mb: int,
         n_processes: Optional[int] = 1,
@@ -398,7 +398,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         chunk_store: ChunkStore,
         n_processes: int = 1,
     ):
@@ -417,7 +417,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         n_processes: Optional[int] = 1,
         chunk_store: Optional[ChunkStore] = None,
         bucket_name: Optional[str] = None,
@@ -452,7 +452,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         n_processes: int = 32,
         chunk_store: Optional[ChunkStore] = None,
         bucket_name: Optional[str] = None,
@@ -525,7 +525,7 @@ def __init__(
         input_queue: GatewayQueue,
         output_queue: GatewayQueue,
         error_event,
-        error_queue: GatewayQueue,
+        error_queue: Queue,
         upload_id_map: DictProxy,  # map of upload_id mappings from client
         n_processes: Optional[int] = 32,
         chunk_store: Optional[ChunkStore] = None,

skyplane/planner/planner.py

@@ -1,5 +1,6 @@
 from importlib.resources import path
 from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 from skyplane import compute
 
@@ -118,6 +119,7 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan:
         for job in jobs:
             src_bucket = job.src_iface.bucket()
             src_region_tag = job.src_iface.region_tag()
+            src_provider = src_region_tag.split(":")[0]
 
             # give each job a different partition id, so we can read/write to different buckets
             partition_id = jobs.index(job)
@@ -139,9 +141,16 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan:
                 # can send to any gateway in region
                 mux_or = src_program.add_operator(GatewayMuxOr(), parent_handle=mux_and, partition_id=partition_id)
                 for i in range(self.n_instances):
+                    private_ip = False
+                    if dst_gateways[i].provider == "gcp" and src_provider == "gcp":
+                        print("Using private IP for GCP to GCP transfer", src_region_tag, dst_region_tag)
+                        private_ip = True
                     src_program.add_operator(
                         GatewaySend(
-                            target_gateway_id=dst_gateways[i].gateway_id, region=dst_region_tag, num_connections=self.n_connections
+                            target_gateway_id=dst_gateways[i].gateway_id,
+                            region=dst_region_tag,
+                            num_connections=self.n_connections,
+                            private_ip=private_ip,
                         ),
                         parent_handle=mux_or,
                         partition_id=partition_id,

skyplane/planner/topology.py

@@ -65,8 +65,8 @@ def regions(self) -> List[str]:
 
     def add_gateway(self, region_tag: str):
         """Create gateway in specified region"""
-        gateway_id = region_tag + str(len([gateway for gateway in self.gateways.values() if gateway.region == region_tag]))
-        assert gateway_id not in self.gateways
+        gateway_id = region_tag + str(len([gateway for gateway in self.gateways.values() if gateway.region_tag == region_tag]))
+        assert gateway_id not in self.gateways, f"Gateway id {gateway_id} in {self.gateways}"
         gateway = TopologyPlanGateway(region_tag, gateway_id)
         self.gateways[gateway_id] = gateway
         return gateway