From b101687835981c165bb417f47c9ec52b39a56594 Mon Sep 17 00:00:00 2001 From: huangw Date: Sat, 11 May 2024 10:31:27 +0800 Subject: [PATCH] [feature] add reuse gateway --- skyplane/api/dataplane.py | 57 +++++++++++++++++++++++++++--------- skyplane/api/provisioner.py | 7 ++++- skyplane/cli/cli_transfer.py | 15 +++++++--- 3 files changed, 60 insertions(+), 19 deletions(-) diff --git a/skyplane/api/dataplane.py b/skyplane/api/dataplane.py index 69916093..e4cc4ec8 100644 --- a/skyplane/api/dataplane.py +++ b/skyplane/api/dataplane.py @@ -28,13 +28,15 @@ class DataplaneAutoDeprovision: - def __init__(self, dataplane: "Dataplane"): + def __init__(self, dataplane: "Dataplane", reuse_instances: bool = False): self.dataplane = dataplane + self.reuse_instances = reuse_instances def __enter__(self): return self.dataplane def __exit__(self, exc_type, exc_value, exc_tb): + if self.reuse_instances: return logger.fs.warning("Deprovisioning dataplane") self.dataplane.deprovision() @@ -129,6 +131,7 @@ def _start_gateway( def provision( self, allow_firewall: bool = True, + reuse_instances: bool = False, gateway_docker_image: str = os.environ.get("SKYPLANE_DOCKER_IMAGE", gateway_docker_image()), authorize_ssh_pub_key: Optional[str] = None, max_jobs: int = 16, @@ -160,28 +163,54 @@ def provision( is_scp_used = any(n.region_tag.startswith("scp:") for n in self.topology.get_gateways()) # create VMs from the topology + uuids = [] + reuse_nodes_tmp = {} + if reuse_instances: + for node in self.topology.get_gateways(): + cloud_provider, region = node.region_tag.split(":") + if reuse_nodes_tmp.get(f"{cloud_provider} {region}", []): continue + if cloud_provider == "gcp": + reuse_nodes_tmp[f"{cloud_provider} {region}"] = self.provisioner.gcp.get_instance_list(region=region) + if cloud_provider == "azure": + reuse_nodes_tmp[f"{cloud_provider} {region}"] = self.provisioner.azure.get_instance_list(region=region) + if cloud_provider == "aws": + reuse_nodes_tmp[f"{cloud_provider} {region}"] = self.provisioner.aws.get_instance_list(region=region) + if cloud_provider == "ibmcloud": + reuse_nodes_tmp[f"{cloud_provider} {region}"] = self.provisioner.ibmcloud.get_instance_list(region=region) + if cloud_provider == "scp": + reuse_nodes_tmp[f"{cloud_provider} {region}"] = self.provisioner.scp.get_instance_list(region=region) for node in self.topology.get_gateways(): cloud_provider, region = node.region_tag.split(":") - 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, - ) + if reuse_instances: + reuse_nodes = reuse_nodes_tmp.get(f"{cloud_provider} {region}", []) + + if reuse_instances and reuse_nodes: + reuse_server = reuse_nodes[0] + reuse_nodes.remove(reuse_server) + uid = self.provisioner.add_node(reuse_server) + uuids.append(uid) + else: + 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, + ) # initialize clouds self.provisioner.init_global(aws=is_aws_used, azure=is_azure_used, gcp=is_gcp_used, ibmcloud=is_ibmcloud_used, scp=is_scp_used) # provision VMs - uuids = self.provisioner.provision( + uuids_tmp = self.provisioner.provision( authorize_firewall=allow_firewall, max_jobs=max_jobs, spinner=spinner, ) + uuids.extend(uuids_tmp) # bind VMs to nodes servers = [self.provisioner.get_node(u) for u in uuids] @@ -291,9 +320,9 @@ def get_error_logs(args): errors[instance] = result return errors - def auto_deprovision(self) -> DataplaneAutoDeprovision: + def auto_deprovision(self, reuse_instances: bool = False) -> DataplaneAutoDeprovision: """Returns a context manager that will automatically call deprovision upon exit.""" - return DataplaneAutoDeprovision(self) + return DataplaneAutoDeprovision(self, reuse_instances) def source_gateways(self) -> List[compute.Server]: """Returns a list of source gateway nodes""" diff --git a/skyplane/api/provisioner.py b/skyplane/api/provisioner.py index 7e098174..0350a64e 100644 --- a/skyplane/api/provisioner.py +++ b/skyplane/api/provisioner.py @@ -161,7 +161,12 @@ def get_node(self, uuid: str) -> compute.Server: :type uuid: str """ return self.provisioned_vms[uuid] - + + def add_node(self, server) -> str: + uid = uuid.uuid4().int + self.provisioned_vms[uid] = server + return uid + def _provision_task(self, task: ProvisionerTask): with Timer() as t: if task.cloud_provider == "aws": diff --git a/skyplane/cli/cli_transfer.py b/skyplane/cli/cli_transfer.py index ebb38592..d914d83f 100644 --- a/skyplane/cli/cli_transfer.py +++ b/skyplane/cli/cli_transfer.py @@ -310,6 +310,7 @@ def run_transfer( src: str, dst: str, recursive: bool, + reuse_gateways: bool, debug: bool, multipart: bool, confirm: bool, @@ -392,15 +393,17 @@ def run_transfer( # dataplane must be created after transfers are queued dp = pipeline.create_dataplane(debug=debug) - with dp.auto_deprovision(): + with dp.auto_deprovision(reuse_instances=reuse_gateways): try: - dp.provision(spinner=True) + dp.provision(reuse_instances=reuse_gateways, spinner=True) dp.run(pipeline.jobs_to_dispatch, hooks=ProgressBarTransferHook(dp.topology.dest_region_tags)) + if reuse_gateways: return 0 except KeyboardInterrupt: logger.fs.warning("Transfer cancelled by user (KeyboardInterrupt).") console.print("\n[red]Transfer cancelled by user. Copying gateway logs and exiting.[/red]") try: dp.copy_gateway_logs() + if reuse_gateways: return 0 force_deprovision(dp) except Exception as e: logger.fs.exception(e) @@ -414,12 +417,14 @@ def run_transfer( console.print(f"[bright_black]{traceback.format_exc()}[/bright_black]") console.print(e.pretty_print_str()) UsageClient.log_exception("cli_query_objstore", e, args, cli.src_region_tag, [cli.dst_region_tag]) + if reuse_gateways: return 0 force_deprovision(dp) except Exception as e: logger.fs.exception(e) console.print(f"[bright_black]{traceback.format_exc()}[/bright_black]") console.print(e) UsageClient.log_exception("cli_query_objstore", e, args, cli.src_region_tag, [cli.dst_region_tag]) + if reuse_gateways: return 0 force_deprovision(dp) @@ -427,6 +432,7 @@ def cp( src: str, dst: str, recursive: bool = typer.Option(False, "--recursive", "-r", help="If true, will copy objects at folder prefix recursively"), + reuse_gateways: bool = typer.Option(False, help="If true, will leave provisioned instances running to be reused"), debug: bool = typer.Option(False, help="If true, will write debug information to debug directory."), multipart: bool = typer.Option(cloud_config.get_flag("multipart_enabled"), help="If true, will use multipart uploads."), # transfer flags @@ -468,12 +474,13 @@ def cp( :param solver: The solver to use for the transfer (default: direct) :type solver: str """ - return run_transfer(src, dst, recursive, debug, multipart, confirm, max_instances, max_connections, solver, "cp") + return run_transfer(src, dst, recursive, reuse_gateways, debug, multipart, confirm, max_instances, max_connections, solver, "cp") def sync( src: str, dst: str, + reuse_gateways: bool = typer.Option(False, help="If true, will leave provisioned instances running to be reused"), debug: bool = typer.Option(False, help="If true, will write debug information to debug directory."), multipart: bool = typer.Option(cloud_config.get_flag("multipart_enabled"), help="If true, will use multipart uploads."), # transfer flags @@ -517,4 +524,4 @@ def sync( :param solver: The solver to use for the transfer (default: direct) :type solver: str """ - return run_transfer(src, dst, False, debug, multipart, confirm, max_instances, max_connections, solver, "sync") + return run_transfer(src, dst, False, debug, reuse_gateways, multipart, confirm, max_instances, max_connections, solver, "sync")