wehub-resource-sync
134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from typing import List, Optional, Tuple
|
|
|
|
from ray._private.accelerators.accelerator import AcceleratorManager
|
|
from ray._private.ray_constants import env_bool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
NEURON_RT_VISIBLE_CORES_ENV_VAR = "NEURON_RT_VISIBLE_CORES"
|
|
NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR = (
|
|
"RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES"
|
|
)
|
|
|
|
# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/inf2-arch.html#aws-inf2-arch
|
|
# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/trn1-arch.html#aws-trn1-arch
|
|
# Subject to removal after the information is available via public API
|
|
AWS_NEURON_INSTANCE_MAP = {
|
|
"trn1.2xlarge": 2,
|
|
"trn1.32xlarge": 32,
|
|
"trn1n.32xlarge": 32,
|
|
"inf2.xlarge": 2,
|
|
"inf2.8xlarge": 2,
|
|
"inf2.24xlarge": 12,
|
|
"inf2.48xlarge": 24,
|
|
}
|
|
|
|
|
|
class NeuronAcceleratorManager(AcceleratorManager):
|
|
"""AWS Inferentia and Trainium accelerators."""
|
|
|
|
@staticmethod
|
|
def get_resource_name() -> str:
|
|
return "neuron_cores"
|
|
|
|
@staticmethod
|
|
def get_visible_accelerator_ids_env_var() -> str:
|
|
return NEURON_RT_VISIBLE_CORES_ENV_VAR
|
|
|
|
@staticmethod
|
|
def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
|
|
neuron_visible_cores = os.environ.get(
|
|
NeuronAcceleratorManager.get_visible_accelerator_ids_env_var(), None
|
|
)
|
|
|
|
if neuron_visible_cores is None:
|
|
return None
|
|
|
|
if neuron_visible_cores == "":
|
|
return []
|
|
|
|
return list(neuron_visible_cores.split(","))
|
|
|
|
@staticmethod
|
|
def get_current_node_num_accelerators() -> int:
|
|
"""
|
|
Attempt to detect the number of Neuron cores on this machine.
|
|
|
|
Returns:
|
|
The number of Neuron cores if any were detected, otherwise 0.
|
|
"""
|
|
nc_count: int = 0
|
|
neuron_path = "/opt/aws/neuron/bin/"
|
|
if sys.platform.startswith("linux") and os.path.isdir(neuron_path):
|
|
result = subprocess.run(
|
|
[os.path.join(neuron_path, "neuron-ls"), "--json-output"],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
if result.returncode == 0 and result.stdout:
|
|
neuron_devices = json.loads(result.stdout)
|
|
for neuron_device in neuron_devices:
|
|
nc_count += neuron_device.get("nc_count", 0)
|
|
return nc_count
|
|
|
|
@staticmethod
|
|
def get_current_node_accelerator_type() -> Optional[str]:
|
|
from ray.util.accelerators import AWS_NEURON_CORE
|
|
|
|
return AWS_NEURON_CORE
|
|
|
|
@staticmethod
|
|
def validate_resource_request_quantity(
|
|
quantity: float,
|
|
) -> Tuple[bool, Optional[str]]:
|
|
if isinstance(quantity, float) and not quantity.is_integer():
|
|
return (
|
|
False,
|
|
f"{NeuronAcceleratorManager.get_resource_name()} resource quantity"
|
|
" must be whole numbers. "
|
|
f"The specified quantity {quantity} is invalid.",
|
|
)
|
|
else:
|
|
return (True, None)
|
|
|
|
@staticmethod
|
|
def set_current_process_visible_accelerator_ids(
|
|
visible_neuron_core_ids: List[str],
|
|
) -> None:
|
|
"""Set the NEURON_RT_VISIBLE_CORES environment variable based on
|
|
given visible_neuron_core_ids.
|
|
|
|
Args:
|
|
visible_neuron_core_ids: List of str representing core IDs.
|
|
"""
|
|
if env_bool(NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR, False):
|
|
return
|
|
|
|
os.environ[
|
|
NeuronAcceleratorManager.get_visible_accelerator_ids_env_var()
|
|
] = ",".join([str(i) for i in visible_neuron_core_ids])
|
|
|
|
@staticmethod
|
|
def get_ec2_instance_num_accelerators(
|
|
instance_type: str, instances: dict
|
|
) -> Optional[int]:
|
|
# TODO: AWS SDK (public API) doesn't yet expose the NeuronCore
|
|
# information. It will be available (work-in-progress)
|
|
# as xxAcceleratorInfo in InstanceTypeInfo.
|
|
# https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_InstanceTypeInfo.html
|
|
# See https://github.com/ray-project/ray/issues/38473
|
|
return AWS_NEURON_INSTANCE_MAP.get(instance_type.lower(), None)
|
|
|
|
@staticmethod
|
|
def get_ec2_instance_accelerator_type(
|
|
instance_type: str, instances: dict
|
|
) -> Optional[str]:
|
|
from ray.util.accelerators import AWS_NEURON_CORE
|
|
|
|
return AWS_NEURON_CORE
|