skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/skylet/providers/oci/node_provider.py

@@ -1,488 +0,0 @@
-"""OCI Node Provider.
-
-Node provider is called by the Ray Autoscaler to provision new compute
-resources (head / worker nodes).
-
-To show debug messages, export SKYPILOT_DEBUG=1 
-
-History:
- - Hysun He (hysun.he@oracle.com) @ Apr, 2023: Initial implementation
- 
-"""
-
-import copy
-from datetime import datetime
-import logging
-import threading
-import time
-
-from ray.autoscaler.node_provider import NodeProvider
-from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME
-from ray.autoscaler.tags import TAG_RAY_LAUNCH_CONFIG
-from ray.autoscaler.tags import TAG_RAY_NODE_KIND
-from ray.autoscaler.tags import TAG_RAY_USER_NODE_TYPE
-
-from sky.adaptors import oci as oci_adaptor
-from sky.clouds.utils import oci_utils
-from sky.skylet.providers.oci import utils
-from sky.skylet.providers.oci.query_helper import oci_query_helper
-
-logger = logging.getLogger(__name__)
-
-
-def synchronized(f):
-
-    def wrapper(self, *args, **kwargs):
-        self.lock.acquire()
-        try:
-            return f(self, *args, **kwargs)
-        finally:
-            self.lock.release()
-
-    return wrapper
-
-
-class OCINodeProvider(NodeProvider):
-    """Node Provider for OracleCloud (OCI)."""
-
-    def __init__(self, provider_config, cluster_name):
-        NodeProvider.__init__(self, provider_config, cluster_name)
-        self.lock = threading.RLock()
-        self.cached_nodes = {}
-        self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes",
-                                                       True)
-        self.region = provider_config["region"]
-
-        # Do a read-ahead cache loading to improve performance.
-        self._get_filtered_nodes({})
-
-    @synchronized
-    def _get_filtered_nodes(self, tag_filters, force=False):
-        # Make sure the cluster_name is always an criterion
-        tag_filters = {**tag_filters, TAG_RAY_CLUSTER_NAME: self.cluster_name}
-
-        return_nodes = {}
-        if not force:
-            # Query cache first to reduce API call.
-            cache_hit = False
-            for k, node in self.cached_nodes.items():
-                tags = node["tags"]
-                unmatched_tags = [
-                    k for k, v in tag_filters.items()
-                    if k not in tags or v != tags[k]
-                ]
-                if len(unmatched_tags) == 0:
-                    return_nodes[k] = node
-                    cache_hit |= True
-
-            if cache_hit:
-                return return_nodes
-
-        insts = oci_query_helper.query_instances_by_tags(
-            tag_filters, self.region)
-        for inst in insts:
-            inst_id = inst.identifier
-            if inst_id in self.cached_nodes:
-                del self.cached_nodes[inst_id]
-
-            item = self.get_inst_obj({
-                "inst_id": inst_id,
-                "ad": inst.availability_domain,
-                "compartment": inst.compartment_id,
-                "lifecycle_state": inst.lifecycle_state,
-                "oci_tags": inst.freeform_tags,
-            })
-            return_nodes[inst_id] = item
-            self.cached_nodes[inst_id] = item
-
-        return return_nodes
-
-    @utils.debug_enabled(logger=logger)
-    def non_terminated_nodes(self, tag_filters):
-        """Return a list of node ids filtered by the specified tags dict.
-
-        This list must not include terminated nodes. For performance reasons,
-        providers are allowed to cache the result of a call to
-        non_terminated_nodes() to serve single-node queries
-        (e.g. is_running(node_id)). This means that non_terminated_nodes()
-        must be called again to refresh results.
-        """
-        VALIDITY_TAGS = [
-            TAG_RAY_CLUSTER_NAME,
-            TAG_RAY_NODE_KIND,
-            TAG_RAY_USER_NODE_TYPE,
-            TAG_RAY_LAUNCH_CONFIG,
-        ]
-        filters = {
-            tag: tag_filters[tag] for tag in VALIDITY_TAGS if tag in tag_filters
-        }
-
-        nodes = self._get_filtered_nodes(tag_filters=filters)
-        return [k for k, v in nodes.items() if v["status"] == "RUNNING"]
-
-    @utils.debug_enabled(logger=logger)
-    def is_running(self, node_id):
-        """Return whether the specified node is running."""
-        node = self._get_cached_node(node_id=node_id)
-        check_result = node is None or node["status"] == "RUNNING"
-
-        return check_result
-
-    @utils.debug_enabled(logger=logger)
-    def is_terminated(self, node_id):
-        """Return whether the specified node is terminated."""
-        node = self._get_cached_node(node_id=node_id)
-        check_result = ((node is None) or (node["status"] == "TERMINATED") or
-                        (node["status"] == "TERMINATING"))
-
-        return check_result
-
-    @utils.debug_enabled(logger=logger)
-    def node_tags(self, node_id):
-        return self.cached_nodes[node_id]["tags"]
-
-    @utils.debug_enabled(logger=logger)
-    def external_ip(self, node_id):
-        """Returns the external ip of the given node."""
-        return self._get_cached_node(node_id=node_id)["external_ip"]
-
-    @utils.debug_enabled(logger=logger)
-    def internal_ip(self, node_id):
-        """Returns the internal ip (Ray ip) of the given node."""
-        return self._get_cached_node(node_id=node_id)["internal_ip"]
-
-    @synchronized
-    @utils.debug_enabled(logger=logger)
-    def create_node(self, node_config, tags, count):
-        """Creates a number of nodes within the namespace."""
-        start_time = round(time.time() * 1000)
-        starting_insts = []
-        # Check first if it neccessary to create new nodes / start stopped nodes
-        VALIDITY_TAGS = [
-            TAG_RAY_CLUSTER_NAME,
-            TAG_RAY_NODE_KIND,
-            TAG_RAY_USER_NODE_TYPE,
-        ]
-        filters = {tag: tags[tag] for tag in VALIDITY_TAGS if tag in tags}
-
-        # Starting stopped nodes if cache_stopped_nodes=True
-        if self.cache_stopped_nodes:
-            logger.debug("Checking existing stopped nodes.")
-
-            filters_with_launch_config = copy.copy(filters)
-            if TAG_RAY_LAUNCH_CONFIG in tags:
-                filters_with_launch_config[TAG_RAY_LAUNCH_CONFIG] = tags[
-                    TAG_RAY_LAUNCH_CONFIG]
-
-            nodes_matching_launch_config = self.stopped_nodes(
-                filters_with_launch_config)
-            logger.debug(f"Found stopped nodes (with same launch config): "
-                         f"{len(nodes_matching_launch_config)}")
-
-            reuse_nodes = []
-            if len(nodes_matching_launch_config) >= count:
-                reuse_nodes = nodes_matching_launch_config[:count]
-            else:
-                nodes_all = self.stopped_nodes(filters)
-                logger.debug(f"Found stopped nodes (regardless launch config): "
-                             f"{len(nodes_all)}")
-                nodes_matching_launch_config_ids = [
-                    n["id"] for n in nodes_matching_launch_config
-                ]
-                nodes_non_matching_launch_config = [
-                    n for n in nodes_all
-                    if n["id"] not in nodes_matching_launch_config_ids
-                ]
-                reuse_nodes = (nodes_matching_launch_config +
-                               nodes_non_matching_launch_config)
-                reuse_nodes = reuse_nodes[:count]
-
-            logger.info(
-                f"Reusing nodes {len(reuse_nodes)}: {list(reuse_nodes)}. "
-                "To disable reuse, set `cache_stopped_nodes: False` "
-                "under `provider` in the cluster configuration.",)
-
-            for reuse_node in reuse_nodes:
-                if reuse_node["status"] == "STOPPING":
-                    get_instance_response = oci_adaptor.get_core_client(
-                        self.region,
-                        oci_utils.oci_config.get_profile()).get_instance(
-                            instance_id=reuse_node["id"])
-                    oci_adaptor.oci.wait_until(
-                        oci_adaptor.get_core_client(
-                            self.region, oci_utils.oci_config.get_profile()),
-                        get_instance_response,
-                        "lifecycle_state",
-                        "STOPPED",
-                    )
-
-            start_time1 = round(time.time() * 1000)
-            for matched_node in reuse_nodes:
-                matched_node_id = matched_node["id"]
-                instance_action_response = oci_adaptor.get_core_client(
-                    self.region,
-                    oci_utils.oci_config.get_profile()).instance_action(
-                        instance_id=matched_node_id, action="START")
-
-                starting_inst = instance_action_response.data
-                starting_insts.append({
-                    "inst_id": starting_inst.id,
-                    "ad": starting_inst.availability_domain,
-                    "compartment": starting_inst.compartment_id,
-                    "lifecycle_state": starting_inst.lifecycle_state,
-                    "oci_tags": starting_inst.freeform_tags,
-                })
-            count -= len(reuse_nodes)
-
-            launch_stopped_time = round(time.time() * 1000) - start_time1
-            logger.debug(
-                "Time elapsed(Launch stopped): {0} milli-seconds.".format(
-                    launch_stopped_time))
-        # end if self.cache_stopped_nodes:...
-
-        # Let's create additional new nodes (if neccessary)
-        if count > 0:
-            compartment = oci_query_helper.find_compartment(self.region)
-            vcn = oci_query_helper.find_create_vcn_subnet(self.region)
-            if vcn is None:
-                raise RuntimeError("VcnSubnetNotFound Error!")
-
-            ocpu_count = 0
-            vcpu_str = node_config["VCPUs"]
-            instance_type_str = node_config["InstanceType"]
-
-            if vcpu_str is not None and vcpu_str != "None":
-                if instance_type_str.startswith(
-                        f"{oci_utils.oci_config.VM_PREFIX}.A"):
-                    # For ARM cpu, 1*ocpu = 1*vcpu
-                    ocpu_count = round(float(vcpu_str))
-                else:
-                    # For Intel / AMD cpu, 1*ocpu = 2*vcpu
-                    ocpu_count = round(float(vcpu_str) / 2)
-            ocpu_count = 1 if (ocpu_count > 0 and
-                               ocpu_count < 1) else ocpu_count
-
-            machine_shape_config = None
-            if ocpu_count > 0:
-                mem = node_config["MemoryInGbs"]
-                if mem is not None and mem != "None":
-                    machine_shape_config = (oci_adaptor.oci.core.models.
-                                            LaunchInstanceShapeConfigDetails(
-                                                ocpus=ocpu_count,
-                                                memory_in_gbs=mem))
-                else:
-                    machine_shape_config = (oci_adaptor.oci.core.models.
-                                            LaunchInstanceShapeConfigDetails(
-                                                ocpus=ocpu_count))
-
-            preempitible_config = (
-                oci_adaptor.oci.core.models.PreemptibleInstanceConfigDetails(
-                    preemption_action=oci_adaptor.oci.core.models.
-                    TerminatePreemptionAction(type="TERMINATE",
-                                              preserve_boot_volume=False))
-                if node_config["Preemptible"] else None)
-
-            logger.debug(f"Shape: {instance_type_str}, ocpu: {ocpu_count}")
-            logger.debug(f"Shape config is {machine_shape_config}")
-            logger.debug(f"Spot config is {preempitible_config}")
-
-            vm_tags = {
-                **tags,
-                TAG_RAY_CLUSTER_NAME: self.cluster_name,
-                "sky_spot_flag": str(node_config["Preemptible"]).lower(),
-            }
-            # Use UTC time so that header & worker nodes use same rule
-            batch_id = datetime.utcnow().strftime("%Y%m%d%H%M%S")
-            node_type = tags[TAG_RAY_NODE_KIND]
-
-            oci_query_helper.subscribe_image(
-                compartment_id=compartment,
-                listing_id=node_config["AppCatalogListingId"],
-                resource_version=node_config["ResourceVersion"],
-                region=self.region,
-            )
-
-            start_time1 = round(time.time() * 1000)
-            for seq in range(1, count + 1):
-                launch_instance_response = oci_adaptor.get_core_client(
-                    self.region, oci_utils.oci_config.get_profile()
-                ).launch_instance(
-                    launch_instance_details=oci_adaptor.oci.core.models.
-                    LaunchInstanceDetails(
-                        availability_domain=node_config["AvailabilityDomain"],
-                        compartment_id=compartment,
-                        shape=instance_type_str,
-                        display_name=
-                        f"{self.cluster_name}_{node_type}_{batch_id}_{seq}",
-                        freeform_tags=vm_tags,
-                        metadata={
-                            "ssh_authorized_keys": node_config["AuthorizedKey"]
-                        },
-                        source_details=oci_adaptor.oci.core.models.
-                        InstanceSourceViaImageDetails(
-                            source_type="image",
-                            image_id=node_config["ImageId"],
-                            boot_volume_size_in_gbs=node_config[
-                                "BootVolumeSize"],
-                            boot_volume_vpus_per_gb=int(
-                                node_config["BootVolumePerf"]),
-                        ),
-                        create_vnic_details=oci_adaptor.oci.core.models.
-                        CreateVnicDetails(
-                            assign_public_ip=True,
-                            subnet_id=vcn,
-                        ),
-                        shape_config=machine_shape_config,
-                        preemptible_instance_config=preempitible_config,
-                    ))
-
-                new_inst = launch_instance_response.data
-                starting_insts.append({
-                    "inst_id": new_inst.id,
-                    "ad": new_inst.availability_domain,
-                    "compartment": new_inst.compartment_id,
-                    "lifecycle_state": new_inst.lifecycle_state,
-                    "oci_tags": new_inst.freeform_tags,
-                })
-            # end for loop
-
-            launch_new_time = round(time.time() * 1000) - start_time1
-            logger.debug("Time elapsed(Launch): {0} milli-seconds.".format(
-                launch_new_time))
-        # end if count > 0:...
-
-        for ninst in starting_insts:
-            # Waiting for the instance to be RUNNING state
-            get_instance_response = oci_adaptor.get_core_client(
-                self.region, oci_utils.oci_config.get_profile()).get_instance(
-                    instance_id=ninst["inst_id"])
-            oci_adaptor.oci.wait_until(
-                oci_adaptor.get_core_client(self.region,
-                                            oci_utils.oci_config.get_profile()),
-                get_instance_response,
-                "lifecycle_state",
-                "RUNNING",
-            )
-            ninst["lifecycle_state"] = "RUNNING"
-            self.cached_nodes[ninst["inst_id"]] = self.get_inst_obj(ninst)
-
-        total_time = round(time.time() * 1000) - start_time
-        logger.debug(
-            "Total time elapsed: {0} milli-seconds.".format(total_time))
-
-    def get_inst_obj(self, inst_info):
-        list_vnic_attachments_response = oci_adaptor.get_core_client(
-            self.region,
-            oci_utils.oci_config.get_profile()).list_vnic_attachments(
-                availability_domain=inst_info["ad"],
-                compartment_id=inst_info["compartment"],
-                instance_id=inst_info["inst_id"],
-            )
-
-        vnic = list_vnic_attachments_response.data[0]
-        get_vnic_response = (oci_adaptor.get_net_client(
-            self.region, oci_utils.oci_config.get_profile()).get_vnic(
-                vnic_id=vnic.vnic_id).data)
-
-        internal_ip = get_vnic_response.private_ip
-        external_ip = get_vnic_response.public_ip
-        if external_ip is None:
-            external_ip = internal_ip
-
-        return {
-            "id": inst_info["inst_id"],
-            "external_ip": external_ip,
-            "internal_ip": internal_ip,
-            "tags": inst_info["oci_tags"],
-            "status": inst_info["lifecycle_state"],
-        }
-
-    @synchronized
-    @utils.debug_enabled(logger=logger)
-    def set_node_tags(self, node_id, tags):
-        existing_tags = self._get_cached_node(node_id)["tags"]
-        combined_tags = dict(existing_tags, **tags)
-
-        self.cached_nodes[node_id]["tags"] = combined_tags
-        retry_count = 0
-        while retry_count < oci_utils.oci_config.MAX_RETRY_COUNT:
-            try:
-                oci_adaptor.get_core_client(
-                    self.region,
-                    oci_utils.oci_config.get_profile()).update_instance(
-                        instance_id=node_id,
-                        update_instance_details=oci_adaptor.oci.core.models.
-                        UpdateInstanceDetails(freeform_tags=combined_tags),
-                    )
-                logger.info(f"Tags are well set for node {node_id}")
-                break
-            except Exception as e:
-                retry_count = retry_count + 1
-                wait_seconds = oci_utils.oci_config.RETRY_INTERVAL_BASE_SECONDS * retry_count
-                logger.warn(
-                    f"Not ready yet, wait {wait_seconds} seconds & retry!")
-                logger.warn(f"Exception message is {str(e)}")
-                time.sleep(wait_seconds)
-
-    @synchronized
-    def terminate_node(self, node_id):
-        """Terminates the specified node."""
-        logger.info(f"terminate_node {node_id}...")
-        node = self._get_cached_node(node_id)
-        if node is None:
-            logger.info(f"The node is not existed: {node_id}..")
-            return  # Node not exists yet.
-
-        logger.debug(f"sky_spot_flag: {node['tags']['sky_spot_flag']}")
-        preemptibleFlag = (True if node and
-                           (str(node["tags"]["sky_spot_flag"]) == "true") else
-                           False)
-
-        if self.cache_stopped_nodes and not preemptibleFlag:
-            logger.info(f"Stopping instance {node_id}"
-                        "(to fully terminate instead, "
-                        "set `cache_stopped_nodes: False` "
-                        "under `provider` in the cluster configuration)")
-            instance_action_response = oci_adaptor.get_core_client(
-                self.region,
-                oci_utils.oci_config.get_profile()).instance_action(
-                    instance_id=node_id, action="STOP")
-            logger.info(
-                f"Stopped the instance {instance_action_response.data.id}")
-            if node_id in self.cached_nodes:
-                self.cached_nodes[node_id]["status"] = "STOPPED"
-            state_word = "Stopped"
-        else:
-            terminate_instance_response = oci_adaptor.get_core_client(
-                self.region,
-                oci_utils.oci_config.get_profile()).terminate_instance(node_id)
-            logger.debug(terminate_instance_response.data)
-            if node_id in self.cached_nodes:
-                del self.cached_nodes[node_id]
-            state_word = "Terminated"
-
-        logger.info(
-            f"{state_word} {node_id} w/ sky_spot_flag: {preemptibleFlag}.")
-
-    def _get_node(self, node_id):
-        self._get_filtered_nodes({},
-                                 force=True)  # All except for those terminated.
-        return self.cached_nodes.get(node_id, None)
-
-    def _get_cached_node(self, node_id):
-        if node_id in self.cached_nodes:
-            return self.cached_nodes[node_id]
-        return self._get_node(node_id=node_id)
-
-    def stopped_nodes(self, tag_filters):
-        """Return a list of stopped nodes filtered by the specified tags dict."""
-        nodes = self._get_filtered_nodes(tag_filters=tag_filters, force=True)
-        return [
-            v for _, v in nodes.items()
-            if v["status"] in ("STOPPED", "STOPPING")
-        ]
-
-    def running_nodes(self, tag_filters):
-        """Return a list of running node ids filtered by the specified tags dict."""
-        nodes = self._get_filtered_nodes(tag_filters=tag_filters)
-        return [k for k, v in nodes.items() if v["status"] == "RUNNING"]