chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/instance_manager/ray_installer.py

@@ -0,0 +1,96 @@
+import logging
+import subprocess
+
+from ray.autoscaler._private.updater import (
+    STATUS_UP_TO_DATE,
+    TAG_RAY_NODE_STATUS,
+    NodeUpdater,
+)
+from ray.autoscaler._private.util import with_envs, with_head_node_ip
+from ray.autoscaler.node_provider import NodeProvider as NodeProviderV1
+from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
+from ray.core.generated.instance_manager_pb2 import Instance
+
+logger = logging.getLogger(__name__)
+
+
+class RayInstaller(object):
+    """
+    RayInstaller is responsible for installing ray on the target instance.
+    """
+
+    def __init__(
+        self,
+        provider: NodeProviderV1,
+        config: AutoscalingConfig,
+        process_runner=subprocess,
+    ) -> None:
+        self._provider = provider
+        self._config = config
+        self._process_runner = process_runner
+
+    def install_ray(self, instance: Instance, head_node_ip: str) -> None:
+        """
+        Install ray on the target instance synchronously.
+        TODO:(rickyx): This runs in another thread, and errors are silently
+        ignored. We should propagate the error to the main thread.
+        """
+        setup_commands = self._config.get_worker_setup_commands(instance.instance_type)
+        ray_start_commands = self._config.get_worker_start_ray_commands()
+        docker_config = self._config.get_docker_config(instance.instance_type)
+
+        logger.info(
+            f"Creating new (spawn_updater) updater thread for node"
+            f" {instance.cloud_instance_id}."
+        )
+        provider_instance_type_name = self._config.get_provider_instance_type(
+            instance.instance_type
+        )
+        updater = NodeUpdater(
+            node_id=instance.cloud_instance_id,
+            provider_config=self._config.get_config("provider"),
+            provider=self._provider,
+            auth_config=self._config.get_config("auth"),
+            cluster_name=self._config.get_config("cluster_name"),
+            file_mounts=self._config.get_config("file_mounts"),
+            initialization_commands=with_head_node_ip(
+                self._config.get_initialization_commands(instance.instance_type),
+                head_node_ip,
+            ),
+            setup_commands=with_head_node_ip(setup_commands, head_node_ip),
+            # This will prepend envs to the begin of the ray start commands, e.g.
+            # `RAY_HEAD_IP=<head_node_ip> \
+            #  RAY_CLOUD_INSTANCE_ID=<instance_id> \
+            #  ray start --head ...`
+            #  See src/ray/common/constants.h for ENV name definitions.
+            ray_start_commands=with_envs(
+                ray_start_commands,
+                {
+                    "RAY_HEAD_IP": head_node_ip,
+                    "RAY_CLOUD_INSTANCE_ID": instance.cloud_instance_id,
+                    "RAY_NODE_TYPE_NAME": instance.instance_type,
+                    "RAY_CLOUD_INSTANCE_TYPE_NAME": provider_instance_type_name,
+                },
+            ),
+            runtime_hash=self._config.runtime_hash,
+            file_mounts_contents_hash=self._config.file_mounts_contents_hash,
+            is_head_node=False,
+            cluster_synced_files=self._config.get_config("cluster_synced_files"),
+            rsync_options={
+                "rsync_exclude": self._config.get_config("rsync_exclude"),
+                "rsync_filter": self._config.get_config("rsync_filter"),
+            },
+            use_internal_ip=True,
+            docker_config=docker_config,
+            node_resources=self._config.get_node_resources(instance.instance_type),
+            node_labels=self._config.get_node_labels(instance.instance_type),
+            process_runner=self._process_runner,
+        )
+        updater.run()
+        # check if the updater was successful by checking the node tags
+        # since the updater could hide exceptions and just set the status tag
+        tags = self._provider.node_tags(instance.cloud_instance_id)
+        if tags.get(TAG_RAY_NODE_STATUS) != STATUS_UP_TO_DATE:
+            raise Exception(
+                f"Ray installation failed with unexpected status: {tags.get(TAG_RAY_NODE_STATUS)}"
+            )