ray-project--ray/python/ray/tests/gcp/test_gcp_tpu_command_runner.py

import hashlib
import os
import sys
from getpass import getuser
from unittest.mock import patch

import pytest

from ray._private import ray_constants
from ray.autoscaler._private.command_runner import SSHCommandRunner
from ray.autoscaler._private.gcp.tpu_command_runner import TPUCommandRunner
from ray.tests.test_autoscaler import MockProcessRunner, MockProvider

_MOCK_TPU_NAME = "my-tpu"
_MOCK_ACCELERATOR_TYPE = "v4-16"

auth_config = {
    "ssh_user": "ray",
    "ssh_private_key": "8265.pem",
}


class MockTpuInstance:
    def __init__(self, num_workers: int = 1):
        self.num_workers = num_workers

    def get_internal_ip(self, worker_index: int) -> str:
        return "0.0.0.0"

    def get_external_ip(self, worker_index: int) -> str:
        return "1.2.3.4"

    def get(self, key) -> str:
        if key == "name":
            return _MOCK_TPU_NAME
        elif key == "acceleratorType":
            return _MOCK_ACCELERATOR_TYPE
        return ""


def test_tpu_ssh_command_runner():
    num_workers = 2
    process_runner = MockProcessRunner()
    provider = MockProvider()
    instance = MockTpuInstance(num_workers=num_workers)
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
        ssh_user_hash[:10], ssh_control_hash[:10]
    )
    args = {
        "instance": instance,
        "log_prefix": "prefix",
        "node_id": "abc",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
    }
    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
    cmd_runner = TPUCommandRunner(**args)
    cmd_runner.run(
        "echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
    )

    expected = [
        "ssh",
        "-tt",
        "-L",
        "8265:localhost:8265",
        "-i",
        "8265.pem",
        "-o",
        "StrictHostKeyChecking=no",
        "-o",
        "UserKnownHostsFile=/dev/null",
        "-o",
        "IdentitiesOnly=yes",
        "-o",
        "ExitOnForwardFailure=yes",
        "-o",
        "ServerAliveInterval=5",
        "-o",
        "ServerAliveCountMax=3",
        "-o",
        "ControlMaster=auto",
        "-o",
        "ControlPath={}/%C".format(ssh_control_path),
        "-o",
        "ControlPersist=10s",
        "-o",
        "ConnectTimeout=120s",
        "ray@1.2.3.4",
        "bash",
        "--login",
        "-c",
        "-i",
        """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"quote between this \\" and this"'"'"';export var2='"'"'"123"'"'"';echo helloo)'""",  # noqa: E501
    ]

    calls = process_runner.calls

    # Asserts that we do make the call once per worker in the TPU pod.
    assert len(process_runner.calls) == num_workers

    # Much easier to debug this loop than the function call.
    for i in range(num_workers):
        for x, y in zip(calls[i], expected):
            assert x == y


def test_tpu_docker_command_runner():
    num_workers = 4
    process_runner = MockProcessRunner()
    provider = MockProvider()
    instance = MockTpuInstance(num_workers=num_workers)
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
        ssh_user_hash[:10], ssh_control_hash[:10]
    )
    docker_config = {"container_name": "container"}
    args = {
        "instance": instance,
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = TPUCommandRunner(**args)
    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
    cmd_runner.run("echo hello", environment_variables=env_vars)

    # This string is insane because there are an absurd number of embedded
    # quotes. While this is a ridiculous string, the escape behavior is
    # important and somewhat difficult to get right for environment variables.
    cmd = """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (docker exec -it  container /bin/bash -c '"'"'bash --login -c -i '"'"'"'"'"'"'"'"'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"quote between this \\" and this"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';export var2='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"123"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';echo hello)'"'"'"'"'"'"'"'"''"'"' )'"""  # noqa: E501

    expected = [
        "ssh",
        "-tt",
        "-i",
        "8265.pem",
        "-o",
        "StrictHostKeyChecking=no",
        "-o",
        "UserKnownHostsFile=/dev/null",
        "-o",
        "IdentitiesOnly=yes",
        "-o",
        "ExitOnForwardFailure=yes",
        "-o",
        "ServerAliveInterval=5",
        "-o",
        "ServerAliveCountMax=3",
        "-o",
        "ControlMaster=auto",
        "-o",
        "ControlPath={}/%C".format(ssh_control_path),
        "-o",
        "ControlPersist=10s",
        "-o",
        "ConnectTimeout=120s",
        "ray@1.2.3.4",
        "bash",
        "--login",
        "-c",
        "-i",
        cmd,
    ]

    calls = process_runner.calls

    # Asserts that we do make the call once per worker in the TPU pod.
    assert len(process_runner.calls) == num_workers

    # Much easier to debug this loop than the function call.
    for i in range(num_workers):
        for x, y in zip(calls[i], expected):
            assert x == y


def test_tpu_docker_run_init():
    num_workers = 1
    process_runner = MockProcessRunner()
    provider = MockProvider()
    instance = MockTpuInstance(num_workers=num_workers)
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    docker_config = {
        "container_name": "container",
        "image": "rayproject/ray:latest",
    }
    args = {
        "instance": instance,
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = TPUCommandRunner(**args)

    # Taken from tests/test_command_runner.py
    # This mocks the response of 'docker inspect' command to return an empty JSON array.
    # This simulates the scenario where the Docker image has no set environment
    # variables, allowing us to test the subsequent code for handling this case.
    process_runner.respond_to_call("json .Config.Env", 2 * ["[]"])
    cmd_runner.run_init(as_head=True, file_mounts={}, sync_run_yet=True)
    process_runner.assert_has_call("1.2.3.4", pattern="docker")


def test_max_active_connections_env_var():
    num_workers = 2
    process_runner = MockProcessRunner()
    provider = MockProvider()
    instance = MockTpuInstance(num_workers=num_workers)
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    docker_config = {"container_name": "container"}
    args = {
        "instance": instance,
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = TPUCommandRunner(**args)
    os.environ[ray_constants.RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR] = "1"
    num_connections = cmd_runner.num_connections
    assert type(num_connections) is int
    assert num_connections == 1


def test_tpu_pod_resources():
    num_workers = 2
    process_runner = MockProcessRunner()
    provider = MockProvider()
    instance = MockTpuInstance(num_workers=num_workers)
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    args = {
        "instance": instance,
        "log_prefix": "prefix",
        "node_id": "abc",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
    }
    env_vars = {
        ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: {
            "TPU": 4,
            f"TPU-{_MOCK_ACCELERATOR_TYPE}-head": 1,
        },
    }

    def test_command_run(self, environment_variables, **kwargs):
        resources = environment_variables[ray_constants.RESOURCES_ENVIRONMENT_VARIABLE]
        if self._worker_id == 0:
            assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" in resources
        else:
            assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" not in resources

    with patch.object(SSHCommandRunner, "run", new=test_command_run):
        cmd_runner = TPUCommandRunner(**args)
        cmd_runner.run(
            "echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
        )


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))