diff --git a/lib/python3.12/site-packages/ray/air/__init__.py b/lib/python3.12/site-packages/ray/air/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..67f41865b669ca092c3a4f63791375a1b1b3f25b --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/__init__.py @@ -0,0 +1,21 @@ +from ray.air.config import ( + CheckpointConfig, + FailureConfig, + RunConfig, + ScalingConfig, +) +from ray.air.data_batch_type import DataBatchType +from ray.air.execution.resources.request import AcquiredResources, ResourceRequest +from ray.air.result import Result +import ray.data # noqa: F401 # TODO: This is a hack to avoid circular import + +__all__ = [ + "DataBatchType", + "RunConfig", + "Result", + "ScalingConfig", + "FailureConfig", + "CheckpointConfig", + "AcquiredResources", + "ResourceRequest", +] diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5afdd99cba222933b4a4e16725d1aeff1de587a Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..866e9ee69124b6a1f2de38cef0a07392be2d5f6b Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5b5c30b4efdc66e598fdc7badafb3013f091ae1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/data_batch_type.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/data_batch_type.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a058fcce373b7edad50f69654ef92c24801451d0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/data_batch_type.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/result.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/result.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57d9c02419741fc888873a6c35b2c74c9071f48d Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/result.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/__pycache__/session.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/__pycache__/session.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cf113547a53b1ae50f90bc36415742009a11dd4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/__pycache__/session.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__init__.py b/lib/python3.12/site-packages/ray/air/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..02d9328cba6bf0e165815a63b760fde02ea32e30 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bb6d26e6cd92073b0a40a4f7354cee736d8cf02 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/filelock.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/filelock.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18ac4c62349c45a78e5e84cc3ab5d9fa58d6190e Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/filelock.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/json.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/json.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..109b256b661ceab48cfd48dc28b50cce64218896 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/json.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/mlflow.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/mlflow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86b6c71daed497b61d8e6b8f88854b939a6e8d71 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/mlflow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/tensorflow_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/tensorflow_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c48ff58460aaabbb9a6c6de85f51ca9ae92cd878 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/tensorflow_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/torch_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/torch_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3548cf51f99d27da06abdde53ef577051f2d0ab Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/torch_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/uri_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/uri_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8e06de81040560d2c47e4da28073a1096bcc854 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/uri_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/usage.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/usage.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17a644d22af782e6c776b1cab17b2636469634fa Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/usage.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55ddc8771288cbcc309a9e9634943647ac69c463 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/config.py b/lib/python3.12/site-packages/ray/air/_internal/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7518c7ccc0411b7f39781dc3497f24a566e5bb5c --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/config.py @@ -0,0 +1,47 @@ +import dataclasses +from typing import Iterable + + +def ensure_only_allowed_dataclass_keys_updated( + dataclass: dataclasses.dataclass, + allowed_keys: Iterable[str], +): + """ + Validate dataclass by raising an exception if any key not included in + ``allowed_keys`` differs from the default value. + + A ``ValueError`` will also be raised if any of the ``allowed_keys`` + is not present in ``dataclass.__dict__``. + + Args: + dataclass: Dict or dataclass to check. + allowed_keys: dataclass attribute keys that can have a value different than + the default one. + """ + default_data = dataclass.__class__() + + allowed_keys = set(allowed_keys) + + # TODO: split keys_not_in_dict validation to a separate function. + keys_not_in_dict = [key for key in allowed_keys if key not in default_data.__dict__] + if keys_not_in_dict: + raise ValueError( + f"Key(s) {keys_not_in_dict} are not present in " + f"{dataclass.__class__.__name__}. " + "Remove them from `allowed_keys`. " + f"Valid keys: {list(default_data.__dict__.keys())}" + ) + + # These keys should not have been updated in the `dataclass` object + prohibited_keys = set(default_data.__dict__) - allowed_keys + + bad_keys = [ + key + for key in prohibited_keys + if dataclass.__dict__[key] != default_data.__dict__[key] + ] + if bad_keys: + raise ValueError( + f"Key(s) {bad_keys} are not allowed to be updated in the current context. " + "Remove them from the dataclass." + ) diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__init__.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..824833df03ffe9e03f583856b9336d7610386d91 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__init__.py @@ -0,0 +1,92 @@ +import logging +import threading +from typing import Optional + +import ray +import ray._private.ray_constants as ray_constants +from ray.air._internal.device_manager.cpu import CPUTorchDeviceManager +from ray.air._internal.device_manager.hpu import HPUTorchDeviceManager +from ray.air._internal.device_manager.npu import NPUTorchDeviceManager +from ray.air._internal.device_manager.nvidia_gpu import CUDATorchDeviceManager +from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager + +logger = logging.getLogger(__name__) + + +DEFAULT_TORCH_DEVICE_MANAGER_CLS = CPUTorchDeviceManager + + +SUPPORTED_ACCELERATOR_TORCH_DEVICE_MANAGER = { + ray_constants.GPU: CUDATorchDeviceManager, + ray_constants.HPU: HPUTorchDeviceManager, + ray_constants.NPU: NPUTorchDeviceManager, +} + + +def register_custom_torch_dist_backend(backend: Optional[str] = None) -> None: + if backend == "hccl": + # The name for the communication backend of Habana and torch-npu is the same. + HPUTorchDeviceManager.register_custom_torch_dist_backend() + + NPUTorchDeviceManager.register_custom_torch_dist_backend() + + +_torch_device_manager = None +_torch_device_manager_lock = threading.Lock() + + +def get_torch_device_manager_by_context() -> TorchDeviceManager: + global _torch_device_manager + + with _torch_device_manager_lock: + if not _torch_device_manager: + existing_device_manager_cls = None + resources = ray.get_runtime_context().get_accelerator_ids() + + # select correct accelerator type from resources + for resource_type, resource_value in resources.items(): + device_manager_cls = SUPPORTED_ACCELERATOR_TORCH_DEVICE_MANAGER.get( + resource_type, None + ) + if resource_value and device_manager_cls: + # An error will raise when multiple accelerators are specified. + if existing_device_manager_cls: + raise RuntimeError( + "Unable to determine the appropriate DeviceManager " + f"for the specified resources {resources}." + ) + else: + existing_device_manager_cls = device_manager_cls + + device_manager_cls = ( + existing_device_manager_cls or DEFAULT_TORCH_DEVICE_MANAGER_CLS + ) + + _torch_device_manager = device_manager_cls() + + return _torch_device_manager + + +def get_torch_device_manager_by_device_type(device_type: str): + if device_type.lower() == ray_constants.GPU.lower() or device_type == "cuda": + return CUDATorchDeviceManager() + elif device_type.lower() == ray_constants.NPU.lower(): + return NPUTorchDeviceManager() + elif device_type.lower() == ray_constants.HPU.lower(): + return HPUTorchDeviceManager() + elif device_type.lower() == "cpu": + return CPUTorchDeviceManager() + + raise RuntimeError(f"Device type {device_type} cannot be recognized.") + + +__all__ = [ + TorchDeviceManager, + CPUTorchDeviceManager, + CUDATorchDeviceManager, + HPUTorchDeviceManager, + NPUTorchDeviceManager, + register_custom_torch_dist_backend, + get_torch_device_manager_by_context, + get_torch_device_manager_by_device_type, +] diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c194ee313c6fcd7ce3d552be8f16cdd547ec412 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/cpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/cpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f683ffe1dc45d135e3094f971aec89bff9483124 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/cpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/hpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/hpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0064df4c760dea68b5559b5a06d458db846eb4b Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/hpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/npu.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/npu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0dfcd4093af9bcf67dfc29a091507672c597f28 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/npu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/nvidia_gpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/nvidia_gpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b346cf4f28c36169f0bdfd3addd0e0daa8013a7d Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/nvidia_gpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/torch_device_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/torch_device_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40997fdd95ad9912549144be235942fbc1fc4073 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/_internal/device_manager/__pycache__/torch_device_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/cpu.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..76fa73765287c6bf68a0b2f6e7d4130297b799df --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/cpu.py @@ -0,0 +1,30 @@ +from contextlib import contextmanager +from typing import List + +import torch + +from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager + + +class CPUTorchDeviceManager(TorchDeviceManager): + """CPU device manager""" + + def is_available(self) -> bool(): + return True + + def get_devices(self) -> List[torch.device]: + """Gets the correct torch device list configured for this process.""" + return [torch.device("cpu")] + + def supports_stream(self) -> bool: + """Validate if the device type support create a stream""" + return False + + def get_stream_context(self, stream): + """Return empty context mananger for CPU.""" + + @contextmanager + def default_context_manager(): + yield + + return default_context_manager() diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/hpu.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/hpu.py new file mode 100644 index 0000000000000000000000000000000000000000..bb402ea65b0d9378cb10d70b50bab3588be79102 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/hpu.py @@ -0,0 +1,50 @@ +from contextlib import contextmanager +from typing import List, Union + +import torch + +from ray._private.accelerators.hpu import HPU_PACKAGE_AVAILABLE +from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager + +if HPU_PACKAGE_AVAILABLE: + import habana_frameworks.torch.hpu as torch_hpu + + +class HPUTorchDeviceManager(TorchDeviceManager): + """HPU device manager""" + + @staticmethod + def register_custom_torch_dist_backend(): + if HPU_PACKAGE_AVAILABLE: + import habana_frameworks.torch.core # noqa: F401 + import habana_frameworks.torch.distributed.hccl # noqa: F401 + + def is_available(self) -> bool(): + if not HPU_PACKAGE_AVAILABLE: + return False + + return torch_hpu.is_available() + + def get_devices(self) -> List[torch.device]: + if not self.is_available(): + raise RuntimeError( + "Using HPUTorchDeviceManager but torch hpu is not available." + ) + + return [torch.device("hpu")] + + def set_device(self, device: Union[torch.device, int, str, None]): + torch_hpu.set_device(device) + + def supports_stream(self) -> bool: + """Validate if the device type support create a stream""" + return False + + def get_stream_context(self, stream): + """Get HPU stream context manager, empty so far.""" + + @contextmanager + def default_context_manager(): + yield + + return default_context_manager() diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/npu.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/npu.py new file mode 100644 index 0000000000000000000000000000000000000000..0a40594e14f1e97b6e77bdede6da03bd9ea92c2d --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/npu.py @@ -0,0 +1,104 @@ +import os +from importlib.util import find_spec +from typing import List, Union + +import torch + +import ray +import ray._private.ray_constants as ray_constants +from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR +from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager + + +def is_package_present(package_name: str) -> bool: + try: + return find_spec(package_name) is not None + except ModuleNotFoundError: + return False + + +NPU_TORCH_PACKAGE_AVAILABLE = is_package_present("torch_npu") + + +if NPU_TORCH_PACKAGE_AVAILABLE: + import torch_npu # noqa: F401 + + +class NPUTorchDeviceManager(TorchDeviceManager): + """Ascend NPU device manager""" + + @staticmethod + def register_custom_torch_dist_backend(): + if NPU_TORCH_PACKAGE_AVAILABLE: + import torch_npu # noqa: F401, F811 + + def is_available(self) -> bool: + if not NPU_TORCH_PACKAGE_AVAILABLE: + return False + + return torch.npu.is_available() + + def get_devices(self) -> List[torch.device]: + """Gets the correct torch device list configured for this process. + + Returns a list of torch NPU devices allocated for the current worker. + If no NPUs are assigned, then it returns a list with a single CPU device. + """ + if NPU_TORCH_PACKAGE_AVAILABLE and torch.npu.is_available(): + npu_ids = [ + str(id) + for id in ray.get_runtime_context().get_accelerator_ids()[ + ray_constants.NPU + ] + ] + + device_ids = [] + + if len(npu_ids) > 0: + npu_visible_str = os.environ.get(ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, "") + if npu_visible_str and npu_visible_str != "NoDevFiles": + npu_visible_list = npu_visible_str.split(",") + else: + npu_visible_list = [] + + for npu_id in npu_ids: + try: + device_ids.append(npu_visible_list.index(npu_id)) + except IndexError: + raise RuntimeError( + "ASCEND_RT_VISIBLE_DEVICES set incorrectly. " + f"Got {npu_visible_str}, expected to include {npu_id}. " + "Did you override the `ASCEND_RT_VISIBLE_DEVICES` " + "environment variable?" + ) + else: + # If called on the driver or outside of Ray Train, return the + # 0th device. + device_ids.append(0) + + devices = [torch.device(f"npu:{device_id}") for device_id in device_ids] + else: + raise RuntimeError( + "Using NPUTorchDeviceManager but torch npu is not available." + ) + + return devices + + def set_device(self, device: Union[torch.device, int]): + torch.npu.set_device(device) + + def supports_stream(self) -> bool: + """Validate if the device type support to create a stream""" + return True + + def create_stream(self, device): + """Create a stream on NPU device""" + return torch.npu.Stream(device) + + def get_stream_context(self, stream): + """Get a torch.stream context on NPU device""" + return torch.npu.stream(stream) + + def get_current_stream(self): + """Get current stream for NPU device""" + return torch.npu.current_stream() diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/nvidia_gpu.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/nvidia_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..f4bb1b54097e153cc7784c849c46331478d87988 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/nvidia_gpu.py @@ -0,0 +1,79 @@ +import os +from typing import List, Union + +import torch + +import ray +from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager + + +class CUDATorchDeviceManager(TorchDeviceManager): + """CUDA device manager""" + + def is_available(self) -> bool(): + return torch.cuda.is_available() + + def get_devices(self) -> List[torch.device]: + """Gets the correct torch device list configured for this process. + + Returns a list of torch CUDA devices allocated for the current worker. + If no GPUs are assigned, then it returns a list with a single CPU device. + + Assumes that `CUDA_VISIBLE_DEVICES` is set and is a + superset of the `ray.get_gpu_ids()`. + """ + + # GPU IDs are assigned by Ray after you specify "use_gpu" + # GPU `ray.get_gpu_ids()` may return ints or may return strings. + # We should always convert to strings. + gpu_ids = [str(id) for id in ray.get_gpu_ids()] + + device_ids = [] + + if len(gpu_ids) > 0: + cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if cuda_visible_str and cuda_visible_str != "NoDevFiles": + cuda_visible_list = cuda_visible_str.split(",") + else: + cuda_visible_list = [] + + # By default, there should only be one GPU ID if `use_gpu=True`. + # If there are multiple GPUs, return a list of devices. + # If using fractional GPUs, these IDs are not guaranteed + # to be unique across different processes. + for gpu_id in gpu_ids: + try: + device_ids.append(cuda_visible_list.index(gpu_id)) + except IndexError: + raise RuntimeError( + "CUDA_VISIBLE_DEVICES set incorrectly. " + f"Got {cuda_visible_str}, expected to include {gpu_id}. " + "Did you override the `CUDA_VISIBLE_DEVICES` environment" + " variable? If not, please help file an issue on Github." + ) + + else: + # If called on the driver or outside of Ray Train, return the + # 0th device. + device_ids.append(0) + + return [torch.device(f"cuda:{device_id}") for device_id in device_ids] + + def set_device(self, device: Union[torch.device, int, str, None]): + torch.cuda.set_device(device) + + def supports_stream(self) -> bool: + """Validate if the device type support create a stream""" + return True + + def create_stream(self, device: torch.device) -> torch.cuda.Stream: + """Create a stream on cuda device""" + return torch.cuda.Stream(device) + + def get_stream_context(self, stream): + """Get a stream context for cuda device""" + return torch.cuda.stream(stream) + + def get_current_stream(self) -> torch.cuda.Stream: + """Get current stream for cuda device""" + return torch.cuda.current_stream() diff --git a/lib/python3.12/site-packages/ray/air/_internal/device_manager/torch_device_manager.py b/lib/python3.12/site-packages/ray/air/_internal/device_manager/torch_device_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d522a477ef58a5021300eab3415abe99d3079213 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/device_manager/torch_device_manager.py @@ -0,0 +1,40 @@ +from abc import ABC +from typing import List, Union + +import torch + + +class TorchDeviceManager(ABC): + """This class contains the function needed for supporting + an acclerator family in Ray AI Library. + """ + + def is_available(self) -> bool: + """Validate if device is available.""" + ... + + def get_devices(self) -> List[torch.device]: + """Gets the correct torch device configured for this process""" + ... + + def set_device(self, device: Union[torch.device, int, str, None]): + """Set the correct device for this process""" + ... + + def supports_stream(self) -> bool: + """Validate if the device type support create a stream""" + ... + + def create_stream(self, device: torch.device): + """Create a device stream""" + ... + + def get_stream_context(self, stream): + """Get a stream context of device. If device didn't support stream, + this should return a empty context manager instead of None. + """ + ... + + def get_current_stream(self): + """Get current stream on accelerators like torch.cuda.current_stream""" + ... diff --git a/lib/python3.12/site-packages/ray/air/_internal/filelock.py b/lib/python3.12/site-packages/ray/air/_internal/filelock.py new file mode 100644 index 0000000000000000000000000000000000000000..832cc37ee4506fdc1b4d286299b355e09bfecbd9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/filelock.py @@ -0,0 +1,46 @@ +import hashlib +import os +from pathlib import Path + +from filelock import FileLock + +import ray + +RAY_LOCKFILE_DIR = "_ray_lockfiles" + + +class TempFileLock: + """FileLock wrapper that uses temporary file locks. + + The temporary directory that these locks are saved to can be configured via + the `RAY_TMPDIR` environment variable. + + Args: + path: The file path that this temporary file lock is used for. + This will be used to generate the lockfile filename. + Ex: For concurrent writes to a file, this is the common filepath + that multiple processes are writing to. + **kwargs: Additional keyword arguments to pass to the underlying `FileLock`. + """ + + def __init__(self, path: str, **kwargs): + self.path = path + temp_dir = Path(ray._common.utils.get_user_temp_dir()).resolve() + self._lock_dir = temp_dir / RAY_LOCKFILE_DIR + self._path_hash = hashlib.sha1( + str(Path(self.path).resolve()).encode("utf-8") + ).hexdigest() + self._lock_path = self._lock_dir / f"{self._path_hash}.lock" + + os.makedirs(str(self._lock_dir), exist_ok=True) + self._lock = FileLock(self._lock_path, **kwargs) + + def __enter__(self): + self._lock.acquire() + return self + + def __exit__(self, type, value, traceback): + self._lock.release() + + def __getattr__(self, name): + return getattr(self._lock, name) diff --git a/lib/python3.12/site-packages/ray/air/_internal/json.py b/lib/python3.12/site-packages/ray/air/_internal/json.py new file mode 100644 index 0000000000000000000000000000000000000000..2e88824e7109106a54020e3d5eec2629043f3ee1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/json.py @@ -0,0 +1,31 @@ +import json +import numbers + +import numpy as np + + +class SafeFallbackEncoder(json.JSONEncoder): + def __init__(self, nan_str="null", **kwargs): + super(SafeFallbackEncoder, self).__init__(**kwargs) + self.nan_str = nan_str + + def default(self, value): + try: + if type(value).__module__ == np.__name__ and isinstance(value, np.ndarray): + return value.tolist() + + if isinstance(value, np.bool_): + return bool(value) + + if np.isnan(value): + return self.nan_str + + if issubclass(type(value), numbers.Integral): + return int(value) + if issubclass(type(value), numbers.Number): + return float(value) + + return super(SafeFallbackEncoder, self).default(value) + + except Exception: + return str(value) # give up, just stringify it (ok for logs) diff --git a/lib/python3.12/site-packages/ray/air/_internal/mlflow.py b/lib/python3.12/site-packages/ray/air/_internal/mlflow.py new file mode 100644 index 0000000000000000000000000000000000000000..727318ce839ff0bf0a78bca5346c66e44b245b75 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/mlflow.py @@ -0,0 +1,342 @@ +import logging +import os +from copy import deepcopy +from typing import TYPE_CHECKING, Dict, Optional + +from packaging import version + +from ray._private.dict import flatten_dict + +if TYPE_CHECKING: + from mlflow.entities import Run + from mlflow.tracking import MlflowClient + +logger = logging.getLogger(__name__) + + +class _MLflowLoggerUtil: + """Util class for setting up and logging to MLflow. + + Use this util for any library that needs MLflow logging/tracking logic + such as Ray Tune or Ray Train. + """ + + def __init__(self): + import mlflow + + self._mlflow = mlflow + self.experiment_id = None + + def __deepcopy__(self, memo=None): + # mlflow is a module, and thus cannot be copied + _mlflow = self._mlflow + self.__dict__.pop("_mlflow") + dict_copy = deepcopy(self.__dict__, memo) + copied_object = _MLflowLoggerUtil() + copied_object.__dict__.update(dict_copy) + self._mlflow = _mlflow + copied_object._mlflow = _mlflow + return copied_object + + def setup_mlflow( + self, + tracking_uri: Optional[str] = None, + registry_uri: Optional[str] = None, + experiment_id: Optional[str] = None, + experiment_name: Optional[str] = None, + tracking_token: Optional[str] = None, + artifact_location: Optional[str] = None, + create_experiment_if_not_exists: bool = True, + ): + """ + Sets up MLflow. + + Sets the Mlflow tracking uri & token, and registry URI. Also sets + the MLflow experiment that the logger should use, and possibly + creates new experiment if it does not exist. + + Args: + tracking_uri: The tracking URI for the MLflow tracking + server. + registry_uri: The registry URI for the MLflow model registry. + experiment_id: The id of an already existing MLflow + experiment to use for logging. If None is passed in + here and the MFLOW_EXPERIMENT_ID is not set, or the + experiment with this id does not exist, + ``experiment_name`` will be used instead. This argument takes + precedence over ``experiment_name`` if both are passed in. + experiment_name: The experiment name to use for logging. + If None is passed in here, the MLFLOW_EXPERIMENT_NAME environment + variable is used to determine the experiment name. + If the experiment with the name already exists with MLflow, + it will be reused. If not, a new experiment will be created + with the provided name if + ``create_experiment_if_not_exists`` is set to True. + artifact_location: The location to store run artifacts. + If not provided, MLFlow picks an appropriate default. + Ignored if experiment already exists. + tracking_token: Tracking token used to authenticate with MLflow. + create_experiment_if_not_exists: Whether to create an + experiment with the provided name if it does not already + exist. Defaults to True. + + Returns: + Whether setup is successful. + """ + if tracking_token: + os.environ["MLFLOW_TRACKING_TOKEN"] = tracking_token + + self._mlflow.set_tracking_uri(tracking_uri) + self._mlflow.set_registry_uri(registry_uri) + + # First check experiment_id. + experiment_id = ( + experiment_id + if experiment_id is not None + else os.environ.get("MLFLOW_EXPERIMENT_ID") + ) + if experiment_id is not None: + from mlflow.exceptions import MlflowException + + try: + self._mlflow.get_experiment(experiment_id=experiment_id) + logger.debug( + f"Experiment with provided id {experiment_id} " + "exists. Setting that as the experiment." + ) + self.experiment_id = experiment_id + return + except MlflowException: + pass + + # Then check experiment_name. + experiment_name = ( + experiment_name + if experiment_name is not None + else os.environ.get("MLFLOW_EXPERIMENT_NAME") + ) + if experiment_name is not None and self._mlflow.get_experiment_by_name( + name=experiment_name + ): + logger.debug( + f"Experiment with provided name {experiment_name} " + "exists. Setting that as the experiment." + ) + self.experiment_id = self._mlflow.get_experiment_by_name( + experiment_name + ).experiment_id + return + + # An experiment with the provided id or name does not exist. + # Create a new experiment if applicable. + if experiment_name and create_experiment_if_not_exists: + logger.debug( + "Existing experiment not found. Creating new " + f"experiment with name: {experiment_name}" + ) + self.experiment_id = self._mlflow.create_experiment( + name=experiment_name, artifact_location=artifact_location + ) + return + + if create_experiment_if_not_exists: + raise ValueError( + f"Experiment with the provided experiment_id: " + f"{experiment_id} does not exist and no " + f"experiment_name provided. At least one of " + f"these has to be provided." + ) + else: + raise ValueError( + f"Experiment with the provided experiment_id: " + f"{experiment_id} or experiment_name: " + f"{experiment_name} does not exist. Please " + f"create an MLflow experiment and provide " + f"either its id or name." + ) + + def _parse_dict(self, dict_to_log: Dict) -> Dict: + """Parses provided dict to convert all values to float. + + MLflow can only log metrics that are floats. This does not apply to + logging parameters or artifacts. + + Args: + dict_to_log: The dictionary containing the metrics to log. + + Returns: + A dictionary containing the metrics to log with all values being + converted to floats, or skipped if not able to be converted. + """ + new_dict = {} + for key, value in dict_to_log.items(): + try: + value = float(value) + new_dict[key] = value + except (ValueError, TypeError): + logger.debug( + "Cannot log key {} with value {} since the " + "value cannot be converted to float.".format(key, value) + ) + continue + + return new_dict + + def start_run( + self, + run_name: Optional[str] = None, + tags: Optional[Dict] = None, + set_active: bool = False, + ) -> "Run": + """Starts a new run and possibly sets it as the active run. + + Args: + tags: Tags to set for the new run. + set_active: Whether to set the new run as the active run. + If an active run already exists, then that run is returned. + + Returns: + The newly created MLflow run. + """ + import mlflow + from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME + + if tags is None: + tags = {} + + if set_active: + return self._start_active_run(run_name=run_name, tags=tags) + + client = self._get_client() + # If `mlflow==1.30.0` and we don't use `run_name`, then MLflow might error. For + # more information, see #29749. + if version.parse(mlflow.__version__) >= version.parse("1.30.0"): + run = client.create_run( + run_name=run_name, experiment_id=self.experiment_id, tags=tags + ) + else: + tags[MLFLOW_RUN_NAME] = run_name + run = client.create_run(experiment_id=self.experiment_id, tags=tags) + + return run + + def _start_active_run( + self, run_name: Optional[str] = None, tags: Optional[Dict] = None + ) -> "Run": + """Starts a run and sets it as the active run if one does not exist. + + If an active run already exists, then returns it. + """ + active_run = self._mlflow.active_run() + if active_run: + return active_run + + return self._mlflow.start_run( + run_name=run_name, experiment_id=self.experiment_id, tags=tags + ) + + def _run_exists(self, run_id: str) -> bool: + """Check if run with the provided id exists.""" + from mlflow.exceptions import MlflowException + + try: + self._mlflow.get_run(run_id=run_id) + return True + except MlflowException: + return False + + def _get_client(self) -> "MlflowClient": + """Returns an ml.tracking.MlflowClient instance to use for logging.""" + tracking_uri = self._mlflow.get_tracking_uri() + registry_uri = self._mlflow.get_registry_uri() + + from mlflow.tracking import MlflowClient + + return MlflowClient(tracking_uri=tracking_uri, registry_uri=registry_uri) + + def log_params(self, params_to_log: Dict, run_id: Optional[str] = None): + """Logs the provided parameters to the run specified by run_id. + + If no ``run_id`` is passed in, then logs to the current active run. + If there is not active run, then creates a new run and sets it as + the active run. + + Args: + params_to_log: Dictionary of parameters to log. + run_id (Optional[str]): The ID of the run to log to. + """ + params_to_log = flatten_dict(params_to_log) + + if run_id and self._run_exists(run_id): + client = self._get_client() + for key, value in params_to_log.items(): + client.log_param(run_id=run_id, key=key, value=value) + + else: + for key, value in params_to_log.items(): + self._mlflow.log_param(key=key, value=value) + + def log_metrics(self, step, metrics_to_log: Dict, run_id: Optional[str] = None): + """Logs the provided metrics to the run specified by run_id. + + + If no ``run_id`` is passed in, then logs to the current active run. + If there is not active run, then creates a new run and sets it as + the active run. + + Args: + metrics_to_log: Dictionary of metrics to log. + run_id (Optional[str]): The ID of the run to log to. + """ + metrics_to_log = flatten_dict(metrics_to_log) + metrics_to_log = self._parse_dict(metrics_to_log) + + if run_id and self._run_exists(run_id): + client = self._get_client() + for key, value in metrics_to_log.items(): + client.log_metric(run_id=run_id, key=key, value=value, step=step) + + else: + for key, value in metrics_to_log.items(): + self._mlflow.log_metric(key=key, value=value, step=step) + + def save_artifacts(self, dir: str, run_id: Optional[str] = None): + """Saves directory as artifact to the run specified by run_id. + + If no ``run_id`` is passed in, then saves to the current active run. + If there is not active run, then creates a new run and sets it as + the active run. + + Args: + dir: Path to directory containing the files to save. + run_id (Optional[str]): The ID of the run to log to. + """ + if run_id and self._run_exists(run_id): + client = self._get_client() + client.log_artifacts(run_id=run_id, local_dir=dir) + else: + self._mlflow.log_artifacts(local_dir=dir) + + def end_run(self, status: Optional[str] = None, run_id=None): + """Terminates the run specified by run_id. + + If no ``run_id`` is passed in, then terminates the + active run if one exists. + + Args: + status (Optional[str]): The status to set when terminating the run. + run_id (Optional[str]): The ID of the run to terminate. + + """ + if ( + run_id + and self._run_exists(run_id) + and not ( + self._mlflow.active_run() + and self._mlflow.active_run().info.run_id == run_id + ) + ): + client = self._get_client() + client.set_terminated(run_id=run_id, status=status) + else: + self._mlflow.end_run(status=status) diff --git a/lib/python3.12/site-packages/ray/air/_internal/tensorflow_utils.py b/lib/python3.12/site-packages/ray/air/_internal/tensorflow_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..46b3c2d1d1b7d5a4b86ffdf46fa6cd8598cca63a --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/tensorflow_utils.py @@ -0,0 +1,137 @@ +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import numpy as np +import pyarrow +import tensorflow as tf + +from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed +from ray.air.util.tensor_extensions.arrow import get_arrow_extension_tensor_types + +if TYPE_CHECKING: + from ray.data._internal.pandas_block import PandasBlockSchema + + +def convert_ndarray_to_tf_tensor( + ndarray: np.ndarray, + dtype: Optional[tf.dtypes.DType] = None, + type_spec: Optional[tf.TypeSpec] = None, +) -> tf.Tensor: + """Convert a NumPy ndarray to a TensorFlow Tensor. + + Args: + ndarray: A NumPy ndarray that we wish to convert to a TensorFlow Tensor. + dtype: A TensorFlow dtype for the created tensor; if None, the dtype will be + inferred from the NumPy ndarray data. + type_spec: A type spec that specifies the shape and dtype of the returned + tensor. If you specify ``dtype``, the dtype stored in the type spec is + ignored. + + Returns: A TensorFlow Tensor. + """ + if dtype is None and type_spec is not None: + dtype = type_spec.dtype + + is_ragged = isinstance(type_spec, tf.RaggedTensorSpec) + ndarray = _unwrap_ndarray_object_type_if_needed(ndarray) + if is_ragged: + return tf.ragged.constant(ndarray, dtype=dtype) + else: + return tf.convert_to_tensor(ndarray, dtype=dtype) + + +def convert_ndarray_batch_to_tf_tensor_batch( + ndarrays: Union[np.ndarray, Dict[str, np.ndarray]], + dtypes: Optional[Union[tf.dtypes.DType, Dict[str, tf.dtypes.DType]]] = None, +) -> Union[tf.Tensor, Dict[str, tf.Tensor]]: + """Convert a NumPy ndarray batch to a TensorFlow Tensor batch. + + Args: + ndarray: A (dict of) NumPy ndarray(s) that we wish to convert to a TensorFlow + Tensor. + dtype: A (dict of) TensorFlow dtype(s) for the created tensor; if None, the + dtype will be inferred from the NumPy ndarray data. + + Returns: A (dict of) TensorFlow Tensor(s). + """ + if isinstance(ndarrays, np.ndarray): + # Single-tensor case. + if isinstance(dtypes, dict): + if len(dtypes) != 1: + raise ValueError( + "When constructing a single-tensor batch, only a single dtype " + f"should be given, instead got: {dtypes}" + ) + dtypes = next(iter(dtypes.values())) + batch = convert_ndarray_to_tf_tensor(ndarrays, dtypes) + else: + # Multi-tensor case. + batch = { + col_name: convert_ndarray_to_tf_tensor( + col_ndarray, + dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, + ) + for col_name, col_ndarray in ndarrays.items() + } + + return batch + + +def get_type_spec( + schema: Union["pyarrow.lib.Schema", "PandasBlockSchema"], + columns: Union[str, List[str]], +) -> Union[tf.TypeSpec, Dict[str, tf.TypeSpec]]: + import pyarrow as pa + + from ray.data.extensions import TensorDtype + + tensor_extension_types = get_arrow_extension_tensor_types() + + assert not isinstance(schema, type) + + dtypes: Dict[str, Union[np.dtype, pa.DataType]] = dict( + zip(schema.names, schema.types) + ) + + def get_dtype(dtype: Union[np.dtype, pa.DataType]) -> tf.dtypes.DType: + if isinstance(dtype, pa.ListType): + dtype = dtype.value_type + if isinstance(dtype, pa.DataType): + dtype = dtype.to_pandas_dtype() + if isinstance(dtype, TensorDtype): + dtype = dtype.element_dtype + res = tf.dtypes.as_dtype(dtype) + return res + + def get_shape(dtype: Union[np.dtype, pa.DataType]) -> Tuple[int, ...]: + shape = (None,) + if isinstance(dtype, tensor_extension_types): + dtype = dtype.to_pandas_dtype() + if isinstance(dtype, pa.ListType): + shape += (None,) + elif isinstance(dtype, TensorDtype): + shape += dtype.element_shape + return shape + + def get_tensor_spec( + dtype: Union[np.dtype, pa.DataType], *, name: str + ) -> tf.TypeSpec: + + shape, dtype = get_shape(dtype), get_dtype(dtype) + # Batch dimension is always `None`. So, if there's more than one `None`-valued + # dimension, then the tensor is ragged. + is_ragged = sum(dim is None for dim in shape) > 1 + if is_ragged: + type_spec = tf.RaggedTensorSpec(shape, dtype=dtype) + else: + type_spec = tf.TensorSpec(shape, dtype=dtype, name=name) + return type_spec + + if isinstance(columns, str): + name, dtype = columns, dtypes[columns] + return get_tensor_spec(dtype, name=name) + + return { + name: get_tensor_spec(dtype, name=name) + for name, dtype in dtypes.items() + if name in columns + } diff --git a/lib/python3.12/site-packages/ray/air/_internal/torch_utils.py b/lib/python3.12/site-packages/ray/air/_internal/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..96fe7bd84c749b27014499e04bd51f11f01a4475 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/torch_utils.py @@ -0,0 +1,548 @@ +import warnings +from typing import Any, Dict, List, Optional, Sequence, Union + +import numpy as np +import pandas as pd +import pyarrow +import torch + +from ray._private.ray_constants import env_bool +from ray.air._internal.device_manager import get_torch_device_manager_by_context +from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed +from ray.data.collate_fn import ( + TensorBatchReturnType, + TensorBatchType, + _is_nested_tensor_sequence, + _is_tensor, + _is_tensor_mapping, + _is_tensor_sequence, + _is_tensor_sequence_mapping, +) + +# Default non-blocking transfer for tensors. +DEFAULT_TENSOR_NON_BLOCKING_TRANSFER = env_bool( + "RAY_AIR_DEFAULT_TENSOR_NON_BLOCKING_TRANSFER", + True, +) + + +def get_devices() -> List[torch.device]: + """Gets the correct torch device list configured for this process. + + Returns a list of torch accelerator (GPU, HPU, NPU...) devices allocated for + the current worker. + If no accelerators are assigned, then it returns a list with a single CPU device. + """ + return get_torch_device_manager_by_context().get_devices() + + +def convert_pandas_to_torch_tensor( + data_batch: pd.DataFrame, + columns: Optional[Union[List[str], List[List[str]]]] = None, + column_dtypes: Optional[Union[torch.dtype, List[torch.dtype]]] = None, + unsqueeze: bool = True, +) -> Union[torch.Tensor, List[torch.Tensor]]: + """Converts a Pandas dataframe to a torch Tensor or list of torch Tensors. + + The format of the return type will match the format of ``columns``. If a + list of columns is provided, the return type will be a single tensor. If + ``columns`` is a list of lists, then the return type will be a list of + tensors. + + Args: + data_batch: The pandas dataframe to convert to a + torch tensor. + columns: + The names of the columns in the dataframe to include in the + torch tensor. If this arg is a List[List[str]], then the return + type will be a List of tensors. This is useful for multi-input + models. If None, then use all columns in the ``data_batch``. + column_dtype: The + torch dtype to use for the tensor. If set to None, + then automatically infer the dtype. + unsqueeze: If set to True, the tensors + will be unsqueezed (reshaped to (N, 1)) before being concatenated into + the final tensor. Otherwise, they will be left as is, that is + (N, ). Defaults to True. + + Returns: + Either a torch tensor of size (N, len(columns)) where N is the + number of rows in the ``data_batch`` Dataframe, or a list of + tensors, where the size of item i is (N, len(columns[i])). + + """ + + multi_input = columns and (isinstance(columns[0], (list, tuple))) + + if not multi_input and column_dtypes and not isinstance(column_dtypes, torch.dtype): + raise TypeError( + "If `columns` is a list of strings, " + "`column_dtypes` must be None or a single `torch.dtype`." + f"Got {type(column_dtypes)} instead." + ) + + columns = columns if columns else [] + + def tensorize(vals, dtype): + """This recursive function allows to convert pyarrow List dtypes + to multi-dimensional tensors.""" + if isinstance(vals, pd.api.extensions.ExtensionArray): + # torch.as_tensor() does not yet support the __array__ protocol, so we need + # to convert extension arrays to ndarrays manually before converting to a + # Torch tensor. + # See https://github.com/pytorch/pytorch/issues/51156. + vals = vals.to_numpy() + + if vals.dtype.type is np.object_: + # Column has an object dtype which Torch can't handle, so we try to + # tensorize each column element and then stack the resulting tensors. + tensors = [tensorize(x, dtype) for x in vals] + try: + return torch.stack(tensors) + except RuntimeError: + # NOTE: RuntimeError is raised when trying to stack ragged tensors. + # Try to coerce the tensor to a nested tensor, if possible. + # If this fails, the exception will be propagated up to the caller. + return torch.nested_tensor(tensors) + else: + return torch.as_tensor(vals, dtype=dtype) + + def get_tensor_for_columns(columns, dtype): + feature_tensors = [] + + if columns: + batch = data_batch[columns] + else: + batch = data_batch + + for col in batch.columns: + col_vals = batch[col].values + try: + t = tensorize(col_vals, dtype=dtype) + except Exception as e: + raise ValueError( + f"Failed to convert column {col} to a Torch Tensor of dtype " + f"{dtype}. See above exception chain for the exact failure." + ) from e + if unsqueeze: + t = t.unsqueeze(1) + feature_tensors.append(t) + + if len(feature_tensors) > 1: + feature_tensor = torch.cat(feature_tensors, dim=1) + else: + feature_tensor = feature_tensors[0] + return feature_tensor + + if multi_input: + if type(column_dtypes) not in [list, tuple]: + column_dtypes = [column_dtypes] * len(columns) + return [ + get_tensor_for_columns(columns=subcolumns, dtype=dtype) + for subcolumns, dtype in zip(columns, column_dtypes) + ] + else: + return get_tensor_for_columns(columns=columns, dtype=column_dtypes) + + +def convert_ndarray_to_torch_tensor( + ndarray: np.ndarray, + dtype: Optional[torch.dtype] = None, + device: Optional[Union[str, "torch.device"]] = None, + pin_memory: bool = False, +) -> torch.Tensor: + """Convert a NumPy ndarray to a Torch Tensor. + + Args: + ndarray: A NumPy ndarray that we wish to convert to a Torch Tensor. + dtype: A Torch dtype for the created tensor; if None, the dtype will be + inferred from the NumPy ndarray data. + device: The device on which the tensor(s) should be placed; if None, the Torch + tensor(s) will be constructed on the CPU. + pin_memory: Whether to pin the memory of the created tensors. + + Returns: A Torch Tensor. + """ + ndarray = _unwrap_ndarray_object_type_if_needed(ndarray) + + # Object dtype cannot be converted into PyTorch Tensor. + if ndarray.dtype.type is np.object_: + raise RuntimeError( + "Numpy array of object dtype cannot be converted to a Torch Tensor. This " + "may because the numpy array is a ragged tensor--it contains items of " + "different sizes. If using `iter_torch_batches()` API, you can pass in a " + "`collate_fn` argument to specify custom logic to convert the Numpy array " + "batch to a Torch tensor batch." + ) + + # The numpy array is not always writeable as it can come from the Ray object store. + # Numpy will throw a verbose warning here, which we suppress, as we don't write + # to the tensors. We also don't want to copy the array to avoid memory overhead. + # Original warning: https://github.com/pytorch/pytorch/blob/v1.13.0/ + # torch/csrc/utils/tensor_numpy.cpp#L198-L206 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = torch.as_tensor(ndarray, dtype=dtype, device=device) + + if pin_memory: + assert result.device.type == "cpu", ( + "Pin memory is only supported for CPU tensors. " + f"Got device: {result.device} and pin_memory: {pin_memory}." + ) + result = result.pin_memory() + + return result + + +def convert_ndarray_batch_to_torch_tensor_batch( + ndarrays: Union[np.ndarray, Dict[str, np.ndarray]], + dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, + device: Optional[Union[str, "torch.device"]] = None, + pin_memory: bool = False, +) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + """Convert a NumPy ndarray batch to a Torch Tensor batch. + + Args: + ndarray: A (dict of) NumPy ndarray(s) that we wish to convert to a Torch Tensor. + dtype: A (dict of) Torch dtype(s) for the created tensor; if None, the dtype + will be inferred from the NumPy ndarray data. + device: The device on which the tensor(s) should be placed; if None, the Torch + tensor(s) will be constructed on the CPU. + pin_memory: Whether to pin the memory of the created tensors. + + Returns: A (dict of) Torch Tensor(s). + """ + if isinstance(ndarrays, np.ndarray): + # Single-tensor case. + if isinstance(dtypes, dict): + if len(dtypes) != 1: + raise ValueError( + "When constructing a single-tensor batch, only a single dtype " + f"should be given, instead got: {dtypes}" + ) + dtypes = next(iter(dtypes.values())) + batch = convert_ndarray_to_torch_tensor( + ndarrays, + dtype=dtypes, + device=device, + pin_memory=pin_memory, + ) + else: + # Multi-tensor case. + batch = { + col_name: convert_ndarray_to_torch_tensor( + col_ndarray, + dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, + device=device, + pin_memory=pin_memory, + ) + for col_name, col_ndarray in ndarrays.items() + } + + return batch + + +def load_torch_model( + saved_model: Union[torch.nn.Module, Dict], + model_definition: Optional[torch.nn.Module] = None, +) -> torch.nn.Module: + """Loads a PyTorch model from the provided ``saved_model``. + + ``model_definition`` is only used when ``saved_model`` is + a torch state dict, which will be loaded into ``model_definition``. + Otherwise, ``model_definition`` is discarded. + """ + if isinstance(saved_model, torch.nn.Module): + return saved_model + elif isinstance(saved_model, dict): + if not model_definition: + raise ValueError( + "Attempting to load torch model from a " + "state_dict, but no `model_definition` was " + "provided." + ) + model_definition.load_state_dict(saved_model) + return model_definition + else: + raise ValueError( + f"Saved model is of type {type(saved_model)}. " + f"The model saved in the checkpoint is expected " + f"to be of type `torch.nn.Module`, or a model " + f"state dict of type dict." + ) + + +def contains_tensor(obj): + if isinstance(obj, torch.Tensor): + return True + elif isinstance(obj, dict): + for k, v in obj.items(): + if contains_tensor(k): + return True + if contains_tensor(v): + return True + elif isinstance(obj, (list, tuple)): + for v in obj: + if contains_tensor(v): + return True + return False + + +# Not present in torch<=1.7.0 +# Adapted from https://github.com/pytorch/pytorch/blob/\ +# c18da597e0bb1c1aecc97c77a73fed1849057fa4/torch/nn/modules/utils.py +def consume_prefix_in_state_dict_if_present_not_in_place( + state_dict: Dict[str, Any], prefix: str +) -> Dict[str, Any]: + """Strip the prefix in state_dict, if any and return a new dict. + + Adapted from https://github.com/pytorch/pytorch/blob/\ +c18da597e0bb1c1aecc97c77a73fed1849057fa4/torch/nn/modules/utils.py + The original method modified the dict in-place. + + Args: + state_dict: a state-dict to be loaded to the model. + prefix: prefix. + + """ + copied = False + + for key in state_dict: + if key.startswith(prefix): + newkey = key[len(prefix) :] + if not copied: + # We are doing shallow copies here, so the performance + # impact should be negligible anyway, but this is + # a simple optimization. + state_dict = state_dict.copy() + copied = True + state_dict[newkey] = state_dict.pop(key) + + if "_metadata" in state_dict: + state_dict["_metadata"] = state_dict["_metadata"].copy() + metadata = state_dict["_metadata"] + for key in metadata: + if len(key) == 0: + continue + newkey = key[len(prefix) :] + metadata[newkey] = metadata.pop(key) + + return state_dict + + +def convert_ndarray_list_to_torch_tensor_list( + ndarrays: Dict[str, List[np.ndarray]], + dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, + device: Optional[Union[str, "torch.device"]] = None, + pin_memory: bool = False, +) -> Dict[str, List[torch.Tensor]]: + """Convert a dict mapping column names to lists of ndarrays to Torch Tensors. + + Args: + ndarrays: A dict mapping column names to lists of ndarrays that we wish to convert + to Torch Tensors. + dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype + will be inferred from the NumPy ndarray data. + device: The device on which the tensor(s) should be placed; if None, the Torch + tensor(s) will be constructed on the CPU. + pin_memory: Whether to pin the memory of the created tensors. + + Returns: + A dict mapping column names to lists of Tensors. + """ + return { + col_name: [ + convert_ndarray_batch_to_torch_tensor_batch( + ndarray, + dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, + device=device, + pin_memory=pin_memory, + ) + for ndarray in col_ndarrays + ] + for col_name, col_ndarrays in ndarrays.items() + } + + +def arrow_batch_to_tensors( + batch: pyarrow.Table, + dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, + combine_chunks: bool = False, + pin_memory: bool = False, +) -> Dict[str, List[torch.Tensor]]: + """Convert PyArrow batch to PyTorch tensors. + + Args: + batch: PyArrow batch to convert + dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype + will be inferred from the NumPy ndarray data. + combine_chunks: If True, combine chunks in Arrow batch before converting to + tensors. + pin_memory: Whether to pin the memory of the created tensors. + + Returns: + A dictionary of column name to list of tensors. For non-chunked columns, + the list will contain a single tensor. + """ + from ray.data._internal.arrow_block import ArrowBlockAccessor + from ray.data._internal.arrow_ops import transform_pyarrow + + if combine_chunks: + numpy_batch = ArrowBlockAccessor(batch).to_batch_format("numpy") + return { + col_name: convert_ndarray_batch_to_torch_tensor_batch( + col_array, + dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, + pin_memory=pin_memory, + ) + for col_name, col_array in numpy_batch.items() + } + else: + numpy_list = transform_pyarrow.table_to_numpy_dict_chunked( + batch, + ) + return convert_ndarray_list_to_torch_tensor_list( + numpy_list, + dtypes=dtypes, + pin_memory=pin_memory, + ) + + +@torch.no_grad() +def concat_tensors_to_device( + tensor_sequence: Sequence[torch.Tensor], + device: Optional[Union[str, "torch.device"]] = None, + non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER, +) -> torch.Tensor: + """Stack sequence of tensors into a contiguous GPU tensor. + + Args: + tensor_sequence: Sequence of tensors to stack + device: The device to move tensors to + non_blocking: If True, perform device transfer without forcing a + synchronization. + + Returns: + A contiguous tensor on the target device + """ + # Assumes tensors have the same shape/dtype + assert ( + tensor_sequence + ), f"Cannot stack empty sequence of tensors. Received: {tensor_sequence}" + + assert all( + isinstance(t, torch.Tensor) for t in tensor_sequence + ), "All items must be torch.Tensor. Found invalid types: " + str( + [type(t) for t in tensor_sequence if not isinstance(t, torch.Tensor)] + ) + + first_dtype = tensor_sequence[0].dtype + assert all(t.dtype == first_dtype for t in tensor_sequence), ( + "All tensors must have the same dtype. " + f"Expected: {first_dtype}, got: {[t.dtype for t in tensor_sequence]}" + ) + + first_shape = tensor_sequence[0].shape[1:] + assert all(t.shape[1:] == first_shape for t in tensor_sequence), ( + "All tensors must have the same shape[1:]. " + f"Expected: {first_shape}, got: {[t.shape[1:] for t in tensor_sequence]}" + ) + + first = tensor_sequence[0] + dtype = first.dtype + shape_tail = first.shape[1:] + total_rows = sum(t.shape[0] for t in tensor_sequence) + + # Allocate an empty Tensor on device + result = torch.empty((total_rows, *shape_tail), dtype=dtype, device=device) + + row_start = 0 + for t in tensor_sequence: + row_end = row_start + t.shape[0] + result[row_start:row_end].copy_(t, non_blocking=non_blocking) + row_start = row_end + + return result + + +def _get_type_str(batch: Any) -> str: + """Get a string representation of the possibly nested type of the batch. + + >>> import torch + >>> _get_type_str([1, 2, "???"]) + 'list[int | str]' + >>> _get_type_str({"a": [1, 2, 3], "b": 4}) + 'dict[str, int | list[int]]' + >>> _get_type_str({"a": torch.tensor(1), "b": [torch.tensor(2)]}) + 'dict[str, Tensor | list[Tensor]]' + >>> _get_type_str({"a": torch.tensor(1), "b": {"c": torch.tensor(2)}}) + 'dict[str, Tensor | dict[str, Tensor]]' + """ + curr_type = type(batch).__name__ + if isinstance(batch, (list, tuple)): + val_types = " | ".join(sorted({_get_type_str(v) for v in batch})) + invalid_type_str = f"{curr_type}[{val_types}]" + elif isinstance(batch, dict): + val_types = " | ".join(sorted({_get_type_str(v) for v in batch.values()})) + invalid_type_str = f"{curr_type}[str, {val_types}]" + else: + invalid_type_str = curr_type + return invalid_type_str + + +@torch.no_grad() +def move_tensors_to_device( + batch: TensorBatchType, + device: Optional[Union[str, "torch.device"]] = None, + non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER, +) -> TensorBatchReturnType: + """Move tensors to the specified device. + + Concatenate nested lists/tuples of tensors along the first (batch) dimension. + For example, for the input + ((feature_0_chunk_0,), (feature_1_chunk_0, feature_1_chunk_1)) + the output will be (feature_0_chunk_0, feature_1_chunk_0+1) + where each feature is concatenated along the batch dimension. + + Args: + batch: A tensor or collection of tensors to move to device. Can be: + - A single tensor + - A sequence of tensors + - A sequence of sequences of tensors. The inner sequence of tensors is + combined during GPU transfer. + - A mapping (e.g., dict) of keys to tensors or sequences of tensors. The + sequence of tensors is combined during GPU transfer. + device: The device to move tensors to. If None, tensors are not moved. + non_blocking: If True, perform device transfer without forcing a + synchronization. + + Returns: + The input tensors moved to the specified device + """ + if device is None: + return batch + + if _is_tensor(batch): + return batch.to(device, non_blocking=non_blocking) + elif _is_tensor_sequence(batch): + return type(batch)([t.to(device, non_blocking=non_blocking) for t in batch]) + elif _is_nested_tensor_sequence(batch): + return type(batch)( + [concat_tensors_to_device(t, device, non_blocking) for t in batch] + ) + elif _is_tensor_mapping(batch): + return {k: t.to(device, non_blocking=non_blocking) for k, t in batch.items()} + elif _is_tensor_sequence_mapping(batch): + return { + k: concat_tensors_to_device(v, device, non_blocking) + for k, v in batch.items() + } + else: + raise ValueError( + f"Invalid input type: {_get_type_str(batch)}.\n" + "Expected one of the following: " + "torch.Tensor, " + "List/Tuple[torch.Tensor], " + "Dict[str, torch.Tensor], " + "Mapping[str, List/Tuple[torch.Tensor]]" + ) diff --git a/lib/python3.12/site-packages/ray/air/_internal/uri_utils.py b/lib/python3.12/site-packages/ray/air/_internal/uri_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..33c9ae7c10a29c57bac8f33a2abfd6678eb4a8de --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/uri_utils.py @@ -0,0 +1,101 @@ +import os +import urllib.parse +from pathlib import Path +from typing import Union + + +class URI: + """Represents a URI, supporting path appending and retrieving parent URIs. + + Example Usage: + + >>> s3_uri = URI("s3://bucket/a?scheme=http¶m=1") + >>> s3_uri + URI + >>> str(s3_uri / "b" / "c") + 's3://bucket/a/b/c?scheme=http¶m=1' + >>> str(s3_uri.parent) + 's3://bucket?scheme=http¶m=1' + >>> str(s3_uri) + 's3://bucket/a?scheme=http¶m=1' + >>> s3_uri.parent.name, s3_uri.name + ('bucket', 'a') + >>> local_path = URI("/tmp/local") + >>> str(local_path) + '/tmp/local' + >>> str(local_path.parent) + '/tmp' + >>> str(local_path / "b" / "c") + '/tmp/local/b/c' + + Args: + uri: The URI to represent. + Ex: s3://bucket?scheme=http&endpoint_override=localhost%3A900 + Ex: file:///a/b/c/d + """ + + def __init__(self, uri: str): + self._parsed = urllib.parse.urlparse(uri) + if not self._parsed.scheme: + # Just treat this as a regular path + self._path = Path(uri) + else: + self._path = Path(os.path.normpath(self._parsed.netloc + self._parsed.path)) + + def rstrip_subpath(self, subpath: Path) -> "URI": + """Returns a new URI that strips the given subpath from the end of this URI. + + Example: + >>> uri = URI("s3://bucket/a/b/c/?param=1") + >>> str(uri.rstrip_subpath(Path("b/c"))) + 's3://bucket/a?param=1' + + >>> uri = URI("/tmp/a/b/c/") + >>> str(uri.rstrip_subpath(Path("/b/c/.//"))) + '/tmp/a' + + """ + assert str(self._path).endswith(str(subpath)), (self._path, subpath) + stripped_path = str(self._path).replace(str(subpath), "") + return URI(self._get_str_representation(self._parsed, stripped_path)) + + @property + def name(self) -> str: + return self._path.name + + @property + def parent(self) -> "URI": + assert self._path.parent != ".", f"{str(self)} has no valid parent URI" + return URI(self._get_str_representation(self._parsed, self._path.parent)) + + @property + def scheme(self) -> str: + return self._parsed.scheme + + @property + def path(self) -> str: + return str(self._path) + + def __truediv__(self, path_to_append): + assert isinstance(path_to_append, str) + return URI( + self._get_str_representation(self._parsed, self._path / path_to_append) + ) + + @classmethod + def _get_str_representation( + cls, parsed_uri: urllib.parse.ParseResult, path: Union[str, Path] + ) -> str: + if not parsed_uri.scheme: + return str(path) + return parsed_uri._replace(netloc=str(path), path="").geturl() + + def __repr__(self): + return f"URI<{str(self)}>" + + def __str__(self): + return self._get_str_representation(self._parsed, self._path) + + +def is_uri(path: str) -> bool: + return bool(urllib.parse.urlparse(path).scheme) diff --git a/lib/python3.12/site-packages/ray/air/_internal/usage.py b/lib/python3.12/site-packages/ray/air/_internal/usage.py new file mode 100644 index 0000000000000000000000000000000000000000..4933a75176311252e691abc0e52da57010928b13 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/usage.py @@ -0,0 +1,274 @@ +import collections +import json +import os +from enum import Enum +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union + +from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag + +if TYPE_CHECKING: + from ray.train._internal.storage import StorageContext + from ray.train.trainer import BaseTrainer + from ray.tune import Callback + from ray.tune.schedulers import TrialScheduler + from ray.tune.search import BasicVariantGenerator, Searcher + + +AIR_TRAINERS = { + "HorovodTrainer", + "LightGBMTrainer", + "TensorflowTrainer", + "TorchTrainer", + "XGBoostTrainer", +} + +TRAIN_V2_TRAINERS = { + "DataParallelTrainer", + "JaxTrainer", + "LightGBMTrainer", + "TensorflowTrainer", + "TorchTrainer", + "XGBoostTrainer", +} + +# searchers implemented by Ray Tune. +TUNE_SEARCHERS = { + "AxSearch", + "BayesOptSearch", + "TuneBOHB", + "HEBOSearch", + "HyperOptSearch", + "NevergradSearch", + "OptunaSearch", + "ZOOptSearch", +} + +# These are just wrappers around real searchers. +# We don't want to double tag in this case, otherwise, the real tag +# will be overwritten. +TUNE_SEARCHER_WRAPPERS = { + "ConcurrencyLimiter", + "Repeater", +} + +TUNE_SCHEDULERS = { + "FIFOScheduler", + "AsyncHyperBandScheduler", + "MedianStoppingRule", + "HyperBandScheduler", + "HyperBandForBOHB", + "PopulationBasedTraining", + "PopulationBasedTrainingReplay", + "PB2", + "ResourceChangingScheduler", +} + + +class AirEntrypoint(Enum): + TUNER = "Tuner.fit" + TRAINER = "Trainer.fit" + TUNE_RUN = "tune.run" + TUNE_RUN_EXPERIMENTS = "tune.run_experiments" + + +def _find_class_name(obj, allowed_module_path_prefix: str, whitelist: Set[str]): + """Find the class name of the object. If the object is not + under `allowed_module_path_prefix` or if its class is not in the whitelist, + return "Custom". + + Args: + obj: The object under inspection. + allowed_module_path_prefix: If the `obj`'s class is not under + the `allowed_module_path_prefix`, its class name will be anonymized. + whitelist: If the `obj`'s class is not in the `whitelist`, + it will be anonymized. + Returns: + The class name to be tagged with telemetry. + """ + module_path = obj.__module__ + cls_name = obj.__class__.__name__ + if module_path.startswith(allowed_module_path_prefix) and cls_name in whitelist: + return cls_name + else: + return "Custom" + + +def tag_air_trainer(trainer: "BaseTrainer"): + from ray.train.trainer import BaseTrainer + + assert isinstance(trainer, BaseTrainer) + trainer_name = _find_class_name(trainer, "ray.train", AIR_TRAINERS) + record_extra_usage_tag(TagKey.AIR_TRAINER, trainer_name) + + +def tag_train_v2_trainer(trainer): + from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer + + assert isinstance(trainer, DataParallelTrainer) + trainer_name = _find_class_name(trainer, "ray.train", TRAIN_V2_TRAINERS) + record_extra_usage_tag(TagKey.TRAIN_TRAINER, trainer_name) + + +def tag_searcher(searcher: Union["BasicVariantGenerator", "Searcher"]): + from ray.tune.search import BasicVariantGenerator, Searcher + + if isinstance(searcher, BasicVariantGenerator): + # Note this could be highly inflated as all train flows are treated + # as using BasicVariantGenerator. + record_extra_usage_tag(TagKey.TUNE_SEARCHER, "BasicVariantGenerator") + elif isinstance(searcher, Searcher): + searcher_name = _find_class_name( + searcher, "ray.tune.search", TUNE_SEARCHERS.union(TUNE_SEARCHER_WRAPPERS) + ) + if searcher_name in TUNE_SEARCHER_WRAPPERS: + # ignore to avoid double tagging with wrapper name. + return + record_extra_usage_tag(TagKey.TUNE_SEARCHER, searcher_name) + else: + assert False, ( + "Not expecting a non-BasicVariantGenerator, " + "non-Searcher type passed in for `tag_searcher`." + ) + + +def tag_scheduler(scheduler: "TrialScheduler"): + from ray.tune.schedulers import TrialScheduler + + assert isinstance(scheduler, TrialScheduler) + scheduler_name = _find_class_name(scheduler, "ray.tune.schedulers", TUNE_SCHEDULERS) + record_extra_usage_tag(TagKey.TUNE_SCHEDULER, scheduler_name) + + +def tag_setup_wandb(): + record_extra_usage_tag(TagKey.AIR_SETUP_WANDB_INTEGRATION_USED, "1") + + +def tag_setup_mlflow(): + record_extra_usage_tag(TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED, "1") + + +def _count_callbacks(callbacks: Optional[List["Callback"]]) -> Dict[str, int]: + """Creates a map of callback class name -> count given a list of callbacks.""" + from ray.air.integrations.comet import CometLoggerCallback + from ray.air.integrations.mlflow import MLflowLoggerCallback + from ray.air.integrations.wandb import WandbLoggerCallback + from ray.tune import Callback + from ray.tune.logger import LoggerCallback + from ray.tune.logger.aim import AimLoggerCallback + from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES + + built_in_callbacks = ( + WandbLoggerCallback, + MLflowLoggerCallback, + CometLoggerCallback, + AimLoggerCallback, + ) + DEFAULT_CALLBACK_CLASSES + + callback_names = [callback_cls.__name__ for callback_cls in built_in_callbacks] + callback_counts = collections.defaultdict(int) + + callbacks = callbacks or [] + for callback in callbacks: + if not isinstance(callback, Callback): + # This will error later, but don't include this as custom usage. + continue + + callback_name = callback.__class__.__name__ + + if callback_name in callback_names: + callback_counts[callback_name] += 1 + elif isinstance(callback, LoggerCallback): + callback_counts["CustomLoggerCallback"] += 1 + else: + callback_counts["CustomCallback"] += 1 + + return callback_counts + + +def tag_callbacks(callbacks: Optional[List["Callback"]]) -> bool: + """Records built-in callback usage via a JSON str representing a + dictionary mapping callback class name -> counts. + + User-defined callbacks will increment the count under the `CustomLoggerCallback` + or `CustomCallback` key depending on which of the provided interfaces they subclass. + NOTE: This will NOT track the name of the user-defined callback, + nor its implementation. + + This will NOT report telemetry if no callbacks are provided by the user. + + Returns: + bool: True if usage was recorded, False otherwise. + """ + if not callbacks: + # User didn't pass in any callbacks -> no usage recorded. + return False + + callback_counts = _count_callbacks(callbacks) + + if callback_counts: + callback_counts_str = json.dumps(callback_counts) + record_extra_usage_tag(TagKey.AIR_CALLBACKS, callback_counts_str) + + +def tag_storage_type(storage: "StorageContext"): + """Records the storage configuration of an experiment. + + The storage configuration is set by `RunConfig(storage_path, storage_filesystem)`. + + The possible storage types (defined by `pyarrow.fs.FileSystem.type_name`) are: + - 'local' = pyarrow.fs.LocalFileSystem. This includes NFS usage. + - 'mock' = pyarrow.fs._MockFileSystem. This is used for testing. + - ('s3', 'gcs', 'abfs', 'hdfs'): Various remote storage schemes + with default implementations in pyarrow. + - 'custom' = All other storage schemes, which includes ALL cases where a + custom `storage_filesystem` is provided. + - 'other' = catches any other cases not explicitly handled above. + """ + whitelist = {"local", "mock", "s3", "gcs", "abfs", "hdfs"} + + if storage.custom_fs_provided: + storage_config_tag = "custom" + elif storage.storage_filesystem.type_name in whitelist: + storage_config_tag = storage.storage_filesystem.type_name + else: + storage_config_tag = "other" + + record_extra_usage_tag(TagKey.AIR_STORAGE_CONFIGURATION, storage_config_tag) + + +def tag_ray_air_env_vars() -> bool: + """Records usage of environment variables exposed by the Ray AIR libraries. + + NOTE: This does not track the values of the environment variables, nor + does this track environment variables not explicitly included in the + `all_ray_air_env_vars` allow-list. + + Returns: + bool: True if at least one environment var is supplied by the user. + """ + from ray.air.constants import AIR_ENV_VARS + from ray.train.constants import TRAIN_ENV_VARS + from ray.tune.constants import TUNE_ENV_VARS + + all_ray_air_env_vars = sorted( + set().union(AIR_ENV_VARS, TUNE_ENV_VARS, TRAIN_ENV_VARS) + ) + + user_supplied_env_vars = [] + + for env_var in all_ray_air_env_vars: + if env_var in os.environ: + user_supplied_env_vars.append(env_var) + + if user_supplied_env_vars: + env_vars_str = json.dumps(user_supplied_env_vars) + record_extra_usage_tag(TagKey.AIR_ENV_VARS, env_vars_str) + return True + + return False + + +def tag_air_entrypoint(entrypoint: AirEntrypoint) -> None: + """Records the entrypoint to an AIR training run.""" + assert entrypoint in AirEntrypoint + record_extra_usage_tag(TagKey.AIR_ENTRYPOINT, entrypoint.value) diff --git a/lib/python3.12/site-packages/ray/air/_internal/util.py b/lib/python3.12/site-packages/ray/air/_internal/util.py new file mode 100644 index 0000000000000000000000000000000000000000..55be6c58f64daaf9621023770b9b69355fc0dc79 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/_internal/util.py @@ -0,0 +1,125 @@ +import copy +import logging +import os +import queue +import threading +from typing import Optional + +import numpy as np + +from ray.air.constants import _ERROR_REPORT_TIMEOUT + +logger = logging.getLogger(__name__) + + +def is_nan(value): + return np.isnan(value) + + +def is_nan_or_inf(value): + return is_nan(value) or np.isinf(value) + + +class StartTraceback(Exception): + """These exceptions (and their tracebacks) can be skipped with `skip_exceptions`""" + + pass + + +class StartTracebackWithWorkerRank(StartTraceback): + def __init__(self, worker_rank: int) -> None: + super().__init__() + self.worker_rank = worker_rank + + def __reduce__(self): + return (self.__class__, (self.worker_rank,)) + + +def skip_exceptions(exc: Optional[Exception]) -> Exception: + """Skip all contained `StartTracebacks` to reduce traceback output. + + Returns a shallow copy of the exception with all `StartTracebacks` removed. + + If the RAY_AIR_FULL_TRACEBACKS environment variable is set, + the original exception (not a copy) is returned. + """ + should_not_shorten = bool(int(os.environ.get("RAY_AIR_FULL_TRACEBACKS", "0"))) + + if should_not_shorten: + return exc + + if isinstance(exc, StartTraceback): + # If this is a StartTraceback, skip + return skip_exceptions(exc.__cause__) + + # Perform a shallow copy to prevent recursive __cause__/__context__. + new_exc = copy.copy(exc).with_traceback(exc.__traceback__) + + # Make sure nested exceptions are properly skipped. + cause = getattr(exc, "__cause__", None) + if cause: + new_exc.__cause__ = skip_exceptions(cause) + + return new_exc + + +def exception_cause(exc: Optional[Exception]) -> Optional[Exception]: + if not exc: + return None + + return getattr(exc, "__cause__", None) + + +class RunnerThread(threading.Thread): + """Supervisor thread that runs your script.""" + + def __init__(self, *args, error_queue, **kwargs): + threading.Thread.__init__(self, *args, **kwargs) + self._error_queue = error_queue + self._ret = None + + def _propagate_exception(self, e: BaseException): + try: + # report the error but avoid indefinite blocking which would + # prevent the exception from being propagated in the unlikely + # case that something went terribly wrong + self._error_queue.put(e, block=True, timeout=_ERROR_REPORT_TIMEOUT) + except queue.Full: + logger.critical( + ( + "Runner Thread was unable to report error to main " + "function runner thread. This means a previous error " + "was not processed. This should never happen." + ) + ) + + def run(self): + try: + self._ret = self._target(*self._args, **self._kwargs) + except StopIteration: + logger.debug( + ( + "Thread runner raised StopIteration. Interpreting it as a " + "signal to terminate the thread without error." + ) + ) + except SystemExit as e: + # Do not propagate up for graceful termination. + if e.code == 0: + logger.debug( + ( + "Thread runner raised SystemExit with error code 0. " + "Interpreting it as a signal to terminate the thread " + "without error." + ) + ) + else: + # If non-zero exit code, then raise exception to main thread. + self._propagate_exception(e) + except BaseException as e: + # Propagate all other exceptions to the main thread. + self._propagate_exception(e) + + def join(self, timeout=None): + super(RunnerThread, self).join(timeout) + return self._ret diff --git a/lib/python3.12/site-packages/ray/air/config.py b/lib/python3.12/site-packages/ray/air/config.py new file mode 100644 index 0000000000000000000000000000000000000000..01c93e3c354cc0c04f33eae894cbf2b03137773f --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/config.py @@ -0,0 +1,695 @@ +import logging +import os +import warnings +from collections import Counter, defaultdict +from dataclasses import _MISSING_TYPE, dataclass, fields +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Tuple, + Union, +) + +import pyarrow.fs + +import ray +from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX +from ray._private.thirdparty.tabulate.tabulate import tabulate +from ray.util.annotations import PublicAPI, RayDeprecationWarning +from ray.widgets import Template, make_table_html_repr + +if TYPE_CHECKING: + import ray.tune.progress_reporter + from ray.tune.callback import Callback + from ray.tune.execution.placement_groups import PlacementGroupFactory + from ray.tune.experimental.output import AirVerbosity + from ray.tune.search.sample import Domain + from ray.tune.stopper import Stopper + from ray.tune.utils.log import Verbosity + + +# Dict[str, List] is to support `tune.grid_search`: +# TODO(sumanthratna/matt): Upstream this to Tune. +SampleRange = Union["Domain", Dict[str, List]] + + +MAX = "max" +MIN = "min" +_DEPRECATED_VALUE = "DEPRECATED" + + +logger = logging.getLogger(__name__) + + +def _repr_dataclass(obj, *, default_values: Optional[Dict[str, Any]] = None) -> str: + """A utility function to elegantly represent dataclasses. + + In contrast to the default dataclass `__repr__`, which shows all parameters, this + function only shows parameters with non-default values. + + Args: + obj: The dataclass to represent. + default_values: An optional dictionary that maps field names to default values. + Use this parameter to specify default values that are generated dynamically + (e.g., in `__post_init__` or by a `default_factory`). If a default value + isn't specified in `default_values`, then the default value is inferred from + the `dataclass`. + + Returns: + A representation of the dataclass. + """ + if default_values is None: + default_values = {} + + non_default_values = {} # Maps field name to value. + + def equals(value, default_value): + # We need to special case None because of a bug in pyarrow: + # https://github.com/apache/arrow/issues/38535 + if value is None and default_value is None: + return True + if value is None or default_value is None: + return False + return value == default_value + + for field in fields(obj): + value = getattr(obj, field.name) + default_value = default_values.get(field.name, field.default) + is_required = isinstance(field.default, _MISSING_TYPE) + if is_required or not equals(value, default_value): + non_default_values[field.name] = value + + string = f"{obj.__class__.__name__}(" + string += ", ".join( + f"{name}={value!r}" for name, value in non_default_values.items() + ) + string += ")" + + return string + + +@dataclass +@PublicAPI(stability="stable") +class ScalingConfig: + """Configuration for scaling training. + + For more details, see :ref:`train_scaling_config`. + + Args: + trainer_resources: Resources to allocate for the training coordinator. + The training coordinator launches the worker group and executes + the training function per worker, and this process does NOT require + GPUs. The coordinator is always scheduled on the same node as the + rank 0 worker, so one example use case is to set a minimum amount + of resources (e.g. CPU memory) required by the rank 0 node. + By default, this assigns 1 CPU to the training coordinator. + num_workers: The number of workers (Ray actors) to launch. + Each worker will reserve 1 CPU by default. The number of CPUs + reserved by each worker can be overridden with the + ``resources_per_worker`` argument. + use_gpu: If True, training will be done on GPUs (1 per worker). + Defaults to False. The number of GPUs reserved by each + worker can be overridden with the ``resources_per_worker`` + argument. + resources_per_worker: If specified, the resources + defined in this Dict is reserved for each worker. + Define the ``"CPU"`` key (case-sensitive) to + override the number of CPUs used by each worker. + This can also be used to request :ref:`custom resources `. + placement_strategy: The placement strategy to use for the + placement group of the Ray actors. See :ref:`Placement Group + Strategies ` for the possible options. + accelerator_type: [Experimental] If specified, Ray Train will launch the + training coordinator and workers on the nodes with the specified type + of accelerators. + See :ref:`the available accelerator types `. + Ensure that your cluster has instances with the specified accelerator type + or is able to autoscale to fulfill the request. + + Example: + + .. code-block:: python + + from ray.train import ScalingConfig + scaling_config = ScalingConfig( + # Number of distributed workers. + num_workers=2, + # Turn on/off GPU. + use_gpu=True, + # Assign extra CPU/GPU/custom resources per worker. + resources_per_worker={"GPU": 1, "CPU": 1, "memory": 1e9, "custom": 1.0}, + # Try to schedule workers on different nodes. + placement_strategy="SPREAD", + ) + + """ + + trainer_resources: Optional[Union[Dict, SampleRange]] = None + num_workers: Union[int, SampleRange] = 1 + use_gpu: Union[bool, SampleRange] = False + resources_per_worker: Optional[Union[Dict, SampleRange]] = None + placement_strategy: Union[str, SampleRange] = "PACK" + accelerator_type: Optional[str] = None + + def __post_init__(self): + if self.resources_per_worker: + if not self.use_gpu and self.num_gpus_per_worker > 0: + raise ValueError( + "`use_gpu` is False but `GPU` was found in " + "`resources_per_worker`. Either set `use_gpu` to True or " + "remove `GPU` from `resources_per_worker." + ) + + if self.use_gpu and self.num_gpus_per_worker == 0: + raise ValueError( + "`use_gpu` is True but `GPU` is set to 0 in " + "`resources_per_worker`. Either set `use_gpu` to False or " + "request a positive number of `GPU` in " + "`resources_per_worker." + ) + + def __repr__(self): + return _repr_dataclass(self) + + def _repr_html_(self) -> str: + return make_table_html_repr(obj=self, title=type(self).__name__) + + def __eq__(self, o: "ScalingConfig") -> bool: + if not isinstance(o, type(self)): + return False + return self.as_placement_group_factory() == o.as_placement_group_factory() + + @property + def _resources_per_worker_not_none(self): + if self.resources_per_worker is None: + if self.use_gpu: + # Note that we don't request any CPUs, which avoids possible + # scheduling contention. Generally nodes have many more CPUs than + # GPUs, so not requesting a CPU does not lead to oversubscription. + resources_per_worker = {"GPU": 1} + else: + resources_per_worker = {"CPU": 1} + else: + resources_per_worker = { + k: v for k, v in self.resources_per_worker.items() if v != 0 + } + + if self.use_gpu: + resources_per_worker.setdefault("GPU", 1) + + if self.accelerator_type: + accelerator = f"{RESOURCE_CONSTRAINT_PREFIX}{self.accelerator_type}" + resources_per_worker.setdefault(accelerator, 0.001) + return resources_per_worker + + @property + def _trainer_resources_not_none(self): + if self.trainer_resources is None: + if self.num_workers: + # For Google Colab, don't allocate resources to the base Trainer. + # Colab only has 2 CPUs, and because of this resource scarcity, + # we have to be careful on where we allocate resources. Since Colab + # is not distributed, the concern about many parallel Ray Tune trials + # leading to all Trainers being scheduled on the head node if we set + # `trainer_resources` to 0 is no longer applicable. + try: + import google.colab # noqa: F401 + + trainer_num_cpus = 0 + except ImportError: + trainer_num_cpus = 1 + else: + # If there are no additional workers, then always reserve 1 CPU for + # the Trainer. + trainer_num_cpus = 1 + + trainer_resources = {"CPU": trainer_num_cpus} + else: + trainer_resources = { + k: v for k, v in self.trainer_resources.items() if v != 0 + } + + return trainer_resources + + @property + def total_resources(self): + """Map of total resources required for the trainer.""" + total_resource_map = defaultdict(float, self._trainer_resources_not_none) + for k, value in self._resources_per_worker_not_none.items(): + total_resource_map[k] += value * self.num_workers + return dict(total_resource_map) + + @property + def num_cpus_per_worker(self): + """The number of CPUs to set per worker.""" + return self._resources_per_worker_not_none.get("CPU", 0) + + @property + def num_gpus_per_worker(self): + """The number of GPUs to set per worker.""" + return self._resources_per_worker_not_none.get("GPU", 0) + + @property + def additional_resources_per_worker(self): + """Resources per worker, not including CPU or GPU resources.""" + return { + k: v + for k, v in self._resources_per_worker_not_none.items() + if k not in ["CPU", "GPU"] + } + + def as_placement_group_factory(self) -> "PlacementGroupFactory": + """Returns a PlacementGroupFactory to specify resources for Tune.""" + from ray.tune.execution.placement_groups import PlacementGroupFactory + + trainer_bundle = self._trainer_resources_not_none + worker_bundle = self._resources_per_worker_not_none + + # Colocate Trainer and rank0 worker by merging their bundles + # Note: This empty bundle is required so that the Tune actor manager schedules + # the Trainable onto the combined bundle while taking none of its resources, + # rather than a non-empty head bundle. + combined_bundle = dict(Counter(trainer_bundle) + Counter(worker_bundle)) + bundles = [{}, combined_bundle] + [worker_bundle] * (self.num_workers - 1) + return PlacementGroupFactory(bundles, strategy=self.placement_strategy) + + @classmethod + def from_placement_group_factory( + cls, pgf: "PlacementGroupFactory" + ) -> "ScalingConfig": + """Create a ScalingConfig from a Tune's PlacementGroupFactory + + Note that this is only needed for ResourceChangingScheduler, which + modifies a trial's PlacementGroupFactory but doesn't propagate + the changes to ScalingConfig. TrainTrainable needs to reconstruct + a ScalingConfig from on the trial's PlacementGroupFactory. + """ + + # pgf.bundles = [{trainer + worker}, {worker}, ..., {worker}] + num_workers = len(pgf.bundles) + combined_resources = pgf.bundles[0] + resources_per_worker = pgf.bundles[-1] + use_gpu = bool(resources_per_worker.get("GPU", False)) + placement_strategy = pgf.strategy + + # In `as_placement_group_factory`, we merged the trainer resource into the + # first worker resources bundle. We need to calculate the resources diff to + # get the trainer resources. + # Note: If there's only one worker, we won't be able to calculate the diff. + # We'll have empty trainer bundle and assign all resources to the worker. + trainer_resources = dict( + Counter(combined_resources) - Counter(resources_per_worker) + ) + + return ScalingConfig( + trainer_resources=trainer_resources, + num_workers=num_workers, + use_gpu=use_gpu, + resources_per_worker=resources_per_worker, + placement_strategy=placement_strategy, + ) + + +@dataclass +@PublicAPI(stability="stable") +class FailureConfig: + """Configuration related to failure handling of each training/tuning run. + + Args: + max_failures: Tries to recover a run at least this many times. + Will recover from the latest checkpoint if present. + Setting to -1 will lead to infinite recovery retries. + Setting to 0 will disable retries. Defaults to 0. + fail_fast: Whether to fail upon the first error. + If fail_fast='raise' provided, the original error during training will be + immediately raised. fail_fast='raise' can easily leak resources and + should be used with caution. + """ + + max_failures: int = 0 + fail_fast: Union[bool, str] = False + + def __post_init__(self): + # Same check as in TuneController + if not (isinstance(self.fail_fast, bool) or self.fail_fast.upper() == "RAISE"): + raise ValueError( + "fail_fast must be one of {bool, 'raise'}. " f"Got {self.fail_fast}." + ) + + # Same check as in tune.run + if self.fail_fast and self.max_failures != 0: + raise ValueError( + f"max_failures must be 0 if fail_fast={repr(self.fail_fast)}." + ) + + def __repr__(self): + return _repr_dataclass(self) + + def _repr_html_(self): + return Template("scrollableTable.html.j2").render( + table=tabulate( + { + "Setting": ["Max failures", "Fail fast"], + "Value": [self.max_failures, self.fail_fast], + }, + tablefmt="html", + showindex=False, + headers="keys", + ), + max_height="none", + ) + + +@dataclass +@PublicAPI(stability="stable") +class CheckpointConfig: + """Configurable parameters for defining the checkpointing strategy. + + Default behavior is to persist all checkpoints to disk. If + ``num_to_keep`` is set, the default retention policy is to keep the + checkpoints with maximum timestamp, i.e. the most recent checkpoints. + + Args: + num_to_keep: The number of checkpoints to keep + on disk for this run. If a checkpoint is persisted to disk after + there are already this many checkpoints, then an existing + checkpoint will be deleted. If this is ``None`` then checkpoints + will not be deleted. Must be >= 1. + checkpoint_score_attribute: The attribute that will be used to + score checkpoints to determine which checkpoints should be kept + on disk when there are greater than ``num_to_keep`` checkpoints. + This attribute must be a key from the checkpoint + dictionary which has a numerical value. Per default, the last + checkpoints will be kept. + checkpoint_score_order: Either "max" or "min". + If "max", then checkpoints with highest values of + ``checkpoint_score_attribute`` will be kept. + If "min", then checkpoints with lowest values of + ``checkpoint_score_attribute`` will be kept. + checkpoint_frequency: Number of iterations between checkpoints. If 0 + this will disable checkpointing. + Please note that most trainers will still save one checkpoint at + the end of training. + This attribute is only supported + by trainers that don't take in custom training loops. + checkpoint_at_end: If True, will save a checkpoint at the end of training. + This attribute is only supported by trainers that don't take in + custom training loops. Defaults to True for trainers that support it + and False for generic function trainables. + _checkpoint_keep_all_ranks: This experimental config is deprecated. + This behavior is now controlled by reporting `checkpoint=None` + in the workers that shouldn't persist a checkpoint. + For example, if you only want the rank 0 worker to persist a checkpoint + (e.g., in standard data parallel training), then you should save and + report a checkpoint if `ray.train.get_context().get_world_rank() == 0` + and `None` otherwise. + _checkpoint_upload_from_workers: This experimental config is deprecated. + Uploading checkpoint directly from the worker is now the default behavior. + """ + + num_to_keep: Optional[int] = None + checkpoint_score_attribute: Optional[str] = None + checkpoint_score_order: Optional[str] = MAX + checkpoint_frequency: Optional[int] = 0 + checkpoint_at_end: Optional[bool] = None + _checkpoint_keep_all_ranks: Optional[bool] = _DEPRECATED_VALUE + _checkpoint_upload_from_workers: Optional[bool] = _DEPRECATED_VALUE + + def __post_init__(self): + if self._checkpoint_keep_all_ranks != _DEPRECATED_VALUE: + raise DeprecationWarning( + "The experimental `_checkpoint_keep_all_ranks` config is deprecated. " + "This behavior is now controlled by reporting `checkpoint=None` " + "in the workers that shouldn't persist a checkpoint. " + "For example, if you only want the rank 0 worker to persist a " + "checkpoint (e.g., in standard data parallel training), " + "then you should save and report a checkpoint if " + "`ray.train.get_context().get_world_rank() == 0` " + "and `None` otherwise." + ) + + if self._checkpoint_upload_from_workers != _DEPRECATED_VALUE: + raise DeprecationWarning( + "The experimental `_checkpoint_upload_from_workers` config is " + "deprecated. Uploading checkpoint directly from the worker is " + "now the default behavior." + ) + + if self.num_to_keep is not None and self.num_to_keep <= 0: + raise ValueError( + f"Received invalid num_to_keep: " + f"{self.num_to_keep}. " + f"Must be None or an integer >= 1." + ) + if self.checkpoint_score_order not in (MAX, MIN): + raise ValueError( + f"checkpoint_score_order must be either " f'"{MAX}" or "{MIN}".' + ) + + if self.checkpoint_frequency < 0: + raise ValueError( + f"checkpoint_frequency must be >=0, got {self.checkpoint_frequency}" + ) + + def __repr__(self): + return _repr_dataclass(self) + + def _repr_html_(self) -> str: + if self.num_to_keep is None: + num_to_keep_repr = "All" + else: + num_to_keep_repr = self.num_to_keep + + if self.checkpoint_score_attribute is None: + checkpoint_score_attribute_repr = "Most recent" + else: + checkpoint_score_attribute_repr = self.checkpoint_score_attribute + + if self.checkpoint_at_end is None: + checkpoint_at_end_repr = "" + else: + checkpoint_at_end_repr = self.checkpoint_at_end + + return Template("scrollableTable.html.j2").render( + table=tabulate( + { + "Setting": [ + "Number of checkpoints to keep", + "Checkpoint score attribute", + "Checkpoint score order", + "Checkpoint frequency", + "Checkpoint at end", + ], + "Value": [ + num_to_keep_repr, + checkpoint_score_attribute_repr, + self.checkpoint_score_order, + self.checkpoint_frequency, + checkpoint_at_end_repr, + ], + }, + tablefmt="html", + showindex=False, + headers="keys", + ), + max_height="none", + ) + + @property + def _tune_legacy_checkpoint_score_attr(self) -> Optional[str]: + """Same as ``checkpoint_score_attr`` in ``tune.run``. + + Only used for Legacy API compatibility. + """ + if self.checkpoint_score_attribute is None: + return self.checkpoint_score_attribute + prefix = "" + if self.checkpoint_score_order == MIN: + prefix = "min-" + return f"{prefix}{self.checkpoint_score_attribute}" + + +@dataclass +@PublicAPI(stability="stable") +class RunConfig: + """Runtime configuration for training and tuning runs. + + Upon resuming from a training or tuning run checkpoint, + Ray Train/Tune will automatically apply the RunConfig from + the previously checkpointed run. + + Args: + name: Name of the trial or experiment. If not provided, will be deduced + from the Trainable. + storage_path: [Beta] Path where all results and checkpoints are persisted. + Can be a local directory or a destination on cloud storage. + For multi-node training/tuning runs, this must be set to a + shared storage location (e.g., S3, NFS). + This defaults to the local ``~/ray_results`` directory. + storage_filesystem: [Beta] A custom filesystem to use for storage. + If this is provided, `storage_path` should be a path with its + prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`). + failure_config: Failure mode configuration. + checkpoint_config: Checkpointing configuration. + sync_config: Configuration object for syncing. See train.SyncConfig. + verbose: 0, 1, or 2. Verbosity mode. + 0 = silent, 1 = default, 2 = verbose. Defaults to 1. + If the ``RAY_AIR_NEW_OUTPUT=1`` environment variable is set, + uses the old verbosity settings: + 0 = silent, 1 = only status updates, 2 = status and brief + results, 3 = status and detailed results. + stop: Stop conditions to consider. Refer to ray.tune.stopper.Stopper + for more info. Stoppers should be serializable. + callbacks: [DeveloperAPI] Callbacks to invoke. + Refer to ray.tune.callback.Callback for more info. + Callbacks should be serializable. + Currently only stateless callbacks are supported for resumed runs. + (any state of the callback will not be checkpointed by Tune + and thus will not take effect in resumed runs). + progress_reporter: [DeveloperAPI] Progress reporter for reporting + intermediate experiment progress. Defaults to CLIReporter if + running in command-line, or JupyterNotebookReporter if running in + a Jupyter notebook. + log_to_file: [DeveloperAPI] Log stdout and stderr to files in + trial directories. If this is `False` (default), no files + are written. If `true`, outputs are written to `trialdir/stdout` + and `trialdir/stderr`, respectively. If this is a single string, + this is interpreted as a file relative to the trialdir, to which + both streams are written. If this is a Sequence (e.g. a Tuple), + it has to have length 2 and the elements indicate the files to + which stdout and stderr are written, respectively. + + """ + + name: Optional[str] = None + storage_path: Optional[str] = None + storage_filesystem: Optional[pyarrow.fs.FileSystem] = None + failure_config: Optional[FailureConfig] = None + checkpoint_config: Optional[CheckpointConfig] = None + sync_config: Optional["ray.train.SyncConfig"] = None + verbose: Optional[Union[int, "AirVerbosity", "Verbosity"]] = None + stop: Optional[Union[Mapping, "Stopper", Callable[[str, Mapping], bool]]] = None + callbacks: Optional[List["Callback"]] = None + progress_reporter: Optional["ray.tune.progress_reporter.ProgressReporter"] = None + log_to_file: Union[bool, str, Tuple[str, str]] = False + + # Deprecated + local_dir: Optional[str] = None + + def __post_init__(self): + from ray.train import SyncConfig + from ray.train.constants import DEFAULT_STORAGE_PATH + from ray.tune.experimental.output import AirVerbosity, get_air_verbosity + + if self.local_dir is not None: + raise DeprecationWarning( + "The `RunConfig(local_dir)` argument is deprecated. " + "You should set the `RunConfig(storage_path)` instead." + "See the docs: https://docs.ray.io/en/latest/train/user-guides/" + "persistent-storage.html#setting-the-local-staging-directory" + ) + + if self.storage_path is None: + self.storage_path = DEFAULT_STORAGE_PATH + + # TODO(justinvyu): [Deprecated] + ray_storage_uri: Optional[str] = os.environ.get("RAY_STORAGE") + if ray_storage_uri is not None: + logger.info( + "Using configured Ray Storage URI as the `storage_path`: " + f"{ray_storage_uri}" + ) + warnings.warn( + "The `RAY_STORAGE` environment variable is deprecated. " + "Please use `RunConfig(storage_path)` instead.", + RayDeprecationWarning, + stacklevel=2, + ) + self.storage_path = ray_storage_uri + + if not self.failure_config: + self.failure_config = FailureConfig() + + if not self.sync_config: + self.sync_config = SyncConfig() + + if not self.checkpoint_config: + self.checkpoint_config = CheckpointConfig() + + # Save the original verbose value to check for deprecations + self._verbose = self.verbose + if self.verbose is None: + # Default `verbose` value. For new output engine, + # this is AirVerbosity.DEFAULT. + # For old output engine, this is Verbosity.V3_TRIAL_DETAILS + # Todo (krfricke): Currently uses number to pass test_configs::test_repr + self.verbose = get_air_verbosity(AirVerbosity.DEFAULT) or 3 + + if isinstance(self.storage_path, Path): + self.storage_path = self.storage_path.as_posix() + + def __repr__(self): + from ray.train import SyncConfig + + return _repr_dataclass( + self, + default_values={ + "failure_config": FailureConfig(), + "sync_config": SyncConfig(), + "checkpoint_config": CheckpointConfig(), + }, + ) + + def _repr_html_(self) -> str: + reprs = [] + if self.failure_config is not None: + reprs.append( + Template("title_data_mini.html.j2").render( + title="Failure Config", data=self.failure_config._repr_html_() + ) + ) + if self.sync_config is not None: + reprs.append( + Template("title_data_mini.html.j2").render( + title="Sync Config", data=self.sync_config._repr_html_() + ) + ) + if self.checkpoint_config is not None: + reprs.append( + Template("title_data_mini.html.j2").render( + title="Checkpoint Config", data=self.checkpoint_config._repr_html_() + ) + ) + + # Create a divider between each displayed repr + subconfigs = [Template("divider.html.j2").render()] * (2 * len(reprs) - 1) + subconfigs[::2] = reprs + + settings = Template("scrollableTable.html.j2").render( + table=tabulate( + { + "Name": self.name, + "Local results directory": self.local_dir, + "Verbosity": self.verbose, + "Log to file": self.log_to_file, + }.items(), + tablefmt="html", + headers=["Setting", "Value"], + showindex=False, + ), + max_height="300px", + ) + + return Template("title_data.html.j2").render( + title="RunConfig", + data=Template("run_config.html.j2").render( + subconfigs=subconfigs, + settings=settings, + ), + ) diff --git a/lib/python3.12/site-packages/ray/air/constants.py b/lib/python3.12/site-packages/ray/air/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..0e79f848b9f1cbfa7c61f9903abee4e1c4390a4c --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/constants.py @@ -0,0 +1,94 @@ +# Key to denote the preprocessor in the checkpoint dict. +PREPROCESSOR_KEY = "_preprocessor" + +# Key to denote the model in the checkpoint dict. +MODEL_KEY = "model" + +# Key to denote which dataset is the evaluation dataset. +# Only used in trainers which do not support multiple +# evaluation datasets. +EVALUATION_DATASET_KEY = "evaluation" + +# Key to denote which dataset is the training dataset. +# This is the dataset that the preprocessor is fit on. +TRAIN_DATASET_KEY = "train" + +# Name to use for the column when representing tensors in table format. +TENSOR_COLUMN_NAME = "__value__" + +# The maximum length of strings returned by `__repr__` for AIR objects constructed with +# default values. +MAX_REPR_LENGTH = int(80 * 1.5) + +# Timeout used when putting exceptions raised by runner thread into the queue. +_ERROR_REPORT_TIMEOUT = 10 + +# Timeout when fetching new results after signaling the training function to continue. +_RESULT_FETCH_TIMEOUT = 0.2 + +# Timeout for fetching exceptions raised by the training function. +_ERROR_FETCH_TIMEOUT = 1 + +# The key used to identify whether we have already warned about ray.air.session +# functions being used outside of the session +SESSION_MISUSE_LOG_ONCE_KEY = "air_warn_session_misuse" + +# Name of attribute in Checkpoint storing current Tune ID for restoring +# training with Ray Train +CHECKPOINT_ID_ATTR = "_current_checkpoint_id" + +# Name of the marker dropped by the Trainable. If a worker detects +# the presence of the marker in the trial dir, it will use lazy +# checkpointing. +LAZY_CHECKPOINT_MARKER_FILE = ".lazy_checkpoint_marker" + + +# The timestamp of when the result is generated. +# Default to when the result is processed by tune. +TIMESTAMP = "timestamp" + +# (Auto-filled) Time in seconds this iteration took to run. +# This may be overridden to override the system-computed time difference. +TIME_THIS_ITER_S = "time_this_iter_s" + +# (Auto-filled) The index of this training iteration. +TRAINING_ITERATION = "training_iteration" + +# File that stores parameters of the trial. +EXPR_PARAM_FILE = "params.json" + +# Pickle File that stores parameters of the trial. +EXPR_PARAM_PICKLE_FILE = "params.pkl" + +# File that stores the progress of the trial. +EXPR_PROGRESS_FILE = "progress.csv" + +# File that stores results of the trial. +EXPR_RESULT_FILE = "result.json" + +# File that stores the pickled error file +EXPR_ERROR_PICKLE_FILE = "error.pkl" + +# File that stores the error file +EXPR_ERROR_FILE = "error.txt" + +# File that stores the checkpoint metadata +CHECKPOINT_TUNE_METADATA_FILE = ".tune_metadata" + +# ================================================== +# Environment Variables +# ================================================== + +# Integer value which if set will copy files in reported AIR directory +# checkpoints instead of moving them (if worker is on the same node as Trainable) +COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV = ( + "TRAIN_COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING" +) + +# NOTE: When adding a new environment variable, please track it in this list. +# TODO(ml-team): Most env var constants should get moved here. +AIR_ENV_VARS = { + COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV, + "RAY_AIR_FULL_TRACEBACKS", + "RAY_AIR_NEW_OUTPUT", +} diff --git a/lib/python3.12/site-packages/ray/air/data_batch_type.py b/lib/python3.12/site-packages/ray/air/data_batch_type.py new file mode 100644 index 0000000000000000000000000000000000000000..5d5d09b3218ee11299dfc1ff6aa5d44b2fff67d4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/data_batch_type.py @@ -0,0 +1,11 @@ +from typing import TYPE_CHECKING, Dict, Union + +if TYPE_CHECKING: + import numpy + import pandas # noqa: F401 + import pyarrow + +# TODO de-dup with ray.data.block.DataBatch +DataBatchType = Union[ + "numpy.ndarray", "pyarrow.Table" "pandas.DataFrame", Dict[str, "numpy.ndarray"] +] diff --git a/lib/python3.12/site-packages/ray/air/examples/__init__.py b/lib/python3.12/site-packages/ray/air/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ab0409ca292f1e62ee5463bd44940aea2bc7d5f Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/__init__.py b/lib/python3.12/site-packages/ray/air/execution/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..991a5edbf194c675203a6115d66132f62f9e38e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/__init__.py @@ -0,0 +1,12 @@ +from ray.air.execution.resources.fixed import FixedResourceManager +from ray.air.execution.resources.placement_group import PlacementGroupResourceManager +from ray.air.execution.resources.request import AcquiredResources, ResourceRequest +from ray.air.execution.resources.resource_manager import ResourceManager + +__all__ = [ + "ResourceRequest", + "AcquiredResources", + "ResourceManager", + "FixedResourceManager", + "PlacementGroupResourceManager", +] diff --git a/lib/python3.12/site-packages/ray/air/execution/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..924e0bf095fa99507282012e080b8cbdc7b2f791 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__init__.py b/lib/python3.12/site-packages/ray/air/execution/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db48d7854eec8683013753f9d54a667826ee4071 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/__init__.py @@ -0,0 +1,5 @@ +from ray.air.execution._internal.actor_manager import RayActorManager +from ray.air.execution._internal.barrier import Barrier +from ray.air.execution._internal.tracked_actor import TrackedActor + +__all__ = ["Barrier", "RayActorManager", "TrackedActor"] diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdb9166772d61850d75d3020c9a2a60bcfb683a7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/actor_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/actor_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2694f37fb05793171daed98b408ce305484f3d2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/actor_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/barrier.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/barrier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71b48ed75027e400615c9d415eddb9b17ceffe38 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/barrier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/event_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/event_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dcb48d86551e4bc0563ceb946bc820aee34f2bfe Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/event_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9f70616e34f6595fc81194de9231d36875e1f7c Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor_task.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor_task.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fedcdb06930ea06811649179469ab8d50b96373d Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/_internal/__pycache__/tracked_actor_task.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/actor_manager.py b/lib/python3.12/site-packages/ray/air/execution/_internal/actor_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..33cf6d9ec9f28ca72197c176c5049c19a7239675 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/actor_manager.py @@ -0,0 +1,894 @@ +import logging +import random +import time +import uuid +from collections import Counter, defaultdict +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union + +import ray +from ray.air.execution._internal.event_manager import RayEventManager +from ray.air.execution._internal.tracked_actor import TrackedActor +from ray.air.execution._internal.tracked_actor_task import TrackedActorTask +from ray.air.execution.resources import ( + AcquiredResources, + ResourceManager, + ResourceRequest, +) +from ray.exceptions import RayActorError, RayTaskError + +logger = logging.getLogger(__name__) + + +class RayActorManager: + """Management class for Ray actors and actor tasks. + + This class provides an event-based management interface for actors, and + actor tasks. + + The manager can be used to start actors, stop actors, and schedule and + track task futures on these actors. + The manager will then invoke callbacks related to the tracked entities. + + For instance, when an actor is added with + :meth:`add_actor() `, + a :ref:`TrackedActor `, + an ``on_result`` callback can be specified that is invoked when the task + successfully resolves, and an ``on_error`` callback will resolve when the + task fails. + + The RayActorManager does not implement any true asynchronous processing. Control + has to be explicitly yielded to the event manager via :meth:`RayActorManager.next`. + Callbacks will only be invoked when control is with the RayActorManager, and + callbacks will always be executed sequentially in order of arriving events. + + Args: + resource_manager: Resource manager used to request resources for the actors. + + Example: + + .. code-block:: python + + from ray.air.execution import ResourceRequest + from ray.air.execution._internal import RayActorManager + + actor_manager = RayActorManager() + + # Request an actor + tracked_actor = actor_manager.add_actor( + ActorClass, + kwargs={}, + resource_request=ResourceRequest([{"CPU": 1}]), + on_start=actor_start_callback, + on_stop=actor_stop_callback, + on_error=actor_error_callback + ) + + # Yield control to event manager to start actor + actor_manager.next() + + # Start task on the actor (ActorClass.foo.remote()) + tracked_actor_task = actor_manager.schedule_actor_task( + tracked_actor, + method_name="foo", + on_result=task_result_callback, + on_error=task_error_callback + ) + + # Again yield control to event manager to process task futures + actor_manager.wait() + + """ + + def __init__(self, resource_manager: ResourceManager): + self._resource_manager: ResourceManager = resource_manager + + self._actor_state_events = RayEventManager() + self._actor_task_events = RayEventManager() + + # --- + # Tracked actor futures. + + # This maps TrackedActor objects to their futures. We use this to see if an + # actor has any futures scheduled and to remove them when we terminate an actor. + + # Actors to actor task futures + self._tracked_actors_to_task_futures: Dict[ + TrackedActor, Set[ray.ObjectRef] + ] = defaultdict(set) + + # Actors to actor state futures (start/terminate) + self._tracked_actors_to_state_futures: Dict[ + TrackedActor, Set[ray.ObjectRef] + ] = defaultdict(set) + + # --- + # Pending actors. + # We use three dicts for actors that are requested but not yet started. + + # This dict keeps a list of actors associated with each resource request. + # We use this to start actors in the correct order when their resources + # become available. + self._resource_request_to_pending_actors: Dict[ + ResourceRequest, List[TrackedActor] + ] = defaultdict(list) + + # This dict stores the actor class, kwargs, and resource request of + # pending actors. Once the resources are available, we start the remote + # actor class with its args. We need the resource request to cancel it + # if needed. + self._pending_actors_to_attrs: Dict[ + TrackedActor, Tuple[Type, Dict[str, Any], ResourceRequest] + ] = {} + + # This dict keeps track of cached actor tasks. We can't schedule actor + # tasks before the actor is actually scheduled/live. So when the caller + # tries to schedule a task, we cache it here, and schedule it once the + # actor is started. + self._pending_actors_to_enqueued_actor_tasks: Dict[ + TrackedActor, List[Tuple[TrackedActorTask, str, Tuple[Any], Dict[str, Any]]] + ] = defaultdict(list) + + # --- + # Live actors. + # We keep one dict for actors that are currently running and a set of + # actors that we should forcefully kill. + + # This dict associates the TrackedActor object with the Ray actor handle + # and the resources associated to the actor. We use it to schedule the + # actual ray tasks, and to return the resources when the actor stopped. + self._live_actors_to_ray_actors_resources: Dict[ + TrackedActor, Tuple[ray.actor.ActorHandle, AcquiredResources] + ] = {} + self._live_resource_cache: Optional[Dict[str, Any]] = None + + # This dict contains all actors that should be killed (after calling + # `remove_actor()`). Kill requests will be handled in wait(). + self._live_actors_to_kill: Set[TrackedActor] = set() + + # Track failed actors + self._failed_actor_ids: Set[int] = set() + + def next(self, timeout: Optional[Union[int, float]] = None) -> bool: + """Yield control to event manager to await the next event and invoke callbacks. + + Calling this method will wait for up to ``timeout`` seconds for the next + event to arrive. + + When events arrive, callbacks relating to the events will be + invoked. A timeout of ``None`` will block until the next event arrives. + + Note: + If an actor task fails with a ``RayActorError``, this is one event, + but it may trigger _two_ `on_error` callbacks: One for the actor, + and one for the task. + + Note: + The ``timeout`` argument is used for pure waiting time for events. It does + not include time spent on processing callbacks. Depending on the processing + time of the callbacks, it can take much longer for this function to + return than the specified timeout. + + Args: + timeout: Timeout in seconds to wait for next event. + + Returns: + True if at least one event was processed. + + """ + # First issue any pending forceful actor kills + actor_killed = self._try_kill_actor() + + # We always try to start actors as this won't trigger an event callback + self._try_start_actors() + + # If an actor was killed, this was our event, and we return. + if actor_killed: + return True + + # Otherwise, collect all futures and await the next. + resource_futures = self._resource_manager.get_resource_futures() + actor_state_futures = self._actor_state_events.get_futures() + actor_task_futures = self._actor_task_events.get_futures() + + # Shuffle state futures + shuffled_state_futures = list(actor_state_futures) + random.shuffle(shuffled_state_futures) + + # Shuffle task futures + shuffled_task_futures = list(actor_task_futures) + random.shuffle(shuffled_task_futures) + + # Prioritize resource futures over actor state over task futures + all_futures = resource_futures + shuffled_state_futures + shuffled_task_futures + + start_wait = time.monotonic() + ready, _ = ray.wait(all_futures, num_returns=1, timeout=timeout) + + if not ready: + return False + + [future] = ready + + if future in actor_state_futures: + self._actor_state_events.resolve_future(future) + elif future in actor_task_futures: + self._actor_task_events.resolve_future(future) + else: + self._handle_ready_resource_future() + # Ready resource futures don't count as one event as they don't trigger + # any callbacks. So we repeat until we hit anything that is not a resource + # future. + time_taken = time.monotonic() - start_wait + return self.next( + timeout=max(1e-9, timeout - time_taken) if timeout is not None else None + ) + + self._try_start_actors() + return True + + def _actor_start_resolved(self, tracked_actor: TrackedActor, future: ray.ObjectRef): + """Callback to be invoked when actor started""" + self._tracked_actors_to_state_futures[tracked_actor].remove(future) + + if tracked_actor._on_start: + tracked_actor._on_start(tracked_actor) + + def _actor_stop_resolved(self, tracked_actor: TrackedActor): + """Callback to be invoked when actor stopped""" + self._cleanup_actor(tracked_actor=tracked_actor) + + if tracked_actor._on_stop: + tracked_actor._on_stop(tracked_actor) + + def _actor_start_failed(self, tracked_actor: TrackedActor, exception: Exception): + """Callback to be invoked when actor start/stop failed""" + self._failed_actor_ids.add(tracked_actor.actor_id) + + self._cleanup_actor(tracked_actor=tracked_actor) + + if tracked_actor._on_error: + tracked_actor._on_error(tracked_actor, exception) + + def _actor_task_failed( + self, tracked_actor_task: TrackedActorTask, exception: Exception + ): + """Handle an actor task future that became ready. + + - On actor error, trigger actor error callback AND error task error callback + - On task error, trigger actor task error callback + - On success, trigger actor task result callback + """ + tracked_actor = tracked_actor_task._tracked_actor + + if isinstance(exception, RayActorError): + self._failed_actor_ids.add(tracked_actor.actor_id) + + # Clean up any references to the actor and its futures + self._cleanup_actor(tracked_actor=tracked_actor) + + # Handle actor state callbacks + if tracked_actor._on_error: + tracked_actor._on_error(tracked_actor, exception) + + # Then trigger actor task error callback + if tracked_actor_task._on_error: + tracked_actor_task._on_error(tracked_actor, exception) + + elif isinstance(exception, RayTaskError): + # Otherwise only the task failed. Invoke callback + if tracked_actor_task._on_error: + tracked_actor_task._on_error(tracked_actor, exception) + else: + raise RuntimeError( + f"Caught unexpected exception: {exception}" + ) from exception + + def _actor_task_resolved(self, tracked_actor_task: TrackedActorTask, result: Any): + tracked_actor = tracked_actor_task._tracked_actor + + # Trigger actor task result callback + if tracked_actor_task._on_result: + tracked_actor_task._on_result(tracked_actor, result) + + def _handle_ready_resource_future(self): + """Handle a resource future that became ready. + + - Update state of the resource manager + - Try to start one actor + """ + # Force resource manager to update internal state + self._resource_manager.update_state() + # We handle resource futures one by one, so only try to start 1 actor at a time + self._try_start_actors(max_actors=1) + + def _try_start_actors(self, max_actors: Optional[int] = None) -> int: + """Try to start up to ``max_actors`` actors. + + This function will iterate through all resource requests we collected for + pending actors. As long as a resource request can be fulfilled (resources + are available), we try to start as many actors as possible. + + This will schedule a `Actor.__ray_ready__()` future which, once resolved, + will trigger the `TrackedActor.on_start` callback. + """ + started_actors = 0 + + # Iterate through all resource requests + for resource_request in self._resource_request_to_pending_actors: + if max_actors is not None and started_actors >= max_actors: + break + + # While we have resources ready and there are actors left to schedule + while ( + self._resource_manager.has_resources_ready(resource_request) + and self._resource_request_to_pending_actors[resource_request] + ): + # Acquire resources for actor + acquired_resources = self._resource_manager.acquire_resources( + resource_request + ) + assert acquired_resources + + # Get tracked actor to start + candidate_actors = self._resource_request_to_pending_actors[ + resource_request + ] + assert candidate_actors + + tracked_actor = candidate_actors.pop(0) + + # Get actor class and arguments + actor_cls, kwargs, _ = self._pending_actors_to_attrs.pop(tracked_actor) + + if not isinstance(actor_cls, ray.actor.ActorClass): + actor_cls = ray.remote(actor_cls) + + # Associate to acquired resources + [remote_actor_cls] = acquired_resources.annotate_remote_entities( + [actor_cls] + ) + + # Start Ray actor + actor = remote_actor_cls.remote(**kwargs) + + # Track + self._live_actors_to_ray_actors_resources[tracked_actor] = ( + actor, + acquired_resources, + ) + self._live_resource_cache = None + + # Schedule ready future + future = actor.__ray_ready__.remote() + + self._tracked_actors_to_state_futures[tracked_actor].add(future) + + # We need to create the callbacks in a function so tracked_actors + # are captured correctly. + def create_callbacks( + tracked_actor: TrackedActor, future: ray.ObjectRef + ): + def on_actor_start(result: Any): + self._actor_start_resolved( + tracked_actor=tracked_actor, future=future + ) + + def on_error(exception: Exception): + self._actor_start_failed( + tracked_actor=tracked_actor, exception=exception + ) + + return on_actor_start, on_error + + on_actor_start, on_error = create_callbacks( + tracked_actor=tracked_actor, future=future + ) + + self._actor_state_events.track_future( + future=future, + on_result=on_actor_start, + on_error=on_error, + ) + + self._enqueue_cached_actor_tasks(tracked_actor=tracked_actor) + + started_actors += 1 + + return started_actors + + def _enqueue_cached_actor_tasks(self, tracked_actor: TrackedActor): + assert tracked_actor in self._live_actors_to_ray_actors_resources + + # Enqueue cached futures + cached_tasks = self._pending_actors_to_enqueued_actor_tasks.pop( + tracked_actor, [] + ) + for tracked_actor_task, method_name, args, kwargs in cached_tasks: + self._schedule_tracked_actor_task( + tracked_actor_task=tracked_actor_task, + method_name=method_name, + args=args, + kwargs=kwargs, + ) + + def _try_kill_actor(self) -> bool: + """Try to kill actor scheduled for termination.""" + if not self._live_actors_to_kill: + return False + + tracked_actor = self._live_actors_to_kill.pop() + + # Remove from tracked actors + ( + ray_actor, + acquired_resources, + ) = self._live_actors_to_ray_actors_resources[tracked_actor] + + # Hard kill if requested + ray.kill(ray_actor) + + self._cleanup_actor_futures(tracked_actor) + + self._actor_stop_resolved(tracked_actor) + + return True + + def _cleanup_actor(self, tracked_actor: TrackedActor): + self._cleanup_actor_futures(tracked_actor) + + # Remove from tracked actors + ( + ray_actor, + acquired_resources, + ) = self._live_actors_to_ray_actors_resources.pop(tracked_actor) + self._live_resource_cache = None + + # Return resources + self._resource_manager.free_resources(acquired_resource=acquired_resources) + + @property + def all_actors(self) -> List[TrackedActor]: + """Return all ``TrackedActor`` objects managed by this manager instance.""" + return self.live_actors + self.pending_actors + + @property + def live_actors(self) -> List[TrackedActor]: + """Return all ``TrackedActor`` objects that are currently alive.""" + return list(self._live_actors_to_ray_actors_resources) + + @property + def pending_actors(self) -> List[TrackedActor]: + """Return all ``TrackedActor`` objects that are currently pending.""" + return list(self._pending_actors_to_attrs) + + @property + def num_live_actors(self): + """Return number of started actors.""" + return len(self.live_actors) + + @property + def num_pending_actors(self) -> int: + """Return number of pending (not yet started) actors.""" + return len(self.pending_actors) + + @property + def num_total_actors(self): + """Return number of total actors.""" + return len(self.all_actors) + + @property + def num_actor_tasks(self): + """Return number of pending tasks""" + return self._actor_task_events.num_futures + + def get_live_actors_resources(self): + if self._live_resource_cache: + return self._live_resource_cache + + counter = Counter() + for _, acq in self._live_actors_to_ray_actors_resources.values(): + for bdl in acq.resource_request.bundles: + counter.update(bdl) + self._live_resource_cache = dict(counter) + return self._live_resource_cache + + def add_actor( + self, + cls: Union[Type, ray.actor.ActorClass], + kwargs: Dict[str, Any], + resource_request: ResourceRequest, + *, + on_start: Optional[Callable[[TrackedActor], None]] = None, + on_stop: Optional[Callable[[TrackedActor], None]] = None, + on_error: Optional[Callable[[TrackedActor, Exception], None]] = None, + ) -> TrackedActor: + """Add an actor to be tracked. + + This method will request resources to start the actor. Once the resources + are available, the actor will be started and the + :meth:`TrackedActor.on_start + ` callback + will be invoked. + + Args: + cls: Actor class to schedule. + kwargs: Keyword arguments to pass to actor class on construction. + resource_request: Resources required to start the actor. + on_start: Callback to invoke when the actor started. + on_stop: Callback to invoke when the actor stopped. + on_error: Callback to invoke when the actor failed. + + Returns: + Tracked actor object to reference actor in subsequent API calls. + + """ + tracked_actor = TrackedActor( + uuid.uuid4().int, on_start=on_start, on_stop=on_stop, on_error=on_error + ) + + self._pending_actors_to_attrs[tracked_actor] = cls, kwargs, resource_request + self._resource_request_to_pending_actors[resource_request].append(tracked_actor) + + self._resource_manager.request_resources(resource_request=resource_request) + + return tracked_actor + + def remove_actor( + self, + tracked_actor: TrackedActor, + kill: bool = False, + stop_future: Optional[ray.ObjectRef] = None, + ) -> bool: + """Remove a tracked actor. + + If the actor has already been started, this will stop the actor. This will + trigger the :meth:`TrackedActor.on_stop + ` + callback once the actor stopped. + + If the actor has only been requested, but not started, yet, this will cancel + the actor request. This will not trigger any callback. + + If ``kill=True``, this will use ``ray.kill()`` to forcefully terminate the + actor. Otherwise, graceful actor deconstruction will be scheduled after + all currently tracked futures are resolved. + + This method returns a boolean, indicating if a stop future is tracked and + the ``on_stop`` callback will be invoked. If the actor has been alive, + this will be ``True``. If the actor hasn't been scheduled, yet, or failed + (and triggered the ``on_error`` callback), this will be ``False``. + + Args: + tracked_actor: Tracked actor to be removed. + kill: If set, will forcefully terminate the actor instead of gracefully + scheduling termination. + stop_future: If set, use this future to track actor termination. + Otherwise, schedule a ``__ray_terminate__`` future. + + Returns: + Boolean indicating if the actor was previously alive, and thus whether + a callback will be invoked once it is terminated. + + """ + if tracked_actor.actor_id in self._failed_actor_ids: + logger.debug( + f"Tracked actor already failed, no need to remove: {tracked_actor}" + ) + return False + elif tracked_actor in self._live_actors_to_ray_actors_resources: + # Ray actor is running. + + if not kill: + # Schedule __ray_terminate__ future + ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor] + + # Clear state futures here to avoid resolving __ray_ready__ futures + for future in list( + self._tracked_actors_to_state_futures[tracked_actor] + ): + self._actor_state_events.discard_future(future) + self._tracked_actors_to_state_futures[tracked_actor].remove(future) + + # If the __ray_ready__ future hasn't resolved yet, but we already + # scheduled the actor via Actor.remote(), we just want to stop + # it but not trigger any callbacks. This is in accordance with + # the contract defined in the docstring. + tracked_actor._on_start = None + tracked_actor._on_stop = None + tracked_actor._on_error = None + + def on_actor_stop(*args, **kwargs): + self._actor_stop_resolved(tracked_actor=tracked_actor) + + if stop_future: + # If the stop future was schedule via the actor manager, + # discard (track it as state future instead). + self._actor_task_events.discard_future(stop_future) + else: + stop_future = ray_actor.__ray_terminate__.remote() + + self._actor_state_events.track_future( + future=stop_future, + on_result=on_actor_stop, + on_error=on_actor_stop, + ) + + self._tracked_actors_to_state_futures[tracked_actor].add(stop_future) + else: + # kill = True + self._live_actors_to_kill.add(tracked_actor) + + return True + + elif tracked_actor in self._pending_actors_to_attrs: + # Actor is pending, stop + _, _, resource_request = self._pending_actors_to_attrs.pop(tracked_actor) + self._resource_request_to_pending_actors[resource_request].remove( + tracked_actor + ) + self._resource_manager.cancel_resource_request( + resource_request=resource_request + ) + return False + else: + raise ValueError(f"Unknown tracked actor: {tracked_actor}") + + def is_actor_started(self, tracked_actor: TrackedActor) -> bool: + """Returns True if the actor has been started. + + Args: + tracked_actor: Tracked actor object. + """ + return ( + tracked_actor in self._live_actors_to_ray_actors_resources + and tracked_actor.actor_id not in self._failed_actor_ids + ) + + def is_actor_failed(self, tracked_actor: TrackedActor) -> bool: + return tracked_actor.actor_id in self._failed_actor_ids + + def get_actor_resources( + self, tracked_actor: TrackedActor + ) -> Optional[AcquiredResources]: + """Returns the acquired resources of an actor that has been started. + + This will return ``None`` if the actor has not been started, yet. + + Args: + tracked_actor: Tracked actor object. + """ + if not self.is_actor_started(tracked_actor): + return None + + return self._live_actors_to_ray_actors_resources[tracked_actor][1] + + def schedule_actor_task( + self, + tracked_actor: TrackedActor, + method_name: str, + args: Optional[Tuple] = None, + kwargs: Optional[Dict] = None, + on_result: Optional[Callable[[TrackedActor, Any], None]] = None, + on_error: Optional[Callable[[TrackedActor, Exception], None]] = None, + _return_future: bool = False, + ) -> Optional[ray.ObjectRef]: + """Schedule and track a task on an actor. + + This method will schedule a remote task ``method_name`` on the + ``tracked_actor``. + + This method accepts two optional callbacks that will be invoked when + their respective events are triggered. + + The ``on_result`` callback is triggered when a task resolves successfully. + It should accept two arguments: The actor for which the + task resolved, and the result received from the remote call. + + The ``on_error`` callback is triggered when a task fails. + It should accept two arguments: The actor for which the + task threw an error, and the exception. + + Args: + tracked_actor: Actor to schedule task on. + method_name: Remote method name to invoke on the actor. If this is + e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be + scheduled. + args: Arguments to pass to the task. + kwargs: Keyword arguments to pass to the task. + on_result: Callback to invoke when the task resolves. + on_error: Callback to invoke when the task fails. + + Raises: + ValueError: If the ``tracked_actor`` is not managed by this event manager. + + """ + args = args or tuple() + kwargs = kwargs or {} + + if tracked_actor.actor_id in self._failed_actor_ids: + return + + tracked_actor_task = TrackedActorTask( + tracked_actor=tracked_actor, on_result=on_result, on_error=on_error + ) + + if tracked_actor not in self._live_actors_to_ray_actors_resources: + # Actor is not started, yet + if tracked_actor not in self._pending_actors_to_attrs: + raise ValueError( + f"Tracked actor is not managed by this event manager: " + f"{tracked_actor}" + ) + + # Cache tasks for future execution + self._pending_actors_to_enqueued_actor_tasks[tracked_actor].append( + (tracked_actor_task, method_name, args, kwargs) + ) + else: + res = self._schedule_tracked_actor_task( + tracked_actor_task=tracked_actor_task, + method_name=method_name, + args=args, + kwargs=kwargs, + _return_future=_return_future, + ) + if _return_future: + return res[1] + + def _schedule_tracked_actor_task( + self, + tracked_actor_task: TrackedActorTask, + method_name: str, + *, + args: Optional[Tuple] = None, + kwargs: Optional[Dict] = None, + _return_future: bool = False, + ) -> Union[TrackedActorTask, Tuple[TrackedActorTask, ray.ObjectRef]]: + tracked_actor = tracked_actor_task._tracked_actor + ray_actor, _ = self._live_actors_to_ray_actors_resources[tracked_actor] + + try: + remote_fn = getattr(ray_actor, method_name) + except AttributeError as e: + raise AttributeError( + f"Remote function `{method_name}()` does not exist for this actor." + ) from e + + def on_result(result: Any): + self._actor_task_resolved( + tracked_actor_task=tracked_actor_task, result=result + ) + + def on_error(exception: Exception): + self._actor_task_failed( + tracked_actor_task=tracked_actor_task, exception=exception + ) + + future = remote_fn.remote(*args, **kwargs) + + self._actor_task_events.track_future( + future=future, on_result=on_result, on_error=on_error + ) + + self._tracked_actors_to_task_futures[tracked_actor].add(future) + + if _return_future: + return tracked_actor_task, future + + return tracked_actor_task + + def schedule_actor_tasks( + self, + tracked_actors: List[TrackedActor], + method_name: str, + *, + args: Optional[Union[Tuple, List[Tuple]]] = None, + kwargs: Optional[Union[Dict, List[Dict]]] = None, + on_result: Optional[Callable[[TrackedActor, Any], None]] = None, + on_error: Optional[Callable[[TrackedActor, Exception], None]] = None, + ) -> None: + """Schedule and track tasks on a list of actors. + + This method will schedule a remote task ``method_name`` on all + ``tracked_actors``. + + ``args`` and ``kwargs`` can be a single tuple/dict, in which case the same + (keyword) arguments are passed to all actors. If a list is passed instead, + they are mapped to the respective actors. In that case, the list of + (keyword) arguments must be the same length as the list of actors. + + This method accepts two optional callbacks that will be invoked when + their respective events are triggered. + + The ``on_result`` callback is triggered when a task resolves successfully. + It should accept two arguments: The actor for which the + task resolved, and the result received from the remote call. + + The ``on_error`` callback is triggered when a task fails. + It should accept two arguments: The actor for which the + task threw an error, and the exception. + + Args: + tracked_actors: List of actors to schedule tasks on. + method_name: Remote actor method to invoke on the actors. If this is + e.g. ``foo``, then ``actor.foo.remote(*args, **kwargs)`` will be + scheduled on all actors. + args: Arguments to pass to the task. + kwargs: Keyword arguments to pass to the task. + on_result: Callback to invoke when the task resolves. + on_error: Callback to invoke when the task fails. + + """ + if not isinstance(args, List): + args_list = [args] * len(tracked_actors) + else: + if len(tracked_actors) != len(args): + raise ValueError( + f"Length of args must be the same as tracked_actors " + f"list. Got `len(kwargs)={len(kwargs)}` and " + f"`len(tracked_actors)={len(tracked_actors)}" + ) + args_list = args + + if not isinstance(kwargs, List): + kwargs_list = [kwargs] * len(tracked_actors) + else: + if len(tracked_actors) != len(kwargs): + raise ValueError( + f"Length of kwargs must be the same as tracked_actors " + f"list. Got `len(args)={len(args)}` and " + f"`len(tracked_actors)={len(tracked_actors)}" + ) + kwargs_list = kwargs + + for tracked_actor, args, kwargs in zip(tracked_actors, args_list, kwargs_list): + self.schedule_actor_task( + tracked_actor=tracked_actor, + method_name=method_name, + args=args, + kwargs=kwargs, + on_result=on_result, + on_error=on_error, + ) + + def clear_actor_task_futures(self, tracked_actor: TrackedActor): + """Discard all actor task futures from a tracked actor.""" + futures = self._tracked_actors_to_task_futures.pop(tracked_actor, []) + for future in futures: + self._actor_task_events.discard_future(future) + + def _cleanup_actor_futures(self, tracked_actor: TrackedActor): + # Remove all actor task futures + self.clear_actor_task_futures(tracked_actor=tracked_actor) + + # Remove all actor state futures + futures = self._tracked_actors_to_state_futures.pop(tracked_actor, []) + for future in futures: + self._actor_state_events.discard_future(future) + + def cleanup(self): + for ( + actor, + acquired_resources, + ) in self._live_actors_to_ray_actors_resources.values(): + ray.kill(actor) + self._resource_manager.free_resources(acquired_resources) + + for ( + resource_request, + pending_actors, + ) in self._resource_request_to_pending_actors.items(): + for i in range(len(pending_actors)): + self._resource_manager.cancel_resource_request(resource_request) + + self._resource_manager.clear() + + self.__init__(resource_manager=self._resource_manager) diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/barrier.py b/lib/python3.12/site-packages/ray/air/execution/_internal/barrier.py new file mode 100644 index 0000000000000000000000000000000000000000..8cb0c5a2bf55242281090a1cc6bcac9009face52 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/barrier.py @@ -0,0 +1,93 @@ +from typing import Any, Callable, List, Optional, Tuple + + +class Barrier: + """Barrier to collect results and process them in bulk. + + A barrier can be used to collect multiple results and process them in bulk once + a certain count or a timeout is reached. + + For instance, if ``max_results=N``, the ``on_completion`` callback will be + invoked once :meth:`arrive` has been called ``N`` times. + + The completion callback will only be invoked once, even if more results + arrive after completion. The collected results can be resetted + with :meth:`reset`, after which the callback may be invoked again. + + The completion callback should expect one argument, which is the barrier + object that completed. + + Args: + max_results: Maximum number of results to collect before a call to + :meth:`wait` resolves or the :meth:`on_completion` callback is invoked. + on_completion: Callback to invoke when ``max_results`` results + arrived at the barrier. + + """ + + def __init__( + self, + max_results: int, + *, + on_completion: Optional[Callable[["Barrier"], None]] = None, + ): + self._max_results = max_results + + # on_completion callback + self._completed = False + self._on_completion = on_completion + + # Collect received results + self._results: List[Tuple[Any]] = [] + + def arrive(self, *data): + """Notify barrier that a result successfully arrived. + + This will count against the ``max_results`` limit. The received result + will be included in a call to :meth:`get_results`. + + Args: + *data: Result data to be cached. Can be obtained via :meth:`get_results`. + + """ + if len(data) == 1: + data = data[0] + + self._results.append(data) + self._check_completion() + + def _check_completion(self): + if self._completed: + # Already fired completion callback + return + + if self.num_results >= self._max_results: + # Barrier is complete + self._completed = True + + if self._on_completion: + self._on_completion(self) + + @property + def completed(self) -> bool: + """Returns True if the barrier is completed.""" + return self._completed + + @property + def num_results(self) -> int: + """Number of received (successful) results.""" + return len(self._results) + + def get_results(self) -> List[Tuple[Any]]: + """Return list of received results.""" + return self._results + + def reset(self) -> None: + """Reset barrier, removing all received results. + + Resetting the barrier will reset the completion status. When ``max_results`` + is set and enough new events arrive after resetting, the + :meth:`on_completion` callback will be invoked again. + """ + self._completed = False + self._results = [] diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/event_manager.py b/lib/python3.12/site-packages/ray/air/execution/_internal/event_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..fbea6f72d2925b5435931e58203d206bd0009709 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/event_manager.py @@ -0,0 +1,148 @@ +import random +from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union + +import ray + +_ResultCallback = Callable[[Any], None] +_ErrorCallback = Callable[[Exception], None] + + +class RayEventManager: + """Event manager for Ray futures. + + The event manager can be used to track futures and invoke callbacks when + they resolve. + + Futures are tracked with :meth:`track_future`. Future can then be awaited with + :meth:`wait`. When futures successfully resolve, they trigger an optional + ``on_result`` callback that can be passed to :meth:`track_future`. If they + fail, they trigger an optional ``on_error`` callback. + + Args: + shuffle_futures: If True, futures will be shuffled before awaited. This + will avoid implicit prioritization of futures within Ray. + """ + + def __init__(self, shuffle_futures: bool = True): + self._shuffle_futures = shuffle_futures + + # Map of futures to callbacks (result, error) + self._tracked_futures: Dict[ + ray.ObjectRef, Tuple[Optional[_ResultCallback], Optional[_ErrorCallback]] + ] = {} + + def track_future( + self, + future: ray.ObjectRef, + on_result: Optional[_ResultCallback] = None, + on_error: Optional[_ErrorCallback] = None, + ): + """Track a single future and invoke callbacks on resolution. + + Control has to be yielded to the event manager for the callbacks to + be invoked, either via :meth:`wait` or via :meth:`resolve_future`. + + Args: + future: Ray future to await. + on_result: Callback to invoke when the future resolves successfully. + on_error: Callback to invoke when the future fails. + + """ + self._tracked_futures[future] = (on_result, on_error) + + def track_futures( + self, + futures: Iterable[ray.ObjectRef], + on_result: Optional[_ResultCallback] = None, + on_error: Optional[_ErrorCallback] = None, + ): + """Track multiple futures and invoke callbacks on resolution. + + Control has to be yielded to the event manager for the callbacks to + be invoked, either via :meth:`wait` or via :meth:`resolve_future`. + + Args: + futures: Ray futures to await. + on_result: Callback to invoke when the future resolves successfully. + on_error: Callback to invoke when the future fails. + + """ + for future in futures: + self.track_future(future, on_result=on_result, on_error=on_error) + + def discard_future(self, future: ray.ObjectRef): + """Remove future from tracking. + + The future will not be awaited anymore, and it will not trigger any callbacks. + + Args: + future: Ray futures to discard. + """ + self._tracked_futures.pop(future, None) + + def get_futures(self) -> Set[ray.ObjectRef]: + """Get futures tracked by the event manager.""" + return set(self._tracked_futures) + + @property + def num_futures(self) -> int: + return len(self._tracked_futures) + + def resolve_future(self, future: ray.ObjectRef): + """Resolve a single future. + + This method will block until the future is available. It will then + trigger the callback associated to the future and the event (success + or error), if specified. + + Args: + future: Ray future to resolve. + + """ + try: + on_result, on_error = self._tracked_futures.pop(future) + except KeyError as e: + raise ValueError( + f"Future {future} is not tracked by this RayEventManager" + ) from e + + try: + result = ray.get(future) + except Exception as e: + if on_error: + on_error(e) + else: + raise e + else: + if on_result: + on_result(result) + + def wait( + self, + timeout: Optional[Union[float, int]] = None, + num_results: Optional[int] = 1, + ): + """Wait up to ``timeout`` seconds for ``num_results`` futures to resolve. + + If ``timeout=None``, this method will block until all `num_results`` futures + resolve. If ``num_results=None``, this method will await all tracked futures. + + For every future that resolves, the respective associated callbacks will be + invoked. + + Args: + timeout: Timeout in second to wait for futures to resolve. + num_results: Number of futures to await. If ``None``, will wait for + all tracked futures to resolve. + + """ + futures = list(self.get_futures()) + + if self._shuffle_futures: + random.shuffle(futures) + + num_results = num_results or len(futures) + + ready, _ = ray.wait(list(futures), timeout=timeout, num_returns=num_results) + for future in ready: + self.resolve_future(future) diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor.py b/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..91abd0556e40b197ea82ac8ac55b53613ef635cc --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor.py @@ -0,0 +1,54 @@ +from typing import Callable, Optional + + +class TrackedActor: + """Actor tracked by an actor manager. + + This object is used to reference a Ray actor on an actor manager + + Existence of this object does not mean that the Ray actor has already been started. + Actor state can be inquired from the actor manager tracking the Ray actor. + + Note: + Objects of this class are returned by the :class:`RayActorManager`. + This class should not be instantiated manually. + + Attributes: + actor_id: ID for identification of the actor within the actor manager. This + ID is not related to the Ray actor ID. + + """ + + def __init__( + self, + actor_id: int, + on_start: Optional[Callable[["TrackedActor"], None]] = None, + on_stop: Optional[Callable[["TrackedActor"], None]] = None, + on_error: Optional[Callable[["TrackedActor", Exception], None]] = None, + ): + self.actor_id = actor_id + self._on_start = on_start + self._on_stop = on_stop + self._on_error = on_error + + def set_on_start(self, on_start: Optional[Callable[["TrackedActor"], None]]): + self._on_start = on_start + + def set_on_stop(self, on_stop: Optional[Callable[["TrackedActor"], None]]): + self._on_stop = on_stop + + def set_on_error( + self, on_error: Optional[Callable[["TrackedActor", Exception], None]] + ): + self._on_error = on_error + + def __repr__(self): + return f"" + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + return self.actor_id == other.actor_id + + def __hash__(self): + return hash(self.actor_id) diff --git a/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor_task.py b/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor_task.py new file mode 100644 index 0000000000000000000000000000000000000000..2fb21cbfb514924d8d38267c4e63321a27f9cb41 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/_internal/tracked_actor_task.py @@ -0,0 +1,42 @@ +from typing import Any, Callable, Optional + +from ray.air.execution._internal.tracked_actor import TrackedActor + + +class TrackedActorTask: + """Actor task tracked by a Ray event manager. + + This container class is used to define callbacks to be invoked when + the task resolves, errors, or times out. + + Note: + Objects of this class are returned by the :class:`RayActorManager`. + This class should not be instantiated manually. + + Args: + tracked_actor: Tracked actor object this task is scheduled on. + on_result: Callback to invoke when the task resolves. + on_error: Callback to invoke when the task fails. + + Example: + + .. code-block:: python + + tracked_futures = actor_manager.schedule_actor_tasks( + actor_manager.live_actors, + "foo", + on_result=lambda actor, result: print(result) + ) + + """ + + def __init__( + self, + tracked_actor: TrackedActor, + on_result: Optional[Callable[[TrackedActor, Any], None]] = None, + on_error: Optional[Callable[[TrackedActor, Exception], None]] = None, + ): + self._tracked_actor = tracked_actor + + self._on_result = on_result + self._on_error = on_error diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__init__.py b/lib/python3.12/site-packages/ray/air/execution/resources/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..991a5edbf194c675203a6115d66132f62f9e38e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/resources/__init__.py @@ -0,0 +1,12 @@ +from ray.air.execution.resources.fixed import FixedResourceManager +from ray.air.execution.resources.placement_group import PlacementGroupResourceManager +from ray.air.execution.resources.request import AcquiredResources, ResourceRequest +from ray.air.execution.resources.resource_manager import ResourceManager + +__all__ = [ + "ResourceRequest", + "AcquiredResources", + "ResourceManager", + "FixedResourceManager", + "PlacementGroupResourceManager", +] diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f08fd7d7c8c9dd99a9a1644cc39a4adf284225eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/fixed.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/fixed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82799cd173415e3a00f595b01f2aeb951a5c1b27 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/fixed.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/placement_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/placement_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ce40af2630caf65b1a057ae9432a6e4b81eacd0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/placement_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/request.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/request.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5596977629c4d9353a1b6aaa0d0f0d2decdb3862 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/request.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/resource_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/resource_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25234e64bd959d2ebd6870f3ba0ae22df5f74b53 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/execution/resources/__pycache__/resource_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/fixed.py b/lib/python3.12/site-packages/ray/air/execution/resources/fixed.py new file mode 100644 index 0000000000000000000000000000000000000000..5e969ed7c23cccf9d16af62bc0e421df2cc53ff0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/resources/fixed.py @@ -0,0 +1,147 @@ +from dataclasses import dataclass +from typing import Dict, List, Optional + +import ray +from ray import LOCAL_MODE, SCRIPT_MODE +from ray.air.execution.resources.request import ( + AcquiredResources, + RemoteRayEntity, + ResourceRequest, +) +from ray.air.execution.resources.resource_manager import ResourceManager +from ray.util.annotations import DeveloperAPI + +# Avoid numerical errors by multiplying and subtracting with this number. +# Compare: 0.99 - 0.33 = 0.65999... vs (0.99 * 1000 - 0.33 * 1000) / 1000 = 0.66 +_DIGITS = 100000 + + +@DeveloperAPI +@dataclass +class FixedAcquiredResources(AcquiredResources): + bundles: List[Dict[str, float]] + + def _annotate_remote_entity( + self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int + ) -> RemoteRayEntity: + bundle = bundle.copy() + num_cpus = bundle.pop("CPU", 0) + num_gpus = bundle.pop("GPU", 0) + memory = bundle.pop("memory", 0.0) + + return entity.options( + num_cpus=num_cpus, + num_gpus=num_gpus, + memory=memory, + resources=bundle, + ) + + +@DeveloperAPI +class FixedResourceManager(ResourceManager): + """Fixed budget based resource manager. + + This resource manager keeps track of a fixed set of resources. When resources + are acquired, they are subtracted from the budget. When resources are freed, + they are added back to the budget. + + The resource manager still requires resources to be requested before they become + available. However, because the resource requests are virtual, this will not + trigger autoscaling. + + Additionally, resources are not reserved on request, only on acquisition. Thus, + acquiring a resource can change the availability of other requests. Note that + this behavior may be changed in future implementations. + + The fixed resource manager does not support placement strategies. Using + ``STRICT_SPREAD`` will result in an error. ``STRICT_PACK`` will succeed only + within a placement group bundle. All other placement group arguments will be + ignored. + + Args: + total_resources: Budget of resources to manage. Defaults to all available + resources in the current task or all cluster resources (if outside a task). + + """ + + _resource_cls: AcquiredResources = FixedAcquiredResources + + def __init__(self, total_resources: Optional[Dict[str, float]] = None): + rtc = ray.get_runtime_context() + + if not total_resources: + if rtc.worker.mode in {None, SCRIPT_MODE, LOCAL_MODE}: + total_resources = ray.cluster_resources() + else: + total_resources = rtc.get_assigned_resources() + + # If we are in a placement group, all of our resources will be in a bundle + # and thus fulfill requirements of STRICT_PACK - but only if child tasks + # are captured by the pg. + self._allow_strict_pack = ( + ray.util.get_current_placement_group() is not None + and rtc.should_capture_child_tasks_in_placement_group + ) + + self._total_resources = total_resources + self._requested_resources = [] + self._used_resources = [] + + @property + def _available_resources(self) -> Dict[str, float]: + available_resources = self._total_resources.copy() + + for used_resources in self._used_resources: + all_resources = used_resources.required_resources + for k, v in all_resources.items(): + available_resources[k] = ( + available_resources[k] * _DIGITS - v * _DIGITS + ) / _DIGITS + return available_resources + + def request_resources(self, resource_request: ResourceRequest): + if resource_request.strategy == "STRICT_SPREAD" or ( + not self._allow_strict_pack and resource_request.strategy == "STRICT_PACK" + ): + raise RuntimeError( + f"Requested a resource with placement strategy " + f"{resource_request.strategy}, but this cannot be fulfilled by a " + f"FixedResourceManager. In a nested setting, please set the inner " + f"placement strategy to be less restrictive (i.e. no STRICT_ strategy)." + ) + + self._requested_resources.append(resource_request) + + def cancel_resource_request(self, resource_request: ResourceRequest): + self._requested_resources.remove(resource_request) + + def has_resources_ready(self, resource_request: ResourceRequest) -> bool: + if resource_request not in self._requested_resources: + return False + + available_resources = self._available_resources + all_resources = resource_request.required_resources + for k, v in all_resources.items(): + if available_resources.get(k, 0.0) < v: + return False + return True + + def acquire_resources( + self, resource_request: ResourceRequest + ) -> Optional[AcquiredResources]: + if not self.has_resources_ready(resource_request): + return None + + self._used_resources.append(resource_request) + return self._resource_cls( + bundles=resource_request.bundles, resource_request=resource_request + ) + + def free_resources(self, acquired_resource: AcquiredResources): + resources = acquired_resource.resource_request + self._used_resources.remove(resources) + + def clear(self): + # Reset internal state + self._requested_resources = [] + self._used_resources = [] diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/placement_group.py b/lib/python3.12/site-packages/ray/air/execution/resources/placement_group.py new file mode 100644 index 0000000000000000000000000000000000000000..d804fc3750a38129fe17f718863b10599e4e170b --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/resources/placement_group.py @@ -0,0 +1,214 @@ +import time +from collections import defaultdict +from dataclasses import dataclass +from typing import Dict, List, Optional, Set + +import ray +from ray.air.execution.resources.request import ( + AcquiredResources, + RemoteRayEntity, + ResourceRequest, +) +from ray.air.execution.resources.resource_manager import ResourceManager +from ray.util.annotations import DeveloperAPI +from ray.util.placement_group import PlacementGroup, remove_placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + +@DeveloperAPI +@dataclass +class PlacementGroupAcquiredResources(AcquiredResources): + placement_group: PlacementGroup + + def _annotate_remote_entity( + self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int + ) -> RemoteRayEntity: + bundle = bundle.copy() + num_cpus = bundle.pop("CPU", 0) + num_gpus = bundle.pop("GPU", 0) + memory = bundle.pop("memory", 0.0) + + return entity.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=self.placement_group, + placement_group_bundle_index=bundle_index, + placement_group_capture_child_tasks=True, + ), + num_cpus=num_cpus, + num_gpus=num_gpus, + memory=memory, + resources=bundle, + ) + + +@DeveloperAPI +class PlacementGroupResourceManager(ResourceManager): + """Resource manager using placement groups as the resource backend. + + This manager will use placement groups to fulfill resource requests. Requesting + a resource will schedule the placement group. Acquiring a resource will + return a ``PlacementGroupAcquiredResources`` that can be used to schedule + Ray tasks and actors on the placement group. Freeing an acquired resource + will destroy the associated placement group. + + Ray core does not emit events when resources are available. Instead, the + scheduling state has to be periodically updated. + + Per default, placement group scheduling state is refreshed every time when + resource state is inquired, but not more often than once every ``update_interval_s`` + seconds. Alternatively, staging futures can be retrieved (and awaited) with + ``get_resource_futures()`` and state update can be force with ``update_state()``. + + Args: + update_interval_s: Minimum interval in seconds between updating scheduling + state of placement groups. + + """ + + _resource_cls: AcquiredResources = PlacementGroupAcquiredResources + + def __init__(self, update_interval_s: float = 0.1): + # Internally, the placement group lifecycle is like this: + # - Resources are requested with ``request_resources()`` + # - A placement group is scheduled ("staged") + # - A ``PlacementGroup.ready()`` future is scheduled ("staging future") + # - We update the scheduling state when we need to + # (e.g. when ``has_resources_ready()`` is called) + # - When staging futures resolve, a placement group is moved from "staging" + # to "ready" + # - When a resource request is canceled, we remove a placement group from + # "staging". If there are not staged placement groups + # (because they are already "ready"), we remove one from "ready" instead. + # - When a resource is acquired, the pg is removed from "ready" and moved + # to "acquired" + # - When a resource is freed, the pg is removed from "acquired" and destroyed + + # Mapping of placement group to request + self._pg_to_request: Dict[PlacementGroup, ResourceRequest] = {} + + # PGs that are staged but not "ready", yet (i.e. not CREATED) + self._request_to_staged_pgs: Dict[ + ResourceRequest, Set[PlacementGroup] + ] = defaultdict(set) + + # PGs that are CREATED and can be used by tasks and actors + self._request_to_ready_pgs: Dict[ + ResourceRequest, Set[PlacementGroup] + ] = defaultdict(set) + + # Staging futures used to update internal state. + # We keep a double mapping here for better lookup efficiency. + self._staging_future_to_pg: Dict[ray.ObjectRef, PlacementGroup] = dict() + self._pg_to_staging_future: Dict[PlacementGroup, ray.ObjectRef] = dict() + + # Set of acquired PGs. We keep track of these here to make sure we + # only free PGs that this manager managed. + self._acquired_pgs: Set[PlacementGroup] = set() + + # Minimum time between updates of the internal state + self.update_interval_s = update_interval_s + self._last_update = time.monotonic() - self.update_interval_s - 1 + + def get_resource_futures(self) -> List[ray.ObjectRef]: + return list(self._staging_future_to_pg.keys()) + + def _maybe_update_state(self): + now = time.monotonic() + if now > self._last_update + self.update_interval_s: + self.update_state() + + def update_state(self): + ready, not_ready = ray.wait( + list(self._staging_future_to_pg.keys()), + num_returns=len(self._staging_future_to_pg), + timeout=0, + ) + for future in ready: + # Remove staging future + pg = self._staging_future_to_pg.pop(future) + self._pg_to_staging_future.pop(pg) + # Fetch resource request + request = self._pg_to_request[pg] + # Remove from staging, add to ready + self._request_to_staged_pgs[request].remove(pg) + self._request_to_ready_pgs[request].add(pg) + self._last_update = time.monotonic() + + def request_resources(self, resource_request: ResourceRequest): + pg = resource_request.to_placement_group() + self._pg_to_request[pg] = resource_request + self._request_to_staged_pgs[resource_request].add(pg) + + future = pg.ready() + self._staging_future_to_pg[future] = pg + self._pg_to_staging_future[pg] = future + + def cancel_resource_request(self, resource_request: ResourceRequest): + if self._request_to_staged_pgs[resource_request]: + pg = self._request_to_staged_pgs[resource_request].pop() + + # PG was staging + future = self._pg_to_staging_future.pop(pg) + self._staging_future_to_pg.pop(future) + + # Cancel the pg.ready task. + # Otherwise, it will be pending node assignment forever. + ray.cancel(future) + else: + # PG might be ready + pg = self._request_to_ready_pgs[resource_request].pop() + if not pg: + raise RuntimeError( + "Cannot cancel resource request: No placement group was " + f"staged or is ready. Make sure to not cancel more resource " + f"requests than you've created. Request: {resource_request}" + ) + + self._pg_to_request.pop(pg) + ray.util.remove_placement_group(pg) + + def has_resources_ready(self, resource_request: ResourceRequest) -> bool: + if not bool(len(self._request_to_ready_pgs[resource_request])): + # Only update state if needed + self._maybe_update_state() + + return bool(len(self._request_to_ready_pgs[resource_request])) + + def acquire_resources( + self, resource_request: ResourceRequest + ) -> Optional[PlacementGroupAcquiredResources]: + if not self.has_resources_ready(resource_request): + return None + + pg = self._request_to_ready_pgs[resource_request].pop() + self._acquired_pgs.add(pg) + + return self._resource_cls(placement_group=pg, resource_request=resource_request) + + def free_resources(self, acquired_resource: PlacementGroupAcquiredResources): + pg = acquired_resource.placement_group + + self._acquired_pgs.remove(pg) + remove_placement_group(pg) + self._pg_to_request.pop(pg) + + def clear(self): + if not ray.is_initialized(): + return + + for staged_pgs in self._request_to_staged_pgs.values(): + for staged_pg in staged_pgs: + remove_placement_group(staged_pg) + + for ready_pgs in self._request_to_ready_pgs.values(): + for ready_pg in ready_pgs: + remove_placement_group(ready_pg) + + for acquired_pg in self._acquired_pgs: + remove_placement_group(acquired_pg) + + # Reset internal state + self.__init__(update_interval_s=self.update_interval_s) + + def __del__(self): + self.clear() diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/request.py b/lib/python3.12/site-packages/ray/air/execution/resources/request.py new file mode 100644 index 0000000000000000000000000000000000000000..7777fe297b6e0c7f291349435563e7ac7d3dac2a --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/resources/request.py @@ -0,0 +1,255 @@ +import abc +import json +from copy import deepcopy +from dataclasses import dataclass +from inspect import signature +from typing import Dict, List, Union + +import ray +from ray.util import placement_group +from ray.util.annotations import DeveloperAPI + +RemoteRayEntity = Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass] + + +def _sum_bundles(bundles: List[Dict[str, float]]) -> Dict[str, float]: + """Sum all resources in a list of resource bundles. + + Args: + bundles: List of resource bundles. + + Returns: Dict containing all resources summed up. + """ + resources = {} + for bundle in bundles: + for k, v in bundle.items(): + resources[k] = resources.get(k, 0) + v + return resources + + +@DeveloperAPI +class ResourceRequest: + """Request for resources. + + This class is used to define a resource request. A resource request comprises one + or more bundles of resources and instructions on the scheduling behavior. + + The resource request can be submitted to a resource manager, which will + schedule the resources. Depending on the resource backend, this may instruct + Ray to scale up (autoscaling). + + Resource requests are compatible with the most fine-grained low-level resource + backend, which are Ray placement groups. + + Args: + bundles: A list of bundles which represent the resources requirements. + E.g. ``[{"CPU": 1, "GPU": 1}]``. + strategy: The scheduling strategy to acquire the bundles. + + - "PACK": Packs Bundles into as few nodes as possible. + - "SPREAD": Places Bundles across distinct nodes as even as possible. + - "STRICT_PACK": Packs Bundles into one node. The group is + not allowed to span multiple nodes. + - "STRICT_SPREAD": Packs Bundles across distinct nodes. + *args: Passed to the call of ``placement_group()``, if applicable. + **kwargs: Passed to the call of ``placement_group()``, if applicable. + + """ + + def __init__( + self, + bundles: List[Dict[str, Union[int, float]]], + strategy: str = "PACK", + *args, + **kwargs, + ): + if not bundles: + raise ValueError("Cannot initialize a ResourceRequest with zero bundles.") + + # Remove empty resource keys + self._bundles = [ + {k: float(v) for k, v in bundle.items() if v != 0} for bundle in bundles + ] + + # Check if the head bundle is empty (no resources defined or all resources + # are 0 (and thus removed in the previous step) + if not self._bundles[0]: + # This is when the head bundle doesn't need resources. + self._head_bundle_is_empty = True + self._bundles.pop(0) + + if not self._bundles: + raise ValueError( + "Cannot initialize a ResourceRequest with an empty head " + "and zero worker bundles." + ) + else: + self._head_bundle_is_empty = False + + self._strategy = strategy + self._args = args + self._kwargs = kwargs + + self._hash = None + self._bound = None + + self._bind() + + @property + def head_bundle_is_empty(self): + """Returns True if head bundle is empty while child bundles + need resources. + + This is considered an internal API within Tune. + """ + return self._head_bundle_is_empty + + @property + @DeveloperAPI + def head_cpus(self) -> float: + """Returns the number of cpus in the head bundle.""" + return 0.0 if self._head_bundle_is_empty else self._bundles[0].get("CPU", 0.0) + + @property + @DeveloperAPI + def bundles(self) -> List[Dict[str, float]]: + """Returns a deep copy of resource bundles""" + return deepcopy(self._bundles) + + @property + def required_resources(self) -> Dict[str, float]: + """Returns a dict containing the sums of all resources""" + return _sum_bundles(self._bundles) + + @property + @DeveloperAPI + def strategy(self) -> str: + """Returns the placement strategy""" + return self._strategy + + def _bind(self): + """Bind the args and kwargs to the `placement_group()` signature. + + We bind the args and kwargs, so we can compare equality of two resource + requests. The main reason for this is that the `placement_group()` API + can evolve independently from the ResourceRequest API (e.g. adding new + arguments). Then, `ResourceRequest(bundles, strategy, arg=arg)` should + be the same as `ResourceRequest(bundles, strategy, arg)`. + """ + sig = signature(placement_group) + try: + self._bound = sig.bind( + self._bundles, self._strategy, *self._args, **self._kwargs + ) + except Exception as exc: + raise RuntimeError( + "Invalid definition for resource request. Please check " + "that you passed valid arguments to the ResourceRequest " + "object." + ) from exc + + def to_placement_group(self): + return placement_group(*self._bound.args, **self._bound.kwargs) + + def __eq__(self, other: "ResourceRequest"): + return ( + isinstance(other, ResourceRequest) + and self._bound == other._bound + and self.head_bundle_is_empty == other.head_bundle_is_empty + ) + + def __hash__(self): + if not self._hash: + # Cache hash + self._hash = hash( + json.dumps( + {"args": self._bound.args, "kwargs": self._bound.kwargs}, + sort_keys=True, + indent=0, + ensure_ascii=True, + ) + ) + return self._hash + + def __getstate__(self): + state = self.__dict__.copy() + state.pop("_hash", None) + state.pop("_bound", None) + return state + + def __setstate__(self, state): + self.__dict__.update(state) + self._hash = None + self._bound = None + self._bind() + + def __repr__(self) -> str: + return ( + f"" + ) + + +@DeveloperAPI +@dataclass +class AcquiredResources(abc.ABC): + """Base class for resources that have been acquired. + + Acquired resources can be associated to Ray objects, which can then be + scheduled using these resources. + + Internally this can point e.g. to a placement group, a placement + group bundle index, or just raw resources. + + The main API is the `annotate_remote_entities` method. This will associate + remote Ray objects (tasks and actors) with the acquired resources by setting + the Ray remote options to use the acquired resources. + """ + + resource_request: ResourceRequest + + def annotate_remote_entities( + self, entities: List[RemoteRayEntity] + ) -> List[Union[RemoteRayEntity]]: + """Return remote ray entities (tasks/actors) to use the acquired resources. + + The first entity will be associated with the first bundle, the second + entity will be associated with the second bundle, etc. + + Args: + entities: Remote Ray entities to annotate with the acquired resources. + """ + bundles = self.resource_request.bundles + + # Also count the empty head bundle as a bundle + num_bundles = len(bundles) + int(self.resource_request.head_bundle_is_empty) + + if len(entities) > num_bundles: + raise RuntimeError( + f"The number of callables to annotate ({len(entities)}) cannot " + f"exceed the number of available bundles ({num_bundles})." + ) + + annotated = [] + + if self.resource_request.head_bundle_is_empty: + # The empty head bundle is place on the first bundle index with empty + # resources. + annotated.append( + self._annotate_remote_entity(entities[0], {}, bundle_index=0) + ) + + # Shift the remaining entities + entities = entities[1:] + + for i, (entity, bundle) in enumerate(zip(entities, bundles)): + annotated.append( + self._annotate_remote_entity(entity, bundle, bundle_index=i) + ) + + return annotated + + def _annotate_remote_entity( + self, entity: RemoteRayEntity, bundle: Dict[str, float], bundle_index: int + ) -> RemoteRayEntity: + raise NotImplementedError diff --git a/lib/python3.12/site-packages/ray/air/execution/resources/resource_manager.py b/lib/python3.12/site-packages/ray/air/execution/resources/resource_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..daa2cdd69215157e21a36ea937bb9399d5c759b3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/execution/resources/resource_manager.py @@ -0,0 +1,155 @@ +import abc +from typing import List, Optional + +import ray +from ray.air.execution.resources.request import AcquiredResources, ResourceRequest +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class ResourceManager(abc.ABC): + """Resource manager interface. + + A resource manager can be used to request resources from a Ray cluster and + allocate them to remote Ray tasks or actors. + + Resources have to be requested before they can be acquired. + + Resources managed by the resource manager can be in three states: + + 1. "Requested": The resources have been requested but are not yet available to + schedule remote Ray objects. The resource request may trigger autoscaling, + and can be cancelled if no longer needed. + 2. "Ready": The requested resources are now available to schedule remote Ray + objects. They can be acquired and subsequently used remote Ray objects. + The resource request can still be cancelled if no longer needed. + 3. "Acquired": The resources have been acquired by a caller to use for scheduling + remote Ray objects. Note that it is the responsibility of the caller to + schedule the Ray objects with these resources. + The associated resource request has been completed and can no longer be + cancelled. The acquired resources can be freed by the resource manager when + they are no longer used. + + The flow is as follows: + + .. code-block:: python + + # Create resource manager + resource_manager = ResourceManager() + + # Create resource request + resource_request = ResourceRequest([{"CPU": 4}]) + + # Pass to resource manager + resource_manager.request_resources(resource_request) + + # Wait until ready + while not resource_manager.has_resources_ready(resource_request): + time.sleep(1) + + # Once ready, acquire resources + acquired_resource = resource_manager.acquire_resources(resource_request) + + # Bind to remote task or actor + annotated_remote_fn = acquired_resource.annotate_remote_entities( + [remote_fn]) + + # Run remote function. This will use the acquired resources + ray.get(annotated_remote_fn.remote()) + + # After using the resources, free + resource_manager.free_resources(annotated_resources) + + """ + + def request_resources(self, resource_request: ResourceRequest): + """Request resources. + + Depending on the backend, resources can trigger autoscaling. Requested + resources can be ready or not ready. Once they are "ready", they can + be acquired and used by remote Ray objects. + + Resource requests can be cancelled anytime using ``cancel_resource_request()``. + Once acquired, the resource request is removed. Acquired resources can be + freed with ``free_resources()``. + """ + raise NotImplementedError + + def cancel_resource_request(self, resource_request: ResourceRequest): + """Cancel resource request. + + Resource requests can be cancelled anytime before a resource is acquired. + Acquiring a resource will remove the associated resource request. + Acquired resources can be freed with ``free_resources()``. + """ + raise NotImplementedError + + def has_resources_ready(self, resource_request: ResourceRequest) -> bool: + """Returns True if resources for the given request are ready to be acquired.""" + raise NotImplementedError + + def acquire_resources( + self, resource_request: ResourceRequest + ) -> Optional[AcquiredResources]: + """Acquire resources. Returns None if resources are not ready to be acquired. + + Acquiring resources will remove the associated resource request. + Acquired resources can be returned with ``free_resources()``. + """ + raise NotImplementedError + + def free_resources(self, acquired_resource: AcquiredResources): + """Free acquired resources from usage and return them to the resource manager. + + Freeing resources will return the resources to the manager, but there are + no guarantees about the tasks and actors scheduled on the resources. The caller + should make sure that any references to tasks or actors scheduled on the + resources have been removed before calling ``free_resources()``. + """ + raise NotImplementedError + + def get_resource_futures(self) -> List[ray.ObjectRef]: + """Return futures for resources to await. + + Depending on the backend, we use resource futures to determine availability + of resources (e.g. placement groups) or resolution of requests. + In this case, the futures can be awaited externally by the caller. + + When a resource future resolved, the caller may call ``update_state()`` + to force the resource manager to update its internal state immediately. + """ + return [] + + def update_state(self): + """Update internal state of the resource manager. + + The resource manager may have internal state that needs periodic updating. + For instance, depending on the backend, resource futures can be awaited + externally (with ``get_resource_futures()``). + + If such a future resolved, the caller can instruct the resource + manager to update its internal state immediately. + """ + pass + + def clear(self): + """Reset internal state and clear all resources. + + Calling this method will reset the resource manager to its initialization state. + All resources will be removed. + + Clearing the state will remove tracked resources from the manager, but there are + no guarantees about the tasks and actors scheduled on the resources. The caller + should make sure that any references to tasks or actors scheduled on the + resources have been removed before calling ``clear()``. + """ + raise NotImplementedError + + def __reduce__(self): + """We disallow serialization. + + Shared resource managers should live on an actor. + """ + raise ValueError( + f"Resource managers cannot be serialized. Resource manager: {str(self)}" + ) diff --git a/lib/python3.12/site-packages/ray/air/integrations/__init__.py b/lib/python3.12/site-packages/ray/air/integrations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/integrations/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9d08277bfc82c49493d10c3a4614d3b2024c0ee Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/integrations/__pycache__/comet.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/comet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcd00410c16e2b692acf2de0a00411dde3f35d92 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/comet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/integrations/__pycache__/keras.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/keras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a5a1d004ee2fe7ab6c480d341d76576f8507b0a Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/keras.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/integrations/__pycache__/mlflow.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/mlflow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e94bbae61d4324297a7ec008145d1731cc4a5dd0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/mlflow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/integrations/__pycache__/wandb.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/wandb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6f6c38e09b01d63f281e44ea45e68c74ed6313d Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/integrations/__pycache__/wandb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/integrations/comet.py b/lib/python3.12/site-packages/ray/air/integrations/comet.py new file mode 100644 index 0000000000000000000000000000000000000000..05d1ebcbeb1a71d31ecffd0f0c5fa8c962d0f8f8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/integrations/comet.py @@ -0,0 +1,260 @@ +import os +from pathlib import Path +from typing import Dict, List + +import pyarrow.fs + +from ray.tune.experiment import Trial +from ray.tune.logger import LoggerCallback +from ray.tune.utils import flatten_dict + + +def _import_comet(): + """Try importing comet_ml. + + Used to check if comet_ml is installed and, otherwise, pass an informative + error message. + """ + if "COMET_DISABLE_AUTO_LOGGING" not in os.environ: + os.environ["COMET_DISABLE_AUTO_LOGGING"] = "1" + + try: + import comet_ml # noqa: F401 + except ImportError: + raise RuntimeError("pip install 'comet-ml' to use CometLoggerCallback") + + return comet_ml + + +class CometLoggerCallback(LoggerCallback): + """CometLoggerCallback for logging Tune results to Comet. + + Comet (https://comet.ml/site/) is a tool to manage and optimize the + entire ML lifecycle, from experiment tracking, model optimization + and dataset versioning to model production monitoring. + + This Ray Tune ``LoggerCallback`` sends metrics and parameters to + Comet for tracking. + + In order to use the CometLoggerCallback you must first install Comet + via ``pip install comet_ml`` + + Then set the following environment variables + ``export COMET_API_KEY=`` + + Alternatively, you can also pass in your API Key as an argument to the + CometLoggerCallback constructor. + + ``CometLoggerCallback(api_key=)`` + + Args: + online: Whether to make use of an Online or + Offline Experiment. Defaults to True. + tags: Tags to add to the logged Experiment. + Defaults to None. + save_checkpoints: If ``True``, model checkpoints will be saved to + Comet ML as artifacts. Defaults to ``False``. + **experiment_kwargs: Other keyword arguments will be passed to the + constructor for comet_ml.Experiment (or OfflineExperiment if + online=False). + + Please consult the Comet ML documentation for more information on the + Experiment and OfflineExperiment classes: https://comet.ml/site/ + + Example: + + .. code-block:: python + + from ray.air.integrations.comet import CometLoggerCallback + tune.run( + train, + config=config + callbacks=[CometLoggerCallback( + True, + ['tag1', 'tag2'], + workspace='my_workspace', + project_name='my_project_name' + )] + ) + + """ + + # Do not enable these auto log options unless overridden + _exclude_autolog = [ + "auto_output_logging", + "log_git_metadata", + "log_git_patch", + "log_env_cpu", + "log_env_gpu", + ] + + # Do not log these metrics. + _exclude_results = ["done", "should_checkpoint"] + + # These values should be logged as system info instead of metrics. + _system_results = ["node_ip", "hostname", "pid", "date"] + + # These values should be logged as "Other" instead of as metrics. + _other_results = ["trial_id", "experiment_id", "experiment_tag"] + + _episode_results = ["hist_stats/episode_reward", "hist_stats/episode_lengths"] + + def __init__( + self, + online: bool = True, + tags: List[str] = None, + save_checkpoints: bool = False, + **experiment_kwargs, + ): + _import_comet() + self.online = online + self.tags = tags + self.save_checkpoints = save_checkpoints + self.experiment_kwargs = experiment_kwargs + + # Disable the specific autologging features that cause throttling. + self._configure_experiment_defaults() + + # Mapping from trial to experiment object. + self._trial_experiments = {} + + self._to_exclude = self._exclude_results.copy() + self._to_system = self._system_results.copy() + self._to_other = self._other_results.copy() + self._to_episodes = self._episode_results.copy() + + def _configure_experiment_defaults(self): + """Disable the specific autologging features that cause throttling.""" + for option in self._exclude_autolog: + if not self.experiment_kwargs.get(option): + self.experiment_kwargs[option] = False + + def _check_key_name(self, key: str, item: str) -> bool: + """ + Check if key argument is equal to item argument or starts with item and + a forward slash. Used for parsing trial result dictionary into ignored + keys, system metrics, episode logs, etc. + """ + return key.startswith(item + "/") or key == item + + def log_trial_start(self, trial: "Trial"): + """ + Initialize an Experiment (or OfflineExperiment if self.online=False) + and start logging to Comet. + + Args: + trial: Trial object. + + """ + _import_comet() # is this necessary? + from comet_ml import Experiment, OfflineExperiment + from comet_ml.config import set_global_experiment + + if trial not in self._trial_experiments: + experiment_cls = Experiment if self.online else OfflineExperiment + experiment = experiment_cls(**self.experiment_kwargs) + self._trial_experiments[trial] = experiment + # Set global experiment to None to allow for multiple experiments. + set_global_experiment(None) + else: + experiment = self._trial_experiments[trial] + + experiment.set_name(str(trial)) + experiment.add_tags(self.tags) + experiment.log_other("Created from", "Ray") + + config = trial.config.copy() + config.pop("callbacks", None) + experiment.log_parameters(config) + + def log_trial_result(self, iteration: int, trial: "Trial", result: Dict): + """ + Log the current result of a Trial upon each iteration. + """ + if trial not in self._trial_experiments: + self.log_trial_start(trial) + experiment = self._trial_experiments[trial] + step = result["training_iteration"] + + config_update = result.pop("config", {}).copy() + config_update.pop("callbacks", None) # Remove callbacks + for k, v in config_update.items(): + if isinstance(v, dict): + experiment.log_parameters(flatten_dict({k: v}, "/"), step=step) + + else: + experiment.log_parameter(k, v, step=step) + + other_logs = {} + metric_logs = {} + system_logs = {} + episode_logs = {} + + flat_result = flatten_dict(result, delimiter="/") + for k, v in flat_result.items(): + if any(self._check_key_name(k, item) for item in self._to_exclude): + continue + + if any(self._check_key_name(k, item) for item in self._to_other): + other_logs[k] = v + + elif any(self._check_key_name(k, item) for item in self._to_system): + system_logs[k] = v + + elif any(self._check_key_name(k, item) for item in self._to_episodes): + episode_logs[k] = v + + else: + metric_logs[k] = v + + experiment.log_others(other_logs) + experiment.log_metrics(metric_logs, step=step) + + for k, v in system_logs.items(): + experiment.log_system_info(k, v) + + for k, v in episode_logs.items(): + experiment.log_curve(k, x=range(len(v)), y=v, step=step) + + def log_trial_save(self, trial: "Trial"): + comet_ml = _import_comet() + + if self.save_checkpoints and trial.checkpoint: + experiment = self._trial_experiments[trial] + + artifact = comet_ml.Artifact( + name=f"checkpoint_{(str(trial))}", artifact_type="model" + ) + + checkpoint_root = None + + if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem): + checkpoint_root = trial.checkpoint.path + # Todo: For other filesystems, we may want to use + # artifact.add_remote() instead. However, this requires a full + # URI. We can add this once we have a way to retrieve it. + + # Walk through checkpoint directory and add all files to artifact + if checkpoint_root: + for root, dirs, files in os.walk(checkpoint_root): + rel_root = os.path.relpath(root, checkpoint_root) + for file in files: + local_file = Path(checkpoint_root, rel_root, file).as_posix() + logical_path = Path(rel_root, file).as_posix() + + # Strip leading `./` + if logical_path.startswith("./"): + logical_path = logical_path[2:] + + artifact.add(local_file, logical_path=logical_path) + + experiment.log_artifact(artifact) + + def log_trial_end(self, trial: "Trial", failed: bool = False): + self._trial_experiments[trial].end() + del self._trial_experiments[trial] + + def __del__(self): + for trial, experiment in self._trial_experiments.items(): + experiment.end() + self._trial_experiments = {} diff --git a/lib/python3.12/site-packages/ray/air/integrations/keras.py b/lib/python3.12/site-packages/ray/air/integrations/keras.py new file mode 100644 index 0000000000000000000000000000000000000000..677213e73dedf7fc16a6199c7ac133b35c4b3827 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/integrations/keras.py @@ -0,0 +1,185 @@ +import shutil +from typing import Dict, List, Optional, Union + +from tensorflow.keras.callbacks import Callback as KerasCallback + +import ray +from ray.train.tensorflow import TensorflowCheckpoint +from ray.util.annotations import PublicAPI + + +class _Callback(KerasCallback): + """Base class for Air's Keras callbacks.""" + + _allowed = [ + "epoch_begin", + "epoch_end", + "train_batch_begin", + "train_batch_end", + "test_batch_begin", + "test_batch_end", + "predict_batch_begin", + "predict_batch_end", + "train_begin", + "train_end", + "test_begin", + "test_end", + "predict_begin", + "predict_end", + ] + + def __init__(self, on: Union[str, List[str]] = "validation_end"): + super(_Callback, self).__init__() + + if not isinstance(on, list): + on = [on] + if any(w not in self._allowed for w in on): + raise ValueError( + "Invalid trigger time selected: {}. Must be one of {}".format( + on, self._allowed + ) + ) + self._on = on + + def _handle(self, logs: Dict, when: str): + raise NotImplementedError + + def on_epoch_begin(self, epoch, logs=None): + if "epoch_begin" in self._on: + self._handle(logs, "epoch_begin") + + def on_epoch_end(self, epoch, logs=None): + if "epoch_end" in self._on: + self._handle(logs, "epoch_end") + + def on_train_batch_begin(self, batch, logs=None): + if "train_batch_begin" in self._on: + self._handle(logs, "train_batch_begin") + + def on_train_batch_end(self, batch, logs=None): + if "train_batch_end" in self._on: + self._handle(logs, "train_batch_end") + + def on_test_batch_begin(self, batch, logs=None): + if "test_batch_begin" in self._on: + self._handle(logs, "test_batch_begin") + + def on_test_batch_end(self, batch, logs=None): + if "test_batch_end" in self._on: + self._handle(logs, "test_batch_end") + + def on_predict_batch_begin(self, batch, logs=None): + if "predict_batch_begin" in self._on: + self._handle(logs, "predict_batch_begin") + + def on_predict_batch_end(self, batch, logs=None): + if "predict_batch_end" in self._on: + self._handle(logs, "predict_batch_end") + + def on_train_begin(self, logs=None): + if "train_begin" in self._on: + self._handle(logs, "train_begin") + + def on_train_end(self, logs=None): + if "train_end" in self._on: + self._handle(logs, "train_end") + + def on_test_begin(self, logs=None): + if "test_begin" in self._on: + self._handle(logs, "test_begin") + + def on_test_end(self, logs=None): + if "test_end" in self._on: + self._handle(logs, "test_end") + + def on_predict_begin(self, logs=None): + if "predict_begin" in self._on: + self._handle(logs, "predict_begin") + + def on_predict_end(self, logs=None): + if "predict_end" in self._on: + self._handle(logs, "predict_end") + + +@PublicAPI(stability="alpha") +class ReportCheckpointCallback(_Callback): + """Keras callback for Ray Train reporting and checkpointing. + + .. note:: + Metrics are always reported with checkpoints, even if the event isn't specified + in ``report_metrics_on``. + + Example: + .. code-block:: python + + ############# Using it in TrainSession ############### + from ray.air.integrations.keras import ReportCheckpointCallback + def train_loop_per_worker(): + strategy = tf.distribute.MultiWorkerMirroredStrategy() + with strategy.scope(): + model = build_model() + + model.fit(dataset_shard, callbacks=[ReportCheckpointCallback()]) + + Args: + metrics: Metrics to report. If this is a list, each item describes + the metric key reported to Keras, and it's reported under the + same name. If this is a dict, each key is the name reported + and the respective value is the metric key reported to Keras. + If this is None, all Keras logs are reported. + report_metrics_on: When to report metrics. Must be one of + the Keras event hooks (less the ``on_``), e.g. + "train_start" or "predict_end". Defaults to "epoch_end". + checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks + (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to + "epoch_end". + """ + + def __init__( + self, + checkpoint_on: Union[str, List[str]] = "epoch_end", + report_metrics_on: Union[str, List[str]] = "epoch_end", + metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, + ): + if isinstance(checkpoint_on, str): + checkpoint_on = [checkpoint_on] + if isinstance(report_metrics_on, str): + report_metrics_on = [report_metrics_on] + + on = list(set(checkpoint_on + report_metrics_on)) + super().__init__(on=on) + + self._checkpoint_on: List[str] = checkpoint_on + self._report_metrics_on: List[str] = report_metrics_on + self._metrics = metrics + + def _handle(self, logs: Dict, when: str): + assert when in self._checkpoint_on or when in self._report_metrics_on + + metrics = self._get_reported_metrics(logs) + + should_checkpoint = when in self._checkpoint_on + if should_checkpoint: + checkpoint = TensorflowCheckpoint.from_model(self.model) + ray.train.report(metrics, checkpoint=checkpoint) + # Clean up temporary checkpoint + shutil.rmtree(checkpoint.path, ignore_errors=True) + else: + ray.train.report(metrics, checkpoint=None) + + def _get_reported_metrics(self, logs: Dict) -> Dict: + assert isinstance(self._metrics, (type(None), str, list, dict)) + + if self._metrics is None: + reported_metrics = logs + elif isinstance(self._metrics, str): + reported_metrics = {self._metrics: logs[self._metrics]} + elif isinstance(self._metrics, list): + reported_metrics = {metric: logs[metric] for metric in self._metrics} + elif isinstance(self._metrics, dict): + reported_metrics = { + key: logs[metric] for key, metric in self._metrics.items() + } + + assert isinstance(reported_metrics, dict) + return reported_metrics diff --git a/lib/python3.12/site-packages/ray/air/integrations/mlflow.py b/lib/python3.12/site-packages/ray/air/integrations/mlflow.py new file mode 100644 index 0000000000000000000000000000000000000000..ef76a221133501d4cbc3c162ade51eea715a6d24 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/integrations/mlflow.py @@ -0,0 +1,337 @@ +import logging +from types import ModuleType +from typing import Dict, Optional, Union + +import ray +from ray.air._internal import usage as air_usage +from ray.air._internal.mlflow import _MLflowLoggerUtil +from ray.air.constants import TRAINING_ITERATION +from ray.tune.experiment import Trial +from ray.tune.logger import LoggerCallback +from ray.tune.result import TIMESTEPS_TOTAL +from ray.util.annotations import PublicAPI + +try: + import mlflow +except ImportError: + mlflow = None + + +logger = logging.getLogger(__name__) + + +class _NoopModule: + def __getattr__(self, item): + return _NoopModule() + + def __call__(self, *args, **kwargs): + return None + + +@PublicAPI(stability="alpha") +def setup_mlflow( + config: Optional[Dict] = None, + tracking_uri: Optional[str] = None, + registry_uri: Optional[str] = None, + experiment_id: Optional[str] = None, + experiment_name: Optional[str] = None, + tracking_token: Optional[str] = None, + artifact_location: Optional[str] = None, + run_name: Optional[str] = None, + create_experiment_if_not_exists: bool = False, + tags: Optional[Dict] = None, + rank_zero_only: bool = True, +) -> Union[ModuleType, _NoopModule]: + """Set up a MLflow session. + + This function can be used to initialize an MLflow session in a + (distributed) training or tuning run. The session will be created on the trainable. + + By default, the MLflow experiment ID is the Ray trial ID and the + MLlflow experiment name is the Ray trial name. These settings can be overwritten by + passing the respective keyword arguments. + + The ``config`` dict is automatically logged as the run parameters (excluding the + mlflow settings). + + In distributed training with Ray Train, only the zero-rank worker will initialize + mlflow. All other workers will return a noop client, so that logging is not + duplicated in a distributed run. This can be disabled by passing + ``rank_zero_only=False``, which will then initialize mlflow in every training + worker. + + This function will return the ``mlflow`` module or a noop module for + non-rank zero workers ``if rank_zero_only=True``. By using + ``mlflow = setup_mlflow(config)`` you can ensure that only the rank zero worker + calls the mlflow API. + + Args: + config: Configuration dict to be logged to mlflow as parameters. + tracking_uri: The tracking URI for MLflow tracking. If using + Tune in a multi-node setting, make sure to use a remote server for + tracking. + registry_uri: The registry URI for the MLflow model registry. + experiment_id: The id of an already created MLflow experiment. + All logs from all trials in ``tune.Tuner()`` will be reported to this + experiment. If this is not provided or the experiment with this + id does not exist, you must provide an``experiment_name``. This + parameter takes precedence over ``experiment_name``. + experiment_name: The name of an already existing MLflow + experiment. All logs from all trials in ``tune.Tuner()`` will be + reported to this experiment. If this is not provided, you must + provide a valid ``experiment_id``. + tracking_token: A token to use for HTTP authentication when + logging to a remote tracking server. This is useful when you + want to log to a Databricks server, for example. This value will + be used to set the MLFLOW_TRACKING_TOKEN environment variable on + all the remote training processes. + artifact_location: The location to store run artifacts. + If not provided, MLFlow picks an appropriate default. + Ignored if experiment already exists. + run_name: Name of the new MLflow run that will be created. + If not set, will default to the ``experiment_name``. + create_experiment_if_not_exists: Whether to create an + experiment with the provided name if it does not already + exist. Defaults to False. + tags: Tags to set for the new run. + rank_zero_only: If True, will return an initialized session only for the + rank 0 worker in distributed training. If False, will initialize a + session for all workers. Defaults to True. + + Example: + + Per default, you can just call ``setup_mlflow`` and continue to use + MLflow like you would normally do: + + .. code-block:: python + + from ray.air.integrations.mlflow import setup_mlflow + + def training_loop(config): + mlflow = setup_mlflow(config) + # ... + mlflow.log_metric(key="loss", val=0.123, step=0) + + In distributed data parallel training, you can utilize the return value of + ``setup_mlflow``. This will make sure it is only invoked on the first worker + in distributed training runs. + + .. code-block:: python + + from ray.air.integrations.mlflow import setup_mlflow + + def training_loop(config): + mlflow = setup_mlflow(config) + # ... + mlflow.log_metric(key="loss", val=0.123, step=0) + + + You can also use MlFlow's autologging feature if using a training + framework like Pytorch Lightning, XGBoost, etc. More information can be + found here + (https://mlflow.org/docs/latest/tracking.html#automatic-logging). + + .. code-block:: python + + from ray.air.integrations.mlflow import setup_mlflow + + def train_fn(config): + mlflow = setup_mlflow(config) + mlflow.autolog() + xgboost_results = xgb.train(config, ...) + + """ + if not mlflow: + raise RuntimeError( + "mlflow was not found - please install with `pip install mlflow`" + ) + + try: + train_context = ray.train.get_context() + + # Do a try-catch here if we are not in a train session + if rank_zero_only and train_context.get_world_rank() != 0: + return _NoopModule() + + default_trial_id = train_context.get_trial_id() + default_trial_name = train_context.get_trial_name() + + except RuntimeError: + default_trial_id = None + default_trial_name = None + + _config = config.copy() if config else {} + + experiment_id = experiment_id or default_trial_id + experiment_name = experiment_name or default_trial_name + + # Setup mlflow + mlflow_util = _MLflowLoggerUtil() + mlflow_util.setup_mlflow( + tracking_uri=tracking_uri, + registry_uri=registry_uri, + experiment_id=experiment_id, + experiment_name=experiment_name, + tracking_token=tracking_token, + artifact_location=artifact_location, + create_experiment_if_not_exists=create_experiment_if_not_exists, + ) + + mlflow_util.start_run( + run_name=run_name or experiment_name, + tags=tags, + set_active=True, + ) + mlflow_util.log_params(_config) + + # Record `setup_mlflow` usage when everything has setup successfully. + air_usage.tag_setup_mlflow() + + return mlflow_util._mlflow + + +class MLflowLoggerCallback(LoggerCallback): + """MLflow Logger to automatically log Tune results and config to MLflow. + + MLflow (https://mlflow.org) Tracking is an open source library for + recording and querying experiments. This Ray Tune ``LoggerCallback`` + sends information (config parameters, training results & metrics, + and artifacts) to MLflow for automatic experiment tracking. + + Keep in mind that the callback will open an MLflow session on the driver and + not on the trainable. Therefore, it is not possible to call MLflow functions + like ``mlflow.log_figure()`` inside the trainable as there is no MLflow session + on the trainable. For more fine grained control, use + :func:`ray.air.integrations.mlflow.setup_mlflow`. + + Args: + tracking_uri: The tracking URI for where to manage experiments + and runs. This can either be a local file path or a remote server. + This arg gets passed directly to mlflow + initialization. When using Tune in a multi-node setting, make sure + to set this to a remote server and not a local file path. + registry_uri: The registry URI that gets passed directly to + mlflow initialization. + experiment_name: The experiment name to use for this Tune run. + If the experiment with the name already exists with MLflow, + it will be reused. If not, a new experiment will be created with + that name. + tags: An optional dictionary of string keys and values to set + as tags on the run + tracking_token: Tracking token used to authenticate with MLflow. + save_artifact: If set to True, automatically save the entire + contents of the Tune local_dir as an artifact to the + corresponding run in MlFlow. + log_params_on_trial_end: If set to True, log parameters to MLflow + at the end of the trial instead of at the beginning + + Example: + + .. code-block:: python + + from ray.air.integrations.mlflow import MLflowLoggerCallback + + tags = { "user_name" : "John", + "git_commit_hash" : "abc123"} + + tune.run( + train_fn, + config={ + # define search space here + "parameter_1": tune.choice([1, 2, 3]), + "parameter_2": tune.choice([4, 5, 6]), + }, + callbacks=[MLflowLoggerCallback( + experiment_name="experiment1", + tags=tags, + save_artifact=True, + log_params_on_trial_end=True)]) + + """ + + def __init__( + self, + tracking_uri: Optional[str] = None, + *, + registry_uri: Optional[str] = None, + experiment_name: Optional[str] = None, + tags: Optional[Dict] = None, + tracking_token: Optional[str] = None, + save_artifact: bool = False, + log_params_on_trial_end: bool = False, + ): + + self.tracking_uri = tracking_uri + self.registry_uri = registry_uri + self.experiment_name = experiment_name + self.tags = tags + self.tracking_token = tracking_token + self.should_save_artifact = save_artifact + self.log_params_on_trial_end = log_params_on_trial_end + + self.mlflow_util = _MLflowLoggerUtil() + + if ray.util.client.ray.is_connected(): + logger.warning( + "When using MLflowLoggerCallback with Ray Client, " + "it is recommended to use a remote tracking " + "server. If you are using a MLflow tracking server " + "backed by the local filesystem, then it must be " + "setup on the server side and not on the client " + "side." + ) + + def setup(self, *args, **kwargs): + # Setup the mlflow logging util. + self.mlflow_util.setup_mlflow( + tracking_uri=self.tracking_uri, + registry_uri=self.registry_uri, + experiment_name=self.experiment_name, + tracking_token=self.tracking_token, + ) + + if self.tags is None: + # Create empty dictionary for tags if not given explicitly + self.tags = {} + + self._trial_runs = {} + + def log_trial_start(self, trial: "Trial"): + # Create run if not already exists. + if trial not in self._trial_runs: + + # Set trial name in tags + tags = self.tags.copy() + tags["trial_name"] = str(trial) + + run = self.mlflow_util.start_run(tags=tags, run_name=str(trial)) + self._trial_runs[trial] = run.info.run_id + + run_id = self._trial_runs[trial] + + # Log the config parameters. + config = trial.config + if not self.log_params_on_trial_end: + self.mlflow_util.log_params(run_id=run_id, params_to_log=config) + + def log_trial_result(self, iteration: int, trial: "Trial", result: Dict): + step = result.get(TIMESTEPS_TOTAL) or result[TRAINING_ITERATION] + run_id = self._trial_runs[trial] + self.mlflow_util.log_metrics(run_id=run_id, metrics_to_log=result, step=step) + + def log_trial_end(self, trial: "Trial", failed: bool = False): + run_id = self._trial_runs[trial] + + # Log the artifact if set_artifact is set to True. + if self.should_save_artifact: + self.mlflow_util.save_artifacts(run_id=run_id, dir=trial.local_path) + + # Stop the run once trial finishes. + status = "FINISHED" if not failed else "FAILED" + + # Log the config parameters. + config = trial.config + if self.log_params_on_trial_end: + self.mlflow_util.log_params(run_id=run_id, params_to_log=config) + + self.mlflow_util.end_run(run_id=run_id, status=status) diff --git a/lib/python3.12/site-packages/ray/air/integrations/wandb.py b/lib/python3.12/site-packages/ray/air/integrations/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..f53d63094d3a5ad2ffdabc67a6dc17985d64e8f4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/integrations/wandb.py @@ -0,0 +1,802 @@ +import enum +import os +import pickle +import urllib +import warnings +from numbers import Number +from types import ModuleType +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import numpy as np +import pyarrow.fs + +import ray +from ray import logger +from ray._common.utils import load_class +from ray.air._internal import usage as air_usage +from ray.air.constants import TRAINING_ITERATION +from ray.air.util.node import _force_on_current_node +from ray.train._internal.session import get_session +from ray.train._internal.syncer import DEFAULT_SYNC_TIMEOUT +from ray.tune.experiment import Trial +from ray.tune.logger import LoggerCallback +from ray.tune.utils import flatten_dict +from ray.util import PublicAPI +from ray.util.queue import Queue + +try: + import wandb + from wandb.sdk.data_types.base_types.wb_value import WBValue + from wandb.sdk.data_types.image import Image + from wandb.sdk.data_types.video import Video + from wandb.sdk.lib.disabled import RunDisabled + from wandb.util import json_dumps_safer + from wandb.wandb_run import Run +except ImportError: + wandb = json_dumps_safer = Run = RunDisabled = WBValue = None + + +WANDB_ENV_VAR = "WANDB_API_KEY" +WANDB_PROJECT_ENV_VAR = "WANDB_PROJECT_NAME" +WANDB_GROUP_ENV_VAR = "WANDB_GROUP_NAME" +WANDB_MODE_ENV_VAR = "WANDB_MODE" +# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback +# to populate the API key if it isn't already set when initializing the callback. +# It doesn't take in any arguments and returns the W&B API key. +# Example: "your.module.wandb_setup_api_key_hook". +WANDB_SETUP_API_KEY_HOOK = "WANDB_SETUP_API_KEY_HOOK" +# Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback +# to populate environment variables to specify the location +# (project and group) of the W&B run. +# It doesn't take in any arguments and doesn't return anything, but it does populate +# WANDB_PROJECT_NAME and WANDB_GROUP_NAME. +# Example: "your.module.wandb_populate_run_location_hook". +WANDB_POPULATE_RUN_LOCATION_HOOK = "WANDB_POPULATE_RUN_LOCATION_HOOK" +# Hook that is invoked after running wandb.init in WandbLoggerCallback +# to process information about the W&B run. +# It takes in a W&B run object and doesn't return anything. +# Example: "your.module.wandb_process_run_info_hook". +WANDB_PROCESS_RUN_INFO_HOOK = "WANDB_PROCESS_RUN_INFO_HOOK" + + +@PublicAPI(stability="alpha") +def setup_wandb( + config: Optional[Dict] = None, + api_key: Optional[str] = None, + api_key_file: Optional[str] = None, + rank_zero_only: bool = True, + **kwargs, +) -> Union[Run, RunDisabled]: + """Set up a Weights & Biases session. + + This function can be used to initialize a Weights & Biases session in a + (distributed) training or tuning run. + + By default, the run ID is the trial ID, the run name is the trial name, and + the run group is the experiment name. These settings can be overwritten by + passing the respective arguments as ``kwargs``, which will be passed to + ``wandb.init()``. + + In distributed training with Ray Train, only the zero-rank worker will initialize + wandb. All other workers will return a disabled run object, so that logging is not + duplicated in a distributed run. This can be disabled by passing + ``rank_zero_only=False``, which will then initialize wandb in every training + worker. + + The ``config`` argument will be passed to Weights and Biases and will be logged + as the run configuration. + + If no API key or key file are passed, wandb will try to authenticate + using locally stored credentials, created for instance by running ``wandb login``. + + Keyword arguments passed to ``setup_wandb()`` will be passed to + ``wandb.init()`` and take precedence over any potential default settings. + + Args: + config: Configuration dict to be logged to Weights and Biases. Can contain + arguments for ``wandb.init()`` as well as authentication information. + api_key: API key to use for authentication with Weights and Biases. + api_key_file: File pointing to API key for with Weights and Biases. + rank_zero_only: If True, will return an initialized session only for the + rank 0 worker in distributed training. If False, will initialize a + session for all workers. + kwargs: Passed to ``wandb.init()``. + + Example: + + .. code-block:: python + + from ray.air.integrations.wandb import setup_wandb + + def training_loop(config): + wandb = setup_wandb(config) + # ... + wandb.log({"loss": 0.123}) + + """ + if not wandb: + raise RuntimeError( + "Wandb was not found - please install with `pip install wandb`" + ) + + default_trial_id = None + default_trial_name = None + default_experiment_name = None + + # Do a try-catch here if we are not in a train session + session = get_session() + + if rank_zero_only: + # Check if we are in a train session and if we are not the rank 0 worker + if session and session.world_rank is not None and session.world_rank != 0: + return RunDisabled() + + if session: + default_trial_id = session.trial_id + default_trial_name = session.trial_name + default_experiment_name = session.experiment_name + + # Default init kwargs + wandb_init_kwargs = { + "trial_id": kwargs.get("trial_id") or default_trial_id, + "trial_name": kwargs.get("trial_name") or default_trial_name, + "group": kwargs.get("group") or default_experiment_name, + } + # Passed kwargs take precedence over default kwargs + wandb_init_kwargs.update(kwargs) + + return _setup_wandb( + config=config, api_key=api_key, api_key_file=api_key_file, **wandb_init_kwargs + ) + + +def _setup_wandb( + trial_id: str, + trial_name: str, + config: Optional[Dict] = None, + api_key: Optional[str] = None, + api_key_file: Optional[str] = None, + _wandb: Optional[ModuleType] = None, + **kwargs, +) -> Union[Run, RunDisabled]: + _config = config.copy() if config else {} + + # If key file is specified, set + if api_key_file: + api_key_file = os.path.expanduser(api_key_file) + + _set_api_key(api_key_file, api_key) + project = _get_wandb_project(kwargs.pop("project", None)) + group = kwargs.pop("group", os.environ.get(WANDB_GROUP_ENV_VAR)) + + # Remove unpickleable items. + _config = _clean_log(_config) + + wandb_init_kwargs = dict( + id=trial_id, + name=trial_name, + resume=True, + reinit=True, + allow_val_change=True, + config=_config, + project=project, + group=group, + ) + + # Update config (e.g. set any other parameters in the call to wandb.init) + wandb_init_kwargs.update(**kwargs) + + # On windows, we can't fork + if os.name == "nt": + os.environ["WANDB_START_METHOD"] = "thread" + else: + os.environ["WANDB_START_METHOD"] = "fork" + + _wandb = _wandb or wandb + + run = _wandb.init(**wandb_init_kwargs) + _run_wandb_process_run_info_hook(run) + + # Record `setup_wandb` usage when everything has setup successfully. + air_usage.tag_setup_wandb() + + return run + + +def _is_allowed_type(obj): + """Return True if type is allowed for logging to wandb""" + if isinstance(obj, np.ndarray) and obj.size == 1: + return isinstance(obj.item(), Number) + if isinstance(obj, Sequence) and len(obj) > 0: + return isinstance(obj[0], (Image, Video, WBValue)) + return isinstance(obj, (Number, WBValue)) + + +def _clean_log( + obj: Any, + *, + video_kwargs: Optional[Dict[str, Any]] = None, + image_kwargs: Optional[Dict[str, Any]] = None, +): + # Fixes https://github.com/ray-project/ray/issues/10631 + if video_kwargs is None: + video_kwargs = {} + if image_kwargs is None: + image_kwargs = {} + if isinstance(obj, dict): + return { + k: _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs) + for k, v in obj.items() + } + elif isinstance(obj, (list, set)): + return [ + _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs) + for v in obj + ] + elif isinstance(obj, tuple): + return tuple( + _clean_log(v, video_kwargs=video_kwargs, image_kwargs=image_kwargs) + for v in obj + ) + elif isinstance(obj, np.ndarray) and obj.ndim == 3: + # Must be single image (H, W, C). + return Image(obj, **image_kwargs) + elif isinstance(obj, np.ndarray) and obj.ndim == 4: + # Must be batch of images (N >= 1, H, W, C). + return ( + _clean_log( + [Image(v, **image_kwargs) for v in obj], + video_kwargs=video_kwargs, + image_kwargs=image_kwargs, + ) + if obj.shape[0] > 1 + else Image(obj[0], **image_kwargs) + ) + elif isinstance(obj, np.ndarray) and obj.ndim == 5: + # Must be batch of videos (N >= 1, T, C, W, H). + return ( + _clean_log( + [Video(v, **video_kwargs) for v in obj], + video_kwargs=video_kwargs, + image_kwargs=image_kwargs, + ) + if obj.shape[0] > 1 + else Video(obj[0], **video_kwargs) + ) + elif _is_allowed_type(obj): + return obj + + # Else + + try: + # This is what wandb uses internally. If we cannot dump + # an object using this method, wandb will raise an exception. + json_dumps_safer(obj) + + # This is probably unnecessary, but left here to be extra sure. + pickle.dumps(obj) + + return obj + except Exception: + # give up, similar to _SafeFallBackEncoder + fallback = str(obj) + + # Try to convert to int + try: + fallback = int(fallback) + return fallback + except ValueError: + pass + + # Try to convert to float + try: + fallback = float(fallback) + return fallback + except ValueError: + pass + + # Else, return string + return fallback + + +def _get_wandb_project(project: Optional[str] = None) -> Optional[str]: + """Get W&B project from environment variable or external hook if not passed + as and argument.""" + if ( + not project + and not os.environ.get(WANDB_PROJECT_ENV_VAR) + and os.environ.get(WANDB_POPULATE_RUN_LOCATION_HOOK) + ): + # Try to populate WANDB_PROJECT_ENV_VAR and WANDB_GROUP_ENV_VAR + # from external hook + try: + load_class(os.environ[WANDB_POPULATE_RUN_LOCATION_HOOK])() + except Exception as e: + logger.exception( + f"Error executing {WANDB_POPULATE_RUN_LOCATION_HOOK} to " + f"populate {WANDB_PROJECT_ENV_VAR} and {WANDB_GROUP_ENV_VAR}: {e}", + exc_info=e, + ) + if not project and os.environ.get(WANDB_PROJECT_ENV_VAR): + # Try to get project and group from environment variables if not + # passed through WandbLoggerCallback. + project = os.environ.get(WANDB_PROJECT_ENV_VAR) + return project + + +def _set_api_key(api_key_file: Optional[str] = None, api_key: Optional[str] = None): + """Set WandB API key from `wandb_config`. Will pop the + `api_key_file` and `api_key` keys from `wandb_config` parameter. + + The order of fetching the API key is: + 1) From `api_key` or `api_key_file` arguments + 2) From WANDB_API_KEY environment variables + 3) User already logged in to W&B (wandb.api.api_key set) + 4) From external hook WANDB_SETUP_API_KEY_HOOK + """ + if os.environ.get(WANDB_MODE_ENV_VAR) in {"offline", "disabled"}: + return + + if api_key_file: + if api_key: + raise ValueError("Both WandB `api_key_file` and `api_key` set.") + with open(api_key_file, "rt") as fp: + api_key = fp.readline().strip() + + if not api_key and not os.environ.get(WANDB_ENV_VAR): + # Check if user is already logged into wandb. + try: + wandb.ensure_configured() + if wandb.api.api_key: + logger.info("Already logged into W&B.") + return + except AttributeError: + pass + # Try to get API key from external hook + if WANDB_SETUP_API_KEY_HOOK in os.environ: + try: + api_key = load_class(os.environ[WANDB_SETUP_API_KEY_HOOK])() + except Exception as e: + logger.exception( + f"Error executing {WANDB_SETUP_API_KEY_HOOK} to setup API key: {e}", + exc_info=e, + ) + if api_key: + os.environ[WANDB_ENV_VAR] = api_key + elif not os.environ.get(WANDB_ENV_VAR): + raise ValueError( + "No WandB API key found. Either set the {} environment " + "variable, pass `api_key` or `api_key_file` to the" + "`WandbLoggerCallback` class as arguments, " + "or run `wandb login` from the command line".format(WANDB_ENV_VAR) + ) + + +def _run_wandb_process_run_info_hook(run: Any) -> None: + """Run external hook to process information about wandb run""" + if WANDB_PROCESS_RUN_INFO_HOOK in os.environ: + try: + load_class(os.environ[WANDB_PROCESS_RUN_INFO_HOOK])(run) + except Exception as e: + logger.exception( + f"Error calling {WANDB_PROCESS_RUN_INFO_HOOK}: {e}", exc_info=e + ) + + +class _QueueItem(enum.Enum): + END = enum.auto() + RESULT = enum.auto() + CHECKPOINT = enum.auto() + + +class _WandbLoggingActor: + """ + Wandb assumes that each trial's information should be logged from a + separate process. We use Ray actors as forking multiprocessing + processes is not supported by Ray and spawn processes run into pickling + problems. + + We use a queue for the driver to communicate with the logging process. + The queue accepts the following items: + + - If it's a dict, it is assumed to be a result and will be logged using + ``wandb.log()`` + - If it's a checkpoint object, it will be saved using ``wandb.log_artifact()``. + """ + + def __init__( + self, + logdir: str, + queue: Queue, + exclude: List[str], + to_config: List[str], + *args, + **kwargs, + ): + import wandb + + self._wandb = wandb + + os.chdir(logdir) + self.queue = queue + self._exclude = set(exclude) + self._to_config = set(to_config) + self.args = args + self.kwargs = kwargs + + self._trial_name = self.kwargs.get("name", "unknown") + self._logdir = logdir + + def run(self): + # Since we're running in a separate process already, use threads. + os.environ["WANDB_START_METHOD"] = "thread" + run = self._wandb.init(*self.args, **self.kwargs) + run.config.trial_log_path = self._logdir + + _run_wandb_process_run_info_hook(run) + + while True: + item_type, item_content = self.queue.get() + if item_type == _QueueItem.END: + break + + if item_type == _QueueItem.CHECKPOINT: + self._handle_checkpoint(item_content) + continue + + assert item_type == _QueueItem.RESULT + log, config_update = self._handle_result(item_content) + try: + self._wandb.config.update(config_update, allow_val_change=True) + self._wandb.log(log, step=log.get(TRAINING_ITERATION)) + except urllib.error.HTTPError as e: + # Ignore HTTPError. Missing a few data points is not a + # big issue, as long as things eventually recover. + logger.warning("Failed to log result to w&b: {}".format(str(e))) + except FileNotFoundError as e: + logger.error( + "FileNotFoundError: Did not log result to Weights & Biases. " + "Possible cause: relative file path used instead of absolute path. " + "Error: %s", + e, + ) + self._wandb.finish() + + def _handle_checkpoint(self, checkpoint_path: str): + artifact = self._wandb.Artifact( + name=f"checkpoint_{self._trial_name}", type="model" + ) + artifact.add_dir(checkpoint_path) + self._wandb.log_artifact(artifact) + + def _handle_result(self, result: Dict) -> Tuple[Dict, Dict]: + config_update = result.get("config", {}).copy() + log = {} + flat_result = flatten_dict(result, delimiter="/") + + for k, v in flat_result.items(): + if any(k.startswith(item + "/") or k == item for item in self._exclude): + continue + elif any(k.startswith(item + "/") or k == item for item in self._to_config): + config_update[k] = v + elif not _is_allowed_type(v): + continue + else: + log[k] = v + + config_update.pop("callbacks", None) # Remove callbacks + return log, config_update + + +@PublicAPI(stability="alpha") +class WandbLoggerCallback(LoggerCallback): + """WandbLoggerCallback + + Weights and biases (https://www.wandb.ai/) is a tool for experiment + tracking, model optimization, and dataset versioning. This Ray Tune + ``LoggerCallback`` sends metrics to Wandb for automatic tracking and + visualization. + + Example: + + .. testcode:: + + import random + + from ray import tune + from ray.air.integrations.wandb import WandbLoggerCallback + + + def train_func(config): + offset = random.random() / 5 + for epoch in range(2, config["epochs"]): + acc = 1 - (2 + config["lr"]) ** -epoch - random.random() / epoch - offset + loss = (2 + config["lr"]) ** -epoch + random.random() / epoch + offset + train.report({"acc": acc, "loss": loss}) + + + tuner = tune.Tuner( + train_func, + param_space={ + "lr": tune.grid_search([0.001, 0.01, 0.1, 1.0]), + "epochs": 10, + }, + run_config=tune.RunConfig( + callbacks=[WandbLoggerCallback(project="Optimization_Project")] + ), + ) + results = tuner.fit() + + .. testoutput:: + :hide: + + ... + + Args: + project: Name of the Wandb project. Mandatory. + group: Name of the Wandb group. Defaults to the trainable + name. + api_key_file: Path to file containing the Wandb API KEY. This + file only needs to be present on the node running the Tune script + if using the WandbLogger. + api_key: Wandb API Key. Alternative to setting ``api_key_file``. + excludes: List of metrics and config that should be excluded from + the log. + log_config: Boolean indicating if the ``config`` parameter of + the ``results`` dict should be logged. This makes sense if + parameters will change during training, e.g. with + PopulationBasedTraining. Defaults to False. + upload_checkpoints: If ``True``, model checkpoints will be uploaded to + Wandb as artifacts. Defaults to ``False``. + video_kwargs: Dictionary of keyword arguments passed to wandb.Video() + when logging videos. Videos have to be logged as 5D numpy arrays + to be affected by this parameter. For valid keyword arguments, see + https://docs.wandb.ai/ref/python/data-types/video/. Defaults to ``None``. + image_kwargs: Dictionary of keyword arguments passed to wandb.Image() + when logging images. Images have to be logged as 3D or 4D numpy arrays + to be affected by this parameter. For valid keyword arguments, see + https://docs.wandb.ai/ref/python/data-types/image/. Defaults to ``None``. + **kwargs: The keyword arguments will be passed to ``wandb.init()``. + + Wandb's ``group``, ``run_id`` and ``run_name`` are automatically selected + by Tune, but can be overwritten by filling out the respective configuration + values. + + Please see here for all other valid configuration settings: + https://docs.wandb.ai/ref/python/init/ + """ # noqa: E501 + + # Do not log these result keys + _exclude_results = ["done", "should_checkpoint"] + + AUTO_CONFIG_KEYS = [ + "trial_id", + "experiment_tag", + "node_ip", + "experiment_id", + "hostname", + "pid", + "date", + ] + """Results that are saved with `wandb.config` instead of `wandb.log`.""" + + _logger_actor_cls = _WandbLoggingActor + + def __init__( + self, + project: Optional[str] = None, + group: Optional[str] = None, + api_key_file: Optional[str] = None, + api_key: Optional[str] = None, + excludes: Optional[List[str]] = None, + log_config: bool = False, + upload_checkpoints: bool = False, + save_checkpoints: bool = False, + upload_timeout: int = DEFAULT_SYNC_TIMEOUT, + video_kwargs: Optional[dict] = None, + image_kwargs: Optional[dict] = None, + **kwargs, + ): + if not wandb: + raise RuntimeError( + "Wandb was not found - please install with `pip install wandb`" + ) + + if save_checkpoints: + warnings.warn( + "`save_checkpoints` is deprecated. Use `upload_checkpoints` instead.", + DeprecationWarning, + ) + upload_checkpoints = save_checkpoints + + self.project = project + self.group = group + self.api_key_path = api_key_file + self.api_key = api_key + self.excludes = excludes or [] + self.log_config = log_config + self.upload_checkpoints = upload_checkpoints + self._upload_timeout = upload_timeout + self.video_kwargs = video_kwargs or {} + self.image_kwargs = image_kwargs or {} + self.kwargs = kwargs + + self._remote_logger_class = None + + self._trial_logging_actors: Dict[ + "Trial", ray.actor.ActorHandle[_WandbLoggingActor] + ] = {} + self._trial_logging_futures: Dict["Trial", ray.ObjectRef] = {} + self._logging_future_to_trial: Dict[ray.ObjectRef, "Trial"] = {} + self._trial_queues: Dict["Trial", Queue] = {} + + def setup(self, *args, **kwargs): + self.api_key_file = ( + os.path.expanduser(self.api_key_path) if self.api_key_path else None + ) + _set_api_key(self.api_key_file, self.api_key) + + self.project = _get_wandb_project(self.project) + if not self.project: + raise ValueError( + "Please pass the project name as argument or through " + f"the {WANDB_PROJECT_ENV_VAR} environment variable." + ) + if not self.group and os.environ.get(WANDB_GROUP_ENV_VAR): + self.group = os.environ.get(WANDB_GROUP_ENV_VAR) + + def log_trial_start(self, trial: "Trial"): + config = trial.config.copy() + + config.pop("callbacks", None) # Remove callbacks + + exclude_results = self._exclude_results.copy() + + # Additional excludes + exclude_results += self.excludes + + # Log config keys on each result? + if not self.log_config: + exclude_results += ["config"] + + # Fill trial ID and name + trial_id = trial.trial_id if trial else None + trial_name = str(trial) if trial else None + + # Project name for Wandb + wandb_project = self.project + + # Grouping + wandb_group = self.group or trial.experiment_dir_name if trial else None + + # remove unpickleable items! + config = _clean_log(config) + config = { + key: value for key, value in config.items() if key not in self.excludes + } + + wandb_init_kwargs = dict( + id=trial_id, + name=trial_name, + resume=False, + reinit=True, + allow_val_change=True, + group=wandb_group, + project=wandb_project, + config=config, + ) + wandb_init_kwargs.update(self.kwargs) + + self._start_logging_actor(trial, exclude_results, **wandb_init_kwargs) + + def _start_logging_actor( + self, trial: "Trial", exclude_results: List[str], **wandb_init_kwargs + ): + # Reuse actor if one already exists. + # This can happen if the trial is restarted. + if trial in self._trial_logging_futures: + return + + if not self._remote_logger_class: + env_vars = {} + # API key env variable is not set if authenticating through `wandb login` + if WANDB_ENV_VAR in os.environ: + env_vars[WANDB_ENV_VAR] = os.environ[WANDB_ENV_VAR] + self._remote_logger_class = ray.remote( + num_cpus=0, + **_force_on_current_node(), + runtime_env={"env_vars": env_vars}, + max_restarts=-1, + max_task_retries=-1, + )(self._logger_actor_cls) + + self._trial_queues[trial] = Queue( + actor_options={ + "num_cpus": 0, + **_force_on_current_node(), + "max_restarts": -1, + "max_task_retries": -1, + } + ) + self._trial_logging_actors[trial] = self._remote_logger_class.remote( + logdir=trial.local_path, + queue=self._trial_queues[trial], + exclude=exclude_results, + to_config=self.AUTO_CONFIG_KEYS, + **wandb_init_kwargs, + ) + logging_future = self._trial_logging_actors[trial].run.remote() + self._trial_logging_futures[trial] = logging_future + self._logging_future_to_trial[logging_future] = trial + + def _signal_logging_actor_stop(self, trial: "Trial"): + self._trial_queues[trial].put((_QueueItem.END, None)) + + def log_trial_result(self, iteration: int, trial: "Trial", result: Dict): + if trial not in self._trial_logging_actors: + self.log_trial_start(trial) + + result = _clean_log( + result, video_kwargs=self.video_kwargs, image_kwargs=self.image_kwargs + ) + self._trial_queues[trial].put((_QueueItem.RESULT, result)) + + def log_trial_save(self, trial: "Trial"): + if self.upload_checkpoints and trial.checkpoint: + checkpoint_root = None + if isinstance(trial.checkpoint.filesystem, pyarrow.fs.LocalFileSystem): + checkpoint_root = trial.checkpoint.path + + if checkpoint_root: + self._trial_queues[trial].put((_QueueItem.CHECKPOINT, checkpoint_root)) + + def log_trial_end(self, trial: "Trial", failed: bool = False): + self._signal_logging_actor_stop(trial=trial) + self._cleanup_logging_actors() + + def _cleanup_logging_actor(self, trial: "Trial"): + del self._trial_queues[trial] + del self._trial_logging_futures[trial] + ray.kill(self._trial_logging_actors[trial]) + del self._trial_logging_actors[trial] + + def _cleanup_logging_actors(self, timeout: int = 0, kill_on_timeout: bool = False): + """Clean up logging actors that have finished uploading to wandb. + Waits for `timeout` seconds to collect finished logging actors. + + Args: + timeout: The number of seconds to wait. Defaults to 0 to clean up + any immediate logging actors during the run. + This is set to a timeout threshold to wait for pending uploads + on experiment end. + kill_on_timeout: Whether or not to kill and cleanup the logging actor if + it hasn't finished within the timeout. + """ + + futures = list(self._trial_logging_futures.values()) + done, remaining = ray.wait(futures, num_returns=len(futures), timeout=timeout) + for ready_future in done: + finished_trial = self._logging_future_to_trial.pop(ready_future) + self._cleanup_logging_actor(finished_trial) + + if kill_on_timeout: + for remaining_future in remaining: + trial = self._logging_future_to_trial.pop(remaining_future) + self._cleanup_logging_actor(trial) + + def on_experiment_end(self, trials: List["Trial"], **info): + """Wait for the actors to finish their call to `wandb.finish`. + This includes uploading all logs + artifacts to wandb.""" + self._cleanup_logging_actors(timeout=self._upload_timeout, kill_on_timeout=True) + + def __del__(self): + if ray.is_initialized(): + for trial in list(self._trial_logging_actors): + self._signal_logging_actor_stop(trial=trial) + + self._cleanup_logging_actors(timeout=2, kill_on_timeout=True) + + self._trial_logging_actors = {} + self._trial_logging_futures = {} + self._logging_future_to_trial = {} + self._trial_queues = {} diff --git a/lib/python3.12/site-packages/ray/air/result.py b/lib/python3.12/site-packages/ray/air/result.py new file mode 100644 index 0000000000000000000000000000000000000000..9b911a563233fde34e665694a74564a6094d0396 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/result.py @@ -0,0 +1,284 @@ +import io +import json +import logging +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import pandas as pd +import pyarrow + +import ray +from ray.air.constants import ( + EXPR_ERROR_PICKLE_FILE, + EXPR_PROGRESS_FILE, + EXPR_RESULT_FILE, +) +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +@dataclass +@PublicAPI(stability="stable") +class Result: + """The final result of a ML training run or a Tune trial. + + This is the output produced by ``Trainer.fit``. + ``Tuner.fit`` outputs a :class:`~ray.tune.ResultGrid` that is a collection + of ``Result`` objects. + + This API is the recommended way to access the outputs such as: + - checkpoints (``Result.checkpoint``) + - the history of reported metrics (``Result.metrics_dataframe``, ``Result.metrics``) + - errors encountered during a training run (``Result.error``) + + The constructor is a private API -- use ``Result.from_path`` to create a result + object from a directory. + + Attributes: + metrics: The latest set of reported metrics. + checkpoint: The latest checkpoint. + error: The execution error of the Trainable run, if the trial finishes in error. + path: Path pointing to the result directory on persistent storage. This can + point to a remote storage location (e.g. S3) or to a local location (path + on the head node). The path is accessible via the result's associated + `filesystem`. For instance, for a result stored in S3 at + ``s3://bucket/location``, ``path`` will have the value ``bucket/location``. + metrics_dataframe: The full result dataframe of the Trainable. + The dataframe is indexed by iterations and contains reported + metrics. Note that the dataframe columns are indexed with the + *flattened* keys of reported metrics, so the format of this dataframe + may be slightly different than ``Result.metrics``, which is an unflattened + dict of the latest set of reported metrics. + best_checkpoints: A list of tuples of the best checkpoints and + their associated metrics. The number of + saved checkpoints is determined by :class:`~ray.train.CheckpointConfig` + (by default, all checkpoints will be saved). + """ + + metrics: Optional[Dict[str, Any]] + checkpoint: Optional["ray.tune.Checkpoint"] + error: Optional[Exception] + path: str + metrics_dataframe: Optional["pd.DataFrame"] = None + best_checkpoints: Optional[ + List[Tuple["ray.tune.Checkpoint", Dict[str, Any]]] + ] = None + _storage_filesystem: Optional[pyarrow.fs.FileSystem] = None + _items_to_repr = ["error", "metrics", "path", "filesystem", "checkpoint"] + + @property + def config(self) -> Optional[Dict[str, Any]]: + """The config associated with the result.""" + if not self.metrics: + return None + return self.metrics.get("config", None) + + @property + def filesystem(self) -> pyarrow.fs.FileSystem: + """Return the filesystem that can be used to access the result path. + + Returns: + pyarrow.fs.FileSystem implementation. + """ + return self._storage_filesystem or pyarrow.fs.LocalFileSystem() + + def _repr(self, indent: int = 0) -> str: + """Construct the representation with specified number of space indent.""" + from ray.tune.experimental.output import BLACKLISTED_KEYS + from ray.tune.result import AUTO_RESULT_KEYS + + shown_attributes = {k: getattr(self, k) for k in self._items_to_repr} + if self.error: + shown_attributes["error"] = type(self.error).__name__ + else: + shown_attributes.pop("error") + + shown_attributes["filesystem"] = shown_attributes["filesystem"].type_name + + if self.metrics: + exclude = set(AUTO_RESULT_KEYS) + exclude.update(BLACKLISTED_KEYS) + shown_attributes["metrics"] = { + k: v for k, v in self.metrics.items() if k not in exclude + } + + cls_indent = " " * indent + kws_indent = " " * (indent + 2) + + kws = [ + f"{kws_indent}{key}={value!r}" for key, value in shown_attributes.items() + ] + kws_repr = ",\n".join(kws) + return "{0}{1}(\n{2}\n{0})".format(cls_indent, type(self).__name__, kws_repr) + + def __repr__(self) -> str: + return self._repr(indent=0) + + @staticmethod + def _read_file_as_str( + storage_filesystem: pyarrow.fs.FileSystem, + storage_path: str, + ) -> str: + """Opens a file as an input stream reading all byte content sequentially and + decoding read bytes as utf-8 string. + + Args: + storage_filesystem: The filesystem to use. + storage_path: The source to open for reading. + """ + + with storage_filesystem.open_input_stream(storage_path) as f: + return f.readall().decode() + + @classmethod + def from_path( + cls, + path: Union[str, os.PathLike], + storage_filesystem: Optional[pyarrow.fs.FileSystem] = None, + ) -> "Result": + """Restore a Result object from local or remote trial directory. + + Args: + path: A path of a trial directory on local or remote storage + (ex: s3://bucket/path or /tmp/ray_results). + storage_filesystem: A custom filesystem to use. If not provided, + this will be auto-resolved by pyarrow. If provided, the path + is assumed to be prefix-stripped already, and must be a valid path + on the filesystem. + + Returns: + A :py:class:`Result` object of that trial. + """ + # TODO(justinvyu): Fix circular dependency. + from ray.train import Checkpoint + from ray.train._internal.storage import ( + _exists_at_fs_path, + _list_at_fs_path, + get_fs_and_path, + ) + from ray.train.constants import CHECKPOINT_DIR_NAME + + fs, fs_path = get_fs_and_path(path, storage_filesystem) + if not _exists_at_fs_path(fs, fs_path): + raise RuntimeError(f"Trial folder {fs_path} doesn't exist!") + + # Restore metrics from result.json + result_json_file = Path(fs_path, EXPR_RESULT_FILE).as_posix() + progress_csv_file = Path(fs_path, EXPR_PROGRESS_FILE).as_posix() + if _exists_at_fs_path(fs, result_json_file): + lines = cls._read_file_as_str(fs, result_json_file).split("\n") + json_list = [json.loads(line) for line in lines if line] + metrics_df = pd.json_normalize(json_list, sep="/") + latest_metrics = json_list[-1] if json_list else {} + # Fallback to restore from progress.csv + elif _exists_at_fs_path(fs, progress_csv_file): + metrics_df = pd.read_csv( + io.StringIO(cls._read_file_as_str(fs, progress_csv_file)) + ) + latest_metrics = ( + metrics_df.iloc[-1].to_dict() if not metrics_df.empty else {} + ) + else: + raise RuntimeError( + f"Failed to restore the Result object: Neither {EXPR_RESULT_FILE}" + f" nor {EXPR_PROGRESS_FILE} exists in the trial folder!" + ) + + # Restore all checkpoints from the checkpoint folders + checkpoint_dir_names = sorted( + _list_at_fs_path( + fs, + fs_path, + file_filter=lambda file_info: file_info.type + == pyarrow.fs.FileType.Directory + and file_info.base_name.startswith("checkpoint_"), + ) + ) + + if checkpoint_dir_names: + checkpoints = [ + Checkpoint( + path=Path(fs_path, checkpoint_dir_name).as_posix(), filesystem=fs + ) + for checkpoint_dir_name in checkpoint_dir_names + ] + + metrics = [] + for checkpoint_dir_name in checkpoint_dir_names: + metrics_corresponding_to_checkpoint = metrics_df[ + metrics_df[CHECKPOINT_DIR_NAME] == checkpoint_dir_name + ] + if metrics_corresponding_to_checkpoint.empty: + logger.warning( + "Could not find metrics corresponding to " + f"{checkpoint_dir_name}. These will default to an empty dict." + ) + metrics.append( + {} + if metrics_corresponding_to_checkpoint.empty + else metrics_corresponding_to_checkpoint.iloc[-1].to_dict() + ) + + latest_checkpoint = checkpoints[-1] + # TODO(justinvyu): These are ordered by checkpoint index, since we don't + # know the metric to order these with. + best_checkpoints = list(zip(checkpoints, metrics)) + else: + best_checkpoints = latest_checkpoint = None + + # Restore the trial error if it exists + error = None + error_file_path = Path(fs_path, EXPR_ERROR_PICKLE_FILE).as_posix() + if _exists_at_fs_path(fs, error_file_path): + with fs.open_input_stream(error_file_path) as f: + error = ray.cloudpickle.load(f) + + return Result( + metrics=latest_metrics, + checkpoint=latest_checkpoint, + path=fs_path, + _storage_filesystem=fs, + metrics_dataframe=metrics_df, + best_checkpoints=best_checkpoints, + error=error, + ) + + @PublicAPI(stability="alpha") + def get_best_checkpoint( + self, metric: str, mode: str + ) -> Optional["ray.tune.Checkpoint"]: + """Get the best checkpoint from this trial based on a specific metric. + + Any checkpoints without an associated metric value will be filtered out. + + Args: + metric: The key for checkpoints to order on. + mode: One of ["min", "max"]. + + Returns: + :class:`Checkpoint ` object, or None if there is + no valid checkpoint associated with the metric. + """ + if not self.best_checkpoints: + raise RuntimeError("No checkpoint exists in the trial directory!") + + if mode not in ["max", "min"]: + raise ValueError( + f'Unsupported mode: {mode}. Please choose from ["min", "max"]!' + ) + + op = max if mode == "max" else min + valid_checkpoints = [ + ckpt_info for ckpt_info in self.best_checkpoints if metric in ckpt_info[1] + ] + + if not valid_checkpoints: + raise RuntimeError( + f"Invalid metric name {metric}! " + f"You may choose from the following metrics: {self.metrics.keys()}." + ) + + return op(valid_checkpoints, key=lambda x: x[1][metric])[0] diff --git a/lib/python3.12/site-packages/ray/air/session.py b/lib/python3.12/site-packages/ray/air/session.py new file mode 100644 index 0000000000000000000000000000000000000000..b6a9fba5a6d6c89e2aeed911849e9ab136724fe1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/session.py @@ -0,0 +1 @@ +from ray.train._internal.session import * # noqa: F401,F403 diff --git a/lib/python3.12/site-packages/ray/air/util/__init__.py b/lib/python3.12/site-packages/ray/air/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/util/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a2cf9a456ffd4e78b4cf69389f7d64c6f451091 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/__pycache__/data_batch_conversion.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/__pycache__/data_batch_conversion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..890613d2e242b551eaeb860cb1dc22bb941ae9f2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/__pycache__/data_batch_conversion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/__pycache__/node.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/__pycache__/node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c753b7e873c71ca01badeaee741769f5b25472be Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/__pycache__/node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/__pycache__/transform_pyarrow.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/__pycache__/transform_pyarrow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8dd3e79af387147e1179cca6605aa30ca2c482b Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/__pycache__/transform_pyarrow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/data_batch_conversion.py b/lib/python3.12/site-packages/ray/air/util/data_batch_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..1bf69b4b93989f8d341d2d5a3fa9423435602aa5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/data_batch_conversion.py @@ -0,0 +1,353 @@ +import warnings +from enum import Enum +from typing import TYPE_CHECKING, Dict, List, Union + +import numpy as np + +from ray.air.constants import TENSOR_COLUMN_NAME +from ray.air.data_batch_type import DataBatchType +from ray.util.annotations import Deprecated, DeveloperAPI + +if TYPE_CHECKING: + import pandas as pd + +# TODO: Consolidate data conversion edges for arrow bug workaround. +try: + import pyarrow +except ImportError: + pyarrow = None + +# Lazy import to avoid ray init failures without pandas installed and allow +# dataset to import modules in this file. +_pandas = None + + +def _lazy_import_pandas(): + global _pandas + if _pandas is None: + import pandas + + _pandas = pandas + return _pandas + + +@DeveloperAPI +class BatchFormat(str, Enum): + PANDAS = "pandas" + # TODO: Remove once Arrow is deprecated as user facing batch format + ARROW = "arrow" + NUMPY = "numpy" # Either a single numpy array or a Dict of numpy arrays. + + +@DeveloperAPI +class BlockFormat(str, Enum): + """Internal Dataset block format enum.""" + + PANDAS = "pandas" + ARROW = "arrow" + SIMPLE = "simple" + + +def _convert_batch_type_to_pandas( + data: DataBatchType, + cast_tensor_columns: bool = False, +) -> "pd.DataFrame": + """Convert the provided data to a Pandas DataFrame. + + Args: + data: Data of type DataBatchType + cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays. + + Returns: + A pandas Dataframe representation of the input data. + + """ + pd = _lazy_import_pandas() + + if isinstance(data, np.ndarray): + data = pd.DataFrame({TENSOR_COLUMN_NAME: _ndarray_to_column(data)}) + elif isinstance(data, dict): + tensor_dict = {} + for col_name, col in data.items(): + if not isinstance(col, np.ndarray): + raise ValueError( + "All values in the provided dict must be of type " + f"np.ndarray. Found type {type(col)} for key {col_name} " + f"instead." + ) + tensor_dict[col_name] = _ndarray_to_column(col) + data = pd.DataFrame(tensor_dict) + elif pyarrow is not None and isinstance(data, pyarrow.Table): + data = data.to_pandas() + elif not isinstance(data, pd.DataFrame): + raise ValueError( + f"Received data of type: {type(data)}, but expected it to be one " + f"of {DataBatchType}" + ) + if cast_tensor_columns: + data = _cast_tensor_columns_to_ndarrays(data) + return data + + +def _convert_pandas_to_batch_type( + data: "pd.DataFrame", + type: BatchFormat, + cast_tensor_columns: bool = False, +) -> DataBatchType: + """Convert the provided Pandas dataframe to the provided ``type``. + + Args: + data: A Pandas DataFrame + type: The specific ``BatchFormat`` to convert to. + cast_tensor_columns: Whether tensor columns should be cast to our tensor + extension type. + + Returns: + The input data represented with the provided type. + """ + if cast_tensor_columns: + data = _cast_ndarray_columns_to_tensor_extension(data) + if type == BatchFormat.PANDAS: + return data + + elif type == BatchFormat.NUMPY: + if len(data.columns) == 1: + # If just a single column, return as a single numpy array. + return data.iloc[:, 0].to_numpy() + else: + # Else return as a dict of numpy arrays. + output_dict = {} + for column in data: + output_dict[column] = data[column].to_numpy() + return output_dict + + elif type == BatchFormat.ARROW: + if not pyarrow: + raise ValueError( + "Attempted to convert data to Pyarrow Table but Pyarrow " + "is not installed. Please do `pip install pyarrow` to " + "install Pyarrow." + ) + return pyarrow.Table.from_pandas(data) + + else: + raise ValueError( + f"Received type {type}, but expected it to be one of {DataBatchType}" + ) + + +@Deprecated +def convert_batch_type_to_pandas( + data: DataBatchType, + cast_tensor_columns: bool = False, +): + """Convert the provided data to a Pandas DataFrame. + + This API is deprecated from Ray 2.4. + + Args: + data: Data of type DataBatchType + cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays. + + Returns: + A pandas Dataframe representation of the input data. + + """ + warnings.warn( + "`convert_batch_type_to_pandas` is deprecated as a developer API " + "starting from Ray 2.4. All batch format conversions should be " + "done manually instead of relying on this API.", + PendingDeprecationWarning, + ) + return _convert_batch_type_to_pandas( + data=data, cast_tensor_columns=cast_tensor_columns + ) + + +@Deprecated +def convert_pandas_to_batch_type( + data: "pd.DataFrame", + type: BatchFormat, + cast_tensor_columns: bool = False, +): + """Convert the provided Pandas dataframe to the provided ``type``. + + Args: + data: A Pandas DataFrame + type: The specific ``BatchFormat`` to convert to. + cast_tensor_columns: Whether tensor columns should be cast to our tensor + extension type. + + Returns: + The input data represented with the provided type. + """ + warnings.warn( + "`convert_pandas_to_batch_type` is deprecated as a developer API " + "starting from Ray 2.4. All batch format conversions should be " + "done manually instead of relying on this API.", + PendingDeprecationWarning, + ) + return _convert_pandas_to_batch_type( + data=data, type=type, cast_tensor_columns=cast_tensor_columns + ) + + +def _convert_batch_type_to_numpy( + data: DataBatchType, +) -> Union[np.ndarray, Dict[str, np.ndarray]]: + """Convert the provided data to a NumPy ndarray or dict of ndarrays. + + Args: + data: Data of type DataBatchType + + Returns: + A numpy representation of the input data. + """ + pd = _lazy_import_pandas() + + if isinstance(data, np.ndarray): + return data + elif isinstance(data, dict): + for col_name, col in data.items(): + if not isinstance(col, np.ndarray): + raise ValueError( + "All values in the provided dict must be of type " + f"np.ndarray. Found type {type(col)} for key {col_name} " + f"instead." + ) + return data + elif pyarrow is not None and isinstance(data, pyarrow.Table): + from ray.air.util.tensor_extensions.arrow import ( + get_arrow_extension_fixed_shape_tensor_types, + ) + from ray.data._internal.arrow_ops import transform_pyarrow + + column_values_ndarrays = [] + + for col in data.columns: + # Combine columnar values arrays to make these contiguous + # (making them compatible with numpy format) + combined_array = transform_pyarrow.combine_chunked_array(col) + + column_values_ndarrays.append( + transform_pyarrow.to_numpy(combined_array, zero_copy_only=False) + ) + + arrow_fixed_shape_tensor_types = get_arrow_extension_fixed_shape_tensor_types() + + # NOTE: This branch is here for backwards-compatibility + if data.column_names == [TENSOR_COLUMN_NAME] and ( + isinstance(data.schema.types[0], arrow_fixed_shape_tensor_types) + ): + return column_values_ndarrays[0] + + return dict(zip(data.column_names, column_values_ndarrays)) + elif isinstance(data, pd.DataFrame): + return _convert_pandas_to_batch_type(data, BatchFormat.NUMPY) + else: + raise ValueError( + f"Received data of type: {type(data)}, but expected it to be one " + f"of {DataBatchType}" + ) + + +def _ndarray_to_column(arr: np.ndarray) -> Union["pd.Series", List[np.ndarray]]: + """Convert a NumPy ndarray into an appropriate column format for insertion into a + pandas DataFrame. + + If conversion to a pandas Series fails (e.g. if the ndarray is multi-dimensional), + fall back to a list of NumPy ndarrays. + """ + pd = _lazy_import_pandas() + try: + # Try to convert to Series, falling back to a list conversion if this fails + # (e.g. if the ndarray is multi-dimensional). + return pd.Series(arr) + except ValueError: + return list(arr) + + +def _unwrap_ndarray_object_type_if_needed(arr: np.ndarray) -> np.ndarray: + """Unwrap an object-dtyped NumPy ndarray containing ndarray pointers into a single + contiguous ndarray, if needed/possible. + """ + if arr.dtype.type is np.object_: + try: + # Try to convert the NumPy ndarray to a non-object dtype. + arr = np.array([np.asarray(v) for v in arr]) + except Exception: + # This may fail if the subndarrays are of heterogeneous shape + pass + return arr + + +def _cast_ndarray_columns_to_tensor_extension(df: "pd.DataFrame") -> "pd.DataFrame": + """ + Cast all NumPy ndarray columns in df to our tensor extension type, TensorArray. + """ + pd = _lazy_import_pandas() + try: + SettingWithCopyWarning = pd.core.common.SettingWithCopyWarning + except AttributeError: + # SettingWithCopyWarning was moved to pd.errors in Pandas 1.5.0. + SettingWithCopyWarning = pd.errors.SettingWithCopyWarning + + from ray.air.util.tensor_extensions.pandas import ( + TensorArray, + column_needs_tensor_extension, + ) + + # Try to convert any ndarray columns to TensorArray columns. + # TODO(Clark): Once Pandas supports registering extension types for type + # inference on construction, implement as much for NumPy ndarrays and remove + # this. See https://github.com/pandas-dev/pandas/issues/41848 + # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of + # column names containing tensor columns, to make this an O(# of tensor columns) + # check rather than the current O(# of columns) check. + for col_name, col in df.items(): + if column_needs_tensor_extension(col): + try: + # Suppress Pandas warnings: + # https://github.com/ray-project/ray/issues/29270 + # We actually want in-place operations so we surpress this warning. + # https://stackoverflow.com/a/74193599 + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=FutureWarning) + warnings.simplefilter("ignore", category=SettingWithCopyWarning) + df[col_name] = TensorArray(col) + except Exception as e: + raise ValueError( + f"Tried to cast column {col_name} to the TensorArray tensor " + "extension type but the conversion failed. To disable " + "automatic casting to this tensor extension, set " + "ctx = DataContext.get_current(); " + "ctx.enable_tensor_extension_casting = False." + ) from e + return df + + +def _cast_tensor_columns_to_ndarrays(df: "pd.DataFrame") -> "pd.DataFrame": + """Cast all tensor extension columns in df to NumPy ndarrays.""" + pd = _lazy_import_pandas() + try: + SettingWithCopyWarning = pd.core.common.SettingWithCopyWarning + except AttributeError: + # SettingWithCopyWarning was moved to pd.errors in Pandas 1.5.0. + SettingWithCopyWarning = pd.errors.SettingWithCopyWarning + from ray.air.util.tensor_extensions.pandas import TensorDtype + + # Try to convert any tensor extension columns to ndarray columns. + # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of + # column names containing tensor columns, to make this an O(# of tensor columns) + # check rather than the current O(# of columns) check. + for col_name, col in df.items(): + if isinstance(col.dtype, TensorDtype): + # Suppress Pandas warnings: + # https://github.com/ray-project/ray/issues/29270 + # We actually want in-place operations so we surpress this warning. + # https://stackoverflow.com/a/74193599 + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=FutureWarning) + warnings.simplefilter("ignore", category=SettingWithCopyWarning) + df[col_name] = list(col.to_numpy()) + return df diff --git a/lib/python3.12/site-packages/ray/air/util/node.py b/lib/python3.12/site-packages/ray/air/util/node.py new file mode 100644 index 0000000000000000000000000000000000000000..60ea3558878f3b220f1ef97210193d5b0e626ef4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/node.py @@ -0,0 +1,69 @@ +from typing import Dict, Optional, Union + +import ray + + +def _get_node_id_from_node_ip(node_ip: str) -> Optional[str]: + """Returns the node ID for the first alive node with the input IP.""" + for node in ray.nodes(): + if node["Alive"] and node["NodeManagerAddress"] == node_ip: + return node["NodeID"] + + return None + + +def _force_on_node( + node_id: str, + remote_func_or_actor_class: Optional[ + Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass] + ] = None, +) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]: + """Schedule a remote function or actor class on a given node. + + Args: + node_id: The node to schedule on. + remote_func_or_actor_class: A Ray remote function or actor class + to schedule on the input node. If None, this function will directly + return the options dict to pass to another remote function or actor class + as remote options. + Returns: + The provided remote function or actor class, but with options modified to force + placement on the input node. If remote_func_or_actor_class is None, + the options dict to pass to another remote function or + actor class as remote options kwargs. + """ + + scheduling_strategy = ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, soft=False + ) + + options = {"scheduling_strategy": scheduling_strategy} + + if remote_func_or_actor_class is None: + return options + + return remote_func_or_actor_class.options(**options) + + +def _force_on_current_node( + remote_func_or_actor_class: Optional[ + Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass] + ] = None +) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]: + """Schedule a remote function or actor class on the current node. + + If using Ray Client, the current node is the client server node. + + Args: + remote_func_or_actor_class: A Ray remote function or actor class + to schedule on the current node. If None, this function will directly + return the options dict to pass to another remote function or actor class + as remote options. + Returns: + The provided remote function or actor class, but with options modified to force + placement on the current node. If remote_func_or_actor_class is None, + the options dict to pass to another remote function or + actor class as remote options kwargs. + """ + current_node_id = ray.get_runtime_context().get_node_id() + return _force_on_node(current_node_id, remote_func_or_actor_class) diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/__init__.py b/lib/python3.12/site-packages/ray/air/util/object_extensions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa4fb399dd01e49aa3c6ac596bdfdd6c8ac35b88 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/arrow.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/arrow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb5538c8a7ec1a6baf28c4e338ab944d1e069f34 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/arrow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/pandas.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/pandas.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6b33a1475823fb1eef879dfc5b9d33e4fdedeb0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/object_extensions/__pycache__/pandas.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/arrow.py b/lib/python3.12/site-packages/ray/air/util/object_extensions/arrow.py new file mode 100644 index 0000000000000000000000000000000000000000..1f78d72e03468e6f646ff39715d6c12227e522e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/object_extensions/arrow.py @@ -0,0 +1,130 @@ +import pickle +import typing + +import numpy as np +import pyarrow as pa +from packaging.version import parse as parse_version + +import ray.air.util.object_extensions.pandas +from ray._common.serialization import pickle_dumps +from ray._private.arrow_utils import _check_pyarrow_version, get_pyarrow_version +from ray.util.annotations import PublicAPI + +# First, assert Arrow version is w/in expected bounds +_check_pyarrow_version() + + +MIN_PYARROW_VERSION_SCALAR_SUBCLASS = parse_version("9.0.0") + +PYARROW_VERSION = get_pyarrow_version() + + +# TODO delete, since min supported pyarrow >= 9.0 +def _object_extension_type_allowed() -> bool: + return ( + PYARROW_VERSION is not None + and PYARROW_VERSION >= MIN_PYARROW_VERSION_SCALAR_SUBCLASS + ) + + +# Please see https://arrow.apache.org/docs/python/extending_types.html for more info +@PublicAPI(stability="alpha") +class ArrowPythonObjectType(pa.ExtensionType): + """Defines a new Arrow extension type for Python objects. + We do not require a parametrized type, so the constructor does not + take any arguments + """ + + def __init__(self) -> None: + # Defines the underlying storage type as the PyArrow LargeBinary type + super().__init__(pa.large_binary(), "ray.data.arrow_pickled_object") + + def __arrow_ext_serialize__(self) -> bytes: + # Since there are no type parameters, we are free to return empty + return b"" + + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: pa.DataType, serialized: bytes + ) -> "ArrowPythonObjectType": + return ArrowPythonObjectType() + + def __arrow_ext_scalar_class__(self) -> type: + """Returns the scalar class of the extension type. Indexing out of the + PyArrow extension array will return instances of this type. + """ + return ArrowPythonObjectScalar + + def __arrow_ext_class__(self) -> type: + """Returns the array type of the extension type. Selecting one array + out of the ChunkedArray that makes up a column in a Table with + this custom type will return an instance of this type. + """ + return ArrowPythonObjectArray + + def to_pandas_dtype(self): + """Pandas interoperability type. This describes the Pandas counterpart + to the Arrow type. See https://pandas.pydata.org/docs/development/extending.html + for more information. + """ + return ray.air.util.object_extensions.pandas.PythonObjectDtype() + + def __reduce__(self): + # Earlier PyArrow versions require custom pickling behavior. + return self.__arrow_ext_deserialize__, ( + self.storage_type, + self.__arrow_ext_serialize__(), + ) + + def __hash__(self) -> int: + return hash((type(self), self.storage_type.id, self.extension_name)) + + +@PublicAPI(stability="alpha") +class ArrowPythonObjectScalar(pa.ExtensionScalar): + """Scalar class for ArrowPythonObjectType""" + + def as_py(self, **kwargs) -> typing.Any: + # Handle None/null values + if self.value is None: + return None + + if not isinstance(self.value, pa.LargeBinaryScalar): + raise RuntimeError( + f"{type(self.value)} is not the expected LargeBinaryScalar" + ) + return pickle.load(pa.BufferReader(self.value.as_buffer())) + + +@PublicAPI(stability="alpha") +class ArrowPythonObjectArray(pa.ExtensionArray): + """Array class for ArrowPythonObjectType""" + + def from_objects( + objects: typing.Union[np.ndarray, typing.Iterable[typing.Any]] + ) -> "ArrowPythonObjectArray": + if isinstance(objects, np.ndarray): + objects = objects.tolist() + type_ = ArrowPythonObjectType() + all_dumped_bytes = [] + for obj in objects: + dumped_bytes = pickle_dumps( + obj, "Error pickling object to convert to Arrow" + ) + all_dumped_bytes.append(dumped_bytes) + arr = pa.array(all_dumped_bytes, type=type_.storage_type) + return type_.wrap_array(arr) + + def to_numpy( + self, zero_copy_only: bool = False, writable: bool = False + ) -> np.ndarray: + arr = np.empty(len(self), dtype=object) + arr[:] = self.to_pylist() + return arr + + +try: + pa.register_extension_type(ArrowPythonObjectType()) +except pa.ArrowKeyError: + # Already registered + pass diff --git a/lib/python3.12/site-packages/ray/air/util/object_extensions/pandas.py b/lib/python3.12/site-packages/ray/air/util/object_extensions/pandas.py new file mode 100644 index 0000000000000000000000000000000000000000..dbc5732f350b8c89d314b8a634c8809ad7e817b3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/object_extensions/pandas.py @@ -0,0 +1,146 @@ +import collections.abc +import typing + +import numpy as np +import pandas as pd +import pyarrow as pa +from pandas._libs import lib +from pandas._typing import ArrayLike, Dtype, PositionalIndexer, TakeIndexer, npt + +import ray.air.util.object_extensions.arrow +from ray.util.annotations import PublicAPI + + +# See https://pandas.pydata.org/docs/development/extending.html for more information. +@PublicAPI(stability="alpha") +class PythonObjectArray(pd.api.extensions.ExtensionArray): + """Implements the Pandas extension array interface for the Arrow object array""" + + def __init__(self, values: collections.abc.Iterable[typing.Any]): + vals = list(values) + self.values = np.empty(len(vals), dtype=object) + self.values[:] = vals + + @classmethod + def _from_sequence( + cls, + scalars: collections.abc.Sequence[typing.Any], + *, + dtype: typing.Union[Dtype, None] = None, + copy: bool = False, + ) -> "PythonObjectArray": + return PythonObjectArray(scalars) + + @classmethod + def _from_factorized( + cls, values: collections.abc.Sequence[typing.Any], original: "PythonObjectArray" + ) -> "PythonObjectArray": + return PythonObjectArray(values) + + def __getitem__(self, item: PositionalIndexer) -> typing.Any: + return self.values[item] + + def __setitem__(self, key, value) -> None: + self.values[key] = value + + def __len__(self) -> int: + return len(self.values) + + def __eq__(self, other: object) -> ArrayLike: + if isinstance(other, PythonObjectArray): + return self.values == other.values + elif isinstance(other, np.ndarray): + return self.values == other + else: + return NotImplemented + + def to_numpy( + self, + dtype: typing.Union["npt.DTypeLike", None] = None, + copy: bool = False, + na_value: object = lib.no_default, + ) -> np.ndarray: + result = self.values + if copy or na_value is not lib.no_default: + result = result.copy() + if na_value is not lib.no_default: + result[self.isna()] = na_value + return result + + @property + def dtype(self) -> pd.api.extensions.ExtensionDtype: + return PythonObjectDtype() + + @property + def nbytes(self) -> int: + return self.values.nbytes + + def __arrow_array__(self, type=None): + return ray.air.util.object_extensions.arrow.ArrowPythonObjectArray.from_objects( + self.values + ) + + def isna(self) -> np.ndarray: + return pd.isnull(self.values) + + def take( + self, + indices: TakeIndexer, + *, + allow_fill: bool = False, + fill_value: typing.Any = None, + ) -> "PythonObjectArray": + if allow_fill and fill_value is None: + fill_value = self.dtype.na_value + + result = pd.core.algorithms.take( + self.values, indices, allow_fill=allow_fill, fill_value=fill_value + ) + return self._from_sequence(result, dtype=self.dtype) + + def copy(self) -> "PythonObjectArray": + return PythonObjectArray(self.values) + + @classmethod + def _concat_same_type( + cls, to_concat: collections.abc.Sequence["PythonObjectArray"] + ) -> "PythonObjectArray": + values_to_concat = [element.values for element in to_concat] + return cls(np.concatenate(values_to_concat)) + + +@PublicAPI(stability="alpha") +@pd.api.extensions.register_extension_dtype +class PythonObjectDtype(pd.api.extensions.ExtensionDtype): + @classmethod + def construct_from_string(cls, string: str): + if string != "python_object()": + raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'") + return cls() + + @property + def type(self): + """ + The scalar type for the array, e.g. ``int`` + It's expected ``ExtensionArray[item]`` returns an instance + of ``ExtensionDtype.type`` for scalar ``item``, assuming + that value is valid (not NA). NA values do not need to be + instances of `type`. + """ + return object + + @property + def name(self) -> str: + return "python_object()" + + @classmethod + def construct_array_type(cls: type) -> type: + """ + Return the array type associated with this dtype. + """ + return PythonObjectArray + + def __from_arrow__( + self, array: typing.Union[pa.Array, pa.ChunkedArray] + ) -> PythonObjectArray: + return PythonObjectArray(array.to_pylist()) diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__init__.py b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0f7f74fdf58a51b1f1e607ed785b7cb2a779b4f Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/arrow.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/arrow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6266f24ea26af94796aae596e326cc77db4f186a Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/arrow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/pandas.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/pandas.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de734d4841c4b82cb11a871778267b357a7a27ca Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/pandas.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72d79b55ffae5f8af1482277edf56f5958db40ef Binary files /dev/null and b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/arrow.py b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/arrow.py new file mode 100644 index 0000000000000000000000000000000000000000..a101ddd61cd5de8f9092f2160d1bc84b24f63dfe --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/arrow.py @@ -0,0 +1,1576 @@ +import abc +import functools +import itertools +import json +import logging +import sys +import threading +from abc import abstractmethod +from datetime import datetime +from enum import Enum +from typing import Any, Collection, Dict, Iterable, List, Optional, Tuple, Union + +import numpy as np +import pyarrow as pa +from packaging.version import parse as parse_version + +import ray.cloudpickle as cloudpickle +from ray._private.arrow_utils import _check_pyarrow_version, get_pyarrow_version +from ray._private.ray_constants import env_integer +from ray.air.util.object_extensions.arrow import ( + MIN_PYARROW_VERSION_SCALAR_SUBCLASS, + ArrowPythonObjectArray, + _object_extension_type_allowed, +) +from ray.air.util.tensor_extensions.utils import ( + ArrayLike, + _is_ndarray_variable_shaped_tensor, + _should_convert_to_tensor, + create_ragged_ndarray, +) +from ray.data._internal.numpy_support import ( + _convert_datetime_to_np_datetime, + convert_to_numpy, +) +from ray.util import log_once +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.common import INT32_MAX + +# First, assert Arrow version is w/in expected bounds +_check_pyarrow_version() + + +PYARROW_VERSION = get_pyarrow_version() + + +# Minimum version supporting `zero_copy_only` flag in `ChunkedArray.to_numpy` +MIN_PYARROW_VERSION_CHUNKED_ARRAY_TO_NUMPY_ZERO_COPY_ONLY = parse_version("13.0.0") +# Min version supporting ``ExtensionArray``s in ``pyarrow.concat`` +MIN_PYARROW_VERSION_EXT_ARRAY_CONCAT_SUPPORTED = parse_version("12.0.0") + + +NUM_BYTES_PER_UNICODE_CHAR = 4 + + +class _SerializationFormat(Enum): + # JSON format is legacy and inefficient, only kept for backward compatibility + JSON = 0 + CLOUDPICKLE = 1 + + +# Set the default serialization format for Arrow extension types. +ARROW_EXTENSION_SERIALIZATION_FORMAT = _SerializationFormat( + _SerializationFormat.JSON # legacy + if env_integer("RAY_DATA_ARROW_EXTENSION_SERIALIZATION_LEGACY_JSON_FORMAT", 0) == 1 + else _SerializationFormat.CLOUDPICKLE # default +) + +# 100,000 entries, about 10MB in memory. +# Most users tables should have less than 100K columns. +ARROW_EXTENSION_SERIALIZATION_CACHE_MAXSIZE = env_integer( + "RAY_EXTENSION_SERIALIZATION_CACHE_MAXSIZE", 10**5 +) + +logger = logging.getLogger(__name__) + + +def _extension_array_concat_supported() -> bool: + return get_pyarrow_version() >= MIN_PYARROW_VERSION_EXT_ARRAY_CONCAT_SUPPORTED + + +def _deserialize_with_fallback(serialized: bytes, field_name: str = "data"): + """Deserialize data with cloudpickle first, fallback to JSON.""" + try: + # Try cloudpickle first (new format) + return cloudpickle.loads(serialized) + except Exception: + # Fallback to JSON format (legacy) + try: + return json.loads(serialized) + except json.JSONDecodeError: + raise ValueError( + f"Unable to deserialize {field_name} from {type(serialized)}" + ) + + +@DeveloperAPI(stability="beta") +class ArrowExtensionSerializeDeserializeCache(abc.ABC): + """Base class for caching Arrow extension type serialization and deserialization. + + The deserialization and serialization of Arrow extension types is frequent, + so we cache the results here to improve performance. + + The deserialization cache uses functools.lru_cache as a classmethod. There is + a single cache instance shared across all subclasses, but the cache key includes + the class (cls parameter) as the first argument, so different subclasses get + different cache entries even when called with the same parameters. The cache is + thread-safe and has a maximum size limit to control memory usage. The cache key + is (cls, *args) where args are the parameters returned by _get_deserialize_parameter(). + + Attributes: + _serialize_cache: Instance-level cache for serialization results. + This is a simple cached value (bytes) that is computed once per + instance and reused. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Initialize the extension type with caching support. + + Args: + *args: Positional arguments passed to the parent class. + **kwargs: Keyword arguments passed to the parent class. + """ + # Instance-level cache for serialization results, no TTL + self._serialize_cache = None + self._cache_lock = threading.RLock() + super().__init__(*args, **kwargs) + + def __arrow_ext_serialize__(self) -> bytes: + """Serialize the extension type using caching if enabled.""" + with self._cache_lock: + if self._serialize_cache is None: + self._serialize_cache = self._arrow_ext_serialize_compute() + return self._serialize_cache + + @abstractmethod + def _arrow_ext_serialize_compute(self) -> bytes: + """Subclasses must implement this method to compute serialization.""" + ... + + @classmethod + @functools.lru_cache(maxsize=ARROW_EXTENSION_SERIALIZATION_CACHE_MAXSIZE) + def _arrow_ext_deserialize_cache(cls: type, *args: Any, **kwargs: Any) -> Any: + """Deserialize the extension type using the class-level cache. + + This method is cached using functools.lru_cache to improve performance + when deserializing extension types. The cache key includes the class (cls) + as the first argument, ensuring different subclasses get separate cache entries. + + Args: + *args: Positional arguments passed to _arrow_ext_deserialize_compute. + **kwargs: Keyword arguments passed to _arrow_ext_deserialize_compute. + + Returns: + The deserialized extension type instance. + """ + return cls._arrow_ext_deserialize_compute(*args, **kwargs) + + @classmethod + @abstractmethod + def _arrow_ext_deserialize_compute(cls, *args: Any, **kwargs: Any) -> Any: + """Subclasses must implement this method to compute deserialization.""" + ... + + @classmethod + @abstractmethod + def _get_deserialize_parameter(cls, storage_type, serialized) -> Tuple: + """Subclasses must implement this method to return the parameters for the deserialization cache.""" + ... + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized) -> Any: + """Deserialize the extension type using caching if enabled.""" + return cls._arrow_ext_deserialize_cache( + *cls._get_deserialize_parameter(storage_type, serialized) + ) + + +@DeveloperAPI +class ArrowConversionError(Exception): + """Error raised when there is an issue converting data to Arrow.""" + + MAX_DATA_STR_LEN = 200 + + def __init__(self, data_str: str): + if len(data_str) > self.MAX_DATA_STR_LEN: + data_str = data_str[: self.MAX_DATA_STR_LEN] + "..." + message = f"Error converting data to Arrow: {data_str}" + super().__init__(message) + + +@DeveloperAPI +def pyarrow_table_from_pydict( + pydict: Dict[str, Union[List[Any], pa.Array]], +) -> pa.Table: + """ + Convert a Python dictionary to a pyarrow Table. + + Raises: + ArrowConversionError: if the conversion fails. + """ + try: + return pa.Table.from_pydict(pydict) + except Exception as e: + raise ArrowConversionError(str(pydict)) from e + + +@DeveloperAPI(stability="alpha") +def convert_to_pyarrow_array( + column_values: Union[List[Any], np.ndarray, ArrayLike], column_name: str +) -> pa.Array: + """Converts provided NumPy `ndarray` into PyArrow's `array` while utilizing + both Arrow's natively supported types as well as custom extension types: + + - ArrowTensorArray (for tensors) + - ArrowPythonObjectArray (for user-defined python class objects, as well as + any python object that aren't represented by a corresponding Arrow's native + scalar type) + """ + + try: + # Since Arrow does NOT support tensors (aka multidimensional arrays) natively, + # we have to make sure that we handle this case utilizing `ArrowTensorArray` + # extension type + if len(column_values) > 0 and _should_convert_to_tensor( + column_values, column_name + ): + from ray.data.extensions.tensor_extension import ArrowTensorArray + + # Convert to Numpy before creating instance of `ArrowTensorArray` to + # align tensor shapes falling back to ragged ndarray only if necessary + return ArrowTensorArray.from_numpy( + convert_to_numpy(column_values), column_name + ) + else: + return _convert_to_pyarrow_native_array(column_values, column_name) + + except ArrowConversionError as ace: + from ray.data.context import DataContext + + enable_fallback_config: Optional[ + bool + ] = DataContext.get_current().enable_fallback_to_arrow_object_ext_type + + if not _object_extension_type_allowed(): + object_ext_type_fallback_allowed = False + object_ext_type_detail = ( + "skipping fallback to serialize as pickled python" + f" objects (due to unsupported Arrow version {PYARROW_VERSION}, " + f"min required version is {MIN_PYARROW_VERSION_SCALAR_SUBCLASS})" + ) + else: + # NOTE: By default setting is unset which (for compatibility reasons) + # is allowing the fallback + object_ext_type_fallback_allowed = ( + enable_fallback_config is None or enable_fallback_config + ) + + if object_ext_type_fallback_allowed: + object_ext_type_detail = ( + "falling back to serialize as pickled python objects" + ) + else: + object_ext_type_detail = ( + "skipping fallback to serialize as pickled python objects " + "(due to DataContext.enable_fallback_to_arrow_object_ext_type " + "= False)" + ) + + # To avoid logging following warning for every block it's + # only going to be logged in following cases + # - It's being logged for the first time, and + # - When config enabling fallback is not set explicitly (in this case + # fallback will still occur by default for compatibility reasons), or + # - Fallback is disallowed (either explicitly or due to use of incompatible + # Pyarrow version) + if ( + enable_fallback_config is None or not object_ext_type_fallback_allowed + ) and log_once("_fallback_to_arrow_object_extension_type_warning"): + logger.warning( + f"Failed to convert column '{column_name}' into pyarrow " + f"array due to: {ace}; {object_ext_type_detail}", + exc_info=ace, + ) + + if not object_ext_type_fallback_allowed: + # If `ArrowPythonObjectType` is not supported raise original exception + raise + + # Otherwise, attempt to fall back to serialize as python objects + return ArrowPythonObjectArray.from_objects(column_values) + + +def _convert_to_pyarrow_native_array( + column_values: Union[List[Any], np.ndarray], column_name: str +) -> pa.Array: + """Converts provided NumPy `ndarray` into PyArrow's `array` while only utilizing + Arrow's natively supported types (ie no custom extension types)""" + + try: + # NOTE: Python's `datetime` only supports precision up to us and could + # inadvertently lose precision when handling Pandas `Timestamp` type. + # To avoid that we convert provided list of `datetime` objects into + # ndarray of `np.datetime64` + if len(column_values) > 0 and isinstance(column_values[0], datetime): + column_values = _convert_datetime_to_np_datetime(column_values) + + # To avoid deserialization penalty of converting Arrow arrays (`Array` and `ChunkedArray`) + # to Python objects and then back to Arrow, we instead combine them into ListArray manually + if len(column_values) > 0 and isinstance( + column_values[0], (pa.Array, pa.ChunkedArray) + ): + return _combine_as_list_array(column_values) + + # NOTE: We explicitly infer PyArrow `DataType` so that + # we can perform upcasting to be able to accommodate + # blocks that are larger than 2Gb in size (limited + # by int32 offsets used by Arrow internally) + pa_type = _infer_pyarrow_type(column_values) + + if pa_type and pa.types.is_timestamp(pa_type): + # NOTE: Quirky Arrow behavior will coerce unsupported Numpy `datetime64` + # precisions that are nested inside a list type, but won't do it, + # if these are top-level ndarray. To work this around we have to cast + # ndarray values manually + if isinstance(column_values, np.ndarray): + column_values = _coerce_np_datetime_to_pa_timestamp_precision( + column_values, pa_type, column_name + ) + + logger.log( + logging.getLevelName("TRACE"), + f"Inferred dtype of '{pa_type}' for column '{column_name}'", + ) + + # NOTE: Pyarrow 19.0 is not able to properly handle `ListScalar(None)` when + # creating native array and hence we have to manually replace any such + # cases w/ an explicit null value + # + # See for more details https://github.com/apache/arrow/issues/45682 + if len(column_values) > 0 and isinstance(column_values[0], pa.ListScalar): + column_values = [v if v.is_valid else None for v in column_values] + + return pa.array(column_values, type=pa_type) + except Exception as e: + raise ArrowConversionError(str(column_values)) from e + + +def _combine_as_list_array(column_values: List[Union[pa.Array, pa.ChunkedArray]]): + """Combines list of Arrow arrays into a single `ListArray`""" + + # First, compute respective offsets in the resulting array + lens = [len(v) for v in column_values] + offsets = pa.array(np.concatenate([[0], np.cumsum(lens)]), type=pa.int32()) + + # Concat all the chunks into a single contiguous array + combined = pa.concat_arrays( + itertools.chain( + *[ + v.chunks if isinstance(v, pa.ChunkedArray) else [v] + for v in column_values + ] + ) + ) + + # TODO support null masking + return pa.ListArray.from_arrays(offsets, combined, pa.list_(combined.type)) + + +def _coerce_np_datetime_to_pa_timestamp_precision( + column_values: np.ndarray, dtype: pa.TimestampType, column_name: str +): + assert np.issubdtype(column_values.dtype, np.datetime64) + + numpy_precision, _ = np.datetime_data(column_values.dtype) + arrow_precision = dtype.unit + + if arrow_precision != numpy_precision: + # Arrow supports fewer timestamp resolutions than NumPy. So, if Arrow + # doesn't support the resolution, we need to cast the NumPy array to a + # different type. This can be a lossy conversion. + column_values = column_values.astype(f"datetime64[{arrow_precision}]") + + if log_once(f"column_{column_name}_timestamp_warning"): + logger.warning( + f"Converting a {numpy_precision!r} precision datetime NumPy " + f"array to '{arrow_precision}' precision Arrow timestamp. This " + "conversion occurs because Arrow supports fewer precisions " + "than Arrow and might result in a loss of precision or " + "unrepresentable values." + ) + + return column_values + + +def _infer_pyarrow_type( + column_values: Union[List[Any], np.ndarray] +) -> Optional[pa.DataType]: + """Infers target Pyarrow `DataType` based on the provided + columnar values. + + NOTE: This is a wrapper on top of `pa.infer_type(...)` utility + performing up-casting of `binary` and `string` types to + corresponding `large_binary` and `large_string` types in case + any of the array elements exceeds 2Gb in size therefore + making it impossible for original types to accommodate such + values. + + Unfortunately, for unknown reasons PA doesn't perform + that upcasting itself henceforth we have to do perform + it manually + + Args: + column_values: List of columnar values + + Returns: + Instance of PyArrow's `DataType` based on the provided + column values + """ + + if len(column_values) == 0: + return None + + # `pyarrow.infer_type` leaks memory if you pass an array with a datetime64 dtype. + # To avoid this, we handle datetime64 dtypes separately. + # See https://github.com/apache/arrow/issues/45493. + dtype_with_timestamp_type = _try_infer_pa_timestamp_type(column_values) + + if dtype_with_timestamp_type is not None: + return dtype_with_timestamp_type + + inferred_pa_dtype = pa.infer_type(column_values) + + def _len_gt_overflow_threshold(obj: Any) -> bool: + # NOTE: This utility could be seeing objects other than strings or bytes in + # cases when column contains non-scalar non-homogeneous object types as + # column values, therefore making Arrow unable to infer corresponding + # column type appropriately, therefore falling back to assume the type + # of the first element in the list. + # + # Check out test cases for this method for an additional context. + if isinstance(obj, (str, bytes)): + return len(obj) > INT32_MAX + + return False + + if pa.types.is_binary(inferred_pa_dtype) and any( + [_len_gt_overflow_threshold(v) for v in column_values] + ): + return pa.large_binary() + elif pa.types.is_string(inferred_pa_dtype) and any( + [_len_gt_overflow_threshold(v) for v in column_values] + ): + return pa.large_string() + + return inferred_pa_dtype + + +_NUMPY_TO_ARROW_PRECISION_MAP = { + # Coarsest timestamp precision in Arrow is seconds + "Y": "s", + "D": "s", + "M": "s", + "W": "s", + "h": "s", + "m": "s", + "s": "s", + "ms": "ms", + "us": "us", + "ns": "ns", + # Finest timestamp precision in Arrow is nanoseconds + "ps": "ns", + "fs": "ns", + "as": "ns", +} + + +def _try_infer_pa_timestamp_type( + column_values: Union[List[Any], np.ndarray] +) -> Optional[pa.DataType]: + if isinstance(column_values, list) and len(column_values) > 0: + # In case provided column values is a list of elements, this + # utility assumes homogeneity (in line with the behavior of Arrow + # type inference utils) + element_type = _try_infer_pa_timestamp_type(column_values[0]) + return pa.list_(element_type) if element_type else None + + if isinstance(column_values, np.ndarray) and np.issubdtype( + column_values.dtype, np.datetime64 + ): + np_precision, _ = np.datetime_data(column_values.dtype) + return pa.timestamp(_NUMPY_TO_ARROW_PRECISION_MAP[np_precision]) + + else: + return None + + +@DeveloperAPI +def get_arrow_extension_tensor_types(): + """Returns list of extension types of Arrow Array holding + multidimensional tensors + """ + return ( + *get_arrow_extension_fixed_shape_tensor_types(), + *get_arrow_extension_variable_shape_tensor_types(), + ) + + +@DeveloperAPI +def get_arrow_extension_fixed_shape_tensor_types(): + """Returns list of Arrow extension types holding multidimensional + tensors of *fixed* shape + """ + return ArrowTensorType, ArrowTensorTypeV2 + + +@DeveloperAPI +def get_arrow_extension_variable_shape_tensor_types(): + """Returns list of Arrow extension types holding multidimensional + tensors of *fixed* shape + """ + return (ArrowVariableShapedTensorType,) + + +# ArrowExtensionSerializeDeserializeCache needs to be first in the MRO to ensure the cache is used +class _BaseFixedShapeArrowTensorType( + ArrowExtensionSerializeDeserializeCache, pa.ExtensionType +): + """ + Arrow ExtensionType for an array of fixed-shaped, homogeneous-typed + tensors. + + This is the Arrow side of TensorDtype. + + See Arrow extension type docs: + https://arrow.apache.org/docs/python/extending_types.html#defining-extension-types-user-defined-types + """ + + def __init__( + self, shape: Tuple[int, ...], tensor_dtype: pa.DataType, ext_type_id: str + ): + self._shape = shape + super().__init__(tensor_dtype, ext_type_id) + + @property + def shape(self) -> Tuple[int, ...]: + """ + Shape of contained tensors. + """ + return self._shape + + @property + def scalar_type(self) -> pa.DataType: + """Returns the type of the underlying tensor elements.""" + return self.storage_type.value_type + + def to_pandas_dtype(self): + """ + Convert Arrow extension type to corresponding Pandas dtype. + + Returns: + An instance of pd.api.extensions.ExtensionDtype. + """ + from ray.air.util.tensor_extensions.pandas import TensorDtype + + return TensorDtype(self._shape, self.scalar_type.to_pandas_dtype()) + + def __reduce__(self): + return self.__arrow_ext_deserialize__, ( + self.storage_type, + self.__arrow_ext_serialize__(), + ) + + def _arrow_ext_serialize_compute(self): + if ARROW_EXTENSION_SERIALIZATION_FORMAT == _SerializationFormat.CLOUDPICKLE: + return cloudpickle.dumps(self._shape) + elif ARROW_EXTENSION_SERIALIZATION_FORMAT == _SerializationFormat.JSON: + return json.dumps(self._shape).encode() + else: + raise ValueError( + f"Invalid serialization format: {ARROW_EXTENSION_SERIALIZATION_FORMAT}" + ) + + def __arrow_ext_class__(self): + """ + ExtensionArray subclass with custom logic for this array of tensors + type. + + Returns: + A subclass of pd.api.extensions.ExtensionArray. + """ + return ArrowTensorArray + + def __arrow_ext_scalar_class__(self): + """ + ExtensionScalar subclass with custom logic for this array of tensors type. + """ + return ArrowTensorScalar + + def _extension_scalar_to_ndarray(self, scalar: "pa.ExtensionScalar") -> np.ndarray: + """ + Convert an ExtensionScalar to a tensor element. + """ + # Handle None/null values + if scalar.value is None: + return None + + raw_values = scalar.value.values + shape = scalar.type.shape + value_type = raw_values.type + offset = raw_values.offset + data_buffer = raw_values.buffers()[1] + return _to_ndarray_helper(shape, value_type, offset, data_buffer) + + def __str__(self) -> str: + return f"{self.__class__.__name__}(shape={self.shape}, dtype={self.storage_type.value_type})" + + def __repr__(self) -> str: + return str(self) + + def __eq__(self, other): + return ( + isinstance(other, type(self)) + and other.extension_name == self.extension_name + and other.shape == self.shape + and other.scalar_type == self.scalar_type + ) + + def __ne__(self, other): + # NOTE: We override ``__ne__`` to override base class' method + return not self.__eq__(other) + + def __hash__(self) -> int: + return hash((self.extension_name, self.scalar_type, self._shape)) + + +@PublicAPI(stability="beta") +class ArrowTensorType(_BaseFixedShapeArrowTensorType): + """Arrow ExtensionType (v1) for tensors. + + NOTE: This type does *NOT* support tensors larger than 4Gb (due to + overflow of int32 offsets utilized inside Pyarrow `ListType`) + """ + + OFFSET_DTYPE = pa.int32() + + def __init__(self, shape: Tuple[int, ...], dtype: pa.DataType): + """ + Construct the Arrow extension type for array of fixed-shaped tensors. + + Args: + shape: Shape of contained tensors. + dtype: pyarrow dtype of tensor elements. + """ + + super().__init__(shape, pa.list_(dtype), "ray.data.arrow_tensor") + + @classmethod + def _get_deserialize_parameter(cls, storage_type, serialized): + return (serialized, storage_type.value_type) + + @classmethod + def _arrow_ext_deserialize_compute(cls, serialized, value_type): + shape = tuple(_deserialize_with_fallback(serialized, "shape")) + return cls(shape, value_type) + + +@PublicAPI(stability="alpha") +class ArrowTensorTypeV2(_BaseFixedShapeArrowTensorType): + """Arrow ExtensionType (v2) for tensors (supporting tensors > 4Gb).""" + + OFFSET_DTYPE = pa.int64() + + def __init__(self, shape: Tuple[int, ...], dtype: pa.DataType): + """ + Construct the Arrow extension type for array of fixed-shaped tensors. + + Args: + shape: Shape of contained tensors. + dtype: pyarrow dtype of tensor elements. + """ + + super().__init__(shape, pa.large_list(dtype), "ray.data.arrow_tensor_v2") + + @classmethod + def _get_deserialize_parameter(cls, storage_type, serialized): + return (serialized, storage_type.value_type) + + @classmethod + def _arrow_ext_deserialize_compute(cls, serialized, value_type): + shape = tuple(_deserialize_with_fallback(serialized, "shape")) + return cls(shape, value_type) + + +@PublicAPI(stability="beta") +class ArrowTensorScalar(pa.ExtensionScalar): + def as_py(self, **kwargs) -> np.ndarray: + return self.__array__() + + def __array__(self) -> np.ndarray: + return self.type._extension_scalar_to_ndarray(self) + + +@PublicAPI(stability="beta") +class ArrowTensorArray(pa.ExtensionArray): + """ + An array of fixed-shape, homogeneous-typed tensors. + + This is the Arrow side of TensorArray. + + See Arrow docs for customizing extension arrays: + https://arrow.apache.org/docs/python/extending_types.html#custom-extension-array-class + """ + + @classmethod + def from_numpy( + cls, + arr: Union[np.ndarray, Iterable[np.ndarray]], + column_name: Optional[str] = None, + ) -> Union["ArrowTensorArray", "ArrowVariableShapedTensorArray"]: + """ + Convert an ndarray or an iterable of ndarrays to an array of homogeneous-typed + tensors. If given fixed-shape tensor elements, this will return an + ``ArrowTensorArray``; if given variable-shape tensor elements, this will return + an ``ArrowVariableShapedTensorArray``. + + Args: + arr: An ndarray or an iterable of ndarrays. + column_name: Optional. Used only in logging outputs to provide + additional details. + + Returns: + - If fixed-shape tensor elements, an ``ArrowTensorArray`` containing + ``len(arr)`` tensors of fixed shape. + - If variable-shaped tensor elements, an ``ArrowVariableShapedTensorArray`` + containing ``len(arr)`` tensors of variable shape. + - If scalar elements, a ``pyarrow.Array``. + """ + if not isinstance(arr, np.ndarray) and isinstance(arr, Iterable): + arr = list(arr) + + if isinstance(arr, (list, tuple)) and arr and isinstance(arr[0], np.ndarray): + # Stack ndarrays and pass through to ndarray handling logic below. + try: + arr = np.stack(arr, axis=0) + except ValueError as ve: + logger.warning( + f"Failed to stack lists due to: {ve}; " + f"falling back to using np.array(..., dtype=object)", + exc_info=ve, + ) + + # ndarray stacking may fail if the arrays are heterogeneously-shaped. + arr = np.array(arr, dtype=object) + + if not isinstance(arr, np.ndarray): + raise ValueError( + f"Must give ndarray or iterable of ndarrays, got {type(arr)} {arr}" + ) + + try: + timestamp_dtype = _try_infer_pa_timestamp_type(arr) + + if timestamp_dtype: + # NOTE: Quirky Arrow behavior will coerce unsupported Numpy `datetime64` + # precisions that are nested inside a list type, but won't do it, + # if these are top-level ndarray. To work this around we have to cast + # ndarray values manually + arr = _coerce_np_datetime_to_pa_timestamp_precision( + arr, timestamp_dtype, column_name + ) + + return cls._from_numpy(arr) + except Exception as e: + data_str = "" + if column_name: + data_str += f"column: '{column_name}', " + data_str += f"shape: {arr.shape}, dtype: {arr.dtype}, data: {arr}" + raise ArrowConversionError(data_str) from e + + @classmethod + def _from_numpy( + cls, + arr: np.ndarray, + ) -> Union["ArrowTensorArray", "ArrowVariableShapedTensorArray"]: + if len(arr) > 0 and np.isscalar(arr[0]): + # Elements are scalar so a plain Arrow Array will suffice. + return pa.array(arr) + + if _is_ndarray_variable_shaped_tensor(arr): + # Tensor elements have variable shape, so we delegate to + # ArrowVariableShapedTensorArray. + return ArrowVariableShapedTensorArray.from_numpy(arr) + + if not arr.flags.c_contiguous: + # We only natively support C-contiguous ndarrays. + arr = np.ascontiguousarray(arr) + + scalar_dtype = pa.from_numpy_dtype(arr.dtype) + + if pa.types.is_string(scalar_dtype): + if arr.dtype.byteorder == ">" or ( + arr.dtype.byteorder == "=" and sys.byteorder == "big" + ): + raise ValueError( + "Only little-endian string tensors are supported, " + f"but got: {arr.dtype}", + ) + scalar_dtype = pa.binary(arr.dtype.itemsize) + + outer_len = arr.shape[0] + element_shape = arr.shape[1:] + total_num_items = arr.size + num_items_per_element = np.prod(element_shape) if element_shape else 1 + + # Shape up data buffer + if pa.types.is_boolean(scalar_dtype): + # NumPy doesn't represent boolean arrays as bit-packed, so we manually + # bit-pack the booleans before handing the buffer off to Arrow. + # NOTE: Arrow expects LSB bit-packed ordering. + # NOTE: This creates a copy. + arr = np.packbits(arr, bitorder="little") + + data_buffer = pa.py_buffer(arr) + data_array = pa.Array.from_buffers( + scalar_dtype, total_num_items, [None, data_buffer] + ) + + from ray.data import DataContext + + if DataContext.get_current().use_arrow_tensor_v2: + pa_type_ = ArrowTensorTypeV2(element_shape, scalar_dtype) + else: + pa_type_ = ArrowTensorType(element_shape, scalar_dtype) + + offset_dtype = pa_type_.OFFSET_DTYPE.to_pandas_dtype() + + # Create offsets buffer + if num_items_per_element == 0: + offsets = np.zeros(outer_len + 1, dtype=offset_dtype) + else: + offsets = np.arange( + 0, + (outer_len + 1) * num_items_per_element, + num_items_per_element, + dtype=offset_dtype, + ) + offset_buffer = pa.py_buffer(offsets) + + storage = pa.Array.from_buffers( + pa_type_.storage_type, + outer_len, + [None, offset_buffer], + children=[data_array], + ) + + return pa_type_.wrap_array(storage) + + def to_numpy(self, zero_copy_only: bool = True): + """ + Convert the entire array of tensors into a single ndarray. + + Args: + zero_copy_only: If True, an exception will be raised if the + conversion to a NumPy array would require copying the + underlying data (e.g. in presence of nulls, or for + non-primitive types). This argument is currently ignored, so + zero-copy isn't enforced even if this argument is true. + + Returns: + A single ndarray representing the entire array of tensors. + """ + + # Buffers layout: [None, offset_buffer, None, data_buffer] + buffers = self.buffers() + data_buffer = buffers[3] + storage_list_type = self.storage.type + value_type = storage_list_type.value_type + shape = self.type.shape + + # Batch type checks + is_boolean = pa.types.is_boolean(value_type) + + # Calculate buffer item width once + if is_boolean: + # Arrow boolean array buffers are bit-packed, with 8 entries per byte, + # and are accessed via bit offsets. + buffer_item_width = value_type.bit_width + else: + # We assume all other array types are accessed via byte array + # offsets. + buffer_item_width = value_type.bit_width // 8 + + # Number of items per inner ndarray. + num_items_per_element = np.prod(shape) if shape else 1 + # Base offset into data buffer, e.g. due to zero-copy slice. + buffer_offset = self.offset * num_items_per_element + # Offset of array data in buffer. + offset = buffer_item_width * buffer_offset + # Update the shape for ndarray + shape = (len(self),) + shape + + if is_boolean: + # Special handling for boolean arrays, since Arrow bit-packs boolean arrays + # while NumPy does not. + # Cast as uint8 array and let NumPy unpack into a boolean view. + # Offset into uint8 array, where each element is a bucket for 8 booleans. + byte_bucket_offset = offset // 8 + # Offset for a specific boolean, within a uint8 array element. + bool_offset = offset % 8 + # The number of uint8 array elements (buckets) that our slice spans. + # Note that, due to the offset for a specific boolean, the slice can span + # byte boundaries even if it contains less than 8 booleans. + num_boolean_byte_buckets = 1 + ((bool_offset + np.prod(shape) - 1) // 8) + # Construct the uint8 array view on the buffer. + arr = np.ndarray( + (num_boolean_byte_buckets,), + dtype=np.uint8, + buffer=data_buffer, + offset=byte_bucket_offset, + ) + # Unpack into a byte per boolean, using LSB bit-packed ordering. + arr = np.unpackbits(arr, bitorder="little") + # Interpret buffer as boolean array. + return np.ndarray(shape, dtype=np.bool_, buffer=arr, offset=bool_offset) + + # Special handling of binary/string types. Assumes unicode string tensor columns + if pa.types.is_fixed_size_binary(value_type): + ext_dtype = np.dtype( + f" "ArrowVariableShapedTensorArray": + """ + Convert this tensor array to a variable-shaped tensor array. + """ + + shape = self.type.shape + if ndim < len(shape): + raise ValueError( + f"Can't convert {self.type} to var-shaped tensor type with {ndim=}" + ) + + # NOTE: For ``ArrowTensorTypeV2`` we can construct variable-shaped + # tensor directly w/o modifying its internal storage. + # + # For (deprecated) ``ArrowTensorType`` we fallback to converting to Numpy, + # and reconstructing. + if not isinstance(self.type, ArrowTensorTypeV2): + return ArrowVariableShapedTensorArray.from_numpy(self.to_numpy()) + + # Pad target shape with singleton axis to match target number of + # dimensions + # TODO avoid padding + target_shape = _pad_shape_with_singleton_axes(shape, ndim) + # Construct shapes array + shape_array = pa.nulls( + len(self.storage), + type=ArrowVariableShapedTensorArray.SHAPES_ARRAY_TYPE, + ).fill_null(target_shape) + + storage = pa.StructArray.from_arrays( + [self.storage, shape_array], + ["data", "shape"], + ) + + target_type = ArrowVariableShapedTensorType( + self.type.scalar_type, + ndim=ndim, + ) + + return target_type.wrap_array(storage) + + +# ArrowExtensionSerializeDeserializeCache needs to be first in the MRO to ensure the cache is used +@PublicAPI(stability="alpha") +class ArrowVariableShapedTensorType( + ArrowExtensionSerializeDeserializeCache, pa.ExtensionType +): + """ + Arrow ExtensionType for an array of heterogeneous-shaped, homogeneous-typed + tensors. + + This is the Arrow side of ``TensorDtype`` for tensor elements with different shapes. + + NOTE: This extension only supports tensor elements with non-ragged, well-defined + shapes; i.e. every tensor element must have a well-defined shape and all of their + shapes have to have same number of dimensions (ie ``len(shape)`` has to be the + same for all of them). + + See Arrow extension type docs: + https://arrow.apache.org/docs/python/extending_types.html#defining-extension-types-user-defined-types + """ + + OFFSET_DTYPE = pa.int64() + + def __init__(self, dtype: pa.DataType, ndim: int): + """ + Construct the Arrow extension type for array of heterogeneous-shaped tensors. + + Args: + dtype: pyarrow dtype of tensor elements. + ndim: The number of dimensions in the tensor elements. + """ + self._ndim = ndim + super().__init__( + pa.struct( + [("data", pa.large_list(dtype)), ("shape", pa.list_(self.OFFSET_DTYPE))] + ), + "ray.data.arrow_variable_shaped_tensor", + ) + + def to_pandas_dtype(self): + """ + Convert Arrow extension type to corresponding Pandas dtype. + + Returns: + An instance of pd.api.extensions.ExtensionDtype. + """ + from ray.air.util.tensor_extensions.pandas import TensorDtype + + return TensorDtype( + self.shape, + self.storage_type["data"].type.value_type.to_pandas_dtype(), + ) + + @property + def ndim(self) -> int: + """Return the number of dimensions in the tensor elements.""" + return self._ndim + + @property + def shape(self) -> Tuple[None, ...]: + return (None,) * self.ndim + + @property + def scalar_type(self) -> pa.DataType: + """Returns the type of the underlying tensor elements.""" + data_field_index = self.storage_type.get_field_index("data") + return self.storage_type[data_field_index].type.value_type + + def __reduce__(self): + return self.__arrow_ext_deserialize__, ( + self.storage_type, + self.__arrow_ext_serialize__(), + ) + + def _arrow_ext_serialize_compute(self): + if ARROW_EXTENSION_SERIALIZATION_FORMAT == _SerializationFormat.CLOUDPICKLE: + return cloudpickle.dumps(self._ndim) + elif ARROW_EXTENSION_SERIALIZATION_FORMAT == _SerializationFormat.JSON: + return json.dumps(self._ndim).encode() + else: + raise ValueError( + f"Invalid serialization format: {ARROW_EXTENSION_SERIALIZATION_FORMAT}" + ) + + @classmethod + def _get_deserialize_parameter(cls, storage_type, serialized): + return (serialized, storage_type["data"].type.value_type) + + @classmethod + def _arrow_ext_deserialize_compute(cls, serialized, value_type): + ndim = _deserialize_with_fallback(serialized, "ndim") + return cls(value_type, ndim) + + def __arrow_ext_class__(self): + """ + ExtensionArray subclass with custom logic for this array of tensors + type. + + Returns: + A subclass of pd.api.extensions.ExtensionArray. + """ + return ArrowVariableShapedTensorArray + + def __arrow_ext_scalar_class__(self): + """ + ExtensionScalar subclass with custom logic for this array of tensors type. + """ + return ArrowTensorScalar + + def __str__(self) -> str: + dtype = self.storage_type["data"].type.value_type + return f"ArrowVariableShapedTensorType(ndim={self.ndim}, dtype={dtype})" + + def __repr__(self) -> str: + return str(self) + + def __eq__(self, other): + # NOTE: This check is deliberately not comparing the ``ndim`` since + # we allow tensor types w/ varying ``ndim``s to be combined + return ( + isinstance(other, ArrowVariableShapedTensorType) + and other.extension_name == self.extension_name + and other.scalar_type == self.scalar_type + ) + + def __ne__(self, other): + # NOTE: We override ``__ne__`` to override base class' method + return not self.__eq__(other) + + def __hash__(self) -> int: + return hash((self.extension_name, self.scalar_type)) + + def _extension_scalar_to_ndarray(self, scalar: "pa.ExtensionScalar") -> np.ndarray: + """ + Convert an ExtensionScalar to a tensor element. + """ + + # Handle None/null values + if scalar.value is None: + return None + + data = scalar.value.get("data") + raw_values = data.values + value_type = raw_values.type + offset = raw_values.offset + data_buffer = raw_values.buffers()[1] + + shape = tuple(scalar.value.get("shape").as_py()) + + return _to_ndarray_helper(shape, value_type, offset, data_buffer) + + +@PublicAPI(stability="alpha") +class ArrowVariableShapedTensorArray(pa.ExtensionArray): + """ + An array of heterogeneous-shaped, homogeneous-typed tensors. + + This is the Arrow side of TensorArray for tensor elements that have differing + shapes. Note that this extension only supports non-ragged tensor elements; i.e., + when considering each tensor element in isolation, they must have a well-defined + shape. This extension also only supports tensor elements that all have the same + number of dimensions. + + See Arrow docs for customizing extension arrays: + https://arrow.apache.org/docs/python/extending_types.html#custom-extension-array-class + """ + + SHAPES_ARRAY_TYPE = pa.list_(pa.int64()) + + @classmethod + def from_numpy( + cls, arr: Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray]] + ) -> "ArrowVariableShapedTensorArray": + """ + Convert an ndarray or an iterable of heterogeneous-shaped ndarrays to an array + of heterogeneous-shaped, homogeneous-typed tensors. + + Args: + arr: An ndarray or an iterable of heterogeneous-shaped ndarrays. + + Returns: + An ArrowVariableShapedTensorArray containing len(arr) tensors of + heterogeneous shape. + """ + # Implementation note - Arrow representation of ragged tensors: + # + # We represent an array of ragged tensors using a struct array containing two + # fields: + # - data: a variable-sized list array, where each element in the array is a + # tensor element stored in a 1D (raveled) variable-sized list of the + # underlying scalar data type. + # - shape: a variable-sized list array containing the shapes of each tensor + # element. + if not isinstance(arr, (list, tuple, np.ndarray)): + raise ValueError( + "ArrowVariableShapedTensorArray can only be constructed from an " + f"ndarray or a list/tuple of ndarrays, but got: {type(arr)}" + ) + + if len(arr) == 0: + # Empty ragged tensor arrays are not supported. + raise ValueError("Creating empty ragged tensor arrays is not supported.") + + # Pre-allocate lists for better performance + raveled = np.empty(len(arr), dtype=np.object_) + shapes = np.empty(len(arr), dtype=np.object_) + + sizes = np.arange(len(arr), dtype=np.int64) + + ndim = None + + for i, a in enumerate(arr): + a = np.asarray(a) + + if ndim is not None and a.ndim != ndim: + raise ValueError( + "ArrowVariableShapedTensorArray only supports tensor elements that " + "all have the same number of dimensions, but got tensor elements " + f"with dimensions: {ndim}, {a.ndim}" + ) + + ndim = a.ndim + shapes[i] = a.shape + sizes[i] = a.size + # Convert to 1D array view; this should be zero-copy in the common case. + # NOTE: If array is not in C-contiguous order, this will convert it to + # C-contiguous order, incurring a copy. + raveled[i] = np.ravel(a, order="C") + + # Get size offsets and total size. + size_offsets = np.cumsum(sizes) + total_size = size_offsets[-1] + + # An optimized zero-copy path if raveled tensor elements are already + # contiguous in memory, e.g. if this tensor array has already done a + # roundtrip through our Arrow representation. + data_buffer = _concat_ndarrays(raveled) + + dtype = data_buffer.dtype + pa_scalar_type = pa.from_numpy_dtype(dtype) + + if pa.types.is_string(pa_scalar_type): + if dtype.byteorder == ">" or ( + dtype.byteorder == "=" and sys.byteorder == "big" + ): + raise ValueError( + "Only little-endian string tensors are supported, " + f"but got: {dtype}" + ) + pa_scalar_type = pa.binary(dtype.itemsize) + + if dtype.type is np.bool_ and data_buffer.size > 0: + # NumPy doesn't represent boolean arrays as bit-packed, so we manually + # bit-pack the booleans before handing the buffer off to Arrow. + # NOTE: Arrow expects LSB bit-packed ordering. + # NOTE: This creates a copy. + data_buffer = np.packbits(data_buffer, bitorder="little") + + # Use foreign_buffer for better performance when possible + data_buffer = pa.py_buffer(data_buffer) + # Construct underlying data array. + data_array = pa.Array.from_buffers( + pa_scalar_type, total_size, [None, data_buffer] + ) + + # Construct array for offsets into the 1D data array, where each offset + # corresponds to a tensor element. + size_offsets = np.insert(size_offsets, 0, 0) + offset_array = pa.array(size_offsets) + data_storage_array = pa.LargeListArray.from_arrays(offset_array, data_array) + # We store the tensor element shapes so we can reconstruct each tensor when + # converting back to NumPy ndarrays. + shape_array = pa.array(shapes) + + # Build storage array containing tensor data and the tensor element shapes. + storage = pa.StructArray.from_arrays( + [data_storage_array, shape_array], + ["data", "shape"], + ) + + type_ = ArrowVariableShapedTensorType(pa_scalar_type, ndim) + return type_.wrap_array(storage) + + def to_numpy(self, zero_copy_only: bool = True): + """ + Convert the entire array of tensors into a single ndarray. + + Args: + zero_copy_only: If True, an exception will be raised if the conversion to a + NumPy array would require copying the underlying data (e.g. in presence + of nulls, or for non-primitive types). This argument is currently + ignored, so zero-copy isn't enforced even if this argument is true. + + Returns: + A single ndarray representing the entire array of tensors. + """ + + data_array = self.storage.field("data") + shapes_array = self.storage.field("shape") + + data_value_type = data_array.type.value_type + data_array_buffer = data_array.buffers()[3] + + shapes = shapes_array.to_pylist() + offsets = data_array.offsets.to_pylist() + + return create_ragged_ndarray( + [ + _to_ndarray_helper(shape, data_value_type, offset, data_array_buffer) + for shape, offset in zip(shapes, offsets) + ] + ) + + def to_var_shaped_tensor_array(self, ndim: int) -> "ArrowVariableShapedTensorArray": + if ndim == self.type.ndim: + return self + elif ndim < self.type.ndim: + raise ValueError( + f"Can't convert {self.type} to var-shaped tensor type with {ndim=}" + ) + + target_type = ArrowVariableShapedTensorType(self.type.scalar_type, ndim) + + # Unpack source tensor array into internal data storage and shapes + # array + data_array = self.storage.field("data") + shapes_array = self.storage.field("shape") + # Pad individual shapes with singleton axes to match target number of + # dimensions + # + # TODO avoid python loop + expanded_shapes_array = pa.array( + [_pad_shape_with_singleton_axes(s, ndim) for s in shapes_array.to_pylist()] + ) + + storage = pa.StructArray.from_arrays([data_array, expanded_shapes_array]) + + return target_type.wrap_array(storage) + + +def _pad_shape_with_singleton_axes( + shape: Tuple[int, ...], ndim: int +) -> Tuple[int, ...]: + assert ndim >= len(shape) + + return (1,) * (ndim - len(shape)) + shape + + +AnyArrowExtTensorType = Union[ + ArrowTensorType, ArrowTensorTypeV2, ArrowVariableShapedTensorType +] + + +@DeveloperAPI(stability="alpha") +def unify_tensor_types( + types: Collection[AnyArrowExtTensorType], +) -> AnyArrowExtTensorType: + """Unifies provided tensor types if compatible. + + Otherwise raises a ``ValueError``. + """ + + assert types, "List of tensor types may not be empty" + + if len(types) == 1: + return types[0] + + shapes = {t.shape for t in types} + scalar_types = {t.scalar_type for t in types} + + # Only tensors with homogenous scalar types and shape dimensions + # are currently supported + if len(scalar_types) > 1: + raise pa.lib.ArrowTypeError( + f"Can't unify tensor types with divergent scalar types: {types}" + ) + + # If all shapes are identical, it's a single tensor type + if len(shapes) == 1: + return next(iter(types)) + + return ArrowVariableShapedTensorType( + dtype=scalar_types.pop(), + # NOTE: Cardinality of variable-shaped tensor type's (``ndims``) is + # derived as the max length of the shapes that are making it up + ndim=max(len(s) for s in shapes), + ) + + +@DeveloperAPI(stability="alpha") +def unify_tensor_arrays( + arrs: List[Union[ArrowTensorArray, ArrowVariableShapedTensorArray]] +) -> List[Union[ArrowTensorArray, ArrowVariableShapedTensorArray]]: + supported_tensor_types = get_arrow_extension_tensor_types() + + # Derive number of distinct tensor types + distinct_types_ = set() + + for arr in arrs: + if isinstance(arr.type, supported_tensor_types): + distinct_types_.add(arr.type) + else: + raise ValueError( + f"Trying to unify unsupported tensor type: {arr.type} (supported types: {supported_tensor_types})" + ) + + if len(distinct_types_) == 1: + return arrs + + # Verify provided tensor arrays could be unified + # + # NOTE: If there's more than 1 distinct tensor types, then unified + # type will be variable-shaped + unified_tensor_type = unify_tensor_types(distinct_types_) + + assert isinstance(unified_tensor_type, ArrowVariableShapedTensorType) + + unified_arrs = [] + for arr in arrs: + unified_arrs.append( + arr.to_var_shaped_tensor_array(ndim=unified_tensor_type.ndim) + ) + + return unified_arrs + + +@DeveloperAPI(stability="alpha") +def concat_tensor_arrays( + arrays: List[Union["ArrowTensorArray", "ArrowVariableShapedTensorArray"]], + ensure_copy: bool = False, +) -> Union["ArrowTensorArray", "ArrowVariableShapedTensorArray"]: + """ + Concatenates multiple tensor arrays. + + NOTE: If one or more of the tensor arrays are variable-shaped and/or any + of the tensor arrays have a different shape than the others, a variable-shaped + tensor array will be returned. + + Args: + arrays: Tensor arrays to concat + ensure_copy: Skip copying when ensure_copy is False and there is exactly 1 chunk. + + Returns: + Either ``ArrowTensorArray`` or ``ArrowVariableShapedTensorArray`` holding + all of the given tensor arrays concatenated. + """ + + assert arrays, "List of tensor arrays may not be empty" + + if len(arrays) == 1 and not ensure_copy: + # Short-circuit + return arrays[0] + + # First, unify provided tensor arrays + unified_arrays = unify_tensor_arrays(arrays) + # Then, simply concat underlying internal storage + storage = pa.concat_arrays([c.storage for c in unified_arrays]) + + unified_array_type = unified_arrays[0].type + return unified_array_type.wrap_array(storage) + + +def _concat_ndarrays(arrs: Union[np.ndarray, List[np.ndarray]]) -> np.ndarray: + """Concatenates provided collection of ``np.ndarray``s in either of the following + ways: + + - If provided ndarrays are contiguous, 1D views sharing the same dtype, + living w/in the same base view, these will be concatenated zero-copy + by reusing underlying view + + - Otherwise, ``np.concatenate(arrays)`` will be invoked + """ + + assert len(arrs) > 0, "Provided collection of ndarrays may not be empty" + + if len(arrs) == 1: + # Short-circuit + return arrs[0] + elif not _are_contiguous_1d_views(arrs): + return np.concatenate(arrs) + + dtype = arrs[0].dtype + base = _get_root_base(arrs[0]) + + base_ptr = _get_buffer_address(base) + start_byte = _get_buffer_address(arrs[0]) - base_ptr + end_byte = start_byte + sum(a.nbytes for a in arrs) + + # Build the view from the base, using byte offsets for generality + byte_view = base.view(np.uint8).reshape(-1) + out = byte_view[start_byte:end_byte].view(dtype) + + return out + + +def _are_contiguous_1d_views(arrs: Union[np.ndarray, List[np.ndarray]]) -> bool: + dtype = arrs[0].dtype + base = _get_root_base(arrs[0]) + expected_addr = _get_base_ptr(arrs[0]) + + for a in arrs: + # Assert all provided arrays are + # - Raveled (1D) + # - Share dtype + # - Contiguous + # - Share the same `base` view (this is crucial to make sure + # that all provided ndarrays live w/in the same allocation and + # share its lifecycle) + if ( + a.ndim != 1 + or a.dtype != dtype + or not a.flags.c_contiguous + or _get_root_base(a) is not base + ): + return False + # Skip empty ndarrays + if a.size == 0: + continue + + buffer_addr = _get_base_ptr(a) + if buffer_addr != expected_addr: + return False + + expected_addr = buffer_addr + a.size * dtype.itemsize + + return True + + +def _get_base_ptr(a: np.ndarray) -> int: + # same as a.ctypes.data, but robust for views + return _get_buffer_address(a) + + +def _get_root_base(a: np.ndarray) -> np.ndarray: + b = a + while isinstance(b.base, np.ndarray): + b = b.base + return b if b.base is not None else b # owner if base is None + + +def _get_buffer_address(arr: np.ndarray) -> int: + """Get the address of the buffer underlying the provided NumPy ndarray.""" + return arr.__array_interface__["data"][0] + + +def _to_ndarray_helper(shape, value_type, offset, data_buffer): + if pa.types.is_boolean(value_type): + # Arrow boolean array buffers are bit-packed, with 8 entries per byte, + # and are accessed via bit offsets. + buffer_item_width = value_type.bit_width + else: + # We assume all other array types are accessed via byte array + # offsets. + buffer_item_width = value_type.bit_width // 8 + data_offset = buffer_item_width * offset + + if pa.types.is_boolean(value_type): + # Special handling for boolean arrays, since Arrow + # bit-packs boolean arrays while NumPy does not. + # Cast as uint8 array and let NumPy unpack into a boolean view. + # Offset into uint8 array, where each element is + # a bucket for 8 booleans. + byte_bucket_offset = data_offset // 8 + # Offset for a specific boolean, within a uint8 array element. + bool_offset = data_offset % 8 + # The number of uint8 array elements (buckets) that our slice spans. + # Note that, due to the offset for a specific boolean, + # the slice can span byte boundaries even if it contains + # less than 8 booleans. + num_boolean_byte_buckets = 1 + ((bool_offset + np.prod(shape) - 1) // 8) + # Construct the uint8 array view on the buffer. + arr = np.ndarray( + (num_boolean_byte_buckets,), + dtype=np.uint8, + buffer=data_buffer, + offset=byte_bucket_offset, + ) + # Unpack into a byte per boolean, using LSB bit-packed ordering. + arr = np.unpackbits(arr, bitorder="little") + # Interpret buffer as boolean array. + return np.ndarray(shape, dtype=np.bool_, buffer=arr, offset=bool_offset) + ext_dtype = value_type.to_pandas_dtype() + # Special handling of ragged string tensors + if pa.types.is_fixed_size_binary(value_type): + ext_dtype = np.dtype(f" List[str]: + from pandas.core.construction import extract_array + from pandas.io.formats.format import format_array + + if not isinstance(self.values, TensorArray): + return self._format_strings_orig() + + values = extract_array(self.values, extract_numpy=True) + array = np.asarray(values) + + if array.ndim == 1: + return self._format_strings_orig() + + def format_array_wrap(array_, formatter_): + fmt_values = format_array( + array_, + formatter_, + float_format=self.float_format, + na_rep=self.na_rep, + digits=self.digits, + space=self.space, + justify=self.justify, + decimal=self.decimal, + leading_space=self.leading_space, + quoting=self.quoting, + ) + return fmt_values + + flat_formatter = self.formatter + if flat_formatter is None: + flat_formatter = values._formatter(boxed=True) + + # Flatten array, call function, reshape (use ravel_compat in v1.3.0) + flat_array = array.ravel("K") + fmt_flat_array = np.asarray(format_array_wrap(flat_array, flat_formatter)) + order = "F" if array.flags.f_contiguous else "C" + fmt_array = fmt_flat_array.reshape(array.shape, order=order) + + # Format the array of nested strings, use default formatter + return format_array_wrap(fmt_array, None) + + +def _format_strings_patched_v1_0_0(self) -> List[str]: + from functools import partial + + from pandas.core.construction import extract_array + from pandas.io.formats.format import format_array + from pandas.io.formats.printing import pprint_thing + + if not isinstance(self.values, TensorArray): + return self._format_strings_orig() + + values = extract_array(self.values, extract_numpy=True) + array = np.asarray(values) + + if array.ndim == 1: + return self._format_strings_orig() + + def format_array_wrap(array_, formatter_): + fmt_values = format_array( + array_, + formatter_, + float_format=self.float_format, + na_rep=self.na_rep, + digits=self.digits, + space=self.space, + justify=self.justify, + decimal=self.decimal, + leading_space=self.leading_space, + ) + return fmt_values + + flat_formatter = self.formatter + if flat_formatter is None: + flat_formatter = values._formatter(boxed=True) + + # Flatten array, call function, reshape (use ravel_compat in v1.3.0) + flat_array = array.ravel("K") + fmt_flat_array = np.asarray(format_array_wrap(flat_array, flat_formatter)) + order = "F" if array.flags.f_contiguous else "C" + fmt_array = fmt_flat_array.reshape(array.shape, order=order) + + # Slimmed down version of GenericArrayFormatter due to: + # https://github.com/pandas-dev/pandas/issues/33770 + def format_strings_slim(array_, leading_space): + formatter = partial( + pprint_thing, + escape_chars=("\t", "\r", "\n"), + ) + + def _format(x): + return str(formatter(x)) + + fmt_values = [] + for v in array_: + tpl = "{v}" if leading_space is False else " {v}" + fmt_values.append(tpl.format(v=_format(v))) + return fmt_values + + return format_strings_slim(fmt_array, self.leading_space) + + +_FORMATTER_ENABLED_ENV_VAR = "TENSOR_COLUMN_EXTENSION_FORMATTER_ENABLED" + +if os.getenv(_FORMATTER_ENABLED_ENV_VAR, "1") == "1": + if Version(pd.__version__) < Version("2.2.0"): + from pandas.io.formats.format import ExtensionArrayFormatter + + formatter_cls = ExtensionArrayFormatter + else: + from pandas.io.formats.format import _ExtensionArrayFormatter + + formatter_cls = _ExtensionArrayFormatter + formatter_cls._format_strings_orig = formatter_cls._format_strings + if Version("1.1.0") <= Version(pd.__version__) < Version("1.3.0"): + formatter_cls._format_strings = _format_strings_patched + else: + formatter_cls._format_strings = _format_strings_patched_v1_0_0 + formatter_cls._patched_by_ray_datasets = True + +########################################### +# End patching of ExtensionArrayFormatter # +########################################### + + +@PublicAPI(stability="beta") +@pd.api.extensions.register_extension_dtype +class TensorDtype(pd.api.extensions.ExtensionDtype): + """ + Pandas extension type for a column of homogeneous-typed tensors. + + This extension supports tensors in which the elements have different shapes. + However, each tensor element must be non-ragged, i.e. each tensor element must have + a well-defined, non-ragged shape. + + See: + https://github.com/pandas-dev/pandas/blob/master/pandas/core/dtypes/base.py + for up-to-date interface documentation and the subclassing contract. The + docstrings of the below properties and methods were copied from the base + ExtensionDtype. + + Examples: + >>> # Create a DataFrame with a list of ndarrays as a column. + >>> import pandas as pd + >>> import numpy as np + >>> import ray + >>> df = pd.DataFrame({ + ... "one": [1, 2, 3], + ... "two": list(np.arange(24).reshape((3, 2, 2, 2)))}) + >>> # Note the opaque np.object dtype for this column. + >>> df.dtypes # doctest: +SKIP + one int64 + two object + dtype: object + >>> # Cast column to our TensorDtype extension type. + >>> from ray.data.extensions import TensorDtype + >>> df["two"] = df["two"].astype(TensorDtype(np.int64, (3, 2, 2, 2))) + >>> # Note that the column dtype is now TensorDtype instead of + >>> # np.object. + >>> df.dtypes # doctest: +SKIP + one int64 + two TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + dtype: object + >>> # Pandas is now aware of this tensor column, and we can do the + >>> # typical DataFrame operations on this column. + >>> col = 2 * (df["two"] + 10) + >>> # The ndarrays underlying the tensor column will be manipulated, + >>> # but the column itself will continue to be a Pandas type. + >>> type(col) # doctest: +SKIP + pandas.core.series.Series + >>> col # doctest: +SKIP + 0 [[[ 2 4] + [ 6 8]] + [[10 12] + [14 16]]] + 1 [[[18 20] + [22 24]] + [[26 28] + [30 32]]] + 2 [[[34 36] + [38 40]] + [[42 44] + [46 48]]] + Name: two, dtype: TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + >>> # Once you do an aggregation on that column that returns a single + >>> # row's value, you get back our TensorArrayElement type. + >>> tensor = col.mean() + >>> type(tensor) # doctest: +SKIP + ray.data.extensions.tensor_extension.TensorArrayElement + >>> tensor # doctest: +SKIP + array([[[18., 20.], + [22., 24.]], + [[26., 28.], + [30., 32.]]]) + >>> # This is a light wrapper around a NumPy ndarray, and can easily + >>> # be converted to an ndarray. + >>> type(tensor.to_numpy()) # doctest: +SKIP + numpy.ndarray + >>> # In addition to doing Pandas operations on the tensor column, + >>> # you can now put the DataFrame into a Dataset. + >>> ds = ray.data.from_pandas(df) # doctest: +SKIP + >>> # Internally, this column is represented the corresponding + >>> # Arrow tensor extension type. + >>> ds.schema() # doctest: +SKIP + one: int64 + two: extension> + >>> # You can write the dataset to Parquet. + >>> ds.write_parquet("/some/path") # doctest: +SKIP + >>> # And you can read it back. + >>> read_ds = ray.data.read_parquet("/some/path") # doctest: +SKIP + >>> read_ds.schema() # doctest: +SKIP + one: int64 + two: extension> + >>> read_df = ray.get(read_ds.to_pandas_refs())[0] # doctest: +SKIP + >>> read_df.dtypes # doctest: +SKIP + one int64 + two TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + dtype: object + >>> # The tensor extension type is preserved along the + >>> # Pandas --> Arrow --> Parquet --> Arrow --> Pandas + >>> # conversion chain. + >>> read_df.equals(df) # doctest: +SKIP + True + """ + + # NOTE(Clark): This is apparently required to prevent integer indexing + # errors, but is an undocumented ExtensionDtype attribute. See issue: + # https://github.com/CODAIT/text-extensions-for-pandas/issues/166 + base = None + + def __init__(self, shape: Tuple[Optional[int], ...], dtype: np.dtype): + self._shape = shape + self._dtype = dtype + + @property + def type(self): + """ + The scalar type for the array, e.g. ``int`` + It's expected ``ExtensionArray[item]`` returns an instance + of ``ExtensionDtype.type`` for scalar ``item``, assuming + that value is valid (not NA). NA values do not need to be + instances of `type`. + """ + return TensorArrayElement + + @property + def element_dtype(self): + """ + The dtype of the underlying tensor elements. + """ + return self._dtype + + @property + def element_shape(self): + """ + The shape of the underlying tensor elements. This will be a tuple of Nones if + the corresponding TensorArray for this TensorDtype holds variable-shaped tensor + elements. + """ + return self._shape + + @property + def is_variable_shaped(self): + """ + Whether the corresponding TensorArray for this TensorDtype holds variable-shaped + tensor elements. + """ + return all(dim_size is None for dim_size in self.shape) + + @property + def name(self) -> str: + """ + A string identifying the data type. + Will be used for display in, e.g. ``Series.dtype`` + """ + return f"TensorDtype(shape={self._shape}, dtype={self._dtype})" + + @classmethod + def construct_from_string(cls, string: str): + r""" + Construct this type from a string. + + This is useful mainly for data types that accept parameters. + For example, a period dtype accepts a frequency parameter that + can be set as ``period[H]`` (where H means hourly frequency). + + By default, in the abstract class, just the name of the type is + expected. But subclasses can overwrite this method to accept + parameters. + + Parameters + ---------- + string : str + The name of the type, for example ``category``. + + Returns + ------- + ExtensionDtype + Instance of the dtype. + + Raises + ------ + TypeError + If a class cannot be constructed from this 'string'. + + Examples + -------- + For extension dtypes with arguments the following may be an + adequate implementation. + + >>> import re + >>> @classmethod + ... def construct_from_string(cls, string): + ... pattern = re.compile(r"^my_type\[(?P.+)\]$") + ... match = pattern.match(string) + ... if match: + ... return cls(**match.groupdict()) + ... else: + ... raise TypeError( + ... f"Cannot construct a '{cls.__name__}' from '{string}'" + ... ) + """ + import ast + import re + + if not isinstance(string, str): + raise TypeError( + f"'construct_from_string' expects a string, got {type(string)}" + ) + # Upstream code uses exceptions as part of its normal control flow and + # will pass this method bogus class names. + regex = ( + r"^(TensorDtype|numpy.ndarray)" + r"\(shape=(\((?:(?:\d+|None),?\s?)*\)), dtype=(\w+)\)$" + ) + m = re.search(regex, string) + err_msg = ( + f"Cannot construct a '{cls.__name__}' from '{string}'; expected a string " + "like 'TensorDtype(shape=(1, 2, 3), dtype=int64)'." + ) + if m is None: + raise TypeError(err_msg) + groups = m.groups() + if len(groups) != 3: + raise TypeError(err_msg) + _, shape, dtype = groups + shape = ast.literal_eval(shape) + dtype = np.dtype(dtype) + return cls(shape, dtype) + + @classmethod + def construct_array_type(cls): + """ + Return the array type associated with this dtype. + + Returns + ------- + type + """ + return TensorArray + + def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]): + """ + Convert a pyarrow (chunked) array to a TensorArray. + + This and TensorArray.__arrow_array__ make up the + Pandas extension type + array <--> Arrow extension type + array + interoperability protocol. See + https://pandas.pydata.org/pandas-docs/stable/development/extending.html#compatibility-with-apache-arrow + for more information. + """ + if isinstance(array, pa.ChunkedArray): + if array.num_chunks > 1: + # TODO(Clark): Remove concat and construct from list with + # shape. + values = np.concatenate( + [chunk.to_numpy() for chunk in array.iterchunks()] + ) + else: + values = array.chunk(0).to_numpy() + else: + values = array.to_numpy() + + return TensorArray(values) + + def __str__(self) -> str: + return self.name + + def __repr__(self) -> str: + return str(self) + + @property + def _is_boolean(self): + """ + Whether this extension array should be considered boolean. + + By default, ExtensionArrays are assumed to be non-numeric. + Setting this to True will affect the behavior of several places, + e.g. + + * is_bool + * boolean indexing + + Returns + ------- + bool + """ + # This is needed to support returning a TensorArray from .isnan(). + from pandas.core.dtypes.common import is_bool_dtype + + return is_bool_dtype(self._dtype) + + +class _TensorOpsMixin(pd.api.extensions.ExtensionScalarOpsMixin): + """ + Mixin for TensorArray operator support, applying operations on the + underlying ndarrays. + """ + + @classmethod + def _create_method(cls, op, coerce_to_dtype=True, result_dtype=None): + """ + Add support for binary operators by unwrapping, applying, and + rewrapping. + """ + + # NOTE(Clark): This overrides, but coerce_to_dtype, result_dtype might + # not be needed + + def _binop(self, other): + lvalues = self._tensor + + if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndex)): + # Rely on Pandas to unbox and dispatch to us. + return NotImplemented + + # divmod returns a tuple + if op_name in ["__divmod__", "__rdivmod__"]: + # TODO(Clark): Add support for divmod and rdivmod. + # div, mod = result + raise NotImplementedError + + if isinstance(other, (TensorArray, TensorArrayElement)): + rvalues = other._tensor + else: + rvalues = other + + result = op(lvalues, rvalues) + + # Force a TensorArray if rvalue is not a scalar. + if isinstance(self, TensorArrayElement) and ( + not isinstance(other, TensorArrayElement) or not np.isscalar(other) + ): + result_wrapped = TensorArray(result) + else: + result_wrapped = cls(result) + + return result_wrapped + + op_name = f"__{op.__name__}__" + return set_function_name(_binop, op_name, cls) + + @classmethod + def _create_logical_method(cls, op): + return cls._create_method(op) + + +class _TensorScalarCastMixin: + """ + Mixin for casting scalar tensors to a particular numeric type. + """ + + def _scalarfunc(self, func: Callable[[Any], Any]): + return func(self._tensor) + + def __complex__(self): + return self._scalarfunc(complex) + + def __float__(self): + return self._scalarfunc(float) + + def __int__(self): + return self._scalarfunc(int) + + def __hex__(self): + return self._scalarfunc(hex) + + def __oct__(self): + return self._scalarfunc(oct) + + +@PublicAPI(stability="beta") +class TensorArrayElement(_TensorOpsMixin, _TensorScalarCastMixin): + """ + Single element of a TensorArray, wrapping an underlying ndarray. + """ + + def __init__(self, values: np.ndarray): + """ + Construct a TensorArrayElement from a NumPy ndarray. + + Args: + values: ndarray that underlies this TensorArray element. + """ + self._tensor = values + + def __repr__(self): + return self._tensor.__repr__() + + def __str__(self): + return self._tensor.__str__() + + @property + def numpy_dtype(self): + """ + Get the dtype of the tensor. + :return: The numpy dtype of the backing ndarray + """ + return self._tensor.dtype + + @property + def numpy_ndim(self): + """ + Get the number of tensor dimensions. + :return: integer for the number of dimensions + """ + return self._tensor.ndim + + @property + def numpy_shape(self): + """ + Get the shape of the tensor. + :return: A tuple of integers for the numpy shape of the backing ndarray + """ + return self._tensor.shape + + @property + def numpy_size(self): + """ + Get the size of the tensor. + :return: integer for the number of elements in the tensor + """ + return self._tensor.size + + def to_numpy(self): + """ + Return the values of this element as a NumPy ndarray. + """ + return np.asarray(self._tensor) + + def __array__(self, dtype: np.dtype = None, **kwargs) -> np.ndarray: + return np.asarray(self._tensor, dtype=dtype, **kwargs) + + +@PublicAPI(stability="beta") +class TensorArray( + pd.api.extensions.ExtensionArray, + _TensorOpsMixin, + _TensorScalarCastMixin, +): + """ + Pandas `ExtensionArray` representing a tensor column, i.e. a column + consisting of ndarrays as elements. + + This extension supports tensors in which the elements have different shapes. + However, each tensor element must be non-ragged, i.e. each tensor element must have + a well-defined, non-ragged shape. + + Examples: + >>> # Create a DataFrame with a list of ndarrays as a column. + >>> import pandas as pd + >>> import numpy as np + >>> import ray + >>> from ray.data.extensions import TensorArray + >>> df = pd.DataFrame({ + ... "one": [1, 2, 3], + ... "two": TensorArray(np.arange(24).reshape((3, 2, 2, 2)))}) + >>> # Note that the column dtype is TensorDtype. + >>> df.dtypes # doctest: +SKIP + one int64 + two TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + dtype: object + >>> # Pandas is aware of this tensor column, and we can do the + >>> # typical DataFrame operations on this column. + >>> col = 2 * (df["two"] + 10) + >>> # The ndarrays underlying the tensor column will be manipulated, + >>> # but the column itself will continue to be a Pandas type. + >>> type(col) # doctest: +SKIP + pandas.core.series.Series + >>> col # doctest: +SKIP + 0 [[[ 2 4] + [ 6 8]] + [[10 12] + [14 16]]] + 1 [[[18 20] + [22 24]] + [[26 28] + [30 32]]] + 2 [[[34 36] + [38 40]] + [[42 44] + [46 48]]] + Name: two, dtype: TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + >>> # Once you do an aggregation on that column that returns a single + >>> # row's value, you get back our TensorArrayElement type. + >>> tensor = col.mean() # doctest: +SKIP + >>> type(tensor) # doctest: +SKIP + ray.data.extensions.tensor_extension.TensorArrayElement + >>> tensor # doctest: +SKIP + array([[[18., 20.], + [22., 24.]], + [[26., 28.], + [30., 32.]]]) + >>> # This is a light wrapper around a NumPy ndarray, and can easily + >>> # be converted to an ndarray. + >>> type(tensor.to_numpy()) # doctest: +SKIP + numpy.ndarray + >>> # In addition to doing Pandas operations on the tensor column, + >>> # you can now put the DataFrame into a Dataset. + >>> ds = ray.data.from_pandas(df) # doctest: +SKIP + >>> # Internally, this column is represented the corresponding + >>> # Arrow tensor extension type. + >>> ds.schema() # doctest: +SKIP + one: int64 + two: extension> + >>> # You can write the dataset to Parquet. + >>> ds.write_parquet("/some/path") # doctest: +SKIP + >>> # And you can read it back. + >>> read_ds = ray.data.read_parquet("/some/path") # doctest: +SKIP + >>> read_ds.schema() # doctest: +SKIP + one: int64 + two: extension> + + >>> read_df = ray.get(read_ds.to_pandas_refs())[0] # doctest: +SKIP + >>> read_df.dtypes # doctest: +SKIP + one int64 + two TensorDtype(shape=(3, 2, 2, 2), dtype=int64) + dtype: object + >>> # The tensor extension type is preserved along the + >>> # Pandas --> Arrow --> Parquet --> Arrow --> Pandas + >>> # conversion chain. + >>> read_df.equals(df) # doctest: +SKIP + True + """ + + SUPPORTED_REDUCERS = { + "sum": np.sum, + "all": np.all, + "any": np.any, + "min": np.min, + "max": np.max, + "mean": np.mean, + "median": np.median, + "prod": np.prod, + "std": np.std, + "var": np.var, + } + + # See https://github.com/pandas-dev/pandas/blob/master/pandas/core/arrays/base.py + # for interface documentation and the subclassing contract. + def __init__( + self, + values: Union[ + np.ndarray, + ABCSeries, + Sequence[Union[np.ndarray, TensorArrayElement]], + TensorArrayElement, + Any, + ], + ): + """ + Args: + values: A NumPy ndarray or sequence of NumPy ndarrays of equal + shape. + """ + # Try to convert some well-known objects to ndarrays before handing off to + # ndarray handling logic. + if isinstance(values, ABCSeries): + values = _create_possibly_ragged_ndarray(values) + elif isinstance(values, Sequence): + values = [ + np.asarray(v) if isinstance(v, TensorArrayElement) else v + for v in values + ] + values = _create_possibly_ragged_ndarray(values) + elif isinstance(values, TensorArrayElement): + values = np.array([np.asarray(values)], copy=False) + + if isinstance(values, np.ndarray): + if values.dtype.type is np.object_: + if len(values) == 0: + # Tensor is empty, pass through to create empty TensorArray. + pass + elif all( + isinstance(v, (np.ndarray, TensorArrayElement, Sequence)) + and not isinstance(v, str) + for v in values + ): + values = [np.asarray(v) for v in values] + # Try to convert ndarrays of ndarrays/TensorArrayElements with an + # opaque object type to a properly typed ndarray of ndarrays. + values = _create_possibly_ragged_ndarray(values) + else: + raise TypeError( + "Expected a well-typed ndarray or an object-typed ndarray of " + "ndarray pointers, but got an object-typed ndarray whose " + f"subndarrays are of type {type(values[0])}." + ) + elif isinstance(values, TensorArray): + raise TypeError("Use the copy() method to create a copy of a TensorArray.") + else: + raise TypeError( + "Expected a numpy.ndarray or sequence of numpy.ndarray, " + f"but received {values} of type {type(values).__name__} instead." + ) + assert isinstance(values, np.ndarray) + self._tensor = values + self._is_variable_shaped = None + + @classmethod + def _from_sequence( + cls, scalars, *, dtype: Optional[Dtype] = None, copy: bool = False + ): + """ + Construct a new ExtensionArray from a sequence of scalars. + + Parameters + ---------- + scalars : Sequence + Each element will be an instance of the scalar type for this + array, ``cls.dtype.type`` or be converted into this type in this + method. + dtype : dtype, optional + Construct for this particular dtype. This should be a Dtype + compatible with the ExtensionArray. + copy : bool, default False + If True, copy the underlying data. + + Returns + ------- + ExtensionArray + """ + if copy and isinstance(scalars, np.ndarray): + scalars = scalars.copy() + elif isinstance(scalars, TensorArray): + scalars = scalars._tensor.copy() if copy else scalars._tensor + return TensorArray(scalars) + + @classmethod + def _from_factorized( + cls, values: np.ndarray, original: pd.api.extensions.ExtensionArray + ): + """ + Reconstruct an ExtensionArray after factorization. + + Parameters + ---------- + values : ndarray + An integer ndarray with the factorized values. + original : ExtensionArray + The original ExtensionArray that factorize was called on. + + See Also + -------- + factorize : Top-level factorize method that dispatches here. + ExtensionArray.factorize : Encode the extension array as an enumerated + type. + """ + raise NotImplementedError + + def __getitem__( + self, item: Union[int, slice, np.ndarray] + ) -> Union["TensorArray", "TensorArrayElement"]: + """ + Select a subset of self. + + Parameters + ---------- + item : int, slice, or ndarray + * int: The position in 'self' to get. + * slice: A slice object, where 'start', 'stop', and 'step' are + integers or None + * ndarray: A 1-d boolean NumPy ndarray the same length as 'self' + + Returns + ------- + item : scalar or ExtensionArray + + Notes + ----- + For scalar ``item``, return a scalar value suitable for the array's + type. This should be an instance of ``self.dtype.type``. + For slice ``key``, return an instance of ``ExtensionArray``, even + if the slice is length 0 or 1. + For a boolean mask, return an instance of ``ExtensionArray``, filtered + to the values where ``item`` is True. + """ + # Return scalar if single value is selected, a TensorArrayElement for + # single array element, or TensorArray for slice. + if isinstance(item, int): + value = self._tensor[item] + if np.isscalar(value): + return value + else: + return TensorArrayElement(value) + else: + # BEGIN workaround for Pandas issue #42430 + if isinstance(item, tuple) and len(item) > 1 and item[0] == Ellipsis: + if len(item) > 2: + # Hopefully this case is not possible, but can't be sure + raise ValueError( + "Workaround Pandas issue #42430 not " + "implemented for tuple length > 2" + ) + item = item[1] + # END workaround for issue #42430 + if isinstance(item, TensorArray): + item = np.asarray(item) + item = check_array_indexer(self, item) + return TensorArray(self._tensor[item]) + + def __len__(self) -> int: + """ + Length of this array. + + Returns + ------- + length : int + """ + return len(self._tensor) + + @property + def dtype(self) -> pd.api.extensions.ExtensionDtype: + """ + An instance of 'ExtensionDtype'. + """ + if self.is_variable_shaped: + # A tensor is only considered variable-shaped if it's non-empty, so no + # non-empty check is needed here. + dtype = self._tensor[0].dtype + shape = (None,) * self._tensor[0].ndim + else: + dtype = self.numpy_dtype + shape = self.numpy_shape[1:] + return TensorDtype(shape, dtype) + + @property + def is_variable_shaped(self): + """ + Whether this TensorArray holds variable-shaped tensor elements. + """ + if self._is_variable_shaped is None: + self._is_variable_shaped = _is_ndarray_variable_shaped_tensor(self._tensor) + return self._is_variable_shaped + + @property + def nbytes(self) -> int: + """ + The number of bytes needed to store this object in memory. + """ + return self._tensor.nbytes + + def isna(self) -> "TensorArray": + """ + A 1-D array indicating if each value is missing. + + Returns + ------- + na_values : Union[np.ndarray, ExtensionArray] + In most cases, this should return a NumPy ndarray. For + exceptional cases like ``SparseArray``, where returning + an ndarray would be expensive, an ExtensionArray may be + returned. + + Notes + ----- + If returning an ExtensionArray, then + + * ``na_values._is_boolean`` should be True + * `na_values` should implement :func:`ExtensionArray._reduce` + * ``na_values.any`` and ``na_values.all`` should be implemented + """ + if self._tensor.dtype.type is np.object_: + # Avoid comparing with __eq__ because the elements of the tensor + # may do something funny with that operation. + return np.array( + [self._tensor[i] is None for i in range(len(self))], dtype=bool + ) + elif self._tensor.dtype.type is np.str_: + return np.all(self._tensor == "", axis=tuple(range(1, self._tensor.ndim))) + else: + return np.all( + np.isnan(self._tensor), axis=tuple(range(1, self._tensor.ndim)) + ) + + def take( + self, indices: Sequence[int], allow_fill: bool = False, fill_value: Any = None + ) -> "TensorArray": + """ + Take elements from an array. + + Parameters + ---------- + indices : sequence of int + Indices to be taken. + allow_fill : bool, default False + How to handle negative values in `indices`. + + * False: negative values in `indices` indicate positional indices + from the right (the default). This is similar to + :func:`numpy.take`. + + * True: negative values in `indices` indicate + missing values. These values are set to `fill_value`. Any other + other negative values raise a ``ValueError``. + + fill_value : any, optional + Fill value to use for NA-indices when `allow_fill` is True. + This may be ``None``, in which case the default NA value for + the type, ``self.dtype.na_value``, is used. + + For many ExtensionArrays, there will be two representations of + `fill_value`: a user-facing "boxed" scalar, and a low-level + physical NA value. `fill_value` should be the user-facing version, + and the implementation should handle translating that to the + physical version for processing the take if necessary. + + Returns + ------- + ExtensionArray + + Raises + ------ + IndexError + When the indices are out of bounds for the array. + ValueError + When `indices` contains negative values other than ``-1`` + and `allow_fill` is True. + + See Also + -------- + numpy.take : Take elements from an array along an axis. + api.extensions.take : Take elements from an array. + + Notes + ----- + ExtensionArray.take is called by ``Series.__getitem__``, ``.loc``, + ``iloc``, when `indices` is a sequence of values. Additionally, + it's called by :meth:`Series.reindex`, or any other method + that causes realignment, with a `fill_value`. + + Examples + -------- + Here's an example implementation, which relies on casting the + extension array to object dtype. This uses the helper method + :func:`pandas.api.extensions.take`. + + .. code-block:: python + + def take(self, indices, allow_fill=False, fill_value=None): + from pandas.core.algorithms import take + + # If the ExtensionArray is backed by an ndarray, then + # just pass that here instead of coercing to object. + data = self.astype(object) + + if allow_fill and fill_value is None: + fill_value = self.dtype.na_value + + # fill value should always be translated from the scalar + # type for the array, to the physical storage type for + # the data, before passing to take. + + result = take(data, indices, fill_value=fill_value, + allow_fill=allow_fill) + return self._from_sequence(result, dtype=self.dtype) + """ + if allow_fill: + # With allow_fill being True, negative values in `indices` indicate + # missing values and should be set to `fill_value`. + indices = np.asarray(indices, dtype=np.intp) + validate_indices(indices, len(self._tensor)) + + # Check if there are missing indices to fill, otherwise we can + # delegate to NumPy ndarray .take(). + has_missing = np.any(indices < 0) + if has_missing: + if fill_value is None: + fill_value = np.nan + + # Create an array populated with fill value. + values = np.full((len(indices),) + self._tensor.shape[1:], fill_value) + + # Put tensors at the given positive indices into array. + is_nonneg = indices >= 0 + np.put(values, np.where(is_nonneg)[0], self._tensor[indices[is_nonneg]]) + + return TensorArray(values) + + # Delegate take to NumPy array. + values = self._tensor.take(indices, axis=0) + + return TensorArray(values) + + def copy(self) -> "TensorArray": + """ + Return a copy of the array. + + Returns + ------- + ExtensionArray + """ + # TODO(Clark): Copy cached properties. + return TensorArray(self._tensor.copy()) + + @classmethod + def _concat_same_type(cls, to_concat: Sequence["TensorArray"]) -> "TensorArray": + """ + Concatenate multiple array of this dtype. + + Parameters + ---------- + to_concat : sequence of this type + + Returns + ------- + ExtensionArray + """ + should_flatten = False + shape = None + for a in to_concat: + if shape is None: + shape = a.dtype.element_shape + if a.is_variable_shaped or a.dtype.element_shape != shape: + should_flatten = True + break + if should_flatten: + concated = TensorArray( + np.array([e for a in to_concat for e in a._tensor], dtype=object) + ) + else: + concated = TensorArray(np.concatenate([a._tensor for a in to_concat])) + return concated + + def __setitem__(self, key: Union[int, np.ndarray], value: Any) -> None: + """ + Set one or more values inplace. + + This method is not required to satisfy the pandas extension array + interface. + + Parameters + ---------- + key : int, ndarray, or slice + When called from, e.g. ``Series.__setitem__``, ``key`` will be + one of + + * scalar int + * ndarray of integers. + * boolean ndarray + * slice object + + value : ExtensionDtype.type, Sequence[ExtensionDtype.type], or object + value or values to be set of ``key``. + + Returns + ------- + None + """ + key = check_array_indexer(self, key) + if isinstance(value, TensorArrayElement) or np.isscalar(value): + value = np.asarray(value) + if isinstance(value, list): + value = [ + np.asarray(v) if isinstance(v, TensorArrayElement) else v for v in value + ] + if isinstance(value, ABCSeries) and isinstance(value.dtype, TensorDtype): + value = value.values + if value is None or isinstance(value, Sequence) and len(value) == 0: + self._tensor[key] = np.full_like(self._tensor[key], np.nan) + elif isinstance(key, (int, slice, np.ndarray)): + self._tensor[key] = value + else: + raise NotImplementedError( + f"__setitem__ with key type '{type(key)}' not implemented" + ) + + def __contains__(self, item) -> bool: + """ + Return for `item in self`. + """ + if isinstance(item, TensorArrayElement): + np_item = np.asarray(item) + if np_item.size == 1 and np.isnan(np_item).all(): + return self.isna().any() + return super().__contains__(item) + + def __repr__(self): + return self._tensor.__repr__() + + def __str__(self): + return self._tensor.__str__() + + def _values_for_factorize(self) -> Tuple[np.ndarray, Any]: + # TODO(Clark): return self._tensor, np.nan + raise NotImplementedError + + def _reduce(self, name: str, skipna: bool = True, **kwargs): + """ + Return a scalar result of performing the reduction operation. + + Parameters + ---------- + name : str + Name of the function, supported values are: + { any, all, min, max, sum, mean, median, prod, + std, var, sem, kurt, skew }. + skipna : bool, default True + If True, skip NaN values. + **kwargs + Additional keyword arguments passed to the reduction function. + Currently, `ddof` is the only supported kwarg. + + Returns + ------- + scalar + + Raises + ------ + TypeError : subclass does not define reductions + """ + supported_kwargs = ["ddof"] + reducer_kwargs = {} + for kw in supported_kwargs: + try: + reducer_kwargs[kw] = kwargs[kw] + except KeyError: + pass + try: + return TensorArrayElement( + self.SUPPORTED_REDUCERS[name](self._tensor, axis=0, **reducer_kwargs) + ) + except KeyError: + raise NotImplementedError(f"'{name}' aggregate not implemented.") from None + + def __array__(self, dtype: np.dtype = None, **kwargs) -> np.ndarray: + return np.asarray(self._tensor, dtype=dtype, **kwargs) + + def __array_ufunc__(self, ufunc: Callable, method: str, *inputs, **kwargs): + """ + Supports NumPy ufuncs without requiring sloppy coercion to an + ndarray. + """ + out = kwargs.get("out", ()) + for x in inputs + out: + if not isinstance(x, (TensorArray, np.ndarray, numbers.Number)): + return NotImplemented + + # Defer to the implementation of the ufunc on unwrapped values. + inputs = tuple(x._tensor if isinstance(x, TensorArray) else x for x in inputs) + if out: + kwargs["out"] = tuple( + x._tensor if isinstance(x, TensorArray) else x for x in out + ) + result = getattr(ufunc, method)(*inputs, **kwargs) + + if type(result) is tuple: + # Multiple return values. + return tuple(type(self)(x) for x in result) + elif method == "at": + # No return value. + return None + else: + # One return value. + return type(self)(result) + + def to_numpy( + self, + dtype: np.dtype = None, + copy: bool = False, + na_value: Any = pd.api.extensions.no_default, + ): + """ + Convert to a NumPy ndarray. + + .. versionadded:: 1.0.0 + + This is similar to :meth:`numpy.asarray`, but may provide additional + control over how the conversion is done. + + Parameters + ---------- + dtype : str or numpy.dtype, optional + The dtype to pass to :meth:`numpy.asarray`. + copy : bool, default False + Whether to ensure that the returned value is a not a view on + another array. Note that ``copy=False`` does not *ensure* that + ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that + a copy is made, even if not strictly necessary. + na_value : Any, optional + The value to use for missing values. The default value depends + on `dtype` and the type of the array. + + Returns + ------- + numpy.ndarray + """ + if dtype is not None: + dtype = pd.api.types.pandas_dtype(dtype) + if copy: + values = np.array(self._tensor, dtype=dtype, copy=True) + else: + values = self._tensor.astype(dtype) + elif copy: + values = self._tensor.copy() + else: + values = self._tensor + return values + + @property + def numpy_dtype(self): + """ + Get the dtype of the tensor. + :return: The numpy dtype of the backing ndarray + """ + return self._tensor.dtype + + @property + def numpy_ndim(self): + """ + Get the number of tensor dimensions. + :return: integer for the number of dimensions + """ + return self._tensor.ndim + + @property + def numpy_shape(self): + """ + Get the shape of the tensor. + :return: A tuple of integers for the numpy shape of the backing ndarray + """ + return self._tensor.shape + + @property + def numpy_size(self): + """ + Get the size of the tensor. + :return: integer for the number of elements in the tensor + """ + return self._tensor.size + + def astype(self, dtype, copy=True): + """ + Cast to a NumPy array with 'dtype'. + + Parameters + ---------- + dtype : str or dtype + Typecode or data-type to which the array is cast. + copy : bool, default True + Whether to copy the data, even if not necessary. If False, + a copy is made only if the old dtype does not match the + new dtype. + + Returns + ------- + array : ndarray + NumPy ndarray with 'dtype' for its dtype. + """ + dtype = pd.api.types.pandas_dtype(dtype) + + if isinstance(dtype, TensorDtype): + values = TensorArray(self._tensor.copy()) if copy else self + elif not ( + pd.api.types.is_object_dtype(dtype) and pd.api.types.is_string_dtype(dtype) + ): + values = np.array([str(t) for t in self._tensor]) + if isinstance(dtype, pd.StringDtype): + return dtype.construct_array_type()._from_sequence(values, copy=False) + else: + return values + elif pd.api.types.is_object_dtype(dtype): + # Interpret astype(object) as "cast to an array of numpy arrays" + values = np.empty(len(self), dtype=object) + for i in range(len(self)): + values[i] = self._tensor[i] + else: + values = self._tensor.astype(dtype, copy=copy) + return values + + def any(self, axis=None, out=None, keepdims=False): + """ + Test whether any array element along a given axis evaluates to True. + + See numpy.any() documentation for more information + https://numpy.org/doc/stable/reference/generated/numpy.any.html#numpy.any + + :param axis: Axis or axes along which a logical OR reduction is + performed. + :param out: Alternate output array in which to place the result. + :param keepdims: If this is set to True, the axes which are reduced are + left in the result as dimensions with size one. + :return: single boolean unless axis is not None else TensorArray + """ + result = self._tensor.any(axis=axis, out=out, keepdims=keepdims) + return result if axis is None else TensorArray(result) + + def all(self, axis=None, out=None, keepdims=False): + """ + Test whether all array elements along a given axis evaluate to True. + + :param axis: Axis or axes along which a logical AND reduction is + performed. + :param out: Alternate output array in which to place the result. + :param keepdims: If this is set to True, the axes which are reduced are + left in the result as dimensions with size one. + :return: single boolean unless axis is not None else TensorArray + """ + result = self._tensor.all(axis=axis, out=out, keepdims=keepdims) + return result if axis is None else TensorArray(result) + + def __arrow_array__(self, type=None): + """ + Convert this TensorArray to an ArrowTensorArray extension array. + + This and TensorDtype.__from_arrow__ make up the + Pandas extension type + array <--> Arrow extension type + array + interoperability protocol. See + https://pandas.pydata.org/pandas-docs/stable/development/extending.html#compatibility-with-apache-arrow + for more information. + """ + from ray.air.util.tensor_extensions.arrow import ( + ArrowTensorArray, + ArrowVariableShapedTensorArray, + ) + + if self.is_variable_shaped: + return ArrowVariableShapedTensorArray.from_numpy(self._tensor) + else: + return ArrowTensorArray.from_numpy(self._tensor) + + @property + def _is_boolean(self): + """ + Whether this extension array should be considered boolean. + + By default, ExtensionArrays are assumed to be non-numeric. + Setting this to True will affect the behavior of several places, + e.g. + + * is_bool + * boolean indexing + + Returns + ------- + bool + """ + # This is needed to support returning a TensorArray from .isnan(). + return self.dtype._is_boolean() + + +# Add operators from the mixin to the TensorArrayElement and TensorArray +# classes. +TensorArrayElement._add_arithmetic_ops() +TensorArrayElement._add_comparison_ops() +TensorArrayElement._add_logical_ops() +TensorArray._add_arithmetic_ops() +TensorArray._add_comparison_ops() +TensorArray._add_logical_ops() + + +@PublicAPI(stability="beta") +def column_needs_tensor_extension(s: pd.Series) -> bool: + """Return whether the provided pandas Series column needs a tensor extension + representation. This tensor extension representation provides more efficient slicing + and interop with ML frameworks. + + Args: + s: The pandas Series column that may need to be represented using the tensor + extension. + + Returns: + Whether the provided Series needs a tensor extension representation. + """ + # NOTE: This is an O(1) check. + return ( + s.dtype.type is np.object_ and not s.empty and isinstance(s.iloc[0], np.ndarray) + ) diff --git a/lib/python3.12/site-packages/ray/air/util/tensor_extensions/utils.py b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..003fd533112eba7c9ba02ac6490834555b53721e --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/tensor_extensions/utils.py @@ -0,0 +1,213 @@ +import warnings +from typing import ( + TYPE_CHECKING, + Any, + List, + Protocol, + Sequence, + Union, +) + +import numpy as np + +from ray.air.constants import TENSOR_COLUMN_NAME +from ray.util import PublicAPI +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + from pandas.core.dtypes.generic import ABCSeries + + +@DeveloperAPI(stability="beta") +class ArrayLike(Protocol): + """Protocol matching ndarray-like objects (like torch.Tensor)""" + + def __array__(self): + ... + + def __len__(self): + ... + + +@DeveloperAPI(stability="beta") +def is_ndarray_like(value: ArrayLike) -> bool: + """Checks whether objects are ndarray-like (for ex, torch.Tensor) + but NOT and ndarray itself""" + return ( + hasattr(value, "__array__") + and hasattr(value, "__len__") + and not isinstance(value, np.ndarray) + ) + + +def _is_arrow_array(value: ArrayLike) -> bool: + import pyarrow as pa + + return isinstance(value, (pa.Array, pa.ChunkedArray)) + + +def _should_convert_to_tensor( + column_values: Union[List[Any], np.ndarray, ArrayLike], column_name: str +) -> bool: + assert len(column_values) > 0 + + # We convert passed in column values into a tensor representation (involving + # Arrow/Pandas extension types) in either of the following cases: + return ( + # - Column name is `TENSOR_COLUMN_NAME` (for compatibility) + column_name == TENSOR_COLUMN_NAME + # - Provided column values are already represented by a Numpy tensor (ie + # ndarray with ndim > 1) + or _is_ndarray_tensor(column_values) + # - Provided collection is already implementing Ndarray protocol like + # `torch.Tensor`, `pd.Series`, etc (but *excluding* `pyarrow.Array`, + # `pyarrow.ChunkedArray`) + or _is_ndarray_like_not_pyarrow_array(column_values) + # - Provided column values is a list of a) ndarrays or b) ndarray-like objects + # (excluding Pyarrow arrays). This is done for compatibility with previous + # existing behavior where all column values were blindly converted to Numpy + # leading to list of ndarrays being converted a tensor): + or isinstance(column_values[0], np.ndarray) + or _is_ndarray_like_not_pyarrow_array(column_values[0]) + ) + + +def _is_ndarray_like_not_pyarrow_array(column_values): + return is_ndarray_like(column_values) and not _is_arrow_array(column_values) + + +def _is_ndarray_tensor(t: Any) -> bool: + """Return whether provided ndarray is a tensor (ie ndim > 1). + + NOTE: Tensor is defined as a NumPy array such that `len(arr.shape) > 1` + """ + + if not isinstance(t, np.ndarray): + return False + + # Case of uniform-shaped (ie non-ragged) tensor + if t.ndim > 1: + return True + + # Case of ragged tensor (as produced by `create_ragged_ndarray` utility) + elif t.dtype.type is np.object_ and len(t) > 0 and isinstance(t[0], np.ndarray): + return True + + return False + + +def _is_ndarray_variable_shaped_tensor(arr: np.ndarray) -> bool: + """Return whether the provided NumPy ndarray is comprised of variable-shaped + tensors. + + NOTE: This is an O(rows) check. + """ + if arr.dtype.type is not np.object_: + return False + if len(arr) == 0: + return False + if not isinstance(arr[0], np.ndarray): + return False + shape = arr[0].shape + for a in arr[1:]: + if not isinstance(a, np.ndarray): + return False + if a.shape != shape: + return True + return True + + +def _create_possibly_ragged_ndarray( + values: Union[np.ndarray, "ABCSeries", Sequence[Any]] +) -> np.ndarray: + """ + Create a possibly ragged ndarray. + Using the np.array() constructor will fail to construct a ragged ndarray that has a + uniform first dimension (e.g. uniform channel dimension in imagery). This function + catches this failure and tries a create-and-fill method to construct the ragged + ndarray. + """ + try: + with warnings.catch_warnings(): + # For NumPy < 1.24, constructing a ragged ndarray directly via + # `np.array(...)` without the `dtype=object` parameter will raise a + # VisibleDeprecationWarning which we suppress. + # More details: https://stackoverflow.com/q/63097829 + if np.lib.NumpyVersion(np.__version__) >= "2.0.0": + copy_if_needed = None + warning_type = np.exceptions.VisibleDeprecationWarning + else: + copy_if_needed = False + warning_type = np.VisibleDeprecationWarning + + warnings.simplefilter("ignore", category=warning_type) + arr = np.array(values, copy=copy_if_needed) + return arr + except ValueError as e: + # Constructing a ragged ndarray directly via `np.array(...)` + # without the `dtype=object` parameter will raise a ValueError. + # For NumPy < 1.24, the message is of the form: + # "could not broadcast input array from shape..." + # For NumPy >= 1.24, the message is of the form: + # "The requested array has an inhomogeneous shape..." + # More details: https://github.com/numpy/numpy/pull/22004 + error_str = str(e) + if ( + "could not broadcast input array from shape" in error_str + or "The requested array has an inhomogeneous shape" in error_str + ): + # Fall back to strictly creating a ragged ndarray. + return create_ragged_ndarray(values) + else: + # Re-raise original error if the failure wasn't a broadcast error. + raise e from None + + +@PublicAPI(stability="alpha") +def create_ragged_ndarray(values: Sequence[Any]) -> np.ndarray: + """Create an array that contains arrays of different length + + If you're working with variable-length arrays like images, use this function to + create ragged arrays instead of ``np.array``. + + .. note:: + ``np.array`` fails to construct ragged arrays if the input arrays have a uniform + first dimension: + + .. testsetup:: + + import numpy as np + from ray.air.util.tensor_extensions.utils import create_ragged_ndarray + + .. doctest:: + + >>> values = [np.zeros((3, 1)), np.zeros((3, 2))] + >>> np.array(values, dtype=object) + Traceback (most recent call last): + ... + ValueError: could not broadcast input array from shape (3,1) into shape (3,) + >>> create_ragged_ndarray(values) + array([array([[0.], + [0.], + [0.]]), array([[0., 0.], + [0., 0.], + [0., 0.]])], dtype=object) + + Or if you're creating a ragged array from a single array: + + .. doctest:: + + >>> values = [np.zeros((3, 1))] + >>> np.array(values, dtype=object)[0].dtype + dtype('O') + >>> create_ragged_ndarray(values)[0].dtype + dtype('float64') + + ``create_ragged_ndarray`` avoids the limitations of ``np.array`` by creating an + empty array and filling it with pointers to the variable-length arrays. + """ # noqa: E501 + # Create an empty object-dtyped 1D array. + arr = np.empty(len(values), dtype=object) + # Try to fill the 1D array of pointers with the (ragged) tensors. + arr[:] = list(values) + return arr diff --git a/lib/python3.12/site-packages/ray/air/util/transform_pyarrow.py b/lib/python3.12/site-packages/ray/air/util/transform_pyarrow.py new file mode 100644 index 0000000000000000000000000000000000000000..1617f04a6bfc0deac13d55b60b51b1c35caa6eb1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/air/util/transform_pyarrow.py @@ -0,0 +1,49 @@ +from ray.air.util.tensor_extensions.arrow import concat_tensor_arrays + +try: + import pyarrow +except ImportError: + pyarrow = None + + +def _is_pa_extension_type(pa_type: "pyarrow.lib.DataType") -> bool: + """Whether the provided Arrow Table column is an extension array, using an Arrow + extension type. + """ + return isinstance(pa_type, pyarrow.ExtensionType) + + +def _concatenate_extension_column( + ca: "pyarrow.ChunkedArray", ensure_copy: bool = False +) -> "pyarrow.Array": + """Concatenate chunks of an extension column into a contiguous array. + + This concatenation is required for creating copies and for .take() to work on + extension arrays. + See https://issues.apache.org/jira/browse/ARROW-16503. + + Args: + ca: The chunked array representing the extension column to be concatenated. + ensure_copy: Skip copying when ensure_copy is False and there is exactly 1 chunk. + """ + from ray.air.util.tensor_extensions.arrow import ( + get_arrow_extension_tensor_types, + ) + + if not _is_pa_extension_type(ca.type): + raise ValueError("Chunked array isn't an extension array: {ca}") + + tensor_extension_types = get_arrow_extension_tensor_types() + + if ca.num_chunks == 0: + # Create empty storage array. + storage = pyarrow.array([], type=ca.type.storage_type) + elif not ensure_copy and len(ca.chunks) == 1: + # Skip copying + return ca.chunks[0] + elif isinstance(ca.type, tensor_extension_types): + return concat_tensor_arrays(ca.chunks, ensure_copy) + else: + storage = pyarrow.concat_arrays([c.storage for c in ca.chunks]) + + return ca.type.wrap_array(storage) diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__init__.py b/lib/python3.12/site-packages/ray/cloudpickle/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8677c85b2c2666f45896d685c7477c287ab4f2fd --- /dev/null +++ b/lib/python3.12/site-packages/ray/cloudpickle/__init__.py @@ -0,0 +1,47 @@ +from __future__ import absolute_import + +import os +from pickle import PicklingError + +from ray.cloudpickle.cloudpickle import * # noqa +from ray.cloudpickle.cloudpickle_fast import CloudPickler, dumps, dump # noqa + + +# Conform to the convention used by python serialization libraries, which +# expose their Pickler subclass at top-level under the "Pickler" name. +Pickler = CloudPickler + +__version__ = '3.0.0' + + +def _warn_msg(obj, method, exc): + return ( + f"{method}({str(obj)}) failed." + "\nTo check which non-serializable variables are captured " + "in scope, re-run the ray script with 'RAY_PICKLE_VERBOSE_DEBUG=1'.") + + +def dump_debug(obj, *args, **kwargs): + try: + return dump(obj, *args, **kwargs) + except (TypeError, PicklingError) as exc: + if os.environ.get("RAY_PICKLE_VERBOSE_DEBUG"): + from ray.util.check_serialize import inspect_serializability + inspect_serializability(obj) + raise + else: + msg = _warn_msg(obj, "ray.cloudpickle.dump", exc) + raise type(exc)(msg) + + +def dumps_debug(obj, *args, **kwargs): + try: + return dumps(obj, *args, **kwargs) + except (TypeError, PicklingError) as exc: + if os.environ.get("RAY_PICKLE_VERBOSE_DEBUG"): + from ray.util.check_serialize import inspect_serializability + inspect_serializability(obj) + raise + else: + msg = _warn_msg(obj, "ray.cloudpickle.dumps", exc) + raise type(exc)(msg) diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c368805f7c8300e132e96c93a7f9fd9b8276e21 Binary files /dev/null and b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40989f974bf791ee020a8c257f20ace8b3ce0539 Binary files /dev/null and b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57b7b88cfecc22df2bd7e9600e4d2eb0c9c6a762 Binary files /dev/null and b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..966135842b29b25714e40480d8dd07c7a2082758 Binary files /dev/null and b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/py_pickle.cpython-312.pyc b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/py_pickle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee51ce56a8f3cd5a2885bcd7f30d88ffa15134ac Binary files /dev/null and b/lib/python3.12/site-packages/ray/cloudpickle/__pycache__/py_pickle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle.py b/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle.py new file mode 100644 index 0000000000000000000000000000000000000000..8f00b5f3f9769f8c549f25e8affe0198e5a65655 --- /dev/null +++ b/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle.py @@ -0,0 +1,1487 @@ +"""Pickler class to extend the standard pickle.Pickler functionality + +The main objective is to make it natural to perform distributed computing on +clusters (such as PySpark, Dask, Ray...) with interactively defined code +(functions, classes, ...) written in notebooks or console. + +In particular this pickler adds the following features: +- serialize interactively-defined or locally-defined functions, classes, + enums, typevars, lambdas and nested functions to compiled byte code; +- deal with some other non-serializable objects in an ad-hoc manner where + applicable. + +This pickler is therefore meant to be used for the communication between short +lived Python processes running the same version of Python and libraries. In +particular, it is not meant to be used for long term storage of Python objects. + +It does not include an unpickler, as standard Python unpickling suffices. + +This module was extracted from the `cloud` package, developed by `PiCloud, Inc. +`_. + +Copyright (c) 2012-now, CloudPickle developers and contributors. +Copyright (c) 2012, Regents of the University of California. +Copyright (c) 2009 `PiCloud, Inc. `_. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the University of California, Berkeley nor the + names of its contributors may be used to endorse or promote + products derived from this software without specific prior written + permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +""" + +import _collections_abc +from collections import ChainMap, OrderedDict +import abc +import builtins +import copyreg +import dataclasses +import dis +from enum import Enum +import io +import itertools +import logging +import opcode +import pickle +from pickle import _getattribute +import platform +import struct +import sys +import threading +import types +import typing +import uuid +import warnings +import weakref + +# The following import is required to be imported in the cloudpickle +# namespace to be able to load pickle files generated with older versions of +# cloudpickle. See: tests/test_backward_compat.py +from types import CellType # noqa: F401 + + +# cloudpickle is meant for inter process communication: we expect all +# communicating processes to run the same Python version hence we favor +# communication speed over compatibility: +DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL + +# Names of modules whose resources should be treated as dynamic. +_PICKLE_BY_VALUE_MODULES = set() + +# Track the provenance of reconstructed dynamic classes to make it possible to +# reconstruct instances from the matching singleton class definition when +# appropriate and preserve the usual "isinstance" semantics of Python objects. +_DYNAMIC_CLASS_TRACKER_BY_CLASS = weakref.WeakKeyDictionary() +_DYNAMIC_CLASS_TRACKER_BY_ID = weakref.WeakValueDictionary() +_DYNAMIC_CLASS_TRACKER_LOCK = threading.Lock() + +PYPY = platform.python_implementation() == "PyPy" + +builtin_code_type = None +if PYPY: + # builtin-code objects only exist in pypy + builtin_code_type = type(float.__new__.__code__) + +_extract_code_globals_cache = weakref.WeakKeyDictionary() + + +def _get_or_create_tracker_id(class_def): + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_tracker_id = _DYNAMIC_CLASS_TRACKER_BY_CLASS.get(class_def) + if class_tracker_id is None: + class_tracker_id = uuid.uuid4().hex + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + _DYNAMIC_CLASS_TRACKER_BY_ID[class_tracker_id] = class_def + return class_tracker_id + + +def _lookup_class_or_track(class_tracker_id, class_def): + if class_tracker_id is not None: + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_def = _DYNAMIC_CLASS_TRACKER_BY_ID.setdefault( + class_tracker_id, class_def + ) + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + return class_def + + +def register_pickle_by_value(module): + """Register a module to make it functions and classes picklable by value. + + By default, functions and classes that are attributes of an importable + module are to be pickled by reference, that is relying on re-importing + the attribute from the module at load time. + + If `register_pickle_by_value(module)` is called, all its functions and + classes are subsequently to be pickled by value, meaning that they can + be loaded in Python processes where the module is not importable. + + This is especially useful when developing a module in a distributed + execution environment: restarting the client Python process with the new + source code is enough: there is no need to re-install the new version + of the module on all the worker nodes nor to restart the workers. + + Note: this feature is considered experimental. See the cloudpickle + README.md file for more details and limitations. + """ + if not isinstance(module, types.ModuleType): + raise ValueError(f"Input should be a module object, got {str(module)} instead") + # In the future, cloudpickle may need a way to access any module registered + # for pickling by value in order to introspect relative imports inside + # functions pickled by value. (see + # https://github.com/cloudpipe/cloudpickle/pull/417#issuecomment-873684633). + # This access can be ensured by checking that module is present in + # sys.modules at registering time and assuming that it will still be in + # there when accessed during pickling. Another alternative would be to + # store a weakref to the module. Even though cloudpickle does not implement + # this introspection yet, in order to avoid a possible breaking change + # later, we still enforce the presence of module inside sys.modules. + if module.__name__ not in sys.modules: + raise ValueError( + f"{module} was not imported correctly, have you used an " + "`import` statement to access it?" + ) + _PICKLE_BY_VALUE_MODULES.add(module.__name__) + + +def unregister_pickle_by_value(module): + """Unregister that the input module should be pickled by value.""" + if not isinstance(module, types.ModuleType): + raise ValueError(f"Input should be a module object, got {str(module)} instead") + if module.__name__ not in _PICKLE_BY_VALUE_MODULES: + raise ValueError(f"{module} is not registered for pickle by value") + else: + _PICKLE_BY_VALUE_MODULES.remove(module.__name__) + + +def list_registry_pickle_by_value(): + return _PICKLE_BY_VALUE_MODULES.copy() + + +def _is_registered_pickle_by_value(module): + module_name = module.__name__ + if module_name in _PICKLE_BY_VALUE_MODULES: + return True + while True: + parent_name = module_name.rsplit(".", 1)[0] + if parent_name == module_name: + break + if parent_name in _PICKLE_BY_VALUE_MODULES: + return True + module_name = parent_name + return False + + +def _whichmodule(obj, name): + """Find the module an object belongs to. + + This function differs from ``pickle.whichmodule`` in two ways: + - it does not mangle the cases where obj's module is __main__ and obj was + not found in any module. + - Errors arising during module introspection are ignored, as those errors + are considered unwanted side effects. + """ + module_name = getattr(obj, "__module__", None) + + if module_name is not None: + return module_name + # Protect the iteration by using a copy of sys.modules against dynamic + # modules that trigger imports of other modules upon calls to getattr or + # other threads importing at the same time. + for module_name, module in sys.modules.copy().items(): + # Some modules such as coverage can inject non-module objects inside + # sys.modules + if ( + module_name == "__main__" + or module is None + or not isinstance(module, types.ModuleType) + ): + continue + try: + if _getattribute(module, name)[0] is obj: + return module_name + except Exception: + pass + return None + + +def _should_pickle_by_reference(obj, name=None): + """Test whether an function or a class should be pickled by reference + + Pickling by reference means by that the object (typically a function or a + class) is an attribute of a module that is assumed to be importable in the + target Python environment. Loading will therefore rely on importing the + module and then calling `getattr` on it to access the function or class. + + Pickling by reference is the only option to pickle functions and classes + in the standard library. In cloudpickle the alternative option is to + pickle by value (for instance for interactively or locally defined + functions and classes or for attributes of modules that have been + explicitly registered to be pickled by value. + """ + if isinstance(obj, types.FunctionType) or issubclass(type(obj), type): + module_and_name = _lookup_module_and_qualname(obj, name=name) + if module_and_name is None: + return False + module, name = module_and_name + return not _is_registered_pickle_by_value(module) + + elif isinstance(obj, types.ModuleType): + # We assume that sys.modules is primarily used as a cache mechanism for + # the Python import machinery. Checking if a module has been added in + # is sys.modules therefore a cheap and simple heuristic to tell us + # whether we can assume that a given module could be imported by name + # in another Python process. + if _is_registered_pickle_by_value(obj): + return False + return obj.__name__ in sys.modules + else: + raise TypeError( + "cannot check importability of {} instances".format(type(obj).__name__) + ) + + +def _lookup_module_and_qualname(obj, name=None): + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, "__name__", None) + + module_name = _whichmodule(obj, name) + + if module_name is None: + # In this case, obj.__module__ is None AND obj was not found in any + # imported module. obj is thus treated as dynamic. + return None + + if module_name == "__main__": + return None + + # Note: if module_name is in sys.modules, the corresponding module is + # assumed importable at unpickling time. See #357 + module = sys.modules.get(module_name, None) + if module is None: + # The main reason why obj's module would not be imported is that this + # module has been dynamically created, using for example + # types.ModuleType. The other possibility is that module was removed + # from sys.modules after obj was created/imported. But this case is not + # supported, as the standard pickle does not support it either. + return None + + try: + obj2, parent = _getattribute(module, name) + except AttributeError: + # obj was not found inside the module it points to + return None + if obj2 is not obj: + return None + return module, name + + +def _extract_code_globals(co): + """Find all globals names read or written to by codeblock co.""" + out_names = _extract_code_globals_cache.get(co) + if out_names is None: + # We use a dict with None values instead of a set to get a + # deterministic order and avoid introducing non-deterministic pickle + # bytes as a results. + out_names = {name: None for name in _walk_global_ops(co)} + + # Declaring a function inside another one using the "def ..." syntax + # generates a constant code object corresponding to the one of the + # nested function's As the nested function may itself need global + # variables, we need to introspect its code, extract its globals, (look + # for code object in it's co_consts attribute..) and add the result to + # code_globals + if co.co_consts: + for const in co.co_consts: + if isinstance(const, types.CodeType): + out_names.update(_extract_code_globals(const)) + + _extract_code_globals_cache[co] = out_names + + return out_names + + +def _find_imported_submodules(code, top_level_dependencies): + """Find currently imported submodules used by a function. + + Submodules used by a function need to be detected and referenced for the + function to work correctly at depickling time. Because submodules can be + referenced as attribute of their parent package (``package.submodule``), we + need a special introspection technique that does not rely on GLOBAL-related + opcodes to find references of them in a code object. + + Example: + ``` + import concurrent.futures + import cloudpickle + def func(): + x = concurrent.futures.ThreadPoolExecutor + if __name__ == '__main__': + cloudpickle.dumps(func) + ``` + The globals extracted by cloudpickle in the function's state include the + concurrent package, but not its submodule (here, concurrent.futures), which + is the module used by func. Find_imported_submodules will detect the usage + of concurrent.futures. Saving this module alongside with func will ensure + that calling func once depickled does not fail due to concurrent.futures + not being imported + """ + + subimports = [] + # check if any known dependency is an imported package + for x in top_level_dependencies: + if ( + isinstance(x, types.ModuleType) + and hasattr(x, "__package__") + and x.__package__ + ): + # check if the package has any currently loaded sub-imports + prefix = x.__name__ + "." + # A concurrent thread could mutate sys.modules, + # make sure we iterate over a copy to avoid exceptions + for name in list(sys.modules): + # Older versions of pytest will add a "None" module to + # sys.modules. + if name is not None and name.startswith(prefix): + # check whether the function can address the sub-module + tokens = set(name[len(prefix) :].split(".")) + if not tokens - set(code.co_names): + subimports.append(sys.modules[name]) + return subimports + + +# relevant opcodes +STORE_GLOBAL = opcode.opmap["STORE_GLOBAL"] +DELETE_GLOBAL = opcode.opmap["DELETE_GLOBAL"] +LOAD_GLOBAL = opcode.opmap["LOAD_GLOBAL"] +GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL) +HAVE_ARGUMENT = dis.HAVE_ARGUMENT +EXTENDED_ARG = dis.EXTENDED_ARG + + +_BUILTIN_TYPE_NAMES = {} +for k, v in types.__dict__.items(): + if type(v) is type: + _BUILTIN_TYPE_NAMES[v] = k + + +def _builtin_type(name): + if name == "ClassType": # pragma: no cover + # Backward compat to load pickle files generated with cloudpickle + # < 1.3 even if loading pickle files from older versions is not + # officially supported. + return type + return getattr(types, name) + + +def _walk_global_ops(code): + """Yield referenced name for global-referencing instructions in code.""" + for instr in dis.get_instructions(code): + op = instr.opcode + if op in GLOBAL_OPS: + yield instr.argval + + +def _extract_class_dict(cls): + """Retrieve a copy of the dict of a class without the inherited method.""" + clsdict = dict(cls.__dict__) # copy dict proxy to a dict + if len(cls.__bases__) == 1: + inherited_dict = cls.__bases__[0].__dict__ + else: + inherited_dict = {} + for base in reversed(cls.__bases__): + inherited_dict.update(base.__dict__) + to_remove = [] + for name, value in clsdict.items(): + try: + base_value = inherited_dict[name] + if value is base_value: + to_remove.append(name) + except KeyError: + pass + for name in to_remove: + clsdict.pop(name) + return clsdict + + +def is_tornado_coroutine(func): + """Return whether `func` is a Tornado coroutine function. + + Running coroutines are not supported. + """ + warnings.warn( + "is_tornado_coroutine is deprecated in cloudpickle 3.0 and will be " + "removed in cloudpickle 4.0. Use tornado.gen.is_coroutine_function " + "directly instead.", + category=DeprecationWarning, + ) + if "tornado.gen" not in sys.modules: + return False + gen = sys.modules["tornado.gen"] + if not hasattr(gen, "is_coroutine_function"): + # Tornado version is too old + return False + return gen.is_coroutine_function(func) + + +def subimport(name): + # We cannot do simply: `return __import__(name)`: Indeed, if ``name`` is + # the name of a submodule, __import__ will return the top-level root module + # of this submodule. For instance, __import__('os.path') returns the `os` + # module. + __import__(name) + return sys.modules[name] + + +def dynamic_subimport(name, vars): + mod = types.ModuleType(name) + mod.__dict__.update(vars) + mod.__dict__["__builtins__"] = builtins.__dict__ + return mod + + +def _get_cell_contents(cell): + try: + return cell.cell_contents + except ValueError: + # Handle empty cells explicitly with a sentinel value. + return _empty_cell_value + + +def instance(cls): + """Create a new instance of a class. + + Parameters + ---------- + cls : type + The class to create an instance of. + + Returns + ------- + instance : cls + A new instance of ``cls``. + """ + return cls() + + +@instance +class _empty_cell_value: + """Sentinel for empty closures.""" + + @classmethod + def __reduce__(cls): + return cls.__name__ + + +def _make_function(code, globals, name, argdefs, closure): + # Setting __builtins__ in globals is needed for nogil CPython. + globals["__builtins__"] = __builtins__ + return types.FunctionType(code, globals, name, argdefs, closure) + + +def _make_empty_cell(): + if False: + # trick the compiler into creating an empty cell in our lambda + cell = None + raise AssertionError("this route should not be executed") + + return (lambda: cell).__closure__[0] + + +def _make_cell(value=_empty_cell_value): + cell = _make_empty_cell() + if value is not _empty_cell_value: + cell.cell_contents = value + return cell + + +def _make_skeleton_class( + type_constructor, name, bases, type_kwargs, class_tracker_id, extra +): + """Build dynamic class with an empty __dict__ to be filled once memoized + + If class_tracker_id is not None, try to lookup an existing class definition + matching that id. If none is found, track a newly reconstructed class + definition under that id so that other instances stemming from the same + class id will also reuse this class definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + skeleton_class = types.new_class( + name, bases, {"metaclass": type_constructor}, lambda ns: ns.update(type_kwargs) + ) + return _lookup_class_or_track(class_tracker_id, skeleton_class) + + +def _make_skeleton_enum( + bases, name, qualname, members, module, class_tracker_id, extra +): + """Build dynamic enum with an empty __dict__ to be filled once memoized + + The creation of the enum class is inspired by the code of + EnumMeta._create_. + + If class_tracker_id is not None, try to lookup an existing enum definition + matching that id. If none is found, track a newly reconstructed enum + definition under that id so that other instances stemming from the same + class id will also reuse this enum definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + # enums always inherit from their base Enum class at the last position in + # the list of base classes: + enum_base = bases[-1] + metacls = enum_base.__class__ + classdict = metacls.__prepare__(name, bases) + + for member_name, member_value in members.items(): + classdict[member_name] = member_value + enum_class = metacls.__new__(metacls, name, bases, classdict) + enum_class.__module__ = module + enum_class.__qualname__ = qualname + + return _lookup_class_or_track(class_tracker_id, enum_class) + + +def _make_typevar(name, bound, constraints, covariant, contravariant, class_tracker_id): + tv = typing.TypeVar( + name, + *constraints, + bound=bound, + covariant=covariant, + contravariant=contravariant, + ) + return _lookup_class_or_track(class_tracker_id, tv) + + +def _decompose_typevar(obj): + return ( + obj.__name__, + obj.__bound__, + obj.__constraints__, + obj.__covariant__, + obj.__contravariant__, + _get_or_create_tracker_id(obj), + ) + + +def _typevar_reduce(obj): + # TypeVar instances require the module information hence why we + # are not using the _should_pickle_by_reference directly + module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__) + + if module_and_name is None: + return (_make_typevar, _decompose_typevar(obj)) + elif _is_registered_pickle_by_value(module_and_name[0]): + return (_make_typevar, _decompose_typevar(obj)) + + return (getattr, module_and_name) + + +def _get_bases(typ): + if "__orig_bases__" in getattr(typ, "__dict__", {}): + # For generic types (see PEP 560) + # Note that simply checking `hasattr(typ, '__orig_bases__')` is not + # correct. Subclasses of a fully-parameterized generic class does not + # have `__orig_bases__` defined, but `hasattr(typ, '__orig_bases__')` + # will return True because it's defined in the base class. + bases_attr = "__orig_bases__" + else: + # For regular class objects + bases_attr = "__bases__" + return getattr(typ, bases_attr) + + +def _make_dict_keys(obj, is_ordered=False): + if is_ordered: + return OrderedDict.fromkeys(obj).keys() + else: + return dict.fromkeys(obj).keys() + + +def _make_dict_values(obj, is_ordered=False): + if is_ordered: + return OrderedDict((i, _) for i, _ in enumerate(obj)).values() + else: + return {i: _ for i, _ in enumerate(obj)}.values() + + +def _make_dict_items(obj, is_ordered=False): + if is_ordered: + return OrderedDict(obj).items() + else: + return obj.items() + + +# COLLECTION OF OBJECTS __getnewargs__-LIKE METHODS +# ------------------------------------------------- + + +def _class_getnewargs(obj): + type_kwargs = {} + if "__module__" in obj.__dict__: + type_kwargs["__module__"] = obj.__module__ + + __dict__ = obj.__dict__.get("__dict__", None) + if isinstance(__dict__, property): + type_kwargs["__dict__"] = __dict__ + + return ( + type(obj), + obj.__name__, + _get_bases(obj), + type_kwargs, + _get_or_create_tracker_id(obj), + None, + ) + + +def _enum_getnewargs(obj): + members = {e.name: e.value for e in obj} + return ( + obj.__bases__, + obj.__name__, + obj.__qualname__, + members, + obj.__module__, + _get_or_create_tracker_id(obj), + None, + ) + + +# COLLECTION OF OBJECTS RECONSTRUCTORS +# ------------------------------------ +def _file_reconstructor(retval): + return retval + + +# COLLECTION OF OBJECTS STATE GETTERS +# ----------------------------------- + + +def _function_getstate(func): + # - Put func's dynamic attributes (stored in func.__dict__) in state. These + # attributes will be restored at unpickling time using + # f.__dict__.update(state) + # - Put func's members into slotstate. Such attributes will be restored at + # unpickling time by iterating over slotstate and calling setattr(func, + # slotname, slotvalue) + slotstate = { + "__name__": func.__name__, + "__qualname__": func.__qualname__, + "__annotations__": func.__annotations__, + "__kwdefaults__": func.__kwdefaults__, + "__defaults__": func.__defaults__, + "__module__": func.__module__, + "__doc__": func.__doc__, + "__closure__": func.__closure__, + } + + f_globals_ref = _extract_code_globals(func.__code__) + f_globals = {k: func.__globals__[k] for k in f_globals_ref if k in func.__globals__} + + if func.__closure__ is not None: + closure_values = list(map(_get_cell_contents, func.__closure__)) + else: + closure_values = () + + # Extract currently-imported submodules used by func. Storing these modules + # in a smoke _cloudpickle_subimports attribute of the object's state will + # trigger the side effect of importing these modules at unpickling time + # (which is necessary for func to work correctly once depickled) + slotstate["_cloudpickle_submodules"] = _find_imported_submodules( + func.__code__, itertools.chain(f_globals.values(), closure_values) + ) + slotstate["__globals__"] = f_globals + + state = func.__dict__ + return state, slotstate + + +def _class_getstate(obj): + clsdict = _extract_class_dict(obj) + clsdict.pop("__weakref__", None) + + if issubclass(type(obj), abc.ABCMeta): + # If obj is an instance of an ABCMeta subclass, don't pickle the + # cache/negative caches populated during isinstance/issubclass + # checks, but pickle the list of registered subclasses of obj. + clsdict.pop("_abc_cache", None) + clsdict.pop("_abc_negative_cache", None) + clsdict.pop("_abc_negative_cache_version", None) + registry = clsdict.pop("_abc_registry", None) + if registry is None: + # The abc caches and registered subclasses of a + # class are bundled into the single _abc_impl attribute + clsdict.pop("_abc_impl", None) + (registry, _, _, _) = abc._get_dump(obj) + + clsdict["_abc_impl"] = [subclass_weakref() for subclass_weakref in registry] + else: + # In the above if clause, registry is a set of weakrefs -- in + # this case, registry is a WeakSet + clsdict["_abc_impl"] = [type_ for type_ in registry] + + if "__slots__" in clsdict: + # pickle string length optimization: member descriptors of obj are + # created automatically from obj's __slots__ attribute, no need to + # save them in obj's state + if isinstance(obj.__slots__, str): + clsdict.pop(obj.__slots__) + else: + for k in obj.__slots__: + clsdict.pop(k, None) + + clsdict.pop("__dict__", None) # unpicklable property object + + return (clsdict, {}) + + +def _enum_getstate(obj): + clsdict, slotstate = _class_getstate(obj) + + members = {e.name: e.value for e in obj} + # Cleanup the clsdict that will be passed to _make_skeleton_enum: + # Those attributes are already handled by the metaclass. + for attrname in [ + "_generate_next_value_", + "_member_names_", + "_member_map_", + "_member_type_", + "_value2member_map_", + ]: + clsdict.pop(attrname, None) + for member in members: + clsdict.pop(member) + # Special handling of Enum subclasses + return clsdict, slotstate + + +# COLLECTIONS OF OBJECTS REDUCERS +# ------------------------------- +# A reducer is a function taking a single argument (obj), and that returns a +# tuple with all the necessary data to re-construct obj. Apart from a few +# exceptions (list, dict, bytes, int, etc.), a reducer is necessary to +# correctly pickle an object. +# While many built-in objects (Exceptions objects, instances of the "object" +# class, etc), are shipped with their own built-in reducer (invoked using +# obj.__reduce__), some do not. The following methods were created to "fill +# these holes". + + +def _code_reduce(obj): + """code object reducer.""" + # If you are not sure about the order of arguments, take a look at help + # of the specific type from types, for example: + # >>> from types import CodeType + # >>> help(CodeType) + if hasattr(obj, "co_exceptiontable"): + # Python 3.11 and later: there are some new attributes + # related to the enhanced exceptions. + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + obj.co_names, + obj.co_varnames, + obj.co_filename, + obj.co_name, + obj.co_qualname, + obj.co_firstlineno, + obj.co_linetable, + obj.co_exceptiontable, + obj.co_freevars, + obj.co_cellvars, + ) + elif hasattr(obj, "co_linetable"): + # Python 3.10 and later: obj.co_lnotab is deprecated and constructor + # expects obj.co_linetable instead. + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + obj.co_names, + obj.co_varnames, + obj.co_filename, + obj.co_name, + obj.co_firstlineno, + obj.co_linetable, + obj.co_freevars, + obj.co_cellvars, + ) + elif hasattr(obj, "co_nmeta"): # pragma: no cover + # "nogil" Python: modified attributes from 3.9 + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_framesize, + obj.co_ndefaultargs, + obj.co_nmeta, + obj.co_flags, + obj.co_code, + obj.co_consts, + obj.co_varnames, + obj.co_filename, + obj.co_name, + obj.co_firstlineno, + obj.co_lnotab, + obj.co_exc_handlers, + obj.co_jump_table, + obj.co_freevars, + obj.co_cellvars, + obj.co_free2reg, + obj.co_cell2reg, + ) + else: + # Backward compat for 3.8 and 3.9 + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + obj.co_names, + obj.co_varnames, + obj.co_filename, + obj.co_name, + obj.co_firstlineno, + obj.co_lnotab, + obj.co_freevars, + obj.co_cellvars, + ) + return types.CodeType, args + + +def _cell_reduce(obj): + """Cell (containing values of a function's free variables) reducer.""" + try: + obj.cell_contents + except ValueError: # cell is empty + return _make_empty_cell, () + else: + return _make_cell, (obj.cell_contents,) + + +def _classmethod_reduce(obj): + orig_func = obj.__func__ + return type(obj), (orig_func,) + + +def _file_reduce(obj): + """Save a file.""" + import io + + if not hasattr(obj, "name") or not hasattr(obj, "mode"): + raise pickle.PicklingError( + "Cannot pickle files that do not map to an actual file" + ) + if obj is sys.stdout: + return getattr, (sys, "stdout") + if obj is sys.stderr: + return getattr, (sys, "stderr") + if obj is sys.stdin: + raise pickle.PicklingError("Cannot pickle standard input") + if obj.closed: + raise pickle.PicklingError("Cannot pickle closed files") + if hasattr(obj, "isatty") and obj.isatty(): + raise pickle.PicklingError("Cannot pickle files that map to tty objects") + if "r" not in obj.mode and "+" not in obj.mode: + raise pickle.PicklingError( + "Cannot pickle files that are not opened for reading: %s" % obj.mode + ) + + name = obj.name + + retval = io.StringIO() + + try: + # Read the whole file + curloc = obj.tell() + obj.seek(0) + contents = obj.read() + obj.seek(curloc) + except OSError as e: + raise pickle.PicklingError( + "Cannot pickle file %s as it cannot be read" % name + ) from e + retval.write(contents) + retval.seek(curloc) + + retval.name = name + return _file_reconstructor, (retval,) + + +def _getset_descriptor_reduce(obj): + return getattr, (obj.__objclass__, obj.__name__) + + +def _mappingproxy_reduce(obj): + return types.MappingProxyType, (dict(obj),) + + +def _memoryview_reduce(obj): + return bytes, (obj.tobytes(),) + + +def _module_reduce(obj): + if _should_pickle_by_reference(obj): + return subimport, (obj.__name__,) + else: + # Some external libraries can populate the "__builtins__" entry of a + # module's `__dict__` with unpicklable objects (see #316). For that + # reason, we do not attempt to pickle the "__builtins__" entry, and + # restore a default value for it at unpickling time. + state = obj.__dict__.copy() + state.pop("__builtins__", None) + return dynamic_subimport, (obj.__name__, state) + + +def _method_reduce(obj): + return (types.MethodType, (obj.__func__, obj.__self__)) + + +def _logger_reduce(obj): + return logging.getLogger, (obj.name,) + + +def _root_logger_reduce(obj): + return logging.getLogger, () + + +def _property_reduce(obj): + return property, (obj.fget, obj.fset, obj.fdel, obj.__doc__) + + +def _weakset_reduce(obj): + return weakref.WeakSet, (list(obj),) + + +def _dynamic_class_reduce(obj): + """Save a class that can't be referenced as a module attribute. + + This method is used to serialize classes that are defined inside + functions, or that otherwise can't be serialized as attribute lookups + from importable modules. + """ + if Enum is not None and issubclass(obj, Enum): + return ( + _make_skeleton_enum, + _enum_getnewargs(obj), + _enum_getstate(obj), + None, + None, + _class_setstate, + ) + else: + return ( + _make_skeleton_class, + _class_getnewargs(obj), + _class_getstate(obj), + None, + None, + _class_setstate, + ) + + +def _class_reduce(obj): + """Select the reducer depending on the dynamic nature of the class obj.""" + if obj is type(None): # noqa + return type, (None,) + elif obj is type(Ellipsis): + return type, (Ellipsis,) + elif obj is type(NotImplemented): + return type, (NotImplemented,) + elif obj in _BUILTIN_TYPE_NAMES: + return _builtin_type, (_BUILTIN_TYPE_NAMES[obj],) + elif not _should_pickle_by_reference(obj): + return _dynamic_class_reduce(obj) + return NotImplemented + + +def _dict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj),) + + +def _dict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj),) + + +def _dict_items_reduce(obj): + return _make_dict_items, (dict(obj),) + + +def _odict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj), True) + + +def _odict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj), True) + + +def _odict_items_reduce(obj): + return _make_dict_items, (dict(obj), True) + + +def _dataclass_field_base_reduce(obj): + return _get_dataclass_field_type_sentinel, (obj.name,) + + +# COLLECTIONS OF OBJECTS STATE SETTERS +# ------------------------------------ +# state setters are called at unpickling time, once the object is created and +# it has to be updated to how it was at unpickling time. + + +def _function_setstate(obj, state): + """Update the state of a dynamic function. + + As __closure__ and __globals__ are readonly attributes of a function, we + cannot rely on the native setstate routine of pickle.load_build, that calls + setattr on items of the slotstate. Instead, we have to modify them inplace. + """ + state, slotstate = state + obj.__dict__.update(state) + + obj_globals = slotstate.pop("__globals__") + obj_closure = slotstate.pop("__closure__") + # _cloudpickle_subimports is a set of submodules that must be loaded for + # the pickled function to work correctly at unpickling time. Now that these + # submodules are depickled (hence imported), they can be removed from the + # object's state (the object state only served as a reference holder to + # these submodules) + slotstate.pop("_cloudpickle_submodules") + + obj.__globals__.update(obj_globals) + obj.__globals__["__builtins__"] = __builtins__ + + if obj_closure is not None: + for i, cell in enumerate(obj_closure): + try: + value = cell.cell_contents + except ValueError: # cell is empty + continue + obj.__closure__[i].cell_contents = value + + for k, v in slotstate.items(): + setattr(obj, k, v) + + +def _class_setstate(obj, state): + state, slotstate = state + registry = None + for attrname, attr in state.items(): + if attrname == "_abc_impl": + registry = attr + else: + setattr(obj, attrname, attr) + if registry is not None: + for subclass in registry: + obj.register(subclass) + + return obj + + +# COLLECTION OF DATACLASS UTILITIES +# --------------------------------- +# There are some internal sentinel values whose identity must be preserved when +# unpickling dataclass fields. Each sentinel value has a unique name that we can +# use to retrieve its identity at unpickling time. + + +_DATACLASSE_FIELD_TYPE_SENTINELS = { + dataclasses._FIELD.name: dataclasses._FIELD, + dataclasses._FIELD_CLASSVAR.name: dataclasses._FIELD_CLASSVAR, + dataclasses._FIELD_INITVAR.name: dataclasses._FIELD_INITVAR, +} + + +def _get_dataclass_field_type_sentinel(name): + return _DATACLASSE_FIELD_TYPE_SENTINELS[name] + + +class Pickler(pickle.Pickler): + # set of reducers defined and used by cloudpickle (private) + _dispatch_table = {} + _dispatch_table[classmethod] = _classmethod_reduce + _dispatch_table[io.TextIOWrapper] = _file_reduce + _dispatch_table[logging.Logger] = _logger_reduce + _dispatch_table[logging.RootLogger] = _root_logger_reduce + _dispatch_table[memoryview] = _memoryview_reduce + _dispatch_table[property] = _property_reduce + _dispatch_table[staticmethod] = _classmethod_reduce + _dispatch_table[CellType] = _cell_reduce + _dispatch_table[types.CodeType] = _code_reduce + _dispatch_table[types.GetSetDescriptorType] = _getset_descriptor_reduce + _dispatch_table[types.ModuleType] = _module_reduce + _dispatch_table[types.MethodType] = _method_reduce + _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce + _dispatch_table[weakref.WeakSet] = _weakset_reduce + _dispatch_table[typing.TypeVar] = _typevar_reduce + _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce + _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce + _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce + _dispatch_table[type(OrderedDict().keys())] = _odict_keys_reduce + _dispatch_table[type(OrderedDict().values())] = _odict_values_reduce + _dispatch_table[type(OrderedDict().items())] = _odict_items_reduce + _dispatch_table[abc.abstractmethod] = _classmethod_reduce + _dispatch_table[abc.abstractclassmethod] = _classmethod_reduce + _dispatch_table[abc.abstractstaticmethod] = _classmethod_reduce + _dispatch_table[abc.abstractproperty] = _property_reduce + _dispatch_table[dataclasses._FIELD_BASE] = _dataclass_field_base_reduce + + dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table) + + # function reducers are defined as instance methods of cloudpickle.Pickler + # objects, as they rely on a cloudpickle.Pickler attribute (globals_ref) + def _dynamic_function_reduce(self, func): + """Reduce a function that is not pickleable via attribute lookup.""" + newargs = self._function_getnewargs(func) + state = _function_getstate(func) + return (_make_function, newargs, state, None, None, _function_setstate) + + def _function_reduce(self, obj): + """Reducer for function objects. + + If obj is a top-level attribute of a file-backed module, this reducer + returns NotImplemented, making the cloudpickle.Pickler fall back to + traditional pickle.Pickler routines to save obj. Otherwise, it reduces + obj using a custom cloudpickle reducer designed specifically to handle + dynamic functions. + """ + if _should_pickle_by_reference(obj): + return NotImplemented + else: + return self._dynamic_function_reduce(obj) + + def _function_getnewargs(self, func): + code = func.__code__ + + # base_globals represents the future global namespace of func at + # unpickling time. Looking it up and storing it in + # cloudpickle.Pickler.globals_ref allow functions sharing the same + # globals at pickling time to also share them once unpickled, at one + # condition: since globals_ref is an attribute of a cloudpickle.Pickler + # instance, and that a new cloudpickle.Pickler is created each time + # cloudpickle.dump or cloudpickle.dumps is called, functions also need + # to be saved within the same invocation of + # cloudpickle.dump/cloudpickle.dumps (for example: + # cloudpickle.dumps([f1, f2])). There is no such limitation when using + # cloudpickle.Pickler.dump, as long as the multiple invocations are + # bound to the same cloudpickle.Pickler instance. + base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) + + if base_globals == {}: + # Add module attributes used to resolve relative imports + # instructions inside func. + for k in ["__package__", "__name__", "__path__", "__file__"]: + if k in func.__globals__: + base_globals[k] = func.__globals__[k] + + # Do not bind the free variables before the function is created to + # avoid infinite recursion. + if func.__closure__ is None: + closure = None + else: + closure = tuple(_make_empty_cell() for _ in range(len(code.co_freevars))) + + return code, base_globals, None, None, closure + + def dump(self, obj): + try: + return super().dump(obj) + except RuntimeError as e: + if len(e.args) > 0 and "recursion" in e.args[0]: + msg = "Could not pickle object as excessively deep recursion required." + raise pickle.PicklingError(msg) from e + else: + raise + + def __init__(self, file, protocol=None, buffer_callback=None): + if protocol is None: + protocol = DEFAULT_PROTOCOL + super().__init__(file, protocol=protocol, buffer_callback=buffer_callback) + # map functions __globals__ attribute ids, to ensure that functions + # sharing the same global namespace at pickling time also share + # their global namespace at unpickling time. + self.globals_ref = {} + self.proto = int(protocol) + + if not PYPY: + # pickle.Pickler is the C implementation of the CPython pickler and + # therefore we rely on reduce_override method to customize the pickler + # behavior. + + # `cloudpickle.Pickler.dispatch` is only left for backward + # compatibility - note that when using protocol 5, + # `cloudpickle.Pickler.dispatch` is not an extension of + # `pickle._Pickler.dispatch` dictionary, because `cloudpickle.Pickler` + # subclasses the C-implemented `pickle.Pickler`, which does not expose + # a `dispatch` attribute. Earlier versions of `cloudpickle.Pickler` + # used `cloudpickle.Pickler.dispatch` as a class-level attribute + # storing all reducers implemented by cloudpickle, but the attribute + # name was not a great choice given because it would collide with a + # similarly named attribute in the pure-Python `pickle._Pickler` + # implementation in the standard library. + dispatch = dispatch_table + + # Implementation of the reducer_override callback, in order to + # efficiently serialize dynamic functions and classes by subclassing + # the C-implemented `pickle.Pickler`. + # TODO: decorrelate reducer_override (which is tied to CPython's + # implementation - would it make sense to backport it to pypy? - and + # pickle's protocol 5 which is implementation agnostic. Currently, the + # availability of both notions coincide on CPython's pickle, but it may + # not be the case anymore when pypy implements protocol 5. + + def reducer_override(self, obj): + """Type-agnostic reducing callback for function and classes. + + For performance reasons, subclasses of the C `pickle.Pickler` class + cannot register custom reducers for functions and classes in the + dispatch_table attribute. Reducers for such types must instead + implemented via the special `reducer_override` method. + + Note that this method will be called for any object except a few + builtin-types (int, lists, dicts etc.), which differs from reducers + in the Pickler's dispatch_table, each of them being invoked for + objects of a specific type only. + + This property comes in handy for classes: although most classes are + instances of the ``type`` metaclass, some of them can be instances + of other custom metaclasses (such as enum.EnumMeta for example). In + particular, the metaclass will likely not be known in advance, and + thus cannot be special-cased using an entry in the dispatch_table. + reducer_override, among other things, allows us to register a + reducer that will be called for any class, independently of its + type. + + Notes: + + * reducer_override has the priority over dispatch_table-registered + reducers. + * reducer_override can be used to fix other limitations of + cloudpickle for other types that suffered from type-specific + reducers, such as Exceptions. See + https://github.com/cloudpipe/cloudpickle/issues/248 + """ + t = type(obj) + try: + is_anyclass = issubclass(t, type) + except TypeError: # t is not a class (old Boost; see SF #502085) + is_anyclass = False + + if is_anyclass: + return _class_reduce(obj) + elif isinstance(obj, types.FunctionType): + return self._function_reduce(obj) + else: + # fallback to save_global, including the Pickler's + # dispatch_table + return NotImplemented + + else: + # When reducer_override is not available, hack the pure-Python + # Pickler's types.FunctionType and type savers. Note: the type saver + # must override Pickler.save_global, because pickle.py contains a + # hard-coded call to save_global when pickling meta-classes. + dispatch = pickle.Pickler.dispatch.copy() + + def _save_reduce_pickle5( + self, + func, + args, + state=None, + listitems=None, + dictitems=None, + state_setter=None, + obj=None, + ): + save = self.save + write = self.write + self.save_reduce( + func, + args, + state=None, + listitems=listitems, + dictitems=dictitems, + obj=obj, + ) + # backport of the Python 3.8 state_setter pickle operations + save(state_setter) + save(obj) # simple BINGET opcode as obj is already memoized. + save(state) + write(pickle.TUPLE2) + # Trigger a state_setter(obj, state) function call. + write(pickle.REDUCE) + # The purpose of state_setter is to carry-out an + # inplace modification of obj. We do not care about what the + # method might return, so its output is eventually removed from + # the stack. + write(pickle.POP) + + def save_global(self, obj, name=None, pack=struct.pack): + """Main dispatch method. + + The name of this method is somewhat misleading: all types get + dispatched here. + """ + if obj is type(None): # noqa + return self.save_reduce(type, (None,), obj=obj) + elif obj is type(Ellipsis): + return self.save_reduce(type, (Ellipsis,), obj=obj) + elif obj is type(NotImplemented): + return self.save_reduce(type, (NotImplemented,), obj=obj) + elif obj in _BUILTIN_TYPE_NAMES: + return self.save_reduce( + _builtin_type, (_BUILTIN_TYPE_NAMES[obj],), obj=obj + ) + + if name is not None: + super().save_global(obj, name=name) + elif not _should_pickle_by_reference(obj, name=name): + self._save_reduce_pickle5(*_dynamic_class_reduce(obj), obj=obj) + else: + super().save_global(obj, name=name) + + dispatch[type] = save_global + + def save_function(self, obj, name=None): + """Registered with the dispatch to handle all function types. + + Determines what kind of function obj is (e.g. lambda, defined at + interactive prompt, etc) and handles the pickling appropriately. + """ + if _should_pickle_by_reference(obj, name=name): + return super().save_global(obj, name=name) + elif PYPY and isinstance(obj.__code__, builtin_code_type): + return self.save_pypy_builtin_func(obj) + else: + return self._save_reduce_pickle5( + *self._dynamic_function_reduce(obj), obj=obj + ) + + def save_pypy_builtin_func(self, obj): + """Save pypy equivalent of builtin functions. + + PyPy does not have the concept of builtin-functions. Instead, + builtin-functions are simple function instances, but with a + builtin-code attribute. + Most of the time, builtin functions should be pickled by attribute. + But PyPy has flaky support for __qualname__, so some builtin + functions such as float.__new__ will be classified as dynamic. For + this reason only, we created this special routine. Because + builtin-functions are not expected to have closure or globals, + there is no additional hack (compared the one already implemented + in pickle) to protect ourselves from reference cycles. A simple + (reconstructor, newargs, obj.__dict__) tuple is save_reduced. Note + also that PyPy improved their support for __qualname__ in v3.6, so + this routing should be removed when cloudpickle supports only PyPy + 3.6 and later. + """ + rv = ( + types.FunctionType, + (obj.__code__, {}, obj.__name__, obj.__defaults__, obj.__closure__), + obj.__dict__, + ) + self.save_reduce(*rv, obj=obj) + + dispatch[types.FunctionType] = save_function + + +# Shorthands similar to pickle.dump/pickle.dumps + + +def dump(obj, file, protocol=None, buffer_callback=None): + """Serialize obj as bytes streamed into file + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python (although this is not always + guaranteed to work because cloudpickle relies on some internal + implementation details that can change from one Python version to the + next). + """ + Pickler(file, protocol=protocol, buffer_callback=buffer_callback).dump(obj) + + +def dumps(obj, protocol=None, buffer_callback=None): + """Serialize obj as a string of bytes allocated in memory + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python (although this is not always + guaranteed to work because cloudpickle relies on some internal + implementation details that can change from one Python version to the + next). + """ + with io.BytesIO() as file: + cp = Pickler(file, protocol=protocol, buffer_callback=buffer_callback) + cp.dump(obj) + return file.getvalue() + + +# Include pickles unloading functions in this namespace for convenience. +load, loads = pickle.load, pickle.loads + +# Backward compat alias. +CloudPickler = Pickler \ No newline at end of file diff --git a/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle_fast.py b/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..52d6732e44ebcc0053b24969943f7c3b742268bb --- /dev/null +++ b/lib/python3.12/site-packages/ray/cloudpickle/cloudpickle_fast.py @@ -0,0 +1,13 @@ +"""Compatibility module. + +It can be necessary to load files generated by previous versions of cloudpickle +that rely on symbols being defined under the `cloudpickle.cloudpickle_fast` +namespace. + +See: tests/test_backward_compat.py +""" +from . import cloudpickle + + +def __getattr__(name): + return getattr(cloudpickle, name) diff --git a/lib/python3.12/site-packages/ray/cloudpickle/compat.py b/lib/python3.12/site-packages/ray/cloudpickle/compat.py new file mode 100644 index 0000000000000000000000000000000000000000..a2eb12e5080b1dd9e0bb7e16e079fbb262678215 --- /dev/null +++ b/lib/python3.12/site-packages/ray/cloudpickle/compat.py @@ -0,0 +1,19 @@ +import logging +import os + +logger = logging.getLogger(__name__) + +RAY_PICKLE_VERBOSE_DEBUG = os.environ.get("RAY_PICKLE_VERBOSE_DEBUG") +verbose_level = int(RAY_PICKLE_VERBOSE_DEBUG) if RAY_PICKLE_VERBOSE_DEBUG else 0 + +if verbose_level > 1: + logger.warning( + "Environmental variable RAY_PICKLE_VERBOSE_DEBUG is set to " + f"'{verbose_level}', this enabled python-based serialization backend " + f"instead of C-Pickle. Serialization would be very slow." + ) + from ray.cloudpickle import py_pickle as pickle + from ray.cloudpickle.py_pickle import Pickler +else: + import pickle # noqa: F401 + from _pickle import Pickler # noqa: F401 diff --git a/lib/python3.12/site-packages/ray/cloudpickle/py_pickle.py b/lib/python3.12/site-packages/ray/cloudpickle/py_pickle.py new file mode 100644 index 0000000000000000000000000000000000000000..b21840702a5167b38a0ff167eccacc8377453677 --- /dev/null +++ b/lib/python3.12/site-packages/ray/cloudpickle/py_pickle.py @@ -0,0 +1,30 @@ +from pickle import ( + _Pickler, + _Unpickler as Unpickler, + _loads as loads, + _load as load, + PickleError, + PicklingError, + UnpicklingError, + HIGHEST_PROTOCOL, +) + +__all__ = [ + "PickleError", + "PicklingError", + "UnpicklingError", + "Pickler", + "Unpickler", + "load", + "loads", + "HIGHEST_PROTOCOL", +] + + +class Pickler(_Pickler): + def __init__(self, file, protocol=None, *, fix_imports=True, buffer_callback=None): + super().__init__( + file, protocol, fix_imports=fix_imports, buffer_callback=buffer_callback + ) + # avoid being overrided by cloudpickle + self.dispatch = _Pickler.dispatch.copy() diff --git a/lib/python3.12/site-packages/ray/core/src/plasma/__init__.py b/lib/python3.12/site-packages/ray/core/src/plasma/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/core/src/plasma/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/core/src/plasma/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..756ee7633612f00c1477e6855982b7fd5ac5894b Binary files /dev/null and b/lib/python3.12/site-packages/ray/core/src/plasma/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/core/src/ray/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/core/src/ray/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76039a985f0d672ba8673ba8097287baf33b7497 Binary files /dev/null and b/lib/python3.12/site-packages/ray/core/src/ray/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/core/src/ray/raylet/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/core/src/ray/raylet/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0db9d4a259329c142ee0123bf6b8c28a939abc4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/core/src/ray/raylet/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/__init__.py b/lib/python3.12/site-packages/ray/train/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2d3dd6bd328a30cf5af1a489977b76468942c551 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/__init__.py @@ -0,0 +1,119 @@ +# Try import ray[train] core requirements (defined in setup.py) +# isort: off +try: + import fsspec # noqa: F401 + import pandas # noqa: F401 + import pyarrow # noqa: F401 + import requests # noqa: F401 +except ImportError as exc: + raise ImportError( + "Can't import ray.train as some dependencies are missing. " + 'Run `pip install "ray[train]"` to fix.' + ) from exc +# isort: on + + +from ray.air.config import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig +from ray.air.result import Result + +# Import this first so it can be used in other modules +from ray.train._checkpoint import Checkpoint +from ray.train._internal.data_config import DataConfig +from ray.train._internal.session import get_checkpoint, get_dataset_shard, report +from ray.train._internal.syncer import SyncConfig +from ray.train.backend import BackendConfig +from ray.train.base_trainer import TrainingFailedError +from ray.train.constants import TRAIN_DATASET_KEY +from ray.train.context import TrainContext, get_context +from ray.train.v2._internal.constants import is_v2_enabled + +if is_v2_enabled(): + try: + import pydantic # noqa: F401 + except (ImportError, ModuleNotFoundError) as exc: + raise ImportError( + "`ray.train.v2` requires the pydantic package, which is missing. " + "Run the following command to fix this: `pip install pydantic`" + ) from exc + from ray.train.v2.api.callback import UserCallback # noqa: F811 + from ray.train.v2.api.config import ( # noqa: F811 + CheckpointConfig, + FailureConfig, + RunConfig, + ScalingConfig, + ) + from ray.train.v2.api.context import TrainContext # noqa: F811 + from ray.train.v2.api.exceptions import ( # noqa: F811 + ControllerError, + TrainingFailedError, + WorkerGroupError, + ) + from ray.train.v2.api.report_config import CheckpointUploadMode # noqa: F811 + from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint # noqa: F811 + from ray.train.v2.api.result import Result # noqa: F811 + from ray.train.v2.api.train_fn_utils import ( # noqa: F811 + get_all_reported_checkpoints, + get_checkpoint, + get_context, + get_dataset_shard, + report, + ) + + +__all__ = [ + "get_checkpoint", + "get_context", + "get_dataset_shard", + "report", + "BackendConfig", + "Checkpoint", + "CheckpointConfig", + "DataConfig", + "FailureConfig", + "Result", + "RunConfig", + "ScalingConfig", + "SyncConfig", + "TrainContext", + "TrainingFailedError", + "TRAIN_DATASET_KEY", +] + +get_checkpoint.__module__ = "ray.train" +get_context.__module__ = "ray.train" +get_dataset_shard.__module__ = "ray.train" +report.__module__ = "ray.train" +BackendConfig.__module__ = "ray.train" +Checkpoint.__module__ = "ray.train" +CheckpointConfig.__module__ = "ray.train" +DataConfig.__module__ = "ray.train" +FailureConfig.__module__ = "ray.train" +Result.__module__ = "ray.train" +RunConfig.__module__ = "ray.train" +ScalingConfig.__module__ = "ray.train" +SyncConfig.__module__ = "ray.train" +TrainContext.__module__ = "ray.train" +TrainingFailedError.__module__ = "ray.train" + +# TODO: consider implementing these in v1 and raising ImportError instead. +if is_v2_enabled(): + __all__.extend( + [ + "CheckpointUploadMode", + "ControllerError", + "ReportedCheckpoint", + "UserCallback", + "WorkerGroupError", + "get_all_reported_checkpoints", + ] + ) + + CheckpointUploadMode.__module__ = "ray.train" + ControllerError.__module__ = "ray.train" + ReportedCheckpoint.__module__ = "ray.train" + UserCallback.__module__ = "ray.train" + WorkerGroupError.__module__ = "ray.train" + get_all_reported_checkpoints.__module__ = "ray.train" + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/_checkpoint.py b/lib/python3.12/site-packages/ray/train/_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..5ee65be4f20fe17126b315726bcfb825e9c89e45 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_checkpoint.py @@ -0,0 +1,424 @@ +import contextlib +import glob +import json +import logging +import os +import platform +import shutil +import tempfile +import traceback +import uuid +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Union + +import pyarrow.fs + +from ray.air._internal.filelock import TempFileLock +from ray.train._internal.storage import _download_from_fs_path, _exists_at_fs_path +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + +# The filename of the file that stores user metadata set on the checkpoint. +_METADATA_FILE_NAME = ".metadata.json" + +# The prefix of the temp checkpoint directory that `to_directory` downloads to +# on the local filesystem. +_CHECKPOINT_TEMP_DIR_PREFIX = "checkpoint_tmp_" + + +class _CheckpointMetaClass(type): + def __getattr__(self, item): + try: + return super().__getattribute__(item) + except AttributeError as exc: + if item in { + "from_dict", + "to_dict", + "from_bytes", + "to_bytes", + "get_internal_representation", + }: + raise _get_migration_error(item) from exc + elif item in { + "from_uri", + "to_uri", + "uri", + }: + raise _get_uri_error(item) from exc + elif item in {"get_preprocessor", "set_preprocessor"}: + raise _get_preprocessor_error(item) from exc + + raise exc + + +@PublicAPI(stability="beta") +class Checkpoint(metaclass=_CheckpointMetaClass): + """A reference to data persisted as a directory in local or remote storage. + + Access the checkpoint contents locally using ``checkpoint.to_directory()`` + or ``checkpoint.as_directory``. + + Attributes + ---------- + path: A path on the filesystem containing the checkpoint contents. + filesystem: PyArrow FileSystem that can be used to access data at the `path`. + + See Also + -------- + ray.train.report : Report a checkpoint during training (with Ray Train/Tune). + ray.train.get_checkpoint : Get the latest checkpoint during training + (for restoration). + + :ref:`train-checkpointing` + :ref:`persistent-storage-guide` + + Examples + -------- + + Creating a checkpoint using ``Checkpoint.from_directory``: + + >>> from ray.train import Checkpoint + >>> checkpoint = Checkpoint.from_directory("/tmp/example_checkpoint_dir") + >>> checkpoint.filesystem # doctest: +ELLIPSIS + >> checkpoint.path + '/tmp/example_checkpoint_dir' + + Creating a checkpoint from a remote URI: + + >>> checkpoint = Checkpoint("s3://bucket/path/to/checkpoint") + >>> checkpoint.filesystem # doctest: +ELLIPSIS + >> checkpoint.path + 'bucket/path/to/checkpoint' + + Creating a checkpoint with a custom filesystem: + + >>> checkpoint = Checkpoint( + ... path="bucket/path/to/checkpoint", + ... filesystem=pyarrow.fs.S3FileSystem(), + ... ) + >>> checkpoint.filesystem # doctest: +ELLIPSIS + >> checkpoint.path + 'bucket/path/to/checkpoint' + + Accessing a checkpoint's contents: + + >>> import os # doctest: +SKIP + >>> with checkpoint.as_directory() as local_checkpoint_dir: # doctest: +SKIP + ... print(os.listdir(local_checkpoint_dir)) # doctest: +SKIP + ['model.pt', 'optimizer.pt', 'misc.pt'] + """ + + def __init__( + self, + path: Union[str, os.PathLike], + filesystem: Optional["pyarrow.fs.FileSystem"] = None, + ): + """Construct a Checkpoint. + + Args: + path: A local path or remote URI containing the checkpoint data. + If a filesystem is provided, then this path must NOT be a URI. + It should be a path on the filesystem with the prefix already stripped. + filesystem: PyArrow FileSystem to use to access data at the path. + If not specified, this is inferred from the URI scheme. + """ + self.path = str(path) + self.filesystem = filesystem + + if path and not filesystem: + self.filesystem, self.path = pyarrow.fs.FileSystem.from_uri(path) + + # This random UUID is used to create a temporary directory name on the + # local filesystem, which will be used for downloading checkpoint data. + # This ensures that if multiple processes download the same checkpoint object + # only one process performs the actual download while the others wait. + # This prevents duplicated download efforts and data. + # NOTE: Calling `to_directory` from multiple `Checkpoint` objects + # that point to the same (fs, path) will still download the data multiple times. + # This only ensures a canonical temp directory name for a single `Checkpoint`. + self._uuid = uuid.uuid4() + + def __repr__(self): + return f"Checkpoint(filesystem={self.filesystem.type_name}, path={self.path})" + + def get_metadata(self) -> Dict[str, Any]: + """Return the metadata dict stored with the checkpoint. + + If no metadata is stored, an empty dict is returned. + """ + metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix() + if not _exists_at_fs_path(self.filesystem, metadata_path): + return {} + + with self.filesystem.open_input_file(metadata_path) as f: + return json.loads(f.readall().decode("utf-8")) + + def set_metadata(self, metadata: Dict[str, Any]) -> None: + """Set the metadata stored with this checkpoint. + + This will overwrite any existing metadata stored with this checkpoint. + """ + metadata_path = Path(self.path, _METADATA_FILE_NAME).as_posix() + with self.filesystem.open_output_stream(metadata_path) as f: + f.write(json.dumps(metadata).encode("utf-8")) + + def update_metadata(self, metadata: Dict[str, Any]) -> None: + """Update the metadata stored with this checkpoint. + + This will update any existing metadata stored with this checkpoint. + """ + existing_metadata = self.get_metadata() + existing_metadata.update(metadata) + self.set_metadata(existing_metadata) + + @classmethod + def from_directory(cls, path: Union[str, os.PathLike]) -> "Checkpoint": + """Create checkpoint object from a local directory. + + Args: + path: Local directory containing checkpoint data. + + Returns: + A ray.train.Checkpoint object. + """ + return cls(path, filesystem=pyarrow.fs.LocalFileSystem()) + + def to_directory(self, path: Optional[Union[str, os.PathLike]] = None) -> str: + """Write checkpoint data to a local directory. + + *If multiple processes on the same node call this method simultaneously,* + only a single process will perform the download, while the others + wait for the download to finish. Once the download finishes, all processes + receive the same local directory to read from. + + Args: + path: Target directory to download data to. If not specified, + this method will use a temporary directory. + + Returns: + str: Directory containing checkpoint data. + """ + user_provided_path = path is not None + local_path = ( + path if user_provided_path else self._get_temporary_checkpoint_dir() + ) + local_path = os.path.normpath(os.path.expanduser(str(local_path))) + os.makedirs(local_path, exist_ok=True) + + try: + # Timeout 0 means there will be only one attempt to acquire + # the file lock. If it cannot be acquired, throw a TimeoutError + with TempFileLock(local_path, timeout=0): + _download_from_fs_path( + fs=self.filesystem, fs_path=self.path, local_path=local_path + ) + except TimeoutError: + # if the directory is already locked, then wait but do not do anything. + with TempFileLock(local_path, timeout=-1): + pass + if not os.path.exists(local_path): + raise RuntimeError( + f"Checkpoint directory {local_path} does not exist, " + "even though it should have been created by " + "another process. Please raise an issue on GitHub: " + "https://github.com/ray-project/ray/issues" + ) + + return local_path + + @contextlib.contextmanager + def as_directory(self) -> Iterator[str]: + """Returns checkpoint contents in a local directory as a context. + + This function makes checkpoint data available as a directory while avoiding + unnecessary copies and left-over temporary data. + + *If the checkpoint points to a local directory*, this method just returns the + local directory path without making a copy, and nothing will be cleaned up + after exiting the context. + + *If the checkpoint points to a remote directory*, this method will download the + checkpoint to a local temporary directory and return the path + to the temporary directory. + + *If multiple processes on the same node call this method simultaneously,* + only a single process will perform the download, while the others + wait for the download to finish. Once the download finishes, all processes + receive the same local (temporary) directory to read from. + + Once all processes have finished working with the checkpoint, + the temporary directory is cleaned up. + + Users should treat the returned checkpoint directory as read-only and avoid + changing any data within it, as it may be deleted when exiting the context. + + Example: + + .. testcode:: + :hide: + + from pathlib import Path + import tempfile + + from ray.train import Checkpoint + + temp_dir = tempfile.mkdtemp() + (Path(temp_dir) / "example.txt").write_text("example checkpoint data") + checkpoint = Checkpoint.from_directory(temp_dir) + + .. testcode:: + + with checkpoint.as_directory() as checkpoint_dir: + # Do some read-only processing of files within checkpoint_dir + pass + + # At this point, if a temporary directory was created, it will have + # been deleted. + + """ + if isinstance(self.filesystem, pyarrow.fs.LocalFileSystem): + yield self.path + else: + del_lock_path = _get_del_lock_path(self._get_temporary_checkpoint_dir()) + open(del_lock_path, "a").close() + + temp_dir = self.to_directory() + try: + yield temp_dir + finally: + # Always cleanup the del lock after we're done with the directory. + # This avoids leaving a lock file behind in the case of an exception + # in the user code. + try: + os.remove(del_lock_path) + except Exception: + logger.warning( + f"Could not remove {del_lock_path} deletion file lock. " + f"Traceback:\n{traceback.format_exc()}" + ) + + # If there are no more lock files, that means there are no more + # readers of this directory, and we can safely delete it. + # In the edge case (process crash before del lock file is removed), + # we do not remove the directory at all. + # Since it's in /tmp, this is not that big of a deal. + # check if any lock files are remaining + remaining_locks = _list_existing_del_locks(temp_dir) + if not remaining_locks: + try: + # Timeout 0 means there will be only one attempt to acquire + # the file lock. If it cannot be acquired, a TimeoutError + # will be thrown. + with TempFileLock(temp_dir, timeout=0): + shutil.rmtree(temp_dir, ignore_errors=True) + except TimeoutError: + pass + + def _get_temporary_checkpoint_dir(self) -> str: + """Return the name for the temporary checkpoint dir that this checkpoint + will get downloaded to, if accessing via `to_directory` or `as_directory`. + """ + tmp_dir_path = tempfile.gettempdir() + checkpoint_dir_name = _CHECKPOINT_TEMP_DIR_PREFIX + self._uuid.hex + if platform.system() == "Windows": + # Max path on Windows is 260 chars, -1 for joining \ + # Also leave a little for the del lock + del_lock_name = _get_del_lock_path("") + checkpoint_dir_name = ( + _CHECKPOINT_TEMP_DIR_PREFIX + + self._uuid.hex[ + -259 + + len(_CHECKPOINT_TEMP_DIR_PREFIX) + + len(tmp_dir_path) + + len(del_lock_name) : + ] + ) + if not checkpoint_dir_name.startswith(_CHECKPOINT_TEMP_DIR_PREFIX): + raise RuntimeError( + "Couldn't create checkpoint directory due to length " + "constraints. Try specifying a shorter checkpoint path." + ) + return Path(tmp_dir_path, checkpoint_dir_name).as_posix() + + def __fspath__(self): + raise TypeError( + "You cannot use `Checkpoint` objects directly as paths. " + "Use `Checkpoint.to_directory()` or `Checkpoint.as_directory()` instead." + ) + + +def _get_del_lock_path(path: str, suffix: str = None) -> str: + """Get the path to the deletion lock file for a file/directory at `path`. + + Example: + + >>> _get_del_lock_path("/tmp/checkpoint_tmp") # doctest: +ELLIPSIS + '/tmp/checkpoint_tmp.del_lock_... + >>> _get_del_lock_path("/tmp/checkpoint_tmp/") # doctest: +ELLIPSIS + '/tmp/checkpoint_tmp.del_lock_... + >>> _get_del_lock_path("/tmp/checkpoint_tmp.txt") # doctest: +ELLIPSIS + '/tmp/checkpoint_tmp.txt.del_lock_... + + """ + suffix = suffix if suffix is not None else str(os.getpid()) + return f"{path.rstrip('/')}.del_lock_{suffix}" + + +def _list_existing_del_locks(path: str) -> List[str]: + """List all the deletion lock files for a file/directory at `path`. + + For example, if 2 checkpoints are being read via `as_directory`, + then this should return a list of 2 deletion lock files. + """ + return list(glob.glob(f"{_get_del_lock_path(path, suffix='*')}")) + + +def _get_migration_error(name: str): + return AttributeError( + f"The new `ray.train.Checkpoint` class does not support `{name}()`. " + f"Instead, only directories are supported.\n\n" + f"Example to store a dictionary in a checkpoint:\n\n" + f"import os, tempfile\n" + f"import ray.cloudpickle as pickle\n" + f"from ray import train\n" + f"from ray.train import Checkpoint\n\n" + f"with tempfile.TemporaryDirectory() as checkpoint_dir:\n" + f" with open(os.path.join(checkpoint_dir, 'data.pkl'), 'wb') as fp:\n" + f" pickle.dump({{'data': 'value'}}, fp)\n\n" + f" checkpoint = Checkpoint.from_directory(checkpoint_dir)\n" + f" train.report(..., checkpoint=checkpoint)\n\n" + f"Example to load a dictionary from a checkpoint:\n\n" + f"if train.get_checkpoint():\n" + f" with train.get_checkpoint().as_directory() as checkpoint_dir:\n" + f" with open(os.path.join(checkpoint_dir, 'data.pkl'), 'rb') as fp:\n" + f" data = pickle.load(fp)" + ) + + +def _get_uri_error(name: str): + return AttributeError( + f"The new `ray.train.Checkpoint` class does not support `{name}()`. " + f"To create a checkpoint from remote storage, create a `Checkpoint` using its " + f"constructor instead of `from_directory`.\n" + f'Example: `Checkpoint(path="s3://a/b/c")`.\n' + f"Then, access the contents of the checkpoint with " + f"`checkpoint.as_directory()` / `checkpoint.to_directory()`.\n" + f"To upload data to remote storage, use e.g. `pyarrow.fs.FileSystem` " + f"or your client of choice." + ) + + +def _get_preprocessor_error(name: str): + return AttributeError( + f"The new `ray.train.Checkpoint` class does not support `{name}()`. " + f"To include preprocessor information in checkpoints, " + f"pass it as metadata in the Trainer constructor.\n" + f"Example: `TorchTrainer(..., metadata={{...}})`.\n" + f"After training, access it in the checkpoint via `checkpoint.get_metadata()`. " + f"See here: https://docs.ray.io/en/master/train/user-guides/" + f"data-loading-preprocessing.html#preprocessing-structured-data" + ) diff --git a/lib/python3.12/site-packages/ray/train/_internal/__init__.py b/lib/python3.12/site-packages/ray/train/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..809b4315ef2f5ca975eb07428ac6229af4486479 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/accelerator.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/accelerator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d43e4450361cec54c0db3cce6fe141a9db84889 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/accelerator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/backend_executor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/backend_executor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc6fa9b7f8114f6498b78798c0121d13b9a6f7d5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/backend_executor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/base_worker_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/base_worker_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..046e225e47d9214fa9c2bb5548d42b8d20baa29d Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/base_worker_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/checkpoint_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/checkpoint_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67ca683b4b318ffff57805bd39aadbf81e88a7fc Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/checkpoint_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/data_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/data_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40af3e80b0c103a6716cf0d1b8aa3dd658075607 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/data_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/dl_predictor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/dl_predictor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7cce6539222ccbc9c35c7d8bc50d5bb56270b73c Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/dl_predictor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/framework_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/framework_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..deed1ccbbc57f3b4ee0a3a8a16e7b6f8f7e4d3cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/framework_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/session.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/session.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..625dd23c5970218596765bd2a65e5dbef8f0be7d Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/session.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/storage.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/storage.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..887f7abb30661d9a206c7083043e327332689fab Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/storage.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/syncer.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/syncer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e62abbf126ff668f554b2fe1d0c1345a9116ab69 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/syncer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b63167466124fb8f912b7b5cf5e55916ba4024a2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/__pycache__/worker_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/worker_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..376c7de66edb18970e18bfa0cc7994aa3219dc95 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/__pycache__/worker_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/accelerator.py b/lib/python3.12/site-packages/ray/train/_internal/accelerator.py new file mode 100644 index 0000000000000000000000000000000000000000..1c31098d86defc704ddef75172543813361119da --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/accelerator.py @@ -0,0 +1,5 @@ +import abc + + +class Accelerator(abc.ABC): + """A utility that contains methods to accelerate training.""" diff --git a/lib/python3.12/site-packages/ray/train/_internal/backend_executor.py b/lib/python3.12/site-packages/ray/train/_internal/backend_executor.py new file mode 100644 index 0000000000000000000000000000000000000000..0e4b7e86e06214dd6c06a15ea27de35c481e6490 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/backend_executor.py @@ -0,0 +1,833 @@ +import logging +import os +import time +from collections import defaultdict +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar + +import ray +import ray._private.ray_constants as ray_constants +from ray._private.accelerators.amd_gpu import HIP_VISIBLE_DEVICES_ENV_VAR +from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR +from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR +from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR +from ray._private.ray_constants import env_integer +from ray.data import Dataset +from ray.exceptions import RayActorError +from ray.train import Checkpoint, DataConfig +from ray.train._internal.session import ( + TrialInfo, + _TrainingResult, + get_session, + init_session, + shutdown_session, +) +from ray.train._internal.storage import StorageContext +from ray.train._internal.utils import check_for_failure +from ray.train._internal.worker_group import WorkerGroup +from ray.train.backend import BackendConfig +from ray.train.constants import ( + ENABLE_DETAILED_AUTOFILLED_METRICS_ENV, + ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV, + ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV, + ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, + ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV, + RAY_TRAIN_ENABLE_STATE_TRACKING, + TRAIN_ENABLE_WORKER_SPREAD_ENV, + TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV, +) +from ray.util.placement_group import get_current_placement_group, remove_placement_group + +T = TypeVar("T") + +logger = logging.getLogger(__name__) + + +class TrainBackendError(Exception): + """Errors with BackendExecutor that should not be exposed to user.""" + + +class TrainingWorkerError(Exception): + """Raised if a worker fails during training.""" + + +@dataclass +class ResourceConfig: + """ + Resource configuration for resource_ids to share between workers. + + Args: + resource_name: The name of the resource to configure + (Example: "neuron_cores" or "gpu"). + resource_enable_sharing_env_var: The environment variable to + check if the resource should be shared. + share_resource_ids_env_var: The environment variable to configure for + sharing the resources with other workers. + """ + + resource_name: str + resource_enable_sharing_env_var: str + share_resource_ids_env_var: str + + +class BackendExecutor: + """Main execution class for training backends. + + This class holds a worker group and is responsible for executing the + training function on the workers, and collecting intermediate results + from ``session.report()``. + + Args: + backend_config: The configurations for this + specific backend. + num_workers: Number of workers to use for training. + resources_per_worker (Optional[Dict[str, float]]): + Dictionary specifying the resources that will be + requested for each worker. Defaults to {"CPU": 1}. + max_retries: Number of retries when Ray actors fail. + Defaults to 3. Set to -1 for unlimited retries. + """ + + def __init__( + self, + backend_config: BackendConfig, + # TODO(xwjiang): Legacy Ray Train trainer clean up! + trial_info: Optional[TrialInfo] = None, + num_workers: int = 1, + resources_per_worker: Optional[Dict[str, float]] = None, + max_retries: int = 3, + ): + if resources_per_worker is None: + self._resources_per_worker = {"CPU": 1} + else: + self._resources_per_worker = resources_per_worker.copy() + + self._backend_config = backend_config + self._backend = backend_config.backend_cls() + self._num_workers = num_workers + self._max_failures = max_retries + if self._max_failures < 0: + self._max_failures = float("inf") + self._num_failures = 0 + self._last_failure = None + self._initialization_hook = None + self._placement_group = None + + self._trial_info = trial_info + + self.worker_group = InactiveWorkerGroup() + self.dataset_shards = None + + self._resource_configs = [ + ResourceConfig( + ray_constants.NEURON_CORES, + ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, + NEURON_RT_VISIBLE_CORES_ENV_VAR, + ), + ResourceConfig( + ray_constants.NPU, + ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV, + ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, + ), + # For AMD GPUs, they are using HIP_VISIBLE_DEVICES env var. + ResourceConfig( + ray_constants.GPU, + ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV, + HIP_VISIBLE_DEVICES_ENV_VAR, + ), + ] + + # Record the initialization time of BackendExecutor, which is + # after trainer.fit() and before worker_group executes the training function. + self._start_time_ms = int(time.time() * 1000) + + self.state_tracking_enabled = env_integer(RAY_TRAIN_ENABLE_STATE_TRACKING, 0) + + def start( + self, + initialization_hook: Optional[Callable[[], None]] = None, + train_cls: Optional[Type] = None, + train_cls_args: Optional[Tuple] = None, + train_cls_kwargs: Optional[Dict] = None, + ): + """Starts the worker group.""" + self._create_placement_group() + placement_group = self._placement_group or "default" + self.worker_group = WorkerGroup( + num_workers=self._num_workers, + resources_per_worker=self._resources_per_worker, + actor_cls=train_cls, + actor_cls_args=train_cls_args, + actor_cls_kwargs=train_cls_kwargs, + placement_group=placement_group, + ) + # Hack to avoid OOMs. + # This is just a temporary solution for Train loading entire checkpoints + # into memory by ensuring that the rank 0 worker is on the same node as + # trainable, thus allowing for lazy checkpoint transfer to be used. + # See https://github.com/ray-project/ray/issues/33073 + # for more context. + # TODO remove passing in trial_driver_ip. + + trial_driver_node_id = ( + self._trial_info.driver_node_id if self._trial_info else None + ) + self.worker_group.sort_workers_by_node_id_and_gpu_id(trial_driver_node_id) + + try: + if initialization_hook: + self._initialization_hook = initialization_hook + self.worker_group.execute(initialization_hook) + + # Always propagate the driver's DataContext to each worker in the group. + from ray.data import DataContext + + def _set_driver_dataset_context(ctx: DataContext): + DataContext._set_current(ctx) + + self.worker_group.execute( + _set_driver_dataset_context, + DataContext.get_current(), + ) + + share_cuda_visible_devices_enabled = bool( + env_integer( + ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV, + self._backend.share_cuda_visible_devices, + ) + ) + + if ( + self._resources_per_worker.get("GPU", 0) > 0 + and share_cuda_visible_devices_enabled + ): + self._share_cuda_visible_devices() + for resource_config in self._resource_configs: + if self._is_share_resources_enabled( + resource_config.resource_name, + resource_config.resource_enable_sharing_env_var, + ): + self._share_resource_ids( + resource_config.resource_name, + resource_config.share_resource_ids_env_var, + ) + self._backend.on_start(self.worker_group, self._backend_config) + except RayActorError as exc: + logger.exception(str(exc)) + logger.warning( + "Failure occurred during startup. Restarting all workers and " + "attempting to startup again." + ) + self._increment_failures() + self._restart() + + if self.state_tracking_enabled: + from ray.train._internal.state import TrainRunStateManager + from ray.train._internal.state.state_actor import get_state_actor + + self.state_manager = TrainRunStateManager(state_actor=get_state_actor()) + + def _create_placement_group(self): + """Creates a placement group if it does not exist. + + If a placement group is already detected (Tune) this will be a no-op. + + By default the placement group will be created with PACK strategy. + This is optimized for colocating GPUs on a minimal number of nodes. + This behavior can be overridden to use the SPREAD strategy by defining + ``TRAIN_ENABLE_WORKER_SPREAD_ENV`` + + If a placement group is created it will be stored as + self._placement_group. + """ + current_placement_group = get_current_placement_group() + worker = ray._private.worker.global_worker + should_capture_child_tasks_in_placement_group = ( + worker.should_capture_child_tasks_in_placement_group + ) + should_create_placement_group = ( + current_placement_group is None + or not should_capture_child_tasks_in_placement_group + ) + + if should_create_placement_group: + bundles = [ + self._resources_per_worker.copy() for _ in range(self._num_workers) + ] + + use_spread = bool(env_integer(TRAIN_ENABLE_WORKER_SPREAD_ENV, 0)) + strategy = "SPREAD" if use_spread else "PACK" + + placement_group = ray.util.placement_group(bundles, strategy=strategy) + logger.debug("Waiting for placement group to start.") + timeout = env_integer(TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV, 100) + ready, _ = ray.wait([placement_group.ready()], timeout=timeout) + if ready: + logger.debug("Placement group has started.") + else: + raise TimeoutError( + "Placement group creation timed out. Make sure your " + "cluster either has enough resources or use an " + "autoscaling cluster. If you are running on a cluster, " + "make sure you specify an address in `ray.init()`, for example, " + '`ray.init("auto")`. You can also increase the timeout by setting ' + "the TRAIN_PLACEMENT_GROUP_TIMEOUT_S environment variable. " + "Current resources available: {}, resources requested by the " + "placement group: {}".format( + ray.available_resources(), placement_group.bundle_specs + ) + ) + self._placement_group = placement_group + + def _share_cuda_visible_devices(self): + """Sets CUDA_VISIBLE_DEVICES on all workers. + + For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs + visible to all workers on that worker's node. + + This allows GPU workers on the same node to communicate with one + another. + + Example: + + Setup: + - Node1: + - Worker1: {0, 1} + - Worker2: {2, 3} + - Node2: + - Worker3: {0, 1} + + CUDA_VISIBLE_DEVICES: + - Worker1: "0,1,2,3" + - Worker2: "0,1,2,3" + - Worker3: "0,1" + + """ + self._share_resource_ids(ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR) + + def _share_resource_ids(self, resource: str, env_var: str): + """Sets the given env_var on all workers. + + For each worker, the cores/devices are visible to all the + workers on that worker's node.This allows workers on the + same node to communicate with one another. + + Example: + + Setup: + - Node1: + - Worker1: {0, 1} + - Worker2: {2, 3} + - Node2: + - Worker3: {0, 1} + + NEURON_RT_VISIBLE_CORES/TPU_VISIBLE_CHIPS/...: + - Worker1: "0,1,2,3" + - Worker2: "0,1,2,3" + - Worker2: "0,1" + + Args: + resource: The name of the resource/accelerator. + env_var: The name of the environment variable to set. + """ + node_ids_and_resource_ids = [ + ( + w.metadata.node_id, + w.metadata.resource_ids[resource], + ) + for w in self.worker_group.workers + ] + node_id_to_worker_id = defaultdict(set) + node_id_to_resource_ids = defaultdict(set) + + for worker_id, (node_id, resource_ids) in enumerate(node_ids_and_resource_ids): + node_id_to_worker_id[node_id].add(worker_id) + node_id_to_resource_ids[node_id].update(resource_ids) + + futures = [] + for node_id, resource_ids in node_id_to_resource_ids.items(): + resource_ids = sorted(resource_ids) + all_resource_ids = ",".join(resource_ids) + + def set_resource_ids(): + os.environ[env_var] = all_resource_ids + + for worker_id in node_id_to_worker_id[node_id]: + futures.append( + self.worker_group.execute_single_async(worker_id, set_resource_ids) + ) + ray.get(futures) + + def _is_share_resources_enabled(self, resource_name: str, enable_sharing_env: str): + """Whether to share resource IDs on all workers + based on enable_sharing_env. + + This will return true if resources are requested and greater than 0. + Also, user can disable by configuring the `enable_sharing_env` to "0". + + Args: + resource_name: The name of the resource/accelerator. + enable_sharing_env: The name of the environment variable + to check. + """ + has_resource_requested = self._resources_per_worker.get(resource_name, 0) > 0 + return has_resource_requested and ray_constants.env_bool( + enable_sharing_env, True + ) + + def _create_rank_world_size_mappings(self) -> List[Dict]: + """Create rank and world size mappings for workers. + There are three maps returned: + - local_rank_map, which maps from worker world_rank to local_rank. + - local_world_size_map, which maps from world_rank to local_world_size + - node_rank_map, which maps from world rank to node rank + + Example: + Worker 0: node 0 + Worker 1: node 0 + Worker 2: node 1 + Worker 3: node 0 + Worker 4: node 1 + + Workers 0, 1, 3 are on node 0. + Workers 2, 4 are on node 1. + + Expected local_rank_map: + { + 0 -> 0, + 1 -> 1, + 2 -> 0, + 3 -> 2, + 4 -> 1 + } + + Expected local_world_size_map: + { + 0 -> 3, + 1 -> 3, + 2 -> 2, + 3 -> 3, + 4 -> 2 + } + + Expected node_rank_map: + { + 0 -> 0, + 1 -> 0, + 2 -> 1, + 3 -> 0, + 4 -> 1 + } + + """ + local_rank_map = {} # map from world rank to local rank + local_world_size_map = {} # map from world rank to local world size + node_rank_map = {} # map from world rank to node rank + node_ids = {} # map from node id to node index + node_cnt = 0 # count the number of nodes + + node_id_dict = defaultdict( + int + ) # map from node id to the number of workers on it. + for world_rank in range(len(self.worker_group)): + worker = self.worker_group.workers[world_rank] + node_id = worker.metadata.node_id + local_rank_map[world_rank] = node_id_dict[node_id] + node_id_dict[node_id] += 1 + + if node_id not in node_ids: + node_ids[node_id] = node_cnt + node_cnt += 1 + node_rank_map[world_rank] = node_ids[node_id] + + for world_rank in range(len(self.worker_group)): + worker = self.worker_group.workers[world_rank] + node_id = worker.metadata.node_id + local_world_size_map[world_rank] = node_id_dict[node_id] + + workers_info = "\n".join( + [ + f"- (node_id={w.metadata.node_id}, ip={w.metadata.node_ip}, " + f"pid={w.metadata.pid}) world_rank={i}, " + f"local_rank={local_rank_map[i]}, node_rank={node_rank_map[i]}" + for i, w in enumerate(self.worker_group.workers) + ] + ) + logger.info(f"Started distributed worker processes: \n{workers_info}") + + return local_rank_map, local_world_size_map, node_rank_map + + def start_training( + self, + train_func: Callable[[], T], + datasets: Dict[str, Dataset], + metadata: Dict[str, Any], + data_config: DataConfig, + storage: StorageContext, + checkpoint: Optional[Checkpoint] = None, + ) -> None: + """Executes a training function on all workers in a separate thread. + + ``finish_training`` should be called after this. + + Args: + train_func: The training function to run on each worker. + datasets: The base datasets. + data_config: The config object for creating dataset shards for workers. + checkpoint: The checkpoint data that + should be loaded onto each worker and accessed by the + training function via ``session.get_checkpoint()``. If this + is ``None`` then no checkpoint will be loaded. + """ + use_detailed_autofilled_metrics = env_integer( + ENABLE_DETAILED_AUTOFILLED_METRICS_ENV, 0 + ) + + # First initialize the session. + def initialize_session( + train_func, + world_rank, + local_rank, + node_rank, + local_world_size, + world_size, + trial_info, + checkpoint, + dataset_shard, + metadata, + storage, + ): + try: + init_session( + training_func=train_func, + world_rank=world_rank, + local_rank=local_rank, + node_rank=node_rank, + local_world_size=local_world_size, + world_size=world_size, + trial_info=trial_info, + dataset_shard=dataset_shard, + metadata=metadata, + checkpoint=checkpoint, + detailed_autofilled_metrics=use_detailed_autofilled_metrics, + storage=storage, + ) + except ValueError: + raise TrainBackendError( + "Attempting to start training but a " + "previous training run is still ongoing. " + "You must call `finish_training` before " + "calling `start_training` again." + ) + + if self.dataset_shards is None: + actors = [worker.actor for worker in self.worker_group.workers] + node_ids = [worker.metadata.node_id for worker in self.worker_group.workers] + self.dataset_shards = data_config.configure( + datasets, + world_size=len(self.worker_group), + worker_handles=actors, + worker_node_ids=node_ids, + ) + + ( + local_rank_map, + local_world_size_map, + node_rank_map, + ) = self._create_rank_world_size_mappings() + + futures = [] + for index in range(len(self.worker_group)): + futures.append( + self.worker_group.execute_single_async( + index, + initialize_session, + world_rank=index, + local_rank=local_rank_map[index], + node_rank=node_rank_map[index], + local_world_size=local_world_size_map[index], + world_size=len(self.worker_group), + trial_info=self._trial_info, + train_func=train_func, + dataset_shard=self.dataset_shards[index], + metadata=metadata, + checkpoint=checkpoint, + storage=storage, + ) + ) + + self._backend.on_training_start(self.worker_group, self._backend_config) + + self.get_with_failure_handling(futures) + + # Register Train Run before training starts + if self.state_tracking_enabled: + from ray.train._internal.state.schema import RunStatusEnum + + core_context = ray.runtime_context.get_runtime_context() + + self.state_manager.register_train_run( + run_id=self._trial_info.run_id, + run_name=self._trial_info.experiment_name, + job_id=core_context.get_job_id(), + controller_actor_id=core_context.get_actor_id(), + datasets=datasets, + worker_group=self.worker_group, + start_time_ms=self._start_time_ms, + run_status=RunStatusEnum.RUNNING, + resources=[self._resources_per_worker] * self._num_workers, + ) + + # Run the training function asynchronously in its own thread. + def train_async(): + session = get_session() + session.start() + + self.worker_group.execute_async(train_async) + + def get_next_results(self) -> Optional[List[_TrainingResult]]: + """Fetches the next ``_TrainingResult`` from each worker. + + Each ``_TrainingResult`` is expected to correspond to the same step from + each worker (e.g. the same call to ``train.report()``). + + Returns: + A list of ``_TrainingResult``s or ``None`` if there are no more results + since the training function has exited on all workers. + """ + + def get_next(): + session = _get_session("get_next_results") + try: + result = session.get_next() + except RuntimeError: + # Training thread has not been started yet. + raise TrainBackendError( + "`get_next_results` has been called " + "before `start_training`. Please call " + "`start_training` before " + "`get_next_results`." + ) + + return result + + # Get next result from each worker. + futures = self.worker_group.execute_async(get_next) + results = self.get_with_failure_handling(futures) + + # Check if any worker returned None. + if any(r is None for r in results): + # Either all workers have results or none of them do. + if not all(r is None for r in results): + raise RuntimeError( + "Some workers returned results while " + "others didn't. Make sure that " + "`session.report()` are called the " + "same number of times on all workers." + ) + else: + # Return None if all results are None. + return None + + return results + + def pause_reporting(self): + """Disable workers from enqueuing results from ``session.report()``. + + Note: Already reported results may still be enqueued at this point, + and should be handled appropriately. + """ + + def pause_session_reporting(): + session = _get_session("pause_reporting") + return session.pause_reporting() + + futures = self.worker_group.execute_async(pause_session_reporting) + self.get_with_failure_handling(futures) + + def finish_training(self): + """Finish training and return final results. Propagate any exceptions. + + Blocks until training is finished on all workers. + + Assumes `start_training` has already been called. + + Returns: + A list of return values from calling ``train_func`` on each worker. + Each item corresponds to the return value from a single worker. + """ + + def end_training(): + session = _get_session("finish_training") + try: + # session.finish raises any Exceptions from training. + output = session.finish() + finally: + # Shutdown session even if session.finish() raises an + # Exception. + shutdown_session() + + return output + + futures = self.worker_group.execute_async(end_training) + results = self.get_with_failure_handling(futures) + return results + + def report_final_run_status( + self, + errored: bool = False, + failed_rank: Optional[int] = None, + stack_trace: Optional[str] = None, + ): + """Report the final train run status, error, and end time to TrainStateActor.""" + if self.state_tracking_enabled: + from ray.train._internal.state.schema import ( + MAX_ERROR_STACK_TRACE_LENGTH, + RunStatusEnum, + ) + + if errored: + run_status = RunStatusEnum.ERRORED + status_detail = "" + if failed_rank is not None: + status_detail += f"Rank {failed_rank} worker raised an error. \n" + if stack_trace is not None: + # Keep only the last part of the stack trace if it's too long. + status_detail += stack_trace[-MAX_ERROR_STACK_TRACE_LENGTH:] + else: + run_status = RunStatusEnum.FINISHED + status_detail = "" + + self.state_manager.end_train_run( + run_id=self._trial_info.run_id, + run_status=run_status, + status_detail=status_detail, + end_time_ms=int(time.time() * 1000), + ) + + def get_with_failure_handling(self, remote_values): + """Gets the remote values while handling for worker failures. + + This method should be called instead of ``ray.get()`` directly in + order to handle worker failures. + + If a worker failure is identified, backend specific failure handling + is executed and a ``TrainingWorkerError`` is raised. + + Args: + remote_values: List of object refs representing functions + that may fail in the middle of execution. For example, running + a Train training loop in multiple parallel actor calls. + Returns: + The resolved objects represented by the passed in ObjectRefs. + """ + success, exception = check_for_failure(remote_values) + if success: + return ray.get(remote_values) + else: + self._last_failure = exception + self._increment_failures() + logger.warning( + "Failure identified during training. Restarting all workers and " + "continuing training from latest checkpoint." + ) + self._restart() + raise TrainingWorkerError + + def shutdown(self, graceful_termination: bool = True): + """Shuts down the workers in the worker group. + + Args: + graceful_termination: If set to True, attempt to clean up the backend + before terminating the Ray actors. + + """ + if graceful_termination: + try: + self._backend.on_shutdown(self.worker_group, self._backend_config) + except RayActorError: + logger.warning( + "Graceful shutdown of backend failed. This is " + "expected if one of the workers has crashed." + ) + + if graceful_termination: + self.worker_group.shutdown() + else: + self.worker_group.shutdown(patience_s=0) + self.worker_group = InactiveWorkerGroup() + + if self._placement_group: + remove_placement_group(self._placement_group) + self._placement_group = None + + self.dataset_shards = None + + def is_started(self): + return not isinstance(self.worker_group, InactiveWorkerGroup) + + def _restart(self): + self.worker_group.shutdown() + if self._initialization_hook is not None: + initialization_hook = self._initialization_hook + else: + initialization_hook = None + if self._placement_group: + remove_placement_group(self._placement_group) + self._placement_group = None + self.start(initialization_hook=initialization_hook) + + def _increment_failures(self): + self._num_failures += 1 + if self._num_failures >= self._max_failures: + failure = self._last_failure + self._last_failure = None + if self._max_failures > 0: + exc = RuntimeError( + f"Training has failed after {self._num_failures} attempts." + ) + raise exc.with_traceback(None) from failure + else: + raise failure + + def get_worker_group(self): + return self.worker_group + + def _get_num_failures(self): + return self._num_failures + + +class InactiveWorkerGroupError(Exception): + """Raised when underlying worker group is inactive.""" + + +class InactiveWorkerGroup: + # TODO: fix inheritence. perhaps create WorkerGroupInterface. + + # Need to define getstate and setstate so that getattr does not screwup + # pickling. See https://stackoverflow.com/a/50888571/11249691 + def __getstate__(self): + return vars(self) + + def __setstate__(self, state): + vars(self).update(state) + + def __getattr__(self, name): + raise InactiveWorkerGroupError() + + def __len__(self): + raise InactiveWorkerGroupError() + + +def _get_session(method_name: str): + # Get the session for this worker. + session = get_session() + if not session: + # Session is not initialized yet. + raise TrainBackendError( + f"`{method_name}` has been called " + "before `start_training`. Please call " + "`start_training` before " + f"`{method_name}`." + ) + return session diff --git a/lib/python3.12/site-packages/ray/train/_internal/base_worker_group.py b/lib/python3.12/site-packages/ray/train/_internal/base_worker_group.py new file mode 100644 index 0000000000000000000000000000000000000000..36f40397ca9ea2ee6a2bb2b55e9733055c8f207c --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/base_worker_group.py @@ -0,0 +1,105 @@ +"""Abstract base class for WorkerGroup implementations. + +This module defines the common base class that both V1 and V2 WorkerGroup +implementations should inherit from to ensure backend compatibility. +""" + +import abc +from typing import Callable, List, TypeVar + +from ray.types import ObjectRef +from ray.util.annotations import DeveloperAPI + +T = TypeVar("T") + + +@DeveloperAPI +class BaseWorkerGroup(abc.ABC): + """Abstract base class for WorkerGroup implementations. + + This base class defines the minimal set of methods that backend classes + expect from WorkerGroup implementations. Both V1 and V2 WorkerGroup + classes should inherit from this base class to ensure compatibility with + all backend configurations. + + The interface focuses on the core operations that backends need: + - Executing functions on workers + - Getting worker count and resource allocation + """ + + @abc.abstractmethod + def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]: + """Execute a function on all workers synchronously. + + Args: + func: The function to execute on each worker. + *args: Positional arguments to pass to the function. + **kwargs: Keyword arguments to pass to the function. + + Returns: + A list of results from each worker, in worker rank order. + """ + pass + + @abc.abstractmethod + def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]: + """Execute a function on all workers asynchronously. + + Args: + func: The function to execute on each worker. + *args: Positional arguments to pass to the function. + **kwargs: Keyword arguments to pass to the function. + + Returns: + A list of ObjectRef results from each worker, in worker rank order. + """ + pass + + @abc.abstractmethod + def execute_single( + self, worker_index: int, func: Callable[..., T], *args, **kwargs + ) -> T: + """Execute a function on a single worker synchronously. + + Args: + worker_index: The index of the worker to execute on. + func: The function to execute. + *args: Positional arguments to pass to the function. + **kwargs: Keyword arguments to pass to the function. + + Returns: + The result from the specified worker. + """ + pass + + @abc.abstractmethod + def execute_single_async( + self, worker_index: int, func: Callable[..., T], *args, **kwargs + ) -> ObjectRef: + """Execute a function on a single worker asynchronously. + + Args: + worker_index: The index of the worker to execute on. + func: The function to execute. + *args: Positional arguments to pass to the function. + **kwargs: Keyword arguments to pass to the function. + + Returns: + An ObjectRef to the result from the specified worker. + """ + pass + + @abc.abstractmethod + def __len__(self) -> int: + """Return the number of workers in the group.""" + pass + + @abc.abstractmethod + def get_resources_per_worker(self) -> dict: + """Get the resources allocated per worker. + + Returns: + A dictionary mapping resource names to quantities per worker. + Common keys include "CPU", "GPU", "memory". + """ + pass diff --git a/lib/python3.12/site-packages/ray/train/_internal/checkpoint_manager.py b/lib/python3.12/site-packages/ray/train/_internal/checkpoint_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..705d34843accc845111afde79912d8ee0fc7e029 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/checkpoint_manager.py @@ -0,0 +1,202 @@ +import logging +import numbers +from typing import Any, Callable, Dict, List, Optional, Tuple + +from ray._private import ray_constants +from ray._private.dict import flatten_dict +from ray.air._internal.util import is_nan +from ray.air.config import MAX +from ray.train import Checkpoint, CheckpointConfig +from ray.train._internal.session import _TrainingResult +from ray.train._internal.storage import _delete_fs_path +from ray.train.constants import TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE + +logger = logging.getLogger(__name__) + + +def _insert_into_sorted_list( + list: List[_TrainingResult], + item: _TrainingResult, + key: Callable[[_TrainingResult], Any], + checkpoint_to_report_index: Optional[Dict[Checkpoint, int]] = None, +): + """Insert an item into a sorted list with a custom key function. + + Args: + list: The list to insert the item into. + item: The item to insert. + key: The key function to use to sort the list. + checkpoint_to_report_index: A dictionary mapping checkpoints to report indices. + Used to break ties when scores are equal. + """ + checkpoint_to_report_index = checkpoint_to_report_index or {} + # TODO: optimize this with sortedlist, batching, etc + i = 0 + while i < len(list): + # When scores are equal, later checkpoints are later in the list. + list_item_key, item_key = key(list[i]), key(item) + if list_item_key > item_key or ( + list_item_key == item_key + and checkpoint_to_report_index.get(list[i].checkpoint, 0) + > checkpoint_to_report_index.get(item.checkpoint, 0) + ): + break + i += 1 + list.insert(i, item) + + +class _CheckpointManager: + """Checkpoint manager that handles checkpoint book-keeping for a trial. + + The main purpose of this abstraction is to keep the top K checkpoints based on + recency/a user-provided metric. + + NOTE: This class interacts with `_TrainingResult` objects, which are + (checkpoint, metrics) pairs. This is to order checkpoints by metrics. + + Args: + checkpoint_config: Defines how many and which checkpoints to keep. + """ + + def __init__(self, checkpoint_config: Optional[CheckpointConfig]): + self._checkpoint_config = checkpoint_config or CheckpointConfig() + + # List of checkpoints ordered by ascending score. + self._checkpoint_results: List[_TrainingResult] = [] + + # The latest registered checkpoint. + # This should never be immediately deleted upon registration, + # even if it's not in the top K checkpoints, based on score. + self._latest_checkpoint_result: Optional[_TrainingResult] = None + + if ( + self._checkpoint_config.num_to_keep is not None + and self._checkpoint_config.num_to_keep <= 0 + ): + raise ValueError( + f"`num_to_keep` must >= 1, got: " + f"{self._checkpoint_config.num_to_keep}" + ) + + @property + def checkpoint_config(self): + return self._checkpoint_config + + def register_checkpoint(self, checkpoint_result: _TrainingResult): + """Register new checkpoint and add to bookkeeping. + + This method will register a new checkpoint and add it to the internal + bookkeeping logic. This means the checkpoint manager will decide if + this checkpoint should be kept, and if older or worse performing + checkpoints should be deleted. + + Args: + checkpoint: Tracked checkpoint object to add to bookkeeping. + """ + self._latest_checkpoint_result = checkpoint_result + + score_attr = self._checkpoint_config.checkpoint_score_attribute + if ray_constants.env_bool(TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, False): + metrics = ( + {score_attr: checkpoint_result.metrics[score_attr]} + if score_attr in checkpoint_result.metrics + else {} + ) + checkpoint_result = _TrainingResult( + checkpoint=checkpoint_result.checkpoint, + metrics=metrics, + ) + + if score_attr is not None and score_attr in checkpoint_result.metrics: + # If we're ordering by a score, insert the checkpoint + # so that the list remains sorted. + _insert_into_sorted_list( + self._checkpoint_results, + checkpoint_result, + key=self._get_checkpoint_score, + ) + else: + # If no metric is provided, just append (ordering by time of registration). + self._checkpoint_results.append(checkpoint_result) + + if self._checkpoint_config.num_to_keep is not None: + # Delete the bottom (N - K) checkpoints + worst_results = set( + self._checkpoint_results[: -self._checkpoint_config.num_to_keep] + ) + # Except for the latest checkpoint. + results_to_delete = worst_results - {self._latest_checkpoint_result} + + # Update internal state before actually deleting them. + self._checkpoint_results = [ + checkpoint_result + for checkpoint_result in self._checkpoint_results + if checkpoint_result not in results_to_delete + ] + + for checkpoint_result in results_to_delete: + checkpoint = checkpoint_result.checkpoint + logger.debug("Deleting checkpoint: ", checkpoint) + _delete_fs_path(fs=checkpoint.filesystem, fs_path=checkpoint.path) + + def _get_checkpoint_score( + self, checkpoint: _TrainingResult + ) -> Tuple[bool, numbers.Number]: + """Get the score for a checkpoint, according to checkpoint config. + + If `mode="min"`, the metric is negated so that the lowest score is + treated as the best. + + Returns: + Tuple: A tuple of (not_is_nan: bool, score: numbers.Number). + This score orders: nan values < float("-inf") < valid numeric metrics + """ + checkpoint_score_attribute = self._checkpoint_config.checkpoint_score_attribute + if checkpoint_score_attribute: + flat_metrics = flatten_dict(checkpoint.metrics) + try: + checkpoint_result = flat_metrics[checkpoint_score_attribute] + except KeyError: + valid_keys = list(flat_metrics.keys()) + logger.error( + f"Result dict has no key: {checkpoint_score_attribute}. " + f"checkpoint_score_attr must be set to a key in the " + f"result dict. Valid keys are: {valid_keys}" + ) + checkpoint_result = float("-inf") + else: + checkpoint_result = float("-inf") + + checkpoint_score_order = self._checkpoint_config.checkpoint_score_order + order_factor = 1.0 if checkpoint_score_order == MAX else -1.0 + + checkpoint_score = order_factor * checkpoint_result + + if not isinstance(checkpoint_score, numbers.Number): + raise ValueError( + f"Unable to persist checkpoint for " + f"checkpoint_score_attribute: " + f"{checkpoint_score_attribute} with value " + f"{checkpoint_score}. " + f"This attribute must be numerical." + ) + + return ( + (not is_nan(checkpoint_score), checkpoint_score) + if not is_nan(checkpoint_score) + else (False, float("-inf")) + ) + + @property + def best_checkpoint_result(self) -> Optional[_TrainingResult]: + return self._checkpoint_results[-1] if self._checkpoint_results else None + + @property + def latest_checkpoint_result(self) -> Optional[_TrainingResult]: + return self._latest_checkpoint_result + + @property + def best_checkpoint_results(self) -> List[_TrainingResult]: + if self._checkpoint_config.num_to_keep is None: + return self._checkpoint_results + return self._checkpoint_results[-self._checkpoint_config.num_to_keep :] diff --git a/lib/python3.12/site-packages/ray/train/_internal/data_config.py b/lib/python3.12/site-packages/ray/train/_internal/data_config.py new file mode 100644 index 0000000000000000000000000000000000000000..da52a4b8fbeca0386123d44c537d3b579cadc799 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/data_config.py @@ -0,0 +1,146 @@ +import copy +from typing import Dict, List, Literal, Optional, Union + +import ray +from ray.actor import ActorHandle +from ray.data import DataIterator, Dataset, ExecutionOptions, NodeIdStr +from ray.data._internal.execution.interfaces.execution_options import ExecutionResources +from ray.util.annotations import DeveloperAPI, PublicAPI + + +@PublicAPI(stability="stable") +class DataConfig: + """Class responsible for configuring Train dataset preprocessing. + + For advanced use cases, this class can be subclassed and the `configure()` method + overriden for custom data preprocessing. + """ + + def __init__( + self, + datasets_to_split: Union[Literal["all"], List[str]] = "all", + execution_options: Optional[ExecutionOptions] = None, + enable_shard_locality: bool = True, + ): + """Construct a DataConfig. + + Args: + datasets_to_split: Specifies which datasets should be split among workers. + Can be set to "all" or a list of dataset names. Defaults to "all", + i.e. split all datasets. + execution_options: The execution options to pass to Ray Data. By default, + the options will be optimized for data ingest. When overriding this, + base your options off of `DataConfig.default_ingest_options()`. + enable_shard_locality: If true, when sharding the datasets across Train + workers, locality will be considered to minimize cross-node data transfer. + This is on by default. + """ + if isinstance(datasets_to_split, list) or datasets_to_split == "all": + self._datasets_to_split = datasets_to_split + else: + raise TypeError( + "`datasets_to_split` should be a 'all' or a list of strings of " + "dataset names. Received " + f"{type(datasets_to_split).__name__} with value {datasets_to_split}." + ) + + self._execution_options: ExecutionOptions = ( + execution_options or DataConfig.default_ingest_options() + ) + self._enable_shard_locality = enable_shard_locality + + self._num_train_cpus = 0.0 + self._num_train_gpus = 0.0 + + def set_train_total_resources(self, num_train_cpus: float, num_train_gpus: float): + """Set the total number of CPUs and GPUs used by training. + + If CPU or GPU resource limits are not set, they will be set to the + total cluster resources minus the resources used by training. + """ + # TODO: We may also include other resources besides CPU and GPU. + self._num_train_cpus = num_train_cpus + self._num_train_gpus = num_train_gpus + + @DeveloperAPI + def configure( + self, + datasets: Dict[str, Dataset], + world_size: int, + worker_handles: Optional[List[ActorHandle]], + worker_node_ids: Optional[List[NodeIdStr]], + **kwargs, + ) -> List[Dict[str, DataIterator]]: + """Configure how Train datasets should be assigned to workers. + + Args: + datasets: The datasets dict passed to Train by the user. + world_size: The number of Train workers in total. + worker_handles: The actor handles of the Train workers. + worker_node_ids: The node ids of the Train workers. + kwargs: Forwards compatibility placeholder. + + Returns: + A list of dataset splits for each worker. The size of the list must be + equal to `world_size`. Each element of the list contains the assigned + `DataIterator` instances by name for the worker. + """ + output = [{} for _ in range(world_size)] + + for dataset_name, dataset in datasets.items(): + if dataset.name is None: + dataset.set_name(dataset_name) + + if self._datasets_to_split == "all": + datasets_to_split = set(datasets.keys()) + else: + datasets_to_split = set(self._datasets_to_split) + + locality_hints = worker_node_ids if self._enable_shard_locality else None + for name, ds in datasets.items(): + execution_options = copy.deepcopy(self._execution_options) + + if execution_options.is_resource_limits_default(): + # If "resource_limits" is not overriden by the user, + # add training-reserved resources to Data's exclude_resources. + execution_options.exclude_resources = ( + execution_options.exclude_resources.add( + ExecutionResources( + cpu=self._num_train_cpus, gpu=self._num_train_gpus + ) + ) + ) + + ds = ds.copy(ds) + ds.context.execution_options = execution_options + + if name in datasets_to_split: + for i, split in enumerate( + ds.streaming_split( + world_size, equal=True, locality_hints=locality_hints + ) + ): + output[i][name] = split + else: + for i in range(world_size): + output[i][name] = ds.iterator() + + return output + + @staticmethod + def default_ingest_options() -> ExecutionOptions: + """The default Ray Data options used for data ingest. + + By default, configurations are carried over from what is already set + in DataContext. + """ + ctx = ray.data.DataContext.get_current() + return ExecutionOptions( + # TODO(hchen): Re-enable `locality_with_output` by default after fixing + # https://github.com/ray-project/ray/issues/40607 + locality_with_output=ctx.execution_options.locality_with_output, + resource_limits=ctx.execution_options.resource_limits, + exclude_resources=ctx.execution_options.exclude_resources, + preserve_order=ctx.execution_options.preserve_order, + verbose_progress=ctx.execution_options.verbose_progress, + ) diff --git a/lib/python3.12/site-packages/ray/train/_internal/dl_predictor.py b/lib/python3.12/site-packages/ray/train/_internal/dl_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..a159cfe7b9e141f1d25cfdf160285505efeb80b8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/dl_predictor.py @@ -0,0 +1,103 @@ +import abc +from typing import Dict, Optional, TypeVar, Union + +import numpy as np +import pandas as pd + +from ray.air.util.data_batch_conversion import ( + BatchFormat, + _convert_batch_type_to_pandas, + _convert_pandas_to_batch_type, +) +from ray.train.predictor import Predictor +from ray.util.annotations import DeveloperAPI + +TensorType = TypeVar("TensorType") +TensorDtype = TypeVar("TensorDtype") + + +class DLPredictor(Predictor): + @abc.abstractmethod + def _arrays_to_tensors( + self, + numpy_arrays: Union[np.ndarray, Dict[str, np.ndarray]], + dtype: Optional[Union[TensorDtype, Dict[str, TensorDtype]]], + ) -> Union[TensorType, Dict[str, TensorType]]: + """Converts a NumPy ndarray batch to the tensor type for the DL framework. + + Args: + numpy_array: The numpy array to convert to a tensor. + dtype: The tensor dtype to use when creating the DL tensor. + ndarray: A (dict of) NumPy ndarray(s) that we wish to convert to a (dict of) + tensor(s). + dtype: A (dict of) tensor dtype(s) to use when creating the DL tensor; if + None, the dtype will be inferred from the NumPy ndarray data. + + Returns: + A deep learning framework specific tensor. + """ + raise NotImplementedError + + @abc.abstractmethod + def _tensor_to_array(self, tensor: TensorType) -> np.ndarray: + """Converts tensor framework specific tensor to a numpy array. + + Args: + tensor: A framework specific tensor. + + Returns: + A numpy array representing the input tensor. + """ + + raise NotImplementedError + + @abc.abstractmethod + @DeveloperAPI + def call_model( + self, inputs: Union[TensorType, Dict[str, TensorType]] + ) -> Union[TensorType, Dict[str, TensorType]]: + """Inputs the tensor to the model for this Predictor and returns the result. + + Args: + inputs: The tensor to input to the model. + + Returns: + A tensor or dictionary of tensors containing the model output. + """ + raise NotImplementedError + + @classmethod + @DeveloperAPI + def preferred_batch_format(cls) -> BatchFormat: + return BatchFormat.NUMPY + + def _predict_pandas( + self, + data: pd.DataFrame, + dtype: Optional[Union[TensorDtype, Dict[str, TensorDtype]]], + ) -> pd.DataFrame: + numpy_input = _convert_pandas_to_batch_type( + data, + BatchFormat.NUMPY, + self._cast_tensor_columns, + ) + numpy_output = self._predict_numpy(numpy_input, dtype) + return _convert_batch_type_to_pandas(numpy_output) + + def _predict_numpy( + self, + data: Union[np.ndarray, Dict[str, np.ndarray]], + dtype: Optional[Union[TensorDtype, Dict[str, TensorDtype]]], + ) -> Union[np.ndarray, Dict[str, np.ndarray]]: + # Single column selection return numpy array so preprocessors can be + # reused in both training and prediction + if isinstance(data, dict) and len(data) == 1: + data = next(iter(data.values())) + model_input = self._arrays_to_tensors(data, dtype) + model_output = self.call_model(model_input) + # TODO (jiaodong): Investigate perf implication of this. + # Move DL Tensor to CPU and convert to numpy. + if isinstance(model_output, dict): + return {k: self._tensor_to_array(v) for k, v in model_output.items()} + else: + return {"predictions": self._tensor_to_array(model_output)} diff --git a/lib/python3.12/site-packages/ray/train/_internal/framework_checkpoint.py b/lib/python3.12/site-packages/ray/train/_internal/framework_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..424dc0dc84be78609384668e4f95fdd6ee7d3077 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/framework_checkpoint.py @@ -0,0 +1,45 @@ +from typing import Optional + +import ray.cloudpickle as ray_pickle +from ray._common.utils import binary_to_hex, hex_to_binary +from ray.data.preprocessor import Preprocessor +from ray.train._checkpoint import Checkpoint + +PREPROCESSOR_KEY = "preprocessor_pkl" + + +class FrameworkCheckpoint(Checkpoint): + """A checkpoint to preserve the functionality of legacy + framework-specific checkpoints. + + Example: + + >>> import tempfile + >>> checkpoint = FrameworkCheckpoint(tempfile.mkdtemp()) + >>> checkpoint.get_preprocessor() is None + True + >>> preprocessor = Preprocessor() + >>> preprocessor._attr = 1234 + >>> checkpoint.set_preprocessor(preprocessor) + >>> checkpoint.get_preprocessor()._attr + 1234 + """ + + def get_preprocessor(self) -> Optional[Preprocessor]: + """Return the preprocessor stored in the checkpoint. + + Returns: + The preprocessor stored in the checkpoint, or ``None`` if no + preprocessor was stored. + """ + metadata = self.get_metadata() + preprocessor_bytes = metadata.get(PREPROCESSOR_KEY) + if preprocessor_bytes is None: + return None + return ray_pickle.loads(hex_to_binary(preprocessor_bytes)) + + def set_preprocessor(self, preprocessor: Preprocessor): + """Store a preprocessor with the checkpoint.""" + self.update_metadata( + {PREPROCESSOR_KEY: binary_to_hex(ray_pickle.dumps(preprocessor))} + ) diff --git a/lib/python3.12/site-packages/ray/train/_internal/session.py b/lib/python3.12/site-packages/ray/train/_internal/session.py new file mode 100644 index 0000000000000000000000000000000000000000..ef5e098efedc257649e30cd62392429a19281b6a --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/session.py @@ -0,0 +1,1204 @@ +import functools +import logging +import os +import platform +import queue +import sys +import threading +import time +import warnings +from dataclasses import dataclass +from datetime import datetime +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type + +import ray +from ray.air._internal.util import RunnerThread, StartTraceback +from ray.air.constants import ( + _ERROR_FETCH_TIMEOUT, + _RESULT_FETCH_TIMEOUT, + SESSION_MISUSE_LOG_ONCE_KEY, + TIME_THIS_ITER_S, + TIMESTAMP, +) +from ray.data import Dataset +from ray.train import Checkpoint +from ray.train._internal.accelerator import Accelerator +from ray.train._internal.storage import StorageContext +from ray.train.constants import ( + CHECKPOINT_DIR_NAME, + DETAILED_AUTOFILLED_KEYS, + RAY_CHDIR_TO_TRIAL_DIR, + TIME_TOTAL_S, + WORKER_HOSTNAME, + WORKER_NODE_IP, + WORKER_PID, + _v2_migration_warnings_enabled, +) +from ray.train.error import SessionMisuseError +from ray.train.utils import _log_deprecation_warning +from ray.util import queue as ray_queue +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.debug import log_once +from ray.util.placement_group import _valid_resource_shape +from ray.util.scheduling_strategies import ( + PlacementGroupSchedulingStrategy, + SchedulingStrategyT, +) + +if TYPE_CHECKING: + from ray.data import DataIterator + from ray.tune.execution.placement_groups import PlacementGroupFactory + + +logger = logging.getLogger(__name__) + + +@dataclass +class TrialInfo: + """The trial information to propagate to TrainSession.""" + + name: str + id: str + resources: Dict[str, float] + logdir: str + driver_ip: str + driver_node_id: str + experiment_name: Optional[str] = None + run_id: Optional[str] = None + + +class _FutureTrainingResult: + """A future that will be resolved to a `_TrainingResult`. + + This is needed for specific schedulers such as PBT that schedule saves. + + This wrapper should be removed after refactoring PBT to not schedule saves anymore. + """ + + def __init__(self, future: ray.ObjectRef): + self.future = future + + def resolve(self, block: bool = True) -> Optional["_TrainingResult"]: + """Resolve into ``_TrainingResult``. + + This will return None for function trainables if no checkpoint has been + saved before. + """ + if block: + timeout = None + else: + timeout = 1e-9 + try: + return ray.get(self.future, timeout=timeout) + except TimeoutError: + # Not ready, yet + pass + except Exception as exc: + logger.error(f"Error resolving result: {exc}") + + +class _TrainingResult: + """A (checkpoint, metrics) result reported by the user.""" + + def __init__(self, checkpoint: Optional[Checkpoint], metrics: Dict[str, Any]): + self.checkpoint = checkpoint + self.metrics = metrics + + def __repr__(self) -> str: + return f"TrainingResult(checkpoint={self.checkpoint}, metrics={self.metrics})" + + +# TODO(xwjiang): This needs a better name. +@DeveloperAPI +class _TrainSession: + """Holds information for training on each worker.""" + + def __init__( + self, + training_func: Callable, + world_rank: Optional[int], + local_rank: Optional[int], + node_rank: Optional[int], + local_world_size: Optional[int], + world_size: Optional[int], + trial_info: Optional[TrialInfo] = None, + dataset_shard: Optional[Dict[str, Dataset]] = None, + metadata: Dict[str, Any] = None, + checkpoint: Optional[Checkpoint] = None, + detailed_autofilled_metrics: bool = False, + storage: Optional[StorageContext] = None, + synchronous_result_reporting: bool = False, + ): + # `synchronous_result_reporting` refers to whether or not the + # training function is immediately unblocked to continue running + # after the main thread receives its result. + # Ex 1: For 2 Ray Train workers with synchronous_result_reporting=True, + # the worker that produces a result first will immediately will continue + # onto the next iteration. + # Ex 2: For a Tune function Trainable with `synchronous_result_reporting=False`, + # training will only continue with an explicit call to `session.get_next`. + # Synchronous reporting in example 2 is needed for Tune schedulers to + # be able to stop the execution of the training function at will, + # for advanced pausing schedulers (PBT, BOHB) and actor reuse. + self.synchronous_result_reporting = synchronous_result_reporting + + # Ray Train worker properties + # Note: These are set to None for Tune function Trainables. + self.dataset_shard = dataset_shard + self.metadata = metadata + + self.world_rank = world_rank + self.local_rank = local_rank + self.node_rank = node_rank + self.local_world_size = local_world_size + self.world_size = world_size + + assert storage + logger.debug(f"StorageContext on SESSION (rank={world_rank}):\n{storage}") + + # NOTE: `reset` will initialize many properties needed to start running the + # training_func as a thread. + self.reset( + training_func=training_func, + trial_info=trial_info, + storage=storage, + loaded_checkpoint=checkpoint, + ) + + # Autofilled metrics attributes. + self.detailed_autofilled_metrics = detailed_autofilled_metrics + self.last_report_time = time.time() + self.iteration = 0 + self.time_total = 0.0 + self.local_ip = self.get_current_ip() + + self.accelerator = None + self._state = {} + + def get_state(self, key: str) -> Any: + return self._state.get(key) + + def set_state(self, key: str, value: Any): + self._state[key] = value + + def get_current_ip(self): + self.local_ip = ray.util.get_node_ip_address() + return self.local_ip + + def start(self): + """Starts the training thread.""" + self.training_started = True + self.training_thread.start() + + def reset( + self, + training_func: Callable, + trial_info: TrialInfo, + storage: StorageContext, + loaded_checkpoint=None, + ): + # This lock is used to control the execution of the training thread. + self.continue_lock = threading.Semaphore(0) + + # This event is used to signal the training thread to stop. + self.stop_event = threading.Event() + + # Queue for sending results across threads. + self.result_queue = queue.Queue(1) + + # Queue for sending results from training actor to main thread. + self._inter_actor_queue: Optional[ray_queue.Queue[Dict]] = None + + # Queue for raising exceptions from runner thread to main thread. + # The error queue has a max size of one to prevent stacking error and force + # error reporting to block until finished. + self.error_queue = queue.Queue(1) + + # The Thread object that is running the training function. + self.training_thread = RunnerThread( + target=training_func, daemon=True, error_queue=self.error_queue + ) + + # Possibly override with new state + self.trial_info = trial_info + self.storage = storage + self.loaded_checkpoint = loaded_checkpoint + + # Reset state + self._state = {} + self.ignore_report = False + self.training_started = False + self._first_report = True + + # Change the working directory to a special trial folder. + # This is to ensure that all Ray Train workers have a common working directory. + os.makedirs(storage.trial_working_directory, exist_ok=True) + if bool(int(os.environ.get(RAY_CHDIR_TO_TRIAL_DIR, "1"))): + logger.debug( + f"Changing the working directory to: {storage.trial_working_directory}" + ) + os.chdir(storage.trial_working_directory) + + def pause_reporting(self): + """Ignore all future ``session.report()`` calls.""" + self.ignore_report = True + + def finish(self, timeout: Optional[float] = None) -> Optional[Any]: + """Finishes the training thread. + + Raises any Exception from training. + """ + # Set the stop event for the training thread to gracefully exit. + self.stop_event.set() + + # Release the lock so that training thread can process this event. + self.continue_lock.release() + + # Force a final (blocking) sync of artifacts in the trial path to storage. + self.storage.persist_artifacts(force=True) + + # Wait for training to finish. + # This will raise any errors that occur during training, including SystemError + # This returns the result of the training function. + output = None + if self.training_started: + output = self.training_thread.join(timeout=timeout) + + return output + + def get_next(self) -> Optional[_TrainingResult]: + """Gets the next ``_TrainingResult`` from the result queue. + + If the result queue is empty, then this function returns ``None``. + """ + if not self.training_started: + raise RuntimeError("Please call start before calling get_next.") + + if self.synchronous_result_reporting: + # There's no need to release the lock on the first report + # since `start` already started the training thread. + if not self._first_report: + # Release the lock to trigger training to continue, + # until the next call to report. + self.continue_lock.release() + self._first_report = False + + result = None + # While training is still ongoing, attempt to get the result. + while result is None and self.training_thread.is_alive(): + result = self._get_result_from_queues(block=True) + + # If no result was found, then the runner must no longer be alive. + if result is None: + # Try one last time to fetch results in case results were + # reported in between the time of the last check and the + # termination of the thread runner. + result = self._get_result_from_queues(block=False) + + # check if error occurred inside the thread runner. + if result is None: + # only raise an error from the runner if all results are consumed + self._report_thread_runner_error(block=True) + else: + if not self.error_queue.empty(): + logger.debug( + ( + "Runner error waiting to be raised in main thread. " + "Logging all available results first." + ) + ) + + if not self.synchronous_result_reporting: + # At this point, the training thread has reached + # the `train.report` and is blocked there. + # If performing asynchronous result reporting, + # release the lock to allow each worker to keep training + # immediately after the coordinator fetches their result. + self.continue_lock.release() + + # Return None if there are no more results to fetch. + return result + + def _get_or_create_inter_actor_queue(self): + """Get or create the inter-actor queue.""" + if self._inter_actor_queue is None: + self._inter_actor_queue = ray_queue.Queue(1, actor_options={"num_cpus": 0}) + return self._inter_actor_queue + + def _get_result_from_queues(self, block: bool) -> Optional[_TrainingResult]: + """Get result from result queue. Pass result from training actor result queue if needed.""" + result = None + if self._inter_actor_queue is not None: + try: + inter_actor_item = self._inter_actor_queue.get( + block=block, timeout=_RESULT_FETCH_TIMEOUT + ) + if inter_actor_item: + # Must release continue_lock to allow report to work. + self.continue_lock.release() + self.report(inter_actor_item) + except ray_queue.Empty: + pass + try: + result = self.result_queue.get(block=block, timeout=_RESULT_FETCH_TIMEOUT) + except queue.Empty: + pass + return result + + def _auto_fill_metrics(self, result: dict) -> dict: + """Add autofilled metrics and update attributes.""" + current_time = time.time() + current_datetime = datetime.now() + if TIME_THIS_ITER_S in result: + time_this_iter = result[TIME_THIS_ITER_S] + else: + time_this_iter = current_time - self.last_report_time + self.iteration += 1 + self.time_total += time_this_iter + self.last_report_time = current_time + + auto_filled_metrics = { + TIMESTAMP: int(time.mktime(current_datetime.timetuple())), + TIME_TOTAL_S: self.time_total, + WORKER_PID: os.getpid(), + WORKER_HOSTNAME: platform.node(), + WORKER_NODE_IP: self.local_ip, + } + + if not self.detailed_autofilled_metrics: + auto_filled_metrics = { + k: v + for k, v in auto_filled_metrics.items() + if k not in DETAILED_AUTOFILLED_KEYS + } + + result = result.copy() + result.update(auto_filled_metrics) + return result + + def _auto_fill_checkpoint_metrics(self, result: dict) -> dict: + """Add autofilled metrics and update attributes.""" + current_datetime = datetime.now() + + auto_filled_metrics = { + TIMESTAMP: int(time.mktime(current_datetime.timetuple())) + } + result = result.copy() + result.update(auto_filled_metrics) + return result + + def _report_thread_runner_error(self, block=False): + try: + e = self.error_queue.get(block=block, timeout=_ERROR_FETCH_TIMEOUT) + raise StartTraceback from e + except queue.Empty: + pass + + def _report_training_result(self, training_result: _TrainingResult) -> None: + """Place a training result on the result queue for the main thread to process, + then block until the main thread signals that training should continue. + + NOTE: This is used internally to report results from Train to Tune + without persisting checkpoints to storage 2 times. + `report` is the public API that directly persists to storage, which + should only be called by user code. + """ + if training_result.checkpoint: + # NOTE: This populates `train.get_checkpoint` + self.loaded_checkpoint = training_result.checkpoint + + # Add result to a thread-safe queue. + self.result_queue.put(training_result, block=True) + + # Acquire lock to stop the training thread until main thread + # triggers resume. + self.continue_lock.acquire() + + # If the trial should be terminated, exit gracefully. + # NOTE: This is only really useful if `synchronous_result_reporting=True`. + # Otherwise, the lock is immediately released on reporting, and this + # check is skipped before the main thread decides to set the stop event. + if self.stop_event.is_set(): + self.stop_event.clear() + sys.exit(0) + + def report(self, metrics: Dict, checkpoint: Optional[Checkpoint] = None) -> None: + # Special case: early fail for Torch tensors + if "torch" in sys.modules: + from ray.air._internal.torch_utils import contains_tensor + + if contains_tensor(metrics): + raise ValueError( + "Passing objects containg Torch tensors as metrics " + "is not supported as it will throw an exception on " + "deserialization. You can either convert the tensors " + "to Python objects or report a `train.Checkpoint` " + "with `ray.train.report` to store your Torch objects." + ) + + if self.ignore_report: + return + + metrics = self._auto_fill_metrics(metrics) + + persisted_checkpoint = None + if checkpoint: + self.storage._update_checkpoint_index(metrics) + + # Persist the reported checkpoint files to storage. + persisted_checkpoint = self.storage.persist_current_checkpoint(checkpoint) + + metrics[CHECKPOINT_DIR_NAME] = self.storage.checkpoint_dir_name + else: + metrics[CHECKPOINT_DIR_NAME] = None + + # Persist trial artifacts to storage. + force_artifact_sync = ( + persisted_checkpoint + and self.storage.sync_config.sync_artifacts_on_checkpoint + ) + self.storage.persist_artifacts(force=force_artifact_sync) + + # Set additional user metadata from the Trainer. + if persisted_checkpoint and self.metadata: + user_metadata = persisted_checkpoint.get_metadata() + for k, v in self.metadata.items(): + # Update keys not already set by the user. This gives user-set keys + # precedence over keys set at the Trainer level. + if k not in user_metadata: + user_metadata[k] = v + persisted_checkpoint.set_metadata(user_metadata) + + result = _TrainingResult(checkpoint=persisted_checkpoint, metrics=metrics) + + self._report_training_result(result) + + @property + def experiment_name(self) -> str: + return self.trial_info.experiment_name + + @property + def trial_name(self) -> str: + return self.trial_info.name + + @property + def trial_id(self) -> str: + return self.trial_info.id + + @property + def run_id(self) -> str: + return self.trial_info.run_id + + @property + def trial_resources(self) -> "PlacementGroupFactory": + return self.trial_info.resources + + @property + def trial_dir(self) -> str: + return self.trial_info.logdir + + def get_dataset_shard( + self, + dataset_name: Optional[str] = None, + ) -> Optional["DataIterator"]: + shard = self.dataset_shard + if shard is None: + warnings.warn( + "No dataset passed in. Returning None. Make sure to " + "pass in a Dataset to Trainer.run to use this " + "function." + ) + elif isinstance(shard, dict): + if not dataset_name: + raise RuntimeError( + "Multiple datasets were passed into ``Trainer``, " + "but no ``dataset_name`` is passed into " + "``get_dataset_shard``. Please specify which " + "dataset shard to retrieve." + ) + return shard.get(dataset_name) + return shard + + +# Cache of resource dicts that have been checked by the launch hook already. +_checked_resources: Set[frozenset] = set() + +# Global _TrainSession object initialized by Ray Tune function trainables +# and Ray Train V1 workers. +_session: Optional[_TrainSession] = None + + +def _tune_task_and_actor_launch_hook( + fn, resources: Dict[str, float], strategy: Optional[SchedulingStrategyT] +): + """Launch hook to catch nested tasks that can't fit in the placement group. + + This gives users a nice warning in case they launch a nested task in a Tune trial + without reserving resources in the trial placement group to fit it. + """ + + # Already checked, skip for performance reasons. + key = frozenset({(k, v) for k, v in resources.items() if v > 0}) + if not key or key in _checked_resources: + return + + # No need to check if placement group is None. + if ( + not isinstance(strategy, PlacementGroupSchedulingStrategy) + or strategy.placement_group is None + ): + return + + # Check if the resource request is targeting the current placement group. + cur_pg = ray.util.get_current_placement_group() + if not cur_pg or strategy.placement_group.id != cur_pg.id: + return + + _checked_resources.add(key) + + # Check if the request can be fulfilled by the current placement group. + pgf = get_trial_resources() + + if pgf.head_bundle_is_empty: + available_bundles = cur_pg.bundle_specs[0:] + else: + available_bundles = cur_pg.bundle_specs[1:] + + # Check if the request can be fulfilled by the current placement group. + if _valid_resource_shape(resources, available_bundles): + return + + if fn.class_name: + submitted = "actor" + name = fn.module_name + "." + fn.class_name + "." + fn.function_name + else: + submitted = "task" + name = fn.module_name + "." + fn.function_name + + # Normalize the resource spec so it looks the same as the placement group bundle. + main_resources = cur_pg.bundle_specs[0] + resources = {k: float(v) for k, v in resources.items() if v > 0} + + raise RuntimeError( + f"No trial resources are available for launching the {submitted} `{name}`. " + "To resolve this, specify the Tune option:\n\n" + "> resources_per_trial=tune.PlacementGroupFactory(\n" + f"> [{main_resources}] + [{resources}] * N\n" + "> )\n\n" + f"Where `N` is the number of slots to reserve for trial {submitted}s. " + "If you are using a Ray training library, there might be a utility function " + "to set this automatically for you. For more information, refer to " + "https://docs.ray.io/en/latest/tune/tutorials/tune-resources.html" + ) + + +def init_session(*args, **kwargs) -> None: + global _session + if _session: + raise ValueError( + "A Train session is already in use. Do not call " + "`init_session()` manually." + ) + + # Setup hooks for generating placement group resource deadlock warnings. + from ray import actor, remote_function + + if "TUNE_DISABLE_RESOURCE_CHECKS" not in os.environ: + actor._actor_launch_hook = _tune_task_and_actor_launch_hook + remote_function._task_launch_hook = _tune_task_and_actor_launch_hook + + _session = _TrainSession(*args, **kwargs) + + +def get_session() -> Optional[_TrainSession]: + return _session + + +def shutdown_session(): + """Shuts down the initialized session.""" + global _session + _session = None + + +def _raise_accelerator_session_misuse(): + """Raises a SessionMisuseError because a utility function was used improperly.""" + raise SessionMisuseError( + "prepare/accelerate utility functions should be called inside a training " + "function executed by `Trainer.run`" + ) + + +def get_accelerator(default_accelerator_cls: Type[Accelerator]) -> Accelerator: + """The accelerator for this training session. + + If an accelerator has not been set, then this method will construct an + accelerator using the provided accelerator class. + + Raises: + SessionMisuseError: if the session is uninitialized. + """ + session = get_session() + if session is None: + _raise_accelerator_session_misuse() + if session.accelerator is None: + session.accelerator = default_accelerator_cls() + return session.accelerator + + +def set_accelerator(accelerator: Accelerator) -> None: + """Sets the accelerator for this training session. + + Args: + accelerator: The accelerator to use for training. + + Raises: + SessionMisuseError: if the session is unitialized. + RuntimeError: if the accelerator has already been set. + """ + session = get_session() + if session is None: + _raise_accelerator_session_misuse() + if session.accelerator is not None: + raise RuntimeError("Cannot change accelerator once set.") + session.accelerator = accelerator + + +def _warn_session_misuse(default_value: Any = None): + """Warns if fn is being used outside of session and returns ``default_value``.""" + + def inner(fn: Callable): + fn_name = fn.__name__ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + session = get_session() + if not session: + if log_once(f"{SESSION_MISUSE_LOG_ONCE_KEY}-{fn_name}"): + warnings.warn( + f"`{fn_name}` is meant to only be " + "called inside a function that is executed by a Tuner" + f" or Trainer. Returning `{default_value}`." + ) + return default_value + return fn(*args, **kwargs) + + return wrapper + + return inner + + +@PublicAPI(stability="stable") +@_warn_session_misuse() +def report( + metrics: Dict, + *, + checkpoint: Optional[Checkpoint] = None, + checkpoint_dir_name: Optional[str] = None, +) -> None: + """Report metrics and optionally save a checkpoint. + + If a checkpoint is provided, it will be + :ref:`persisted to storage `. + + If this is called in multiple distributed training workers: + + - Only the metrics reported by the rank 0 worker will be tracked by Ray Train. + See :ref:`the metrics logging guide `. + - A checkpoint will be registered as long as one or more workers reports + checkpoint that is not None. + See the :ref:`checkpointing guide `. + - Checkpoints from multiple workers will be merged into one directory + in persistent storage. + See :ref:`the distributed checkpointing guide `. + + .. note:: + + Each invocation of this method will automatically increment the underlying + ``training_iteration`` number. The physical meaning of this "iteration" is + defined by user depending on how often they call ``report``. + It does not necessarily map to one epoch. + + .. warning:: + + All workers must call `ray.train.report` the same number of times + so that Ray Train can properly synchronize the training state across + workers. Otherwise, your training will hang. + + .. warning:: + + This method does NOT act as a barrier for distributed training workers. + Workers will upload their checkpoint, then continue training immediately. + If you need to synchronize workers, you can use a framework-native barrier + such as `torch.distributed.barrier()`. + + Example: + + .. testcode:: + + import tempfile + + from ray import train + from ray.train import Checkpoint + from ray.train.torch import TorchTrainer + + + def train_func(config): + start_epoch = 0 + checkpoint = train.get_checkpoint() + if checkpoint: + with checkpoint.as_directory() as checkpoint_dir: + # Load back training state + ... + + for epoch in range(start_epoch, config.get("num_epochs", 10)): + # Do training... + + metrics = {"loss": ...} + + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + # Save the checkpoint... + # torch.save(...) + + checkpoint = Checkpoint.from_directory(temp_checkpoint_dir) + + # Example: Only the rank 0 worker uploads the checkpoint. + if ray.train.get_context().get_world_rank() == 0: + train.report(metrics, checkpoint=checkpoint) + else: + train.report(metrics, checkpoint=None) + + trainer = TorchTrainer( + train_func, scaling_config=train.ScalingConfig(num_workers=2) + ) + + Args: + metrics: The metrics you want to report. + checkpoint: The optional checkpoint you want to report. + """ + if checkpoint_dir_name is not None: + logger.warning( + "`checkpoint_dir_name` is only supported in the new Ray Train " + "implementation, which can be enabled with `RAY_TRAIN_V2_ENABLED=1`. " + "This argument will be ignored." + ) + + # If we are running in a Tune function, switch to `ray.tune.report`. + from ray.tune.trainable.trainable_fn_utils import _in_tune_session + + if _in_tune_session(): + import ray.tune + + if _v2_migration_warnings_enabled(): + _log_deprecation_warning( + "`ray.train.report` should be switched to " + "`ray.tune.report` when running in a function " + "passed to Ray Tune. This will be an error in the future. " + "See this issue for more context: " + "https://github.com/ray-project/ray/issues/49454" + ) + return ray.tune.report(metrics, checkpoint=checkpoint) + + get_session().report(metrics, checkpoint=checkpoint) + + +@PublicAPI(stability="stable") +@_warn_session_misuse() +def get_checkpoint() -> Optional[Checkpoint]: + """Access the latest reported checkpoint to resume from if one exists. + + Example: + + .. testcode:: + + import tempfile + + from ray import train + from ray.train import Checkpoint + from ray.train.torch import TorchTrainer + + + def train_func(config): + start_epoch = 0 + checkpoint = train.get_checkpoint() + if checkpoint: + with checkpoint.as_directory() as checkpoint_dir: + # Load back training state + ... + + for epoch in range(start_epoch, config.get("num_epochs", 10)): + # Do training... + + metrics = {"loss": ...} + + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + # Save the checkpoint... + + checkpoint = Checkpoint.from_directory(temp_checkpoint_dir) + train.report(metrics, checkpoint=checkpoint) + + trainer = TorchTrainer( + train_func, scaling_config=train.ScalingConfig(num_workers=2) + ) + + Returns: + Checkpoint object if the session is currently being resumed. + Otherwise, return None. + """ + # If we are running in a Tune function, switch to `ray.tune.get_checkpoint`. + from ray.tune.trainable.trainable_fn_utils import _in_tune_session + + if _in_tune_session(): + import ray.tune + + if _v2_migration_warnings_enabled(): + _log_deprecation_warning( + "`ray.train.get_checkpoint` should be switched to " + "`ray.tune.get_checkpoint` when running in a function " + "passed to Ray Tune. This will be an error in the future. " + "See this issue for more context: " + "https://github.com/ray-project/ray/issues/49454" + ) + return ray.tune.get_checkpoint() + + return get_session().loaded_checkpoint + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_metadata() -> Dict[str, Any]: + """User metadata dict passed to the Trainer constructor.""" + return get_session().metadata + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_experiment_name() -> str: + """Experiment name for the corresponding trial.""" + return get_session().experiment_name + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_trial_name() -> str: + """Trial name for the corresponding trial.""" + return get_session().trial_name + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_trial_id() -> str: + """Trial id for the corresponding trial.""" + return get_session().trial_id + + +@PublicAPI(stability="alpha") +@_warn_session_misuse() +def get_run_id() -> str: + """Unique Train Run id for the corresponding trial.""" + return get_session().run_id + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_trial_resources() -> "PlacementGroupFactory": + """Trial resources for the corresponding trial.""" + return get_session().trial_resources + + +@PublicAPI(stability="beta") +@_warn_session_misuse() +def get_trial_dir() -> str: + """Log directory corresponding to the trial directory for a Tune session. + If calling from a Train session, this will give the trial directory of its parent + Tune session. + + .. testcode:: + + import ray.tune + + def train_func(config): + print(ray.tune.get_context().get_trial_dir()) + + tuner = ray.tune.Tuner(train_func) + tuner.fit() + + .. testoutput:: + :options: +MOCK + + /Users/root/ray_results/train_func_2023-07-19_15-01-37/train_func_d620c_00000_0_2023-07-19_15-01-40 + """ + return get_session().trial_dir + + +@PublicAPI(stability="beta") +@_warn_session_misuse(default_value=1) +def get_world_size() -> int: + """Get the current world size (i.e. total number of workers) for this run. + + .. testcode:: + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.tensorflow import TensorflowTrainer + + NUM_WORKERS = 2 + + def train_loop_per_worker(config): + assert train.get_context().get_world_size() == NUM_WORKERS + + train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv") + trainer = TensorflowTrainer( + train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=NUM_WORKERS), + datasets={"train": train_dataset} + ) + trainer.fit() + + .. testoutput:: + :hide: + + ... + """ + session = get_session() + if not hasattr(session, "world_size"): + raise RuntimeError( + "`get_world_size` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.world_size + + +@PublicAPI(stability="beta") +@_warn_session_misuse(default_value=0) +def get_world_rank() -> int: + """Get the world rank of this worker. + + .. testcode:: + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.tensorflow import TensorflowTrainer + + def train_loop_per_worker(config): + if train.get_context().get_world_rank() == 0: + print("Worker 0") + + train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv") + trainer = TensorflowTrainer( + train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=2), + datasets={"train": train_dataset} + ) + trainer.fit() + + .. testoutput:: + :hide: + + ... + """ + session = get_session() + if not hasattr(session, "world_rank"): + raise RuntimeError( + "`get_world_rank` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.world_rank + + +@PublicAPI(stability="beta") +@_warn_session_misuse(default_value=0) +def get_local_rank() -> int: + """Get the local rank of this worker (rank of the worker on its node). + + .. testcode:: + + import torch + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.torch import TorchTrainer + + def train_loop_per_worker(config): + if torch.cuda.is_available(): + torch.cuda.set_device(train.get_context().get_local_rank()) + ... + + train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv") + trainer = TorchTrainer( + train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=2, use_gpu=True), + datasets={"train": train_dataset} + ) + trainer.fit() + + .. testoutput:: + :hide: + + ... + """ + session = get_session() + if not hasattr(session, "local_rank"): + raise RuntimeError( + "`get_local_rank` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.local_rank + + +@PublicAPI(stability="beta") +@_warn_session_misuse(default_value=0) +def get_local_world_size() -> int: + """Get the local world size of this node (i.e. number of workers on this node). + + Example: + + .. testcode:: + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.torch import TorchTrainer + + def train_loop_per_worker(): + print(train.get_context().get_local_world_size()) + + train_dataset = ray.data.from_items( + [{"x": x, "y": x + 1} for x in range(32)]) + trainer = TorchTrainer(train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=1), + datasets={"train": train_dataset}) + trainer.fit() + + .. testoutput:: + :hide: + + ... + """ + session = get_session() + if not hasattr(session, "local_world_size"): + raise RuntimeError( + "`get_local_world_size` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.local_world_size + + +@PublicAPI(stability="beta") +@_warn_session_misuse(default_value=0) +def get_node_rank() -> int: + """Get the rank of this node. + + Example: + + .. testcode:: + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.torch import TorchTrainer + + def train_loop_per_worker(): + print(train.get_context().get_node_rank()) + + train_dataset = ray.data.from_items( + [{"x": x, "y": x + 1} for x in range(32)]) + trainer = TorchTrainer(train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=1), + datasets={"train": train_dataset}) + trainer.fit() + + .. testoutput:: + :hide: + + ... + """ + session = get_session() + if not hasattr(session, "node_rank"): + raise RuntimeError( + "`get_node_rank` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.node_rank + + +@PublicAPI(stability="stable") +@_warn_session_misuse() +def get_dataset_shard( + dataset_name: Optional[str] = None, +) -> Optional["DataIterator"]: + """Returns the :class:`ray.data.DataIterator` shard for this worker. + + Call :meth:`~ray.data.DataIterator.iter_torch_batches` or + :meth:`~ray.data.DataIterator.to_tf` on this shard to convert it to the + appropriate framework-specific data type. + + .. testcode:: + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.torch import TorchTrainer + + def train_loop_per_worker(config): + ... + for epoch in range(2): + # Trainer will automatically handle sharding. + data_shard = train.get_dataset_shard("train") + for batch in data_shard.iter_torch_batches(): + ... + + train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv") + trainer = TorchTrainer( + train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=2), + datasets={"train": train_dataset} + ) + trainer.fit() + + .. testoutput:: + :hide: + + ... + + Args: + dataset_name: If a Dictionary of Datasets was passed to ``Trainer``, then + specifies which dataset shard to return. + + Returns: + The ``DataIterator`` shard to use for this worker. + If no dataset is passed into Trainer, then return None. + """ + session = get_session() + if not hasattr(session, "get_dataset_shard"): + raise RuntimeError( + "`get_dataset_shard` can only be called for TrainSession! " + "Make sure you only use that in `train_loop_per_worker` function" + "that is passed into `DataParallelTrainer`." + ) + return session.get_dataset_shard(dataset_name) + + +@DeveloperAPI +@_warn_session_misuse() +def get_storage() -> StorageContext: + """Returns the :class:`~ray.train._internal.storage.StorageContext` storage + context which gives advanced access to the filesystem and paths + configured through `RunConfig`. + + NOTE: This is a developer API, and the `StorageContext` interface may change + without notice between minor versions. + """ + return get_session().storage + + +def _in_ray_train_worker() -> bool: + """Check if the current process is a Ray Train V1 worker.""" + return bool(get_session()) and get_session().world_rank is not None diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__init__.py b/lib/python3.12/site-packages/ray/train/_internal/state/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..604a4fa3932930fa1d728ceb6e9bdd4619449bf6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/state/__init__.py @@ -0,0 +1,14 @@ +from ray.train._internal.state.state_manager import TrainRunStateManager + +try: + import pydantic # noqa: F401 +except ImportError: + raise ModuleNotFoundError( + "pydantic isn't installed." + "To install pydantic, please run 'pip install pydantic'" + ) + + +__all__ = [ + "TrainRunStateManager", +] diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb6a9517184c3fc1bea845d5a23ee967d1c58988 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/export.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/export.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5eacb6e869bd777c88815862c4cef53d17e2f9f1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/export.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/schema.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c6eb6797c714b40c69ebdb0b79fc8bcdae14943 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_actor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d7feb3688f0b5b5320e3a138eaf57783fb28ed0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_actor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f77a9cd265e7447b781e92e6d45af563a91ce1b Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/_internal/state/__pycache__/state_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/export.py b/lib/python3.12/site-packages/ray/train/_internal/state/export.py new file mode 100644 index 0000000000000000000000000000000000000000..5ae562dd9f814af464c7a33681b047332277cc64 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/state/export.py @@ -0,0 +1,105 @@ +from typing import Optional + +from ray.core.generated.export_train_state_pb2 import ( + ExportTrainRunAttemptEventData as ProtoTrainRunAttempt, + ExportTrainRunEventData as ProtoTrainRun, +) +from ray.train._internal.state.schema import ( + ActorStatusEnum, + RunStatusEnum, + TrainRunInfo, + TrainWorkerInfo, +) + +TRAIN_SCHEMA_VERSION = 1 +RAY_TRAIN_VERSION = 1 + +# Status mapping dictionaries +_ACTOR_STATUS_MAP = { + ActorStatusEnum.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE, + ActorStatusEnum.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD, +} + +_RUN_ATTEMPT_STATUS_MAP = { + RunStatusEnum.STARTED: ProtoTrainRunAttempt.RunAttemptStatus.PENDING, + RunStatusEnum.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING, + RunStatusEnum.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED, + RunStatusEnum.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED, + RunStatusEnum.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED, +} + +_RUN_STATUS_MAP = { + RunStatusEnum.STARTED: ProtoTrainRun.RunStatus.INITIALIZING, + RunStatusEnum.RUNNING: ProtoTrainRun.RunStatus.RUNNING, + RunStatusEnum.FINISHED: ProtoTrainRun.RunStatus.FINISHED, + RunStatusEnum.ERRORED: ProtoTrainRun.RunStatus.ERRORED, + RunStatusEnum.ABORTED: ProtoTrainRun.RunStatus.ABORTED, +} + + +def _ms_to_ns(ms: Optional[int]) -> Optional[int]: + if ms is None: + return None + return ms * 1000000 + + +# Helper conversion functions +def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources: + """Convert resources dictionary to protobuf TrainResources.""" + return ProtoTrainRunAttempt.TrainResources(resources=resources) + + +def _to_proto_worker(worker: TrainWorkerInfo) -> ProtoTrainRunAttempt.TrainWorker: + """Convert TrainWorker to protobuf format.""" + proto_worker = ProtoTrainRunAttempt.TrainWorker( + world_rank=worker.world_rank, + local_rank=worker.local_rank, + node_rank=worker.node_rank, + actor_id=bytes.fromhex(worker.actor_id), + node_id=bytes.fromhex(worker.node_id), + node_ip=worker.node_ip, + pid=worker.pid, + gpu_ids=worker.gpu_ids, + status=_ACTOR_STATUS_MAP[worker.status], + resources=_to_proto_resources(worker.resources), + ) + + return proto_worker + + +# Main conversion functions +def train_run_info_to_proto_run(run_info: TrainRunInfo) -> ProtoTrainRun: + """Convert TrainRunInfo to TrainRun protobuf format.""" + proto_run = ProtoTrainRun( + schema_version=TRAIN_SCHEMA_VERSION, + ray_train_version=RAY_TRAIN_VERSION, + id=run_info.id, + name=run_info.name, + job_id=bytes.fromhex(run_info.job_id), + controller_actor_id=bytes.fromhex(run_info.controller_actor_id), + status=_RUN_STATUS_MAP[run_info.run_status], + status_detail=run_info.status_detail, + start_time_ns=_ms_to_ns(run_info.start_time_ms), + end_time_ns=_ms_to_ns(run_info.end_time_ms), + ) + + return proto_run + + +def train_run_info_to_proto_attempt(run_info: TrainRunInfo) -> ProtoTrainRunAttempt: + """Convert TrainRunInfo to TrainRunAttempt protobuf format.""" + + proto_attempt = ProtoTrainRunAttempt( + schema_version=TRAIN_SCHEMA_VERSION, + ray_train_version=RAY_TRAIN_VERSION, + run_id=run_info.id, + attempt_id=run_info.id, # Same as run_id + status=_RUN_ATTEMPT_STATUS_MAP[run_info.run_status], + status_detail=run_info.status_detail, + start_time_ns=_ms_to_ns(run_info.start_time_ms), + end_time_ns=_ms_to_ns(run_info.end_time_ms), + resources=[_to_proto_resources(r) for r in run_info.resources], + workers=[_to_proto_worker(worker) for worker in run_info.workers], + ) + + return proto_attempt diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/schema.py b/lib/python3.12/site-packages/ray/train/_internal/state/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..87f9f51a6c3d1bfaa8b5d333888ad971a9ce3c7f --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/state/schema.py @@ -0,0 +1,164 @@ +from enum import Enum +from typing import Dict, List, Optional + +from ray._common.pydantic_compat import BaseModel, Field +from ray.dashboard.modules.job.pydantic_models import JobDetails +from ray.util.annotations import DeveloperAPI + +MAX_ERROR_STACK_TRACE_LENGTH = 50000 + + +@DeveloperAPI +class RunStatusEnum(str, Enum): + """Enumeration for the status of a train run.""" + + # (Deprecated) Replaced by RUNNING. + # The train run has started + STARTED = "STARTED" + # The train run is running + RUNNING = "RUNNING" + # The train run was terminated as expected + FINISHED = "FINISHED" + # The train run was terminated early due to errors in the training function + ERRORED = "ERRORED" + # The train run was terminated early due to system errors or controller errors + ABORTED = "ABORTED" + + +@DeveloperAPI +class ActorStatusEnum(str, Enum): + DEAD = "DEAD" + ALIVE = "ALIVE" + + +@DeveloperAPI +class TrainWorkerInfo(BaseModel): + """Metadata of a Ray Train worker.""" + + actor_id: str = Field(description="Actor ID of the worker.") + world_rank: int = Field(description="World rank of the worker.") + local_rank: int = Field(description="Local rank of the worker.") + node_rank: int = Field(description="Node rank of the worker.") + node_id: str = Field(description="ID of the node that the worker is running on.") + node_ip: str = Field( + description="IP address of the node that the worker is running on." + ) + pid: int = Field(description="Process ID of the worker.") + gpu_ids: List[int] = Field( + description="A list of GPU ids allocated to that worker." + ) + status: ActorStatusEnum = Field( + description="The status of the train worker actor. It can be ALIVE or DEAD." + ) + resources: Dict[str, float] = Field( + description="The resources allocated to the worker." + ) + + +@DeveloperAPI +class MemoryInfo(BaseModel): + rss: int + vms: int + pfaults: Optional[int] + pageins: Optional[int] + + +@DeveloperAPI +class ProcessStats(BaseModel): + cpuPercent: float + # total memory, free memory, memory used ratio + mem: Optional[List[int]] + memoryInfo: MemoryInfo + + +class ProcessGPUUsage(BaseModel): + # This gpu usage stats from a process + pid: int + gpuMemoryUsage: int + + +@DeveloperAPI +class GPUStats(BaseModel): + uuid: str + index: int + name: str + utilizationGpu: Optional[float] + memoryUsed: float + memoryTotal: float + processInfo: ProcessGPUUsage + + +@DeveloperAPI +class TrainWorkerInfoWithDetails(TrainWorkerInfo): + """Metadata of a Ray Train worker.""" + + processStats: Optional[ProcessStats] = Field( + None, description="Process stats of the worker." + ) + gpus: List[GPUStats] = Field( + default_factory=list, + description=( + "GPU stats of the worker. " + "Only returns GPUs that are attached to the worker process." + ), + ) + + +@DeveloperAPI +class TrainDatasetInfo(BaseModel): + name: str = Field( + description="The key of the dataset dict specified in Ray Train Trainer." + ) + dataset_uuid: str = Field(description="The uuid of the dataset.") + dataset_name: Optional[str] = Field(description="The name of the dataset.") + + +@DeveloperAPI +class TrainRunInfo(BaseModel): + """Metadata for a Ray Train run and information about its workers.""" + + name: str = Field(description="The name of the Train run.") + id: str = Field(description="The unique identifier for each Train run.") + job_id: str = Field(description="The Ray Job ID.") + controller_actor_id: str = Field(description="Actor Id of the Train controller.") + workers: List[TrainWorkerInfo] = Field( + description="A List of Train workers sorted by global ranks." + ) + datasets: List[TrainDatasetInfo] = Field( + description="A List of dataset info for this Train run." + ) + run_status: RunStatusEnum = Field( + description="The current status of the train run. It can be one of the " + "following: RUNNING, FINISHED, ERRORED, or ABORTED." + ) + status_detail: str = Field( + description="Detailed information about the current run status, " + "such as error messages." + ) + start_time_ms: int = Field( + description="The UNIX timestamp of the start time of this Train run." + ) + end_time_ms: Optional[int] = Field( + description="The UNIX timestamp of the end time of this Train run. " + "If null, the Train run has not ended yet." + ) + resources: List[Dict[str, float]] = Field( + description="The resources allocated to the worker." + ) + + +@DeveloperAPI +class TrainRunInfoWithDetails(TrainRunInfo): + """Metadata for a Ray Train run and information about its workers.""" + + workers: List[TrainWorkerInfoWithDetails] = Field( + description="A List of Train workers sorted by global ranks." + ) + job_details: Optional[JobDetails] = Field( + None, description="Details of the job that started this Train run." + ) + + +@DeveloperAPI +class TrainRunsResponse(BaseModel): + train_runs: List[TrainRunInfoWithDetails] diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/state_actor.py b/lib/python3.12/site-packages/ray/train/_internal/state/state_actor.py new file mode 100644 index 0000000000000000000000000000000000000000..4001d31129c1eba07a8b38d5b485db0f75647e7f --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/state/state_actor.py @@ -0,0 +1,152 @@ +import logging +import os +import threading +from typing import Dict, Optional + +import ray +from ray._private.event.export_event_logger import ( + EventLogType, + check_export_api_enabled, + get_export_event_logger, +) +from ray.actor import ActorHandle +from ray.train._internal.state.schema import TrainRunInfo + +logger = logging.getLogger(__name__) + + +@ray.remote(num_cpus=0) +class TrainStateActor: + def __init__(self): + self._run_infos: Dict[str, TrainRunInfo] = {} + + ( + self._export_logger, + self._is_train_run_export_api_enabled, + self._is_train_run_attempt_export_api_enabled, + ) = self._init_export_logger() + + def register_train_run(self, run_info: TrainRunInfo) -> None: + # Register a new train run. + self._run_infos[run_info.id] = run_info + + self._maybe_export_train_run(run_info) + self._maybe_export_train_run_attempt(run_info) + + def get_train_run(self, run_id: str) -> Optional[TrainRunInfo]: + # Retrieve a registered run with its id + return self._run_infos.get(run_id, None) + + def get_all_train_runs(self) -> Dict[str, TrainRunInfo]: + # Retrieve all registered train runs + return self._run_infos + + # ============================ + # Export API + # ============================ + + def is_export_api_enabled(self) -> bool: + return self._export_logger is not None + + def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]: + """Initialize the export logger and check if the export API is enabled. + + Returns: + A tuple containing: + - The export logger (or None if export API is not enabled). + - A boolean indicating if the export API is enabled for train runs. + - A boolean indicating if the export API is enabled for train run attempts. + """ + # Proto schemas should be imported within the scope of TrainStateActor to + # prevent serialization errors. + from ray.core.generated.export_event_pb2 import ExportEvent + + is_train_run_export_api_enabled = check_export_api_enabled( + ExportEvent.SourceType.EXPORT_TRAIN_RUN + ) + is_train_run_attempt_export_api_enabled = check_export_api_enabled( + ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT + ) + export_api_enabled = ( + is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled + ) + + if not export_api_enabled: + return None, False, False + + log_directory = os.path.join( + ray._private.worker._global_node.get_session_dir_path(), "logs" + ) + logger = None + try: + logger = get_export_event_logger( + EventLogType.TRAIN_STATE, + log_directory, + ) + except Exception: + logger.exception( + "Unable to initialize the export event logger, so no Train export " + "events will be written." + ) + + if logger is None: + return None, False, False + + return ( + logger, + is_train_run_export_api_enabled, + is_train_run_attempt_export_api_enabled, + ) + + def _maybe_export_train_run(self, run_info: TrainRunInfo) -> None: + if not self._is_train_run_export_api_enabled: + return + + from ray.train._internal.state.export import train_run_info_to_proto_run + + run_proto = train_run_info_to_proto_run(run_info) + self._export_logger.send_event(run_proto) + + def _maybe_export_train_run_attempt(self, run_info: TrainRunInfo) -> None: + if not self._is_train_run_attempt_export_api_enabled: + return + + from ray.train._internal.state.export import train_run_info_to_proto_attempt + + run_attempt_proto = train_run_info_to_proto_attempt(run_info) + self._export_logger.send_event(run_attempt_proto) + + +TRAIN_STATE_ACTOR_NAME = "train_state_actor" +TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor" + +_state_actor_lock: threading.RLock = threading.RLock() + + +def get_or_create_state_actor() -> ActorHandle: + """Get or create a `TrainStateActor` on the head node.""" + with _state_actor_lock: + state_actor = TrainStateActor.options( + name=TRAIN_STATE_ACTOR_NAME, + namespace=TRAIN_STATE_ACTOR_NAMESPACE, + get_if_exists=True, + lifetime="detached", + resources={"node:__internal_head__": 0.001}, + # Escape from the parent's placement group + scheduling_strategy="DEFAULT", + ).remote() + + # Ensure the state actor is ready + ray.get(state_actor.__ray_ready__.remote()) + return state_actor + + +def get_state_actor() -> Optional[ActorHandle]: + """Get the `TrainStateActor` if exists, otherwise return None.""" + try: + return ray.get_actor( + name=TRAIN_STATE_ACTOR_NAME, + namespace=TRAIN_STATE_ACTOR_NAMESPACE, + ) + except ValueError: + return None diff --git a/lib/python3.12/site-packages/ray/train/_internal/state/state_manager.py b/lib/python3.12/site-packages/ray/train/_internal/state/state_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..5684606c601ac5df77ce40e7a22d6ce86ee86998 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/state/state_manager.py @@ -0,0 +1,130 @@ +import logging +import os +from collections import defaultdict +from typing import Any, Dict, List + +import ray +from ray.data import Dataset +from ray.train._internal.state.schema import ( + ActorStatusEnum, + RunStatusEnum, + TrainDatasetInfo, + TrainRunInfo, + TrainWorkerInfo, +) +from ray.train._internal.utils import check_for_failure +from ray.train._internal.worker_group import WorkerGroup + +logger = logging.getLogger(__name__) + + +class TrainRunStateManager: + """A class that aggregates and reports train run info to TrainStateActor. + + This manager class is created on the train controller layer for each run. + """ + + def __init__(self, state_actor) -> None: + self.state_actor = state_actor + self.train_run_info_dict = defaultdict(dict) + + def register_train_run( + self, + run_id: str, + job_id: str, + run_name: str, + run_status: str, + controller_actor_id: str, + datasets: Dict[str, Dataset], + worker_group: WorkerGroup, + start_time_ms: float, + resources: List[Dict[str, float]], + status_detail: str = "", + ) -> None: + """Collect Train Run Info and report to StateActor.""" + + if not self.state_actor: + logger.warning( + "Unable to register train run since `TrainStateActor` is not started." + ) + return + + def collect_train_worker_info(): + train_context = ray.train.get_context() + core_context = ray.runtime_context.get_runtime_context() + return TrainWorkerInfo( + world_rank=train_context.get_world_rank(), + local_rank=train_context.get_local_rank(), + node_rank=train_context.get_node_rank(), + actor_id=core_context.get_actor_id(), + node_id=core_context.get_node_id(), + node_ip=ray.util.get_node_ip_address(), + gpu_ids=ray.get_gpu_ids(), + pid=os.getpid(), + resources=resources[0], + status=ActorStatusEnum.ALIVE, + ) + + futures = [ + worker_group.execute_single_async(index, collect_train_worker_info) + for index in range(len(worker_group)) + ] + success, exception = check_for_failure(futures) + + if not success: + logger.error( + "Failed to collect run information from the Ray Train " + f"workers:\n{exception}" + ) + return + + worker_info_list = ray.get(futures) + worker_info_list = sorted(worker_info_list, key=lambda info: info.world_rank) + + dataset_info_list = [ + TrainDatasetInfo( + name=ds_name, + dataset_name=ds._plan._dataset_name, + dataset_uuid=ds._plan._dataset_uuid, + ) + for ds_name, ds in datasets.items() + ] + + updates = dict( + id=run_id, + job_id=job_id, + name=run_name, + controller_actor_id=controller_actor_id, + workers=worker_info_list, + datasets=dataset_info_list, + start_time_ms=start_time_ms, + run_status=run_status, + status_detail=status_detail, + resources=resources, + ) + + # Clear the cached info to avoid registering the same run twice + self.train_run_info_dict[run_id] = {} + self._update_train_run_info(run_id, updates) + + def end_train_run( + self, + run_id: str, + run_status: RunStatusEnum, + status_detail: str, + end_time_ms: int, + ): + """Update the train run status when the training is finished.""" + updates = dict( + run_status=run_status, + status_detail=status_detail, + end_time_ms=end_time_ms, + ) + self._update_train_run_info(run_id, updates) + + def _update_train_run_info(self, run_id: str, updates: Dict[str, Any]) -> None: + """Update specific fields of a registered TrainRunInfo instance.""" + if run_id in self.train_run_info_dict: + self.train_run_info_dict[run_id].update(updates) + train_run_info = TrainRunInfo(**self.train_run_info_dict[run_id]) + ray.get(self.state_actor.register_train_run.remote(train_run_info)) diff --git a/lib/python3.12/site-packages/ray/train/_internal/storage.py b/lib/python3.12/site-packages/ray/train/_internal/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..35332963f46bb2918ac85fedee27708d6ca88793 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/storage.py @@ -0,0 +1,724 @@ +# Try import ray[train] core requirements (defined in setup.py) +# isort: off +try: + import fsspec # noqa + from fsspec.implementations.local import LocalFileSystem + +except (ImportError, ModuleNotFoundError) as e: + raise RuntimeError( + "fsspec is a required dependency of Ray Train and Ray Tune. " + "Please install with: `pip install fsspec`" + ) from e + +try: + import pyarrow + import pyarrow.fs + +except (ImportError, ModuleNotFoundError) as e: + raise RuntimeError( + "pyarrow is a required dependency of Ray Train and Ray Tune. " + "Please install with: `pip install pyarrow`" + ) from e + +try: + # check if Arrow has S3 support + from pyarrow.fs import S3FileSystem +except ImportError: + S3FileSystem = None +# isort: on + +import fnmatch +import logging +import os +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union + +from ray.air._internal.filelock import TempFileLock +from ray.train._internal.syncer import SyncConfig, Syncer, _BackgroundSyncer +from ray.train.constants import _get_ray_train_session_dir +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + from ray.train._checkpoint import Checkpoint + + +logger = logging.getLogger(__name__) + + +_VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker" + + +class _ExcludingLocalFilesystem(LocalFileSystem): + """LocalFileSystem wrapper to exclude files according to patterns. + + Args: + root_path: Root path to strip when matching with the exclude pattern. + Ex: root_path="/tmp/a/b/c", exclude=["*a*"], will exclude + /tmp/a/b/c/_a_.txt but not ALL of /tmp/a/*. + exclude: List of patterns that are applied to files returned by + ``self.find()``. If a file path matches this pattern, it will + be excluded. + + """ + + def __init__(self, root_path: Path, exclude: List[str], **kwargs): + super().__init__(**kwargs) + self._exclude = exclude + self._root_path = root_path + + @property + def fsid(self): + return "_excluding_local" + + def _should_exclude(self, path: str) -> bool: + """Return True if `path` (relative to `root_path`) matches any of the + `self._exclude` patterns.""" + path = Path(path) + relative_path = path.relative_to(self._root_path).as_posix() + match_candidates = [relative_path] + if path.is_dir(): + # Everything is in posix path format ('/') + match_candidates.append(relative_path + "/") + + for excl in self._exclude: + if any(fnmatch.fnmatch(candidate, excl) for candidate in match_candidates): + return True + return False + + def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs): + """Call parent find() and exclude from result.""" + paths = super().find( + path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, **kwargs + ) + if detail: + return { + path: out + for path, out in paths.items() + if not self._should_exclude(path) + } + else: + return [path for path in paths if not self._should_exclude(path)] + + +def _pyarrow_fs_copy_files( + source, destination, source_filesystem=None, destination_filesystem=None, **kwargs +): + if S3FileSystem and isinstance(destination_filesystem, pyarrow.fs.S3FileSystem): + # Workaround multi-threading issue with pyarrow. Note that use_threads=True + # is safe for download, just not for uploads, see: + # https://github.com/apache/arrow/issues/32372 + kwargs.setdefault("use_threads", False) + + # Use a large chunk size to speed up large checkpoint transfers. + kwargs.setdefault("chunk_size", 64 * 1024 * 1024) + + return pyarrow.fs.copy_files( + source, + destination, + source_filesystem=source_filesystem, + destination_filesystem=destination_filesystem, + **kwargs, + ) + + +# TODO(justinvyu): Add unit tests for all these utils. + + +def _delete_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str): + is_dir = _is_directory(fs, fs_path) + + try: + if is_dir: + fs.delete_dir(fs_path) + else: + fs.delete_file(fs_path) + except Exception: + logger.exception(f"Caught exception when deleting path at ({fs}, {fs_path}):") + + +def _download_from_fs_path( + fs: pyarrow.fs.FileSystem, + fs_path: str, + local_path: str, + filelock: bool = True, +): + """Downloads a directory or file from (fs, fs_path) to a local path. + + If fs_path points to a directory: + - The full directory contents are downloaded directly into `local_path`, + rather than to a subdirectory of `local_path`. + + If fs_path points to a file: + - The file is downloaded to `local_path`, which is expected to be a file path. + + If the download fails, the `local_path` contents are + cleaned up before raising, if the directory did not previously exist. + + NOTE: This method creates `local_path`'s parent directories if they do not + already exist. If the download fails, this does NOT clean up all the parent + directories that were created. + + Args: + fs: The filesystem to download from. + fs_path: The filesystem path (either a directory or a file) to download. + local_path: The local path to download to. + filelock: Whether to require a file lock before downloading, useful for + multiple downloads to the same directory that may be happening in parallel. + + Raises: + FileNotFoundError: if (fs, fs_path) doesn't exist. + """ + + _local_path = Path(local_path).resolve() + exists_before = _local_path.exists() + if _is_directory(fs=fs, fs_path=fs_path): + _local_path.mkdir(parents=True, exist_ok=True) + else: + _local_path.parent.mkdir(parents=True, exist_ok=True) + + try: + if filelock: + with TempFileLock(f"{os.path.normpath(local_path)}.lock"): + _pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs) + else: + _pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs) + except Exception as e: + # Clean up the directory if downloading was unsuccessful + if not exists_before: + shutil.rmtree(local_path, ignore_errors=True) + raise e + + +def _upload_to_fs_path( + local_path: str, + fs: pyarrow.fs.FileSystem, + fs_path: str, + exclude: Optional[List[str]] = None, +) -> None: + """Uploads a local directory or file to (fs, fs_path). + + NOTE: This will create all necessary parent directories at the destination. + + Args: + local_path: The local path to upload. + fs: The filesystem to upload to. + fs_path: The filesystem path where the dir/file will be uploaded to. + exclude: A list of filename matches to exclude from upload. This includes + all files under subdirectories as well. + This pattern will match with the relative paths of all files under + `local_path`. + Ex: ["*.png"] to exclude all .png images. + """ + + if not exclude: + # TODO(justinvyu): uploading a single file doesn't work + # (since we always create a directory at fs_path) + _create_directory(fs=fs, fs_path=fs_path) + _pyarrow_fs_copy_files(local_path, fs_path, destination_filesystem=fs) + return + + _upload_to_uri_with_exclude_fsspec( + local_path=local_path, fs=fs, fs_path=fs_path, exclude=exclude + ) + + +def _upload_to_uri_with_exclude_fsspec( + local_path: str, fs: "pyarrow.fs", fs_path: str, exclude: Optional[List[str]] +) -> None: + local_fs = _ExcludingLocalFilesystem(root_path=local_path, exclude=exclude) + handler = pyarrow.fs.FSSpecHandler(local_fs) + source_fs = pyarrow.fs.PyFileSystem(handler) + + _create_directory(fs=fs, fs_path=fs_path) + _pyarrow_fs_copy_files( + local_path, fs_path, source_filesystem=source_fs, destination_filesystem=fs + ) + + +def _list_at_fs_path( + fs: pyarrow.fs.FileSystem, + fs_path: str, + file_filter: Optional[Callable[[pyarrow.fs.FileInfo], bool]] = None, +) -> List[str]: + """Returns the list of filenames at (fs, fs_path), similar to os.listdir. + + If the path doesn't exist, returns an empty list. + """ + if file_filter is None: + file_filter = lambda x: True # noqa: E731 + + selector = pyarrow.fs.FileSelector(fs_path, allow_not_found=True, recursive=False) + return [ + os.path.relpath(file_info.path.lstrip("/"), start=fs_path.lstrip("/")) + for file_info in fs.get_file_info(selector) + if file_filter(file_info) + ] + + +def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool: + """Returns True if (fs, fs_path) exists.""" + + valid = fs.get_file_info(fs_path) + return valid.type != pyarrow.fs.FileType.NotFound + + +def _is_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool: + """Checks if (fs, fs_path) is a directory or a file. + + Raises: + FileNotFoundError: if (fs, fs_path) doesn't exist. + """ + + file_info = fs.get_file_info(fs_path) + if file_info.type == pyarrow.fs.FileType.NotFound: + raise FileNotFoundError(f"Path not found: ({fs}, {fs_path})") + + return not file_info.is_file + + +def _create_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> None: + """Create directory at (fs, fs_path). + + Some external filesystems require directories to already exist, or at least + the `netloc` to be created (e.g. PyArrows ``mock://`` filesystem). + + Generally this should be done before and outside of Ray applications. This + utility is thus primarily used in testing, e.g. of ``mock://` URIs. + """ + try: + fs.create_dir(fs_path) + except Exception: + logger.exception( + f"Caught exception when creating directory at ({fs}, {fs_path}):" + ) + + +def get_fs_and_path( + storage_path: Union[str, os.PathLike], + storage_filesystem: Optional[pyarrow.fs.FileSystem] = None, +) -> Tuple[pyarrow.fs.FileSystem, str]: + """Returns the fs and path from a storage path and an optional custom fs. + + Args: + storage_path: A storage path or URI. (ex: s3://bucket/path or /tmp/ray_results) + storage_filesystem: A custom filesystem to use. If not provided, + this will be auto-resolved by pyarrow. If provided, the storage_path + is assumed to be prefix-stripped already, and must be a valid path + on the filesystem. + """ + storage_path = str(storage_path) + + if storage_filesystem: + return storage_filesystem, storage_path + + return pyarrow.fs.FileSystem.from_uri(storage_path) + + +class _FilesystemSyncer(_BackgroundSyncer): + """Syncer between local filesystem and a `storage_filesystem`.""" + + def __init__(self, storage_filesystem: Optional["pyarrow.fs.FileSystem"], **kwargs): + self.storage_filesystem = storage_filesystem + super().__init__(**kwargs) + + def _sync_up_command( + self, local_path: str, uri: str, exclude: Optional[List] = None + ) -> Tuple[Callable, Dict]: + # TODO(justinvyu): Defer this cleanup up as part of the + # external-facing Syncer deprecation. + fs_path = uri + return ( + _upload_to_fs_path, + dict( + local_path=local_path, + fs=self.storage_filesystem, + fs_path=fs_path, + exclude=exclude, + ), + ) + + def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]: + fs_path = uri + return ( + _download_from_fs_path, + dict( + fs=self.storage_filesystem, + fs_path=fs_path, + local_path=local_path, + ), + ) + + def _delete_command(self, uri: str) -> Tuple[Callable, Dict]: + fs_path = uri + return _delete_fs_path, dict(fs=self.storage_filesystem, fs_path=fs_path) + + +@DeveloperAPI +class StorageContext: + """Shared context that holds the source of truth for all paths and + storage utilities, passed along from the driver to workers. + + This object defines a few types of paths: + 1. *_fs_path: A path on the `storage_filesystem`. This is a regular path + which has been prefix-stripped by pyarrow.fs.FileSystem.from_uri and + can be joined with `Path(...).as_posix()`. + 2. *_driver_staging_path: The temporary staging directory on the local filesystem + where driver artifacts are saved to before persisting them to storage. + 3. trial_working_directory: The local filesystem path that the remote + actors' working directories are moved to by default. + This is separated from the driver staging path so that driver syncing + does not implicitly upload the trial working directory, for trials on the + driver node. + + Example with storage_path="mock:///bucket/path?param=1": + + >>> import ray + >>> from ray.train._internal.storage import StorageContext + >>> import os + >>> _ = ray.init() + >>> storage = StorageContext( + ... storage_path="mock://netloc/bucket/path?param=1", + ... experiment_dir_name="exp_name", + ... ) + >>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS + >> storage.experiment_fs_path + 'bucket/path/exp_name' + >>> storage.experiment_driver_staging_path # doctest: +ELLIPSIS + '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts' + >>> storage.trial_dir_name = "trial_dir" + >>> storage.trial_fs_path + 'bucket/path/exp_name/trial_dir' + >>> storage.trial_driver_staging_path # doctest: +ELLIPSIS + '/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts/trial_dir' + >>> storage.trial_working_directory # doctest: +ELLIPSIS + '/tmp/ray/session_.../artifacts/.../exp_name/working_dirs/trial_dir' + >>> storage.current_checkpoint_index = 1 + >>> storage.checkpoint_fs_path + 'bucket/path/exp_name/trial_dir/checkpoint_000001' + >>> ray.shutdown() + + Example with storage_path="/tmp/ray_results": + + >>> from ray.train._internal.storage import StorageContext + >>> storage = StorageContext( + ... storage_path="/tmp/ray_results", + ... experiment_dir_name="exp_name", + ... ) + >>> storage.storage_fs_path + '/tmp/ray_results' + >>> storage.experiment_fs_path + '/tmp/ray_results/exp_name' + >>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS + " + ) + + def _create_validation_file(self): + """On the creation of a storage context, create a validation file at the + storage path to verify that the storage path can be written to. + This validation file is also used to check whether the storage path is + accessible by all nodes in the cluster.""" + valid_file = Path( + self.experiment_fs_path, _VALIDATE_STORAGE_MARKER_FILENAME + ).as_posix() + self.storage_filesystem.create_dir(self.experiment_fs_path) + with self.storage_filesystem.open_output_stream(valid_file): + pass + + def _check_validation_file(self): + """Checks that the validation file exists at the storage path.""" + valid_file = Path( + self.experiment_fs_path, _VALIDATE_STORAGE_MARKER_FILENAME + ).as_posix() + if not _exists_at_fs_path(fs=self.storage_filesystem, fs_path=valid_file): + raise RuntimeError( + f"Unable to set up cluster storage with the following settings:\n{self}" + "\nCheck that all nodes in the cluster have read/write access " + "to the configured storage path. `RunConfig(storage_path)` should be " + "set to a cloud storage URI or a shared filesystem path accessible " + "by all nodes in your cluster ('s3://bucket' or '/mnt/nfs'). " + "A local path on the head node is not accessible by worker nodes. " + "See: https://docs.ray.io/en/latest/train/user-guides/persistent-storage.html" # noqa: E501 + ) + + def _update_checkpoint_index(self, metrics: Dict): + # Per default, increase by 1. This can be overwritten to customize checkpoint + # directories. + self.current_checkpoint_index += 1 + + def persist_current_checkpoint(self, checkpoint: "Checkpoint") -> "Checkpoint": + """Persists a given checkpoint to the current checkpoint path on the filesystem. + + "Current" is defined by the `current_checkpoint_index` attribute of the + storage context. + + This method copies the checkpoint files to the storage location. + It's up to the user to delete the original checkpoint files if desired. + + For example, the original directory is typically a local temp directory. + + Args: + checkpoint: The checkpoint to persist to (fs, checkpoint_fs_path). + + Returns: + Checkpoint: A Checkpoint pointing to the persisted checkpoint location. + """ + # TODO(justinvyu): Fix this cyclical import. + + logger.debug( + "Copying checkpoint files to storage path:\n" + "({source_fs}, {source}) -> ({dest_fs}, {destination})".format( + source=checkpoint.path, + destination=self.checkpoint_fs_path, + source_fs=checkpoint.filesystem, + dest_fs=self.storage_filesystem, + ) + ) + + # Raise an error if the storage path is not accessible when + # attempting to upload a checkpoint from a remote worker. + # Ex: If storage_path is a local path, then a validation marker + # will only exist on the head node but not the worker nodes. + self._check_validation_file() + + self.storage_filesystem.create_dir(self.checkpoint_fs_path) + _pyarrow_fs_copy_files( + source=checkpoint.path, + destination=self.checkpoint_fs_path, + source_filesystem=checkpoint.filesystem, + destination_filesystem=self.storage_filesystem, + ) + + persisted_checkpoint = checkpoint.__class__( + filesystem=self.storage_filesystem, + path=self.checkpoint_fs_path, + ) + logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}") + return persisted_checkpoint + + def persist_artifacts(self, force: bool = False) -> None: + """Persists all artifacts within `trial_local_dir` to storage. + + This method possibly launches a background task to sync the trial dir, + depending on the `sync_period` + `sync_artifacts_on_checkpoint` + settings of `SyncConfig`. + + `(local_fs, trial_working_dir) -> (storage_filesystem, trial_fs_path)` + + Args: + force: If True, wait for a previous sync to finish, launch a new one, + and wait for that one to finish. By the end of a `force=True` call, the + latest version of the trial artifacts will be persisted. + """ + if not self.sync_config.sync_artifacts: + return + + # Skip if there are no artifacts to sync + is_empty = not any(os.scandir(self.trial_working_directory)) + if is_empty: + return + + if force: + self.syncer.wait() + self.syncer.sync_up( + local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path + ) + self.syncer.wait() + else: + self.syncer.sync_up_if_needed( + local_dir=self.trial_working_directory, remote_dir=self.trial_fs_path + ) + + @property + def experiment_fs_path(self) -> str: + """The path on the `storage_filesystem` to the experiment directory. + + NOTE: This does not have a URI prefix anymore, since it has been stripped + by pyarrow.fs.FileSystem.from_uri already. The URI scheme information is + kept in `storage_filesystem` instead. + """ + return Path(self.storage_fs_path, self.experiment_dir_name).as_posix() + + def _get_session_path(self) -> str: + """The Ray Train/Tune session local directory used to stage files + before persisting to the storage filesystem.""" + return Path( + _get_ray_train_session_dir(), self._timestamp, self.experiment_dir_name + ).as_posix() + + @property + def experiment_driver_staging_path(self) -> str: + """The local filesystem path of the experiment directory on the driver node. + + The driver is the node where `Trainer.fit`/`Tuner.fit` is being called. + + This path is of the form: + `/tmp/ray/session_/artifacts// + /driver_artifacts` + + This should be used as the temporary staging location for files *on the driver* + before syncing them to `experiment_fs_path`. + For example, the search algorithm should dump its state to this directory. + See `trial_driver_staging_path` for writing trial-specific artifacts. + + The directory is synced to + `{storage_path}/{experiment_dir_name}` periodically. + See `_ExperimentCheckpointManager.checkpoint` for where that happens. + """ + return Path(self._get_session_path(), "driver_artifacts").as_posix() + + @property + def trial_fs_path(self) -> str: + """The trial directory path on the `storage_filesystem`. + + Raises a ValueError if `trial_dir_name` is not set beforehand. + """ + if self.trial_dir_name is None: + raise RuntimeError( + "Should not access `trial_fs_path` without setting `trial_dir_name`" + ) + return Path(self.experiment_fs_path, self.trial_dir_name).as_posix() + + @property + def trial_driver_staging_path(self) -> str: + """The local filesystem path of the trial directory on the driver. + + The driver is the node where `Trainer.fit`/`Tuner.fit` is being called. + + This path is of the form: + `/tmp/ray/session_/artifacts// + /driver_artifacts/` + + This should be used as the temporary location for files on the driver + before persisting them to `trial_fs_path`. + + For example, callbacks (e.g., JsonLoggerCallback) should write trial-specific + logfiles within this directory. + """ + if self.trial_dir_name is None: + raise RuntimeError( + "Should not access `trial_driver_staging_path` " + "without setting `trial_dir_name`" + ) + return Path(self.experiment_driver_staging_path, self.trial_dir_name).as_posix() + + @property + def trial_working_directory(self) -> str: + """The local filesystem path to trial working directory. + + This path is of the form: + `/tmp/ray/session_/artifacts// + /working_dirs/` + + Ray Train/Tune moves the remote actor's working directory to this path + by default, unless disabled by `RAY_CHDIR_TO_TRIAL_DIR` environment variable. + + Writing files to this directory allows users to persist training artifacts + if `SyncConfig(sync_artifacts=True)` is set. + """ + if self.trial_dir_name is None: + raise RuntimeError( + "Cannot access `trial_working_directory` without " + "setting `trial_dir_name`" + ) + return Path( + self._get_session_path(), "working_dirs", self.trial_dir_name + ).as_posix() + + @property + def checkpoint_fs_path(self) -> str: + """The current checkpoint directory path on the `storage_filesystem`. + + "Current" refers to the checkpoint that is currently being created/persisted. + The user of this class is responsible for setting the `current_checkpoint_index` + (e.g., incrementing when needed). + """ + return Path(self.trial_fs_path, self.checkpoint_dir_name).as_posix() + + @property + def checkpoint_dir_name(self) -> str: + """The current checkpoint directory name, based on the checkpoint index.""" + return StorageContext._make_checkpoint_dir_name(self.current_checkpoint_index) + + @staticmethod + def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str: + from ray.tune.experiment import Experiment + from ray.tune.utils import date_str + + run_identifier = Experiment.get_trainable_name(run_obj) + + if bool(int(os.environ.get("TUNE_DISABLE_DATED_SUBDIR", 0))): + dir_name = run_identifier + else: + dir_name = "{}_{}".format(run_identifier, date_str()) + return dir_name + + @staticmethod + def _make_checkpoint_dir_name(index: int): + """Get the name of the checkpoint directory, given an index.""" + return f"checkpoint_{index:06d}" diff --git a/lib/python3.12/site-packages/ray/train/_internal/syncer.py b/lib/python3.12/site-packages/ray/train/_internal/syncer.py new file mode 100644 index 0000000000000000000000000000000000000000..8321552e5fbaa3e13bea6648df5ee960079dc0a1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/syncer.py @@ -0,0 +1,423 @@ +import abc +import logging +import threading +import time +import traceback +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +from ray._private.thirdparty.tabulate.tabulate import tabulate +from ray.util.annotations import Deprecated, DeveloperAPI +from ray.widgets import Template + +logger = logging.getLogger(__name__) + +# Syncing period for syncing checkpoints between nodes or to cloud. +DEFAULT_SYNC_PERIOD = 300 + +# Default sync timeout after which syncing processes are aborted +DEFAULT_SYNC_TIMEOUT = 1800 + + +@Deprecated +@dataclass +class SyncConfig: + sync_period: int = DEFAULT_SYNC_PERIOD + sync_timeout: int = DEFAULT_SYNC_TIMEOUT + sync_artifacts: bool = False + sync_artifacts_on_checkpoint: bool = True + + def _repr_html_(self) -> str: + """Generate an HTML representation of the SyncConfig.""" + return Template("scrollableTable.html.j2").render( + table=tabulate( + { + "Setting": ["Sync period", "Sync timeout"], + "Value": [self.sync_period, self.sync_timeout], + }, + tablefmt="html", + showindex=False, + headers="keys", + ), + max_height="none", + ) + + +class _BackgroundProcess: + def __init__(self, fn: Callable): + self._fn = fn + self._process = None + self._result = {} + self._start_time = float("-inf") + + @property + def is_running(self): + return self._process and self._process.is_alive() + + @property + def start_time(self): + return self._start_time + + def start(self, *args, **kwargs): + if self.is_running: + return False + + self._result = {} + + def entrypoint(): + try: + result = self._fn(*args, **kwargs) + except Exception as e: + self._result["exception"] = e + return + + self._result["result"] = result + + self._process = threading.Thread(target=entrypoint) + self._process.daemon = True + self._process.start() + self._start_time = time.time() + + def wait(self, timeout: Optional[float] = None) -> Any: + """Waits for the background process to finish running. Waits until the + background process has run for at least `timeout` seconds, counting from + the time when the process was started.""" + if not self._process: + return None + + time_remaining = None + if timeout: + elapsed = time.time() - self.start_time + time_remaining = max(timeout - elapsed, 0) + + self._process.join(timeout=time_remaining) + + if self._process.is_alive(): + self._process = None + raise TimeoutError( + f"{getattr(self._fn, '__name__', str(self._fn))} did not finish " + f"running within the timeout of {timeout} seconds." + ) + + self._process = None + + exception = self._result.get("exception") + if exception: + raise exception + + result = self._result.get("result") + + self._result = {} + return result + + +@DeveloperAPI +class Syncer(abc.ABC): + """Syncer class for synchronizing data between Ray nodes and remote (cloud) storage. + + This class handles data transfer for two cases: + + 1. Synchronizing data such as experiment state snapshots from the driver to + cloud storage. + 2. Synchronizing data such as trial checkpoints from remote trainables to + cloud storage. + + Synchronizing tasks are usually asynchronous and can be awaited using ``wait()``. + The base class implements a ``wait_or_retry()`` API that will retry a failed + sync command. + + The base class also exposes an API to only kick off syncs every ``sync_period`` + seconds. + + Args: + sync_period: The minimum time in seconds between sync operations, as + used by ``sync_up/down_if_needed``. + sync_timeout: The maximum time to wait for a sync process to finish before + issuing a new sync operation. Ex: should be used by ``wait`` if launching + asynchronous sync tasks. + """ + + def __init__( + self, + sync_period: float = DEFAULT_SYNC_PERIOD, + sync_timeout: float = DEFAULT_SYNC_TIMEOUT, + ): + self.sync_period = sync_period + self.sync_timeout = sync_timeout + self.last_sync_up_time = float("-inf") + self.last_sync_down_time = float("-inf") + + @abc.abstractmethod + def sync_up( + self, local_dir: str, remote_dir: str, exclude: Optional[List] = None + ) -> bool: + """Synchronize local directory to remote directory. + + This function can spawn an asynchronous process that can be awaited in + ``wait()``. + + Args: + local_dir: Local directory to sync from. + remote_dir: Remote directory to sync up to. This is an URI + (``protocol://remote/path``). + exclude: Pattern of files to exclude, e.g. + ``["*/checkpoint_*]`` to exclude trial checkpoints. + + Returns: + True if sync process has been spawned, False otherwise. + + """ + raise NotImplementedError + + @abc.abstractmethod + def sync_down( + self, remote_dir: str, local_dir: str, exclude: Optional[List] = None + ) -> bool: + """Synchronize remote directory to local directory. + + This function can spawn an asynchronous process that can be awaited in + ``wait()``. + + Args: + remote_dir: Remote directory to sync down from. This is an URI + (``protocol://remote/path``). + local_dir: Local directory to sync to. + exclude: Pattern of files to exclude, e.g. + ``["*/checkpoint_*]`` to exclude trial checkpoints. + + Returns: + True if sync process has been spawned, False otherwise. + + """ + raise NotImplementedError + + @abc.abstractmethod + def delete(self, remote_dir: str) -> bool: + """Delete directory on remote storage. + + This function can spawn an asynchronous process that can be awaited in + ``wait()``. + + Args: + remote_dir: Remote directory to delete. This is an URI + (``protocol://remote/path``). + + Returns: + True if sync process has been spawned, False otherwise. + + """ + raise NotImplementedError + + def retry(self): + """Retry the last sync up, sync down, or delete command. + + You should implement this method if you spawn asynchronous syncing + processes. + """ + pass + + def wait(self, timeout: Optional[float] = None): + """Wait for asynchronous sync command to finish. + + You should implement this method if you spawn asynchronous syncing + processes. This method should timeout after the asynchronous command + has run for `sync_timeout` seconds and raise a `TimeoutError`. + """ + pass + + def sync_up_if_needed( + self, local_dir: str, remote_dir: str, exclude: Optional[List] = None + ) -> bool: + """Syncs up if time since last sync up is greater than sync_period. + + Args: + local_dir: Local directory to sync from. + remote_dir: Remote directory to sync up to. This is an URI + (``protocol://remote/path``). + exclude: Pattern of files to exclude, e.g. + ``["*/checkpoint_*]`` to exclude trial checkpoints. + """ + now = time.time() + if now - self.last_sync_up_time >= self.sync_period: + result = self.sync_up( + local_dir=local_dir, remote_dir=remote_dir, exclude=exclude + ) + self.last_sync_up_time = now + return result + + def sync_down_if_needed( + self, remote_dir: str, local_dir: str, exclude: Optional[List] = None + ): + """Syncs down if time since last sync down is greater than sync_period. + + Args: + remote_dir: Remote directory to sync down from. This is an URI + (``protocol://remote/path``). + local_dir: Local directory to sync to. + exclude: Pattern of files to exclude, e.g. + ``["*/checkpoint_*]`` to exclude trial checkpoints. + """ + now = time.time() + if now - self.last_sync_down_time >= self.sync_period: + result = self.sync_down( + remote_dir=remote_dir, local_dir=local_dir, exclude=exclude + ) + self.last_sync_down_time = now + return result + + def wait_or_retry(self, max_retries: int = 2, backoff_s: int = 5): + assert max_retries > 0 + last_error_traceback = None + for i in range(max_retries + 1): + try: + self.wait() + except Exception as e: + attempts_remaining = max_retries - i + + # If we're out of retries, then save the full traceback of the last + # error and show it when raising an exception. + if attempts_remaining == 0: + last_error_traceback = traceback.format_exc() + break + + logger.error( + f"The latest sync operation failed with the following error: " + f"{repr(e)}\n" + f"Retrying {attempts_remaining} more time(s) after sleeping " + f"for {backoff_s} seconds..." + ) + time.sleep(backoff_s) + self.retry() + continue + # Succeeded! + return + raise RuntimeError( + f"Failed sync even after {max_retries} retries. " + f"The latest sync failed with the following error:\n{last_error_traceback}" + ) + + def reset(self): + self.last_sync_up_time = float("-inf") + self.last_sync_down_time = float("-inf") + + def close(self): + pass + + def _repr_html_(self) -> str: + return + + +class _BackgroundSyncer(Syncer): + """Syncer using a background process for asynchronous file transfer.""" + + def __init__( + self, + sync_period: float = DEFAULT_SYNC_PERIOD, + sync_timeout: float = DEFAULT_SYNC_TIMEOUT, + ): + super(_BackgroundSyncer, self).__init__( + sync_period=sync_period, sync_timeout=sync_timeout + ) + self._sync_process = None + self._current_cmd = None + + def _should_continue_existing_sync(self): + """Returns whether a previous sync is still running within the timeout.""" + return ( + self._sync_process + and self._sync_process.is_running + and time.time() - self._sync_process.start_time < self.sync_timeout + ) + + def _launch_sync_process(self, sync_command: Tuple[Callable, Dict]): + """Waits for the previous sync process to finish, + then launches a new process that runs the given command.""" + if self._sync_process: + try: + self.wait() + except Exception: + logger.warning( + f"Last sync command failed with the following error:\n" + f"{traceback.format_exc()}" + ) + + self._current_cmd = sync_command + self.retry() + + def sync_up( + self, local_dir: str, remote_dir: str, exclude: Optional[List] = None + ) -> bool: + if self._should_continue_existing_sync(): + logger.debug( + f"Last sync still in progress, " + f"skipping sync up of {local_dir} to {remote_dir}" + ) + return False + + sync_up_cmd = self._sync_up_command( + local_path=local_dir, uri=remote_dir, exclude=exclude + ) + self._launch_sync_process(sync_up_cmd) + + return True + + def _sync_up_command( + self, local_path: str, uri: str, exclude: Optional[List] = None + ) -> Tuple[Callable, Dict]: + raise NotImplementedError + + def sync_down( + self, remote_dir: str, local_dir: str, exclude: Optional[List] = None + ) -> bool: + if self._should_continue_existing_sync(): + logger.warning( + f"Last sync still in progress, " + f"skipping sync down of {remote_dir} to {local_dir}" + ) + return False + + sync_down_cmd = self._sync_down_command(uri=remote_dir, local_path=local_dir) + self._launch_sync_process(sync_down_cmd) + + return True + + def _sync_down_command(self, uri: str, local_path: str) -> Tuple[Callable, Dict]: + raise NotImplementedError + + def delete(self, remote_dir: str) -> bool: + if self._should_continue_existing_sync(): + logger.warning( + f"Last sync still in progress, skipping deletion of {remote_dir}" + ) + return False + + delete_cmd = self._delete_command(uri=remote_dir) + self._launch_sync_process(delete_cmd) + + return True + + def _delete_command(self, uri: str) -> Tuple[Callable, Dict]: + raise NotImplementedError + + def wait(self, timeout: Optional[float] = None): + if self._sync_process: + try: + self._sync_process.wait(timeout=timeout or self.sync_timeout) + except Exception as e: + raise e + finally: + # Regardless of whether the sync process succeeded within the timeout, + # clear the sync process so a new one can be created. + self._sync_process = None + + def retry(self): + if not self._current_cmd: + raise RuntimeError("No sync command set, cannot retry.") + cmd, kwargs = self._current_cmd + self._sync_process = _BackgroundProcess(cmd) + self._sync_process.start(**kwargs) + + def __getstate__(self): + state = self.__dict__.copy() + state["_sync_process"] = None + return state diff --git a/lib/python3.12/site-packages/ray/train/_internal/utils.py b/lib/python3.12/site-packages/ray/train/_internal/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0524bb5623e7cdd94b7bf2310f8ee4f792149ca2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/utils.py @@ -0,0 +1,223 @@ +import abc +import functools +import inspect +import logging +import os +import socket +from typing import ( + Any, + Callable, + ContextManager, + Dict, + List, + Optional, + Tuple, + TypeVar, + Union, +) + +import ray +from ray._common.network_utils import find_free_port, is_ipv6 +from ray.actor import ActorHandle +from ray.air._internal.util import ( + StartTraceback, + StartTracebackWithWorkerRank, +) +from ray.exceptions import RayActorError +from ray.types import ObjectRef + +T = TypeVar("T") + +logger = logging.getLogger(__name__) + + +def check_for_failure( + remote_values: List[ObjectRef], +) -> Tuple[bool, Optional[Exception]]: + """Check for actor failure when retrieving the remote values. + + Args: + remote_values: List of object references from Ray actor methods. + + Returns: + A tuple of (bool, Exception). The bool is + True if evaluating all object references is successful, False otherwise. + """ + unfinished = remote_values.copy() + + while len(unfinished) > 0: + finished, unfinished = ray.wait(unfinished) + + # If a failure occurs the ObjectRef will be marked as finished. + # Calling ray.get will expose the failure as a RayActorError. + for object_ref in finished: + # Everything in finished has either failed or completed + # successfully. + try: + ray.get(object_ref) + except RayActorError as exc: + failed_actor_rank = remote_values.index(object_ref) + logger.info(f"Worker {failed_actor_rank} has failed.") + return False, exc + except Exception as exc: + # Other (e.g. training) errors should be directly raised + failed_worker_rank = remote_values.index(object_ref) + raise StartTracebackWithWorkerRank( + worker_rank=failed_worker_rank + ) from exc + + return True, None + + +def get_address_and_port() -> Tuple[str, int]: + """Returns the IP address and a free port on this node.""" + addr = ray.util.get_node_ip_address() + port = find_free_port(socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET) + return addr, port + + +def update_env_vars(env_vars: Dict[str, Any]): + """Updates the environment variables on this worker process. + + Args: + env_vars: Environment variables to set. + """ + sanitized = {k: str(v) for k, v in env_vars.items()} + os.environ.update(sanitized) + + +def count_required_parameters(fn: Callable) -> int: + """Counts the number of required parameters of a function. + + NOTE: *args counts as 1 required parameter. + + Examples + -------- + + >>> def fn(a, b, /, c, *args, d=1, e=2, **kwargs): + ... pass + >>> count_required_parameters(fn) + 4 + + >>> fn = lambda: 1 + >>> count_required_parameters(fn) + 0 + + >>> def fn(config, a, b=1, c=2): + ... pass + >>> from functools import partial + >>> count_required_parameters(partial(fn, a=0)) + 1 + """ + params = inspect.signature(fn).parameters.values() + + positional_param_kinds = { + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.VAR_POSITIONAL, + } + return len( + [ + p + for p in params + if p.default == inspect.Parameter.empty and p.kind in positional_param_kinds + ] + ) + + +def construct_train_func( + train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]], + config: Optional[Dict[str, Any]], + train_func_context: ContextManager, + fn_arg_name: Optional[str] = "train_func", + discard_returns: bool = False, +) -> Callable[[], T]: + """Validates and constructs the training function to execute. + Args: + train_func: The training function to execute. + This can either take in no arguments or a ``config`` dict. + config (Optional[Dict]): Configurations to pass into + ``train_func``. If None then an empty Dict will be created. + train_func_context: Context manager for user's `train_func`, which executes + backend-specific logic before and after the training function. + fn_arg_name (Optional[str]): The name of training function to use for error + messages. + discard_returns: Whether to discard any returns from train_func or not. + Returns: + A valid training function. + Raises: + ValueError: if the input ``train_func`` is invalid. + """ + num_required_params = count_required_parameters(train_func) + + if discard_returns: + # Discard any returns from the function so that + # BackendExecutor doesn't try to deserialize them. + # Those returns are inaccesible with AIR anyway. + @functools.wraps(train_func) + def discard_return_wrapper(*args, **kwargs): + try: + train_func(*args, **kwargs) + except Exception as e: + raise StartTraceback from e + + wrapped_train_func = discard_return_wrapper + else: + wrapped_train_func = train_func + + if num_required_params > 1: + err_msg = ( + f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts " + f"{num_required_params} required arguments instead." + ) + raise ValueError(err_msg) + elif num_required_params == 1: + config = {} if config is None else config + + @functools.wraps(wrapped_train_func) + def train_fn(): + try: + with train_func_context(): + return wrapped_train_func(config) + except Exception as e: + raise StartTraceback from e + + else: # num_params == 0 + + @functools.wraps(wrapped_train_func) + def train_fn(): + try: + with train_func_context(): + return wrapped_train_func() + except Exception as e: + raise StartTraceback from e + + return train_fn + + +class Singleton(abc.ABCMeta): + """Singleton Abstract Base Class + + https://stackoverflow.com/questions/33364070/implementing + -singleton-as-metaclass-but-for-abstract-classes + """ + + _instances = {} + + def __call__(cls, *args, **kwargs): + if cls not in cls._instances: + cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs) + return cls._instances[cls] + + +class ActorWrapper: + """Wraps an actor to provide same API as using the base class directly.""" + + def __init__(self, actor: ActorHandle): + self.actor = actor + + def __getattr__(self, item): + # The below will fail if trying to access an attribute (not a method) from the + # actor. + actor_method = getattr(self.actor, item) + return lambda *args, **kwargs: ray.get(actor_method.remote(*args, **kwargs)) diff --git a/lib/python3.12/site-packages/ray/train/_internal/worker_group.py b/lib/python3.12/site-packages/ray/train/_internal/worker_group.py new file mode 100644 index 0000000000000000000000000000000000000000..8a992898475db544c3ca5f2abb375b47055d72e6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/_internal/worker_group.py @@ -0,0 +1,437 @@ +import copy +import logging +import os +import socket +from collections import defaultdict +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union + +import ray +from ray.actor import ActorHandle +from ray.air._internal.util import exception_cause, skip_exceptions +from ray.train._internal.base_worker_group import BaseWorkerGroup +from ray.types import ObjectRef +from ray.util.placement_group import PlacementGroup + +T = TypeVar("T") + +logger = logging.getLogger(__name__) + + +class RayTrainWorker: + """A class to execute arbitrary functions. Does not hold any state.""" + + def __execute(self, func: Callable[..., T], *args, **kwargs) -> T: + """Executes the input function and returns the output. + + Args: + func: The function to execute. + args, kwargs: The arguments to pass into func. + """ + try: + return func(*args, **kwargs) + except Exception as e: + skipped = skip_exceptions(e) + raise skipped from exception_cause(skipped) + + +@dataclass +class WorkerMetadata: + """Metadata for each worker/actor. + + This information is expected to stay the same throughout the lifetime of + actor. + + Args: + node_id: ID of the node this worker is on. + node_ip: IP address of the node this worker is on. + hostname: Hostname that this worker is on. + resource_ids: Map of accelerator resources + ("GPU", "neuron_cores", ..) to their IDs. + pid: Process ID of this worker. + """ + + node_id: str + node_ip: str + hostname: str + resource_ids: Dict[str, List[str]] + pid: int + + +@dataclass +class Worker: + """Class representing a Worker.""" + + actor: ActorHandle + metadata: WorkerMetadata + + +def create_executable_class(executable_cls: Optional[Type] = None) -> Type: + """Create the executable class to use as the Ray actors.""" + if not executable_cls: + return RayTrainWorker + elif issubclass(executable_cls, RayTrainWorker): + return executable_cls + else: + + class _WrappedExecutable(executable_cls, RayTrainWorker): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + return _WrappedExecutable + + +def construct_metadata() -> WorkerMetadata: + """Creates metadata for this worker. + + This function is expected to be run on the actor. + """ + node_id = ray.get_runtime_context().get_node_id() + node_ip = ray.util.get_node_ip_address() + hostname = socket.gethostname() + accelerator_ids = ray.get_runtime_context().get_accelerator_ids() + pid = os.getpid() + + return WorkerMetadata( + node_id=node_id, + node_ip=node_ip, + hostname=hostname, + resource_ids=accelerator_ids, + pid=pid, + ) + + +class WorkerGroup(BaseWorkerGroup): + """Group of Ray Actors that can execute arbitrary functions. + + ``WorkerGroup`` launches Ray actors according to the given + specification. It can then execute arbitrary Python functions in each of + these workers. + + If not enough resources are available to launch the actors, the Ray + cluster will automatically scale up if autoscaling is enabled. + + Args: + num_workers: The number of workers (Ray actors) to launch. + Defaults to 1. + resources_per_worker (Optional[Dict[str, float]]): + Dictionary specifying the resources that will be + requested for each worker. Defaults to {"CPU": 1}. + actor_cls (Optional[Type]): If specified use this class as the + remote actors. + remote_cls_args, remote_cls_kwargs: If ``remote_cls`` is provided, + these args will be used for the worker initialization. + placement_group (PlacementGroup|str): The placement group that workers + should be created in. Defaults to "default" which will inherit the + parent placement group (if child tasks should be captured). + + + Example: + + .. code_block:: python + + worker_group = WorkerGroup(num_workers=2) + output = worker_group.execute(lambda: 1) + assert len(output) == 2 + assert all(o == 1 for o in output) + """ + + def __init__( + self, + num_workers: int = 1, + resources_per_worker: Optional[Dict[str, float]] = None, + actor_cls: Type = None, + actor_cls_args: Optional[Tuple] = None, + actor_cls_kwargs: Optional[Dict] = None, + placement_group: Union[PlacementGroup, str] = "default", + ): + if resources_per_worker is None: + resources_per_worker = {"CPU": 1} + else: + resources_per_worker = resources_per_worker.copy() + + if num_workers <= 0: + raise ValueError( + "The provided `num_workers` must be greater " + f"than 0. Received num_workers={num_workers} " + f"instead." + ) + + if any(v < 0 for v in resources_per_worker.values()): + raise ValueError( + "The number of resources per worker must not be negative. " + f"Received resources_per_worker={resources_per_worker}." + ) + + if (actor_cls_args or actor_cls_kwargs) and not actor_cls: + raise ValueError( + "`actor_cls_args` or `actor_class_kwargs` are " + "passed in but no `actor_cls` is passed in." + ) + + self.num_workers = num_workers + self.resources_per_worker = resources_per_worker + + _resources_per_worker = copy.deepcopy(resources_per_worker) + self.num_cpus_per_worker = _resources_per_worker.pop("CPU", 0) + self.num_gpus_per_worker = _resources_per_worker.pop("GPU", 0) + self.memory_per_worker = _resources_per_worker.pop("memory", 0) + self.workers = [] + self._base_cls = create_executable_class(actor_cls) + assert issubclass(self._base_cls, RayTrainWorker) + + self._actor_cls_args = actor_cls_args or [] + self._actor_cls_kwargs = actor_cls_kwargs or {} + + self._placement_group = placement_group + + # TODO(matt): Validate resources. Fast-fail if it is impossible to + # handle the request, rather than hang indefinitely. + self._remote_cls = ray.remote( + num_cpus=self.num_cpus_per_worker, + num_gpus=self.num_gpus_per_worker, + memory=self.memory_per_worker, + resources=_resources_per_worker, + )(self._base_cls) + self.start() + + def start(self): + """Starts all the workers in this worker group.""" + if self.workers and len(self.workers) > 0: + raise RuntimeError( + "The workers have already been started. " + "Please call `shutdown` first if you want to " + "restart them." + ) + + logger.debug(f"Starting {self.num_workers} workers.") + self.add_workers(self.num_workers) + logger.debug(f"{len(self.workers)} workers have successfully started.") + + def shutdown(self, patience_s: float = 5): + """Shutdown all the workers in this worker group. + + Args: + patience_s: Attempt a graceful shutdown + of the workers for this many seconds. Fallback to force kill + if graceful shutdown is not complete after this time. If + this is less than or equal to 0, immediately force kill all + workers. + """ + logger.debug(f"Shutting down {len(self.workers)} workers.") + if patience_s <= 0: + for worker in self.workers: + ray.kill(worker.actor) + else: + done_refs = [w.actor.__ray_terminate__.remote() for w in self.workers] + # Wait for actors to die gracefully. + done, not_done = ray.wait(done_refs, timeout=patience_s) + if not_done: + logger.debug("Graceful termination failed. Falling back to force kill.") + # If all actors are not able to die gracefully, then kill them. + for worker in self.workers: + ray.kill(worker.actor) + + logger.debug("Shutdown successful.") + self.workers = [] + + def execute_async(self, func: Callable[..., T], *args, **kwargs) -> List[ObjectRef]: + """Execute ``func`` on each worker and return the futures. + + Args: + func: A function to call on each worker. + args, kwargs: Passed directly into func. + + Returns: + (List[ObjectRef]) A list of ``ObjectRef`` representing the + output of ``func`` from each worker. The order is the same + as ``self.workers``. + + """ + if len(self.workers) <= 0: + raise RuntimeError( + "There are no active workers. This worker " + "group has most likely been shut down. Please" + "create a new WorkerGroup or restart this one." + ) + + return [ + w.actor._RayTrainWorker__execute.options( + name=f"_RayTrainWorker__execute.{func.__name__}" + ).remote(func, *args, **kwargs) + for w in self.workers + ] + + def execute(self, func: Callable[..., T], *args, **kwargs) -> List[T]: + """Execute ``func`` on each worker and return the outputs of ``func``. + + Args: + func: A function to call on each worker. + args, kwargs: Passed directly into func. + + Returns: + (List[T]) A list containing the output of ``func`` from each + worker. The order is the same as ``self.workers``. + + """ + # TODO: Add a timeout in the case of a hang, particularly + # relevant when func is TorchConfig.on_shutdown + return ray.get(self.execute_async(func, *args, **kwargs)) + + def execute_single_async( + self, worker_index: int, func: Callable[..., T], *args, **kwargs + ) -> ObjectRef: + """Execute ``func`` on worker ``worker_index`` and return futures. + + Args: + worker_index: The index to execute func on. + func: A function to call on the first worker. + args, kwargs: Passed directly into func. + + Returns: + (ObjectRef) An ObjectRef representing the output of func. + + """ + if worker_index >= len(self.workers): + raise ValueError( + f"The provided worker_index {worker_index} is " + f"not valid for {self.num_workers} workers." + ) + return ( + self.workers[worker_index] + .actor._RayTrainWorker__execute.options( + name=f"_RayTrainWorker__execute.{func.__name__}" + ) + .remote(func, *args, **kwargs) + ) + + def execute_single( + self, worker_index: int, func: Callable[..., T], *args, **kwargs + ) -> T: + """Execute ``func`` on worker with index ``worker_index``. + + Args: + worker_index: The index to execute func on. + func: A function to call on the first worker. + args, kwargs: Passed directly into func. + + Returns: + (T) The output of func. + + """ + + return ray.get(self.execute_single_async(worker_index, func, *args, **kwargs)) + + def remove_workers(self, worker_indexes: List[int]): + """Removes the workers with the specified indexes. + + The removed workers will go out of scope and their actor processes + will be terminated. + + Args: + worker_indexes (List[int]): The indexes of the workers to remove. + """ + new_workers = [] + for i in range(len(self.workers)): + if i not in worker_indexes: + new_workers.append(self.workers[i]) + self.workers = new_workers + + def add_workers(self, num_workers: int): + """Adds ``num_workers`` to this WorkerGroup. + + Note: Adding workers when the cluster/placement group is at capacity + may lead to undefined hanging behavior. If you are attempting to + replace existing workers in the WorkerGroup, remove_workers() should + be called first. + + Args: + num_workers: The number of workers to add. + """ + new_actors = [] + new_actor_metadata = [] + for _ in range(num_workers): + actor = self._remote_cls.options( + placement_group=self._placement_group + ).remote(*self._actor_cls_args, **self._actor_cls_kwargs) + new_actors.append(actor) + new_actor_metadata.append( + actor._RayTrainWorker__execute.options( + name="_RayTrainWorker__execute.construct_metadata" + ).remote(construct_metadata) + ) + + # Get metadata from all actors. + metadata = ray.get(new_actor_metadata) + + for i in range(len(new_actors)): + self.workers.append(Worker(actor=new_actors[i], metadata=metadata[i])) + + def sort_workers_by_node_id_and_gpu_id(self, _first_node_id: Optional[str] = None): + """Reorder the workers by their node id and the lowest GPU id. + + This is useful for collocating workers on the same node. + + Example: + Given workers with the following attributes: + worker_0: node_id=1, gpu_ids=[1] + worker_1: node_id=0, gpu_ids=[0] + worker_2: node_id=1, gpu_ids=[0] + worker_3: node_id=0, gpu_ids=[1] + + The function will perform the following steps: + 1. Group by node ID: + node_id=0: worker_1, worker_3 + node_id=1: worker_0, worker_2 + + 2. Sort each group by GPU ID: + node_id=0: worker_1 (gpu_id=0), worker_3 (gpu_id=1) + node_id=1: worker_2 (gpu_id=0), worker_0 (gpu_id=1) + + Resulting in the order: [worker_1, worker_3, worker_2, worker_0] + + Args: + _first_node_id: The first ID to group by. + Set this to the node ID of the trainer coordinator to ensure that the + rank 0 worker is on the same node, allowing additional resources to + be specified for rank 0 workers via + `ScalingConfig(trainer_resources=)`. + """ + node_id_to_workers = defaultdict(list) + + if _first_node_id is not None: + node_id_to_workers[_first_node_id] = [] + + for worker in self.workers: + node_id_to_workers[worker.metadata.node_id].append(worker) + + # Sort workers on the same node by the lowest GPU id + # More details: https://github.com/ray-project/ray/issues/40803 + def get_lowest_gpu_id(worker) -> int: + gpu_ids = worker.metadata.resource_ids.get("GPU", []) + # If there are no GPU IDs, return 0 as a default + if not gpu_ids: + return 0 + + # Attempt to convert GPU IDs to integers and find the minimum ID. + # Fallback to return the minimum string-based ID + try: + return min(int(gpu_id) for gpu_id in gpu_ids) + except ValueError: + return min(gpu_ids) + + for node_id in node_id_to_workers: + node_id_to_workers[node_id].sort(key=get_lowest_gpu_id) + + sorted_workers = [] + for workers in node_id_to_workers.values(): + sorted_workers.extend(workers) + + self.workers = sorted_workers + + def __len__(self): + return len(self.workers) + + def get_resources_per_worker(self) -> dict: + """Get the resources allocated per worker.""" + return copy.deepcopy(self.resources_per_worker) diff --git a/lib/python3.12/site-packages/ray/train/backend.py b/lib/python3.12/site-packages/ray/train/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..48f6c45b8193bd0ac28103099667b6f9738fdbcc --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/backend.py @@ -0,0 +1,59 @@ +import logging +from contextlib import nullcontext +from typing import TypeVar + +from ray.train._internal.base_worker_group import BaseWorkerGroup +from ray.train._internal.utils import Singleton +from ray.util.annotations import DeveloperAPI +from ray.widgets import make_table_html_repr + +EncodedData = TypeVar("EncodedData") + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class BackendConfig: + """Parent class for configurations of training backend.""" + + @property + def backend_cls(self): + return Backend + + @property + def train_func_context(self): + return nullcontext + + def _repr_html_(self) -> str: + return make_table_html_repr(obj=self, title=type(self).__name__) + + +@DeveloperAPI +class Backend(metaclass=Singleton): + """Singleton for distributed communication backend. + + Attributes: + share_cuda_visible_devices: If True, each worker + process will have CUDA_VISIBLE_DEVICES set as the visible device + IDs of all workers on the same node for this training instance. + If False, each worker will have CUDA_VISIBLE_DEVICES set to the + device IDs allocated by Ray for that worker. + """ + + share_cuda_visible_devices: bool = False + + def on_start(self, worker_group: BaseWorkerGroup, backend_config: BackendConfig): + """Logic for starting this backend.""" + pass + + def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: BackendConfig): + """Logic for shutting down the backend.""" + pass + + def on_training_start( + self, worker_group: BaseWorkerGroup, backend_config: BackendConfig + ): + """Logic ran right before training is started. + + Session API is available at this point.""" + pass diff --git a/lib/python3.12/site-packages/ray/train/base_trainer.py b/lib/python3.12/site-packages/ray/train/base_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..07e7f0e4073a5a7f72f26317e2a972f74e2d2dc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/base_trainer.py @@ -0,0 +1,917 @@ +import abc +import copy +import inspect +import json +import logging +import os +import warnings +from functools import partial +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union + +import pyarrow.fs + +import ray +import ray.cloudpickle as pickle +from ray._common.usage import usage_lib +from ray._private.dict import deep_update +from ray.air._internal import usage as air_usage +from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated +from ray.air._internal.usage import AirEntrypoint +from ray.air.config import RunConfig, ScalingConfig +from ray.air.result import Result +from ray.train import Checkpoint +from ray.train._internal.session import get_session +from ray.train._internal.storage import ( + StorageContext, + _exists_at_fs_path, + get_fs_and_path, +) +from ray.train.constants import ( + V2_MIGRATION_GUIDE_MESSAGE, + _v2_migration_warnings_enabled, +) +from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE +from ray.train.utils import _log_deprecation_warning +from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI + +if TYPE_CHECKING: + from ray.data import Dataset + from ray.tune import Trainable + +_TRAINER_PKL = "trainer.pkl" + +# A type representing either a ray.data.Dataset or a function that returns a +# ray.data.Dataset and accepts no arguments. +GenDataset = Union["Dataset", Callable[[], "Dataset"]] + + +logger = logging.getLogger(__name__) + +PREPROCESSOR_DEPRECATION_MESSAGE = ( + "The `preprocessor` argument to Trainers is deprecated as of Ray 2.7. " + "Instead, use the Preprocessor `fit` and `transform` APIs directly on the Ray " + "Dataset. For any state that needs to be saved to the trained checkpoint, pass it " + "in using the `metadata` argument of the `Trainer`. " + "For a full example, see " + "https://docs.ray.io/en/master/train/user-guides/data-loading-preprocessing.html#preprocessing-structured-data " # noqa:E501 +) + +_TRAINER_RESTORE_DEPRECATION_WARNING = ( + "The `restore` and `can_restore` APIs are deprecated and " + f"will be removed in a future release. {V2_MIGRATION_GUIDE_MESSAGE}" +) + +_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING = ( + "`resume_from_checkpoint` is deprecated and will be removed in an upcoming " + f"release. {V2_MIGRATION_GUIDE_MESSAGE}" +) + + +@PublicAPI(stability="beta") +class TrainingFailedError(RuntimeError): + """An error indicating that training has failed.""" + + _RESTORE_MSG = ( + "The Ray Train run failed. Please inspect the previous error messages for a " + "cause. After fixing the issue (assuming that the error is not caused by " + "your own application logic, but rather an error such as OOM), you can restart " + "the run from scratch or continue this run.\n" + "To continue this run, you can use: " + '`trainer = {trainer_cls_name}.restore("{path}")`.' + ) + + _FAILURE_CONFIG_MSG = ( + "To start a new run that will retry on training failures, set " + "`train.RunConfig(failure_config=train.FailureConfig(max_failures))` " + "in the Trainer's `run_config` with `max_failures > 0`, or `max_failures = -1` " + "for unlimited retries." + ) + + +def _train_coordinator_fn( + config: dict, trainer_cls: Type["BaseTrainer"], metadata: dict +): + """This is the function that defines the logic of the Ray Train coordinator. + This is responsible for setting up a remote instance of the `trainer_cls` + (a different instance than the one calling `trainer.fit` on the driver!) + and running the training loop. + """ + assert metadata is not None, metadata + # Propagate user metadata from the Trainer constructor. + get_session().metadata = metadata + + # config already contains merged values. + # Instantiate new Trainer in Trainable. + trainer = trainer_cls(**config) + + # Get the checkpoint from Tune and pass it to workers later on. + checkpoint = ray.tune.get_checkpoint() + if checkpoint: + # Set `starting_checkpoint` for auto-recovery fault-tolerance + # as well as manual restoration. + trainer.starting_checkpoint = checkpoint + # else: Train will restore from the user-provided + # `resume_from_checkpoint` == `starting_checkpoint`. + + # Evaluate datasets if they are wrapped in a factory. + trainer.datasets = { + k: d() if callable(d) else d for k, d in trainer.datasets.items() + } + + trainer.setup() + trainer.training_loop() + + +@DeveloperAPI +class BaseTrainer(abc.ABC): + """Defines interface for distributed training on Ray. + + Note: The base ``BaseTrainer`` class cannot be instantiated directly. Only + one of its subclasses can be used. + + Note to developers: If a new trainer is added, please update + `air/_internal/usage.py`. + + **How does a trainer work?** + + - First, initialize the Trainer. The initialization runs locally, + so heavyweight setup should not be done in ``__init__``. + - Then, when you call ``trainer.fit()``, the Trainer is serialized + and copied to a remote Ray actor. The following methods are then + called in sequence on the remote actor. + - ``trainer.setup()``: Any heavyweight Trainer setup should be + specified here. + - ``trainer.training_loop()``: Executes the main training logic. + - Calling ``trainer.fit()`` will return a ``ray.result.Result`` + object where you can access metrics from your training run, as well + as any checkpoints that may have been saved. + + **How do I create a new Trainer?** + + Subclass ``ray.train.trainer.BaseTrainer``, and override the ``training_loop`` + method, and optionally ``setup``. + + .. testcode:: + :skipif: True + + import torch + + from ray.train.trainer import BaseTrainer + from ray import train, tune + + + class MyPytorchTrainer(BaseTrainer): + def setup(self): + self.model = torch.nn.Linear(1, 1) + self.optimizer = torch.optim.SGD( + self.model.parameters(), lr=0.1) + + def training_loop(self): + # You can access any Trainer attributes directly in this method. + # self.datasets["train"] has already been + dataset = self.datasets["train"] + + torch_ds = dataset.iter_torch_batches(dtypes=torch.float) + loss_fn = torch.nn.MSELoss() + + for epoch_idx in range(10): + loss = 0 + num_batches = 0 + torch_ds = dataset.iter_torch_batches( + dtypes=torch.float, batch_size=2 + ) + for batch in torch_ds: + X = torch.unsqueeze(batch["x"], 1) + y = torch.unsqueeze(batch["y"], 1) + # Compute prediction error + pred = self.model(X) + batch_loss = loss_fn(pred, y) + + # Backpropagation + self.optimizer.zero_grad() + batch_loss.backward() + self.optimizer.step() + + loss += batch_loss.item() + num_batches += 1 + loss /= num_batches + + # Use Tune functions to report intermediate + # results. + train.report({"loss": loss, "epoch": epoch_idx}) + + + # Initialize the Trainer, and call Trainer.fit() + import ray + train_dataset = ray.data.from_items( + [{"x": i, "y": i} for i in range(10)]) + my_trainer = MyPytorchTrainer(datasets={"train": train_dataset}) + result = my_trainer.fit() + + Args: + scaling_config: Configuration for how to scale training. + run_config: Configuration for the execution of the training run. + datasets: Any Datasets to use for training. Use the key "train" + to denote which dataset is the training dataset. + metadata: Dict that should be made available via + `train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + resume_from_checkpoint: A checkpoint to resume training from. + """ + + _scaling_config_allowed_keys: List[str] = [ + "trainer_resources", + ] + _handles_checkpoint_freq: bool = False + _handles_checkpoint_at_end: bool = False + + # fields to propagate to Tuner param_space. + # See `BaseTrainer._extract_fields_for_tuner_param_space` for more details. + _fields_for_tuner_param_space = [] + + def __init__( + self, + *, + scaling_config: Optional[ScalingConfig] = None, + run_config: Optional[RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + metadata: Optional[Dict[str, Any]] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + ): + self.scaling_config = ( + scaling_config if scaling_config is not None else ScalingConfig() + ) + self.run_config = ( + copy.copy(run_config) if run_config is not None else RunConfig() + ) + self.metadata = metadata + self.datasets = datasets if datasets is not None else {} + self.starting_checkpoint = resume_from_checkpoint + + if _v2_migration_warnings_enabled(): + if metadata is not None: + _log_deprecation_warning(_GET_METADATA_DEPRECATION_MESSAGE) + if resume_from_checkpoint is not None: + _log_deprecation_warning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING) + + # These attributes should only be set through `BaseTrainer.restore` + self._restore_path = None + self._restore_storage_filesystem = None + + self._validate_attributes() + + usage_lib.record_library_usage("train") + air_usage.tag_air_trainer(self) + + @classmethod + @Deprecated(message=_TRAINER_RESTORE_DEPRECATION_WARNING) + def restore( + cls: Type["BaseTrainer"], + path: Union[str, os.PathLike], + storage_filesystem: Optional[pyarrow.fs.FileSystem] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + scaling_config: Optional[ScalingConfig] = None, + **kwargs, + ) -> "BaseTrainer": + """Restores a Train experiment from a previously interrupted/failed run. + + Restore should be used for experiment-level fault tolerance in the event + that the head node crashes (e.g., OOM or some other runtime error) or the + entire cluster goes down (e.g., network error affecting all nodes). + + A run that has already completed successfully will not be resumed from this API. + To continue training from a successful run, launch a new run with the + ``Trainer(resume_from_checkpoint)`` API instead, passing in a + checkpoint from the previous run to start with. + + .. note:: + + Restoring an experiment from a path that's pointing to a *different* + location than the original experiment path is supported. However, Ray Train + assumes that the full experiment directory is available + (including checkpoints) so that it's possible to resume trials from their + latest state. + + For example, if the original experiment path was run locally, then the + results are uploaded to cloud storage, Ray Train expects the full contents + to be available in cloud storage if attempting to resume + via ``Trainer.restore("s3://...")``. The restored run will + continue writing results to the same cloud storage location. + + The following example can be paired with implementing job retry using + :ref:`Ray Jobs ` to produce a Train experiment that will + attempt to resume on both experiment-level and trial-level failures: + + .. testcode:: + :skipif: True + + import os + import ray + from ray import train + from ray.train.trainer import BaseTrainer + + experiment_name = "unique_experiment_name" + storage_path = os.path.expanduser("~/ray_results") + experiment_dir = os.path.join(storage_path, experiment_name) + + # Define some dummy inputs for demonstration purposes + datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])} + + class CustomTrainer(BaseTrainer): + def training_loop(self): + pass + + if CustomTrainer.can_restore(experiment_dir): + trainer = CustomTrainer.restore( + experiment_dir, datasets=datasets + ) + else: + trainer = CustomTrainer( + datasets=datasets, + run_config=train.RunConfig( + name=experiment_name, + storage_path=storage_path, + # Tip: You can also enable retries on failure for + # worker-level fault tolerance + failure_config=train.FailureConfig(max_failures=3), + ), + ) + + result = trainer.fit() + + Args: + path: The path to the experiment directory of the training run to restore. + This can be a local path or a remote URI if the experiment was + uploaded to the cloud. + storage_filesystem: Custom ``pyarrow.fs.FileSystem`` + corresponding to the ``path``. This may be necessary if the original + experiment passed in a custom filesystem. + datasets: Re-specified datasets used in the original training run. + This must include all the datasets that were passed in the + original trainer constructor. + scaling_config: Optionally re-specified scaling config. This can be + modified to be different from the original spec. + **kwargs: Other optionally re-specified arguments, passed in by subclasses. + + Raises: + ValueError: If all datasets were not re-supplied on restore. + + Returns: + BaseTrainer: A restored instance of the class that is calling this method. + """ + if _v2_migration_warnings_enabled(): + _log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING) + + if not cls.can_restore(path, storage_filesystem): + raise ValueError( + f"Invalid restore path: {path}. Make sure that this path exists and " + "is the experiment directory that results from a call to " + "`trainer.fit()`." + ) + fs, fs_path = get_fs_and_path(path, storage_filesystem) + trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix() + with fs.open_input_file(trainer_pkl_path) as f: + trainer_cls, param_dict = pickle.loads(f.readall()) + + if trainer_cls is not cls: + warnings.warn( + f"Invalid trainer type. You are attempting to restore a trainer of type" + f" {trainer_cls} with `{cls.__name__}.restore`, " + "which will most likely fail. " + f"Use `{trainer_cls.__name__}.restore` instead." + ) + + original_datasets = param_dict.pop("datasets", {}) + if original_datasets and not datasets: + raise ValueError( + "The following datasets need to be provided again on restore: " + f"{list(original_datasets.keys())}\n" + f"Use {cls.__name__}.restore(..., datasets=datasets) " + "with the datasets that were provided to the original trainer." + ) + datasets = datasets or {} + if set(original_datasets) != set(datasets): + raise ValueError( + "The provided datasets don't match the original dataset keys.\n" + f" Expected datasets for the keys: {list(original_datasets.keys())}\n" + f" Actual datasets provided: {list(datasets.keys())}" + ) + param_dict["datasets"] = datasets + + if scaling_config: + param_dict["scaling_config"] = scaling_config + + for param_name, val in kwargs.items(): + # Overwrite the old value if something is passed into restore + if val is not None: + param_dict[param_name] = val + + try: + trainer = cls(**param_dict) + except Exception as e: + raise ValueError( + "Trainer restoration failed (see above for the stack trace). " + "Make sure that you use the right trainer class to restore: " + f"`{cls.__name__}.restore`\n" + ) from e + trainer._restore_path = path + trainer._restore_storage_filesystem = storage_filesystem + return trainer + + @classmethod + @Deprecated( + message=_TRAINER_RESTORE_DEPRECATION_WARNING, + warning=_v2_migration_warnings_enabled(), + ) + def can_restore( + cls: Type["BaseTrainer"], + path: Union[str, os.PathLike], + storage_filesystem: Optional[pyarrow.fs.FileSystem] = None, + ) -> bool: + """Checks whether a given directory contains a restorable Train experiment. + + Args: + path: The path to the experiment directory of the Train experiment. + This can be either a local directory (e.g., ~/ray_results/exp_name) + or a remote URI (e.g., s3://bucket/exp_name). + + Returns: + bool: Whether this path exists and contains the trainer state to resume from + """ + if _v2_migration_warnings_enabled(): + _log_deprecation_warning(_TRAINER_RESTORE_DEPRECATION_WARNING) + + fs, fs_path = get_fs_and_path(path, storage_filesystem) + trainer_pkl_path = Path(fs_path, _TRAINER_PKL).as_posix() + return _exists_at_fs_path(fs, trainer_pkl_path) + + def __repr__(self): + # A dictionary that maps parameters to their default values. + default_values: Dict[str, Any] = { + "scaling_config": ScalingConfig(), + "run_config": RunConfig(), + "datasets": {}, + "starting_checkpoint": None, + } + + non_default_arguments = [] + for parameter, default_value in default_values.items(): + value = getattr(self, parameter) + if value != default_value: + non_default_arguments.append(f"{parameter}={value!r}") + + if non_default_arguments: + return f"<{self.__class__.__name__} {' '.join(non_default_arguments)}>" + + return f"<{self.__class__.__name__}>" + + def __new__(cls, *args, **kwargs): + # Store the init args as attributes so this can be merged with Tune hparams. + trainer = super(BaseTrainer, cls).__new__(cls) + parameters = inspect.signature(cls.__init__).parameters + parameters = list(parameters.keys()) + # Remove self. + parameters = parameters[1:] + arg_dict = dict(zip(parameters, args)) + trainer._param_dict = {**arg_dict, **kwargs} + return trainer + + def _validate_attributes(self): + """Called on __init()__ to validate trainer attributes.""" + # Run config + if not isinstance(self.run_config, RunConfig): + raise ValueError( + f"`run_config` should be an instance of `ray.train.RunConfig`, " + f"found {type(self.run_config)} with value `{self.run_config}`." + ) + # Scaling config + if not isinstance(self.scaling_config, ScalingConfig): + raise ValueError( + "`scaling_config` should be an instance of `ScalingConfig`, " + f"found {type(self.scaling_config)} with value `{self.scaling_config}`." + ) + # Datasets + if not isinstance(self.datasets, dict): + raise ValueError( + f"`datasets` should be a dict mapping from a string to " + f"`ray.data.Dataset` objects, " + f"found {type(self.datasets)} with value `{self.datasets}`." + ) + else: + for key, dataset in self.datasets.items(): + if not isinstance(dataset, ray.data.Dataset) and not callable(dataset): + raise ValueError( + f"The Dataset under '{key}' key is not a " + "`ray.data.Dataset`. " + f"Received {dataset} instead." + ) + # Metadata. + self.metadata = self.metadata or {} + if not isinstance(self.metadata, dict): + raise TypeError( + f"The provided metadata must be a dict, was {type(self.metadata)}." + ) + try: + self.metadata = json.loads(json.dumps(self.metadata)) + except Exception as e: + raise ValueError( + "The provided metadata must be JSON-serializable: " + f"{self.metadata}: {e}" + ) + + if self.starting_checkpoint is not None and not isinstance( + self.starting_checkpoint, Checkpoint + ): + raise ValueError( + f"`resume_from_checkpoint` should be an instance of " + f"`ray.train.Checkpoint`, found {type(self.starting_checkpoint)} " + f"with value `{self.starting_checkpoint}`." + ) + + self._log_v2_deprecation_warnings() + + def _log_v2_deprecation_warnings(self): + """Logs deprecation warnings for v2 migration. + + Log them here in the Ray Train case rather than in the configuration + constructors to avoid logging incorrect deprecation warnings when + `ray.train.RunConfig` is passed to Ray Tune. + """ + from ray.train.v2._internal.constants import V2_ENABLED_ENV_VAR, is_v2_enabled + + if is_v2_enabled(): + raise DeprecationWarning( + f"Detected use of a deprecated Trainer import from `{self.__class__.__module__}`. " + "This Trainer class is not compatible with Ray Train V2.\n" + "To fix this:\n" + " - Update to use the new import path. For example, " + "`from ray.train.torch.torch_trainer import TorchTrainer` -> " + "`from ray.train.torch import TorchTrainer`\n" + f" - Or, explicitly disable V2 by setting: {V2_ENABLED_ENV_VAR}=0\n" + "See this issue for more context: " + "https://github.com/ray-project/ray/issues/49454" + ) + + if not _v2_migration_warnings_enabled(): + return + + from ray.train.v2._internal.migration_utils import ( + CALLBACKS_DEPRECATION_MESSAGE, + FAIL_FAST_DEPRECATION_MESSAGE, + LOG_TO_FILE_DEPRECATION_MESSAGE, + PROGRESS_REPORTER_DEPRECATION_MESSAGE, + STOP_DEPRECATION_MESSAGE, + SYNC_CONFIG_DEPRECATION_MESSAGE, + TRAINER_RESOURCES_DEPRECATION_MESSAGE, + VERBOSE_DEPRECATION_MESSAGE, + ) + + # ScalingConfig deprecations + if self.scaling_config.trainer_resources is not None: + _log_deprecation_warning(TRAINER_RESOURCES_DEPRECATION_MESSAGE) + + # FailureConfig deprecations + if self.run_config.failure_config.fail_fast: + _log_deprecation_warning(FAIL_FAST_DEPRECATION_MESSAGE) + + # RunConfig deprecations + # NOTE: _verbose is the original verbose value passed by the user + if self.run_config._verbose is not None: + _log_deprecation_warning(VERBOSE_DEPRECATION_MESSAGE) + + if self.run_config.log_to_file: + _log_deprecation_warning(LOG_TO_FILE_DEPRECATION_MESSAGE) + + if self.run_config.stop is not None: + _log_deprecation_warning(STOP_DEPRECATION_MESSAGE) + + if self.run_config.callbacks is not None: + _log_deprecation_warning(CALLBACKS_DEPRECATION_MESSAGE) + + if self.run_config.progress_reporter is not None: + _log_deprecation_warning(PROGRESS_REPORTER_DEPRECATION_MESSAGE) + + if self.run_config.sync_config != ray.train.SyncConfig(): + _log_deprecation_warning(SYNC_CONFIG_DEPRECATION_MESSAGE) + + @classmethod + def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig: + """Returns scaling config dataclass after validating updated keys.""" + ensure_only_allowed_dataclass_keys_updated( + dataclass=scaling_config, + allowed_keys=cls._scaling_config_allowed_keys, + ) + return scaling_config + + def setup(self) -> None: + """Called during fit() to perform initial setup on the Trainer. + + .. note:: This method is run on a remote process. + + This method will not be called on the driver, so any expensive setup + operations should be placed here and not in ``__init__``. + + This method is called prior to ``preprocess_datasets`` and + ``training_loop``. + """ + pass + + def preprocess_datasets(self) -> None: + """Deprecated.""" + raise DeprecationWarning( + "`preprocess_datasets` is no longer used, since preprocessors " + f"are no longer accepted by Trainers.\n{PREPROCESSOR_DEPRECATION_MESSAGE}" + ) + + @abc.abstractmethod + def training_loop(self) -> None: + """Loop called by fit() to run training and report results to Tune. + + .. note:: This method runs on a remote process. + + ``self.datasets`` have already been evaluated if they were wrapped in a factory. + + You can use the :ref:`Ray Train utilities ` + (:func:`train.report() ` and + :func:`train.get_checkpoint() `) inside + this training loop. + + Example: + + .. testcode:: + + from ray.train.trainer import BaseTrainer + from ray import train + + class MyTrainer(BaseTrainer): + def training_loop(self): + for epoch_idx in range(5): + ... + train.report({"epoch": epoch_idx}) + + """ + raise NotImplementedError + + @PublicAPI(stability="beta") + def fit(self) -> Result: + """Runs training. + + Returns: + A Result object containing the training result. + + Raises: + ray.train.base_trainer.TrainingFailedError: If any failures during the execution + of ``self.as_trainable()``, or during the Tune execution loop. + """ + from ray.tune import ResumeConfig, TuneError + from ray.tune.tuner import Tuner + + trainable = self.as_trainable() + param_space = self._extract_fields_for_tuner_param_space() + + self.run_config.name = ( + self.run_config.name or StorageContext.get_experiment_dir_name(trainable) + ) + # The storage context here is only used to access the resolved + # storage fs and experiment path, in order to avoid duplicating that logic. + # This is NOT the storage context object that gets passed to remote workers. + storage = StorageContext( + storage_path=self.run_config.storage_path, + experiment_dir_name=self.run_config.name, + storage_filesystem=self.run_config.storage_filesystem, + ) + + if self._restore_path: + tuner = Tuner.restore( + path=self._restore_path, + trainable=trainable, + param_space=param_space, + _resume_config=ResumeConfig( + finished=ResumeConfig.ResumeType.RESUME, + unfinished=ResumeConfig.ResumeType.RESUME, + errored=ResumeConfig.ResumeType.RESUME, + ), + storage_filesystem=self._restore_storage_filesystem, + ) + else: + tuner = Tuner( + trainable=trainable, + param_space=param_space, + run_config=self.run_config, + _entrypoint=AirEntrypoint.TRAINER, + ) + + self._save(storage.storage_filesystem, storage.experiment_fs_path) + + restore_msg = TrainingFailedError._RESTORE_MSG.format( + trainer_cls_name=self.__class__.__name__, + path=str(storage.experiment_fs_path), + ) + + try: + result_grid = tuner.fit() + except TuneError as e: + # Catch any `TuneError`s raised by the `Tuner.fit` call. + # Unwrap the `TuneError` if needed. + parent_error = e.__cause__ or e + + # Raise it to the user as a `TrainingFailedError` with a message to restore. + raise TrainingFailedError(restore_msg) from parent_error + # Other exceptions get passed through directly (ex: on `fail_fast='raise'`) + + assert len(result_grid) == 1 + result = result_grid[0] + if result.error: + # Raise trainable errors to the user with a message to restore + # or configure `FailureConfig` in a new run. + raise TrainingFailedError( + "\n".join([restore_msg, TrainingFailedError._FAILURE_CONFIG_MSG]) + ) from result.error + return result + + def _save(self, fs: pyarrow.fs.FileSystem, experiment_path: str): + """Saves the current trainer's class along with the `param_dict` of + parameters passed to this trainer's constructor. + + This is used to recreate the trainer on restore. + Unless a parameter is re-specified during restoration (only a subset + of parameters can be passed in again), that parameter will be loaded + from the saved copy. + + Datasets should not be saved as part of the state. Instead, we save the + keys and replace the dataset values with dummy functions that will + raise an error if invoked. The error only serves as a guardrail for + misuse (e.g., manually unpickling and constructing the Trainer again) + and is not typically surfaced, since datasets must be re-specified + upon restoration. + """ + param_dict = self._param_dict.copy() + datasets = param_dict.pop("datasets", {}) + + def raise_fn(): + raise RuntimeError + + if datasets: + param_dict["datasets"] = { + dataset_name: raise_fn for dataset_name in datasets + } + + cls_and_param_dict = (self.__class__, param_dict) + + fs.create_dir(experiment_path) + with fs.open_output_stream(Path(experiment_path, _TRAINER_PKL).as_posix()) as f: + f.write(pickle.dumps(cls_and_param_dict)) + + def _extract_fields_for_tuner_param_space(self) -> Dict: + """Extracts fields to be included in `Tuner.param_space`. + + This is needed to leverage the full logging/integration offerings from Tune. + For example, `param_space` is logged automatically to wandb integration. + + Currently only done for `train_loop_config`. + + Returns: + A dictionary that should be passed to Tuner.param_space. + """ + result = {} + for key in self._fields_for_tuner_param_space: + if key in self._param_dict.keys(): + result[key] = copy.deepcopy(self._param_dict[key]) + return result + + def _generate_trainable_cls(self) -> Type["Trainable"]: + """Generates the base Trainable class. + + Returns: + A Trainable class to use for training. + """ + + from ray.tune.execution.placement_groups import PlacementGroupFactory + from ray.tune.trainable import wrap_function + + trainer_cls = self.__class__ + scaling_config = self.scaling_config + metadata = self.metadata + + train_coordinator_fn = partial( + _train_coordinator_fn, trainer_cls=trainer_cls, metadata=metadata + ) + # Change the name of the training function to match the name of the Trainer + # class. This will mean the Tune trial name will match the name of Trainer on + # stdout messages and the results directory. + train_coordinator_fn.__name__ = trainer_cls.__name__ + + trainable_cls = wrap_function(train_coordinator_fn) + has_base_dataset = bool(self.datasets) + if has_base_dataset: + from ray.data.context import DataContext + + dataset_context = DataContext.get_current() + else: + dataset_context = None + + class TrainTrainable(trainable_cls): + """Adds default resources to the Trainable.""" + + _handles_checkpoint_freq = trainer_cls._handles_checkpoint_freq + _handles_checkpoint_at_end = trainer_cls._handles_checkpoint_at_end + + @classmethod + def has_base_dataset(cls) -> bool: + """Whether a dataset is provided through the Trainer.""" + return has_base_dataset + + @classmethod + def base_scaling_config(cls) -> ScalingConfig: + """Returns the unchanged scaling config provided through the Trainer.""" + return scaling_config + + def setup(self, config, **kwargs): + base_config = dict(kwargs) + # Merge Tuner param space hyperparameters in `config` into the + # base config passed to the Trainer constructor, which is `base_config`. + # `base_config` is pulled from the object store from the usage of + # tune.with_parameters in `BaseTrainer.as_trainable`. + + # run_config is not a tunable hyperparameter so it does not need to be + # merged. + run_config = base_config.pop("run_config", None) + self._merged_config = deep_update( + base_config, self.config, new_keys_allowed=True + ) + self._merged_config["run_config"] = run_config + merged_scaling_config = self._merged_config.get( + "scaling_config", ScalingConfig() + ) + if isinstance(merged_scaling_config, dict): + merged_scaling_config = ScalingConfig(**merged_scaling_config) + self._merged_config[ + "scaling_config" + ] = self._reconcile_scaling_config_with_trial_resources( + merged_scaling_config + ) + if self.has_base_dataset(): + # Set the DataContext on the Trainer actor to the DataContext + # specified on the driver. + DataContext._set_current(dataset_context) + super(TrainTrainable, self).setup(config) + + def _reconcile_scaling_config_with_trial_resources( + self, scaling_config: ScalingConfig + ) -> ScalingConfig: + """ + ResourceChangingScheduler workaround. + + Ensures that the scaling config matches trial resources. + + This should be replaced with RCS returning a ScalingConfig + in the future. + """ + + trial_resources = self.trial_resources + # This will be false if the resources are default + if not isinstance(trial_resources, PlacementGroupFactory): + return scaling_config + + # Ignore ResourceChangingScheduler workaround when resource bundles + # are unchanged + if self.trial_resources == scaling_config.as_placement_group_factory(): + return scaling_config + + trainer_cls._validate_scaling_config(scaling_config) + + return ScalingConfig.from_placement_group_factory(trial_resources) + + def _trainable_func(self, config): + # We ignore the config passed by Tune and instead use the merged + # config which includes the initial Trainer args. + super()._trainable_func(self._merged_config) + + @classmethod + def default_resource_request(cls, config): + # `config["scaling_config"] is a dataclass when passed via the + # `scaling_config` argument in `Trainer` and is a dict when passed + # via the `scaling_config` key of `param_spec`. + + # Conversion logic must be duplicated in `TrainTrainable.__init__` + # because this is a class method. + updated_scaling_config = config.get("scaling_config", scaling_config) + if isinstance(updated_scaling_config, dict): + updated_scaling_config = ScalingConfig(**updated_scaling_config) + validated_scaling_config = trainer_cls._validate_scaling_config( + updated_scaling_config + ) + return validated_scaling_config.as_placement_group_factory() + + return TrainTrainable + + def as_trainable(self) -> Type["Trainable"]: + """Converts self to a ``tune.Trainable`` class.""" + from ray import tune + + base_config = self._param_dict + trainable_cls = self._generate_trainable_cls() + + # Wrap with `tune.with_parameters` to handle very large values in base_config + return tune.with_parameters(trainable_cls, **base_config) diff --git a/lib/python3.12/site-packages/ray/train/collective/__init__.py b/lib/python3.12/site-packages/ray/train/collective/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6cffcda40d745eb695e24db0959c3e8ab2134d18 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/collective/__init__.py @@ -0,0 +1,19 @@ +from ray.train.v2._internal.constants import is_v2_enabled + +if is_v2_enabled(): + from ray.train.collective.collectives import barrier, broadcast_from_rank_zero + + __all__ = [ + "broadcast_from_rank_zero", + "barrier", + ] + + broadcast_from_rank_zero.__module__ = "ray.train.collective" + barrier.__module__ = "ray.train.collective" +else: + raise ImportError( + "`ray.train.collective` is only available in Ray Train v2. " + "To enable it, please set `RAY_TRAIN_V2_ENABLED=1`." + ) + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/collective/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/collective/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9863d28055dbe70030d31bceeb9b09a8c5fcf680 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/collective/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/collective/__pycache__/collectives.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/collective/__pycache__/collectives.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82c0865cca2846fc6201c989d7d28260ee89bd43 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/collective/__pycache__/collectives.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/collective/collectives.py b/lib/python3.12/site-packages/ray/train/collective/collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..2d77fb659a280b6c5c65171935b5d2f184d45f73 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/collective/collectives.py @@ -0,0 +1,82 @@ +import logging +from typing import Optional, TypeVar + +from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils +from ray.train.v2._internal.util import requires_train_worker +from ray.util.annotations import PublicAPI + +T = TypeVar("T", bound=Optional[object]) + + +logger = logging.getLogger(__file__) + + +@PublicAPI(stability="alpha") +@requires_train_worker() +def broadcast_from_rank_zero(data: T) -> T: + """Broadcast small (<1kb) data from the rank 0 worker to all other workers. + + Serves as a barrier, meaning that all workers must call this method before + the training function can continue. + + Example: + + .. testcode: + + from ray.train import get_context + from ray.train.collective import broadcast_from_rank_zero + from ray.train.torch import TorchTrainer + + def train_func(): + ... + if get_context().get_world_rank() == 0: + data = {"some_key": "some_value"} + else: + data = None + data = broadcast_from_rank_zero(data) + ... + + trainer = TorchTrainer(train_func) + trainer.fit() + + Args: + data: The small (1kb) data to broadcast from the rank 0 worker to all + other workers. + + Returns: + The data broadcasted from the rank 0 worker. + + Raises: + ValueError: If the data is too big. + pickle.PicklingError: If the data is not pickleable. + TypeError: If the data is not pickleable. + """ + return get_train_fn_utils().broadcast_from_rank_zero(data) + + +@PublicAPI(stability="alpha") +@requires_train_worker() +def barrier() -> None: + """Create a barrier across all workers. + + All workers must call this method before the training function can continue. + + Example: + + .. testcode: + + from ray.train import get_context + from ray.train.collective import barrier + from ray.train.torch import TorchTrainer + + def train_func(): + ... + print(f"Rank {get_context().get_world_rank()} is waiting at the barrier.") + barrier() + print(f"Rank {get_context().get_world_rank()} has passed the barrier.") + ... + + trainer = TorchTrainer(train_func) + trainer.fit() + """ + return get_train_fn_utils().barrier() diff --git a/lib/python3.12/site-packages/ray/train/constants.py b/lib/python3.12/site-packages/ray/train/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..a51cc492084525bbf8eb4361f94fb1c031bdfdd3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/constants.py @@ -0,0 +1,158 @@ +from pathlib import Path +from typing import Any + +import ray +from ray._private.ray_constants import env_bool +from ray.air.constants import ( # noqa: F401 + COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV, + EVALUATION_DATASET_KEY, + MODEL_KEY, + PREPROCESSOR_KEY, + TRAIN_DATASET_KEY, +) + + +def _get_ray_train_session_dir() -> str: + assert ray.is_initialized(), "Ray must be initialized to get the session dir." + return Path( + ray._private.worker._global_node.get_session_dir_path(), "artifacts" + ).as_posix() + + +DEFAULT_STORAGE_PATH = Path("~/ray_results").expanduser().as_posix() + +# Autofilled ray.train.report() metrics. Keys should be consistent with Tune. +CHECKPOINT_DIR_NAME = "checkpoint_dir_name" +TIME_TOTAL_S = "_time_total_s" +WORKER_HOSTNAME = "_hostname" +WORKER_NODE_IP = "_node_ip" +WORKER_PID = "_pid" + +# Will not be reported unless ENABLE_DETAILED_AUTOFILLED_METRICS_ENV +# env var is not 0 +DETAILED_AUTOFILLED_KEYS = {WORKER_HOSTNAME, WORKER_NODE_IP, WORKER_PID, TIME_TOTAL_S} + +# Default filename for JSON logger +RESULT_FILE_JSON = "results.json" + +# The name of the subdirectory inside the trainer run_dir to store checkpoints. +TRAIN_CHECKPOINT_SUBDIR = "checkpoints" + +# The key to use to specify the checkpoint id for Tune. +# This needs to be added to the checkpoint dictionary so if the Tune trial +# is restarted, the checkpoint_id can continue to increment. +TUNE_CHECKPOINT_ID = "_current_checkpoint_id" + +# Deprecated configs can use this value to detect if the user has set it. +# This has type Any to allow it to be assigned to any annotated parameter +# without causing type errors. +_DEPRECATED_VALUE: Any = "DEPRECATED" + + +# ================================================== +# Train V2 constants +# ================================================== + +# Set this to 1 to enable deprecation warnings for V2 migration. +ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR = "RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS" + + +V2_MIGRATION_GUIDE_MESSAGE = ( + "See this issue for more context and migration options: " + "https://github.com/ray-project/ray/issues/49454. " + "Disable these warnings by setting the environment variable: " + f"{ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR}=0" +) + + +def _v2_migration_warnings_enabled() -> bool: + return env_bool(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, True) + + +# ================================================== +# Environment Variables +# ================================================== + +ENABLE_DETAILED_AUTOFILLED_METRICS_ENV = ( + "TRAIN_RESULT_ENABLE_DETAILED_AUTOFILLED_METRICS" +) + +# Integer value which if set will override the value of +# Backend.share_cuda_visible_devices. 1 for True, 0 for False. +ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_CUDA_VISIBLE_DEVICES" + +# Integer value which if set will not share HIP accelerator visible devices +# across workers. 1 for True (default), 0 for False. +ENABLE_SHARE_HIP_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_HIP_VISIBLE_DEVICES" + +# Integer value which if set will not share neuron-core accelerator visible cores +# across workers. 1 for True (default), 0 for False. +ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV = ( + "TRAIN_ENABLE_SHARE_NEURON_CORES_ACCELERATOR" +) + +# Integer value which if set will not share npu visible devices +# across workers. 1 for True (default), 0 for False. +ENABLE_SHARE_NPU_RT_VISIBLE_DEVICES_ENV = "TRAIN_ENABLE_SHARE_ASCEND_RT_VISIBLE_DEVICES" + +# Integer value which indicates the number of seconds to wait when creating +# the worker placement group before timing out. +TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV = "TRAIN_PLACEMENT_GROUP_TIMEOUT_S" + +# Integer value which if set will change the placement group strategy from +# PACK to SPREAD. 1 for True, 0 for False. +TRAIN_ENABLE_WORKER_SPREAD_ENV = "TRAIN_ENABLE_WORKER_SPREAD" + +# Set this to 0 to disable changing the working directory of each Tune Trainable +# or Train worker to the trial directory. Defaults to 1. +RAY_CHDIR_TO_TRIAL_DIR = "RAY_CHDIR_TO_TRIAL_DIR" + +# Set this to 1 to count preemption errors toward `FailureConfig(max_failures)`. +# Defaults to 0, which always retries on node preemption failures. +RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE = "RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE" + +# Set this to 1 to start a StateActor and collect information Train Runs +# Defaults to 0 +RAY_TRAIN_ENABLE_STATE_TRACKING = "RAY_TRAIN_ENABLE_STATE_TRACKING" + +# Set this to 1 to only store the checkpoint score attribute with the Checkpoint +# in the CheckpointManager. The Result will only have the checkpoint score attribute +# but files written to disk like result.json will still have all the metrics. +# Defaults to 0. +# TODO: this is a temporary solution to avoid CheckpointManager OOM. +# See https://github.com/ray-project/ray/pull/54642#issue-3234029360 for more details. +TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE = ( + "TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE" +) + +# Seconds to wait for torch process group to shut down. +# Shutting down a healthy torch process group, which we may want to do for reasons +# like restarting a group of workers if an async checkpoint upload fails, can hang. +# This is a workaround until we figure out how to avoid this hang. +TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S = "TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S" +DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S = 30 + +# Seconds to wait for JAX distributed shutdown. +JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S = "JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S" +DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S = 30 + +# NOTE: When adding a new environment variable, please track it in this list. +TRAIN_ENV_VARS = { + ENABLE_DETAILED_AUTOFILLED_METRICS_ENV, + ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV, + ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, + TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV, + TRAIN_ENABLE_WORKER_SPREAD_ENV, + RAY_CHDIR_TO_TRIAL_DIR, + RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE, + RAY_TRAIN_ENABLE_STATE_TRACKING, + TUNE_ONLY_STORE_CHECKPOINT_SCORE_ATTRIBUTE, + TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, + JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S, +} + +# Key for AIR Checkpoint metadata in TrainingResult metadata +CHECKPOINT_METADATA_KEY = "checkpoint_metadata" + +# Key for AIR Checkpoint world rank in TrainingResult metadata +CHECKPOINT_RANK_KEY = "checkpoint_rank" diff --git a/lib/python3.12/site-packages/ray/train/context.py b/lib/python3.12/site-packages/ray/train/context.py new file mode 100644 index 0000000000000000000000000000000000000000..7aefb4b9005544bd8ce3e9c1f8d4e78659dee87c --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/context.py @@ -0,0 +1,144 @@ +import threading +from typing import TYPE_CHECKING, Any, Dict, Optional + +from ray.train._internal import session +from ray.train._internal.storage import StorageContext +from ray.train.constants import ( + V2_MIGRATION_GUIDE_MESSAGE, + _v2_migration_warnings_enabled, +) +from ray.train.utils import _copy_doc, _log_deprecation_warning +from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI + +if TYPE_CHECKING: + from ray.tune.execution.placement_groups import PlacementGroupFactory + + +# The context singleton on this process. +_default_context: "Optional[TrainContext]" = None +_context_lock = threading.Lock() + + +_GET_METADATA_DEPRECATION_MESSAGE = ( + "`get_metadata` was an experimental API that accessed the metadata passed " + "to `Trainer(metadata=...)`. This API can be replaced by passing " + "the metadata directly to the training function (e.g., via `train_loop_config`). " + f"{V2_MIGRATION_GUIDE_MESSAGE}" +) + +_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE = ( + "`{}` is deprecated because the concept of a `Trial` will " + "soon be removed in Ray Train." + "Ray Train will no longer assume that it's running within a Ray Tune `Trial` " + "in the future. " + f"{V2_MIGRATION_GUIDE_MESSAGE}" +) + + +@PublicAPI(stability="stable") +class TrainContext: + """Context containing metadata that can be accessed within Ray Train workers.""" + + @_copy_doc(session.get_experiment_name) + def get_experiment_name(self) -> str: + return session.get_experiment_name() + + @_copy_doc(session.get_world_size) + def get_world_size(self) -> int: + return session.get_world_size() + + @_copy_doc(session.get_world_rank) + def get_world_rank(self) -> int: + return session.get_world_rank() + + @_copy_doc(session.get_local_rank) + def get_local_rank(self) -> int: + return session.get_local_rank() + + @_copy_doc(session.get_local_world_size) + def get_local_world_size(self) -> int: + return session.get_local_world_size() + + @_copy_doc(session.get_node_rank) + def get_node_rank(self) -> int: + return session.get_node_rank() + + @DeveloperAPI + @_copy_doc(session.get_storage) + def get_storage(self) -> StorageContext: + return session.get_storage() + + # Deprecated APIs + + @Deprecated( + message=_GET_METADATA_DEPRECATION_MESSAGE, + warning=_v2_migration_warnings_enabled(), + ) + @_copy_doc(session.get_metadata) + def get_metadata(self) -> Dict[str, Any]: + return session.get_metadata() + + @Deprecated( + message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_name"), + warning=_v2_migration_warnings_enabled(), + ) + @_copy_doc(session.get_trial_name) + def get_trial_name(self) -> str: + return session.get_trial_name() + + @Deprecated( + message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_id"), + warning=_v2_migration_warnings_enabled(), + ) + @_copy_doc(session.get_trial_id) + def get_trial_id(self) -> str: + return session.get_trial_id() + + @Deprecated( + message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format( + "get_trial_resources" + ), + warning=_v2_migration_warnings_enabled(), + ) + @_copy_doc(session.get_trial_resources) + def get_trial_resources(self) -> "PlacementGroupFactory": + return session.get_trial_resources() + + @Deprecated( + message=_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_dir"), + warning=_v2_migration_warnings_enabled(), + ) + @_copy_doc(session.get_trial_dir) + def get_trial_dir(self) -> str: + return session.get_trial_dir() + + +@PublicAPI(stability="stable") +def get_context() -> TrainContext: + """Get or create a singleton training context. + + The context is only available within a function passed to Ray Train. + + See the :class:`~ray.train.TrainContext` API reference to see available methods. + """ + from ray.tune.trainable.trainable_fn_utils import _in_tune_session + + # If we are running in a Tune function, switch to Tune context. + if _in_tune_session(): + from ray.tune import get_context as get_tune_context + + if _v2_migration_warnings_enabled(): + _log_deprecation_warning( + "`ray.train.get_context()` should be switched to " + "`ray.tune.get_context()` when running in a function " + "passed to Ray Tune. This will be an error in the future. " + f"{V2_MIGRATION_GUIDE_MESSAGE}" + ) + return get_tune_context() + + global _default_context + + with _context_lock: + if _default_context is None: + _default_context = TrainContext() + return _default_context diff --git a/lib/python3.12/site-packages/ray/train/data_parallel_trainer.py b/lib/python3.12/site-packages/ray/train/data_parallel_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..ea6a13653717e36c015bb9d016d8b74dfacab4e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/data_parallel_trainer.py @@ -0,0 +1,582 @@ +import logging +import uuid +from typing import Any, Callable, Dict, List, Optional, Type, Union + +import ray +from ray._private.ray_constants import env_integer +from ray._private.thirdparty.tabulate.tabulate import tabulate +from ray.air.config import RunConfig, ScalingConfig +from ray.train import BackendConfig, Checkpoint +from ray.train._internal import session +from ray.train._internal.backend_executor import BackendExecutor, TrialInfo +from ray.train._internal.data_config import DataConfig +from ray.train._internal.session import _TrainingResult, get_session +from ray.train._internal.utils import construct_train_func, count_required_parameters +from ray.train.base_trainer import _TRAINER_RESTORE_DEPRECATION_WARNING +from ray.train.constants import RAY_TRAIN_ENABLE_STATE_TRACKING +from ray.train.trainer import BaseTrainer, GenDataset, TrainingIterator +from ray.util.annotations import Deprecated, DeveloperAPI +from ray.widgets import Template +from ray.widgets.util import repr_with_fallback + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class DataParallelTrainer(BaseTrainer): + """A Trainer for data parallel training. + + You should subclass this Trainer if your Trainer follows SPMD (single program, + multiple data) programming paradigm - you want multiple processes to run the same + function, but on different data. + + This Trainer runs the function ``train_loop_per_worker`` on multiple Ray + Actors. + + The ``train_loop_per_worker`` function is expected to take in either 0 or 1 + arguments: + + .. testcode:: + + def train_loop_per_worker(): + ... + + .. testcode:: + + def train_loop_per_worker(config: Dict): + ... + + If ``train_loop_per_worker`` accepts an argument, then + ``train_loop_config`` will be passed in as the argument. This is useful if you + want to tune the values in ``train_loop_config`` as hyperparameters. + + If the ``datasets`` dict contains a training dataset (denoted by + the "train" key), then it will be split into multiple dataset + shards that can then be accessed by ``train.get_dataset_shard("train")`` inside + ``train_loop_per_worker``. All the other datasets will not be split and + ``train.get_dataset_shard(...)`` will return the entire Dataset. + + Inside the ``train_loop_per_worker`` function, you can use any of the + :ref:`Ray Train loop methods `. + + .. testcode:: + + from ray import train + + def train_loop_per_worker(): + # Report intermediate results for callbacks or logging and + # checkpoint data. + train.report(...) + + # Returns dict of last saved checkpoint. + train.get_checkpoint() + + # Returns the Dataset shard for the given key. + train.get_dataset_shard("my_dataset") + + # Returns the total number of workers executing training. + train.get_context().get_world_size() + + # Returns the rank of this worker. + train.get_context().get_world_rank() + + # Returns the rank of the worker on the current node. + train.get_context().get_local_rank() + + Any returns from the ``train_loop_per_worker`` will be discarded and not + used or persisted anywhere. + + **How do I use DataParallelTrainer or any of its subclasses?** + + Example: + + .. testcode:: + :skipif: True + + import ray + from ray import train + from ray.train import ScalingConfig + from ray.train.data_parallel_trainer import DataParallelTrainer + + def train_loop_for_worker(): + dataset_shard_for_this_worker = train.get_dataset_shard("train") + + # 3 items for 3 workers, each worker gets 1 item + batches = list(dataset_shard_for_this_worker.iter_batches(batch_size=1)) + assert len(batches) == 1 + + train_dataset = ray.data.from_items([1, 2, 3]) + assert train_dataset.count() == 3 + trainer = DataParallelTrainer( + train_loop_for_worker, + scaling_config=ScalingConfig(num_workers=3), + datasets={"train": train_dataset}, + ) + result = trainer.fit() + + **How do I develop on top of DataParallelTrainer?** + + In many cases, using DataParallelTrainer directly is sufficient to execute + functions on multiple actors. + + However, you may want to subclass ``DataParallelTrainer`` and create a custom + Trainer for the following 2 use cases: + + - **Use Case 1:** You want to do data parallel training, but want to have + a predefined ``training_loop_per_worker``. + + - **Use Case 2:** You want to implement a custom + :py:class:`~ray.train.backend.Backend` that automatically handles + additional setup or teardown logic on each actor, so that the users of this + new trainer do not have to implement this logic. For example, a + ``TensorflowTrainer`` can be built on top of ``DataParallelTrainer`` + that automatically handles setting the proper environment variables for + distributed Tensorflow on each actor. + + For 1, you can set a predefined training loop in __init__ + + .. testcode:: + + from ray.train.data_parallel_trainer import DataParallelTrainer + + class MyDataParallelTrainer(DataParallelTrainer): + def __init__(self, *args, **kwargs): + predefined_train_loop_per_worker = lambda: 1 + super().__init__(predefined_train_loop_per_worker, *args, **kwargs) + + + For 2, you can implement the ``ray.train.Backend`` and ``ray.train.BackendConfig`` + interfaces. + + .. testcode:: + + from dataclasses import dataclass + from ray.train.backend import Backend, BackendConfig + + class MyBackend(Backend): + def on_start(self, worker_group, backend_config): + def set_env_var(env_var_value): + import os + os.environ["MY_ENV_VAR"] = env_var_value + + worker_group.execute(set_env_var, backend_config.env_var) + + @dataclass + class MyBackendConfig(BackendConfig): + env_var: str = "default_value" + + def backend_cls(self): + return MyBackend + + class MyTrainer(DataParallelTrainer): + def __init__(self, train_loop_per_worker, my_backend_config: + MyBackendConfig, **kwargs): + + super().__init__( + train_loop_per_worker, + backend_config=my_backend_config, **kwargs) + + Args: + train_loop_per_worker: The training function to execute. + This can either take in no arguments or a ``config`` dict. + train_loop_config: Configurations to pass into + ``train_loop_per_worker`` if it accepts an argument. + backend_config: Configuration for setting up a Backend (e.g. Torch, + Tensorflow, Horovod) on each worker to enable distributed + communication. If no Backend should be set up, then set this to None. + scaling_config: Configuration for how to scale data parallel training. + dataset_config: Configuration for dataset ingest. This is merged with the + default dataset config for the given trainer (`cls._dataset_config`). + run_config: Configuration for the execution of the training run. + datasets: Ray Datasets to use for training and evaluation. + This is a dict where the key is the name of the dataset, which + can be accessed from within the ``train_loop_per_worker`` by calling + ``train.get_dataset_shard(dataset_key)``. + By default, all datasets are sharded equally across workers. + This can be configured via ``dataset_config``. + metadata: Dict that should be made available via + `train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + resume_from_checkpoint: A checkpoint to resume training from. + """ + + # Exposed here for testing purposes. Should never need + # to be overriden. + _backend_executor_cls: Type[BackendExecutor] = BackendExecutor + _training_iterator_cls: Type[TrainingIterator] = TrainingIterator + + _scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [ + "num_workers", + "resources_per_worker", + "use_gpu", + "placement_strategy", + "accelerator_type", + ] + + # For backwards compatibility with the legacy dataset config API. + _dataset_config = None + + _fields_for_tuner_param_space = BaseTrainer._fields_for_tuner_param_space + [ + "train_loop_config" + ] + + def __init__( + self, + train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]], + *, + train_loop_config: Optional[Dict] = None, + backend_config: Optional[BackendConfig] = None, + scaling_config: Optional[ScalingConfig] = None, + dataset_config: Optional[DataConfig] = None, + run_config: Optional[RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + metadata: Optional[Dict[str, Any]] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + ): + self._train_loop_per_worker = train_loop_per_worker + self._train_loop_config = train_loop_config + + if dataset_config is None: + dataset_config = DataConfig() + + if not isinstance(dataset_config, DataConfig): + raise ValueError( + "`dataset_config` must be an instance of ray.train.DataConfig, " + f"was: {dataset_config}" + ) + self._data_config = dataset_config + + backend_config = ( + backend_config if backend_config is not None else BackendConfig() + ) + self._backend_config = backend_config + + super(DataParallelTrainer, self).__init__( + scaling_config=scaling_config, + run_config=run_config, + datasets=datasets, + metadata=metadata, + resume_from_checkpoint=resume_from_checkpoint, + ) + + train_total_resources = self.scaling_config.total_resources + self._data_config.set_train_total_resources( + train_total_resources.get("CPU", 0), + train_total_resources.get("GPU", 0), + ) + + if env_integer(RAY_TRAIN_ENABLE_STATE_TRACKING, 0): + from ray.train._internal.state.state_actor import get_or_create_state_actor + + get_or_create_state_actor() + + @classmethod + @Deprecated(message=_TRAINER_RESTORE_DEPRECATION_WARNING) + def restore( + cls, + path: str, + train_loop_per_worker: Optional[ + Union[Callable[[], None], Callable[[Dict], None]] + ] = None, + train_loop_config: Optional[Dict] = None, + **kwargs, + ): + """Restores a DataParallelTrainer from a previously interrupted/failed run. + + Args: + train_loop_per_worker: Optionally re-specified train loop function. + This should be used to re-specify a function that is not + restorable in a new Ray cluster (e.g., it holds onto outdated + object references). This should be the same training loop + that was passed to the original trainer constructor. + train_loop_config: Optionally re-specified train config. + This should similarly be used if the original `train_loop_config` + contained outdated object references, and it should not be modified + from what was originally passed in. + + See :meth:`BaseTrainer.restore() ` + for descriptions of the other arguments. + + Returns a restored instance of the `DataParallelTrainer`. + """ + return super(DataParallelTrainer, cls).restore( + path=path, + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + **kwargs, + ) + + def _validate_attributes(self): + super()._validate_attributes() + + self._validate_train_loop_per_worker( + self._train_loop_per_worker, "train_loop_per_worker" + ) + + def _validate_train_loop_per_worker( + self, train_loop_per_worker: Callable, fn_name: str + ) -> None: + num_required_params = count_required_parameters(train_loop_per_worker) + if num_required_params > 1: + raise ValueError( + f"{fn_name} should take in 0 or 1 arguments, " + f"but it accepts {num_required_params} arguments instead." + ) + + @classmethod + def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig: + scaling_config = super(DataParallelTrainer, cls)._validate_scaling_config( + scaling_config + ) + + # This validation happens after the scaling config is updated from + # its specification in the Tuner `param_space` + if not scaling_config.use_gpu and "GPU" in ray.available_resources(): + logger.info( + "GPUs are detected in your Ray cluster, but GPU " + "training is not enabled for this trainer. To enable " + "GPU training, make sure to set `use_gpu` to True " + "in your scaling config." + ) + + if scaling_config.num_workers is None: + raise ValueError( + "You must specify the 'num_workers' in `scaling_config` as either an " + f"argument of `{cls.__name__}` or through the `param_space` of a " + "`Tuner` (if performing hyperparameter tuning)." + ) + + if scaling_config.num_workers <= 0: + raise ValueError( + "'num_workers' in `scaling_config` must be a positive " + f"integer. Received {scaling_config.num_workers}" + ) + + return scaling_config + + def _run_training(self, training_iterator: TrainingIterator) -> None: + """This method loops over the `TrainingIterator`: + The actual iteration (for ... in ...) waits for the training function + on each worker to report a result and supplies it as a list of results. + Afterwards (in the body of the loop), it will report the result + to the Tune session. + The iterator ends after the training function on each worker has finished. + """ + for training_results in training_iterator: + # TODO(ml-team): add ability to report results from multiple workers. + self._propagate_results(training_results) + + def _propagate_results(self, training_results: List[_TrainingResult]): + first_worker_result = training_results[0] + assert all(isinstance(result, _TrainingResult) for result in training_results) + + tune_session = get_session() + + # Check if any workers reported a checkpoint. + # If so, report a checkpoint pointing to the persisted location + # to Tune for book-keeping. + # NOTE: This removes the restriction for any individual worker + # (ex: global rank 0 worker) from needing to report a checkpoint. + # All workers reported a checkpoint to the same fs path, so there's + # no need to report multiple checkpoints to Tune. + worker_checkpoints = [ + result.checkpoint + for result in training_results + if result.checkpoint is not None + ] + at_least_one_reported_checkpoint = len(worker_checkpoints) > 0 + + if at_least_one_reported_checkpoint: + # Update the coordinator's checkpoint index to the latest. + # This is what keeps the checkpoint index in line with the workers. + tune_session.storage._update_checkpoint_index(first_worker_result.metrics) + + # Make sure that all workers uploaded to the same location. + assert all( + checkpoint.path == tune_session.storage.checkpoint_fs_path + for checkpoint in worker_checkpoints + ) + + checkpoint = ( + Checkpoint( + filesystem=tune_session.storage.storage_filesystem, + path=tune_session.storage.checkpoint_fs_path, + ) + if at_least_one_reported_checkpoint + else None + ) + + tracked_training_result = _TrainingResult( + checkpoint=checkpoint, + metrics=first_worker_result.metrics, + ) + + logger.debug( + "Report (metrics, checkpoint) to the Tune session:\n" + f" metrics={tracked_training_result.metrics}\n" + f" checkpoint={tracked_training_result.checkpoint}" + ) + + # Report the metrics and checkpoint to Tune. + tune_session._report_training_result(tracked_training_result) + + def training_loop(self) -> None: + scaling_config = self._validate_scaling_config(self.scaling_config) + + train_loop_per_worker = construct_train_func( + self._train_loop_per_worker, + self._train_loop_config, + train_func_context=self._backend_config.train_func_context, + fn_arg_name="train_loop_per_worker", + discard_returns=True, + ) + + trial_info = TrialInfo( + name=session.get_trial_name(), + id=session.get_trial_id(), + resources=session.get_trial_resources(), + logdir=session.get_trial_dir(), + driver_ip=ray.util.get_node_ip_address(), + driver_node_id=ray.get_runtime_context().get_node_id(), + experiment_name=session.get_experiment_name(), + run_id=uuid.uuid4().hex, + ) + + backend_executor = self._backend_executor_cls( + backend_config=self._backend_config, + trial_info=trial_info, + num_workers=scaling_config.num_workers, + resources_per_worker=scaling_config._resources_per_worker_not_none, + max_retries=0, + ) + + # Start the remote actors. + backend_executor.start() + + training_iterator = self._training_iterator_cls( + backend_executor=backend_executor, + backend_config=self._backend_config, + train_func=train_loop_per_worker, + datasets=self.datasets, + metadata=self.metadata, + data_config=self._data_config, + checkpoint=self.starting_checkpoint, + ) + + self._run_training(training_iterator) + + # Shutdown workers. + backend_executor.shutdown() + + def get_dataset_config(self) -> DataConfig: + """Returns a copy of this Trainer's final dataset configs. + + Returns: + The merged default + user-supplied dataset config. + """ + + return self._data_config + + @repr_with_fallback(["ipywidgets", "8"]) + def _repr_mimebundle_(self, **kwargs): + """Returns a mimebundle with an ipywidget repr and a simple text repr. + + Depending on the frontend where the data is being displayed, + different mimetypes will be used from this bundle. + See https://ipython.readthedocs.io/en/stable/config/integrating.html + for information about this method, and + https://ipywidgets.readthedocs.io/en/latest/embedding.html + for more information about the jupyter widget mimetype. + + Returns: + A mimebundle containing an ipywidget repr and a simple text repr. + """ + from ipywidgets import HTML, Layout, Tab, VBox + + title = HTML(f"

{self.__class__.__name__}

") + + children = [] + titles = [] + + if self.datasets: + children.append(self._datasets_repr_()) + titles.append("Datasets") + + children.append(HTML(self._data_config_repr_html_())) + titles.append("Data Config") + + if self._train_loop_config: + children.append(HTML(self._train_loop_config_repr_html_())) + titles.append("Train Loop Config") + + if self.scaling_config: + children.append(HTML(self.scaling_config._repr_html_())) + titles.append("Scaling Config") + + if self.run_config: + children.append(HTML(self.run_config._repr_html_())) + titles.append("Run Config") + + if self._backend_config: + children.append(HTML(self._backend_config._repr_html_())) + titles.append("Backend Config") + + tab = Tab(children, titles=titles) + widget = VBox([title, tab], layout=Layout(width="100%")) + bundle = widget._repr_mimebundle_(**kwargs) + bundle.update( + { + "text/plain": repr(self), + } + ) + return bundle + + def _train_loop_config_repr_html_(self) -> str: + if self._train_loop_config: + table_data = {} + for k, v in self._train_loop_config.items(): + if isinstance(v, str) or str(v).isnumeric(): + table_data[k] = v + elif hasattr(v, "_repr_html_"): + table_data[k] = v._repr_html_() + else: + table_data[k] = str(v) + + return Template("title_data.html.j2").render( + title="Train Loop Config", + data=Template("scrollableTable.html.j2").render( + table=tabulate( + table_data.items(), + headers=["Setting", "Value"], + showindex=False, + tablefmt="unsafehtml", + ), + max_height="none", + ), + ) + else: + return "" + + def _data_config_repr_html_(self) -> str: + # TODO make this rendering nicer. + content = [str(self._data_config)] + return Template("rendered_html_common.html.j2").render(content=content) + + def _datasets_repr_(self) -> str: + from ipywidgets import HTML, Layout, VBox + + content = [] + if self.datasets: + for name, config in self.datasets.items(): + tab = config._tab_repr_() + if tab: + content.append( + HTML( + Template("title_data.html.j2").render( + title=f"Dataset - {name}", data=None + ) + ) + ) + content.append(config._tab_repr_()) + + return VBox(content, layout=Layout(width="100%")) diff --git a/lib/python3.12/site-packages/ray/train/error.py b/lib/python3.12/site-packages/ray/train/error.py new file mode 100644 index 0000000000000000000000000000000000000000..1aa8c82471bbe8c800b2415c8af3b1aef601d00d --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/error.py @@ -0,0 +1,6 @@ +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="beta") +class SessionMisuseError(Exception): + """Indicates a method or function was used outside of a session.""" diff --git a/lib/python3.12/site-packages/ray/train/examples/__init__.py b/lib/python3.12/site-packages/ray/train/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb07580d9798f6bf2ae5e3943a7df5898c96d15b Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/__init__.py b/lib/python3.12/site-packages/ray/train/examples/horovod/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e649c33fab60c3128a0c966cc665e2a1b2aba348 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a355282dda49837b13f1d40f2b49aa5359702fe5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_pytorch_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_pytorch_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94bd14de46804ba64333ee0e859280875496932d Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_pytorch_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_tune_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_tune_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7433a318ed5370b5c628e32e674782f3d01e4a3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/horovod/__pycache__/horovod_tune_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_example.py b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_example.py new file mode 100644 index 0000000000000000000000000000000000000000..236814aa8afccd775c3b67ab96edf8d3da93a0c8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_example.py @@ -0,0 +1,286 @@ +import argparse +import os + +import horovod.torch as hvd +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import torch.utils.data.distributed +from filelock import FileLock +from torchvision import datasets, transforms + +import ray +from ray import train +from ray.train import ScalingConfig +from ray.train.horovod import HorovodTrainer + + +def metric_average(val, name): + tensor = torch.tensor(val) + avg_tensor = hvd.allreduce(tensor, name=name) + return avg_tensor.item() + + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x) + + +def setup(config): + data_dir = config.get("data_dir", None) + seed = config.get("seed", 42) + batch_size = config.get("batch_size", 64) + use_adasum = config.get("use_adasum", False) + lr = config.get("lr", 0.01) + momentum = config.get("momentum", 0.5) + use_cuda = config.get("use_cuda", False) + + # Horovod: initialize library. + hvd.init() + torch.manual_seed(seed) + + if use_cuda: + # Horovod: pin GPU to local rank. + torch.cuda.set_device(hvd.local_rank()) + torch.cuda.manual_seed(seed) + + # Horovod: limit # of CPU threads to be used per worker. + torch.set_num_threads(1) + + kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {} + data_dir = data_dir or "~/data" + with FileLock(os.path.expanduser("~/.horovod_lock")): + train_dataset = datasets.MNIST( + data_dir, + train=True, + download=True, + transform=transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ), + ) + # Horovod: use DistributedSampler to partition the training data. + train_sampler = torch.utils.data.distributed.DistributedSampler( + train_dataset, num_replicas=hvd.size(), rank=hvd.rank() + ) + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs + ) + + model = Net() + + # By default, Adasum doesn't need scaling up learning rate. + lr_scaler = hvd.size() if not use_adasum else 1 + + if use_cuda: + # Move model to GPU. + model.cuda() + # If using GPU Adasum allreduce, scale learning rate by local_size. + if use_adasum and hvd.nccl_built(): + lr_scaler = hvd.local_size() + + # Horovod: scale learning rate by lr_scaler. + optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum) + + # Horovod: wrap optimizer with DistributedOptimizer. + optimizer = hvd.DistributedOptimizer( + optimizer, + named_parameters=model.named_parameters(), + op=hvd.Adasum if use_adasum else hvd.Average, + ) + + return model, optimizer, train_loader, train_sampler + + +def train_epoch( + model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda +): + loss = None + model.train() + # Horovod: set epoch to sampler for shuffling. + train_sampler.set_epoch(epoch) + for batch_idx, (data, target) in enumerate(train_loader): + if use_cuda: + data, target = data.cuda(), target.cuda() + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % log_interval == 0: + # Horovod: use train_sampler to determine the number of + # examples in this worker's partition. + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(train_sampler), + 100.0 * batch_idx / len(train_loader), + loss.item(), + ) + ) + return loss.item() if loss else None + + +# Horovod function API. + + +def train_func(config): + num_epochs = config.get("num_epochs", 10) + log_interval = config.get("log_interval", 10) + use_cuda = config.get("use_cuda", False) + + model, optimizer, train_loader, train_sampler = setup(config) + + for epoch in range(num_epochs): + loss = train_epoch( + model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda + ) + train.report(dict(loss=loss)) + + +def main(num_workers, use_gpu, kwargs): + trainer = HorovodTrainer( + train_func, + train_loop_config=kwargs, + scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=num_workers), + ) + results = trainer.fit() + print(results.metrics) + + +# Horovod Class API. + + +class HorovodTrainClass: + def __init__(self, config): + self.log_interval = config.get("log_interval", 10) + self.use_cuda = config.get("use_cuda", False) + + if self.use_cuda: + torch.cuda.set_device(hvd.local_rank()) + + self.model, self.optimizer, self.train_loader, self.train_sampler = setup( + config + ) + + def train(self, epoch): + loss = train_epoch( + self.model, + self.optimizer, + self.train_sampler, + self.train_loader, + epoch, + self.log_interval, + self.use_cuda, + ) + return loss + + +if __name__ == "__main__": + # Training settings + parser = argparse.ArgumentParser( + description="PyTorch MNIST Example", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--batch-size", + type=int, + default=64, + metavar="N", + help="input batch size for training (default: 64)", + ) + parser.add_argument( + "--num-epochs", + type=int, + default=5, + metavar="N", + help="number of epochs to train (default: 10)", + ) + parser.add_argument( + "--lr", + type=float, + default=0.01, + metavar="LR", + help="learning rate (default: 0.01)", + ) + parser.add_argument( + "--momentum", + type=float, + default=0.5, + metavar="M", + help="SGD momentum (default: 0.5)", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="enables CUDA training" + ) + parser.add_argument( + "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)" + ) + parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", + ) + parser.add_argument( + "--use-adasum", + action="store_true", + default=False, + help="use adasum algorithm to do reduction", + ) + parser.add_argument( + "--num-workers", + type=int, + default=2, + help="Number of Ray workers to use for training.", + ) + parser.add_argument( + "--data-dir", + help="location of the training dataset in the local filesystem (" + "will be downloaded if needed)", + ) + parser.add_argument( + "--address", + required=False, + type=str, + default=None, + help="Address of Ray cluster.", + ) + + args = parser.parse_args() + + if args.address: + ray.init(args.address) + else: + ray.init() + + use_cuda = args.use_gpu if args.use_gpu is not None else False + + kwargs = { + "data_dir": args.data_dir, + "seed": args.seed, + "use_cuda": use_cuda, + "batch_size": args.batch_size, + "use_adasum": args.use_adasum if args.use_adasum else False, + "lr": args.lr, + "momentum": args.momentum, + "num_epochs": args.num_epochs, + "log_interval": args.log_interval, + } + + main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs) diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_pytorch_example.py b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_pytorch_example.py new file mode 100644 index 0000000000000000000000000000000000000000..d20ed51e9ef7e39c467948954816e91e509cff9c --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_pytorch_example.py @@ -0,0 +1,270 @@ +import argparse +import os +import tempfile + +import horovod.torch as hvd +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import torch.utils.data.distributed +from filelock import FileLock +from torchvision import datasets, transforms + +import ray.train.torch +from ray import train +from ray.train import Checkpoint, ScalingConfig +from ray.train.horovod import HorovodTrainer + + +def metric_average(val, name): + tensor = torch.tensor(val) + avg_tensor = hvd.allreduce(tensor, name=name) + return avg_tensor.item() + + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 10, kernel_size=5) + self.conv2 = nn.Conv2d(10, 20, kernel_size=5) + self.conv2_drop = nn.Dropout2d() + self.fc1 = nn.Linear(320, 50) + self.fc2 = nn.Linear(50, 10) + + def forward(self, x): + x = F.relu(F.max_pool2d(self.conv1(x), 2)) + x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) + x = x.view(-1, 320) + x = F.relu(self.fc1(x)) + x = F.dropout(x, training=self.training) + x = self.fc2(x) + return F.log_softmax(x) + + +def setup(config): + data_dir = config.get("data_dir", None) + seed = config.get("seed", 42) + batch_size = config.get("batch_size", 64) + use_adasum = config.get("use_adasum", False) + lr = config.get("lr", 0.01) + momentum = config.get("momentum", 0.5) + use_cuda = config.get("use_cuda", False) + + # Horovod: initialize library. + hvd.init() + torch.manual_seed(seed) + + if use_cuda: + # Horovod: pin GPU to local rank. + torch.cuda.set_device(hvd.local_rank()) + torch.cuda.manual_seed(seed) + + # Horovod: limit # of CPU threads to be used per worker. + torch.set_num_threads(1) + + kwargs = {"pin_memory": True} if use_cuda else {} + data_dir = data_dir or "~/data" + with FileLock(os.path.expanduser("~/.horovod_lock")): + train_dataset = datasets.MNIST( + data_dir, + train=True, + download=True, + transform=transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ), + ) + # Horovod: use DistributedSampler to partition the training data. + train_sampler = torch.utils.data.distributed.DistributedSampler( + train_dataset, num_replicas=hvd.size(), rank=hvd.rank() + ) + # Note, don't set `num_workers` in DataLoader (not even 1), + # as that will separately start multiple processes (each corresponding to 1 worker) + # to load the data. This is known to cause issues with Ray. + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs + ) + + model = Net() + + # By default, Adasum doesn't need scaling up learning rate. + lr_scaler = hvd.size() if not use_adasum else 1 + + if use_cuda: + # Move model to GPU. + model.cuda() + # If using GPU Adasum allreduce, scale learning rate by local_size. + if use_adasum and hvd.nccl_built(): + lr_scaler = hvd.local_size() + + # Horovod: scale learning rate by lr_scaler. + optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum) + + # Horovod: wrap optimizer with DistributedOptimizer. + optimizer = hvd.DistributedOptimizer( + optimizer, + named_parameters=model.named_parameters(), + op=hvd.Adasum if use_adasum else hvd.Average, + ) + + return model, optimizer, train_loader, train_sampler + + +def train_epoch( + model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda +): + loss = None + model.train() + # Horovod: set epoch to sampler for shuffling. + train_sampler.set_epoch(epoch) + for batch_idx, (data, target) in enumerate(train_loader): + if use_cuda: + data, target = data.cuda(), target.cuda() + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % log_interval == 0: + # Horovod: use train_sampler to determine the number of + # examples in this worker's partition. + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(train_sampler), + 100.0 * batch_idx / len(train_loader), + loss.item(), + ) + ) + return loss.item() if loss else None + + +def train_func(config): + num_epochs = config.get("num_epochs", 10) + log_interval = config.get("log_interval", 10) + use_cuda = config.get("use_cuda", False) + + model, optimizer, train_loader, train_sampler = setup(config) + + results = [] + for epoch in range(num_epochs): + loss = train_epoch( + model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda + ) + results.append(loss) + with tempfile.TemporaryDirectory() as tmpdir: + torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt")) + train.report({"loss": loss}, checkpoint=Checkpoint.from_directory(tmpdir)) + + # Only used for testing. + return results + + +def main(num_workers, use_gpu, kwargs): + trainer = HorovodTrainer( + train_loop_per_worker=train_func, + train_loop_config={ + "num_epochs": kwargs["num_epochs"], + "log_interval": kwargs["log_interval"], + "use_cuda": kwargs["use_cuda"], + }, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + ) + result = trainer.fit() + print(result) + + +if __name__ == "__main__": + # Training settings + parser = argparse.ArgumentParser( + description="PyTorch MNIST Example", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--batch-size", + type=int, + default=64, + metavar="N", + help="input batch size for training (default: 64)", + ) + parser.add_argument( + "--num-epochs", + type=int, + default=5, + metavar="N", + help="number of epochs to train (default: 10)", + ) + parser.add_argument( + "--lr", + type=float, + default=0.01, + metavar="LR", + help="learning rate (default: 0.01)", + ) + parser.add_argument( + "--momentum", + type=float, + default=0.5, + metavar="M", + help="SGD momentum (default: 0.5)", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="enables CUDA training" + ) + parser.add_argument( + "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)" + ) + parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", + ) + parser.add_argument( + "--use-adasum", + action="store_true", + default=False, + help="use adasum algorithm to do reduction", + ) + parser.add_argument( + "--num-workers", + type=int, + default=2, + help="Number of Ray workers to use for training.", + ) + parser.add_argument( + "--data-dir", + help="location of the training dataset in the local filesystem (" + "will be downloaded if needed)", + ) + parser.add_argument( + "--address", + required=False, + type=str, + default=None, + help="Address of Ray cluster.", + ) + + args = parser.parse_args() + + if args.address: + ray.init(args.address) + else: + ray.init() + + use_cuda = args.use_gpu if args.use_gpu is not None else False + + kwargs = { + "data_dir": args.data_dir, + "seed": args.seed, + "use_cuda": use_cuda, + "batch_size": args.batch_size, + "use_adasum": args.use_adasum if args.use_adasum else False, + "lr": args.lr, + "momentum": args.momentum, + "num_epochs": args.num_epochs, + "log_interval": args.log_interval, + } + + main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs) diff --git a/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_tune_example.py b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_tune_example.py new file mode 100644 index 0000000000000000000000000000000000000000..9433d50635ad0a8db585aefc454dd6c9649685a3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/horovod/horovod_tune_example.py @@ -0,0 +1,139 @@ +import time + +import numpy as np +import torch + +import ray +import ray.train.torch +from ray import train, tune +from ray.train import ScalingConfig +from ray.train.horovod import HorovodTrainer +from ray.tune.tune_config import TuneConfig +from ray.tune.tuner import Tuner + + +def sq(x): + m2 = 1.0 + m1 = -20.0 + m0 = 50.0 + return m2 * x * x + m1 * x + m0 + + +def qu(x): + m3 = 10.0 + m2 = 5.0 + m1 = -20.0 + m0 = -5.0 + return m3 * x * x * x + m2 * x * x + m1 * x + m0 + + +class Net(torch.nn.Module): + def __init__(self, mode="sq"): + super(Net, self).__init__() + + if mode == "square": + self.mode = 0 + self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0])) + else: + self.mode = 1 + self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0, 1.0])) + + def forward(self, x): + if ~self.mode: + return x * x + self.param[0] * x + self.param[1] + else: + return_val = 10 * x * x * x + return_val += self.param[0] * x * x + return_val += self.param[1] * x + self.param[2] + return return_val + + +def train_loop_per_worker(config): + import horovod.torch as hvd + import torch + + hvd.init() + device = ray.train.torch.get_device() + mode = config["mode"] + net = Net(mode).to(device) + optimizer = torch.optim.SGD( + net.parameters(), + lr=config["lr"], + ) + optimizer = hvd.DistributedOptimizer(optimizer) + + num_steps = 5 + print(hvd.size()) + np.random.seed(1 + hvd.rank()) + torch.manual_seed(1234) + # To ensure consistent initialization across workers, + hvd.broadcast_parameters(net.state_dict(), root_rank=0) + hvd.broadcast_optimizer_state(optimizer, root_rank=0) + + start = time.time() + x_max = config["x_max"] + for step in range(1, num_steps + 1): + features = torch.Tensor(np.random.rand(1) * 2 * x_max - x_max).to(device) + if mode == "square": + labels = sq(features) + else: + labels = qu(features) + optimizer.zero_grad() + outputs = net(features) + loss = torch.nn.MSELoss()(outputs, labels) + loss.backward() + + optimizer.step() + time.sleep(0.1) + train.report(dict(loss=loss.item())) + total = time.time() - start + print(f"Took {total:0.3f} s. Avg: {total / num_steps:0.3f} s.") + + +def tune_horovod(num_workers, num_samples, use_gpu, mode="square", x_max=1.0): + horovod_trainer = HorovodTrainer( + train_loop_per_worker=train_loop_per_worker, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + train_loop_config={"mode": mode, "x_max": x_max}, + ) + + tuner = Tuner( + horovod_trainer, + param_space={"train_loop_config": {"lr": tune.uniform(0.1, 1)}}, + tune_config=TuneConfig(mode="min", metric="loss", num_samples=num_samples), + _tuner_kwargs={"fail_fast": True}, + ) + + result_grid = tuner.fit() + + print("Best hyperparameters found were: ", result_grid.get_best_result().config) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--mode", type=str, default="square", choices=["square", "cubic"] + ) + parser.add_argument( + "--learning_rate", type=float, default=0.1, dest="learning_rate" + ) + parser.add_argument("--x_max", type=float, default=1.0, dest="x_max") + parser.add_argument("--gpu", action="store_true") + parser.add_argument( + "--smoke-test", action="store_true", help=("Finish quickly for testing.") + ) + parser.add_argument("--num-workers", type=int, default=2) + args, _ = parser.parse_known_args() + + if args.smoke_test: + ray.init(num_cpus=3) + + tune_horovod( + num_workers=args.num_workers, + num_samples=2 if args.smoke_test else 10, + use_gpu=args.gpu, + mode=args.mode, + x_max=args.x_max, + ) diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__init__.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..294d9049398b02035e575e907f0ed97c1b0961f8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_fashion_mnist_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_fashion_mnist_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9744e026c026cd0399f170d9a7f8ff3dd90971fd Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_fashion_mnist_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_linear_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_linear_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aba1823bdcc7f38fec95e02d3233b10d8b9cc66d Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_linear_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_quick_start.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_quick_start.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..839c66c75c50294b8a60023545573335e55bd1fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_quick_start.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_regression_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_regression_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e69a2334f5462d24f5a716fa842d99ea14f1e67 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/__pycache__/torch_regression_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__init__.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87e9eb4d91eaed0de4fd44c0ad5198ff56784b1c Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/auto_pipeline_for_host_to_device_data_transfer.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/auto_pipeline_for_host_to_device_data_transfer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4ab748a91a74a1a4ce6f0b34e2b323315db5c96 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/__pycache__/auto_pipeline_for_host_to_device_data_transfer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/auto_pipeline_for_host_to_device_data_transfer.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/auto_pipeline_for_host_to_device_data_transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..28fe7461bc3c30dd1631bc30454b6a5dc64a9696 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_data_prefetch_benchmark/auto_pipeline_for_host_to_device_data_transfer.py @@ -0,0 +1,161 @@ +# The PyTorch data transfer benchmark script. +import argparse +import warnings + +import numpy as np +import torch +import torch.nn as nn + +import ray.train as train +from ray.train import ScalingConfig +from ray.train.torch import TorchTrainer + + +class Net(nn.Module): + def __init__(self, in_d, hidden): + # output dim = 1 + super(Net, self).__init__() + dims = [in_d] + hidden + [1] + self.layers = nn.ModuleList( + [nn.Linear(dims[i - 1], dims[i]) for i in range(len(dims))] + ) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class BenchmarkDataset(torch.utils.data.Dataset): + """Create a naive dataset for the benchmark""" + + def __init__(self, dim, size=1000): + self.x = torch.from_numpy(np.random.normal(size=(size, dim))).float() + self.y = torch.from_numpy(np.random.normal(size=(size, 1))).float() + self.size = size + + def __getitem__(self, index): + return self.x[index, None], self.y[index, None] + + def __len__(self): + return self.size + + +def train_epoch(epoch, dataloader, model, loss_fn, optimizer): + if train.get_context().get_world_size() > 1: + dataloader.sampler.set_epoch(epoch) + + for X, y in dataloader: + # Compute prediction error + pred = model(X) + loss = loss_fn(pred, y) + + # Backpropagation + optimizer.zero_grad() + loss.backward() + optimizer.step() + + +def train_func(config): + data_size = config.get("data_size", 4096 * 50) + batch_size = config.get("batch_size", 4096) + hidden_size = config.get("hidden_size", 1) + use_auto_transfer = config.get("use_auto_transfer", False) + lr = config.get("lr", 1e-2) + epochs = config.get("epochs", 10) + + train_dataset = BenchmarkDataset(4096, size=data_size) + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=batch_size, shuffle=True + ) + + train_loader = train.torch.prepare_data_loader( + data_loader=train_loader, move_to_device=True, auto_transfer=use_auto_transfer + ) + + model = Net(in_d=4096, hidden=[4096] * hidden_size) + model = train.torch.prepare_model(model) + + loss_fn = nn.MSELoss() + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + choice = "with" if use_auto_transfer else "without" + print(f"Starting the torch data prefetch benchmark {choice} auto pipeline...") + + torch.cuda.synchronize() + start.record() + for epoch in range(epochs): + train_epoch(epoch, train_loader, model, loss_fn, optimizer) + end.record() + torch.cuda.synchronize() + + print( + f"Finished the torch data prefetch benchmark {choice} " + f"auto pipeline: {start.elapsed_time(end)} ms." + ) + + return "Experiment done." + + +def train_linear(num_workers=1, num_hidden_layers=1, use_auto_transfer=True, epochs=3): + config = { + "lr": 1e-2, + "hidden_size": num_hidden_layers, + "batch_size": 4096, + "epochs": epochs, + "use_auto_transfer": use_auto_transfer, + } + trainer = TorchTrainer( + train_func, + train_loop_config=config, + scaling_config=ScalingConfig(use_gpu=True, num_workers=num_workers), + ) + results = trainer.fit() + + print(results.metrics) + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--epochs", type=int, default=1, help="Number of epochs to train for." + ) + parser.add_argument( + "--num_hidden_layers", + type=int, + default=1, + help="Number of epochs to train for.", + ) + + args, _ = parser.parse_known_args() + + import ray + + ray.init(address=args.address) + + if not torch.cuda.is_available(): + warnings.warn("GPU is not available. Skip the test using auto pipeline.") + else: + train_linear( + num_workers=1, + num_hidden_layers=args.num_hidden_layers, + use_auto_transfer=True, + epochs=args.epochs, + ) + + torch.cuda.empty_cache() + train_linear( + num_workers=1, + num_hidden_layers=args.num_hidden_layers, + use_auto_transfer=False, + epochs=args.epochs, + ) + + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_fashion_mnist_example.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_fashion_mnist_example.py new file mode 100644 index 0000000000000000000000000000000000000000..dc5de0d9004a8ae0404b5fe561541147a6c4a51e --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_fashion_mnist_example.py @@ -0,0 +1,154 @@ +import os +from typing import Dict + +import torch +from filelock import FileLock +from torch import nn +from torch.utils.data import DataLoader +from torchvision import datasets, transforms +from torchvision.transforms import Normalize, ToTensor +from tqdm import tqdm + +import ray.train +from ray.train import ScalingConfig +from ray.train.torch import TorchTrainer + + +def get_dataloaders(batch_size): + # Transform to normalize the input images + transform = transforms.Compose([ToTensor(), Normalize((0.28604,), (0.32025,))]) + + with FileLock(os.path.expanduser("~/data.lock")): + # Download training data from open datasets + training_data = datasets.FashionMNIST( + root="~/data", + train=True, + download=True, + transform=transform, + ) + + # Download test data from open datasets + test_data = datasets.FashionMNIST( + root="~/data", + train=False, + download=True, + transform=transform, + ) + + # Create data loaders + train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True) + test_dataloader = DataLoader(test_data, batch_size=batch_size) + + return train_dataloader, test_dataloader + + +# Model Definition +class NeuralNetwork(nn.Module): + def __init__(self): + super(NeuralNetwork, self).__init__() + self.flatten = nn.Flatten() + self.linear_relu_stack = nn.Sequential( + nn.Linear(28 * 28, 512), + nn.ReLU(), + nn.Dropout(0.25), + nn.Linear(512, 512), + nn.ReLU(), + nn.Dropout(0.25), + nn.Linear(512, 10), + nn.ReLU(), + ) + + def forward(self, x): + x = self.flatten(x) + logits = self.linear_relu_stack(x) + return logits + + +def train_func_per_worker(config: Dict): + ray.train.torch.enable_reproducibility() + + lr = config["lr"] + epochs = config["epochs"] + batch_size = config["batch_size_per_worker"] + + # Get dataloaders inside the worker training function + train_dataloader, test_dataloader = get_dataloaders(batch_size=batch_size) + + # [1] Prepare Dataloader for distributed training + # Shard the datasets among workers and move batches to the correct device + # ======================================================================= + train_dataloader = ray.train.torch.prepare_data_loader(train_dataloader) + test_dataloader = ray.train.torch.prepare_data_loader(test_dataloader) + + model = NeuralNetwork() + + # [2] Prepare and wrap your model with DistributedDataParallel + # Move the model to the correct GPU/CPU device + # ============================================================ + model = ray.train.torch.prepare_model(model) + + loss_fn = nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) + + # Model training loop + for epoch in range(epochs): + if ray.train.get_context().get_world_size() > 1: + # Required for the distributed sampler to shuffle properly across epochs. + train_dataloader.sampler.set_epoch(epoch) + + model.train() + for X, y in tqdm(train_dataloader, desc=f"Train Epoch {epoch}"): + pred = model(X) + loss = loss_fn(pred, y) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + model.eval() + test_loss, num_correct, num_total = 0, 0, 0 + with torch.no_grad(): + for X, y in tqdm(test_dataloader, desc=f"Test Epoch {epoch}"): + pred = model(X) + loss = loss_fn(pred, y) + + test_loss += loss.item() + num_total += y.shape[0] + num_correct += (pred.argmax(1) == y).sum().item() + + test_loss /= len(test_dataloader) + accuracy = num_correct / num_total + + # [3] Report metrics to Ray Train + # =============================== + ray.train.report(metrics={"loss": test_loss, "accuracy": accuracy}) + + +def train_fashion_mnist(num_workers=2, use_gpu=False): + global_batch_size = 32 + + train_config = { + "lr": 1e-3, + "epochs": 10, + "batch_size_per_worker": global_batch_size // num_workers, + } + + # Configure computation resources + scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu) + + # Initialize a Ray TorchTrainer + trainer = TorchTrainer( + train_loop_per_worker=train_func_per_worker, + train_loop_config=train_config, + scaling_config=scaling_config, + ) + + # [4] Start distributed training + # Run `train_func_per_worker` on all workers + # ============================================= + result = trainer.fit() + print(f"Training result: {result}") + + +if __name__ == "__main__": + train_fashion_mnist(num_workers=4, use_gpu=True) diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_linear_example.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_linear_example.py new file mode 100644 index 0000000000000000000000000000000000000000..19d5848473b91719f488d19a3364ddc55da6d899 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_linear_example.py @@ -0,0 +1,147 @@ +import argparse +import os +import tempfile + +import numpy as np +import torch +import torch.nn as nn + +import ray.train as train +from ray.train import Checkpoint, RunConfig, ScalingConfig +from ray.train.torch import TorchTrainer + + +class LinearDataset(torch.utils.data.Dataset): + """y = a * x + b""" + + def __init__(self, a, b, size=1000): + x = np.arange(0, 10, 10 / size, dtype=np.float32) + self.x = torch.from_numpy(x) + self.y = torch.from_numpy(a * x + b) + + def __getitem__(self, index): + return self.x[index, None], self.y[index, None] + + def __len__(self): + return len(self.x) + + +def train_epoch(epoch, dataloader, model, loss_fn, optimizer): + if train.get_context().get_world_size() > 1: + dataloader.sampler.set_epoch(epoch) + + for X, y in dataloader: + # Compute prediction error + pred = model(X) + loss = loss_fn(pred, y) + + # Backpropagation + optimizer.zero_grad() + loss.backward() + optimizer.step() + + +def validate_epoch(dataloader, model, loss_fn): + num_batches = len(dataloader) + model.eval() + loss = 0 + with torch.no_grad(): + for X, y in dataloader: + pred = model(X) + loss += loss_fn(pred, y).item() + loss /= num_batches + import copy + + model_copy = copy.deepcopy(model) + return model_copy.cpu().state_dict(), loss + + +def train_func(config): + data_size = config.get("data_size", 1000) + val_size = config.get("val_size", 400) + batch_size = config.get("batch_size", 32) + hidden_size = config.get("hidden_size", 1) + lr = config.get("lr", 1e-2) + epochs = config.get("epochs", 3) + + train_dataset = LinearDataset(2, 5, size=data_size) + val_dataset = LinearDataset(2, 5, size=val_size) + train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size) + validation_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size) + + train_loader = train.torch.prepare_data_loader(train_loader) + validation_loader = train.torch.prepare_data_loader(validation_loader) + + model = nn.Linear(1, hidden_size) + model = train.torch.prepare_model(model) + + loss_fn = nn.MSELoss() + + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + + results = [] + for epoch in range(epochs): + train_epoch(epoch, train_loader, model, loss_fn, optimizer) + state_dict, loss = validate_epoch(validation_loader, model, loss_fn) + result = dict(loss=loss) + results.append(result) + + with tempfile.TemporaryDirectory() as tmpdir: + torch.save(state_dict, os.path.join(tmpdir, "model.pt")) + train.report(result, checkpoint=Checkpoint.from_directory(tmpdir)) + + return results + + +def train_linear(num_workers=2, use_gpu=False, epochs=3, storage_path=None): + config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs} + trainer = TorchTrainer( + train_loop_per_worker=train_func, + train_loop_config=config, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + run_config=RunConfig(storage_path=storage_path), + ) + result = trainer.fit() + + print(result.metrics) + return result.metrics + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--use-gpu", action="store_true", help="Whether to use GPU for training." + ) + parser.add_argument( + "--epochs", type=int, default=3, help="Number of epochs to train for." + ) + parser.add_argument( + "--smoke-test", + action="store_true", + default=False, + help="Finish quickly for testing.", + ) + + args, _ = parser.parse_known_args() + + import ray + + if args.smoke_test: + # 2 workers + 1 for trainer. + ray.init(num_cpus=3) + train_linear() + else: + ray.init(address=args.address) + train_linear( + num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs + ) diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_quick_start.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_quick_start.py new file mode 100644 index 0000000000000000000000000000000000000000..df1ae3461bd7aaac9d5c5d98c604115df36d3760 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_quick_start.py @@ -0,0 +1,110 @@ +# ruff: noqa +# fmt: off +# isort: skip_file + +# __torch_setup_begin__ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from torchvision import datasets +from torchvision.transforms import ToTensor + +def get_dataset(): + return datasets.FashionMNIST( + root="/tmp/data", + train=True, + download=True, + transform=ToTensor(), + ) + +class NeuralNetwork(nn.Module): + def __init__(self): + super().__init__() + self.flatten = nn.Flatten() + self.linear_relu_stack = nn.Sequential( + nn.Linear(28 * 28, 512), + nn.ReLU(), + nn.Linear(512, 512), + nn.ReLU(), + nn.Linear(512, 10), + ) + + def forward(self, inputs): + inputs = self.flatten(inputs) + logits = self.linear_relu_stack(inputs) + return logits +# __torch_setup_end__ + +# __torch_single_begin__ +def train_func(): + num_epochs = 3 + batch_size = 64 + + dataset = get_dataset() + dataloader = DataLoader(dataset, batch_size=batch_size) + + model = NeuralNetwork() + + criterion = nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + for epoch in range(num_epochs): + for inputs, labels in dataloader: + optimizer.zero_grad() + pred = model(inputs) + loss = criterion(pred, labels) + loss.backward() + optimizer.step() + print(f"epoch: {epoch}, loss: {loss.item()}") +# __torch_single_end__ + +# __torch_distributed_begin__ +import ray.train.torch + +def train_func_distributed(): + num_epochs = 3 + batch_size = 64 + + dataset = get_dataset() + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) + dataloader = ray.train.torch.prepare_data_loader(dataloader) + + model = NeuralNetwork() + model = ray.train.torch.prepare_model(model) + + criterion = nn.CrossEntropyLoss() + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + for epoch in range(num_epochs): + if ray.train.get_context().get_world_size() > 1: + dataloader.sampler.set_epoch(epoch) + + for inputs, labels in dataloader: + optimizer.zero_grad() + pred = model(inputs) + loss = criterion(pred, labels) + loss.backward() + optimizer.step() + print(f"epoch: {epoch}, loss: {loss.item()}") +# __torch_distributed_end__ + + +if __name__ == "__main__": + # __torch_single_run_begin__ + train_func() + # __torch_single_run_end__ + + # __torch_trainer_begin__ + from ray.train.torch import TorchTrainer + from ray.train import ScalingConfig + + # For GPU Training, set `use_gpu` to True. + use_gpu = False + + trainer = TorchTrainer( + train_func_distributed, + scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu) + ) + + results = trainer.fit() + # __torch_trainer_end__ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_regression_example.py b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_regression_example.py new file mode 100644 index 0000000000000000000000000000000000000000..4586e954f37c980d646257a82a3f2e1e5541af35 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch/torch_regression_example.py @@ -0,0 +1,159 @@ +import argparse +import os +import tempfile +from typing import Tuple + +import pandas as pd +import torch +import torch.nn as nn + +import ray +import ray.train as train +from ray.data import Dataset +from ray.train import Checkpoint, DataConfig, ScalingConfig +from ray.train.torch import TorchTrainer + + +def get_datasets(split: float = 0.7) -> Tuple[Dataset]: + dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv") + + def combine_x(batch): + return pd.DataFrame( + { + "x": batch[[f"x{i:03d}" for i in range(100)]].values.tolist(), + "y": batch["y"], + } + ) + + dataset = dataset.map_batches(combine_x, batch_format="pandas") + train_dataset, validation_dataset = dataset.repartition( + num_blocks=4 + ).train_test_split(split, shuffle=True) + return train_dataset, validation_dataset + + +def train_epoch(iterable_dataset, model, loss_fn, optimizer, device): + model.train() + for X, y in iterable_dataset: + X = X.to(device) + y = y.to(device) + + # Compute prediction error + pred = model(X) + loss = loss_fn(pred, y) + + # Backpropagation + optimizer.zero_grad() + loss.backward() + optimizer.step() + + +def validate_epoch(iterable_dataset, model, loss_fn, device): + num_batches = 0 + model.eval() + loss = 0 + with torch.no_grad(): + for X, y in iterable_dataset: + X = X.to(device) + y = y.to(device) + num_batches += 1 + pred = model(X) + loss += loss_fn(pred, y).item() + loss /= num_batches + result = {"loss": loss} + return result + + +def train_func(config): + batch_size = config.get("batch_size", 32) + hidden_size = config.get("hidden_size", 10) + lr = config.get("lr", 1e-2) + epochs = config.get("epochs", 3) + + train_dataset_shard = train.get_dataset_shard("train") + validation_dataset = train.get_dataset_shard("validation") + + model = nn.Sequential( + nn.Linear(100, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1) + ) + model = train.torch.prepare_model(model) + + loss_fn = nn.L1Loss() + + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + + results = [] + + def create_torch_iterator(shard): + iterator = shard.iter_torch_batches(batch_size=batch_size) + for batch in iterator: + yield batch["x"].float(), batch["y"].float() + + for _ in range(epochs): + train_torch_dataset = create_torch_iterator(train_dataset_shard) + validation_torch_dataset = create_torch_iterator(validation_dataset) + + device = train.torch.get_device() + + train_epoch(train_torch_dataset, model, loss_fn, optimizer, device) + if train.get_context().get_world_rank() == 0: + result = validate_epoch(validation_torch_dataset, model, loss_fn, device) + else: + result = {} + results.append(result) + + with tempfile.TemporaryDirectory() as tmpdir: + torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt")) + train.report(result, checkpoint=Checkpoint.from_directory(tmpdir)) + + return results + + +def train_regression(num_workers=2, use_gpu=False): + train_dataset, val_dataset = get_datasets() + config = {"lr": 1e-2, "hidden_size": 20, "batch_size": 4, "epochs": 3} + + trainer = TorchTrainer( + train_loop_per_worker=train_func, + train_loop_config=config, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + datasets={"train": train_dataset, "validation": val_dataset}, + dataset_config=DataConfig(datasets_to_split=["train"]), + ) + + result = trainer.fit() + return result + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--smoke-test", + action="store_true", + default=False, + help="Finish quickly for testing.", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="Use GPU for training." + ) + + args, _ = parser.parse_known_args() + + if args.smoke_test: + # 2 workers, 1 for trainer, 1 for datasets + ray.init(num_cpus=4) + result = train_regression() + else: + ray.init(address=args.address) + result = train_regression(num_workers=args.num_workers, use_gpu=args.use_gpu) + print(result) diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__init__.py b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1dcb15a633ee2711482e296864571f3bb18a1df Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/distributed_sage_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/distributed_sage_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7602af134c97f2397b2a92c8cdb6f07186345baf Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/__pycache__/distributed_sage_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/distributed_sage_example.py b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/distributed_sage_example.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3b1e126d05a266b43496fa9ec3c097c3aaba19 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/pytorch_geometric/distributed_sage_example.py @@ -0,0 +1,228 @@ +# Adapted from https://github.com/pyg-team/pytorch_geometric/blob/2.1.0 +# /examples/multi_gpu/distributed_sampling.py + +import argparse +import os + +import torch +import torch.nn.functional as F +from filelock import FileLock +from torch_geometric.datasets import FakeDataset, Reddit +from torch_geometric.loader import NeighborSampler +from torch_geometric.nn import SAGEConv +from torch_geometric.transforms import RandomNodeSplit + +from ray import train +from ray.train import ScalingConfig +from ray.train.torch import TorchTrainer + + +class SAGE(torch.nn.Module): + def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2): + super().__init__() + self.num_layers = num_layers + + self.convs = torch.nn.ModuleList() + self.convs.append(SAGEConv(in_channels, hidden_channels)) + for _ in range(self.num_layers - 2): + self.convs.append(SAGEConv(hidden_channels, hidden_channels)) + self.convs.append(SAGEConv(hidden_channels, out_channels)) + + def forward(self, x, adjs): + for i, (edge_index, _, size) in enumerate(adjs): + x_target = x[: size[1]] # Target nodes are always placed first. + x = self.convs[i]((x, x_target), edge_index) + if i != self.num_layers - 1: + x = F.relu(x) + x = F.dropout(x, p=0.5, training=self.training) + return x.log_softmax(dim=-1) + + @torch.no_grad() + def test(self, x_all, subgraph_loader): + for i in range(self.num_layers): + xs = [] + for batch_size, n_id, adj in subgraph_loader: + edge_index, _, size = adj + x = x_all[n_id.to(x_all.device)].to(train.torch.get_device()) + x_target = x[: size[1]] + x = self.convs[i]((x, x_target), edge_index) + if i != self.num_layers - 1: + x = F.relu(x) + xs.append(x.cpu()) + + x_all = torch.cat(xs, dim=0) + + return x_all + + +def train_loop_per_worker(train_loop_config): + dataset = train_loop_config["dataset_fn"]() + batch_size = train_loop_config["batch_size"] + num_epochs = train_loop_config["num_epochs"] + + data = dataset[0] + train_idx = data.train_mask.nonzero(as_tuple=False).view(-1) + train_idx = train_idx.split( + train_idx.size(0) // train.get_context().get_world_size() + )[train.get_context().get_world_rank()] + + train_loader = NeighborSampler( + data.edge_index, + node_idx=train_idx, + sizes=[25, 10], + batch_size=batch_size, + shuffle=True, + ) + + # Disable distributed sampler since the train_loader has already been split above. + train_loader = train.torch.prepare_data_loader(train_loader, add_dist_sampler=False) + + # Do validation on rank 0 worker only. + if train.get_context().get_world_rank() == 0: + subgraph_loader = NeighborSampler( + data.edge_index, node_idx=None, sizes=[-1], batch_size=2048, shuffle=False + ) + subgraph_loader = train.torch.prepare_data_loader( + subgraph_loader, add_dist_sampler=False + ) + + model = SAGE(dataset.num_features, 256, dataset.num_classes) + model = train.torch.prepare_model(model) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + + x, y = data.x.to(train.torch.get_device()), data.y.to(train.torch.get_device()) + + for epoch in range(num_epochs): + model.train() + + # ``batch_size`` is the number of samples in the current batch. + # ``n_id`` are the ids of all the nodes used in the computation. This is + # needed to pull in the necessary features just for the current batch that is + # being trained on. + # ``adjs`` is a list of 3 element tuple consisting of ``(edge_index, e_id, + # size)`` for each sample in the batch, where ``edge_index``represent the + # edges of the sampled subgraph, ``e_id`` are the ids of the edges in the + # sample, and ``size`` holds the shape of the subgraph. + # See ``torch_geometric.loader.neighbor_sampler.NeighborSampler`` for more info. + for batch_size, n_id, adjs in train_loader: + optimizer.zero_grad() + out = model(x[n_id], adjs) + loss = F.nll_loss(out, y[n_id[:batch_size]]) + loss.backward() + optimizer.step() + + if train.get_context().get_world_rank() == 0: + print(f"Epoch: {epoch:03d}, Loss: {loss:.4f}") + + train_accuracy = validation_accuracy = test_accuracy = None + + # Do validation on rank 0 worker only. + if train.get_context().get_world_rank() == 0: + model.eval() + with torch.no_grad(): + out = model.module.test(x, subgraph_loader) + res = out.argmax(dim=-1) == data.y + train_accuracy = int(res[data.train_mask].sum()) / int( + data.train_mask.sum() + ) + validation_accuracy = int(res[data.val_mask].sum()) / int( + data.val_mask.sum() + ) + test_accuracy = int(res[data.test_mask].sum()) / int(data.test_mask.sum()) + + train.report( + dict( + train_accuracy=train_accuracy, + validation_accuracy=validation_accuracy, + test_accuracy=test_accuracy, + ) + ) + + +def gen_fake_dataset(): + """Returns a function to be called on each worker that returns a Fake Dataset.""" + + # For fake dataset, since the dataset is randomized, we create it once on the + # driver, and then send the same dataset to all the training workers. + # Use 10% of nodes for validation and 10% for testing. + fake_dataset = FakeDataset(transform=RandomNodeSplit(num_val=0.1, num_test=0.1)) + + def gen_dataset(): + return fake_dataset + + return gen_dataset + + +def gen_reddit_dataset(): + """Returns a function to be called on each worker that returns Reddit Dataset.""" + + # For Reddit dataset, we have to download the data on each node, so we create the + # dataset on each training worker. + with FileLock(os.path.expanduser("~/.reddit_dataset_lock")): + dataset = Reddit("./data/Reddit") + return dataset + + +def train_gnn( + num_workers=2, use_gpu=False, epochs=3, global_batch_size=32, dataset="reddit" +): + per_worker_batch_size = global_batch_size // num_workers + + trainer = TorchTrainer( + train_loop_per_worker=train_loop_per_worker, + train_loop_config={ + "num_epochs": epochs, + "batch_size": per_worker_batch_size, + "dataset_fn": gen_reddit_dataset + if dataset == "reddit" + else gen_fake_dataset(), + }, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + ) + result = trainer.fit() + print(result.metrics) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--use-gpu", action="store_true", help="Whether to use GPU for training." + ) + parser.add_argument( + "--epochs", type=int, default=3, help="Number of epochs to train for." + ) + parser.add_argument( + "--global-batch-size", + "-b", + type=int, + default=32, + help="Global batch size to use for training.", + ) + parser.add_argument( + "--dataset", + "-d", + type=str, + choices=["reddit", "fake"], + default="reddit", + help="The dataset to use. Either 'reddit' or 'fake' Defaults to 'reddit'.", + ) + + args, _ = parser.parse_known_args() + + train_gnn( + num_workers=args.num_workers, + use_gpu=args.use_gpu, + epochs=args.epochs, + global_batch_size=args.global_batch_size, + dataset=args.dataset, + ) diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/__init__.py b/lib/python3.12/site-packages/ray/train/examples/tf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/__pycache__/tensorflow_regression_example.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/examples/tf/__pycache__/tensorflow_regression_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df7575b42f16a3d46c7752e0cfc992165bb1473b Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/examples/tf/__pycache__/tensorflow_regression_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_autoencoder_example.py b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_autoencoder_example.py new file mode 100644 index 0000000000000000000000000000000000000000..6d976203efdc9fee5bbfbb27a8aa39d080bea994 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_autoencoder_example.py @@ -0,0 +1,174 @@ +# This example showcases how to use Tensorflow with Ray Train. +# Original code: +# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras +# https://blog.keras.io/building-autoencoders-in-keras.html +import argparse + +import numpy as np +import pandas as pd +import tensorflow as tf +import tensorflow_datasets as tfds + +import ray +from ray import train +from ray.air.integrations.keras import ReportCheckpointCallback +from ray.data.datasource import SimpleTensorFlowDatasource +from ray.data.extensions import TensorArray +from ray.train import Result, ScalingConfig +from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard + + +def get_dataset(split_type="train"): + def dataset_factory(): + return tfds.load("mnist", split=[split_type], as_supervised=True)[0].take(128) + + dataset = ray.data.read_datasource( + SimpleTensorFlowDatasource(), dataset_factory=dataset_factory + ) + + def normalize_images(x): + x = np.float32(x.numpy()) / 255.0 + x = np.reshape(x, (-1,)) + return x + + def preprocess_dataset(batch): + return [ + (normalize_images(image), normalize_images(image)) for image, _ in batch + ] + + dataset = dataset.map_batches(preprocess_dataset) + + def convert_batch_to_pandas(batch): + + images = [TensorArray(image) for image, _ in batch] + # because we did autoencoder here + df = pd.DataFrame({"image": images, "label": images}) + return df + + dataset = dataset.map_batches(convert_batch_to_pandas) + return dataset + + +def build_autoencoder_model() -> tf.keras.Model: + model = tf.keras.Sequential( + [ + tf.keras.Input(shape=(784,)), + # encoder + tf.keras.layers.Dense(128, activation="relu"), + tf.keras.layers.Dense(64, activation="relu"), + tf.keras.layers.Dense(32, activation="relu"), + # decoder + tf.keras.layers.Dense(64, activation="relu"), + tf.keras.layers.Dense(128, activation="relu"), + tf.keras.layers.Dense(784, activation="sigmoid"), + ] + ) + return model + + +def train_func(config: dict): + + per_worker_batch_size = config.get("batch_size", 64) + epochs = config.get("epochs", 3) + + dataset_shard = train.get_dataset_shard("train") + + strategy = tf.distribute.MultiWorkerMirroredStrategy() + + with strategy.scope(): + # Model building/compiling need to be within `strategy.scope()`. + multi_worker_model = build_autoencoder_model() + learning_rate = config.get("lr", 0.001) + multi_worker_model.compile( + loss=tf.keras.losses.BinaryCrossentropy(), + optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate), + metrics=[ + "binary_crossentropy", + ], + ) + + def to_tf_dataset(dataset, batch_size): + def to_tensor_iterator(): + for batch in dataset.iter_tf_batches( + batch_size=batch_size, dtypes=tf.float32 + ): + yield batch["image"], batch["label"] + + output_signature = ( + tf.TensorSpec(shape=(None, 784), dtype=tf.float32), + tf.TensorSpec(shape=(None, 784), dtype=tf.float32), + ) + tf_dataset = tf.data.Dataset.from_generator( + to_tensor_iterator, output_signature=output_signature + ) + return prepare_dataset_shard(tf_dataset) + + results = [] + for epoch in range(epochs): + tf_dataset = to_tf_dataset( + dataset=dataset_shard, + batch_size=per_worker_batch_size, + ) + history = multi_worker_model.fit( + tf_dataset, callbacks=[ReportCheckpointCallback()] + ) + results.append(history.history) + return results + + +def train_tensorflow_mnist( + num_workers: int = 2, use_gpu: bool = False, epochs: int = 4 +) -> Result: + train_dataset = get_dataset(split_type="train") + config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs} + scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu) + trainer = TensorflowTrainer( + train_loop_per_worker=train_func, + train_loop_config=config, + datasets={"train": train_dataset}, + scaling_config=scaling_config, + ) + + results = trainer.fit() + print(results.metrics) + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="Enables GPU training" + ) + parser.add_argument( + "--epochs", type=int, default=3, help="Number of epochs to train for." + ) + parser.add_argument( + "--smoke-test", + action="store_true", + default=False, + help="Finish quickly for testing.", + ) + + args, _ = parser.parse_known_args() + + if args.smoke_test: + # 2 workers, 1 for trainer, 1 for datasets + num_gpus = args.num_workers if args.use_gpu else 0 + ray.init(num_cpus=4, num_gpus=num_gpus) + result = train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu) + else: + ray.init(address=args.address) + result = train_tensorflow_mnist( + num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs + ) + print(result) diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_mnist_example.py b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_mnist_example.py new file mode 100644 index 0000000000000000000000000000000000000000..3fd5d7c759df6723a7a16a622d70c743519ed638 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_mnist_example.py @@ -0,0 +1,135 @@ +# This example showcases how to use Tensorflow with Ray Train. +# Original code: +# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras +import argparse +import json +import os + +import numpy as np +import tensorflow as tf +from filelock import FileLock + +from ray.air.integrations.keras import ReportCheckpointCallback +from ray.train import Result, RunConfig, ScalingConfig +from ray.train.tensorflow import TensorflowTrainer + + +def mnist_dataset(batch_size: int) -> tf.data.Dataset: + with FileLock(os.path.expanduser("~/.mnist_lock")): + (x_train, y_train), _ = tf.keras.datasets.mnist.load_data() + # The `x` arrays are in uint8 and have values in the [0, 255] range. + # You need to convert them to float32 with values in the [0, 1] range. + x_train = x_train / np.float32(255) + y_train = y_train.astype(np.int64) + train_dataset = ( + tf.data.Dataset.from_tensor_slices((x_train, y_train)) + .shuffle(60000) + .repeat() + .batch(batch_size) + ) + return train_dataset + + +def build_cnn_model() -> tf.keras.Model: + model = tf.keras.Sequential( + [ + tf.keras.Input(shape=(28, 28)), + tf.keras.layers.Reshape(target_shape=(28, 28, 1)), + tf.keras.layers.Conv2D(32, 3, activation="relu"), + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(128, activation="relu"), + tf.keras.layers.Dense(10), + ] + ) + return model + + +def train_func(config: dict): + per_worker_batch_size = config.get("batch_size", 64) + epochs = config.get("epochs", 3) + steps_per_epoch = config.get("steps_per_epoch", 70) + + tf_config = json.loads(os.environ["TF_CONFIG"]) + num_workers = len(tf_config["cluster"]["worker"]) + + strategy = tf.distribute.MultiWorkerMirroredStrategy() + + global_batch_size = per_worker_batch_size * num_workers + multi_worker_dataset = mnist_dataset(global_batch_size) + + with strategy.scope(): + # Model building/compiling need to be within `strategy.scope()`. + multi_worker_model = build_cnn_model() + learning_rate = config.get("lr", 0.001) + multi_worker_model.compile( + loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), + optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate), + metrics=["accuracy"], + ) + + history = multi_worker_model.fit( + multi_worker_dataset, + epochs=epochs, + steps_per_epoch=steps_per_epoch, + callbacks=[ReportCheckpointCallback()], + ) + results = history.history + return results + + +def train_tensorflow_mnist( + num_workers: int = 2, + use_gpu: bool = False, + epochs: int = 4, + storage_path: str = None, +) -> Result: + config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs} + trainer = TensorflowTrainer( + train_loop_per_worker=train_func, + train_loop_config=config, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + run_config=RunConfig(storage_path=storage_path), + ) + results = trainer.fit() + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="Enables GPU training" + ) + parser.add_argument( + "--epochs", type=int, default=3, help="Number of epochs to train for." + ) + parser.add_argument( + "--smoke-test", + action="store_true", + default=False, + help="Finish quickly for testing.", + ) + + args, _ = parser.parse_known_args() + + import ray + + if args.smoke_test: + # 2 workers, 1 for trainer, 1 for datasets + num_gpus = args.num_workers if args.use_gpu else 0 + ray.init(num_cpus=4, num_gpus=num_gpus) + train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu) + else: + ray.init(address=args.address) + train_tensorflow_mnist( + num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs + ) diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_quick_start.py b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_quick_start.py new file mode 100644 index 0000000000000000000000000000000000000000..4b078675960230461f3c6d493d10a3f56ec0ddea --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_quick_start.py @@ -0,0 +1,87 @@ +# ruff: noqa +# fmt: off +# isort: skip_file + +# __tf_setup_begin__ +import sys +import numpy as np + +if sys.version_info >= (3, 12): + # Tensorflow is not installed for Python 3.12 because of keras compatibility. + sys.exit(0) +else: + import tensorflow as tf + +def mnist_dataset(batch_size): + (x_train, y_train), _ = tf.keras.datasets.mnist.load_data() + # The `x` arrays are in uint8 and have values in the [0, 255] range. + # You need to convert them to float32 with values in the [0, 1] range. + x_train = x_train / np.float32(255) + y_train = y_train.astype(np.int64) + train_dataset = tf.data.Dataset.from_tensor_slices( + (x_train, y_train)).shuffle(60000).repeat().batch(batch_size) + return train_dataset + + +def build_and_compile_cnn_model(): + model = tf.keras.Sequential([ + tf.keras.layers.InputLayer(input_shape=(28, 28)), + tf.keras.layers.Reshape(target_shape=(28, 28, 1)), + tf.keras.layers.Conv2D(32, 3, activation='relu'), + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(128, activation='relu'), + tf.keras.layers.Dense(10) + ]) + model.compile( + loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), + optimizer=tf.keras.optimizers.SGD(learning_rate=0.001), + metrics=['accuracy']) + return model +# __tf_setup_end__ + +# __tf_single_begin__ +def train_func(): + batch_size = 64 + single_worker_dataset = mnist_dataset(batch_size) + single_worker_model = build_and_compile_cnn_model() + single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70) +# __tf_single_end__ + +# __tf_distributed_begin__ +import json +import os + +def train_func_distributed(): + per_worker_batch_size = 64 + # This environment variable will be set by Ray Train. + tf_config = json.loads(os.environ['TF_CONFIG']) + num_workers = len(tf_config['cluster']['worker']) + + strategy = tf.distribute.MultiWorkerMirroredStrategy() + + global_batch_size = per_worker_batch_size * num_workers + multi_worker_dataset = mnist_dataset(global_batch_size) + + with strategy.scope(): + # Model building/compiling need to be within `strategy.scope()`. + multi_worker_model = build_and_compile_cnn_model() + + multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70) +# __tf_distributed_end__ + +if __name__ == "__main__": + # __tf_single_run_begin__ + train_func() + # __tf_single_run_end__ + + # __tf_trainer_begin__ + from ray.train.tensorflow import TensorflowTrainer + from ray.train import ScalingConfig + + # For GPU Training, set `use_gpu` to True. + use_gpu = False + + trainer = TensorflowTrainer(train_func_distributed, scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu)) + + trainer.fit() + # __tf_trainer_end__ diff --git a/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_regression_example.py b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_regression_example.py new file mode 100644 index 0000000000000000000000000000000000000000..7984fa99029171a9448c4f43910889331a7d13fc --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/examples/tf/tensorflow_regression_example.py @@ -0,0 +1,111 @@ +import argparse +import sys + +import ray +from ray import train +from ray.data.preprocessors import Concatenator +from ray.train import Result, ScalingConfig + +if sys.version_info >= (3, 12): + # Skip this test in Python 3.12+ because TensorFlow is not supported. + sys.exit(0) +else: + import tensorflow as tf + + from ray.train.tensorflow import TensorflowTrainer + from ray.train.tensorflow.keras import ReportCheckpointCallback + + +def build_model() -> tf.keras.Model: + model = tf.keras.Sequential( + [ + tf.keras.layers.InputLayer(input_shape=(100,)), + tf.keras.layers.Dense(10), + tf.keras.layers.Dense(1), + ] + ) + return model + + +def train_func(config: dict): + batch_size = config.get("batch_size", 64) + epochs = config.get("epochs", 3) + + strategy = tf.distribute.MultiWorkerMirroredStrategy() + with strategy.scope(): + # Model building/compiling need to be within `strategy.scope()`. + multi_worker_model = build_model() + multi_worker_model.compile( + optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)), + loss=tf.keras.losses.mean_absolute_error, + metrics=[tf.keras.metrics.mean_squared_error], + ) + + dataset = train.get_dataset_shard("train") + + results = [] + for _ in range(epochs): + tf_dataset = dataset.to_tf( + feature_columns="x", label_columns="y", batch_size=batch_size + ) + history = multi_worker_model.fit( + tf_dataset, callbacks=[ReportCheckpointCallback()] + ) + results.append(history.history) + return results + + +def train_tensorflow_regression(num_workers: int = 2, use_gpu: bool = False) -> Result: + dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv") + columns_to_concatenate = [f"x{i:03}" for i in range(100)] + preprocessor = Concatenator(columns=columns_to_concatenate, output_column_name="x") + dataset = preprocessor.fit_transform(dataset) + + config = {"lr": 1e-3, "batch_size": 32, "epochs": 4} + scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu) + trainer = TensorflowTrainer( + train_loop_per_worker=train_func, + train_loop_config=config, + scaling_config=scaling_config, + datasets={"train": dataset}, + ) + results = trainer.fit() + print(results.metrics) + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--address", required=False, type=str, help="the address to use for Ray" + ) + parser.add_argument( + "--num-workers", + "-n", + type=int, + default=2, + help="Sets number of workers for training.", + ) + parser.add_argument( + "--use-gpu", action="store_true", default=False, help="Enables GPU training" + ) + parser.add_argument( + "--smoke-test", + action="store_true", + default=False, + help="Finish quickly for testing.", + ) + + args, _ = parser.parse_known_args() + + if args.smoke_test: + # 2 workers, 1 for trainer, 1 for datasets + num_gpus = args.num_workers if args.use_gpu else 0 + ray.init(num_cpus=4, num_gpus=num_gpus) + result = train_tensorflow_regression(num_workers=2, use_gpu=args.use_gpu) + else: + ray.init(address=args.address) + result = train_tensorflow_regression( + num_workers=args.num_workers, use_gpu=args.use_gpu + ) + print(result) diff --git a/lib/python3.12/site-packages/ray/train/horovod/__init__.py b/lib/python3.12/site-packages/ray/train/horovod/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..73f9e7ab585a1d577f589e9c795a5560d830169a --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/horovod/__init__.py @@ -0,0 +1,22 @@ +# isort: off +try: + import horovod # noqa: F401 +except ModuleNotFoundError: + raise ModuleNotFoundError( + "Horovod isn't installed. To install Horovod with PyTorch support, run 'pip " + "install 'horovod[pytorch]''. To install Horovod with TensorFlow support, " + "run 'pip install 'horovod[tensorflow]''." + ) +# isort: on + +from ray.train.horovod.config import HorovodConfig +from ray.train.horovod.horovod_trainer import HorovodTrainer +from ray.train.v2._internal.constants import is_v2_enabled + +if is_v2_enabled(): + from ray.train.v2.horovod.horovod_trainer import HorovodTrainer # noqa: F811 + +__all__ = ["HorovodConfig", "HorovodTrainer"] + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/horovod/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7445dfff1bd6429ad5b7426a7a44ab20839d6e5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/horovod/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd0584962c0454032b9db330f7db1edb01406497 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/horovod/__pycache__/horovod_trainer.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/horovod_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29d4b95625eb863d9c3e094d0b4b4b4b1d5c6474 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/horovod/__pycache__/horovod_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/horovod/config.py b/lib/python3.12/site-packages/ray/train/horovod/config.py new file mode 100644 index 0000000000000000000000000000000000000000..00f16242dd5a1f7635ab1312e6d4b32e6f336652 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/horovod/config.py @@ -0,0 +1,163 @@ +import os +from dataclasses import dataclass +from typing import Optional, Set + +from horovod.ray.runner import Coordinator +from horovod.ray.utils import detect_nics, nics_to_env_var +from horovod.runner.common.util import secret, timeout + +import ray +from ray.train._internal.utils import update_env_vars +from ray.train._internal.worker_group import Worker, WorkerGroup +from ray.train.backend import Backend, BackendConfig +from ray.util import PublicAPI + + +@PublicAPI(stability="beta") +@dataclass +class HorovodConfig(BackendConfig): + """Configurations for Horovod setup. + + See https://github.com/horovod/horovod/blob/master/horovod/runner/common/util/settings.py # noqa: E501 + + Args: + nics (Optional[Set[str]): Network interfaces that can be used for + communication. + verbose: Horovod logging verbosity. + key (Optional[str]): Secret used for communication between workers. + ssh_port (Optional[int]): Port for SSH server running on worker nodes. + ssh_identity_file (Optional[str]): Path to the identity file to + ssh into different hosts on the cluster. + ssh_str (Optional[str]): CAUTION WHEN USING THIS. Private key + file contents. Writes the private key to ssh_identity_file. + timeout_s: Timeout parameter for Gloo rendezvous. + placement_group_timeout_s: Timeout parameter for Ray + Placement Group creation. Currently unused. + """ + + nics: Optional[Set[str]] = None + verbose: int = 1 + key: Optional[str] = None + ssh_port: Optional[int] = None + ssh_identity_file: Optional[str] = None + ssh_str: Optional[str] = None + timeout_s: int = 300 + placement_group_timeout_s: int = 100 + + @property + def start_timeout(self): + return timeout.Timeout( + self.timeout_s, + message="Timed out waiting for {activity}. Please " + "check connectivity between servers. You " + "may need to increase the --start-timeout " + "parameter if you have too many servers.", + ) + + def __post_init__(self): + if self.ssh_str and not os.path.exists(self.ssh_identity_file): + with open(self.ssh_identity_file, "w") as f: + os.chmod(self.ssh_identity_file, 0o600) + f.write(self.ssh_str) + + if self.key is None: + self.key = secret.make_secret_key() + + @property + def backend_cls(self): + return _HorovodBackend + + +class _HorovodBackend(Backend): + share_cuda_visible_devices: bool = True + + def on_start(self, worker_group: WorkerGroup, backend_config: HorovodConfig): + # NOTE: Horovod backend uses V1 WorkerGroup directly instead of BaseWorkerGroup + # because it requires direct access to worker metadata (node_id, hostname) that is + # specific to the V1 implementation. Horovod does not support V2 WorkerGroup. + + # TODO(matt): Implement placement group strategies in BackendExecutor. + + # Initialize workers with Horovod environment variables + setup_futures = [] + for rank in range(len(worker_group)): + worker_node_id = worker_group.workers[rank].metadata.node_id + setup_futures.append( + worker_group.execute_single_async( + rank, + _init_env_vars, + rank, + len(worker_group), + worker_node_id, + ) + ) + ray.get(setup_futures) + + # Use Horovod Ray Coordinator + # backend_config as settings + self.coordinator = Coordinator(backend_config) + + # Get all the hostnames of all workers + node_ids = [w.metadata.node_id for w in worker_group.workers] + hostnames = [w.metadata.hostname for w in worker_group.workers] + # Register each hostname to the coordinator. assumes the hostname + # ordering is the same. + for rank, (hostname, node_id) in enumerate(zip(hostnames, node_ids)): + self.coordinator.register(hostname, node_id, rank) + all_info = self.coordinator.finalize_registration() + + setup_futures = [] + for rank, local_cross_env_var in all_info.items(): + setup_futures.append( + worker_group.execute_single_async( + rank, update_env_vars, local_cross_env_var + ) + ) + ray.get(setup_futures) + + coordinator_envs = self.coordinator.establish_rendezvous() + + # Get one worker from each host/node. + node_worker_indexes = [node_ids.index(node_id) for node_id in set(node_ids)] + node_workers = [ + _HorovodWorkerWrapper(worker_group.workers[worker_index]) + for worker_index in node_worker_indexes + ] + assert len(node_workers) == len(self.coordinator.hostnames) + + nics = detect_nics( + backend_config, + all_host_names=list(self.coordinator.hostnames), + node_workers=node_workers, + ) + coordinator_envs.update(nics_to_env_var(nics)) + + worker_group.execute(update_env_vars, coordinator_envs) + + +def _init_env_vars(world_rank: int, world_size: int, node_id: str): + """Initialize Horovod environment variables.""" + os.environ["HOROVOD_HOSTNAME"] = node_id + os.environ["HOROVOD_RANK"] = str(world_rank) + os.environ["HOROVOD_SIZE"] = str(world_size) + + +# TODO(tgaddair): temporary workaround for Horovod's worker discovery logic, +# which requires passing in an extra parameter as part of the RayExecutor +# API. This will be removed in the future as we migrate more of the +# RayExecutor utils into Ray Train. +# See: https://github.com/horovod/horovod/blob/v0.23.0/horovod/ray/driver_service.py#L9 # noqa: E501 +@dataclass +class _HorovodWorkerWrapper: + w: Worker + + @property + def execute(self): + w = self.w + + class ExecuteHandle: + def remote(self, func, *args, **kwargs): + _ = None + return w.actor._RayTrainWorker__execute.remote(func, _, *args, **kwargs) + + return ExecuteHandle() diff --git a/lib/python3.12/site-packages/ray/train/horovod/horovod_trainer.py b/lib/python3.12/site-packages/ray/train/horovod/horovod_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..93936cac759ad897a1b1cef9c66998ae1de896ae --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/horovod/horovod_trainer.py @@ -0,0 +1,202 @@ +from typing import Any, Callable, Dict, Optional, Union + +from ray.air.config import RunConfig, ScalingConfig +from ray.train import Checkpoint, DataConfig +from ray.train.data_parallel_trainer import DataParallelTrainer +from ray.train.horovod.config import HorovodConfig +from ray.train.trainer import GenDataset +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="beta") +class HorovodTrainer(DataParallelTrainer): + """A Trainer for data parallel Horovod training. + + This Trainer runs the function ``train_loop_per_worker`` on multiple Ray + Actors. These actors already have the necessary Horovod setup already + configured for distributed Horovod training. + + The ``train_loop_per_worker`` function is expected to take in either 0 or 1 + arguments: + + .. testcode:: + + def train_loop_per_worker(): + ... + + .. testcode:: + + def train_loop_per_worker(config: Dict): + ... + + If ``train_loop_per_worker`` accepts an argument, then + ``train_loop_config`` will be passed in as the argument. This is useful if you + want to tune the values in ``train_loop_config`` as hyperparameters. + + If the ``datasets`` dict contains a training dataset (denoted by + the "train" key), then it will be split into multiple dataset + shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` inside + ``train_loop_per_worker``. All the other datasets will not be split and + ``ray.train.get_dataset_shard(...)`` will return the entire Dataset. + + Inside the ``train_loop_per_worker`` function, you can use any of the + :ref:`Ray Train loop methods `. + + .. testcode:: + + from ray import train + + def train_loop_per_worker(): + # Report intermediate results for callbacks or logging and + # checkpoint data. + train.report(...) + + # Returns dict of last saved checkpoint. + train.get_checkpoint() + + # Returns the Dataset shard for the given key. + train.get_dataset_shard("my_dataset") + + # Returns the total number of workers executing training. + train.get_context().get_world_size() + + # Returns the rank of this worker. + train.get_context().get_world_rank() + + # Returns the rank of the worker on the current node. + train.get_context().get_local_rank() + + Any returns from the ``train_loop_per_worker`` will be discarded and not + used or persisted anywhere. + + You could use ``TensorflowPredictor`` or ``TorchPredictor`` in conjunction with + HorovodTrainer. You must save the model under the "model" kwarg in the + ``Checkpoint`` passed to ``train.report()``, so that it can be used by + corresponding predictors. + + Example: + + + .. testcode:: + :skipif: True + + import os + import tempfile + + import ray + import horovod.torch as hvd + import torch + import torch.nn as nn + + from ray import train + import ray.train.torch # Need this to use `train.torch.get_device()` + from ray.train import Checkpoint, ScalingConfig + from ray.train.horovod import HorovodTrainer + + # If using GPUs, set this to True. + use_gpu = False + + input_size = 1 + layer_size = 15 + output_size = 1 + num_epochs = 3 + + class NeuralNetwork(nn.Module): + def __init__(self): + super(NeuralNetwork, self).__init__() + self.layer1 = nn.Linear(input_size, layer_size) + self.relu = nn.ReLU() + self.layer2 = nn.Linear(layer_size, output_size) + def forward(self, input): + return self.layer2(self.relu(self.layer1(input))) + + def train_loop_per_worker(): + hvd.init() + dataset_shard = train.get_dataset_shard("train") + model = NeuralNetwork() + device = train.torch.get_device() + model.to(device) + loss_fn = nn.MSELoss() + lr_scaler = 1 + optimizer = torch.optim.SGD(model.parameters(), lr=0.1 * lr_scaler) + # Horovod: wrap optimizer with DistributedOptimizer. + optimizer = hvd.DistributedOptimizer( + optimizer, + named_parameters=model.named_parameters(), + op=hvd.Average, + ) + for epoch in range(num_epochs): + model.train() + for batch in dataset_shard.iter_torch_batches( + batch_size=32, dtypes=torch.float + ): + inputs, labels = torch.unsqueeze(batch["x"], 1), batch["y"] + outputs = model(inputs) + loss = loss_fn(outputs, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + print(f"epoch: {epoch}, loss: {loss.item()}") + + # Save a model checkpoint at the end of each epoch + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + ckpt_path = os.path.join(temp_checkpoint_dir, "model.pt") + torch.save(model.state_dict(), ckpt_path) + train.report( + {"loss": loss.item(), "epoch": epoch}, + checkpoint=Checkpoint.from_directory(temp_checkpoint_dir), + ) + + train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)]) + scaling_config = ScalingConfig(num_workers=3, use_gpu=use_gpu) + trainer = HorovodTrainer( + train_loop_per_worker=train_loop_per_worker, + scaling_config=scaling_config, + datasets={"train": train_dataset}, + ) + result = trainer.fit() + + Args: + train_loop_per_worker: The training function to execute. + This can either take in no arguments or a ``config`` dict. + train_loop_config: Configurations to pass into + ``train_loop_per_worker`` if it accepts an argument. + horovod_config: Configuration for setting up the Horovod backend. + If set to None, use the default configuration. This replaces the + ``backend_config`` arg of ``DataParallelTrainer``. + scaling_config: Configuration for how to scale data parallel training. + dataset_config: Configuration for dataset ingest. + run_config: Configuration for the execution of the training run. + datasets: Any Datasets to use for training. Use + the key "train" to denote which dataset is the training + dataset. + resume_from_checkpoint: A checkpoint to resume training from. + metadata: Dict that should be made available via + `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + """ + + def __init__( + self, + train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]], + *, + train_loop_config: Optional[Dict] = None, + horovod_config: Optional[HorovodConfig] = None, + scaling_config: Optional[ScalingConfig] = None, + dataset_config: Optional[DataConfig] = None, + run_config: Optional[RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + metadata: Optional[Dict[str, Any]] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + ): + super().__init__( + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + backend_config=horovod_config or HorovodConfig(), + scaling_config=scaling_config, + dataset_config=dataset_config, + run_config=run_config, + datasets=datasets, + resume_from_checkpoint=resume_from_checkpoint, + metadata=metadata, + ) diff --git a/lib/python3.12/site-packages/ray/train/huggingface/__init__.py b/lib/python3.12/site-packages/ray/train/huggingface/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/huggingface/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/huggingface/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b75ad03ee16dd5073634b78f7af288f1d784df98 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/huggingface/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/huggingface/transformers/__init__.py b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..870fe4ee94b873e077a7bf0d3bf48429e1bf8df0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__init__.py @@ -0,0 +1,12 @@ +from ray.train.huggingface.transformers._transformers_utils import ( + RayTrainReportCallback, + prepare_trainer, +) + +__all__ = [ + "RayTrainReportCallback", + "prepare_trainer", +] + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1827c89f58701c02c3033e017d2a3459a9a37358 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/_transformers_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/_transformers_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..537f868bf90d9692a9463e9466a6884440de7d68 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/huggingface/transformers/__pycache__/_transformers_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/huggingface/transformers/_transformers_utils.py b/lib/python3.12/site-packages/ray/train/huggingface/transformers/_transformers_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7f3eaeefac4af79be05d64d51ad50267666369f0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/huggingface/transformers/_transformers_utils.py @@ -0,0 +1,150 @@ +import logging +import shutil +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Iterator, Optional, Type, Union + +from torch.utils.data import DataLoader, Dataset, IterableDataset + +import ray +from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag +from ray.data.iterator import _IterableFromIterator +from ray.train import Checkpoint +from ray.util import PublicAPI + +logger = logging.getLogger(__name__) + + +TRANSFORMERS_IMPORT_ERROR: Optional[ImportError] = None + +try: + import transformers.trainer + from transformers import Trainer + from transformers.trainer_callback import TrainerCallback +except ImportError as e: + TRANSFORMERS_IMPORT_ERROR = e + TrainerCallback = object + + +@PublicAPI(stability="beta") +class RayTrainReportCallback(TrainerCallback): + """A simple callback to report checkpoints and metrics to Ray Train. + + This callback is a subclass of `transformers.TrainerCallback + `_ + and overrides the `TrainerCallback.on_save()` method. After + a new checkpoint get saved, it fetches the latest metric dictionary + from `TrainerState.log_history` and reports it with the latest checkpoint + to Ray Train. + + Checkpoints will be saved in the following structure:: + + checkpoint_00000*/ Ray Train Checkpoint + └─ checkpoint/ Hugging Face Transformers Checkpoint + + For customized reporting and checkpointing logic, implement your own + `transformers.TrainerCallback` following this user + guide: :ref:`Saving and Loading Checkpoints `. + + Note that users should ensure that the logging, evaluation, and saving frequencies + are properly configured so that the monitoring metric is always up-to-date + when `transformers.Trainer` saves a checkpoint. + + Suppose the monitoring metric is reported from evaluation stage: + + Some valid configurations: + - evaluation_strategy == save_strategy == "epoch" + - evaluation_strategy == save_strategy == "steps", save_steps % eval_steps == 0 + + Some invalid configurations: + - evaluation_strategy != save_strategy + - evaluation_strategy == save_strategy == "steps", save_steps % eval_steps != 0 + + """ + + CHECKPOINT_NAME = "checkpoint" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + record_extra_usage_tag(TagKey.TRAIN_TRANSFORMERS_RAYTRAINREPORTCALLBACK, "1") + + def on_save(self, args, state, control, **kwargs): + """Event called after a checkpoint save.""" + with TemporaryDirectory() as tmpdir: + # Aggregate all the logged metrics + metrics = {} + for log in state.log_history: + metrics.update(log) + + # Copy ckpt files and construct a Ray Train Checkpoint + source_ckpt_path = transformers.trainer.get_last_checkpoint(args.output_dir) + if source_ckpt_path is not None: + target_ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix() + shutil.copytree(source_ckpt_path, target_ckpt_path) + checkpoint = Checkpoint.from_directory(tmpdir) + else: + checkpoint = None + + # Report latest metrics and checkpoint to Ray Train + ray.train.report(metrics=metrics, checkpoint=checkpoint) + + +class RayTorchIterableDataset(IterableDataset): + """Wrapper class for ray data iterables.""" + + def __init__(self, data_iterable) -> None: + super().__init__() + self.data_iterable = data_iterable + + def __iter__(self) -> Iterator: + return iter(self.data_iterable) + + +@PublicAPI(stability="beta") +def prepare_trainer(trainer: "Trainer") -> "Trainer": + """Prepare your HuggingFace Transformer Trainer for Ray Train. + + This utility function enable the trainer integrates with Ray Data Integration. + Internally, it overrides the `get_train_dataloader` and `get_eval_dataloader` + methods and inject the data integration logics if the `train_dataset` and + `eval_dataset` are Ray Data Iterables. + """ + + if TRANSFORMERS_IMPORT_ERROR is not None: + raise TRANSFORMERS_IMPORT_ERROR + + base_trainer_class: Type[transformers.trainer.Trainer] = trainer.__class__ + + class RayTransformersTrainer(base_trainer_class): + """A Wrapper of `transformers.Trainer` for Ray Data Integration.""" + + def get_train_dataloader(self) -> DataLoader: + if isinstance(self.train_dataset, _IterableFromIterator): + dataset = RayTorchIterableDataset(self.train_dataset) + return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0]) + else: + return super().get_train_dataloader() + + def get_eval_dataloader( + self, eval_dataset: Optional[Union[str, Dataset]] = None + ) -> DataLoader: + if eval_dataset is None: + eval_dataset = self.eval_dataset + + if ( + isinstance(eval_dataset, str) + and isinstance(self.eval_dataset, dict) + and isinstance(self.eval_dataset[eval_dataset], _IterableFromIterator) + ): + dataset = RayTorchIterableDataset(self.eval_dataset[eval_dataset]) + return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0]) + elif isinstance(eval_dataset, _IterableFromIterator): + dataset = RayTorchIterableDataset(eval_dataset) + return DataLoader(dataset, batch_size=1, collate_fn=lambda x: x[0]) + else: + return super().get_eval_dataloader(eval_dataset) + + trainer.__class__ = RayTransformersTrainer + + record_extra_usage_tag(TagKey.TRAIN_TRANSFORMERS_PREPARE_TRAINER, "1") + return trainer diff --git a/lib/python3.12/site-packages/ray/train/lightning/__init__.py b/lib/python3.12/site-packages/ray/train/lightning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8be5886a805cb80b5f30bf10ac04d37db696ce23 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/lightning/__init__.py @@ -0,0 +1,33 @@ +# isort: off +try: + import lightning # noqa: F401 +except ModuleNotFoundError: + try: + import pytorch_lightning # noqa: F401 + except ModuleNotFoundError: + raise ModuleNotFoundError( + "PyTorch Lightning isn't installed. To install PyTorch Lightning, " + "please run 'pip install lightning'" + ) +# isort: on + +from ray.train.lightning._lightning_utils import ( + RayDDPStrategy, + RayDeepSpeedStrategy, + RayFSDPStrategy, + RayLightningEnvironment, + RayTrainReportCallback, + prepare_trainer, +) + +__all__ = [ + "prepare_trainer", + "RayDDPStrategy", + "RayFSDPStrategy", + "RayDeepSpeedStrategy", + "RayLightningEnvironment", + "RayTrainReportCallback", +] + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/lightning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/lightning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a6e755aaff4238c934ade1e0a6581a616fe20a5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/lightning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/lightning/__pycache__/_lightning_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/lightning/__pycache__/_lightning_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b9211798e97db3e62e8bf45d44fc22452012c93c Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/lightning/__pycache__/_lightning_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/lightning/_lightning_utils.py b/lib/python3.12/site-packages/ray/train/lightning/_lightning_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2da924a3357c4a3ef2f76ce445dbe676640c7be7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/lightning/_lightning_utils.py @@ -0,0 +1,300 @@ +import logging +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any, Dict + +import torch +from packaging.version import Version + +import ray +import ray.train +from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag +from ray.train import Checkpoint +from ray.util import PublicAPI + + +def import_lightning(): # noqa: F402 + try: + import lightning.pytorch as pl + except ModuleNotFoundError: + import pytorch_lightning as pl + return pl + + +pl = import_lightning() + +_LIGHTNING_GREATER_EQUAL_2_0 = Version(pl.__version__) >= Version("2.0.0") +_LIGHTNING_LESS_THAN_2_1 = Version(pl.__version__) < Version("2.1.0") +_TORCH_GREATER_EQUAL_1_12 = Version(torch.__version__) >= Version("1.12.0") +_TORCH_FSDP_AVAILABLE = _TORCH_GREATER_EQUAL_1_12 and torch.distributed.is_available() + +try: + from lightning.pytorch.plugins.environments import LightningEnvironment +except ModuleNotFoundError: + from pytorch_lightning.plugins.environments import LightningEnvironment + +if _LIGHTNING_GREATER_EQUAL_2_0: + FSDPStrategy = pl.strategies.FSDPStrategy +else: + FSDPStrategy = pl.strategies.DDPFullyShardedStrategy + +if _TORCH_FSDP_AVAILABLE: + from torch.distributed.fsdp import ( + FullStateDictConfig, + FullyShardedDataParallel, + StateDictType, + ) + + +logger = logging.getLogger(__name__) + +LIGHTNING_REPORT_STAGE_KEY = "_report_on" + + +@PublicAPI(stability="beta") +class RayDDPStrategy(pl.strategies.DDPStrategy): + """Subclass of DDPStrategy to ensure compatibility with Ray orchestration. + + For a full list of initialization arguments, please refer to: + https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.DDPStrategy.html + + Note that `process_group_backend`, `timeout`, and `start_method` are disabled here, + please specify these arguments in :class:`~ray.train.torch.TorchConfig` instead. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYDDPSTRATEGY, "1") + + @property + def root_device(self) -> torch.device: + return ray.train.torch.get_device() + + @property + def distributed_sampler_kwargs(self) -> Dict[str, Any]: + return dict( + num_replicas=self.world_size, + rank=self.global_rank, + ) + + +@PublicAPI(stability="beta") +class RayFSDPStrategy(FSDPStrategy): # noqa: F821 + """Subclass of FSDPStrategy to ensure compatibility with Ray orchestration. + + For a full list of initialization arguments, please refer to: + https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html + + .. note:: + It is recommended to upgrade `lightning>=2.1` or above when using FSDP + with Lightning, since Lightning starts to natively support `state_dict_type`, + `sharding_strategy`, `auto_wrap_policy` and other FSDP configurations from 2.1. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYFSDPSTRATEGY, "1") + + @property + def root_device(self) -> torch.device: + return ray.train.torch.get_device() + + @property + def distributed_sampler_kwargs(self) -> Dict[str, Any]: + return dict( + num_replicas=self.world_size, + rank=self.global_rank, + ) + + def lightning_module_state_dict(self) -> Dict[str, Any]: + """Gathers the full state dict to rank 0 on CPU. + + FSDP checkpointing is broken in Lightning 2.0.x. This subclass patches the + behavior to perform a full state dict checkpointing, gathering the checkpoint + shards on rank 0 CPU. Upgrade to `lightning>=2.1` to do sharded state dict + checkpointing. + + See the note in the class docstring for more details. + """ + + assert self.model is not None, "Failed to get the state dict for a None model!" + + if ( + _TORCH_FSDP_AVAILABLE + and _LIGHTNING_GREATER_EQUAL_2_0 + and _LIGHTNING_LESS_THAN_2_1 + ): + with FullyShardedDataParallel.state_dict_type( + module=self.model, + state_dict_type=StateDictType.FULL_STATE_DICT, + state_dict_config=FullStateDictConfig( + offload_to_cpu=True, rank0_only=True + ), + ): + state_dict = self.model.state_dict() + + ckpt_state_dict = {} + prefix_len = len("_forward_module.") + for k, v in state_dict.items(): + if k.startswith("_forward_module."): + non_prefixed_key = k[prefix_len:] + ckpt_state_dict[non_prefixed_key] = v + else: + ckpt_state_dict[k] = v + return ckpt_state_dict + else: + # Otherwise Lightning uses Fairscale FSDP, no need to unshard by ourself. + return super().lightning_module_state_dict() + + +@PublicAPI(stability="beta") +class RayDeepSpeedStrategy(pl.strategies.DeepSpeedStrategy): + """Subclass of DeepSpeedStrategy to ensure compatibility with Ray orchestration. + + For a full list of initialization arguments, please refer to: + https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.DeepSpeedStrategy.html + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYDEEPSPEEDSTRATEGY, "1") + + @property + def root_device(self) -> torch.device: + return ray.train.torch.get_device() + + @property + def distributed_sampler_kwargs(self) -> Dict[str, Any]: + return dict( + num_replicas=self.world_size, + rank=self.global_rank, + ) + + +@PublicAPI(stability="beta") +class RayLightningEnvironment(LightningEnvironment): # noqa: F821 + """Setup Lightning DDP training environment for Ray cluster.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYLIGHTNINGENVIRONMENT, "1") + + def world_size(self) -> int: + return ray.train.get_context().get_world_size() + + def global_rank(self) -> int: + return ray.train.get_context().get_world_rank() + + def local_rank(self) -> int: + return ray.train.get_context().get_local_rank() + + def node_rank(self) -> int: + return ray.train.get_context().get_node_rank() + + def set_world_size(self, size: int) -> None: + # Disable it since `world_size()` directly returns data from Train context. + pass + + def set_global_rank(self, rank: int) -> None: + # Disable it since `global_rank()` directly returns data from Train. + pass + + def teardown(self): + pass + + +@PublicAPI(stability="beta") +def prepare_trainer(trainer: pl.Trainer) -> pl.Trainer: + """Prepare the PyTorch Lightning Trainer for distributed execution.""" + + # Check strategy class + valid_strategy_class = [RayDDPStrategy, RayFSDPStrategy, RayDeepSpeedStrategy] + + if not any(isinstance(trainer.strategy, cls) for cls in valid_strategy_class): + raise RuntimeError( + f"Invalid strategy class: {type(trainer.strategy)}. To use " + "PyTorch Lightning with Ray, the strategy object should be one of " + f"{[cls.__name__ for cls in valid_strategy_class]} class " + "or its subclass." + ) + + # Check cluster environment + cluster_environment = getattr(trainer.strategy, "cluster_environment", None) + if cluster_environment and not isinstance( + cluster_environment, RayLightningEnvironment + ): + raise RuntimeError( + "Invalid cluster environment plugin. The expected class is" + "`ray.train.lightning.RayLightningEnvironment` " + f"but got {type(cluster_environment)}!" + ) + + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_PREPARE_TRAINER, "1") + return trainer + + +@PublicAPI(stability="beta") +class RayTrainReportCallback(pl.callbacks.Callback): + """A simple callback that reports checkpoints to Ray on train epoch end. + + This callback is a subclass of `lightning.pytorch.callbacks.Callback + `_. + + It fetches the latest `trainer.callback_metrics` and reports together with + the checkpoint on each training epoch end. + + Checkpoints will be saved in the following structure:: + + checkpoint_00000*/ Ray Train Checkpoint + └─ checkpoint.ckpt PyTorch Lightning Checkpoint + + For customized reporting and checkpointing logic, implement your own + `lightning.pytorch.callbacks.Callback` following this user + guide: :ref:`Saving and Loading Checkpoints `. + """ + + CHECKPOINT_NAME = "checkpoint.ckpt" + + def __init__(self) -> None: + super().__init__() + job_id = ray.get_runtime_context().get_job_id() + experiment_name = ray.train.get_context().get_experiment_name() + self.local_rank = ray.train.get_context().get_local_rank() + + self.tmpdir_prefix = Path( + tempfile.gettempdir(), + f"lightning_checkpoints-job_id={job_id}-name={experiment_name}", + ).as_posix() + if os.path.isdir(self.tmpdir_prefix) and self.local_rank == 0: + shutil.rmtree(self.tmpdir_prefix) + + record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1") + + def on_train_epoch_end(self, trainer, pl_module) -> None: + # Creates a checkpoint dir with fixed name + tmpdir = Path(self.tmpdir_prefix, str(trainer.current_epoch)).as_posix() + os.makedirs(tmpdir, exist_ok=True) + + # Fetch metrics + metrics = trainer.callback_metrics + metrics = {k: v.item() for k, v in metrics.items()} + + # (Optional) Add customized metrics + metrics["epoch"] = trainer.current_epoch + metrics["step"] = trainer.global_step + + # Save checkpoint to local + ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix() + trainer.save_checkpoint(ckpt_path, weights_only=False) + + # Report to train session + checkpoint = Checkpoint.from_directory(tmpdir) + ray.train.report(metrics=metrics, checkpoint=checkpoint) + + # Add a barrier to ensure all workers finished reporting here + trainer.strategy.barrier() + + if self.local_rank == 0: + shutil.rmtree(tmpdir) diff --git a/lib/python3.12/site-packages/ray/train/predictor.py b/lib/python3.12/site-packages/ray/train/predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..0d4ef100e360b14cdabc71a7fc9a3dc4e25d771b --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/predictor.py @@ -0,0 +1,254 @@ +import abc +from typing import Callable, Dict, Optional, Type, Union + +import numpy as np +import pandas as pd + +from ray.air.data_batch_type import DataBatchType +from ray.air.util.data_batch_conversion import ( + BatchFormat, + _convert_batch_type_to_numpy, + _convert_batch_type_to_pandas, +) +from ray.data import Preprocessor +from ray.train import Checkpoint +from ray.util.annotations import DeveloperAPI, PublicAPI + +try: + import pyarrow + + pa_table = pyarrow.Table +except ImportError: + pa_table = None + +# Reverse mapping from data batch type to batch format. +TYPE_TO_ENUM: Dict[Type[DataBatchType], BatchFormat] = { + np.ndarray: BatchFormat.NUMPY, + dict: BatchFormat.NUMPY, + pd.DataFrame: BatchFormat.PANDAS, +} + + +@PublicAPI(stability="beta") +class PredictorNotSerializableException(RuntimeError): + """Error raised when trying to serialize a Predictor instance.""" + + pass + + +@PublicAPI(stability="beta") +class Predictor(abc.ABC): + """Predictors load models from checkpoints to perform inference. + + .. note:: + The base ``Predictor`` class cannot be instantiated directly. Only one of + its subclasses can be used. + + **How does a Predictor work?** + + Predictors expose a ``predict`` method that accepts an input batch of type + ``DataBatchType`` and outputs predictions of the same type as the input batch. + + When the ``predict`` method is called the following occurs: + + - The input batch is converted into a pandas DataFrame. Tensor input (like a + ``np.ndarray``) will be converted into a single column Pandas Dataframe. + - If there is a :ref:`Preprocessor ` saved in the provided + :class:`Checkpoint `, the preprocessor will be used to + transform the DataFrame. + - The transformed DataFrame will be passed to the model for inference (via the + ``predictor._predict_pandas`` method). + - The predictions will be outputted by ``predict`` in the same type as the + original input. + + **How do I create a new Predictor?** + + To implement a new Predictor for your particular framework, you should subclass + the base ``Predictor`` and implement the following two methods: + + 1. ``_predict_pandas``: Given a pandas.DataFrame input, return a + pandas.DataFrame containing predictions. + 2. ``from_checkpoint``: Logic for creating a Predictor from a + :class:`Checkpoint `. + 3. Optionally ``_predict_numpy`` for better performance when working with + tensor data to avoid extra copies from Pandas conversions. + """ + + def __init__(self, preprocessor: Optional[Preprocessor] = None): + """Subclasseses must call Predictor.__init__() to set a preprocessor.""" + self._preprocessor: Optional[Preprocessor] = preprocessor + # Whether tensor columns should be automatically cast from/to the tensor + # extension type at UDF boundaries. This can be overridden by subclasses. + self._cast_tensor_columns = False + + @classmethod + @abc.abstractmethod + def from_checkpoint(cls, checkpoint: Checkpoint, **kwargs) -> "Predictor": + """Create a specific predictor from a checkpoint. + + Args: + checkpoint: Checkpoint to load predictor data from. + kwargs: Arguments specific to predictor implementations. + + Returns: + Predictor: Predictor object. + """ + raise NotImplementedError + + @classmethod + def from_pandas_udf( + cls, pandas_udf: Callable[[pd.DataFrame], pd.DataFrame] + ) -> "Predictor": + """Create a Predictor from a Pandas UDF. + + Args: + pandas_udf: A function that takes a pandas.DataFrame and other + optional kwargs and returns a pandas.DataFrame. + """ + + class PandasUDFPredictor(Predictor): + @classmethod + def from_checkpoint(cls, checkpoint: Checkpoint, **kwargs) -> "Predictor": + return PandasUDFPredictor() + + def _predict_pandas(self, df, **kwargs) -> "pd.DataFrame": + return pandas_udf(df, **kwargs) + + return PandasUDFPredictor() + + def get_preprocessor(self) -> Optional[Preprocessor]: + """Get the preprocessor to use prior to executing predictions.""" + return self._preprocessor + + def set_preprocessor(self, preprocessor: Optional[Preprocessor]) -> None: + """Set the preprocessor to use prior to executing predictions.""" + self._preprocessor = preprocessor + + @classmethod + @DeveloperAPI + def preferred_batch_format(cls) -> BatchFormat: + """Batch format hint for upstream producers to try yielding best block format. + + The preferred batch format to use if both `_predict_pandas` and + `_predict_numpy` are implemented. Defaults to Pandas. + + Can be overriden by predictor classes depending on the framework type, + e.g. TorchPredictor prefers Numpy and XGBoostPredictor prefers Pandas as + native batch format. + + """ + return BatchFormat.PANDAS + + @classmethod + def _batch_format_to_use(cls) -> BatchFormat: + """Determine the batch format to use for the predictor.""" + has_pandas_implemented = cls._predict_pandas != Predictor._predict_pandas + has_numpy_implemented = cls._predict_numpy != Predictor._predict_numpy + if has_pandas_implemented and has_numpy_implemented: + return cls.preferred_batch_format() + elif has_pandas_implemented: + return BatchFormat.PANDAS + elif has_numpy_implemented: + return BatchFormat.NUMPY + else: + raise NotImplementedError( + f"Predictor {cls.__name__} must implement at least one of " + "`_predict_pandas` and `_predict_numpy`." + ) + + def _set_cast_tensor_columns(self): + """Enable automatic tensor column casting. + + If this is called on a predictor, the predictor will cast tensor columns to + NumPy ndarrays in the input to the preprocessors and cast tensor columns back to + the tensor extension type in the prediction outputs. + """ + self._cast_tensor_columns = True + + def predict(self, data: DataBatchType, **kwargs) -> DataBatchType: + """Perform inference on a batch of data. + + Args: + data: A batch of input data of type ``DataBatchType``. + kwargs: Arguments specific to predictor implementations. These are passed + directly to ``_predict_numpy`` or ``_predict_pandas``. + + Returns: + DataBatchType: + Prediction result. The return type will be the same as the input type. + """ + if not hasattr(self, "_preprocessor"): + raise NotImplementedError( + "Subclasses of Predictor must call Predictor.__init__(preprocessor)." + ) + try: + batch_format = TYPE_TO_ENUM[type(data)] + except KeyError: + raise RuntimeError( + f"Invalid input data type of {type(data)}, supported " + f"types: {list(TYPE_TO_ENUM.keys())}" + ) + + if self._preprocessor: + data = self._preprocessor.transform_batch(data) + + batch_format_to_use = self._batch_format_to_use() + + # We can finish prediction as long as one predict method is implemented. + # For prediction, we have to return back in the same format as the input. + if batch_format == BatchFormat.PANDAS: + if batch_format_to_use == BatchFormat.PANDAS: + return self._predict_pandas( + _convert_batch_type_to_pandas(data), **kwargs + ) + elif batch_format_to_use == BatchFormat.NUMPY: + return _convert_batch_type_to_pandas( + self._predict_numpy(_convert_batch_type_to_numpy(data), **kwargs) + ) + elif batch_format == BatchFormat.NUMPY: + if batch_format_to_use == BatchFormat.PANDAS: + return _convert_batch_type_to_numpy( + self._predict_pandas(_convert_batch_type_to_pandas(data), **kwargs) + ) + elif batch_format_to_use == BatchFormat.NUMPY: + return self._predict_numpy(_convert_batch_type_to_numpy(data), **kwargs) + + @DeveloperAPI + def _predict_pandas(self, data: "pd.DataFrame", **kwargs) -> "pd.DataFrame": + """Perform inference on a Pandas DataFrame. + + Args: + data: A pandas DataFrame to perform predictions on. + kwargs: Arguments specific to the predictor implementation. + + Returns: + A pandas DataFrame containing the prediction result. + + """ + raise NotImplementedError + + @DeveloperAPI + def _predict_numpy( + self, data: Union[np.ndarray, Dict[str, np.ndarray]], **kwargs + ) -> Union[np.ndarray, Dict[str, np.ndarray]]: + """Perform inference on a Numpy data. + + All Predictors working with tensor data (like deep learning predictors) + should implement this method. + + Args: + data: A Numpy ndarray or dictionary of ndarrays to perform predictions on. + kwargs: Arguments specific to the predictor implementation. + + Returns: + A Numpy ndarray or dictionary of ndarray containing the prediction result. + + """ + raise NotImplementedError + + def __reduce__(self): + raise PredictorNotSerializableException( + "Predictor instances are not serializable. Instead, you may want " + "to serialize a checkpoint and initialize the Predictor with " + "Predictor.from_checkpoint." + ) diff --git a/lib/python3.12/site-packages/ray/train/session.py b/lib/python3.12/site-packages/ray/train/session.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/train/torch/__init__.py b/lib/python3.12/site-packages/ray/train/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1774b98cb18a88c7f86d112bd599485b47626b37 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/__init__.py @@ -0,0 +1,58 @@ +# isort: off +try: + import torch # noqa: F401 +except ModuleNotFoundError: + raise ModuleNotFoundError( + "PyTorch isn't installed. To install PyTorch, run 'pip install torch'" + ) +# isort: on + +from ray.train.torch.config import TorchConfig +from ray.train.torch.torch_checkpoint import TorchCheckpoint +from ray.train.torch.torch_detection_predictor import TorchDetectionPredictor +from ray.train.torch.torch_predictor import TorchPredictor +from ray.train.torch.torch_trainer import TorchTrainer +from ray.train.torch.train_loop_utils import ( + accelerate, + backward, + enable_reproducibility, + get_device, + get_devices, + prepare_data_loader, + prepare_model, + prepare_optimizer, +) +from ray.train.v2._internal.constants import is_v2_enabled + +if is_v2_enabled(): + from ray.train.v2.torch.torch_trainer import TorchTrainer # noqa: F811 + from ray.train.v2.torch.train_loop_utils import ( # noqa: F811 + accelerate, + backward, + enable_reproducibility, + get_device, + get_devices, + prepare_data_loader, + prepare_model, + prepare_optimizer, + ) + + +__all__ = [ + "TorchTrainer", + "TorchCheckpoint", + "TorchConfig", + "accelerate", + "get_device", + "get_devices", + "prepare_model", + "prepare_optimizer", + "prepare_data_loader", + "backward", + "enable_reproducibility", + "TorchPredictor", + "TorchDetectionPredictor", +] + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4e72152cc91ac4fe7a8f2ce16feddf2acd321fa Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22631acfcb4ae0b19972d6723d61e7baad01dda1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bfd33e8a9b2349977332db2c3efbfce28f65e94 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_detection_predictor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_detection_predictor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5df0433ad8d9080abaf5d5bf87bac2b641b69139 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_detection_predictor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_predictor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_predictor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5f379ddfd3378c285e3d008f2e890bdf4fddbe2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_predictor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_trainer.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80bbb053f9739bea62f6bd6cb32ad52b07c0bd72 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/torch_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/__pycache__/train_loop_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/__pycache__/train_loop_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e188869336e667414ad9066d6f7e2f681995601a Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/__pycache__/train_loop_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/config.py b/lib/python3.12/site-packages/ray/train/torch/config.py new file mode 100644 index 0000000000000000000000000000000000000000..9042e1d2d1157cf956a227c2a58dbec1cebc36ab --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/config.py @@ -0,0 +1,233 @@ +import logging +import os +from dataclasses import dataclass +from datetime import timedelta +from typing import Optional + +import torch +import torch.distributed as dist +from packaging.version import Version + +import ray +from ray._common.network_utils import build_address +from ray._private import ray_constants +from ray.air._internal.device_manager import register_custom_torch_dist_backend +from ray.exceptions import GetTimeoutError +from ray.train._internal.base_worker_group import BaseWorkerGroup +from ray.train._internal.utils import get_address_and_port +from ray.train.backend import Backend, BackendConfig +from ray.train.constants import ( + DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, + TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, +) +from ray.util import PublicAPI + +logger = logging.getLogger(__name__) + + +class TorchConfigContextManager: + def __enter__(self): + # Set default cuda device + if torch.cuda.is_available(): + device = ray.train.torch.get_device() + if device.type == "cuda": + torch.cuda.set_device(device) + + def __exit__(self, type, value, traceback): + # Propagate exceptions if any + return False + + +@PublicAPI(stability="stable") +@dataclass +class TorchConfig(BackendConfig): + """Configuration for torch process group setup. + + See https://pytorch.org/docs/stable/distributed.html for more info. + + Args: + backend: The backend to use for training. + See ``torch.distributed.init_process_group`` for more info and + valid values. + If set to None, nccl will be used if GPUs are requested, else gloo + will be used. + init_method: The initialization method to use. Either "env" + for environment variable initialization or "tcp" for TCP + initialization. Defaults to "env". + timeout_s: Seconds for process group operations to timeout. + """ + + backend: Optional[str] = None + init_method: str = "env" + timeout_s: int = 1800 + + @property + def backend_cls(self): + return _TorchBackend + + @property + def train_func_context(self): + return TorchConfigContextManager + + +def _setup_torch_process_group( + backend: str, + world_rank: int, + world_size: int, + init_method: str, + timeout_s: int = 1800, +): + """Connects the distributed PyTorch backend. + + Args: + backend: The backend (nccl, gloo, etc.) to use for training. + world_rank: Rank of the current worker. + world_size: Number of workers participating in the job. + init_method: URL specifying how to initialize the process group. + timeout_s: Seconds for process group operations to timeout. + """ + if world_rank == 0: + logger.info( + f"Setting up process group for: {init_method} [rank={world_rank}, " + f"world_size={world_size}]" + ) + else: + logger.debug( + f"Setting up process group for: {init_method} [rank={world_rank}, " + f"world_size={world_size}]" + ) + logger.debug(f"using {backend}") + + if backend == "nccl": + # See https://github.com/pytorch/pytorch/blob/c263bd43e8e8502d4726643bc6fd046f0130ac0e/torch/distributed/distributed_c10d.py#L803-L823 # noqa: E501 + # We do not use TORCH_NCCL_BLOCKING_WAIT due to performance overhead. + if Version(torch.__version__) < Version("2.2.0"): + TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR = "NCCL_ASYNC_ERROR_HANDLING" + TORCH_NCCL_BLOCKING_WAIT_ENV_VAR = "NCCL_BLOCKING_WAIT" + else: + TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR = "TORCH_NCCL_ASYNC_ERROR_HANDLING" + TORCH_NCCL_BLOCKING_WAIT_ENV_VAR = "TORCH_NCCL_BLOCKING_WAIT" + if ( + TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR not in os.environ + and TORCH_NCCL_BLOCKING_WAIT_ENV_VAR not in os.environ + ): + logger.debug( + f"Setting {TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR}=1 to fail if NCCL collective communication operations are timing out. " # noqa: E501 + f"To override this behavior, you can set {TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR}=0." # noqa: E501 + ) + os.environ[TORCH_NCCL_ASYNC_ERROR_HANDLING_ENV_VAR] = "1" + elif backend == "hccl": + register_custom_torch_dist_backend(backend) + + dist.init_process_group( + backend=backend, + init_method=init_method, + rank=world_rank, + world_size=world_size, + timeout=timedelta(seconds=timeout_s), + ) + + +def _shutdown_torch(destroy_process_group=False): + from ray.air._internal.torch_utils import get_devices + + devices = get_devices() + if destroy_process_group: + dist.destroy_process_group() + if torch.cuda.is_available(): + for device in devices: + with torch.cuda.device(device): + torch.cuda.empty_cache() + + +def _set_torch_distributed_env_vars(): + # Same env vars as in + # https://pytorch.org/docs/stable/elastic/run.html#environment-variables + from ray.train.torch import get_device + + context = ray.train.get_context() + os.environ["LOCAL_RANK"] = str(context.get_local_rank()) + os.environ["RANK"] = str(context.get_world_rank()) + os.environ["LOCAL_WORLD_SIZE"] = str(context.get_local_world_size()) + os.environ["WORLD_SIZE"] = str(context.get_world_size()) + os.environ["NODE_RANK"] = str(context.get_node_rank()) + + # Makes sure Hugging Face Accelerate uses the correct device + device = get_device() + os.environ["ACCELERATE_TORCH_DEVICE"] = str(device) + + +class _TorchBackend(Backend): + share_cuda_visible_devices: bool = True + + def on_start(self, worker_group: BaseWorkerGroup, backend_config: TorchConfig): + if dist.is_available(): + # Set the appropriate training backend. + if backend_config.backend is None: + resources = worker_group.get_resources_per_worker() + num_gpus_per_worker = resources.get("GPU", 0) + + if num_gpus_per_worker > 0: + backend = "nccl" + else: + backend = "gloo" + else: + backend = backend_config.backend + + master_addr, master_port = worker_group.execute_single( + 0, get_address_and_port + ) + if backend_config.init_method == "env": + + def set_env_vars(addr, port): + os.environ["MASTER_ADDR"] = addr + os.environ["MASTER_PORT"] = str(port) + + worker_group.execute(set_env_vars, addr=master_addr, port=master_port) + url = "env://" + elif backend_config.init_method == "tcp": + url = f"tcp://{build_address(master_addr, master_port)}" + else: + raise ValueError( + f"The provided init_method (" + f"{backend_config.init_method}) is not supported. Must " + f"be either 'env' or 'tcp'." + ) + + setup_futures = [] + for i in range(len(worker_group)): + setup_futures.append( + worker_group.execute_single_async( + i, + _setup_torch_process_group, + backend=backend, + world_rank=i, + world_size=len(worker_group), + init_method=url, + timeout_s=backend_config.timeout_s, + ) + ) + ray.get(setup_futures) + else: + raise RuntimeError("Distributed torch is not available.") + + def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config): + futures = worker_group.execute_async( + _shutdown_torch, + destroy_process_group=len(worker_group) > 1, + ) + timeout_s = ray_constants.env_integer( + TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, + DEFAULT_TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, + ) + try: + ray.get(futures, timeout=timeout_s) + except GetTimeoutError: + logger.warning( + f"Torch process group shutdown timed out after {timeout_s} seconds" + ) + + def on_training_start( + self, worker_group: BaseWorkerGroup, backend_config: BackendConfig + ): + worker_group.execute(_set_torch_distributed_env_vars) diff --git a/lib/python3.12/site-packages/ray/train/torch/torch_checkpoint.py b/lib/python3.12/site-packages/ray/train/torch/torch_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..7b6eeae305185bdd638042a7c1a0386f8d26bd86 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/torch_checkpoint.py @@ -0,0 +1,182 @@ +import os +import tempfile +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Optional + +import torch + +from ray.air._internal.torch_utils import ( + consume_prefix_in_state_dict_if_present_not_in_place, + load_torch_model, +) +from ray.train._internal.framework_checkpoint import FrameworkCheckpoint +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.data.preprocessor import Preprocessor + +ENCODED_DATA_KEY = "torch_encoded_data" + + +@PublicAPI(stability="beta") +class TorchCheckpoint(FrameworkCheckpoint): + """A :class:`~ray.train.Checkpoint` with Torch-specific functionality.""" + + MODEL_FILENAME = "model.pt" + + @classmethod + def from_state_dict( + cls, + state_dict: Dict[str, Any], + *, + preprocessor: Optional["Preprocessor"] = None, + ) -> "TorchCheckpoint": + """Create a :class:`~ray.train.Checkpoint` that stores a model state dictionary. + + .. tip:: + + This is the recommended method for creating + :class:`TorchCheckpoints`. + + Args: + state_dict: The model state dictionary to store in the checkpoint. + preprocessor: A fitted preprocessor to be applied before inference. + + Returns: + A :class:`TorchCheckpoint` containing the specified state dictionary. + + Examples: + + .. testcode:: + + import torch + import torch.nn as nn + from ray.train.torch import TorchCheckpoint + + # Set manual seed + torch.manual_seed(42) + + # Function to create a NN model + def create_model() -> nn.Module: + model = nn.Sequential(nn.Linear(1, 10), + nn.ReLU(), + nn.Linear(10,1)) + return model + + # Create a TorchCheckpoint from our model's state_dict + model = create_model() + checkpoint = TorchCheckpoint.from_state_dict(model.state_dict()) + + # Now load the model from the TorchCheckpoint by providing the + # model architecture + model_from_chkpt = checkpoint.get_model(create_model()) + + # Assert they have the same state dict + assert str(model.state_dict()) == str(model_from_chkpt.state_dict()) + print("worked") + + .. testoutput:: + :hide: + + ... + """ + tempdir = tempfile.mkdtemp() + + model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix() + stripped_state_dict = consume_prefix_in_state_dict_if_present_not_in_place( + state_dict, "module." + ) + torch.save(stripped_state_dict, model_path) + + checkpoint = cls.from_directory(tempdir) + if preprocessor: + checkpoint.set_preprocessor(preprocessor) + return checkpoint + + @classmethod + def from_model( + cls, + model: torch.nn.Module, + *, + preprocessor: Optional["Preprocessor"] = None, + ) -> "TorchCheckpoint": + """Create a :class:`~ray.train.Checkpoint` that stores a Torch model. + + .. note:: + + PyTorch recommends storing state dictionaries. To create a + :class:`TorchCheckpoint` from a state dictionary, call + :meth:`~ray.train.torch.TorchCheckpoint.from_state_dict`. To learn more + about state dictionaries, read + `Saving and Loading Models `_. # noqa: E501 + + Args: + model: The Torch model to store in the checkpoint. + preprocessor: A fitted preprocessor to be applied before inference. + + Returns: + A :class:`TorchCheckpoint` containing the specified model. + + Examples: + + .. testcode:: + + from ray.train.torch import TorchCheckpoint + import torch + + # Create model identity and send a random tensor to it + model = torch.nn.Identity() + input = torch.randn(2, 2) + output = model(input) + + # Create a checkpoint + checkpoint = TorchCheckpoint.from_model(model) + print(checkpoint) + + .. testoutput:: + :hide: + + ... + """ + tempdir = tempfile.mkdtemp() + + model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix() + torch.save(model, model_path) + + checkpoint = cls.from_directory(tempdir) + if preprocessor: + checkpoint.set_preprocessor(preprocessor) + return checkpoint + + def get_model(self, model: Optional[torch.nn.Module] = None) -> torch.nn.Module: + """Retrieve the model stored in this checkpoint. + + Args: + model: If the checkpoint contains a model state dict, and not + the model itself, then the state dict will be loaded to this + ``model``. Otherwise, the model will be discarded. + """ + with self.as_directory() as tempdir: + model_path = Path(tempdir, self.MODEL_FILENAME).as_posix() + if not os.path.exists(model_path): + raise RuntimeError( + "`model.pt` not found within this checkpoint. Make sure you " + "created this `TorchCheckpoint` from one of its public " + "constructors (`from_state_dict` or `from_model`)." + ) + model_or_state_dict = torch.load(model_path, map_location="cpu") + + if isinstance(model_or_state_dict, torch.nn.Module): + if model: + warnings.warn( + "TorchCheckpoint already contains all information needed. " + "Discarding provided `model` argument. If you are using " + "TorchPredictor directly, you should do " + "`TorchPredictor.from_checkpoint(checkpoint)` by removing kwargs " + "`model=`." + ) + model = load_torch_model( + saved_model=model_or_state_dict, model_definition=model + ) + return model diff --git a/lib/python3.12/site-packages/ray/train/torch/torch_detection_predictor.py b/lib/python3.12/site-packages/ray/train/torch/torch_detection_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..98a708a731ca2976d57c7866dfc0707d5c255545 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/torch_detection_predictor.py @@ -0,0 +1,90 @@ +import collections +from typing import Dict, List, Optional, Union + +import numpy as np +import torch + +from ray.train._internal.dl_predictor import TensorDtype +from ray.train.torch.torch_predictor import TorchPredictor +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class TorchDetectionPredictor(TorchPredictor): + """A predictor for TorchVision detection models. + + Unlike other Torch models, instance segmentation models return + `List[Dict[str, Tensor]]`. This predictor extends :class:`TorchPredictor` to support + the non-standard outputs. + + To learn more about instance segmentation models, read + `Instance segmentation models `_. + + Example: + + .. testcode:: + + import numpy as np + from torchvision import models + + from ray.train.torch import TorchDetectionPredictor + + model = models.detection.fasterrcnn_resnet50_fpn_v2(pretrained=True) + + predictor = TorchDetectionPredictor(model=model) + predictions = predictor.predict(np.zeros((4, 3, 32, 32), dtype=np.float32)) + + print(predictions.keys()) + + .. testoutput:: + + dict_keys(['pred_boxes', 'pred_labels', 'pred_scores']) + + """ # noqa: E501 + + def _predict_numpy( + self, + data: Union[np.ndarray, Dict[str, np.ndarray]], + dtype: Optional[Union[TensorDtype, Dict[str, TensorDtype]]], + ) -> Dict[str, np.ndarray]: + if isinstance(data, dict) and len(data) != 1: + raise ValueError( + f"""Expected input to contain one key, but got {len(data)} instead.""" + ) + + if dtype is not None and not isinstance(dtype, torch.dtype): + raise ValueError( + "Expected `dtype` to be a `torch.dtype`, but got a " + f"{type(dtype).__name__} instead." + ) + + if isinstance(data, dict): + images = next(iter(data.values())) + else: + images = data + + inputs = [ + torch.as_tensor(image, dtype=dtype).to(self.device) for image in images + ] + outputs = self.call_model(inputs) + outputs = _convert_outputs_to_batch(outputs) + outputs = {"pred_" + key: value for key, value in outputs.items()} + + return outputs + + +def _convert_outputs_to_batch( + outputs: List[Dict[str, torch.Tensor]], +) -> Dict[str, List[torch.Tensor]]: + """Batch detection model outputs. + + TorchVision detection models return `List[Dict[Tensor]]`. Each `Dict` contain + 'boxes', 'labels, and 'scores'. + + This function batches values and returns a `Dict[str, List[Tensor]]`. + """ # noqa: E501 + batch = collections.defaultdict(list) + for output in outputs: + for key, value in output.items(): + batch[key].append(value.cpu().detach()) + return batch diff --git a/lib/python3.12/site-packages/ray/train/torch/torch_predictor.py b/lib/python3.12/site-packages/ray/train/torch/torch_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..be81d5b0f68c0a397239940bc9b310fc79a0b864 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/torch_predictor.py @@ -0,0 +1,250 @@ +import logging +from typing import TYPE_CHECKING, Dict, Optional, Union + +import numpy as np +import torch + +from ray.air._internal.torch_utils import convert_ndarray_batch_to_torch_tensor_batch +from ray.train._internal.dl_predictor import DLPredictor +from ray.train.predictor import DataBatchType +from ray.train.torch import TorchCheckpoint +from ray.util import log_once +from ray.util.annotations import DeveloperAPI, PublicAPI + +if TYPE_CHECKING: + from ray.data.preprocessor import Preprocessor + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="beta") +class TorchPredictor(DLPredictor): + """A predictor for PyTorch models. + + Args: + model: The torch module to use for predictions. + preprocessor: A preprocessor used to transform data batches prior + to prediction. + use_gpu: If set, the model will be moved to GPU on instantiation and + prediction happens on GPU. + """ + + def __init__( + self, + model: torch.nn.Module, + preprocessor: Optional["Preprocessor"] = None, + use_gpu: bool = False, + ): + self.model = model + self.model.eval() + self.use_gpu = use_gpu + + if use_gpu: + # TODO (jiaodong): #26249 Use multiple GPU devices with sharded input + self.device = torch.device("cuda") + else: + self.device = torch.device("cpu") + + # Ensure input tensor and model live on the same device + self.model.to(self.device) + + if ( + not use_gpu + and torch.cuda.device_count() > 0 + and log_once("torch_predictor_not_using_gpu") + ): + logger.warning( + "You have `use_gpu` as False but there are " + f"{torch.cuda.device_count()} GPUs detected on host where " + "prediction will only use CPU. Please consider explicitly " + "setting `TorchPredictor(use_gpu=True)` or " + "`batch_predictor.predict(ds, num_gpus_per_worker=1)` to " + "enable GPU prediction." + ) + + super().__init__(preprocessor) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(model={self.model!r}, " + f"preprocessor={self._preprocessor!r}, use_gpu={self.use_gpu!r})" + ) + + @classmethod + def from_checkpoint( + cls, + checkpoint: TorchCheckpoint, + model: Optional[torch.nn.Module] = None, + use_gpu: bool = False, + ) -> "TorchPredictor": + """Instantiate the predictor from a TorchCheckpoint. + + Args: + checkpoint: The checkpoint to load the model and preprocessor from. + model: If the checkpoint contains a model state dict, and not + the model itself, then the state dict will be loaded to this + ``model``. If the checkpoint already contains the model itself, + this model argument will be discarded. + use_gpu: If set, the model will be moved to GPU on instantiation and + prediction happens on GPU. + """ + model = checkpoint.get_model(model) + preprocessor = checkpoint.get_preprocessor() + return cls(model=model, preprocessor=preprocessor, use_gpu=use_gpu) + + @DeveloperAPI + def call_model( + self, inputs: Union[torch.Tensor, Dict[str, torch.Tensor]] + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + """Runs inference on a single batch of tensor data. + + This method is called by `TorchPredictor.predict` after converting the + original data batch to torch tensors. + + Override this method to add custom logic for processing the model input or + output. + + Args: + inputs: A batch of data to predict on, represented as either a single + PyTorch tensor or for multi-input models, a dictionary of tensors. + + Returns: + The model outputs, either as a single tensor or a dictionary of tensors. + + Example: + + .. testcode:: + + import numpy as np + import torch + from ray.train.torch import TorchPredictor + + # List outputs are not supported by default TorchPredictor. + # So let's define a custom TorchPredictor and override call_model + class MyModel(torch.nn.Module): + def forward(self, input_tensor): + return [input_tensor, input_tensor] + + # Use a custom predictor to format model output as a dict. + class CustomPredictor(TorchPredictor): + def call_model(self, inputs): + model_output = super().call_model(inputs) + return { + str(i): model_output[i] for i in range(len(model_output)) + } + + # create our data batch + data_batch = np.array([1, 2]) + # create custom predictor and predict + predictor = CustomPredictor(model=MyModel()) + predictions = predictor.predict(data_batch) + print(f"Predictions: {predictions.get('0')}, {predictions.get('1')}") + + .. testoutput:: + + Predictions: [1 2], [1 2] + + """ + with torch.no_grad(): + output = self.model(inputs) + return output + + def predict( + self, + data: DataBatchType, + dtype: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, + ) -> DataBatchType: + """Run inference on data batch. + + If the provided data is a single array or a dataframe/table with a single + column, it will be converted into a single PyTorch tensor before being + inputted to the model. + + If the provided data is a multi-column table or a dict of numpy arrays, + it will be converted into a dict of tensors before being inputted to the + model. This is useful for multi-modal inputs (for example your model accepts + both image and text). + + Args: + data: A batch of input data of ``DataBatchType``. + dtype: The dtypes to use for the tensors. Either a single dtype for all + tensors or a mapping from column name to dtype. + + Returns: + DataBatchType: Prediction result. The return type will be the same as the + input type. + + Example: + + .. testcode:: + + import numpy as np + import pandas as pd + import torch + import ray + from ray.train.torch import TorchPredictor + + # Define a custom PyTorch module + class CustomModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(1, 1) + self.linear2 = torch.nn.Linear(1, 1) + + def forward(self, input_dict: dict): + out1 = self.linear1(input_dict["A"].unsqueeze(1)) + out2 = self.linear2(input_dict["B"].unsqueeze(1)) + return out1 + out2 + + # Set manul seed so we get consistent output + torch.manual_seed(42) + + # Use Standard PyTorch model + model = torch.nn.Linear(2, 1) + predictor = TorchPredictor(model=model) + # Define our data + data = np.array([[1, 2], [3, 4]]) + predictions = predictor.predict(data, dtype=torch.float) + print(f"Standard model predictions: {predictions}") + print("---") + + # Use Custom PyTorch model with TorchPredictor + predictor = TorchPredictor(model=CustomModule()) + # Define our data and predict Customer model with TorchPredictor + data = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"]) + predictions = predictor.predict(data, dtype=torch.float) + print(f"Custom model predictions: {predictions}") + + .. testoutput:: + + Standard model predictions: {'predictions': array([[1.5487633], + [3.8037925]], dtype=float32)} + --- + Custom model predictions: predictions + 0 [0.61623406] + 1 [2.857038] + """ + return super(TorchPredictor, self).predict(data=data, dtype=dtype) + + def _arrays_to_tensors( + self, + numpy_arrays: Union[np.ndarray, Dict[str, np.ndarray]], + dtype: Optional[Union[torch.dtype, Dict[str, torch.dtype]]], + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + return convert_ndarray_batch_to_torch_tensor_batch( + numpy_arrays, + dtypes=dtype, + device=self.device, + ) + + def _tensor_to_array(self, tensor: torch.Tensor) -> np.ndarray: + if not isinstance(tensor, torch.Tensor): + raise ValueError( + "Expected the model to return either a torch.Tensor or a " + f"dict of torch.Tensor, but got {type(tensor)} instead. " + f"To support models with different output types, subclass " + f"TorchPredictor and override the `call_model` method to " + f"process the output into either torch.Tensor or Dict[" + f"str, torch.Tensor]." + ) + return tensor.cpu().detach().numpy() diff --git a/lib/python3.12/site-packages/ray/train/torch/torch_trainer.py b/lib/python3.12/site-packages/ray/train/torch/torch_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..27296664f122c9db780508ee0cfac3a2128c9ecf --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/torch_trainer.py @@ -0,0 +1,200 @@ +from typing import Any, Callable, Dict, Optional, Union + +from ray.train import Checkpoint, DataConfig, RunConfig, ScalingConfig +from ray.train.data_parallel_trainer import DataParallelTrainer +from ray.train.torch.config import TorchConfig +from ray.train.trainer import GenDataset +from ray.util import PublicAPI + + +@PublicAPI(stability="stable") +class TorchTrainer(DataParallelTrainer): + """A Trainer for data parallel PyTorch training. + + At a high level, this Trainer does the following: + + 1. Launches multiple workers as defined by the ``scaling_config``. + 2. Sets up a distributed PyTorch environment + on these workers as defined by the ``torch_config``. + 3. Ingests the input ``datasets`` based on the ``dataset_config``. + 4. Runs the input ``train_loop_per_worker(train_loop_config)`` + on all workers. + + For more details, see: + + * :ref:`PyTorch Guide ` + * :ref:`PyTorch Lightning Guide ` + * :ref:`Hugging Face Transformers Guide ` + + Example: + + .. testcode:: + :skipif: True + + import os + import tempfile + + import torch + from torch import nn + from torch.nn.parallel import DistributedDataParallel + + import ray + from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig + from ray.train.torch import TorchTrainer + + # If using GPUs, set this to True. + use_gpu = False + # Number of processes to run training on. + num_workers = 4 + + # Define your network structure. + class NeuralNetwork(nn.Module): + def __init__(self): + super(NeuralNetwork, self).__init__() + self.layer1 = nn.Linear(1, 32) + self.relu = nn.ReLU() + self.layer2 = nn.Linear(32, 1) + + def forward(self, input): + return self.layer2(self.relu(self.layer1(input))) + + # Training loop. + def train_loop_per_worker(config): + + # Read configurations. + lr = config["lr"] + batch_size = config["batch_size"] + num_epochs = config["num_epochs"] + + # Fetch training dataset. + train_dataset_shard = ray.train.get_dataset_shard("train") + + # Instantiate and prepare model for training. + model = NeuralNetwork() + model = ray.train.torch.prepare_model(model) + + # Define loss and optimizer. + loss_fn = nn.MSELoss() + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + + # Create data loader. + dataloader = train_dataset_shard.iter_torch_batches( + batch_size=batch_size, dtypes=torch.float + ) + + # Train multiple epochs. + for epoch in range(num_epochs): + + # Train epoch. + for batch in dataloader: + output = model(batch["input"]) + loss = loss_fn(output, batch["label"]) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + # Create checkpoint. + base_model = (model.module + if isinstance(model, DistributedDataParallel) else model) + checkpoint_dir = tempfile.mkdtemp() + torch.save( + {"model_state_dict": base_model.state_dict()}, + os.path.join(checkpoint_dir, "model.pt"), + ) + checkpoint = Checkpoint.from_directory(checkpoint_dir) + + # Report metrics and checkpoint. + ray.train.report({"loss": loss.item()}, checkpoint=checkpoint) + + + # Define configurations. + train_loop_config = {"num_epochs": 20, "lr": 0.01, "batch_size": 32} + scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu) + run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1)) + + # Define datasets. + train_dataset = ray.data.from_items( + [{"input": [x], "label": [2 * x + 1]} for x in range(2000)] + ) + datasets = {"train": train_dataset} + + # Initialize the Trainer. + trainer = TorchTrainer( + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + scaling_config=scaling_config, + run_config=run_config, + datasets=datasets + ) + + # Train the model. + result = trainer.fit() + + # Inspect the results. + final_loss = result.metrics["loss"] + + Args: + + train_loop_per_worker: The training function to execute on each worker. + This function can either take in zero arguments or a single ``Dict`` + argument which is set by defining ``train_loop_config``. + Within this function you can use any of the + :ref:`Ray Train Loop utilities `. + train_loop_config: A configuration ``Dict`` to pass in as an argument to + ``train_loop_per_worker``. + This is typically used for specifying hyperparameters. Passing large + datasets via `train_loop_config` is not recommended and may introduce + large overhead and unknown issues with serialization and deserialization. + torch_config: The configuration for setting up the PyTorch Distributed backend. + If set to None, a default configuration will be used in which + GPU training uses NCCL and CPU training uses Gloo. + scaling_config: The configuration for how to scale data parallel training. + ``num_workers`` determines how many Python processes are used for training, + and ``use_gpu`` determines whether or not each process should use GPUs. + See :class:`~ray.train.ScalingConfig` for more info. + run_config: The configuration for the execution of the training run. + See :class:`~ray.train.RunConfig` for more info. + datasets: The Ray Datasets to ingest for training. + Datasets are keyed by name (``{name: dataset}``). + Each dataset can be accessed from within the ``train_loop_per_worker`` + by calling ``ray.train.get_dataset_shard(name)``. + Sharding and additional configuration can be done by + passing in a ``dataset_config``. + dataset_config: The configuration for ingesting the input ``datasets``. + By default, all the Ray Dataset are split equally across workers. + See :class:`~ray.train.DataConfig` for more details. + resume_from_checkpoint: A checkpoint to resume training from. + This checkpoint can be accessed from within ``train_loop_per_worker`` + by calling ``ray.train.get_checkpoint()``. + metadata: Dict that should be made available via + `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + """ + + def __init__( + self, + train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]], + *, + train_loop_config: Optional[Dict] = None, + torch_config: Optional[TorchConfig] = None, + scaling_config: Optional[ScalingConfig] = None, + run_config: Optional[RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + dataset_config: Optional[DataConfig] = None, + metadata: Optional[Dict[str, Any]] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + ): + if not torch_config: + torch_config = TorchConfig() + + super(TorchTrainer, self).__init__( + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + backend_config=torch_config, + scaling_config=scaling_config, + dataset_config=dataset_config, + run_config=run_config, + datasets=datasets, + resume_from_checkpoint=resume_from_checkpoint, + metadata=metadata, + ) diff --git a/lib/python3.12/site-packages/ray/train/torch/train_loop_utils.py b/lib/python3.12/site-packages/ray/train/torch/train_loop_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bb48cf9fec7beaa8eda5596e2b9f709523f5e481 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/train_loop_utils.py @@ -0,0 +1,776 @@ +import collections +import logging +import os +import random +import types +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from packaging.version import Version +from torch.cuda.amp import GradScaler, autocast +from torch.nn.parallel import DistributedDataParallel +from torch.optim import Optimizer +from torch.utils.data import ( + DataLoader, + DistributedSampler, + IterableDataset, + RandomSampler, + SequentialSampler, +) + +from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag +from ray.air._internal.device_manager import ( + get_torch_device_manager_by_context, + get_torch_device_manager_by_device_type, +) +from ray.train._internal import session +from ray.train._internal.accelerator import Accelerator +from ray.train._internal.session import get_accelerator, set_accelerator +from ray.train.utils import _log_deprecation_warning +from ray.util.annotations import Deprecated, PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="stable") +def get_device() -> torch.device: + """Gets the correct torch device configured for this process. + + Returns the torch device for the current worker. If more than 1 GPU is + requested per worker, returns the device with the minimal device index. + + .. note:: + + If you requested multiple GPUs per worker, and want to get + the full list of torch devices, please use + :meth:`~ray.train.torch.get_devices`. + + Assumes that `CUDA_VISIBLE_DEVICES` is set and is a + superset of the `ray.get_gpu_ids()`. + + Examples: + + Example: Launched 2 workers on the current node, each with 1 GPU + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" + ray.get_gpu_ids() == [2] + torch.cuda.is_available() == True + get_device() == torch.device("cuda:0") + + Example: Launched 4 workers on the current node, each with 1 GPU + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" + ray.get_gpu_ids() == [2] + torch.cuda.is_available() == True + get_device() == torch.device("cuda:2") + + Example: Launched 2 workers on the current node, each with 2 GPUs + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" + ray.get_gpu_ids() == [2,3] + torch.cuda.is_available() == True + get_device() == torch.device("cuda:2") + + + You can move a model to device by: + + .. testcode:: + :skipif: True + + model.to(ray.train.torch.get_device()) + + Instead of manually checking the device type: + + .. testcode:: + :skipif: True + + model.to("cuda" if torch.cuda.is_available() else "cpu") + """ + from ray.air._internal import torch_utils + + record_extra_usage_tag(TagKey.TRAIN_TORCH_GET_DEVICE, "1") + return torch_utils.get_devices()[0] + + +@PublicAPI(stability="beta") +def get_devices() -> List[torch.device]: + """Gets the correct torch device list configured for this process. + + Assumes that `CUDA_VISIBLE_DEVICES` is set and is a + superset of the `ray.get_gpu_ids()`. + + + Examples: + + Example: Launched 2 workers on the current node, each with 1 GPU + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" + ray.get_gpu_ids() == [2] + torch.cuda.is_available() == True + get_devices() == [torch.device("cuda:0")] + + Example: Launched 4 workers on the current node, each with 1 GPU + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + ray.get_gpu_ids() == [2] + torch.cuda.is_available() == True + get_devices() == [torch.device("cuda:2")] + + Example: Launched 2 workers on the current node, each with 2 GPUs + + .. testcode:: + :skipif: True + + os.environ["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + ray.get_gpu_ids() == [2,3] + torch.cuda.is_available() == True + get_devices() == [torch.device("cuda:2"), torch.device("cuda:3")] + """ + + from ray.air._internal import torch_utils + + record_extra_usage_tag(TagKey.TRAIN_TORCH_GET_DEVICES, "1") + return torch_utils.get_devices() + + +@PublicAPI(stability="stable") +def prepare_model( + model: torch.nn.Module, + move_to_device: Union[bool, torch.device] = True, + parallel_strategy: Optional[str] = "ddp", + parallel_strategy_kwargs: Optional[Dict[str, Any]] = None, +) -> torch.nn.Module: + """Prepares the model for distributed execution. + + This allows you to use the same exact code regardless of number of + workers or the device type being used (CPU, GPU). + + Args: + model (torch.nn.Module): A torch model to prepare. + move_to_device: Either a boolean indiciating whether to move + the model to the correct device or an actual device to + move the model to. If set to False, the model needs + to manually be moved to the correct device. + parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models + in ``DistributedDataParallel``, ``FullyShardedDataParallel``, + or neither. + parallel_strategy_kwargs (Dict[str, Any]): Args to pass into + ``DistributedDataParallel`` or ``FullyShardedDataParallel`` + initialization if ``parallel_strategy`` is set to "ddp" + or "fsdp", respectively. + """ + if parallel_strategy == "fsdp" and Version(torch.__version__) < Version("1.11.0"): + raise ImportError( + "FullyShardedDataParallel requires torch>=1.11.0. " + "Run `pip install 'torch>=1.11.0'` to use FullyShardedDataParallel." + ) + + record_extra_usage_tag(TagKey.TRAIN_TORCH_PREPARE_MODEL, "1") + return get_accelerator(_TorchAccelerator).prepare_model( + model, + move_to_device=move_to_device, + parallel_strategy=parallel_strategy, + parallel_strategy_kwargs=parallel_strategy_kwargs, + ) + + +@PublicAPI(stability="stable") +def prepare_data_loader( + data_loader: torch.utils.data.DataLoader, + add_dist_sampler: bool = True, + move_to_device: bool = True, + auto_transfer: bool = True, +) -> torch.utils.data.DataLoader: + """Prepares :class:`~torch.utils.data.DataLoader` for distributed execution. + + This allows you to use the same exact code regardless of number of + workers or the device type being used (CPU, GPU). + + .. note:: + + This method adds a `DistributedSampler` to the `DataLoader` if the + number of training workers is greater than 1. If shuffling is + enabled on the original `DataLoader`, then `shuffle=True` will also + be passed into the `DistributedSampler` constructor. `shuffle=False` + on the original `DataLoader` also means that shuffling is disabled + on the sampler. + + With more than 1 worker, calling the `DistributedSampler.set_epoch` method + at the beginning of each epoch before creating the DataLoader iterator + is necessary to make shuffling work properly across multiple epochs. + Otherwise, the same ordering will be always used. + See: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler # noqa: E501 + + Example: + + .. testcode: + :skipif: True + + import torch + + import ray.train.torch + + train_dataloader = torch.utils.data.DataLoader( + ..., batch_size=..., shuffle=True + ) + train_dataloader = ray.train.torch.prepare_data_loader(train_loader) + + for epoch in range(10): + if ray.train.get_context().get_world_size() > 1: + # Required for the distributed sampler to shuffle properly across epochs + train_dataloader.sampler.set_epoch(epoch) + + for X, y in train_loader: + # No need to move data to GPU, this is done by `prepare_data_loader`! + # X, y = X.to("cuda"), y.to("cuda") + ... + + Args: + data_loader (torch.utils.data.DataLoader): The DataLoader to + prepare. + add_dist_sampler: Whether to add a DistributedSampler to + the provided DataLoader. + move_to_device: If set, automatically move the data + returned by the data loader to the correct device. + auto_transfer: If set and device is GPU, another CUDA stream + is created to automatically copy data from host (CPU) memory + to device (GPU) memory (the default CUDA stream still runs the + training procedure). If device is CPU, it will be disabled + regardless of the setting. This configuration will be ignored + if ``move_to_device`` is False. + """ + record_extra_usage_tag(TagKey.TRAIN_TORCH_PREPARE_DATALOADER, "1") + return get_accelerator(_TorchAccelerator).prepare_data_loader( + data_loader, + add_dist_sampler=add_dist_sampler, + move_to_device=move_to_device, + auto_transfer=auto_transfer, + ) + + +def _log_amp_deprecation_warning(): + # Keep V2 imports out of top-level V1 imports. + from ray.train.v2.torch.train_loop_utils import _TORCH_AMP_DEPRECATION_MESSAGE + + _log_deprecation_warning(_TORCH_AMP_DEPRECATION_MESSAGE) + + +@Deprecated +def accelerate(amp: bool = False) -> None: + """[Deprecated] Enables training optimizations. + + Arguments: + amp: If true, perform training with automatic mixed precision. + Otherwise, use full precision. + + .. warning:: ``train.torch.accelerate`` cannot be called more than once, and it + must be called before any other ``train.torch`` utility function. + """ + _log_amp_deprecation_warning() + try: + set_accelerator(_TorchAccelerator(amp=amp)) + except RuntimeError: + raise RuntimeError( + "An accelerator has already been set. Make sure " + "`train.torch.accelerate()` is not called multiple times, and is called " + "before any of the prepare methods." + ) + + +@Deprecated +def prepare_optimizer(optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer: + """[Deprecated] Wraps optimizer to support automatic mixed precision. + + Args: + optimizer (torch.optim.Optimizer): The DataLoader to prepare. + + Returns: + A wrapped optimizer. + """ + _log_amp_deprecation_warning() + return get_accelerator(_TorchAccelerator).prepare_optimizer(optimizer) + + +@Deprecated +def backward(tensor: torch.Tensor) -> None: + """[Deprecated] Computes the gradient of the specified tensor w.r.t. graph leaves. + + Args: + tensor (torch.Tensor): Tensor of which the derivative will be computed. + """ + _log_amp_deprecation_warning() + get_accelerator(_TorchAccelerator).backward(tensor) + + +@PublicAPI(stability="stable") +def enable_reproducibility(seed: int = 0) -> None: + """Limits sources of nondeterministic behavior. + + This function: + + * Seeds PyTorch, Python, and NumPy. + * Disables CUDA convolution benchmarking. + * Configures PyTorch to use determinstic algorithms. + * Seeds workers spawned for multi-process data loading. + + Args: + seed: The number to seed libraries and data workers with. + + .. warning:: ``train.torch.enable_reproducibility()`` can't guarantee + completely reproducible results across executions. To learn more, read + the `PyTorch notes on randomness + `_. + """ + get_accelerator(_TorchAccelerator).enable_reproducibility(seed) + + +@Deprecated +class TorchWorkerProfiler: + """Utility class for running PyTorch Profiler on a Train worker. + + Args: + trace_dir (Optional[str]): The directory to store traces on the + worker node. If ``None``, this will use a default temporary dir. + """ + + WORKER_TRACE_DIR_NAME = "pytorch_profiler_worker_traces" + + def __init__(self, trace_dir: Optional[str] = None): + raise DeprecationWarning( + "The `ray.train.torch.TorchWorkerProfiler` API is deprecated in Ray 2.0.", + ) + + +class _TorchAccelerator(Accelerator): + """A utility that implements methods to accelerate PyTorch training. + + Arguments: + amp: If true, perform training with automatic mixed precision. + Otherwise, use full precision. + """ + + def __init__(self, amp: bool = False): + self.amp_is_enabled = amp + self.scaler = GradScaler() if amp else None + self._seed = None + self.device_manager = get_torch_device_manager_by_context() + + def prepare_model( + self, + model: torch.nn.Module, + move_to_device: bool = True, + parallel_strategy: Optional[str] = "ddp", + parallel_strategy_kwargs: Optional[Dict[str, Any]] = None, + ) -> torch.nn.Module: + """Prepares the model for distributed execution. + + This allows you to use the same exact code regardless of number of + workers or the device type being used (CPU, GPU). + + Args: + model (torch.nn.Module): A torch model to prepare. + move_to_device: Whether to move the model to the correct + device. If set to False, the model needs to manually be moved + to the correct device. + parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models + in ``DistributedDataParallel``, ``FullyShardedDataParallel`` ( + Experimental), or neither. + parallel_strategy_kwargs (Dict[str, Any]): Args to pass into + ``DistributedDataParallel`` or ``FullyShardedDataParallel`` + initialization if ``parallel_strategy`` is set to "ddp" + or "fsdp", respectively. + """ + parallel_strategy_kwargs = parallel_strategy_kwargs or {} + + rank = session.get_local_rank() + + if isinstance(move_to_device, torch.device): + device = move_to_device + else: + device = get_device() + if isinstance(device, list): + device = device[0] + + if self.device_manager.is_available(): + self.device_manager.set_device(device) + + if move_to_device: + if rank == 0: + logger.info(f"Moving model to device: {device}") + else: + logger.debug(f"Moving model to device: {device}") + model = model.to(device) + + def model_get_state(self): + # `__getstate__` is an special method that informs pickle which attributes + # to serialize. This custom implementation ensures that the wrapped forward + # method and custom `__getstate__` method aren't serialized. + if hasattr(self, "_original_get_state"): + state = self._original_get_state() + state["__getstate__"] = state["_original_get_state"] + del state["_original_get_state"] + else: + # If model does not have a `__getstate__` already defined, use default + # implementation. + state = self.__dict__.copy() + del state["__getstate__"] + state["forward"] = state["_unwrapped_forward"] + del state["_unwrapped_forward"] + + return state + + if self.amp_is_enabled: + # Pickle cannot serialize the wrapped forward method. As a workaround, + # define a custom `__getstate__` method that unwraps the forward method. + model._unwrapped_forward = model.forward + model.forward = autocast()(model.forward) + + # TODO(amogkam): Replace below logic with a generic "unpack model" method. + # Replacing the `model.forward` method makes the model no longer + # serializable. When serializing the model, we have to override the + # `__getstate__` method to set back the original forward method. + if hasattr(model, "__getstate__"): + model._original_get_state = model.__getstate__ + # `__getstate__` must be a bound method rather than an callable attribute. + # See https://stackoverflow.com/questions/972/adding-a-method-to-an-existing-object-instance. # noqa: E501 + model.__getstate__ = types.MethodType(model_get_state, model) + + world_size = session.get_world_size() + + if parallel_strategy and world_size > 1: + if parallel_strategy == "ddp": + DataParallel = DistributedDataParallel + if self.device_manager.is_available() and device.type != "cpu": + parallel_strategy_kwargs = { + "device_ids": [device], + "output_device": device, + **parallel_strategy_kwargs, + } + else: + if not torch.cuda.is_available(): + raise RuntimeError( + "FSDP is only available with GPU-enabled " + "training. Set " + "`use_gpu=True` in your Trainer to train with " + "GPUs." + ) + from torch.distributed.fsdp import FullyShardedDataParallel + + DataParallel = FullyShardedDataParallel + if rank == 0: + logger.info(f"Wrapping provided model in {DataParallel.__name__}.") + else: + logger.debug(f"Wrapping provided model in {DataParallel.__name__}.") + model = DataParallel(model, **parallel_strategy_kwargs) + + return model + + def prepare_data_loader( + self, + data_loader: torch.utils.data.DataLoader, + add_dist_sampler: bool = True, + move_to_device: bool = True, + auto_transfer: bool = False, + ) -> torch.utils.data.DataLoader: + """Prepares DataLoader for distributed execution. + + This allows you to use the same exact code regardless of number of + workers or the device type being used (CPU, GPU). + + Args: + data_loader (torch.utils.data.DataLoader): The DataLoader to + prepare. + add_dist_sampler: Whether to add a DistributedSampler to + the provided DataLoader. + move_to_device: If set, automatically move the data + returned by the data loader to the correct device. + auto_transfer: (Experimental) If set and device is GPU, another CUDA stream + is created to automatically copy data from host (CPU) memory + to device (GPU) memory (the default CUDA stream still runs the + training procedure). If device is CPU, it will be disabled + regardless of the setting. This configuration will be ignored + if ``move_to_device`` is False. + """ + + world_size = session.get_world_size() + world_rank = session.get_world_rank() + + # Only add Distributed Sampler if the following conditions hold: + # 1. More than one training worker is being used. + # 2. A DistributedSampler has not already been added by the user. + # 3. The dataset is not an IterableDataset. Samplers do not worker with + # IterableDatasets. + if ( + world_size > 1 + and not isinstance(data_loader.sampler, DistributedSampler) + and not ( + hasattr(data_loader, "dataset") + and isinstance(data_loader.dataset, IterableDataset) + ) + and add_dist_sampler + ): + + def with_sampler(loader): + # Automatically set the DistributedSampler + + # If you're using a sampler, the DataLoader shuffle flag must be set to + # False. Shuffling is instead determined by the shuffle argument passed + # to the DistributedSampler constructor. + + # If no sampler is passed to the DataLoader constructor, Torch + # constructs a default sampler. The default sampler is a RandomSampler + # if shuffling is enabled and a SequentialSampler otherwise. DataLoader + # does not have a shuffle attribute, so we instead identify whether + # shuffling is enabled by checking the default sampler type. + shuffle = not isinstance(loader.sampler, SequentialSampler) + + def seeded_worker_init_fn( + worker_init_fn: Optional[Callable[[int], None]], + ): + def wrapper(worker_id: int): + worker_seed = torch.initial_seed() % 2**32 + np.random.seed(worker_seed) + random.seed(worker_seed) + if worker_init_fn: + worker_init_fn(worker_id) + + return wrapper + + worker_init_fn: Optional[Callable[[int], None]] = loader.worker_init_fn + generator: Optional[torch.Generator] = loader.generator + if self._seed is not None: + worker_init_fn = seeded_worker_init_fn(worker_init_fn) + generator = torch.Generator() + generator.manual_seed(self._seed) + + using_default_sampler = isinstance( + loader.sampler, (SequentialSampler, RandomSampler) + ) + if not using_default_sampler and world_rank == 0: + logger.warning( + f"The {loader.sampler.__class__.__name__} will be overwritten " + "with a DistributedSampler. You can disable this by setting " + "`with_sampler` to False in `prepare_data_loader`." + ) + + data_loader_args = { + "dataset": loader.dataset, + "batch_size": loader.batch_size, + "shuffle": False, + "num_workers": loader.num_workers, + "collate_fn": loader.collate_fn, + "pin_memory": loader.pin_memory, + "drop_last": loader.drop_last, + "timeout": loader.timeout, + "worker_init_fn": worker_init_fn, + "generator": generator, + "sampler": DistributedSampler(loader.dataset, shuffle=shuffle), + } + return DataLoader(**data_loader_args) + + data_loader = with_sampler(data_loader) + + if move_to_device: + device = get_device() + data_loader = _WrappedDataLoader(data_loader, device, auto_transfer) + + return data_loader + + def prepare_optimizer(self, optimizer: Optimizer) -> Optimizer: + """Wraps optimizer to support automatic mixed precision. + + Args: + optimizer (torch.optim.Optimizer): The DataLoader to prepare. + + Returns: + A wrapped optimizer. + """ + return _WrappedOptimizer(optimizer, scaler=self.scaler) + + def backward(self, tensor: torch.Tensor) -> None: + """Computes the gradient of the specified tensor w.r.t. graph leaves. + + Args: + tensor (torch.Tensor): Tensor of which the derivative will be computed. + """ + if self.amp_is_enabled: + self.scaler.scale(tensor).backward() + else: + tensor.backward() + + def enable_reproducibility(self, seed: int = 0) -> None: + """Limits sources of nondeterministic behavior.""" + self._seed = seed + + torch.manual_seed(seed) + random.seed(seed) + np.random.seed(seed) + + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + + # If you want to use deterministic algorithms with CUDA, then you need to set + # the CUBLAS_WORKSPACE_CONFIG environment variable; otherwise, Torch errors. + # See https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility. + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +class _WrappedDataLoader(DataLoader): + def __init__( + self, base_dataloader: DataLoader, device: torch.device, auto_transfer: bool + ): + self.__dict__.update(getattr(base_dataloader, "__dict__", {})) + self._dataloader = base_dataloader + self.dataloader_iter = None + self.device = device + + self.device_manager = get_torch_device_manager_by_device_type(device.type) + + # disable auto transfer (host->device) if cpu is used + if device.type != "cpu" and self.device_manager.supports_stream(): + self._auto_transfer = auto_transfer + else: + self._auto_transfer = False + # create a new device stream to move data from host to device concurrently + self._memcpy_stream = ( + self.device_manager.create_stream(device) + if device.type != "cpu" and self._auto_transfer + else None + ) + self.next_batch = None + + def _move_to_device(self, item): + if item is None: + return None + + def try_move_device(i): + try: + i = i.to(self.device, non_blocking=self._auto_transfer) + except AttributeError: + logger.debug(f"Item {i} cannot be moved to device " f"{self.device}.") + return i + + with self.device_manager.get_stream_context(self._memcpy_stream): + if isinstance(item, collections.abc.Mapping): + item_on_device = {k: self._move_to_device(v) for k, v in item.items()} + elif isinstance(item, tuple): + item_on_device = tuple(self._move_to_device(i) for i in item) + elif isinstance(item, list): + item_on_device = [self._move_to_device(i) for i in item] + elif isinstance(item, torch.Tensor): + item_on_device = try_move_device(item) + else: + logger.debug( + f"Data type {type(item)} doesn't support being moved to device." + ) + item_on_device = item + + return item_on_device + + def _wait_for_batch(self, item): + if self._memcpy_stream is None: + return + # Reference: + # https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html + # The training stream (current) needs to wait until + # the memory copy stream finishes. + curr_stream = self.device_manager.get_current_stream() + curr_stream.wait_stream(self._memcpy_stream) + # When a tensor is used by CUDA streams different from + # its original allocator, we need to call ``record_stream`` + # to inform the allocator of all these streams. Otherwise, + # the tensor might be freed once it is no longer used by + # the creator stream. + for i in item: + # The Pytorch DataLoader has no restrictions on what is outputted for + # each batch. We should only ``record_stream`` if the item has the + # ability to do so. + try: + i.record_stream(curr_stream) + except AttributeError: + pass + + def __len__(self): + return len(self._dataloader) + + def _prefetch_next_batch(self): + next_batch = next(self.dataloader_iter, None) + self.next_batch = self._move_to_device(next_batch) + + def __iter__(self): + self.dataloader_iter = iter(self._dataloader) + self._prefetch_next_batch() + return self + + def __next__(self): + next_batch = self.next_batch + if next_batch is None: + raise StopIteration + self._wait_for_batch(next_batch) + self._prefetch_next_batch() + return next_batch + + +class _WrappedOptimizer(Optimizer): + def __init__(self, optimizer: Optimizer, scaler: Optional[GradScaler] = None): + self.optimizer = optimizer + self.scaler = scaler + + @property + def state(self): + return self.optimizer.state + + @state.setter + def state(self, state): + self.optimizer.state = state + + @property + def param_groups(self): + return self.optimizer.param_groups + + @param_groups.setter + def param_groups(self, param_groups): + self.optimizer.param_groups = param_groups + + @property + def defaults(self): + return self.optimizer.defaults + + @defaults.setter + def defaults(self, defaults): + self.optimizer.defaults = defaults + + def add_param_group(self, param_group): + self.optimizer.add_param_group(param_group) + + def load_state_dict(self, state_dict): + self.optimizer.load_state_dict(state_dict) + + def state_dict(self): + return self.optimizer.state_dict() + + def zero_grad(self): + self.optimizer.zero_grad() + + def step(self, closure=None): + if self.scaler is not None: + self.scaler.step(self.optimizer, closure) + self.scaler.update() + else: + self.optimizer.step(closure) diff --git a/lib/python3.12/site-packages/ray/train/torch/xla/__init__.py b/lib/python3.12/site-packages/ray/train/torch/xla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ea32abc8c9d7b5514b3c5680e06ee0ab334eb4eb --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/xla/__init__.py @@ -0,0 +1,5 @@ +from ray.train.torch.xla.config import TorchXLAConfig + +__all__ = [ + "TorchXLAConfig", +] diff --git a/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..408d6c4bd9bf7968644e57bf797368946771e9a8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d0a0c65f95de4cd4ec0219f3ef8bd5e750065b9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/torch/xla/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/torch/xla/config.py b/lib/python3.12/site-packages/ray/train/torch/xla/config.py new file mode 100644 index 0000000000000000000000000000000000000000..71208f4a2ddc2a83b9467c9c22d331b8f7237a3d --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/torch/xla/config.py @@ -0,0 +1,171 @@ +import logging +import os +import re +import shutil +import uuid +from dataclasses import dataclass + +import ray +from ray.train._internal.base_worker_group import BaseWorkerGroup +from ray.train._internal.utils import get_address_and_port +from ray.train.backend import Backend +from ray.train.torch import TorchConfig +from ray.util import PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="alpha") +@dataclass +class TorchXLAConfig(TorchConfig): + """ + Configuration for torch XLA setup. + See https://pytorch.org/xla/release/1.13/index.html for more info. + Currently, only "neuron_cores" accelerator (AwsNeuronXLABackend) + is supported with xrt runtime. + """ + + neuron_parallel_compile: bool = False + + @property + def backend_cls(self): + return _TorchAwsNeuronXLABackend + + +def _kill_xrt_server(): + import subprocess + + subprocess.call(["pkill", "-f", "xrt_run_server"]) + + +def _set_xla_env_vars(): + # https://pytorch.org/docs/1.13/elastic/run.html#environment-variables + context = ray.train.get_context() + + os.environ["LOCAL_RANK"] = str(context.get_local_rank()) + os.environ["RANK"] = str(context.get_world_rank()) + os.environ["LOCAL_WORLD_SIZE"] = str(context.get_local_world_size()) + os.environ["WORLD_SIZE"] = str(context.get_world_size()) + os.environ["GROUP_RANK"] = str(context.get_node_rank()) + os.environ["GROUP_WORLD_SIZE"] = str( + context.get_world_size() / context.get_local_world_size() + ) + os.environ["ROLE_RANK"] = str(context.get_world_rank()) + os.environ["ROLE_WORLD_RANK"] = str(context.get_world_rank()) + os.environ["ROLE_WORLD_SIZE"] = str(context.get_world_size()) + + # EFA and XLA setup + # https://github.com/aws/libfabric/blob/master/prov/efa/src/rxr/rxr_init.c + # https://github.com/aws-neuron/aws-neuron-samples/blob/master/torch-neuronx/training/dp_bert_hf_pretrain/run_dp_bert_large_hf_pretrain_bf16_s128.sh # noqa + os.environ["FI_PROVIDER"] = "efa" + os.environ["FI_EFA_USE_DEVICE_RDMA"] = "1" + os.environ["FI_EFA_FORK_SAFE"] = "1" + os.environ["XLA_TRANSFER_SEED_ASYNC"] = "1" + os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1" + + +def _setup_xla_torch_process_group(): + try: + import torch.distributed as dist + import torch_xla.core.xla_model as xm # noqa F401 + import torch_xla.distributed.xla_backend # noqa F401 + + dist.init_process_group("xla") + except ImportError: + raise ImportError("torch_xla must be installed to use torch_xla backend.") + + +# The following env vars enable Neuron graph extraction for parallel compilation +# Note: model outputs are invalid and should be ignored while these env vars are set +def _set_neuron_parallel_compile_env_vars(): + os.environ["NEURON_PARALLEL_COMPILE"] = "1" + os.environ["NEURON_EXTRACT_GRAPHS_ONLY"] = "1" + os.environ["NEURON_FALL_BACK_TO_NULL_NEFF"] = "1" + + +# Compile previously extracted Neuron graphs +def _neuron_compile_extracted_graphs(): + try: + from libneuronxla.neuron_cc_cache import CacheUrl + from libneuronxla.neuron_parallel_compile import parallel_compile + except ImportError: + raise ImportError( + "libneuronxla must be installed to use Neuron parallel compilation." + ) + + # Only 1 worker per node should run parallel_compile() + if os.environ.get("LOCAL_RANK") == "0": + logger.info("Compiling extracted graphs on local rank0 worker") + + parallel_compile_workdir = ( + f"/tmp/{os.environ.get('USER','no-user')}/parallel_compile_workdir/" + ) + if os.path.exists(parallel_compile_workdir): + shutil.rmtree(parallel_compile_workdir) + os.makedirs(parallel_compile_workdir, exist_ok=True) + + # Users can set the cache directory using --cache_dir in NEURON_CC_FLAGS or by + # using NEURON_COMPILE_CACHE_URL. --cache_dir takes precedence. + explicit_cache_dir = None + if neuron_cc_flags := os.environ.get("NEURON_CC_FLAGS"): + if s := re.search(r"--cache_dir[= ](\S+)", neuron_cc_flags): + explicit_cache_dir = s.group(1) + + parallel_compile( + parallel_compile_workdir, + CacheUrl.get_cache_url(explicit_cache_dir), + ) + + +class _TorchAwsNeuronXLABackend(Backend): + unique_run_id: str = str(uuid.uuid4()) + + def on_start(self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig): + """Logic ran right before training is started.""" + + # On previous worker failure, we don't run graceful shutdown on workers. + # This would leak any running xrt server. + worker_group.execute(_kill_xrt_server) + + # Get master address and port from the first worker. + master_addr, master_port = worker_group.execute_single(0, get_address_and_port) + + def set_env_vars(addr, port): + os.environ["MASTER_ADDR"] = addr + os.environ["MASTER_PORT"] = str(port) + # To trigger the xrt server + os.environ["TORCHELASTIC_RUN_ID"] = self.unique_run_id + + # Set the env vars on all workers. + worker_group.execute(set_env_vars, addr=master_addr, port=master_port) + + # Set up env vars for neuron parallel compilation graph extraction + if backend_config.neuron_parallel_compile: + logger.info("Extracting graphs for Neuron parallel compilation") + worker_group.execute(_set_neuron_parallel_compile_env_vars) + + def on_training_start( + self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig + ): + """ + Configure the environment variables for the worker group. + And initialize the xla distributed process group. + TODO: Current setup only supports homogenous cluster with + neuron_cores accelerator and xrt runtime. + """ + worker_group.execute(_set_xla_env_vars) + worker_group.execute(_setup_xla_torch_process_group) + + def on_shutdown( + self, worker_group: BaseWorkerGroup, backend_config: TorchXLAConfig + ): + """ + Logic ran right after training is finished. + This is a sanity cleanup to kill xrt server, and to optionally + run neuron parallel graph compilation + """ + worker_group.execute(_kill_xrt_server) + + # Compile the extracted graphs. This must run at end of training. + if backend_config.neuron_parallel_compile: + worker_group.execute(_neuron_compile_extracted_graphs) diff --git a/lib/python3.12/site-packages/ray/train/trainer.py b/lib/python3.12/site-packages/ray/train/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..725ba029d766c37d838e456f58abd2182e836a8d --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/trainer.py @@ -0,0 +1,194 @@ +import logging +import traceback +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, TypeVar, Union + +from ray.air._internal.util import ( + StartTraceback, + StartTracebackWithWorkerRank, + skip_exceptions, +) +from ray.data import Dataset +from ray.train import Checkpoint, DataConfig +from ray.train._internal.backend_executor import ( + BackendExecutor, + InactiveWorkerGroupError, + TrainBackendError, + TrainingWorkerError, +) +from ray.train._internal.session import _TrainingResult, _TrainSession, get_session +from ray.train._internal.utils import ActorWrapper +from ray.train.backend import BackendConfig +from ray.train.base_trainer import ( # noqa: F401 + BaseTrainer, + GenDataset, + TrainingFailedError, +) +from ray.util.annotations import DeveloperAPI + +T = TypeVar("T") +S = TypeVar("S") + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class TrainingIterator: + """An iterator over Train results. Returned by ``trainer.run_iterator``.""" + + def __init__( + self, + backend_executor: Union[BackendExecutor, ActorWrapper], + backend_config: BackendConfig, + train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]], + datasets: Dict[str, Dataset], + metadata: Dict[str, Any], + data_config: DataConfig, + checkpoint: Optional[Union[Dict, str, Path, Checkpoint]], + ): + self._backend_executor = backend_executor + self._backend = backend_config.backend_cls() + self._train_func = train_func + self._datasets = datasets + self._metadata = metadata + self._data_config = data_config + + self._start_training( + train_func=train_func, + datasets=self._datasets, + metadata=self._metadata, + data_config=self._data_config, + checkpoint=checkpoint, + ) + + self._finished_training = False + + def __iter__(self): + return self + + def _start_training( + self, + train_func, + datasets, + metadata, + data_config, + checkpoint: Optional[Checkpoint] = None, + ): + tune_session: _TrainSession = get_session() + assert tune_session, "`_start_training` should only be called from within Tune" + storage = tune_session.storage + + self._run_with_error_handling( + lambda: self._backend_executor.start_training( + train_func=train_func, + datasets=datasets, + metadata=metadata, + data_config=data_config, + storage=storage, + checkpoint=checkpoint, + ) + ) + + def _run_with_error_handling(self, func: Callable): + try: + return func() + except TrainingWorkerError: + # TODO(ml-team): This Train fault-tolerance code doesn't get used + # since max_retries=0 + # Workers have already been restarted. + logger.info( + "Workers have been successfully restarted. Resuming " + "training from latest checkpoint." + ) + self._start_training( + self._train_func, + self._datasets, + self._metadata, + self._data_config, + ) + return self._run_with_error_handling(func) + except InactiveWorkerGroupError: + raise RuntimeError( + "This Trainer is not active. It is either shutdown " + "already or never started in the first place. " + "Either create a new Trainer or start this one." + ) from None + except TrainBackendError: + raise RuntimeError( + "Training failed. You should not be seeing " + "this error and this is a bug. Please create " + "a new issue at " + "https://github.com/ray-project/ray." + ) from None + + def __next__(self): + if self.is_finished(): + self._backend_executor.report_final_run_status(errored=False) + raise StopIteration + try: + next_results = self._run_with_error_handling(self._fetch_next_result) + if next_results is None: + self._backend_executor.report_final_run_status(errored=False) + self._run_with_error_handling(self._finish_training) + self._finished_training = True + raise StopIteration + else: + return next_results + except StartTraceback as e: + # If this is a StartTraceback, then this is a user error. + # We raise it directly + if isinstance(e, StartTracebackWithWorkerRank): + failed_rank = e.worker_rank + else: + failed_rank = None + + # Extract the stack trace from the exception + e = skip_exceptions(e) + stack_trace = "".join( + traceback.format_exception(type(e), e, e.__traceback__) + ) + + self._backend_executor.report_final_run_status( + errored=True, stack_trace=stack_trace, failed_rank=failed_rank + ) + try: + # Exception raised in at least one training worker. Immediately raise + # this error to the user and do not attempt to terminate gracefully. + self._backend_executor.shutdown(graceful_termination=False) + self._finished_training = True + except Exception: + pass + raise + + def _fetch_next_result(self) -> Optional[List[Dict]]: + """Fetch next results produced by ``session.report()`` from each worker. + + Assumes ``start_training`` has already been called. + + Returns: + A list of dictionaries of values passed to ``session.report()`` from + each worker. Each item corresponds to an intermediate result + a single worker. If there are no more items to fetch, + returns None. + """ + results = self._backend_executor.get_next_results() + if results is None: + return None + assert all(isinstance(result, _TrainingResult) for result in results) + return results + + def _finish_training(self): + """Finish training and return final results. Propagate any exceptions. + + Blocks until training is finished on all workers. + + Assumes `start_training` has already been called. + + Returns: + A list of return values from calling ``train_func`` on each worker. + Each item corresponds to the return value from a single worker. + """ + return self._backend_executor.finish_training() + + def is_finished(self) -> bool: + return self._finished_training diff --git a/lib/python3.12/site-packages/ray/train/utils.py b/lib/python3.12/site-packages/ray/train/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a9163d6852cb62b4331cb701558043cde38b4caf --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/utils.py @@ -0,0 +1,36 @@ +import warnings + +from ray.util.annotations import RayDeprecationWarning + + +def _copy_doc(copy_func): + def wrapped(func): + func.__doc__ = copy_func.__doc__ + return func + + return wrapped + + +def _log_deprecation_warning(message: str): + warnings.warn( + message, + RayDeprecationWarning, + stacklevel=2, + ) + + +def _in_ray_train_worker() -> bool: + from ray.train.v2._internal.constants import is_v2_enabled + + if is_v2_enabled(): + from ray.train.v2._internal.util import ( + _in_ray_train_worker as _in_ray_train_v2_worker, + ) + + return _in_ray_train_v2_worker() + else: + from ray.train._internal.session import ( + _in_ray_train_worker as _in_ray_train_v1_worker, + ) + + return _in_ray_train_v1_worker() diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__init__.py b/lib/python3.12/site-packages/ray/train/xgboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..447515b95b44d380a934237cc51630b5b86168e7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/__init__.py @@ -0,0 +1,21 @@ +from ray.train.v2._internal.constants import is_v2_enabled +from ray.train.xgboost._xgboost_utils import RayTrainReportCallback +from ray.train.xgboost.config import XGBoostConfig +from ray.train.xgboost.xgboost_checkpoint import XGBoostCheckpoint +from ray.train.xgboost.xgboost_predictor import XGBoostPredictor +from ray.train.xgboost.xgboost_trainer import XGBoostTrainer + +if is_v2_enabled(): + from ray.train.v2.xgboost.config import XGBoostConfig # noqa: F811 + from ray.train.v2.xgboost.xgboost_trainer import XGBoostTrainer # noqa: F811 + +__all__ = [ + "RayTrainReportCallback", + "XGBoostCheckpoint", + "XGBoostConfig", + "XGBoostPredictor", + "XGBoostTrainer", +] + + +# DO NOT ADD ANYTHING AFTER THIS LINE. diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0abb798239506ef3a512cc623adbb5ef0050d6ec Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/_xgboost_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/_xgboost_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9665d617c1e4561b9422932b7af0e71f1fb5bdb9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/_xgboost_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2027732aa1e64bb07159f11d532e6952964a918b Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/v2.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/v2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ce419129f04186ee0e78bc7201ee60538cacab6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/v2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d670bbed42eee28579b5da18cb063e0c312ae32 Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_predictor.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_predictor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2599753f1a051e4d2377762c115c817891bbf10c Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_predictor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_trainer.cpython-312.pyc b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88ea64b77f60ae83ab035359628c61e84140c53c Binary files /dev/null and b/lib/python3.12/site-packages/ray/train/xgboost/__pycache__/xgboost_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/train/xgboost/_xgboost_utils.py b/lib/python3.12/site-packages/ray/train/xgboost/_xgboost_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e200e073f0f2d5fa97b7fc8c93fadda7cd168df5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/_xgboost_utils.py @@ -0,0 +1,252 @@ +import tempfile +from abc import abstractmethod +from collections import OrderedDict +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Dict, List, Optional, Union + +from xgboost.core import Booster + +import ray.train +from ray.train import Checkpoint +from ray.tune.utils import flatten_dict +from ray.util.annotations import PublicAPI + +try: + from xgboost.callback import TrainingCallback +except ImportError: + + class TrainingCallback: + pass + + +class RayReportCallback(TrainingCallback): + CHECKPOINT_NAME = "model.ubj" + + def __init__( + self, + metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, + filename: str = CHECKPOINT_NAME, + frequency: int = 0, + checkpoint_at_end: bool = True, + results_postprocessing_fn: Optional[ + Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]] + ] = None, + ): + if isinstance(metrics, str): + metrics = [metrics] + self._metrics = metrics + self._filename = filename + self._frequency = frequency + self._checkpoint_at_end = checkpoint_at_end + self._results_postprocessing_fn = results_postprocessing_fn + + # Keeps track of the eval metrics from the last iteration, + # so that the latest metrics can be reported with the checkpoint + # at the end of training. + self._evals_log = None + # Keep track of the last checkpoint iteration to avoid double-checkpointing + # when using `checkpoint_at_end=True`. + self._last_checkpoint_iteration = None + + @classmethod + def get_model( + cls, + checkpoint: Checkpoint, + filename: str = CHECKPOINT_NAME, + ) -> Booster: + """Retrieve the model stored in a checkpoint reported by this callback. + + Args: + checkpoint: The checkpoint object returned by a training run. + The checkpoint should be saved by an instance of this callback. + filename: The filename to load the model from, which should match + the filename used when creating the callback. + + Returns: + The model loaded from the checkpoint. + """ + with checkpoint.as_directory() as checkpoint_path: + booster = Booster() + booster.load_model(Path(checkpoint_path, filename).as_posix()) + return booster + + def _get_report_dict(self, evals_log): + if isinstance(evals_log, OrderedDict): + # xgboost>=1.3 + result_dict = flatten_dict(evals_log, delimiter="-") + for k in list(result_dict): + result_dict[k] = result_dict[k][-1] + else: + # xgboost<1.3 + result_dict = dict(evals_log) + if not self._metrics: + report_dict = result_dict + else: + report_dict = {} + for key in self._metrics: + if isinstance(self._metrics, dict): + metric = self._metrics[key] + else: + metric = key + report_dict[key] = result_dict[metric] + + if self._results_postprocessing_fn: + report_dict = self._results_postprocessing_fn(report_dict) + + return report_dict + + @abstractmethod + def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]: + """Get checkpoint from model. + + This method needs to be implemented by subclasses. + """ + raise NotImplementedError + + @abstractmethod + def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster): + """Save checkpoint and report metrics corresonding to this checkpoint. + + This method needs to be implemented by subclasses. + """ + raise NotImplementedError + + @abstractmethod + def _report_metrics(self, report_dict: Dict): + """Report Metrics. + + This method needs to be implemented by subclasses. + """ + raise NotImplementedError + + def after_iteration(self, model: Booster, epoch: int, evals_log: Dict): + self._evals_log = evals_log + + checkpointing_disabled = self._frequency == 0 + # Ex: if frequency=2, checkpoint at epoch 1, 3, 5, ... (counting from 0) + should_checkpoint = ( + not checkpointing_disabled and (epoch + 1) % self._frequency == 0 + ) + + report_dict = self._get_report_dict(evals_log) + if should_checkpoint: + self._last_checkpoint_iteration = epoch + self._save_and_report_checkpoint(report_dict, model) + + else: + self._report_metrics(report_dict) + + def after_training(self, model: Booster) -> Booster: + if not self._checkpoint_at_end: + return model + + if ( + self._last_checkpoint_iteration is not None + and model.num_boosted_rounds() - 1 == self._last_checkpoint_iteration + ): + # Avoids a duplicate checkpoint if the checkpoint frequency happens + # to align with the last iteration. + return model + + report_dict = self._get_report_dict(self._evals_log) if self._evals_log else {} + self._save_and_report_checkpoint(report_dict, model) + + return model + + +@PublicAPI(stability="beta") +class RayTrainReportCallback(RayReportCallback): + """XGBoost callback to save checkpoints and report metrics. + + Args: + metrics: Metrics to report. If this is a list, + each item describes the metric key reported to XGBoost, + and it will be reported under the same name. + This can also be a dict of {: }, + which can be used to rename xgboost default metrics. + filename: Customize the saved checkpoint file type by passing + a filename. Defaults to "model.ubj". + frequency: How often to save checkpoints, in terms of iterations. + Defaults to 0 (no checkpoints are saved during training). + checkpoint_at_end: Whether or not to save a checkpoint at the end of training. + results_postprocessing_fn: An optional Callable that takes in + the metrics dict that will be reported (after it has been flattened) + and returns a modified dict. For example, this can be used to + average results across CV fold when using ``xgboost.cv``. + + Examples + -------- + + Reporting checkpoints and metrics to Ray Tune when running many + independent xgboost trials (without data parallelism within a trial). + + .. testcode:: + :skipif: True + + import xgboost + + from ray.tune import Tuner + from ray.train.xgboost import RayTrainReportCallback + + def train_fn(config): + # Report log loss to Ray Tune after each validation epoch. + bst = xgboost.train( + ..., + callbacks=[ + RayTrainReportCallback( + metrics={"loss": "eval-logloss"}, frequency=1 + ) + ], + ) + + tuner = Tuner(train_fn) + results = tuner.fit() + + Loading a model from a checkpoint reported by this callback. + + .. testcode:: + :skipif: True + + from ray.train.xgboost import RayTrainReportCallback + + # Get a `Checkpoint` object that is saved by the callback during training. + result = trainer.fit() + booster = RayTrainReportCallback.get_model(result.checkpoint) + + """ + + def __init__( + self, + metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, + filename: str = RayReportCallback.CHECKPOINT_NAME, + frequency: int = 0, + checkpoint_at_end: bool = True, + results_postprocessing_fn: Optional[ + Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]] + ] = None, + ): + super().__init__( + metrics=metrics, + filename=filename, + frequency=frequency, + checkpoint_at_end=checkpoint_at_end, + results_postprocessing_fn=results_postprocessing_fn, + ) + + @contextmanager + def _get_checkpoint(self, model: Booster) -> Optional[Checkpoint]: + # NOTE: The world rank returns None for Tune usage without Train. + if ray.train.get_context().get_world_rank() in (0, None): + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + model.save_model(Path(temp_checkpoint_dir, self._filename).as_posix()) + yield Checkpoint(temp_checkpoint_dir) + else: + yield None + + def _save_and_report_checkpoint(self, report_dict: Dict, model: Booster): + with self._get_checkpoint(model=model) as checkpoint: + ray.train.report(report_dict, checkpoint=checkpoint) + + def _report_metrics(self, report_dict: Dict): + ray.train.report(report_dict) diff --git a/lib/python3.12/site-packages/ray/train/xgboost/config.py b/lib/python3.12/site-packages/ray/train/xgboost/config.py new file mode 100644 index 0000000000000000000000000000000000000000..0da6648ec705c5c175f1d80c11928205c453d032 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/config.py @@ -0,0 +1,202 @@ +import json +import logging +import os +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Optional + +import xgboost +from packaging.version import Version +from xgboost import RabitTracker +from xgboost.collective import CommunicatorContext + +import ray +from ray.train._internal.base_worker_group import BaseWorkerGroup +from ray.train.backend import Backend, BackendConfig + +logger = logging.getLogger(__name__) + + +@dataclass +class XGBoostConfig(BackendConfig): + """Configuration for xgboost collective communication setup. + + Ray Train will set up the necessary coordinator processes and environment + variables for your workers to communicate with each other. + Additional configuration options can be passed into the + `xgboost.collective.CommunicatorContext` that wraps your own `xgboost.train` code. + + See the `xgboost.collective` module for more information: + https://github.com/dmlc/xgboost/blob/master/python-package/xgboost/collective.py + + Args: + xgboost_communicator: The backend to use for collective communication for + distributed xgboost training. For now, only "rabit" is supported. + """ + + xgboost_communicator: str = "rabit" + + @property + def train_func_context(self): + @contextmanager + def collective_communication_context(): + with CommunicatorContext(**_get_xgboost_args()): + yield + + return collective_communication_context + + @property + def backend_cls(self): + if self.xgboost_communicator == "rabit": + return ( + _XGBoostRabitBackend + if Version(xgboost.__version__) >= Version("2.1.0") + else _XGBoostRabitBackend_pre_xgb210 + ) + + raise NotImplementedError(f"Unsupported backend: {self.xgboost_communicator}") + + +class _XGBoostRabitBackend(Backend): + def __init__(self): + self._tracker: Optional[RabitTracker] = None + self._wait_thread: Optional[threading.Thread] = None + + def _setup_xgboost_distributed_backend(self, worker_group: BaseWorkerGroup): + # Set up the rabit tracker on the Train driver. + num_workers = len(worker_group) + rabit_args = {"n_workers": num_workers} + train_driver_ip = ray.util.get_node_ip_address() + + # NOTE: sortby="task" is needed to ensure that the xgboost worker ranks + # align with Ray Train worker ranks. + # The worker ranks will be sorted by `dmlc_task_id`, + # which is defined below. + self._tracker = RabitTracker( + n_workers=num_workers, host_ip=train_driver_ip, sortby="task" + ) + self._tracker.start() + + # The RabitTracker is started in a separate thread, and the + # `wait_for` method must be called for `worker_args` to return. + self._wait_thread = threading.Thread(target=self._tracker.wait_for, daemon=True) + self._wait_thread.start() + + rabit_args.update(self._tracker.worker_args()) + + start_log = ( + "RabitTracker coordinator started with parameters:\n" + f"{json.dumps(rabit_args, indent=2)}" + ) + logger.debug(start_log) + + def set_xgboost_communicator_args(args): + import ray.train + + args["dmlc_task_id"] = ( + f"[xgboost.ray-rank={ray.train.get_context().get_world_rank():08}]:" + f"{ray.get_runtime_context().get_actor_id()}" + ) + + _set_xgboost_args(args) + + worker_group.execute(set_xgboost_communicator_args, rabit_args) + + def on_training_start( + self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig + ): + assert backend_config.xgboost_communicator == "rabit" + self._setup_xgboost_distributed_backend(worker_group) + + def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig): + timeout = 5 + + if self._wait_thread is not None: + self._wait_thread.join(timeout=timeout) + + if self._wait_thread.is_alive(): + logger.warning( + "During shutdown, the RabitTracker thread failed to join " + f"within {timeout} seconds. " + "The process will still be terminated as part of Ray actor cleanup." + ) + + +class _XGBoostRabitBackend_pre_xgb210(Backend): + def __init__(self): + self._tracker: Optional[RabitTracker] = None + + def _setup_xgboost_distributed_backend(self, worker_group: BaseWorkerGroup): + # Set up the rabit tracker on the Train driver. + num_workers = len(worker_group) + rabit_args = {"DMLC_NUM_WORKER": num_workers} + train_driver_ip = ray.util.get_node_ip_address() + + # NOTE: sortby="task" is needed to ensure that the xgboost worker ranks + # align with Ray Train worker ranks. + # The worker ranks will be sorted by `DMLC_TASK_ID`, + # which is defined below. + self._tracker = RabitTracker( + n_workers=num_workers, host_ip=train_driver_ip, sortby="task" + ) + self._tracker.start(n_workers=num_workers) + + worker_args = self._tracker.worker_envs() + rabit_args.update(worker_args) + + start_log = ( + "RabitTracker coordinator started with parameters:\n" + f"{json.dumps(rabit_args, indent=2)}" + ) + logger.debug(start_log) + + def set_xgboost_env_vars(): + import ray.train + + for k, v in rabit_args.items(): + os.environ[k] = str(v) + + # Ranks are assigned in increasing order of the worker's task id. + # This task id will be sorted by increasing world rank. + os.environ["DMLC_TASK_ID"] = ( + f"[xgboost.ray-rank={ray.train.get_context().get_world_rank():08}]:" + f"{ray.get_runtime_context().get_actor_id()}" + ) + + worker_group.execute(set_xgboost_env_vars) + + def on_training_start( + self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig + ): + assert backend_config.xgboost_communicator == "rabit" + self._setup_xgboost_distributed_backend(worker_group) + + def on_shutdown(self, worker_group: BaseWorkerGroup, backend_config: XGBoostConfig): + if not self._tracker: + return + + timeout = 5 + self._tracker.thread.join(timeout=timeout) + + if self._tracker.thread.is_alive(): + logger.warning( + "During shutdown, the RabitTracker thread failed to join " + f"within {timeout} seconds. " + "The process will still be terminated as part of Ray actor cleanup." + ) + + +_xgboost_args: dict = {} +_xgboost_args_lock = threading.Lock() + + +def _set_xgboost_args(args): + with _xgboost_args_lock: + global _xgboost_args + _xgboost_args = args + + +def _get_xgboost_args() -> dict: + with _xgboost_args_lock: + return _xgboost_args diff --git a/lib/python3.12/site-packages/ray/train/xgboost/v2.py b/lib/python3.12/site-packages/ray/train/xgboost/v2.py new file mode 100644 index 0000000000000000000000000000000000000000..d5e5cb052df45cb856ca0c38d04fe3c0366fe89c --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/v2.py @@ -0,0 +1,128 @@ +import logging +from typing import Any, Callable, Dict, Optional, Union + +import ray.train +from ray.train import Checkpoint +from ray.train.data_parallel_trainer import DataParallelTrainer +from ray.train.trainer import GenDataset +from ray.train.xgboost import XGBoostConfig + +logger = logging.getLogger(__name__) + + +class XGBoostTrainer(DataParallelTrainer): + """A Trainer for distributed data-parallel XGBoost training. + + Example + ------- + + .. testcode:: + :skipif: True + + import xgboost + + import ray.data + import ray.train + from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer + + def train_fn_per_worker(config: dict): + # (Optional) Add logic to resume training state from a checkpoint. + # ray.train.get_checkpoint() + + # 1. Get the dataset shard for the worker and convert to a `xgboost.DMatrix` + train_ds_iter, eval_ds_iter = ( + ray.train.get_dataset_shard("train"), + ray.train.get_dataset_shard("validation"), + ) + train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize() + + train_df, eval_df = train_ds.to_pandas(), eval_ds.to_pandas() + train_X, train_y = train_df.drop("y", axis=1), train_df["y"] + eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"] + + dtrain = xgboost.DMatrix(train_X, label=train_y) + deval = xgboost.DMatrix(eval_X, label=eval_y) + + params = { + "tree_method": "approx", + "objective": "reg:squarederror", + "eta": 1e-4, + "subsample": 0.5, + "max_depth": 2, + } + + # 2. Do distributed data-parallel training. + # Ray Train sets up the necessary coordinator processes and + # environment variables for your workers to communicate with each other. + bst = xgboost.train( + params, + dtrain=dtrain, + evals=[(deval, "validation")], + num_boost_round=10, + callbacks=[RayTrainReportCallback()], + ) + + train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)]) + eval_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(16)]) + trainer = XGBoostTrainer( + train_fn_per_worker, + datasets={"train": train_ds, "validation": eval_ds}, + scaling_config=ray.train.ScalingConfig(num_workers=4), + ) + result = trainer.fit() + booster = RayTrainReportCallback.get_model(result.checkpoint) + + Args: + train_loop_per_worker: The training function to execute on each worker. + This function can either take in zero arguments or a single ``Dict`` + argument which is set by defining ``train_loop_config``. + Within this function you can use any of the + :ref:`Ray Train Loop utilities `. + train_loop_config: A configuration ``Dict`` to pass in as an argument to + ``train_loop_per_worker``. + This is typically used for specifying hyperparameters. + xgboost_config: The configuration for setting up the distributed xgboost + backend. Defaults to using the "rabit" backend. + See :class:`~ray.train.xgboost.XGBoostConfig` for more info. + datasets: The Ray Datasets to use for training and validation. + dataset_config: The configuration for ingesting the input ``datasets``. + By default, all the Ray Datasets are split equally across workers. + See :class:`~ray.train.DataConfig` for more details. + scaling_config: The configuration for how to scale data parallel training. + ``num_workers`` determines how many Python processes are used for training, + and ``use_gpu`` determines whether or not each process should use GPUs. + See :class:`~ray.train.ScalingConfig` for more info. + run_config: The configuration for the execution of the training run. + See :class:`~ray.train.RunConfig` for more info. + resume_from_checkpoint: A checkpoint to resume training from. + This checkpoint can be accessed from within ``train_loop_per_worker`` + by calling ``ray.train.get_checkpoint()``. + metadata: Dict that should be made available via + `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + """ + + def __init__( + self, + train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]], + *, + train_loop_config: Optional[Dict] = None, + xgboost_config: Optional[XGBoostConfig] = None, + scaling_config: Optional[ray.train.ScalingConfig] = None, + run_config: Optional[ray.train.RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + dataset_config: Optional[ray.train.DataConfig] = None, + metadata: Optional[Dict[str, Any]] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + ): + super(XGBoostTrainer, self).__init__( + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + backend_config=xgboost_config or XGBoostConfig(), + scaling_config=scaling_config, + dataset_config=dataset_config, + run_config=run_config, + datasets=datasets, + resume_from_checkpoint=resume_from_checkpoint, + metadata=metadata, + ) diff --git a/lib/python3.12/site-packages/ray/train/xgboost/xgboost_checkpoint.py b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..fc8b70465b5c01f85d8f16f841eda2c0969748b8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_checkpoint.py @@ -0,0 +1,75 @@ +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +import xgboost + +from ray.train._internal.framework_checkpoint import FrameworkCheckpoint +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.data.preprocessor import Preprocessor + + +@PublicAPI(stability="beta") +class XGBoostCheckpoint(FrameworkCheckpoint): + """A :py:class:`~ray.train.Checkpoint` with XGBoost-specific functionality.""" + + MODEL_FILENAME = "model.json" + + @classmethod + def from_model( + cls, + booster: xgboost.Booster, + *, + preprocessor: Optional["Preprocessor"] = None, + path: Optional[str] = None, + ) -> "XGBoostCheckpoint": + """Create a :py:class:`~ray.train.Checkpoint` that stores an XGBoost + model. + + Args: + booster: The XGBoost model to store in the checkpoint. + preprocessor: A fitted preprocessor to be applied before inference. + path: The path to the directory where the checkpoint file will be saved. + This should start as an empty directory, since the *entire* + directory will be treated as the checkpoint when reported. + By default, a temporary directory will be created. + + Returns: + An :py:class:`XGBoostCheckpoint` containing the specified ``Estimator``. + + Examples: + + ... testcode:: + + import numpy as np + import ray + from ray.train.xgboost import XGBoostCheckpoint + import xgboost + + train_X = np.array([[1, 2], [3, 4]]) + train_y = np.array([0, 1]) + + model = xgboost.XGBClassifier().fit(train_X, train_y) + checkpoint = XGBoostCheckpoint.from_model(model.get_booster()) + + """ + checkpoint_path = Path(path or tempfile.mkdtemp()) + + if not checkpoint_path.is_dir(): + raise ValueError(f"`path` must be a directory, but got: {checkpoint_path}") + + booster.save_model(checkpoint_path.joinpath(cls.MODEL_FILENAME).as_posix()) + + checkpoint = cls.from_directory(checkpoint_path.as_posix()) + if preprocessor: + checkpoint.set_preprocessor(preprocessor) + return checkpoint + + def get_model(self) -> xgboost.Booster: + """Retrieve the XGBoost model stored in this checkpoint.""" + with self.as_directory() as checkpoint_path: + booster = xgboost.Booster() + booster.load_model(Path(checkpoint_path, self.MODEL_FILENAME).as_posix()) + return booster diff --git a/lib/python3.12/site-packages/ray/train/xgboost/xgboost_predictor.py b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..1b319b93b299bc02a5b83a2f1cdcfa1e8fab6e8e --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_predictor.py @@ -0,0 +1,160 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import pandas as pd +import xgboost + +from ray.air.constants import TENSOR_COLUMN_NAME +from ray.air.data_batch_type import DataBatchType +from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed +from ray.train.predictor import Predictor +from ray.train.xgboost import XGBoostCheckpoint +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.data.preprocessor import Preprocessor + + +@PublicAPI(stability="beta") +class XGBoostPredictor(Predictor): + """A predictor for XGBoost models. + + Args: + model: The XGBoost booster to use for predictions. + preprocessor: A preprocessor used to transform data batches prior + to prediction. + """ + + def __init__( + self, model: xgboost.Booster, preprocessor: Optional["Preprocessor"] = None + ): + self.model = model + super().__init__(preprocessor) + + def __repr__(self): + return ( + f"{self.__class__.__name__}(model={self.model!r}, " + f"preprocessor={self._preprocessor!r})" + ) + + @classmethod + def from_checkpoint(cls, checkpoint: XGBoostCheckpoint) -> "XGBoostPredictor": + """Instantiate the predictor from a Checkpoint. + + This is a helper constructor that instantiates the predictor from a + framework-specific XGBoost checkpoint. + + Args: + checkpoint: The checkpoint to load the model and preprocessor from. + + """ + model = checkpoint.get_model() + preprocessor = checkpoint.get_preprocessor() + return cls(model=model, preprocessor=preprocessor) + + def predict( + self, + data: DataBatchType, + feature_columns: Optional[Union[List[str], List[int]]] = None, + dmatrix_kwargs: Optional[Dict[str, Any]] = None, + **predict_kwargs, + ) -> DataBatchType: + """Run inference on data batch. + + The data is converted into an XGBoost DMatrix before being inputted to + the model. + + Args: + data: A batch of input data. + feature_columns: The names or indices of the columns in the + data to use as features to predict on. If None, then use + all columns in ``data``. + dmatrix_kwargs: Dict of keyword arguments passed to ``xgboost.DMatrix``. + **predict_kwargs: Keyword arguments passed to ``xgboost.Booster.predict``. + + + Examples: + + .. testcode:: + + import numpy as np + import xgboost as xgb + from ray.train.xgboost import XGBoostPredictor + train_X = np.array([[1, 2], [3, 4]]) + train_y = np.array([0, 1]) + model = xgb.XGBClassifier().fit(train_X, train_y) + predictor = XGBoostPredictor(model=model.get_booster()) + data = np.array([[1, 2], [3, 4]]) + predictions = predictor.predict(data) + # Only use first and second column as the feature + data = np.array([[1, 2, 8], [3, 4, 9]]) + predictions = predictor.predict(data, feature_columns=[0, 1]) + + .. testcode:: + + import pandas as pd + import xgboost as xgb + from ray.train.xgboost import XGBoostPredictor + train_X = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"]) + train_y = pd.Series([0, 1]) + model = xgb.XGBClassifier().fit(train_X, train_y) + predictor = XGBoostPredictor(model=model.get_booster()) + # Pandas dataframe. + data = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"]) + predictions = predictor.predict(data) + # Only use first and second column as the feature + data = pd.DataFrame([[1, 2, 8], [3, 4, 9]], columns=["A", "B", "C"]) + predictions = predictor.predict(data, feature_columns=["A", "B"]) + + + Returns: + Prediction result. + + """ + return Predictor.predict( + self, + data, + feature_columns=feature_columns, + dmatrix_kwargs=dmatrix_kwargs, + **predict_kwargs, + ) + + def _predict_pandas( + self, + data: "pd.DataFrame", + feature_columns: Optional[Union[List[str], List[int]]] = None, + dmatrix_kwargs: Optional[Dict[str, Any]] = None, + **predict_kwargs, + ) -> "pd.DataFrame": + dmatrix_kwargs = dmatrix_kwargs or {} + + feature_names = None + if TENSOR_COLUMN_NAME in data: + data = data[TENSOR_COLUMN_NAME].to_numpy() + data = _unwrap_ndarray_object_type_if_needed(data) + if feature_columns: + # In this case feature_columns is a list of integers + data = data[:, feature_columns] + elif feature_columns: + # feature_columns is a list of integers or strings + data = data[feature_columns].to_numpy() + # Only set the feature names if they are strings + if all(isinstance(fc, str) for fc in feature_columns): + feature_names = feature_columns + else: + feature_columns = data.columns.tolist() + data = data.to_numpy() + + if all(isinstance(fc, str) for fc in feature_columns): + feature_names = feature_columns + + if feature_names: + dmatrix_kwargs["feature_names"] = feature_names + + matrix = xgboost.DMatrix(data, **dmatrix_kwargs) + df = pd.DataFrame(self.model.predict(matrix, **predict_kwargs)) + df.columns = ( + ["predictions"] + if len(df.columns) == 1 + else [f"predictions_{i}" for i in range(len(df.columns))] + ) + return df diff --git a/lib/python3.12/site-packages/ray/train/xgboost/xgboost_trainer.py b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..85dc433584493e0c2f6f6c20ba016d6f3c2fc3c7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/train/xgboost/xgboost_trainer.py @@ -0,0 +1,313 @@ +import logging +from functools import partial +from typing import Any, Callable, Dict, Optional, Union + +import xgboost +from packaging.version import Version + +import ray.train +from ray.train import Checkpoint +from ray.train.constants import TRAIN_DATASET_KEY +from ray.train.trainer import GenDataset +from ray.train.utils import _log_deprecation_warning +from ray.train.xgboost import RayTrainReportCallback, XGBoostConfig +from ray.train.xgboost.v2 import XGBoostTrainer as SimpleXGBoostTrainer +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE = ( + "Passing in `xgboost.train` kwargs such as `params`, `num_boost_round`, " + "`label_column`, etc. to `XGBoostTrainer` is deprecated " + "in favor of the new API which accepts a training function, " + "similar to the other DataParallelTrainer APIs (ex: TorchTrainer). " + "See this issue for more context: " + "https://github.com/ray-project/ray/issues/50042" +) + + +def _xgboost_train_fn_per_worker( + config: dict, + label_column: str, + num_boost_round: int, + dataset_keys: set, + xgboost_train_kwargs: dict, +): + checkpoint = ray.train.get_checkpoint() + starting_model = None + remaining_iters = num_boost_round + if checkpoint: + starting_model = RayTrainReportCallback.get_model(checkpoint) + starting_iter = starting_model.num_boosted_rounds() + remaining_iters = num_boost_round - starting_iter + logger.info( + f"Model loaded from checkpoint will train for " + f"additional {remaining_iters} iterations (trees) in order " + "to achieve the target number of iterations " + f"({num_boost_round=})." + ) + + train_ds_iter = ray.train.get_dataset_shard(TRAIN_DATASET_KEY) + train_df = train_ds_iter.materialize().to_pandas() + + eval_ds_iters = { + k: ray.train.get_dataset_shard(k) + for k in dataset_keys + if k != TRAIN_DATASET_KEY + } + eval_dfs = {k: d.materialize().to_pandas() for k, d in eval_ds_iters.items()} + + train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column] + dtrain = xgboost.DMatrix(train_X, label=train_y) + + # NOTE: Include the training dataset in the evaluation datasets. + # This allows `train-*` metrics to be calculated and reported. + evals = [(dtrain, TRAIN_DATASET_KEY)] + + for eval_name, eval_df in eval_dfs.items(): + eval_X, eval_y = eval_df.drop(label_column, axis=1), eval_df[label_column] + evals.append((xgboost.DMatrix(eval_X, label=eval_y), eval_name)) + + evals_result = {} + xgboost.train( + config, + dtrain=dtrain, + evals=evals, + evals_result=evals_result, + num_boost_round=remaining_iters, + xgb_model=starting_model, + **xgboost_train_kwargs, + ) + + +@PublicAPI(stability="beta") +class XGBoostTrainer(SimpleXGBoostTrainer): + """A Trainer for distributed data-parallel XGBoost training. + + Example + ------- + + .. testcode:: + :skipif: True + + import xgboost + + import ray.data + import ray.train + from ray.train.xgboost import RayTrainReportCallback, XGBoostTrainer + + def train_fn_per_worker(config: dict): + # (Optional) Add logic to resume training state from a checkpoint. + # ray.train.get_checkpoint() + + # 1. Get the dataset shard for the worker and convert to a `xgboost.DMatrix` + train_ds_iter, eval_ds_iter = ( + ray.train.get_dataset_shard("train"), + ray.train.get_dataset_shard("validation"), + ) + train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize() + + train_df, eval_df = train_ds.to_pandas(), eval_ds.to_pandas() + train_X, train_y = train_df.drop("y", axis=1), train_df["y"] + eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"] + + dtrain = xgboost.DMatrix(train_X, label=train_y) + deval = xgboost.DMatrix(eval_X, label=eval_y) + + params = { + "tree_method": "approx", + "objective": "reg:squarederror", + "eta": 1e-4, + "subsample": 0.5, + "max_depth": 2, + } + + # 2. Do distributed data-parallel training. + # Ray Train sets up the necessary coordinator processes and + # environment variables for your workers to communicate with each other. + bst = xgboost.train( + params, + dtrain=dtrain, + evals=[(deval, "validation")], + num_boost_round=10, + callbacks=[RayTrainReportCallback()], + ) + + train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)]) + eval_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(16)]) + trainer = XGBoostTrainer( + train_fn_per_worker, + datasets={"train": train_ds, "validation": eval_ds}, + scaling_config=ray.train.ScalingConfig(num_workers=4), + ) + result = trainer.fit() + booster = RayTrainReportCallback.get_model(result.checkpoint) + + Args: + train_loop_per_worker: The training function to execute on each worker. + This function can either take in zero arguments or a single ``Dict`` + argument which is set by defining ``train_loop_config``. + Within this function you can use any of the + :ref:`Ray Train Loop utilities `. + train_loop_config: A configuration ``Dict`` to pass in as an argument to + ``train_loop_per_worker``. + This is typically used for specifying hyperparameters. + xgboost_config: The configuration for setting up the distributed xgboost + backend. Defaults to using the "rabit" backend. + See :class:`~ray.train.xgboost.XGBoostConfig` for more info. + datasets: The Ray Datasets to use for training and validation. + dataset_config: The configuration for ingesting the input ``datasets``. + By default, all the Ray Datasets are split equally across workers. + See :class:`~ray.train.DataConfig` for more details. + scaling_config: The configuration for how to scale data parallel training. + ``num_workers`` determines how many Python processes are used for training, + and ``use_gpu`` determines whether or not each process should use GPUs. + See :class:`~ray.train.ScalingConfig` for more info. + run_config: The configuration for the execution of the training run. + See :class:`~ray.train.RunConfig` for more info. + resume_from_checkpoint: A checkpoint to resume training from. + This checkpoint can be accessed from within ``train_loop_per_worker`` + by calling ``ray.train.get_checkpoint()``. + metadata: Dict that should be made available via + `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()` + for checkpoints saved from this Trainer. Must be JSON-serializable. + label_column: [Deprecated] Name of the label column. A column with this name + must be present in the training dataset. + params: [Deprecated] XGBoost training parameters. + Refer to `XGBoost documentation `_ + for a list of possible parameters. + num_boost_round: [Deprecated] Target number of boosting iterations (trees in the model). + Note that unlike in ``xgboost.train``, this is the target number + of trees, meaning that if you set ``num_boost_round=10`` and pass a model + that has already been trained for 5 iterations, it will be trained for 5 + iterations more, instead of 10 more. + **train_kwargs: [Deprecated] Additional kwargs passed to ``xgboost.train()`` function. + """ + + _handles_checkpoint_freq = True + _handles_checkpoint_at_end = True + + def __init__( + self, + train_loop_per_worker: Optional[ + Union[Callable[[], None], Callable[[Dict], None]] + ] = None, + *, + train_loop_config: Optional[Dict] = None, + xgboost_config: Optional[XGBoostConfig] = None, + scaling_config: Optional[ray.train.ScalingConfig] = None, + run_config: Optional[ray.train.RunConfig] = None, + datasets: Optional[Dict[str, GenDataset]] = None, + dataset_config: Optional[ray.train.DataConfig] = None, + resume_from_checkpoint: Optional[Checkpoint] = None, + metadata: Optional[Dict[str, Any]] = None, + # TODO(justinvyu): [Deprecated] Legacy XGBoostTrainer API + label_column: Optional[str] = None, + params: Optional[Dict[str, Any]] = None, + num_boost_round: Optional[int] = None, + **train_kwargs, + ): + if Version(xgboost.__version__) < Version("1.7.0"): + raise ImportError( + "`XGBoostTrainer` requires the `xgboost` version to be >= 1.7.0. " + 'Upgrade with: `pip install -U "xgboost>=1.7"`' + ) + + # TODO(justinvyu): [Deprecated] Legacy XGBoostTrainer API + legacy_api = train_loop_per_worker is None + if legacy_api: + train_loop_per_worker = self._get_legacy_train_fn_per_worker( + xgboost_train_kwargs=train_kwargs, + run_config=run_config, + label_column=label_column, + num_boost_round=num_boost_round, + datasets=datasets, + ) + train_loop_config = params or {} + elif train_kwargs: + _log_deprecation_warning( + "Passing `xgboost.train` kwargs to `XGBoostTrainer` is deprecated. " + "In your training function, you can call `xgboost.train(**kwargs)` " + "with arbitrary arguments. " + f"{LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE}" + ) + + super(XGBoostTrainer, self).__init__( + train_loop_per_worker=train_loop_per_worker, + train_loop_config=train_loop_config, + xgboost_config=xgboost_config, + scaling_config=scaling_config, + run_config=run_config, + datasets=datasets, + dataset_config=dataset_config, + resume_from_checkpoint=resume_from_checkpoint, + metadata=metadata, + ) + + def _get_legacy_train_fn_per_worker( + self, + xgboost_train_kwargs: Dict, + run_config: Optional[ray.train.RunConfig], + datasets: Optional[Dict[str, GenDataset]], + label_column: Optional[str], + num_boost_round: Optional[int], + ) -> Callable[[Dict], None]: + """Get the training function for the legacy XGBoostTrainer API.""" + + datasets = datasets or {} + if not datasets.get(TRAIN_DATASET_KEY): + raise ValueError( + "`datasets` must be provided for the XGBoostTrainer API " + "if `train_loop_per_worker` is not provided. " + "This dict must contain the training dataset under the " + f"key: '{TRAIN_DATASET_KEY}'. " + f"Got keys: {list(datasets.keys())}" + ) + if not label_column: + raise ValueError( + "`label_column` must be provided for the XGBoostTrainer API " + "if `train_loop_per_worker` is not provided. " + "This is the column name of the label in the dataset." + ) + + num_boost_round = num_boost_round or 10 + + _log_deprecation_warning(LEGACY_XGBOOST_TRAINER_DEPRECATION_MESSAGE) + + # Initialize a default Ray Train metrics/checkpoint reporting callback if needed + callbacks = xgboost_train_kwargs.get("callbacks", []) + user_supplied_callback = any( + isinstance(callback, RayTrainReportCallback) for callback in callbacks + ) + callback_kwargs = {} + if run_config: + checkpoint_frequency = run_config.checkpoint_config.checkpoint_frequency + checkpoint_at_end = run_config.checkpoint_config.checkpoint_at_end + + callback_kwargs["frequency"] = checkpoint_frequency + # Default `checkpoint_at_end=True` unless the user explicitly sets it. + callback_kwargs["checkpoint_at_end"] = ( + checkpoint_at_end if checkpoint_at_end is not None else True + ) + + if not user_supplied_callback: + callbacks.append(RayTrainReportCallback(**callback_kwargs)) + xgboost_train_kwargs["callbacks"] = callbacks + + train_fn_per_worker = partial( + _xgboost_train_fn_per_worker, + label_column=label_column, + num_boost_round=num_boost_round, + dataset_keys=set(datasets), + xgboost_train_kwargs=xgboost_train_kwargs, + ) + return train_fn_per_worker + + @classmethod + def get_model( + cls, + checkpoint: Checkpoint, + ) -> xgboost.Booster: + """Retrieve the XGBoost model stored in this checkpoint.""" + return RayTrainReportCallback.get_model(checkpoint)