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

sky/skylet/providers/azure/config.py

@@ -1,170 +0,0 @@
-import json
-import logging
-import random
-from hashlib import sha256
-from pathlib import Path
-import time
-from typing import Any, Callable
-
-from azure.common.credentials import get_cli_profile
-from azure.identity import AzureCliCredential
-from azure.mgmt.network import NetworkManagementClient
-from azure.mgmt.resource import ResourceManagementClient
-from azure.mgmt.resource.resources.models import DeploymentMode
-
-from sky.utils import common_utils
-
-UNIQUE_ID_LEN = 4
-_WAIT_NSG_CREATION_NUM_TIMEOUT_SECONDS = 600
-
-logger = logging.getLogger(__name__)
-
-
-def get_azure_sdk_function(client: Any, function_name: str) -> Callable:
-    """Retrieve a callable function from Azure SDK client object.
-
-    Newer versions of the various client SDKs renamed function names to
-    have a begin_ prefix. This function supports both the old and new
-    versions of the SDK by first trying the old name and falling back to
-    the prefixed new name.
-    """
-    func = getattr(
-        client, function_name, getattr(client, f"begin_{function_name}", None)
-    )
-    if func is None:
-        raise AttributeError(
-            "'{obj}' object has no {func} or begin_{func} attribute".format(
-                obj={client.__name__}, func=function_name
-            )
-        )
-    return func
-
-
-def bootstrap_azure(config):
-    config = _configure_key_pair(config)
-    config = _configure_resource_group(config)
-    return config
-
-
-def _configure_resource_group(config):
-    # TODO: look at availability sets
-    # https://docs.microsoft.com/en-us/azure/virtual-machines/windows/tutorial-availability-sets
-    subscription_id = config["provider"].get("subscription_id")
-    if subscription_id is None:
-        subscription_id = get_cli_profile().get_subscription_id()
-    # Increase the timeout to fix the Azure get-access-token (used by ray azure
-    # node_provider) timeout issue.
-    # Tracked in https://github.com/Azure/azure-cli/issues/20404#issuecomment-1249575110
-    credentials = AzureCliCredential(process_timeout=30)
-    resource_client = ResourceManagementClient(credentials, subscription_id)
-    config["provider"]["subscription_id"] = subscription_id
-    logger.info("Using subscription id: %s", subscription_id)
-
-    assert (
-        "resource_group" in config["provider"]
-    ), "Provider config must include resource_group field"
-    resource_group = config["provider"]["resource_group"]
-
-    assert (
-        "location" in config["provider"]
-    ), "Provider config must include location field"
-    params = {"location": config["provider"]["location"]}
-
-    if "tags" in config["provider"]:
-        params["tags"] = config["provider"]["tags"]
-
-    logger.info("Creating/Updating resource group: %s", resource_group)
-    rg_create_or_update = get_azure_sdk_function(
-        client=resource_client.resource_groups, function_name="create_or_update"
-    )
-    rg_create_or_update(resource_group_name=resource_group, parameters=params)
-
-    # load the template file
-    current_path = Path(__file__).parent
-    template_path = current_path.joinpath("azure-config-template.json")
-    with open(template_path, "r") as template_fp:
-        template = json.load(template_fp)
-
-    logger.info("Using cluster name: %s", config["cluster_name"])
-
-    # set unique id for resources in this cluster
-    unique_id = config["provider"].get("unique_id")
-    if unique_id is None:
-        hasher = sha256()
-        hasher.update(config["provider"]["resource_group"].encode("utf-8"))
-        unique_id = hasher.hexdigest()[:UNIQUE_ID_LEN]
-    else:
-        unique_id = str(unique_id)
-    config["provider"]["unique_id"] = unique_id
-    logger.info("Using unique id: %s", unique_id)
-    cluster_id = "{}-{}".format(config["cluster_name"], unique_id)
-
-    subnet_mask = config["provider"].get("subnet_mask")
-    if subnet_mask is None:
-        # choose a random subnet, skipping most common value of 0
-        random.seed(unique_id)
-        subnet_mask = "10.{}.0.0/16".format(random.randint(1, 254))
-    logger.info("Using subnet mask: %s", subnet_mask)
-
-    parameters = {
-        "properties": {
-            "mode": DeploymentMode.incremental,
-            "template": template,
-            "parameters": {
-                "subnet": {"value": subnet_mask},
-                "clusterId": {"value": cluster_id},
-            },
-        }
-    }
-
-    create_or_update = get_azure_sdk_function(
-        client=resource_client.deployments, function_name="create_or_update"
-    )
-    # TODO (skypilot): this takes a long time (> 40 seconds) for stopping an
-    # azure VM, and this can be called twice during ray down.
-    outputs = (
-        create_or_update(
-            resource_group_name=resource_group,
-            deployment_name="ray-config",
-            parameters=parameters,
-        )
-        .result()
-        .properties.outputs
-    )
-
-    # We should wait for the NSG to be created before opening any ports
-    # to avoid overriding the newly-added NSG rules.
-    nsg_id = outputs["nsg"]["value"]
-    nsg_name = nsg_id.split("/")[-1]
-    network_client = NetworkManagementClient(credentials, subscription_id)
-    backoff = common_utils.Backoff(max_backoff_factor=1)
-    start_time = time.time()
-    while True:
-        nsg = network_client.network_security_groups.get(resource_group, nsg_name)
-        if nsg.provisioning_state == "Succeeded":
-            break
-        if time.time() - start_time > _WAIT_NSG_CREATION_NUM_TIMEOUT_SECONDS:
-            raise RuntimeError(
-                f"Fails to create NSG {nsg_name} in {resource_group} within "
-                f"{_WAIT_NSG_CREATION_NUM_TIMEOUT_SECONDS} seconds."
-            )
-        backoff_time = backoff.current_backoff()
-        logger.info(
-            f"NSG {nsg_name} is not created yet. Waiting for "
-            f"{backoff_time} seconds before checking again."
-        )
-        time.sleep(backoff_time)
-
-    # append output resource ids to be used with vm creation
-    config["provider"]["msi"] = outputs["msi"]["value"]
-    config["provider"]["nsg"] = nsg_id
-    config["provider"]["subnet"] = outputs["subnet"]["value"]
-
-    return config
-
-
-def _configure_key_pair(config):
-    # SkyPilot: The original checks and configurations are no longer
-    # needed, since we have already set them up in the upper level
-    # SkyPilot codes. See sky/templates/azure-ray.yml.j2
-    return config

sky/skylet/providers/azure/node_provider.py

@@ -1,466 +0,0 @@
-import copy
-import json
-import logging
-from pathlib import Path
-from threading import RLock
-from uuid import uuid4
-
-from azure.identity import AzureCliCredential
-from azure.mgmt.compute import ComputeManagementClient
-from azure.mgmt.network import NetworkManagementClient
-from azure.mgmt.resource import ResourceManagementClient
-from azure.mgmt.resource.resources.models import DeploymentMode
-
-from sky.skylet.providers.azure.config import (
-    bootstrap_azure,
-    get_azure_sdk_function,
-)
-from sky.skylet import autostop_lib
-from sky.skylet.providers.command_runner import SkyDockerCommandRunner
-from sky.provision import docker_utils
-
-from ray.autoscaler._private.command_runner import SSHCommandRunner
-from ray.autoscaler.node_provider import NodeProvider
-from ray.autoscaler.tags import (
-    TAG_RAY_CLUSTER_NAME,
-    TAG_RAY_LAUNCH_CONFIG,
-    TAG_RAY_NODE_KIND,
-    TAG_RAY_NODE_NAME,
-    TAG_RAY_USER_NODE_TYPE,
-)
-
-VM_NAME_MAX_LEN = 64
-UNIQUE_ID_LEN = 4
-
-logger = logging.getLogger(__name__)
-azure_logger = logging.getLogger("azure.core.pipeline.policies.http_logging_policy")
-azure_logger.setLevel(logging.WARNING)
-
-
-def synchronized(f):
-    def wrapper(self, *args, **kwargs):
-        self.lock.acquire()
-        try:
-            return f(self, *args, **kwargs)
-        finally:
-            self.lock.release()
-
-    return wrapper
-
-
-class AzureNodeProvider(NodeProvider):
-    """Node Provider for Azure
-
-    This provider assumes Azure credentials are set by running ``az login``
-    and the default subscription is configured through ``az account``
-    or set in the ``provider`` field of the autoscaler configuration.
-
-    Nodes may be in one of three states: {pending, running, terminated}. Nodes
-    appear immediately once started by ``create_node``, and transition
-    immediately to terminated when ``terminate_node`` is called.
-    """
-
-    def __init__(self, provider_config, cluster_name):
-        NodeProvider.__init__(self, provider_config, cluster_name)
-        if not autostop_lib.get_is_autostopping():
-            # TODO(suquark): This is a temporary patch for resource group.
-            # By default, Ray autoscaler assumes the resource group is still
-            # here even after the whole cluster is destroyed. However, now we
-            # deletes the resource group after tearing down the cluster. To
-            # comfort the autoscaler, we need to create/update it here, so the
-            # resource group always exists.
-            #
-            # We should not re-configure the resource group again, when it is
-            # running on the remote VM and the autostopping is in progress,
-            # because the VM is running which guarantees the resource group
-            # exists.
-            from sky.skylet.providers.azure.config import _configure_resource_group
-
-            _configure_resource_group(
-                {"cluster_name": cluster_name, "provider": provider_config}
-            )
-        subscription_id = provider_config["subscription_id"]
-        self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
-        # Sky only supports Azure CLI credential for now.
-        # Increase the timeout to fix the Azure get-access-token (used by ray azure
-        # node_provider) timeout issue.
-        # Tracked in https://github.com/Azure/azure-cli/issues/20404#issuecomment-1249575110
-        credential = AzureCliCredential(process_timeout=30)
-        self.compute_client = ComputeManagementClient(credential, subscription_id)
-        self.network_client = NetworkManagementClient(credential, subscription_id)
-        self.resource_client = ResourceManagementClient(credential, subscription_id)
-
-        self.lock = RLock()
-
-        # cache node objects
-        self.cached_nodes = {}
-
-    @synchronized
-    def _get_filtered_nodes(self, tag_filters):
-        # add cluster name filter to only get nodes from this cluster
-        cluster_tag_filters = {**tag_filters, TAG_RAY_CLUSTER_NAME: self.cluster_name}
-
-        def match_tags(vm):
-            for k, v in cluster_tag_filters.items():
-                if vm.tags.get(k) != v:
-                    return False
-            return True
-
-        vms = self.compute_client.virtual_machines.list(
-            resource_group_name=self.provider_config["resource_group"]
-        )
-
-        nodes = [self._extract_metadata(vm) for vm in filter(match_tags, vms)]
-        self.cached_nodes = {node["name"]: node for node in nodes}
-        return self.cached_nodes
-
-    def _extract_metadata(self, vm):
-        # get tags
-        metadata = {"name": vm.name, "tags": vm.tags, "status": ""}
-
-        # get status
-        resource_group = self.provider_config["resource_group"]
-        instance = self.compute_client.virtual_machines.instance_view(
-            resource_group_name=resource_group, vm_name=vm.name
-        ).as_dict()
-        for status in instance["statuses"]:
-            code_state = status["code"].split("/")
-            # It is possible that sometimes the 'code' is empty string, and we
-            # should skip them.
-            if len(code_state) != 2:
-                continue
-            code, state = code_state
-            # skip provisioning status
-            if code == "PowerState":
-                metadata["status"] = state
-                break
-
-        # get ip data
-        nic_id = vm.network_profile.network_interfaces[0].id
-        metadata["nic_name"] = nic_id.split("/")[-1]
-        nic = self.network_client.network_interfaces.get(
-            resource_group_name=resource_group,
-            network_interface_name=metadata["nic_name"],
-        )
-        ip_config = nic.ip_configurations[0]
-
-        if not self.provider_config.get("use_internal_ips", False):
-            public_ip_id = ip_config.public_ip_address.id
-            metadata["public_ip_name"] = public_ip_id.split("/")[-1]
-            public_ip = self.network_client.public_ip_addresses.get(
-                resource_group_name=resource_group,
-                public_ip_address_name=metadata["public_ip_name"],
-            )
-            metadata["external_ip"] = public_ip.ip_address
-
-        metadata["internal_ip"] = ip_config.private_ip_address
-
-        return metadata
-
-    def stopped_nodes(self, tag_filters):
-        """Return a list of stopped 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"].startswith("deallocat")]
-
-    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 nodes() to
-        serve single-node queries (e.g. is_running(node_id)). This means that
-        nodes() must be called again to refresh results.
-
-        Examples:
-            >>> from ray.autoscaler.tags import TAG_RAY_NODE_KIND
-            >>> provider = ... # doctest: +SKIP
-            >>> provider.non_terminated_nodes( # doctest: +SKIP
-            ...     {TAG_RAY_NODE_KIND: "worker"})
-            ["node-1", "node-2"]
-        """
-        nodes = self._get_filtered_nodes(tag_filters=tag_filters)
-        return [k for k, v in nodes.items() if not v["status"].startswith("deallocat")]
-
-    def is_running(self, node_id):
-        """Return whether the specified node is running."""
-        # always get current status
-        node = self._get_node(node_id=node_id)
-        return node["status"] == "running"
-
-    def is_terminated(self, node_id):
-        """Return whether the specified node is terminated."""
-        # always get current status
-        node = self._get_node(node_id=node_id)
-        return node["status"].startswith("deallocat")
-
-    def node_tags(self, node_id):
-        """Returns the tags of the given node (string dict)."""
-        return self._get_cached_node(node_id=node_id)["tags"]
-
-    def external_ip(self, node_id):
-        """Returns the external ip of the given node."""
-        ip = (
-            self._get_cached_node(node_id=node_id)["external_ip"]
-            or self._get_node(node_id=node_id)["external_ip"]
-        )
-        return ip
-
-    def internal_ip(self, node_id):
-        """Returns the internal ip (Ray ip) of the given node."""
-        ip = (
-            self._get_cached_node(node_id=node_id)["internal_ip"]
-            or self._get_node(node_id=node_id)["internal_ip"]
-        )
-        return ip
-
-    def create_node(self, node_config, tags, count):
-        resource_group = self.provider_config["resource_group"]
-
-        if self.cache_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}
-            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
-                ]
-
-            # SkyPilot: We try to use the instances with the same matching launch_config first. If
-            # there is not enough instances with matching launch_config, we then use all the
-            # instances with the same matching launch_config plus some instances with wrong
-            # launch_config.
-            nodes_matching_launch_config = self.stopped_nodes(
-                filters_with_launch_config
-            )
-            nodes_matching_launch_config.sort(reverse=True)
-            if len(nodes_matching_launch_config) >= count:
-                reuse_nodes = nodes_matching_launch_config[:count]
-            else:
-                nodes_all = self.stopped_nodes(filters)
-                nodes_non_matching_launch_config = [
-                    n for n in nodes_all if n not in nodes_matching_launch_config
-                ]
-                # This sort is for backward compatibility, where the user already has
-                # leaked stopped nodes with the different launch config before update
-                # to #1671, and the total number of the leaked nodes is greater than
-                # the number of nodes to be created. With this, we make sure the nodes
-                # are reused in a deterministic order (sorting by str IDs; we cannot
-                # get the launch time info here; otherwise, sort by the launch time
-                # is more accurate.)
-                # This can be removed in the future when we are sure all the users
-                # have updated to #1671.
-                nodes_non_matching_launch_config.sort(reverse=True)
-                reuse_nodes = (
-                    nodes_matching_launch_config + nodes_non_matching_launch_config
-                )
-                # The total number of reusable nodes can be less than the number of nodes to be created.
-                # This `[:count]` is fine, as it will get all the reusable nodes, even if there are
-                # less nodes.
-                reuse_nodes = reuse_nodes[:count]
-
-            logger.info(
-                f"Reusing nodes {list(reuse_nodes)}. "
-                "To disable reuse, set `cache_stopped_nodes: False` "
-                "under `provider` in the cluster configuration.",
-            )
-            start = get_azure_sdk_function(
-                client=self.compute_client.virtual_machines, function_name="start"
-            )
-            for node_id in reuse_nodes:
-                start(resource_group_name=resource_group, vm_name=node_id).wait()
-                self.set_node_tags(node_id, tags)
-            count -= len(reuse_nodes)
-
-        if count:
-            self._create_node(node_config, tags, count)
-
-    def _create_node(self, node_config, tags, count):
-        """Creates a number of nodes within the namespace."""
-        resource_group = self.provider_config["resource_group"]
-
-        # load the template file
-        current_path = Path(__file__).parent
-        template_path = current_path.joinpath("azure-vm-template.json")
-        with open(template_path, "r") as template_fp:
-            template = json.load(template_fp)
-
-        # get the tags
-        config_tags = node_config.get("tags", {}).copy()
-        config_tags.update(tags)
-        config_tags[TAG_RAY_CLUSTER_NAME] = self.cluster_name
-
-        vm_name = "{node}-{unique_id}-{vm_id}".format(
-            node=config_tags.get(TAG_RAY_NODE_NAME, "node"),
-            unique_id=self.provider_config["unique_id"],
-            vm_id=uuid4().hex[:UNIQUE_ID_LEN],
-        )[:VM_NAME_MAX_LEN]
-        use_internal_ips = self.provider_config.get("use_internal_ips", False)
-
-        template_params = node_config["azure_arm_parameters"].copy()
-        template_params["vmName"] = vm_name
-        template_params["provisionPublicIp"] = not use_internal_ips
-        template_params["vmTags"] = config_tags
-        template_params["vmCount"] = count
-        template_params["msi"] = self.provider_config["msi"]
-        template_params["nsg"] = self.provider_config["nsg"]
-        template_params["subnet"] = self.provider_config["subnet"]
-
-        parameters = {
-            "properties": {
-                "mode": DeploymentMode.incremental,
-                "template": template,
-                "parameters": {
-                    key: {"value": value} for key, value in template_params.items()
-                },
-            }
-        }
-
-        # TODO: we could get the private/public ips back directly
-        create_or_update = get_azure_sdk_function(
-            client=self.resource_client.deployments, function_name="create_or_update"
-        )
-        create_or_update(
-            resource_group_name=resource_group,
-            deployment_name=vm_name,
-            parameters=parameters,
-        ).wait()
-
-    @synchronized
-    def set_node_tags(self, node_id, tags):
-        """Sets the tag values (string dict) for the specified node."""
-        node_tags = self._get_cached_node(node_id)["tags"]
-        node_tags.update(tags)
-        update = get_azure_sdk_function(
-            client=self.compute_client.virtual_machines, function_name="update"
-        )
-        update(
-            resource_group_name=self.provider_config["resource_group"],
-            vm_name=node_id,
-            parameters={"tags": node_tags},
-        )
-        self.cached_nodes[node_id]["tags"] = node_tags
-
-    def terminate_node(self, node_id):
-        """Terminates the specified node. This will delete the VM and
-        associated resources (NIC, IP, Storage) for the specified node."""
-
-        resource_group = self.provider_config["resource_group"]
-        try:
-            # get metadata for node
-            metadata = self._get_node(node_id)
-        except KeyError:
-            # node no longer exists
-            return
-
-        if self.cache_stopped_nodes:
-            try:
-                # stop machine and leave all resources
-                logger.info(
-                    f"Stopping instance {node_id}"
-                    "(to fully terminate instead, "
-                    "set `cache_stopped_nodes: False` "
-                    "under `provider` in the cluster configuration)"
-                )
-                stop = get_azure_sdk_function(
-                    client=self.compute_client.virtual_machines,
-                    function_name="deallocate",
-                )
-                stop(resource_group_name=resource_group, vm_name=node_id)
-            except Exception as e:
-                logger.warning("Failed to stop VM: {}".format(e))
-        else:
-            vm = self.compute_client.virtual_machines.get(
-                resource_group_name=resource_group, vm_name=node_id
-            )
-            disks = {d.name for d in vm.storage_profile.data_disks}
-            disks.add(vm.storage_profile.os_disk.name)
-
-            try:
-                # delete machine, must wait for this to complete
-                delete = get_azure_sdk_function(
-                    client=self.compute_client.virtual_machines, function_name="delete"
-                )
-                delete(resource_group_name=resource_group, vm_name=node_id).wait()
-            except Exception as e:
-                logger.warning("Failed to delete VM: {}".format(e))
-
-            try:
-                # delete nic
-                delete = get_azure_sdk_function(
-                    client=self.network_client.network_interfaces,
-                    function_name="delete",
-                )
-                delete(
-                    resource_group_name=resource_group,
-                    network_interface_name=metadata["nic_name"],
-                )
-            except Exception as e:
-                logger.warning("Failed to delete nic: {}".format(e))
-
-            # delete ip address
-            if "public_ip_name" in metadata:
-                try:
-                    delete = get_azure_sdk_function(
-                        client=self.network_client.public_ip_addresses,
-                        function_name="delete",
-                    )
-                    delete(
-                        resource_group_name=resource_group,
-                        public_ip_address_name=metadata["public_ip_name"],
-                    )
-                except Exception as e:
-                    logger.warning("Failed to delete public ip: {}".format(e))
-
-            # delete disks
-            for disk in disks:
-                try:
-                    delete = get_azure_sdk_function(
-                        client=self.compute_client.disks, function_name="delete"
-                    )
-                    delete(resource_group_name=resource_group, disk_name=disk)
-                except Exception as e:
-                    logger.warning("Failed to delete disk: {}".format(e))
-
-    def _get_node(self, node_id):
-        self._get_filtered_nodes({})  # Side effect: updates cache
-        return self.cached_nodes[node_id]
-
-    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)
-
-    @staticmethod
-    def bootstrap_config(cluster_config):
-        return bootstrap_azure(cluster_config)
-
-    def get_command_runner(
-        self,
-        log_prefix,
-        node_id,
-        auth_config,
-        cluster_name,
-        process_runner,
-        use_internal_ip,
-        docker_config=None,
-    ):
-        common_args = {
-            "log_prefix": log_prefix,
-            "node_id": node_id,
-            "provider": self,
-            "auth_config": auth_config,
-            "cluster_name": cluster_name,
-            "process_runner": process_runner,
-            "use_internal_ip": use_internal_ip,
-        }
-        if docker_config and docker_config["container_name"] != "":
-            if "docker_login_config" in self.provider_config:
-                docker_config["docker_login_config"] = docker_utils.DockerLoginConfig(
-                    **self.provider_config["docker_login_config"]
-                )
-            return SkyDockerCommandRunner(docker_config, **common_args)
-        else:
-            return SSHCommandRunner(**common_args)

sky/skylet/providers/lambda_cloud/__init__.py

@@ -1,2 +0,0 @@
-"""Lambda Cloud node provider"""
-from sky.skylet.providers.lambda_cloud.node_provider import LambdaNodeProvider