chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,329 @@
|
||||
"""Command runners specific to TPU VM pods.
|
||||
|
||||
TPU VM pods may contain multiple hosts, each including attached TPU chips and
|
||||
associated internal/external IP addresses.
|
||||
|
||||
To support TPU VM pods, we represent entire TPU pods as "Ray Nodes", meaning
|
||||
that TPU pods will need to run the operations specified in `CommandRunnerInterface`
|
||||
N times, where N denotes the number of hosts that comprise a TPU pod.
|
||||
|
||||
To maintain feature completeness, we simply wrap the existing `SSHCommandRunner` and
|
||||
`DockerCommandRunner` and run them as batched calls.
|
||||
|
||||
"""
|
||||
import copy
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import ModuleType
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray.autoscaler._private.command_runner import DockerCommandRunner, SSHCommandRunner
|
||||
from ray.autoscaler._private.gcp.node import GCPTPUNode
|
||||
from ray.autoscaler.command_runner import CommandRunnerInterface
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
|
||||
|
||||
class TPUVMSSHCommandRunner(SSHCommandRunner):
|
||||
"""An SSH command runner with overwritten IP address calls."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
internal_ip: str,
|
||||
external_ip: str,
|
||||
worker_id: int,
|
||||
accelerator_type: str,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self._internal_ip = internal_ip
|
||||
self._external_ip = external_ip
|
||||
self._worker_id = worker_id
|
||||
self._accelerator_type = accelerator_type
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def _get_node_ip(self) -> str:
|
||||
if self.use_internal_ip:
|
||||
return self._internal_ip
|
||||
else:
|
||||
return self._external_ip
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env="auto", # Unused argument.
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
) -> str:
|
||||
"""Override the SSH run for TPU VM pods.
|
||||
|
||||
Main functionality here we need to inject is to intercept the resources
|
||||
provided by the node_provider TPU node type fillout.
|
||||
|
||||
node_provider will provide a resource "TPU-{TPU_POD_TYPE}-head" which:
|
||||
1) allows application developers to target worker 0 of an arbitary TPU pod, and
|
||||
2) signals to the autoscaler how to address the demand for more TPU pods.
|
||||
|
||||
Without this intercept, then all workers of a TPU pod will have the
|
||||
"TPU-{TPU_POD_TYPE}-head" resource which will violate functionality (1)
|
||||
above.
|
||||
|
||||
"""
|
||||
|
||||
if environment_variables:
|
||||
environment_variables = _maybe_remove_head_resource(
|
||||
environment_variables, self._worker_id, self._accelerator_type
|
||||
)
|
||||
|
||||
return super().run(
|
||||
cmd=cmd,
|
||||
timeout=timeout,
|
||||
exit_on_fail=exit_on_fail,
|
||||
port_forward=port_forward,
|
||||
with_output=with_output,
|
||||
environment_variables=environment_variables,
|
||||
run_env=run_env,
|
||||
ssh_options_override_ssh_key=ssh_options_override_ssh_key,
|
||||
shutdown_after_run=shutdown_after_run,
|
||||
)
|
||||
|
||||
|
||||
class TPUVMDockerCommandRunner(DockerCommandRunner):
|
||||
"""A Docker command runner with overwritten IP addresses."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
docker_config: Dict[str, Any],
|
||||
internal_ip: str,
|
||||
external_ip: str,
|
||||
worker_id: int,
|
||||
accelerator_type: str,
|
||||
**common_args,
|
||||
):
|
||||
super().__init__(docker_config=docker_config, **common_args)
|
||||
self._worker_id = worker_id
|
||||
self._accelerator_type = accelerator_type
|
||||
|
||||
self.ssh_command_runner = TPUVMSSHCommandRunner(
|
||||
internal_ip=internal_ip,
|
||||
external_ip=external_ip,
|
||||
worker_id=worker_id,
|
||||
accelerator_type=accelerator_type,
|
||||
**common_args,
|
||||
)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Optional[Dict[str, object]] = None,
|
||||
run_env="auto",
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
):
|
||||
if environment_variables:
|
||||
environment_variables = _maybe_remove_head_resource(
|
||||
environment_variables, self._worker_id, self._accelerator_type
|
||||
)
|
||||
return super().run(
|
||||
cmd,
|
||||
timeout,
|
||||
exit_on_fail,
|
||||
port_forward,
|
||||
with_output,
|
||||
environment_variables,
|
||||
run_env,
|
||||
ssh_options_override_ssh_key,
|
||||
shutdown_after_run,
|
||||
)
|
||||
|
||||
|
||||
class TPUCommandRunner(CommandRunnerInterface):
|
||||
"""A TPU pod command runner."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
instance: GCPTPUNode,
|
||||
log_prefix: str,
|
||||
node_id: str,
|
||||
auth_config: Dict[str, Any],
|
||||
provider: NodeProvider,
|
||||
cluster_name: str,
|
||||
process_runner: ModuleType,
|
||||
use_internal_ip: bool,
|
||||
docker_config: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
def create_command_runner(
|
||||
worker_id: int, accelerator_type: str, internal_ip: str, external_ip: str
|
||||
) -> CommandRunnerInterface:
|
||||
"""Returns the correct base command runner."""
|
||||
|
||||
common_args = {
|
||||
"internal_ip": internal_ip,
|
||||
"external_ip": external_ip,
|
||||
"worker_id": worker_id,
|
||||
"accelerator_type": accelerator_type,
|
||||
"log_prefix": "[tpu_worker_{}] ".format(worker_id) + log_prefix,
|
||||
"node_id": node_id,
|
||||
"provider": provider,
|
||||
"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"] != "":
|
||||
return TPUVMDockerCommandRunner(
|
||||
docker_config=docker_config, **common_args
|
||||
)
|
||||
else:
|
||||
return TPUVMSSHCommandRunner(**common_args)
|
||||
|
||||
self._command_runners = []
|
||||
self._num_workers = instance.num_workers
|
||||
for i in range(self._num_workers):
|
||||
self._command_runners.append(
|
||||
create_command_runner(
|
||||
worker_id=i,
|
||||
accelerator_type=instance.get("acceleratorType"),
|
||||
internal_ip=instance.get_internal_ip(i),
|
||||
external_ip=instance.get_external_ip(i),
|
||||
)
|
||||
)
|
||||
|
||||
@property
|
||||
def num_connections(self) -> int:
|
||||
"""Return the number of active connections allowed at a time.
|
||||
|
||||
We occasionally see issues where too many concurrent connections may lead to
|
||||
failed SSH connections when there are too many TPU hosts.
|
||||
|
||||
We utilize this property to cap the maximum number of active connections
|
||||
at a time until a proper fix is found.
|
||||
|
||||
"""
|
||||
num_max_concurrent_active_connections = ray_constants.env_integer(
|
||||
ray_constants.RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR, default=16
|
||||
)
|
||||
return min(self._num_workers, num_max_concurrent_active_connections)
|
||||
|
||||
def run(
|
||||
self,
|
||||
cmd,
|
||||
timeout=120,
|
||||
exit_on_fail=False,
|
||||
port_forward=None,
|
||||
with_output=False,
|
||||
environment_variables: Dict[str, object] = None,
|
||||
run_env="auto", # Unused argument.
|
||||
ssh_options_override_ssh_key="",
|
||||
shutdown_after_run=False,
|
||||
) -> str:
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
results = executor.map(
|
||||
lambda i: self._command_runners[i].run(
|
||||
cmd=cmd,
|
||||
timeout=timeout,
|
||||
exit_on_fail=exit_on_fail,
|
||||
port_forward=port_forward,
|
||||
with_output=with_output,
|
||||
environment_variables=copy.deepcopy(environment_variables),
|
||||
run_env=run_env,
|
||||
ssh_options_override_ssh_key=ssh_options_override_ssh_key,
|
||||
shutdown_after_run=shutdown_after_run,
|
||||
),
|
||||
range(self._num_workers),
|
||||
)
|
||||
# Note: the `run` abstract function may return a string representing
|
||||
# representing command output, but this result is rarely used - especially
|
||||
# if the node is a worker (which a TPU pod is).
|
||||
# We return only the results from worker 0 which may not always be expected.
|
||||
return list(results)[0]
|
||||
|
||||
def run_rsync_up(self, *args, **kwargs) -> None:
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
executor.map(
|
||||
lambda i: self._command_runners[i].run_rsync_up(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
|
||||
def run_rsync_down(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Rsync files down from the cluster node.
|
||||
|
||||
Args:
|
||||
*args: Forwarded to each per-worker ``run_rsync_down`` call.
|
||||
Includes the (remote) source path and (local) target path.
|
||||
**kwargs: Forwarded to each per-worker ``run_rsync_down`` call.
|
||||
"""
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
executor.map(
|
||||
lambda i: self._command_runners[i].run_rsync_down(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
|
||||
def remote_shell_command_str(self) -> str:
|
||||
"""Return the command the user can use to open a shell."""
|
||||
# Note: this function is rarely used if the node is a worker.
|
||||
# We return only the results from worker 0 which may not always be expected.
|
||||
return self._command_runners[0].remote_shell_command_str()
|
||||
|
||||
def run_init(self, *args: Any, **kwargs: Any) -> Optional[bool]:
|
||||
"""Used to run extra initialization commands.
|
||||
|
||||
Args:
|
||||
*args: Forwarded to each per-worker ``run_init`` call. Includes
|
||||
``as_head``, ``file_mounts``, and ``sync_run_yet``.
|
||||
**kwargs: Forwarded to each per-worker ``run_init`` call.
|
||||
|
||||
Returns:
|
||||
Whether initialization is necessary on any worker.
|
||||
"""
|
||||
with ThreadPoolExecutor(self.num_connections) as executor:
|
||||
results = executor.map(
|
||||
lambda i: self._command_runners[i].run_init(*args, **kwargs),
|
||||
range(self._num_workers),
|
||||
)
|
||||
# Note: the `run_init` abstract function may return a bool representing
|
||||
# whether initialization is necessary, but this result is rarely used -
|
||||
# especially if the node is a worker (which a TPU pod is).
|
||||
# Here we return whether any workers require initialization, which may not be
|
||||
# the expected result.
|
||||
return any(results)
|
||||
|
||||
|
||||
def _maybe_remove_head_resource(
|
||||
environment_variables: Dict[str, Any], worker_id: int, accelerator_type: str
|
||||
):
|
||||
"""
|
||||
node_provider will provide a resource "TPU-{TPU_POD_TYPE}-head" which:
|
||||
1) allows application developers to target worker 0 of an arbitary TPU pod, and
|
||||
2) signals to the autoscaler how to address the demand for more TPU pods.
|
||||
|
||||
Without this intercept, then all workers of a TPU pod will have the
|
||||
"TPU-{TPU_POD_TYPE}-head" resource which will violate functionality (1)
|
||||
above.
|
||||
"""
|
||||
resources = environment_variables.get(
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE, None
|
||||
)
|
||||
|
||||
if resources:
|
||||
# For TPU pod support, we need to ensure that the
|
||||
# tpu pod resource type only propagates to worker 0.
|
||||
if worker_id != 0:
|
||||
tpu_pod_resource_type = f"TPU-{accelerator_type}-head"
|
||||
if tpu_pod_resource_type in resources:
|
||||
resources = copy.copy(resources)
|
||||
resources.pop(tpu_pod_resource_type, None)
|
||||
environment_variables = {
|
||||
**environment_variables,
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: resources,
|
||||
}
|
||||
|
||||
return environment_variables
|
||||
Reference in New Issue
Block a user