chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,290 @@
|
||||
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__]))
|
||||
Reference in New Issue
Block a user