chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,52 @@
|
||||
import dataclasses
|
||||
from typing import Iterable
|
||||
|
||||
|
||||
def _validate_allowed_keys_exist(
|
||||
dataclass_name: str, data_dict: dict, allowed_keys: set
|
||||
):
|
||||
keys_not_in_dict = allowed_keys.difference(data_dict)
|
||||
if keys_not_in_dict:
|
||||
raise ValueError(
|
||||
f"Key(s) {sorted(keys_not_in_dict)} are not present in {dataclass_name}. "
|
||||
"Remove them from `allowed_keys`. "
|
||||
f"Valid keys: {sorted(data_dict.keys())}"
|
||||
)
|
||||
|
||||
|
||||
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__()
|
||||
default_data_dict = default_data.__dict__
|
||||
|
||||
allowed_keys = set(allowed_keys)
|
||||
_validate_allowed_keys_exist(
|
||||
dataclass.__class__.__name__, default_data_dict, allowed_keys
|
||||
)
|
||||
|
||||
# These keys should not have been updated in the `dataclass` object
|
||||
prohibited_keys = set(default_data_dict) - allowed_keys
|
||||
dataclass_dict = dataclass.__dict__
|
||||
|
||||
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."
|
||||
)
|
||||
@@ -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,
|
||||
]
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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"""
|
||||
...
|
||||
@@ -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_default_system_temp_dir()).resolve()
|
||||
self._lock_dir = temp_dir / RAY_LOCKFILE_DIR
|
||||
self._path_hash = hashlib.sha256(
|
||||
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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,346 @@
|
||||
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,
|
||||
) -> None:
|
||||
"""
|
||||
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.
|
||||
tracking_token: Tracking token used to authenticate with MLflow.
|
||||
artifact_location: The location to store run artifacts.
|
||||
If not provided, MLFlow picks an appropriate default.
|
||||
Ignored if experiment already exists.
|
||||
create_experiment_if_not_exists: Whether to create an
|
||||
experiment with the provided name if it does not already
|
||||
exist. Defaults to True.
|
||||
|
||||
Raises:
|
||||
ValueError: ``experiment_id`` and ``experiment_name`` are both ``None``.
|
||||
"""
|
||||
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:
|
||||
run_name: Name of the new MLflow run to create.
|
||||
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: 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: int, 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:
|
||||
step: Step at which the metrics are logged.
|
||||
metrics_to_log: Dictionary of metrics to log.
|
||||
run_id: 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: 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: Optional[str] = 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: The status to set when terminating the run.
|
||||
run_id: 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)
|
||||
@@ -0,0 +1,73 @@
|
||||
from typing import Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
|
||||
|
||||
|
||||
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:
|
||||
ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a TensorFlow
|
||||
Tensor.
|
||||
dtypes: 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
|
||||
@@ -0,0 +1,105 @@
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from ray.air._internal.device_manager import get_torch_device_manager_by_context
|
||||
|
||||
|
||||
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 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.
|
||||
|
||||
Returns:
|
||||
A new state-dict with the prefix stripped from the keys.
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,106 @@
|
||||
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<s3://bucket/a?scheme=http¶m=1>
|
||||
>>> 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'
|
||||
|
||||
Args:
|
||||
subpath: The subpath to strip from the end of this URI.
|
||||
|
||||
Returns:
|
||||
A new URI with the subpath stripped from the end.
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,277 @@
|
||||
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: object, 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.
|
||||
|
||||
Args:
|
||||
callbacks: List of callbacks supplied by the user. May be ``None``.
|
||||
|
||||
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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user