diff --git a/skyplane/api/dataplane.py b/skyplane/api/dataplane.py index c8253a922..799edfd39 100644 --- a/skyplane/api/dataplane.py +++ b/skyplane/api/dataplane.py @@ -22,6 +22,10 @@ from skyplane.utils.definitions import gateway_docker_image, tmp_log_dir from skyplane.utils.fn import PathLike, do_parallel +from skyplane.compute.aws.aws_server import AWSServer +from skyplane.compute.gcp.gcp_server import GCPServer +from skyplane.compute.azure.azure_server import AzureServer + if TYPE_CHECKING: from skyplane.api.provisioner import Provisioner @@ -123,6 +127,7 @@ def _start_gateway( use_bbr=self.transfer_config.use_bbr, # TODO: remove use_compression=self.transfer_config.use_compression, use_socket_tls=self.transfer_config.use_socket_tls, + instance_path=gateway_node.gateway_instance_path, # TODO: better way of mapping the path of VM src/dst ) def provision( @@ -163,13 +168,16 @@ def provision( assert ( cloud_provider != "cloudflare" ), f"Cannot create VMs in certain cloud providers: check planner output {self.topology.to_dict()}" - self.provisioner.add_task( - cloud_provider=cloud_provider, - region=region, - vm_type=node.vm_type or getattr(self.transfer_config, f"{cloud_provider}_instance_class"), - spot=getattr(self.transfer_config, f"{cloud_provider}_use_spot_instances"), - autoterminate_minutes=self.transfer_config.autoterminate_minutes, - ) + + # Only provision if it is not VM source or destination + if node.gateway_instance_id is None: + self.provisioner.add_task( + cloud_provider=cloud_provider, + region=region, + vm_type=node.vm_type or getattr(self.transfer_config, f"{cloud_provider}_instance_class"), + spot=getattr(self.transfer_config, f"{cloud_provider}_use_spot_instances"), + autoterminate_minutes=self.transfer_config.autoterminate_minutes, + ) # initialize clouds self.provisioner.init_global(aws=is_aws_used, azure=is_azure_used, gcp=is_gcp_used, ibmcloud=is_ibmcloud_used) @@ -186,8 +194,19 @@ def provision( servers_by_region = defaultdict(list) for s in servers: servers_by_region[s.region_tag].append(s) + for node in self.topology.get_gateways(): - instance = servers_by_region[node.region_tag].pop() + if node.region_tag not in servers_by_region: + if node.region_tag.startswith("aws"): + instance = AWSServer(node.region_tag, node.gateway_instance_id, key_path=node.gateway_key_path) + elif node.region_tag.startswith("azure"): + instance = AzureServer(node.gateway_instance_id) + elif node.region_tag.startswith("gcp"): + instance = GCPServer(node.region_tag, node.gateway_instance_id) + else: + raise Exception(f"Invalid region tag: {node.region_tag}") + else: + instance = servers_by_region[node.region_tag].pop() self.bound_nodes[node] = instance # set ip addresses (for gateway program generation) diff --git a/skyplane/api/pipeline.py b/skyplane/api/pipeline.py index 6fa3face4..193f6996e 100644 --- a/skyplane/api/pipeline.py +++ b/skyplane/api/pipeline.py @@ -10,7 +10,14 @@ from skyplane.api.transfer_job import CopyJob, SyncJob, TransferJob from skyplane.api.config import TransferConfig -from skyplane.planner.planner import MulticastDirectPlanner, DirectPlannerSourceOneSided, DirectPlannerDestOneSided +from skyplane.planner.planner import ( + MulticastDirectPlanner, + DirectPlannerSourceOneSided, + DirectPlannerDestOneSided, + DirectPlannerVMSource, + DirectPlannerVMDest, + DirectPlannerVMSourceDest, +) from skyplane.planner.topology import TopologyPlanGateway from skyplane.utils import logger from skyplane.utils.definitions import tmp_log_dir @@ -67,6 +74,12 @@ def __init__( self.planner = DirectPlannerSourceOneSided(self.max_instances, self.n_connections, self.transfer_config) elif self.planning_algorithm == "dst_one_sided": self.planner = DirectPlannerDestOneSided(self.max_instances, self.n_connections, self.transfer_config) + elif self.planning_algorithm == "vm_source": + self.planner = DirectPlannerVMSource(self.max_instances, 64, self.transfer_config) + elif self.planning_algorithm == "vm_dest": + self.planner = DirectPlannerVMDest(self.max_instances, 64, self.transfer_config) + elif self.planning_algorithm == "vm_to_vm": + self.planner = DirectPlannerVMSourceDest(self.max_instances, 64, self.transfer_config) else: raise ValueError(f"No such planning algorithm {planning_algorithm}") diff --git a/skyplane/api/transfer_job.py b/skyplane/api/transfer_job.py index 0155b7c17..57e3333f6 100644 --- a/skyplane/api/transfer_job.py +++ b/skyplane/api/transfer_job.py @@ -14,7 +14,16 @@ from queue import Queue from abc import ABC -from typing import TYPE_CHECKING, Callable, Generator, List, Optional, Tuple, TypeVar, Dict +from typing import ( + TYPE_CHECKING, + Callable, + Generator, + List, + Optional, + Tuple, + TypeVar, + Dict, +) from abc import ABC @@ -27,6 +36,7 @@ from skyplane.chunk import Chunk from skyplane.obj_store.storage_interface import StorageInterface from skyplane.obj_store.object_store_interface import ObjectStoreObject, ObjectStoreInterface +from skyplane.obj_store.vm_interface import VMInterface from skyplane.utils import logger from skyplane.utils.definitions import MB from skyplane.utils.fn import do_parallel @@ -108,10 +118,12 @@ def _run_multipart_chunk_thread( upload_id_mapping = {} for dest_iface in self.dst_ifaces: dest_object = dest_objects[dest_iface.region_tag()] + upload_id = dest_iface.initiate_multipart_upload(dest_object.key, mime_type=mime_type) # print(f"Created upload id for key {dest_object.key} with upload id {upload_id} for bucket {dest_iface.bucket_name}") # store mapping between key and upload id for each region upload_id_mapping[dest_iface.region_tag()] = (src_object.key, upload_id) + out_queue_chunks.put(GatewayMessage(upload_id_mapping=upload_id_mapping)) # send to output queue # get source and destination object and then compute number of chunks @@ -164,7 +176,15 @@ def _run_multipart_chunk_thread( metadata = (block_ids, mime_type) self.multipart_upload_requests.append( - dict(upload_id=upload_id, key=dest_object.key, parts=parts, region=region, bucket=bucket, metadata=metadata) + dict( + upload_id=upload_id, + key=dest_object.key, + parts=parts, + region=region, + bucket=bucket, + metadata=metadata, + vm=True if dest_iface.provider == "vm" else False, + ) ) else: mime_type = None @@ -291,24 +311,34 @@ def transfer_pair_generator( logger.fs.exception(e) raise e from None - if dest_provider == "aws": - from skyplane.obj_store.s3_interface import S3Object - - dest_obj = S3Object(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) - elif dest_provider == "azure": - from skyplane.obj_store.azure_blob_interface import AzureBlobObject + if isinstance(dst_iface, VMInterface): + # VM destination + from skyplane.obj_store.vm_interface import VMFile - dest_obj = AzureBlobObject(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) - elif dest_provider == "gcp": - from skyplane.obj_store.gcs_interface import GCSObject + host_ip = dst_iface.host_ip() - dest_obj = GCSObject(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) - elif dest_provider == "cloudflare": - from skyplane.obj_store.r2_interface import R2Object + dest_obj = VMFile(provider=dest_provider, bucket=host_ip, key=dest_key) - dest_obj = R2Object(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) else: - raise ValueError(f"Invalid dest_region {dest_region}, unknown provider") + # Bucket destination + if dest_provider == "aws": + from skyplane.obj_store.s3_interface import S3Object + + dest_obj = S3Object(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) + elif dest_provider == "azure": + from skyplane.obj_store.azure_blob_interface import AzureBlobObject + + dest_obj = AzureBlobObject(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) + elif dest_provider == "gcp": + from skyplane.obj_store.gcs_interface import GCSObject + + dest_obj = GCSObject(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) + elif dest_provider == "cloudflare": + from skyplane.obj_store.r2_interface import R2Object + + dest_obj = R2Object(provider=dest_provider, bucket=dst_iface.bucket(), key=dest_key) + else: + raise ValueError(f"Invalid dest_region {dest_region}, unknown provider") dest_objs[dst_iface.region_tag()] = dest_obj # assert that all destinations share the same post-fix key @@ -332,7 +362,7 @@ def chunk(self, transfer_pair_generator: Generator[TransferPair, None, None]) -> multipart_chunk_threads = [] # start chunking threads - if self.transfer_config.multipart_enabled: + if self.transfer_config.multipart_enabled: # and not isinstance(self.dst_ifaces[0], VMInterface): for _ in range(self.concurrent_multipart_chunk_threads): t = threading.Thread( target=self._run_multipart_chunk_thread, @@ -346,7 +376,11 @@ def chunk(self, transfer_pair_generator: Generator[TransferPair, None, None]) -> for transfer_pair in transfer_pair_generator: # print("transfer_pair", transfer_pair.src_obj.key, transfer_pair.dst_objs) src_obj = transfer_pair.src_obj - if self.transfer_config.multipart_enabled and src_obj.size > self.transfer_config.multipart_threshold_mb * MB: + if ( + self.transfer_config.multipart_enabled + # and not isinstance(self.dst_ifaces[0], VMInterface) + and src_obj.size > self.transfer_config.multipart_threshold_mb * MB + ): multipart_send_queue.put(transfer_pair) else: if transfer_pair.src_obj.size == 0: @@ -362,12 +396,12 @@ def chunk(self, transfer_pair_generator: Generator[TransferPair, None, None]) -> ) ) - if self.transfer_config.multipart_enabled: + if self.transfer_config.multipart_enabled: # and not isinstance(self.dst_ifaces[0], VMInterface): # drain multipart chunk queue and yield with updated chunk IDs while not multipart_chunk_queue.empty(): yield multipart_chunk_queue.get() - if self.transfer_config.multipart_enabled: + if self.transfer_config.multipart_enabled: # and not isinstance(self.dst_ifaces[0], VMInterface): # wait for processing multipart requests to finish logger.fs.debug("Waiting for multipart threads to finish") # while not multipart_send_queue.empty(): @@ -697,11 +731,14 @@ def finalize(self): for req in self.multipart_transfer_list: if "region" not in req or "bucket" not in req: raise Exception(f"Invalid multipart upload request: {req}") - groups[(req["region"], req["bucket"])].append(req) + groups[(req["region"], req["bucket"], req["vm"])].append(req) for key, group in groups.items(): - region, bucket = key + region, bucket, vm = key batch_len = max(1, len(group) // 128) batches = [group[i : i + batch_len] for i in range(0, len(group), batch_len)] + print(f"region: {region}, bucket: {bucket}") + if vm: + region = "vm:" + region obj_store_interface = StorageInterface.create(region, bucket) def complete_fn(batch): @@ -723,14 +760,19 @@ def verify(self): def verify_region(i): dst_iface = self.dst_ifaces[i] dst_prefix = self.dst_prefixes[i] + print("Dst prefix: ", dst_prefix) # gather destination key mapping for this region dst_keys = {pair.dst_objs[dst_iface.region_tag()].key: pair.src_obj for pair in self.transfer_list} + print(f"Destination key mappings: {dst_keys}") # list and check destination prefix for obj in dst_iface.list_objects(dst_prefix): + print(f"Object listed: {obj.key}") # check metadata (src.size == dst.size) && (src.modified <= dst.modified) src_obj = dst_keys.get(obj.key) + print(f"src_obj: {src_obj}") + print(f"Object: {obj}") if src_obj and src_obj.size == obj.size and src_obj.last_modified <= obj.last_modified: del dst_keys[obj.key] diff --git a/skyplane/cli/cli.py b/skyplane/cli/cli.py index eb1d38f5d..27eac5147 100644 --- a/skyplane/cli/cli.py +++ b/skyplane/cli/cli.py @@ -8,7 +8,7 @@ import skyplane.cli.cli_cloud import skyplane.cli.cli_config -# import skyplane.cli.experiments # disable experiments +import skyplane.cli.experiments # disable experiments from skyplane import compute from skyplane.cli.cli_init import init @@ -30,7 +30,7 @@ name="init", help="Initialize the Skyplane CLI with your cloud credentials", )(init) -# app.add_typer(skyplane.cli.experiments.app, name="experiments") # disable experiments +app.add_typer(skyplane.cli.experiments.app, name="experiments") # disable experiments app.add_typer(skyplane.cli.cli_cloud.app, name="cloud") app.add_typer(skyplane.cli.cli_config.app, name="config") diff --git a/skyplane/cli/cli_transfer.py b/skyplane/cli/cli_transfer.py index c1d6620fd..6d34d3255 100644 --- a/skyplane/cli/cli_transfer.py +++ b/skyplane/cli/cli_transfer.py @@ -313,19 +313,28 @@ def run_transfer( register_exception_handler() print_header() - provider_src, bucket_src, path_src = parse_path(src) - provider_dst, bucket_dst, path_dst = parse_path(dst) + provider_src, transfer_src, path_src = parse_path(src) + provider_dst, transfer_dst, path_dst = parse_path(dst) # update planner for one-sided transfer # somet process for other cloud providers with no VM support - assert provider_src != "cloudflare" or provider_dst != "cloudflare", "Cannot transfer between two Cloudflare buckets" - if provider_src == "cloudflare": - solver = "dst_one_sided" - elif provider_dst == "cloudflare": - solver = "src_one_sided" - - src_region_tag = StorageInterface.create(f"{provider_src}:infer", bucket_src).region_tag() - dst_region_tag = StorageInterface.create(f"{provider_dst}:infer", bucket_dst).region_tag() + if provider_src == "vm" and provider_dst == "vm": + solver = "vm_to_vm" + elif provider_src == "vm": + solver = "vm_source" + elif provider_dst == "vm": + solver = "vm_dest" + else: + # the previous handling for non-VM transfers + assert provider_src != "cloudflare" or provider_dst != "cloudflare", "Cannot transfer between two Cloudflare buckets" + if provider_src == "cloudflare": + solver = "dst_one_sided" + elif provider_dst == "cloudflare": + solver = "src_one_sided" + + src_region_tag = StorageInterface.create(f"{provider_src}:infer", transfer_src).region_tag() + dst_region_tag = StorageInterface.create(f"{provider_dst}:infer", transfer_dst).region_tag() + args = { "cmd": cmd, "recursive": True, @@ -371,7 +380,8 @@ def run_transfer( # fallback option: transfer is too small if cli.args["cmd"] == "cp": job = CopyJob(src, [dst], recursive=recursive) # TODO: rever to using pipeline - if cli.estimate_small_transfer(job, cloud_config.get_flag("native_cmd_threshold_gb") * GB): + if cli.estimate_small_transfer(job, 0.01 * GB): # Test small transfer + # if cli.estimate_small_transfer(job, cloud_config.get_flag("native_cmd_threshold_gb") * GB): small_transfer_status = cli.transfer_cp_small(src, dst, recursive) return 0 if small_transfer_status else 1 else: diff --git a/skyplane/cli/experiments/__init__.py b/skyplane/cli/experiments/__init__.py index dac9869f0..a076b4f04 100644 --- a/skyplane/cli/experiments/__init__.py +++ b/skyplane/cli/experiments/__init__.py @@ -1,7 +1,13 @@ import typer from skyplane.cli.experiments.cli_profile import latency_grid, throughput_grid -from skyplane.cli.experiments.cli_query import get_max_throughput, util_grid_throughput, util_grid_cost, dump_full_util_cost_grid +from skyplane.cli.experiments.cli_query import ( + get_max_throughput, + util_grid_throughput, + util_grid_cost, + dump_full_util_cost_grid, +) +from skyplane.cli.experiments.cli_create_instance import create_instance app = typer.Typer(name="experiments") app.command()(latency_grid) @@ -10,3 +16,4 @@ app.command()(util_grid_throughput) app.command()(util_grid_cost) app.command()(dump_full_util_cost_grid) +app.command()(create_instance) diff --git a/skyplane/cli/experiments/cli_create_instance.py b/skyplane/cli/experiments/cli_create_instance.py new file mode 100644 index 000000000..f541c9f67 --- /dev/null +++ b/skyplane/cli/experiments/cli_create_instance.py @@ -0,0 +1,155 @@ +import typer +from typing import List + +from skyplane import compute +from skyplane.cli.experiments.provision import provision +from skyplane.compute.const_cmds import make_sysctl_tcp_tuning_command +from skyplane.utils import logger +from skyplane.utils.fn import do_parallel + +all_aws_regions = compute.AWSCloudProvider.region_list() +all_azure_regions = compute.AzureCloudProvider.region_list() +all_gcp_regions = compute.GCPCloudProvider.region_list() +all_gcp_regions_standard = compute.GCPCloudProvider.region_list_standard() +all_ibmcloud_regions = compute.IBMCloudProvider.region_list() +from skyplane.compute.aws.aws_auth import AWSAuthentication + + +def aws_credentials(): + auth = AWSAuthentication() + access_key, secret_key = auth.get_credentials() + return access_key, secret_key + + +def create_instance( + # regions + aws_region_list: List[str] = typer.Option(all_aws_regions, "-aws"), + azure_region_list: List[str] = typer.Option(all_azure_regions, "-azure"), + gcp_region_list: List[str] = typer.Option(all_gcp_regions, "-gcp"), + gcp_standard_region_list: List[str] = typer.Option(all_gcp_regions_standard, "-gcp-standard"), + ibmcloud_region_list: List[str] = typer.Option(all_ibmcloud_regions, "-ibmcloud"), + # + enable_aws: bool = typer.Option(True), + enable_azure: bool = typer.Option(False), + enable_gcp: bool = typer.Option(False), + enable_gcp_standard: bool = typer.Option(False), + enable_ibmcloud: bool = typer.Option(False), + # instances to provision + aws_instance_class: str = typer.Option("m5.8xlarge", help="AWS instance class to use"), + azure_instance_class: str = typer.Option("Standard_D32_v5", help="Azure instance class to use"), + gcp_instance_class: str = typer.Option("n2-standard-32", help="GCP instance class to use"), + ibmcloud_instance_class: str = typer.Option("bx2-2x8", help="IBM Cloud instance class to use"), +): + def check_stderr(tup): + assert tup[1].strip() == "", f"Command failed, err: {tup[1]}" + + # validate arguments + # aws_region_list = aws_region_list if enable_aws else [] + # azure_region_list = azure_region_list if enable_azure else [] + # gcp_region_list = gcp_region_list if enable_gcp else [] + + # aws_region_list = ["us-east-1", "us-west-1"] + # gcp_region_list = ['me-west1-a','europe-north1-a',] + aws_region_list = ["us-east-1"] + # gcp_region_list = ["me-west1-a"] + + # validate AWS regions + aws_region_list = aws_region_list if enable_aws else [] + azure_region_list = azure_region_list if enable_azure else [] + gcp_region_list = gcp_region_list if enable_gcp else [] + ibmcloud_region_list = ibmcloud_region_list if enable_ibmcloud else [] + if not enable_aws and not enable_azure and not enable_gcp and not enable_ibmcloud: + logger.error("At least one of -aws, -azure, -gcp, -ibmcloud must be enabled.") + raise typer.Abort() + + # validate AWS regions + if not enable_aws: + aws_region_list = [] + elif not all(r in all_aws_regions for r in aws_region_list): + logger.error(f"Invalid AWS region list: {aws_region_list}") + raise typer.Abort() + + # validate Azure regions + if not enable_azure: + azure_region_list = [] + elif not all(r in all_azure_regions for r in azure_region_list): + logger.error(f"Invalid Azure region list: {azure_region_list}") + raise typer.Abort() + + # validate GCP regions + assert not enable_gcp_standard or enable_gcp, f"GCP is disabled but GCP standard is enabled" + if not enable_gcp: + gcp_region_list = [] + elif not all(r in all_gcp_regions for r in gcp_region_list): + logger.error(f"Invalid GCP region list: {gcp_region_list}") + raise typer.Abort() + + # validate GCP standard instances + if not enable_gcp_standard: + gcp_standard_region_list = [] + if not all(r in all_gcp_regions_standard for r in gcp_standard_region_list): + logger.error(f"Invalid GCP standard region list: {gcp_standard_region_list}") + raise typer.Abort() + + # validate IBM Cloud regions + if not enable_ibmcloud: + ibmcloud_region_list = [] + elif not all(r in all_ibmcloud_regions for r in ibmcloud_region_list): + logger.error(f"Invalid IBM Cloud region list: {ibmcloud_region_list}") + raise typer.Abort() + + # provision servers + aws = compute.AWSCloudProvider() + azure = compute.AzureCloudProvider() + gcp = compute.GCPCloudProvider() + ibmcloud = compute.IBMCloudProvider() + + aws_instances, azure_instances, gcp_instances, ibmcloud_instances = provision( + aws=aws, + azure=azure, + gcp=gcp, + ibmcloud=ibmcloud, + aws_regions_to_provision=aws_region_list, + azure_regions_to_provision=azure_region_list, + gcp_regions_to_provision=gcp_region_list, + ibmcloud_regions_to_provision=ibmcloud_region_list, + aws_instance_class=aws_instance_class, + azure_instance_class=azure_instance_class, + gcp_instance_class=gcp_instance_class, + ibmcloud_instance_class=ibmcloud_instance_class, + aws_instance_os="ubuntu", + gcp_instance_os="ubuntu", + gcp_use_premium_network=True, + ) + instance_list: List[compute.Server] = [i for ilist in aws_instances.values() for i in ilist] + instance_list.extend([i for ilist in azure_instances.values() for i in ilist]) + instance_list.extend([i for ilist in gcp_instances.values() for i in ilist]) + + # setup instances + def setup(server: compute.Server): + check_stderr(server.run_command("echo 'debconf debconf/frontend select Noninteractive' | sudo debconf-set-selections")) + check_stderr( + server.run_command( + "sudo add-apt-repository universe;\ + (sudo apt-get update && sudo apt-get install python3-pip -y && sudo pip3 install awscli)" + ) + ) + check_stderr(server.run_command(make_sysctl_tcp_tuning_command(cc="cubic"))) + server.run_command( + f"aws configure set aws_access_key_id {aws_credentials()[0]}; aws configure set aws_secret_access_key {aws_credentials()[1]}" + ) + + do_parallel(setup, instance_list, spinner=True, n=-1, desc="Setup") + + with open("ssh_cmd.txt", "a") as f: + for instance in instance_list: + print("instance: ", instance.region_tag) + ssh_cmd = instance.get_ssh_cmd() + print(ssh_cmd) + + # Insert the '-o StrictHostKeyChecking=accept-new' option in the middle of the ssh command + ssh_parts = ssh_cmd.split(" ", 1) + modified_ssh_cmd = f"{ssh_parts[0]} -o StrictHostKeyChecking=accept-new {ssh_parts[1]}" + f.write(modified_ssh_cmd + "\n") + + f.close() diff --git a/skyplane/compute/aws/aws_auth.py b/skyplane/compute/aws/aws_auth.py index 532c01e20..b7b88956e 100644 --- a/skyplane/compute/aws/aws_auth.py +++ b/skyplane/compute/aws/aws_auth.py @@ -100,6 +100,18 @@ def secret_key(self): def enabled(self): return self.config.aws_enabled + @imports.inject("boto3", pip_extra="aws") + def get_credentials(boto3, self): + cached_credential = None + + if cached_credential is None: + session = boto3.Session() + credentials = session.get_credentials() + if credentials: + credentials = credentials.get_frozen_credentials() + cached_credential = (credentials.access_key, credentials.secret_key) + return cached_credential if cached_credential else (None, None) + @imports.inject("boto3", pip_extra="aws") def infer_credentials(boto3, self): # todo load temporary credentials from STS diff --git a/skyplane/compute/aws/aws_server.py b/skyplane/compute/aws/aws_server.py index a2c1d2a22..14ec9082e 100644 --- a/skyplane/compute/aws/aws_server.py +++ b/skyplane/compute/aws/aws_server.py @@ -21,13 +21,14 @@ class AWSServer(Server): """AWS Server class to support basic SSH operations""" - def __init__(self, region_tag, instance_id, log_dir=None): + def __init__(self, region_tag, instance_id, key_path=None, log_dir=None): super().__init__(region_tag, log_dir=log_dir) assert self.region_tag.split(":")[0] == "aws" self.auth = AWSAuthentication() self.key_manager = AWSKeyManager(self.auth) self.aws_region = self.region_tag.split(":")[1] self.instance_id = instance_id + self.key_path = key_path @property @functools.lru_cache(maxsize=None) @@ -89,6 +90,9 @@ def instance_state(self): @property @ignore_lru_cache() def local_keyfile(self): + if self.key_path: + return self.key_path + key_name = self.get_boto3_instance_resource().key_name if self.key_manager.key_exists_local(key_name): return self.key_manager.get_key(key_name) diff --git a/skyplane/compute/server.py b/skyplane/compute/server.py index 585773105..fa21547fe 100644 --- a/skyplane/compute/server.py +++ b/skyplane/compute/server.py @@ -292,6 +292,7 @@ def start_gateway( use_compression=False, e2ee_key_bytes=None, use_socket_tls=False, + instance_path=None, ): def check_stderr(tup): assert tup[1].strip() == "", f"Command failed, err: {tup[1]}" @@ -354,6 +355,11 @@ def check_stderr(tup): docker_envs["GATEWAY_INFO_FILE"] = f"/pkg/data/gateway_info.json" docker_run_flags += f" -v /tmp/{gateway_program_file}:/pkg/data/gateway_program.json" docker_run_flags += f" -v /tmp/{gateway_info_file}:/pkg/data/gateway_info.json" + + # Instance path to mount if the source / destination is an instance + if instance_path is not None: + docker_run_flags += f" -v {instance_path}:{instance_path}" + gateway_daemon_cmd = f"/etc/init.d/stunnel4 start && python -u /pkg/skyplane/gateway/gateway_daemon.py --chunk-dir /skyplane/chunks" # update docker flags diff --git a/skyplane/gateway/gateway_daemon.py b/skyplane/gateway/gateway_daemon.py index 08e6562df..beded7312 100644 --- a/skyplane/gateway/gateway_daemon.py +++ b/skyplane/gateway/gateway_daemon.py @@ -20,6 +20,7 @@ GatewaySender, GatewayRandomDataGen, GatewayWriteLocal, + GatewayLocalReadOperator, GatewayObjStoreReadOperator, GatewayObjStoreWriteOperator, GatewayWaitReceiver, @@ -192,6 +193,19 @@ def create_gateway_operators_helper(input_queue, program: List[Dict], partition_ error_queue=self.error_queue, ) total_p += 1 + elif op["op_type"] == "read_local": + # TODO: add support for this + operators[handle] = GatewayLocalReadOperator( + handle=handle, + region=self.region, + input_queue=input_queue, + output_queue=output_queue, + error_queue=self.error_queue, + error_event=self.error_event, + chunk_store=self.chunk_store, + path=op["path"], + ) + total_p += 1 # ? elif op["op_type"] == "read_object_store": operators[handle] = GatewayObjStoreReadOperator( handle=handle, @@ -262,6 +276,7 @@ def create_gateway_operators_helper(input_queue, program: List[Dict], partition_ error_queue=self.error_queue, error_event=self.error_event, chunk_store=self.chunk_store, + path=op["path"], ) total_p += 1 else: diff --git a/skyplane/gateway/gateway_program.py b/skyplane/gateway/gateway_program.py index c27427fd9..661d2d008 100644 --- a/skyplane/gateway/gateway_program.py +++ b/skyplane/gateway/gateway_program.py @@ -72,6 +72,12 @@ def __init__(self, bucket_name: str, bucket_region: str, num_connections: int = self.num_connections = num_connections +class GatewayReadLocal(GatewayOperator): + def __init__(self, path: Optional[str] = None): + super().__init__("read_local") + self.path = path + + class GatewayWriteObjectStore(GatewayOperator): def __init__(self, bucket_name: str, bucket_region: str, num_connections: int = 32, key_prefix: Optional[str] = ""): super().__init__("write_object_store") diff --git a/skyplane/gateway/operators/gateway_operator.py b/skyplane/gateway/operators/gateway_operator.py index c574caebc..f55af7d6c 100644 --- a/skyplane/gateway/operators/gateway_operator.py +++ b/skyplane/gateway/operators/gateway_operator.py @@ -1,4 +1,5 @@ import json +import mmap from pathlib import Path import os from typing import List @@ -412,6 +413,7 @@ def process(self, chunk_req: ChunkRequest): class GatewayWriteLocal(GatewayOperator): def __init__( self, + path: str, handle: str, region: str, input_queue: GatewayQueue, @@ -422,9 +424,29 @@ def __init__( n_processes: int = 1, ): super().__init__(handle, region, input_queue, output_queue, error_event, error_queue, chunk_store, n_processes) + self.path = path def process(self, chunk_req: ChunkRequest): # do nothing (already written locally) + # TODO: sort and reassemble chunks + # Generate the chunk file path + fpath = str(self.chunk_store.get_chunk_file_path(chunk_req.chunk.chunk_id).absolute()) + + # If the final file does not exist, create it + if not os.path.exists(self.path): + open(self.path, "a").close() + + # Open the final file in read-write mode, seek to the correct position, and write the chunk data + with open(self.path, "r+b") as final_file: + # Read the chunk data + with open(fpath, "rb") as chunk_file: + chunk_data = chunk_file.read() + + # Seek to the chunk's offset in the final file and write the chunk data + offset_bytes = chunk_req.chunk.file_offset_bytes if chunk_req.chunk.file_offset_bytes is not None else 0 + final_file.seek(offset_bytes) + final_file.write(chunk_data) + return True @@ -463,6 +485,54 @@ def get_obj_store_interface(self, region: str, bucket: str) -> ObjectStoreInterf return self.obj_store_interfaces[key] +class GatewayLocalReadOperator(GatewayOperator): + def __init__( + self, + path: str, + handle: str, + region: str, + input_queue: GatewayQueue, + output_queue: GatewayQueue, + error_event, + error_queue: Queue, + n_processes: int = 32, + chunk_store: Optional[ChunkStore] = None, + ): + super().__init__( + handle, + region, + input_queue, + output_queue, + error_event, + error_queue, + chunk_store, + n_processes, + ) + self.path = path + + def process(self, chunk_req: ChunkRequest, **args): + # Determine the start and end position of the chunk in the file + chunk_start = chunk_req.chunk.file_offset_bytes if chunk_req.chunk.file_offset_bytes is not None else 0 + chunk_end = chunk_start + chunk_req.chunk.chunk_length_bytes + + # Read the specified part of the file and write the chunk + with open(self.path, "rb") as f: + mmapped_file = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) # memory-mapped file for efficient large file handling + + # Generate the chunk file path + fpath = str(self.chunk_store.get_chunk_file_path(chunk_req.chunk.chunk_id).absolute()) + + # Write the chunk to a new file + with open(fpath, "wb") as chunk_file: + chunk_file.write(mmapped_file[chunk_start:chunk_end]) + + logger.debug(f"[{self.handle}:{self.worker_id}] Local read: Wrote chunk {chunk_req.chunk.chunk_id} to {fpath}") + + mmapped_file.close() # Make sure to close the memory-mapped file + + return True + + class GatewayObjStoreReadOperator(GatewayObjStoreOperator): def __init__( self, @@ -568,6 +638,7 @@ def __init__( self.prefix = prefix def process(self, chunk_req: ChunkRequest): + print(f"Chunk request: {chunk_req}") fpath = str(self.chunk_store.get_chunk_file_path(chunk_req.chunk.chunk_id).absolute()) logger.debug( f"[{self.handle}:{self.worker_id}] Start upload {chunk_req.chunk.chunk_id} to {self.bucket_name}, key {chunk_req.chunk.dest_key}" diff --git a/skyplane/obj_store/storage_interface.py b/skyplane/obj_store/storage_interface.py index a430573f8..e36712660 100644 --- a/skyplane/obj_store/storage_interface.py +++ b/skyplane/obj_store/storage_interface.py @@ -35,40 +35,59 @@ def list_objects(self, prefix="") -> Iterator[Any]: raise NotImplementedError() @staticmethod - def create(region_tag: str, bucket: str): + def create(region_tag: str, transfer_loc: str): # TODO: modify this to also support local file if region_tag.startswith("aws"): from skyplane.obj_store.s3_interface import S3Interface - return S3Interface(bucket) + return S3Interface(transfer_loc) elif region_tag.startswith("gcp"): from skyplane.obj_store.gcs_interface import GCSInterface - return GCSInterface(bucket) + return GCSInterface(transfer_loc) elif region_tag.startswith("azure"): from skyplane.obj_store.azure_blob_interface import AzureBlobInterface - storage_account, container = bucket.split("/", 1) # / + storage_account, container = transfer_loc.split("/", 1) # / return AzureBlobInterface(storage_account, container) elif region_tag.startswith("ibmcloud"): from skyplane.obj_store.cos_interface import COSInterface - return COSInterface(bucket, region_tag) + return COSInterface(transfer_loc, region_tag) elif region_tag.startswith("hdfs"): from skyplane.obj_store.hdfs_interface import HDFSInterface - logger.fs.debug(f"attempting to create hdfs bucket {bucket}") - return HDFSInterface(host=bucket) + logger.fs.debug(f"attempting to create hdfs bucket {transfer_loc}") + return HDFSInterface(host=transfer_loc) elif region_tag.startswith("local"): # from skyplane.obj_store.file_system_interface import FileSystemInterface from skyplane.obj_store.posix_file_interface import POSIXInterface - return POSIXInterface(bucket) + return POSIXInterface(transfer_loc) + + elif region_tag.startswith("vm"): + from skyplane.obj_store.vm_interface import VMInterface + + # transfer_loc should be in format cloud_region@username@host:/path?private_key_path + cloud_region_user_host_path, private_key_path = transfer_loc.split("?") + cloud_region, host_path = cloud_region_user_host_path.split("@", 1) + username, host = host_path.split("@", 1) + host, path = host.split(":", 1) + parent_dir = "/".join(path.split("/")[:-1]) # Get the parent directory of the path + + return VMInterface( + host, + username, + cloud_region, + local_path=parent_dir, + private_key_path=private_key_path.removeprefix("private_key_path="), + ) + elif region_tag.startswith("cloudflare"): from skyplane.obj_store.r2_interface import R2Interface - account, bucket = bucket.split("/", 1) # / + account, bucket = transfer_loc.split("/", 1) # / return R2Interface(account, bucket) else: raise ValueError(f"Invalid region_tag {region_tag} - could not create interface") diff --git a/skyplane/obj_store/vm_interface.py b/skyplane/obj_store/vm_interface.py new file mode 100644 index 000000000..a7e051fc7 --- /dev/null +++ b/skyplane/obj_store/vm_interface.py @@ -0,0 +1,193 @@ +from dataclasses import dataclass +from datetime import datetime, timezone +import json +import mimetypes +import os +from typing import Any, Iterator, List, Optional +import uuid +from dateutil.parser import parse +import paramiko +import pytz +from skyplane.obj_store.object_store_interface import ( + ObjectStoreInterface, + ObjectStoreObject, +) +from skyplane import exceptions + + +@dataclass +class VMFile(ObjectStoreObject): + def full_path(self): + if self.key.startswith("/"): + return f"vm://{self.bucket}{self.key}" + else: + return f"vm://{self.bucket}/{self.key}" + + +class VMInterface(ObjectStoreInterface): + def __init__( + self, + host, + username, + region, + private_key_path, + local_path="/", + ssh_key_password="skyplane", + ): + self.host = host + self.username = username + self.region = region + self.private_key_path = private_key_path + self.local_path = local_path + self.temp_dir = "/tmp/multipart_uploads/" # directory on the VMs + + # Set up SSH + ssh_client = paramiko.SSHClient() + ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + try: + ssh_client.connect( + hostname=host, + username=username, + pkey=paramiko.RSAKey.from_private_key_file(str(private_key_path), password=ssh_key_password), + look_for_keys=False, + banner_timeout=200, + ) + self.client = ssh_client + except paramiko.AuthenticationException as e: + raise exceptions.BadConfigException(f"Failed to connect to Server") from e + + # TODO: check if this works for all CPs + if self.region.startswith("aws"): + _, stdout, _ = self.client.exec_command("curl http://169.254.169.254/latest/meta-data/instance-id") + self.instance_id = stdout.read().decode("utf-8").strip() + elif self.region.startswith("gcp"): + _, stdout, _ = self.client.exec_command( + 'curl "http://metadata.google.internal/computeMetadata/v1/instance/name" -H "Metadata-Flavor: Google"' + ) + self.instance_id = stdout.read().decode("utf-8").strip() + elif self.region.startswith("azure"): + _, stdout, _ = self.client.exec_command( + 'curl -H Metadata:true "http://169.254.169.254/metadata/instance/compute?api-version=2017-08-01"' + ) + metadata = json.loads(stdout.read().decode("utf-8").strip()) + self.instance_id = metadata["name"] + else: + raise exceptions.BadConfigException(f"Invalid region tag: {self.region}") + + @property + def provider(self) -> str: + return "vm" + + def region_tag(self) -> str: + return self.region + + def id(self) -> str: + return self.instance_id + + def path(self) -> str: + return self.local_path + + def key_path(self) -> str: + return str(self.private_key_path) + + def bucket(self) -> str: + return f"{self.region}@{self.username}@{self.host}:{self.local_path}?private_key_path={self.private_key_path}" + + def host_ip(self) -> str: + return self.host + + def list_objects(self, prefix="") -> Iterator[VMFile]: + # List files in directory, recursively + _, stdout, _ = self.client.exec_command(f"find {prefix} -type f") + files = stdout.readlines() + for file_path in files: + file_path = file_path.strip() + _, stdout, _ = self.client.exec_command(f"ls -l --time-style=full-iso {file_path}") + file_info = stdout.readline().split() + file_size = file_info[4] + # Get the last modified time from the ls output + file_datetime_str = " ".join(file_info[5:8]) + datetime_str, _ = file_datetime_str.rsplit(" ", 1) + timestamp, nanosec = datetime_str.split(".") + timestamp_str = f"{timestamp}.{nanosec[:6]}" + + dt_naive = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S.%f") + dt_aware = dt_naive.replace(tzinfo=timezone.utc) + + yield VMFile( + provider="vm", + bucket=self.host, + key=file_path, + size=int(file_size), + last_modified=dt_aware, + ) + + def exists(self, obj_name: str): + _, stdout, _ = self.client.exec_command(f"ls {self.path}/{obj_name}") + return stdout.readline() != "" + + def create_bucket(self, region_tag: str): + return None + + def delete_bucket(self): + return None + + def bucket_exists(self) -> bool: + """We always have a bucket""" + return True + + def download_object(self, src_object_name, dst_file_path): + sftp = self.client.open_sftp() + sftp.get(f"{self.path}/{src_object_name}", dst_file_path) + + def upload_object(self, src_file_path, dst_object_name, part_number=None, upload_id=None): + sftp = self.client.open_sftp() + if part_number and upload_id: + remote_part_path = f"{self.temp_dir}/{upload_id}/{part_number}" + sftp.put(src_file_path, remote_part_path) + else: + sftp.put(src_file_path, f"{self.path}/{dst_object_name}") + + def delete_objects(self, keys: List[str]): + for key in keys: + self.client.exec_command(f"rm {self.path}/{key}") + + def get_obj_size(self, obj_name) -> int: + _, stdout, _ = self.client.exec_command(f"ls -l {self.path}/{obj_name}") + return int(stdout.readline().split()[4]) + + def get_obj_last_modified(self, obj_name): + _, stdout, _ = self.client.exec_command(f"ls -l --time-style=full-iso {self.path}/{obj_name}") + file_info = stdout.readline().split() + file_datetime_str = " ".join(file_info[5:8]) + + # parse datetime string and convert it to UTC + dt_aware = parse(file_datetime_str) + dt_utc = dt_aware.astimezone(pytz.UTC) + + return dt_utc + + def get_obj_mime_type(self, obj_name): + return mimetypes.guess_type(obj_name)[0] + + def initiate_multipart_upload(self, dst_object_name: str, mime_type: Optional[str] = None) -> str: + upload_id = str(uuid.uuid4()) + _, stderr, _ = self.client.exec_command(f"mkdir -p {self.temp_dir}/{upload_id}") + error_message = stderr.read().decode().strip() + if error_message: + raise exceptions.BadConfigException(f"Failed to create directory on VM: {error_message}") + return upload_id + + def complete_multipart_upload(self, dst_object_name, upload_id, metadata: Optional[Any] = None): + _, stdout, _ = self.client.exec_command(f"ls {self.temp_dir}/{upload_id}") + parts = [f"{self.temp_dir}/{upload_id}/{part}" for part in sorted(stdout.read().decode().split(), key=int)] + + # Concatenate all parts together + concatenated_parts = " ".join(parts) + _, stderr, _ = self.client.exec_command(f"cat {concatenated_parts} > {dst_object_name}") + error_message = stderr.read().decode().strip() + if error_message: + raise exceptions.BadConfigException(f"Failed to complete multipart upload on VM: {error_message}") + + # Cleanup + # _, _, _ = self.client.exec_command(f"rm -r {self.temp_dir}/{upload_id}") diff --git a/skyplane/planner/planner.py b/skyplane/planner/planner.py index 2bb21e0ac..f406ce4fb 100644 --- a/skyplane/planner/planner.py +++ b/skyplane/planner/planner.py @@ -14,7 +14,9 @@ GatewayMuxOr, GatewayMuxAnd, GatewayReadObjectStore, + GatewayReadLocal, GatewayWriteObjectStore, + GatewayWriteLocal, GatewayReceive, GatewaySend, ) @@ -24,6 +26,7 @@ from skyplane.utils.fn import do_parallel from skyplane.config_paths import config_path, azure_standardDv5_quota_path, aws_quota_path, gcp_quota_path + from skyplane.config import SkyplaneConfig @@ -376,6 +379,314 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan: return plan +class DirectPlannerVMSource(MulticastDirectPlanner): + def plan(self, jobs: List[TransferJob]) -> TopologyPlan: + src_region_tag = jobs[0].src_iface.region_tag() + dst_region_tags = [iface.region_tag() for iface in jobs[0].dst_ifaces] + # jobs must have same sources and destinations + for job in jobs[1:]: + assert job.src_iface.region_tag() == src_region_tag, "All jobs must have same source region" + assert [iface.region_tag() for iface in job.dst_ifaces] == dst_region_tags, "Add jobs must have same destination set" + + plan = TopologyPlan(src_region_tag=src_region_tag, dest_region_tags=dst_region_tags) + + # Dynammically calculate n_instances based on quota limits + vm_types, n_instances = self._get_vm_type_and_instances(src_region_tag=src_region_tag, dst_region_tags=dst_region_tags) + + for i in range(n_instances): + plan.add_gateway( + src_region_tag, + vm_types[src_region_tag] if vm_types else None, + instance_id=jobs[0].src_iface.id(), + instance_path=jobs[0].src_iface.path(), + ) + for dst_region_tag in dst_region_tags: + plan.add_gateway(dst_region_tag, vm_types[dst_region_tag] if vm_types else None) + + # initialize gateway programs per region + dst_program = {dst_region: GatewayProgram() for dst_region in dst_region_tags} + src_program = GatewayProgram() + + # iterate through all jobs + for job in jobs: + 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 = job.uuid + + # source region gateway program + obj_store_read = src_program.add_operator(GatewayReadLocal(job.src_prefix), partition_id=partition_id) + # send to all destination + mux_and = src_program.add_operator(GatewayMuxAnd(), parent_handle=obj_store_read, partition_id=partition_id) + dst_prefixes = job.dst_prefixes + for i in range(len(job.dst_ifaces)): + dst_iface = job.dst_ifaces[i] + dst_prefix = dst_prefixes[i] + dst_region_tag = dst_iface.region_tag() + dst_bucket = dst_iface.bucket() + dst_gateways = plan.get_region_gateways(dst_region_tag) + + # special case where destination is same region as source + if dst_region_tag == src_region_tag: + src_program.add_operator( + GatewayWriteObjectStore( + dst_bucket, + dst_region_tag, + self.n_connections, + key_prefix=dst_prefix, + ), + parent_handle=mux_and, + partition_id=partition_id, + ) + continue + + # 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(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=int(self.n_connections / len(dst_gateways)), + private_ip=private_ip, + ), + parent_handle=mux_or, + partition_id=partition_id, + ) + + # each gateway also recieves data from source + recv_op = dst_program[dst_region_tag].add_operator(GatewayReceive(), partition_id=partition_id) + dst_program[dst_region_tag].add_operator( + GatewayWriteObjectStore( + dst_bucket, + dst_region_tag, + self.n_connections, + key_prefix=dst_prefix, + ), + parent_handle=recv_op, + partition_id=partition_id, + ) + + # update cost per GB + plan.cost_per_gb += compute.CloudProvider.get_transfer_cost(src_region_tag, dst_region_tag) + + # set gateway programs + plan.set_gateway_program(src_region_tag, src_program) + for dst_region_tag, program in dst_program.items(): + if dst_region_tag != src_region_tag: # don't overwrite + plan.set_gateway_program(dst_region_tag, program) + + return plan + + +class DirectPlannerVMDest(MulticastDirectPlanner): + def plan(self, jobs: List[TransferJob]) -> TopologyPlan: + src_region_tag = jobs[0].src_iface.region_tag() + dst_region_tags = [iface.region_tag() for iface in jobs[0].dst_ifaces] + # jobs must have same sources and destinations + for job in jobs[1:]: + assert job.src_iface.region_tag() == src_region_tag, "All jobs must have same source region" + assert [iface.region_tag() for iface in job.dst_ifaces] == dst_region_tags, "Add jobs must have same destination set" + + plan = TopologyPlan(src_region_tag=src_region_tag, dest_region_tags=dst_region_tags) + + # Dynammically calculate n_instances based on quota limits + vm_types, n_instances = self._get_vm_type_and_instances(src_region_tag=src_region_tag, dst_region_tags=dst_region_tags) + + for i in range(n_instances): + plan.add_gateway(src_region_tag, vm_types[src_region_tag] if vm_types else None) + for iface in jobs[0].dst_ifaces: + dst_region_tag = iface.region_tag() + dst_vm_instance_id = iface.id() + dst_vm_instance_path = iface.path() + dst_vm_key_path = iface.key_path() + plan.add_gateway( + dst_region_tag, + vm_types[dst_region_tag] if vm_types else None, + instance_id=dst_vm_instance_id, + instance_path=dst_vm_instance_path, + instance_key_path=dst_vm_key_path, + ) + + # initialize gateway programs per region + dst_program = {dst_region: GatewayProgram() for dst_region in dst_region_tags} + src_program = GatewayProgram() + + # iterate through all jobs + 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 = job.uuid + + # source region gateway program + obj_store_read = src_program.add_operator( + GatewayReadObjectStore(src_bucket, src_region_tag, self.n_connections), + partition_id=partition_id, + ) + + # send to all destination + mux_and = src_program.add_operator(GatewayMuxAnd(), parent_handle=obj_store_read, partition_id=partition_id) + dst_prefixes = job.dst_prefixes + for i in range(len(job.dst_ifaces)): + dst_iface = job.dst_ifaces[i] + dst_region_tag = dst_iface.region_tag() + dst_prefix = dst_prefixes[i] + dst_gateways = plan.get_region_gateways(dst_region_tag) + + # special case where destination is same region as source + if dst_region_tag == src_region_tag: + src_program.add_operator( + GatewayWriteLocal(dst_prefix), + parent_handle=mux_and, + partition_id=partition_id, + ) + + # 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(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=int(self.n_connections / len(dst_gateways)), + private_ip=private_ip, + ), + parent_handle=mux_or, + partition_id=partition_id, + ) + + # each gateway also recieves data from source + recv_op = dst_program[dst_region_tag].add_operator(GatewayReceive(), partition_id=partition_id) + dst_program[dst_region_tag].add_operator( + GatewayWriteLocal(dst_prefix), + parent_handle=recv_op, + partition_id=partition_id, + ) + + # update cost per GB + plan.cost_per_gb += compute.CloudProvider.get_transfer_cost(src_region_tag, dst_region_tag) + + # set gateway programs + plan.set_gateway_program(src_region_tag, src_program) + for dst_region_tag, program in dst_program.items(): + if dst_region_tag != src_region_tag: # don't overwrite + plan.set_gateway_program(dst_region_tag, program) + return plan + + +class DirectPlannerVMSourceDest(MulticastDirectPlanner): + def plan(self, jobs: List[TransferJob]) -> TopologyPlan: + src_region_tag = jobs[0].src_iface.region_tag() + dst_region_tags = [iface.region_tag() for iface in jobs[0].dst_ifaces] + # jobs must have same sources and destinations + for job in jobs[1:]: + assert job.src_iface.region_tag() == src_region_tag, "All jobs must have same source region" + assert [iface.region_tag() for iface in job.dst_ifaces] == dst_region_tags, "Add jobs must have same destination set" + + plan = TopologyPlan(src_region_tag=src_region_tag, dest_region_tags=dst_region_tags) + + # Dynammically calculate n_instances based on quota limits + vm_types, n_instances = self._get_vm_type_and_instances(src_region_tag=src_region_tag, dst_region_tags=dst_region_tags) + for i in range(n_instances): + plan.add_gateway( + src_region_tag, + vm_types[src_region_tag] if vm_types else None, + instance_id=jobs[0].src_iface.id(), + instance_path=jobs[0].src_iface.path(), + ) + for iface in jobs[0].dst_ifaces: + dst_region_tag = iface.region_tag() + dst_vm_instance_id = iface.id() + dst_vm_instance_path = iface.path() + plan.add_gateway( + dst_region_tag, + vm_types[dst_region_tag] if vm_types else None, + instance_id=dst_vm_instance_id, + instance_path=dst_vm_instance_path, + ) + + # initialize gateway programs per region + dst_program = {dst_region: GatewayProgram() for dst_region in dst_region_tags} + src_program = GatewayProgram() + + # iterate through all jobs + for job in jobs: + 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 = job.uuid + + # source region gateway program + obj_store_read = src_program.add_operator(GatewayReadLocal(job.src_prefix), partition_id=partition_id) + + # send to all destination + mux_and = src_program.add_operator(GatewayMuxAnd(), parent_handle=obj_store_read, partition_id=partition_id) + dst_prefixes = job.dst_prefixes + for i in range(len(job.dst_ifaces)): + dst_iface = job.dst_ifaces[i] + dst_region_tag = dst_iface.region_tag() + dst_prefix = dst_prefixes[i] + dst_gateways = plan.get_region_gateways(dst_region_tag) + + # special case where destination is same region as source + if dst_region_tag == src_region_tag: + src_program.add_operator( + GatewayWriteLocal(dst_prefix), + parent_handle=mux_and, + partition_id=partition_id, + ) + continue + + # 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(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=int(self.n_connections / len(dst_gateways)), + private_ip=private_ip, + ), + parent_handle=mux_or, + partition_id=partition_id, + ) + + # each gateway also recieves data from source + recv_op = dst_program[dst_region_tag].add_operator(GatewayReceive(), partition_id=partition_id) + dst_program[dst_region_tag].add_operator( + GatewayWriteLocal(dst_prefix), + parent_handle=recv_op, + partition_id=partition_id, + ) + + # update cost per GB + plan.cost_per_gb += compute.CloudProvider.get_transfer_cost(src_region_tag, dst_region_tag) + + # set gateway programs + plan.set_gateway_program(src_region_tag, src_program) + for dst_region_tag, program in dst_program.items(): + if dst_region_tag != src_region_tag: # don't overwrite + plan.set_gateway_program(dst_region_tag, program) + return plan + + class DirectPlannerSourceOneSided(MulticastDirectPlanner): """Planner that only creates VMs in the source region""" @@ -410,7 +721,8 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan: # source region gateway program obj_store_read = src_program.add_operator( - GatewayReadObjectStore(src_bucket, src_region_tag, self.n_connections), partition_id=partition_id + GatewayReadObjectStore(src_bucket, src_region_tag, self.n_connections), + partition_id=partition_id, ) # send to all destination mux_and = src_program.add_operator(GatewayMuxAnd(), parent_handle=obj_store_read, partition_id=partition_id) @@ -424,7 +736,12 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan: # special case where destination is same region as source src_program.add_operator( - GatewayWriteObjectStore(dst_bucket, dst_region_tag, self.n_connections, key_prefix=dst_prefix), + GatewayWriteObjectStore( + dst_bucket, + dst_region_tag, + self.n_connections, + key_prefix=dst_prefix, + ), parent_handle=mux_and, partition_id=partition_id, ) @@ -480,11 +797,17 @@ def plan(self, jobs: List[TransferJob]) -> TopologyPlan: # source region gateway program obj_store_read = dst_program[dst_region_tag].add_operator( - GatewayReadObjectStore(src_bucket, src_region_tag, self.n_connections), partition_id=partition_id + GatewayReadObjectStore(src_bucket, src_region_tag, self.n_connections), + partition_id=partition_id, ) dst_program[dst_region_tag].add_operator( - GatewayWriteObjectStore(dst_bucket, dst_region_tag, self.n_connections, key_prefix=dst_prefix), + GatewayWriteObjectStore( + dst_bucket, + dst_region_tag, + self.n_connections, + key_prefix=dst_prefix, + ), parent_handle=obj_store_read, partition_id=partition_id, ) diff --git a/skyplane/planner/topology.py b/skyplane/planner/topology.py index fd5a1898f..2ba50f0b7 100644 --- a/skyplane/planner/topology.py +++ b/skyplane/planner/topology.py @@ -5,6 +5,7 @@ GatewayWriteObjectStore, GatewayGenData, GatewayReadObjectStore, + GatewayReadLocal, ) from typing import List, Dict, Optional @@ -15,12 +16,25 @@ class TopologyPlanGateway: Represents a gateway in the topology plan. """ - def __init__(self, region_tag: str, gateway_id: str, gateway_vm: Optional[str]): + def __init__( + self, + region_tag: str, + gateway_id: str, + gateway_vm: Optional[str], + gateway_instance_id: Optional[str] = None, + gateway_instance_path: Optional[str] = None, + gateway_key_path: Optional[str] = None, + ): self.region_tag = region_tag self.gateway_id = gateway_id self.gateway_vm = gateway_vm self.gateway_program = None + # TODO: hard code instance id and path for now for initializing Server + self.gateway_instance_id = gateway_instance_id + self.gateway_instance_path = gateway_instance_path + self.gateway_key_path = gateway_key_path + # ip addresses self.private_ip_address = None self.public_ip_address = None @@ -79,11 +93,18 @@ def region_tags(self) -> List[str]: """Get all region tags in the topology plan""" return list(set([gateway.region_tag for gateway in self.gateways.values()])) - def add_gateway(self, region_tag: str, vm_type: Optional[str] = None): + def add_gateway( + self, + region_tag: str, + vm_type: Optional[str] = None, + instance_id: Optional[str] = None, + instance_path: Optional[str] = None, + instance_key_path: Optional[str] = None, + ): """Create gateway in specified region""" 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, vm_type) + gateway = TopologyPlanGateway(region_tag, gateway_id, vm_type, instance_id, instance_path, instance_key_path) self.gateways[gateway_id] = gateway return gateway @@ -168,7 +189,11 @@ def source_instances(self): nodes = [] for gateway in self.gateways.values(): for operator in gateway.gateway_program.get_operators(): - if isinstance(operator, GatewayReadObjectStore) or isinstance(operator, GatewayGenData): + if ( + isinstance(operator, GatewayReadObjectStore) + or isinstance(operator, GatewayGenData) + or isinstance(operator, GatewayReadLocal) + ): nodes.append(gateway) break diff --git a/skyplane/utils/path.py b/skyplane/utils/path.py index 9670934f8..8c737bf2c 100644 --- a/skyplane/utils/path.py +++ b/skyplane/utils/path.py @@ -49,6 +49,18 @@ def is_plausible_local_path(path_test: str): raise ValueError(f"Invalid Azure path: {path}") account, container, blob_path = match.groups() return "azure", f"{account}/{container}", blob_path + elif path.startswith("vm://"): + # VM URL with private key path + regex = re.compile(r"vm://([^@]+)@([^@]+)@([^:/]+):([^?]+)\?private_key_path=(.*)") + match = regex.match(path) + if match is None: + raise ValueError(f"Invalid VM path: {path}") + cloud_region, username, host, path, private_key_path = match.groups() + return ( + "vm", + f"{cloud_region}@{username}@{host}:{path}?private_key_path={private_key_path}", + path, + ) elif path.startswith("azure://"): regex = re.compile(r"azure://([^/]+)/([^/]+)/?(.*)") match = regex.match(path) diff --git a/tests/interface_util.py b/tests/interface_util.py index f6d1f85d9..46c015c0b 100644 --- a/tests/interface_util.py +++ b/tests/interface_util.py @@ -11,6 +11,10 @@ def interface_test_framework(region, bucket, multipart: bool, test_delete_bucket: bool = False, file_size_mb: int = 1): interface = ObjectStoreInterface.create(region, bucket) + return interface_test_from_iface(interface, multipart=multipart, test_delete_bucket=test_delete_bucket, file_size_mb=file_size_mb) + + +def interface_test_from_iface(interface, multipart: bool, test_delete_bucket: bool = False, file_size_mb: int = 1): interface.create_bucket(region.split(":")[1]) time.sleep(5) diff --git a/tests/unit_vm/test_vm_interface.py b/tests/unit_vm/test_vm_interface.py new file mode 100644 index 000000000..1f9c26f95 --- /dev/null +++ b/tests/unit_vm/test_vm_interface.py @@ -0,0 +1,28 @@ +import uuid +from skyplane.compute.gcp.gcp_cloud_provider import GCPCloudProvider +from skyplane.obj_store.object_store_interface import ObjectStoreInterface +from tests.interface_util import interface_test_framework +from skyplane.utils import logger + + +def provision_vm(): + vm = GCPCloudProvider().provision_instance("us-east1", "n2-standard-2") + return vm + + +def test_vm_simple(): + # provision vm + vm = provision_vm() + + # create iface + region = vm.region_tag + vm_host = "skyplane" + vm_private_key_path = vm.ssh_private_key + + vm_iface = VMInterface(vm_host, vm.gcp_instance_name, vm_region, vm_private_key_path) + + # test a provisioned vm exists + assert vm_iface.exists() + + # test basic transfer + assert interface_test_from_iface(vm_iface)