chore: import upstream snapshot with attribution

python/ray/autoscaler/_private/providers.py

@@ -0,0 +1,299 @@
+import copy
+import json
+import logging
+import os
+from typing import Any, Dict
+
+import yaml
+
+from ray.autoscaler._private.loader import load_function_or_class
+
+logger = logging.getLogger(__name__)
+
+# For caching provider instantiations across API calls of one python session
+_provider_instances = {}
+
+# Minimal config for compatibility with legacy-style external configs.
+MINIMAL_EXTERNAL_CONFIG = {
+    "available_node_types": {
+        "ray.head.default": {},
+        "ray.worker.default": {},
+    },
+    "head_node_type": "ray.head.default",
+    "head_node": {},
+    "worker_nodes": {},
+}
+
+
+def _import_aws(provider_config):
+    try:
+        # boto3 and botocore are imported in multiple places in the codebase,
+        # so we just import them here to ensure that they are installed.
+        import boto3  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "The Ray AWS VM launcher requires the AWS SDK for Python (Boto3) "
+            "to be installed. You can install it with `pip install boto3`."
+        ) from e
+
+    from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
+
+    return AWSNodeProvider
+
+
+def _import_gcp(provider_config):
+    try:
+        import googleapiclient  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "The Ray GCP VM launcher requires the Google API Client to be installed. "
+            "You can install it with `pip install google-api-python-client`."
+        ) from e
+
+    from ray.autoscaler._private.gcp.node_provider import GCPNodeProvider
+
+    return GCPNodeProvider
+
+
+def _import_azure(provider_config):
+    from ray.autoscaler._private._azure.node_provider import AzureNodeProvider
+
+    return AzureNodeProvider
+
+
+def _import_vsphere(provider_config):
+    from ray.autoscaler._private.vsphere.node_provider import VsphereWcpNodeProvider
+
+    return VsphereWcpNodeProvider
+
+
+def _import_local(provider_config):
+    if "coordinator_address" in provider_config:
+        from ray.autoscaler._private.local.coordinator_node_provider import (
+            CoordinatorSenderNodeProvider,
+        )
+
+        return CoordinatorSenderNodeProvider
+    else:
+        from ray.autoscaler._private.local.node_provider import LocalNodeProvider
+
+        return LocalNodeProvider
+
+
+def _import_readonly(provider_config):
+    from ray.autoscaler._private.readonly.node_provider import ReadOnlyNodeProvider
+
+    return ReadOnlyNodeProvider
+
+
+def _import_fake_multinode(provider_config):
+    from ray.autoscaler._private.fake_multi_node.node_provider import (
+        FakeMultiNodeProvider,
+    )
+
+    return FakeMultiNodeProvider
+
+
+def _import_fake_multinode_docker(provider_config):
+    from ray.autoscaler._private.fake_multi_node.node_provider import (
+        FakeMultiNodeDockerProvider,
+    )
+
+    return FakeMultiNodeDockerProvider
+
+
+def _import_kuberay(provider_config):
+    from ray.autoscaler._private.kuberay.node_provider import KubeRayNodeProvider
+
+    return KubeRayNodeProvider
+
+
+def _import_aliyun(provider_config):
+    from ray.autoscaler._private.aliyun.node_provider import AliyunNodeProvider
+
+    return AliyunNodeProvider
+
+
+def _import_spark(provider_config):
+    from ray.autoscaler._private.spark.node_provider import SparkNodeProvider
+
+    return SparkNodeProvider
+
+
+def _load_fake_multinode_defaults_config():
+    import ray.autoscaler._private.fake_multi_node as ray_fake_multinode
+
+    return os.path.join(os.path.dirname(ray_fake_multinode.__file__), "example.yaml")
+
+
+def _load_read_only_defaults_config():
+    import ray.autoscaler._private.readonly as ray_readonly
+
+    return os.path.join(os.path.dirname(ray_readonly.__file__), "defaults.yaml")
+
+
+def _load_fake_multinode_docker_defaults_config():
+    import ray.autoscaler._private.fake_multi_node as ray_fake_multinode
+
+    return os.path.join(
+        os.path.dirname(ray_fake_multinode.__file__), "example_docker.yaml"
+    )
+
+
+def _load_local_defaults_config():
+    import ray.autoscaler.local as ray_local
+
+    return os.path.join(os.path.dirname(ray_local.__file__), "defaults.yaml")
+
+
+def _load_aws_defaults_config():
+    import ray.autoscaler.aws as ray_aws
+
+    return os.path.join(os.path.dirname(ray_aws.__file__), "defaults.yaml")
+
+
+def _load_vsphere_defaults_config():
+    import ray.autoscaler.vsphere as ray_vsphere
+
+    return os.path.join(os.path.dirname(ray_vsphere.__file__), "defaults.yaml")
+
+
+def _load_gcp_defaults_config():
+    import ray.autoscaler.gcp as ray_gcp
+
+    return os.path.join(os.path.dirname(ray_gcp.__file__), "defaults.yaml")
+
+
+def _load_azure_defaults_config():
+    import ray.autoscaler.azure as ray_azure
+
+    return os.path.join(os.path.dirname(ray_azure.__file__), "defaults.yaml")
+
+
+def _load_aliyun_defaults_config():
+    import ray.autoscaler.aliyun as ray_aliyun
+
+    return os.path.join(os.path.dirname(ray_aliyun.__file__), "defaults.yaml")
+
+
+def _import_external(provider_config):
+    provider_cls = load_function_or_class(path=provider_config["module"])
+    return provider_cls
+
+
+_NODE_PROVIDERS = {
+    "local": _import_local,
+    "fake_multinode": _import_fake_multinode,
+    "fake_multinode_docker": _import_fake_multinode_docker,
+    "readonly": _import_readonly,
+    "aws": _import_aws,
+    "gcp": _import_gcp,
+    "vsphere": _import_vsphere,
+    "azure": _import_azure,
+    "kuberay": _import_kuberay,
+    "aliyun": _import_aliyun,
+    "external": _import_external,  # Import an external module
+    "spark": _import_spark,
+}
+
+_PROVIDER_PRETTY_NAMES = {
+    "readonly": "Readonly (Manual Cluster Setup)",
+    "fake_multinode": "Fake Multinode",
+    "fake_multinode_docker": "Fake Multinode Docker",
+    "local": "Local",
+    "aws": "AWS",
+    "gcp": "GCP",
+    "azure": "Azure",
+    "kuberay": "KubeRay",
+    "aliyun": "Aliyun",
+    "external": "External",
+    "vsphere": "vSphere",
+    "spark": "Spark",
+}
+
+_DEFAULT_CONFIGS = {
+    "fake_multinode": _load_fake_multinode_defaults_config,
+    "fake_multinode_docker": _load_fake_multinode_docker_defaults_config,
+    "local": _load_local_defaults_config,
+    "aws": _load_aws_defaults_config,
+    "gcp": _load_gcp_defaults_config,
+    "azure": _load_azure_defaults_config,
+    "aliyun": _load_aliyun_defaults_config,
+    "vsphere": _load_vsphere_defaults_config,
+    "readonly": _load_read_only_defaults_config,
+}
+
+
+def _get_node_provider_cls(provider_config: Dict[str, Any]):
+    """Get the node provider class for a given provider config.
+
+    Note that this may be used by private node providers that proxy methods to
+    built-in node providers, so we should maintain backwards compatibility.
+
+    Args:
+        provider_config: provider section of the autoscaler config.
+
+    Returns:
+        NodeProvider class
+    """
+    importer = _NODE_PROVIDERS.get(provider_config["type"])
+    if importer is None:
+        raise NotImplementedError(
+            "Unsupported node provider: {}".format(provider_config["type"])
+        )
+    return importer(provider_config)
+
+
+def _get_node_provider(
+    provider_config: Dict[str, Any], cluster_name: str, use_cache: bool = True
+) -> Any:
+    """Get the instantiated node provider for a given provider config.
+
+    Note that this may be used by private node providers that proxy methods to
+    built-in node providers, so we should maintain backwards compatibility.
+
+    Args:
+        provider_config: provider section of the autoscaler config.
+        cluster_name: cluster name from the autoscaler config.
+        use_cache: whether or not to use a cached definition if available. If
+            False, the returned object will also not be stored in the cache.
+
+    Returns:
+        NodeProvider
+    """
+    provider_key = (json.dumps(provider_config, sort_keys=True), cluster_name)
+    if use_cache and provider_key in _provider_instances:
+        return _provider_instances[provider_key]
+
+    provider_cls = _get_node_provider_cls(provider_config)
+    new_provider = provider_cls(provider_config, cluster_name)
+
+    if use_cache:
+        _provider_instances[provider_key] = new_provider
+
+    return new_provider
+
+
+def _clear_provider_cache():
+    global _provider_instances
+    _provider_instances = {}
+
+
+def _get_default_config(provider_config):
+    """Retrieve a node provider.
+
+    This is an INTERNAL API. It is not allowed to call this from any Ray
+    package outside the autoscaler.
+    """
+    if provider_config["type"] == "external":
+        return copy.deepcopy(MINIMAL_EXTERNAL_CONFIG)
+    load_config = _DEFAULT_CONFIGS.get(provider_config["type"])
+    if load_config is None:
+        raise NotImplementedError(
+            "Unsupported node provider: {}".format(provider_config["type"])
+        )
+    path_to_default = load_config()
+    with open(path_to_default) as f:
+        defaults = yaml.safe_load(f)
+
+    return defaults