wehub-resource-sync
533 lines
18 KiB
Python
533 lines
18 KiB
Python
import math
|
|
from collections.abc import Mapping
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
|
|
|
|
from pydantic import field_validator
|
|
|
|
from ray._common.pydantic_compat import BaseModel, Field
|
|
from ray.dashboard.modules.job.pydantic_models import JobDetails
|
|
from ray.train.v2._internal.util import TrainingFramework
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
MAX_ERROR_STACK_TRACE_LENGTH = 50000
|
|
|
|
|
|
def _to_json_serializable_value(value: Any, *, max_depth: int = 3) -> Any:
|
|
"""Recursively coerce a value into a human-readable, JSON serializable representation.
|
|
|
|
If ``value`` is a list or dict, this function walks through it and replaces non-JSON
|
|
serializable fields (e.g. custom objects, modules, tensors, callables, etc.) with a
|
|
human-readable string representation.
|
|
|
|
Args:
|
|
value: Any Python value. Primitives pass through; collections recurse;
|
|
other types are stringified.
|
|
max_depth: Truncates dicts nested beyond ``max_depth`` to ``"..."``.
|
|
Lists do not consume depth.
|
|
|
|
Returns:
|
|
The JSON serializable representation of the value.
|
|
"""
|
|
if max_depth <= 0:
|
|
raise ValueError("max_depth must be greater than 0")
|
|
|
|
def _safe_str(v):
|
|
try:
|
|
return str(v)
|
|
except Exception:
|
|
return type(v).__name__
|
|
|
|
def _walk(value, depth):
|
|
if value is None or isinstance(value, (bool, int, str)):
|
|
return value
|
|
if isinstance(value, float):
|
|
return str(value) if not math.isfinite(value) else value
|
|
if isinstance(value, Mapping):
|
|
if depth <= 0:
|
|
return "..."
|
|
try:
|
|
items = list(value.items())
|
|
except Exception:
|
|
# Custom Mapping subclass with a broken `.items()`.
|
|
return type(value).__name__
|
|
return {_safe_str(k): _walk(v, depth - 1) for k, v in items}
|
|
|
|
# Tuples, sets, and frozensets all become lists in JSON.
|
|
if isinstance(value, (list, tuple, set, frozenset)):
|
|
return [_walk(v, depth) for v in value]
|
|
|
|
cls = type(value)
|
|
# Use class name if no custom string representation is defined.
|
|
if cls.__str__ is object.__str__ and cls.__repr__ is object.__repr__:
|
|
return cls.__name__
|
|
return _safe_str(value)
|
|
|
|
return _walk(value, max_depth)
|
|
|
|
|
|
@DeveloperAPI
|
|
class RunStatus(str, Enum):
|
|
"""Enumeration of the possible statuses for a Train run."""
|
|
|
|
# ====== Active States ======
|
|
# The Train run is currently in the process of initializing.
|
|
INITIALIZING = "INITIALIZING"
|
|
# The Train run is waiting to be scheduled.
|
|
SCHEDULING = "SCHEDULING"
|
|
# The Train run is currently in progress.
|
|
RUNNING = "RUNNING"
|
|
# The Train run is recovering from a failure or restart.
|
|
RESTARTING = "RESTARTING"
|
|
# The Train run is resizing.
|
|
RESIZING = "RESIZING"
|
|
|
|
# ===== Terminal States ======
|
|
# The Train run completed successfully.
|
|
FINISHED = "FINISHED"
|
|
# The Train run failed due to an error in the training workers.
|
|
ERRORED = "ERRORED"
|
|
# The Train run was terminated due to system or controller errors.
|
|
ABORTED = "ABORTED"
|
|
|
|
def is_terminal(self) -> bool:
|
|
return self in [RunStatus.FINISHED, RunStatus.ERRORED, RunStatus.ABORTED]
|
|
|
|
|
|
@DeveloperAPI
|
|
class RunAttemptStatus(str, Enum):
|
|
"""Enumeration of the possible statuses for a Train run attempt."""
|
|
|
|
# ====== Active States ======
|
|
# The run attempt is waiting to be scheduled.
|
|
PENDING = "PENDING"
|
|
# The run attempt is currently in progress.
|
|
RUNNING = "RUNNING"
|
|
|
|
# ===== Terminal States =====
|
|
# The run attempt completed successfully.
|
|
FINISHED = "FINISHED"
|
|
# The run attempt failed due to an error in the training workers.
|
|
ERRORED = "ERRORED"
|
|
# The run attempt was terminated due to system or controller errors.
|
|
ABORTED = "ABORTED"
|
|
|
|
def is_terminal(self) -> bool:
|
|
return self in [
|
|
RunAttemptStatus.FINISHED,
|
|
RunAttemptStatus.ERRORED,
|
|
RunAttemptStatus.ABORTED,
|
|
]
|
|
|
|
|
|
@DeveloperAPI
|
|
class ActorStatus(str, Enum):
|
|
"""Enumeration of the statuses for a Train worker actor."""
|
|
|
|
# The actor is currently active.
|
|
ALIVE = "ALIVE"
|
|
# The actor is no longer active.
|
|
DEAD = "DEAD"
|
|
|
|
|
|
@DeveloperAPI
|
|
class TrainResources(BaseModel):
|
|
"""Resources allocated for a Train worker or run."""
|
|
|
|
resources: Dict[str, float] = Field(
|
|
description="A dictionary specifying the types and amounts of resources "
|
|
"allocated (e.g., CPU, GPU)."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class TrainWorker(BaseModel):
|
|
"""Metadata about a Ray Train worker."""
|
|
|
|
world_rank: int = Field(
|
|
description="The global rank of the worker in the training cluster."
|
|
)
|
|
local_rank: int = Field(description="The local rank of the worker on its node.")
|
|
node_rank: int = Field(description="The rank of the worker's node in the cluster.")
|
|
actor_id: str = Field(description="The unique ID of the worker's actor.")
|
|
node_id: str = Field(
|
|
description="The unique ID of the node where the worker is running."
|
|
)
|
|
node_ip: str = Field(
|
|
description="The IP address of the node where the worker is running."
|
|
)
|
|
pid: int = Field(description="The process ID of the worker.")
|
|
gpu_ids: List[int] = Field(description="A list of GPU IDs allocated to the worker.")
|
|
status: Optional[ActorStatus] = Field(
|
|
None, description="The current status of the worker actor."
|
|
)
|
|
resources: TrainResources = Field(
|
|
description="The resources allocated to this Train worker."
|
|
)
|
|
log_file_path: Optional[str] = Field(
|
|
None, description="The path to the log file for the Train worker."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class MemoryInfo(BaseModel):
|
|
"""Memory usage information for a process."""
|
|
|
|
rss: int = Field(description="The resident set size (RSS) memory usage in bytes.")
|
|
vms: int = Field(description="The virtual memory size (VMS) usage in bytes.")
|
|
pfaults: Optional[int] = Field(None, description="The number of page faults.")
|
|
pageins: Optional[int] = Field(None, description="The number of page-ins.")
|
|
|
|
|
|
@DeveloperAPI
|
|
class ProcessStats(BaseModel):
|
|
"""CPU and memory statistics for a process."""
|
|
|
|
cpuPercent: float = Field(description="The percentage of CPU usage.")
|
|
mem: Optional[List[int]] = Field(
|
|
None,
|
|
description="Memory statistics, including total memory, free memory, "
|
|
"and memory usage ratio.",
|
|
)
|
|
memoryInfo: MemoryInfo = Field(description="Detailed memory usage information.")
|
|
|
|
|
|
class ProcessGPUUsage(BaseModel):
|
|
"""GPU usage statistics for a process."""
|
|
|
|
pid: int = Field(description="The process ID.")
|
|
gpuMemoryUsage: int = Field(description="The GPU memory usage in bytes.")
|
|
|
|
|
|
@DeveloperAPI
|
|
class GPUStats(BaseModel):
|
|
"""Statistics for a GPU."""
|
|
|
|
uuid: str = Field(description="The unique identifier of the GPU.")
|
|
index: int = Field(description="The index of the GPU.")
|
|
name: str = Field(description="The name of the GPU.")
|
|
utilizationGpu: Optional[float] = Field(
|
|
None, description="The percentage utilization of the GPU."
|
|
)
|
|
memoryUsed: float = Field(description="The amount of GPU memory used in bytes.")
|
|
memoryTotal: float = Field(description="The total amount of GPU memory in bytes.")
|
|
processInfo: ProcessGPUUsage = Field(
|
|
description="GPU usage statistics for the associated process."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class DecoratedTrainWorker(TrainWorker):
|
|
"""Detailed metadata for a Ray Train worker, including process and GPU stats."""
|
|
|
|
processStats: Optional[ProcessStats] = Field(
|
|
None, description="CPU and memory statistics for the worker process."
|
|
)
|
|
gpus: List[GPUStats] = Field(
|
|
default_factory=list,
|
|
description="A list of GPUs used by the worker process,"
|
|
" with detailed statistics.",
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class TrainRunAttempt(BaseModel):
|
|
"""Metadata for an individual attempt to execute a Train run."""
|
|
|
|
run_id: str = Field(description="Unique identifier for the parent Train run.")
|
|
attempt_id: str = Field(
|
|
description="Unique identifier for this specific Train run attempt."
|
|
)
|
|
status: RunAttemptStatus = Field(
|
|
description="The current execution status of the Train run attempt."
|
|
)
|
|
status_detail: Optional[str] = Field(
|
|
None,
|
|
description="Additional details about the status,"
|
|
" including error messages if applicable.",
|
|
)
|
|
start_time_ns: int = Field(
|
|
description="The UNIX timestamp (in nanoseconds)"
|
|
" when the Train run attempt started."
|
|
)
|
|
end_time_ns: Optional[int] = Field(
|
|
None,
|
|
description="The UNIX timestamp (in nanoseconds)"
|
|
" when the Train run attempt ended. "
|
|
"If null, the attempt is still ongoing.",
|
|
)
|
|
resources: List[TrainResources] = Field(
|
|
description="The resources (e.g., CPU, GPU) allocated to the Train run attempt."
|
|
)
|
|
workers: List[TrainWorker] = Field(
|
|
description="List of Train workers participating in this attempt, "
|
|
"sorted by global ranks."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class DecoratedTrainRunAttempt(TrainRunAttempt):
|
|
"""Detailed metadata for a Train run attempt, including decorated worker data."""
|
|
|
|
workers: List[DecoratedTrainWorker] = Field(
|
|
description="A list of Train workers with detailed statistics, "
|
|
"sorted by global ranks."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class ExecutionOptions(BaseModel):
|
|
"""ExecutionOptions for a single Ray Data ingest pipeline."""
|
|
|
|
resource_limits: Dict[str, Any] = Field(
|
|
description="The resource limits applied to the Ray Data execution plan."
|
|
)
|
|
exclude_resources: Dict[str, Any] = Field(
|
|
description="The resources excluded from the Ray Data execution plan "
|
|
"(e.g. resources reserved by Ray Train workers)."
|
|
)
|
|
|
|
@field_validator("resource_limits", "exclude_resources", mode="before")
|
|
@classmethod
|
|
def _sanitize_dict(cls, v):
|
|
return _to_json_serializable_value(v)
|
|
|
|
preserve_order: bool = Field(
|
|
description="Whether to preserve the order of outputs across operators."
|
|
)
|
|
actor_locality_enabled: bool = Field(
|
|
description="Whether actor-based locality optimizations are enabled."
|
|
)
|
|
verbose_progress: bool = Field(
|
|
description="Whether verbose progress reporting is enabled."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class DataExecutionOptions(BaseModel):
|
|
"""ExecutionOptions for a Ray Train run, split into defaults and per-dataset overrides."""
|
|
|
|
default: ExecutionOptions = Field(
|
|
description="Execution options applied to any dataset without a per-dataset override."
|
|
)
|
|
per_dataset_execution_options: Dict[str, ExecutionOptions] = Field(
|
|
default_factory=dict,
|
|
description="Per-dataset execution option overrides, keyed by dataset name.",
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class DataConfig(BaseModel):
|
|
"""Configuration for dataset splitting and execution options within Ray Train."""
|
|
|
|
datasets_to_split: Union[Literal["all"], List[str]] = Field(
|
|
description="Which datasets to split; either 'all' or a list of dataset names."
|
|
)
|
|
execution_options: Optional[Dict] = Field(
|
|
default=None,
|
|
deprecated="DEPRECATED: Use data_execution_options instead.",
|
|
)
|
|
data_execution_options: DataExecutionOptions = Field(
|
|
description="Data execution options"
|
|
)
|
|
enable_shard_locality: bool = Field(
|
|
description="Whether to enable shard locality optimization."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class ScalingConfig(BaseModel):
|
|
"""Scaling config for a Train run."""
|
|
|
|
num_workers: Union[int, Tuple[int, int]] = Field(
|
|
description="The number of workers for the Train run."
|
|
)
|
|
use_gpu: bool = Field(description="Whether to use GPUs for the Train run.")
|
|
resources_per_worker: Optional[Dict[str, float]] = Field(
|
|
None, description="The resources per worker for a Train run."
|
|
)
|
|
placement_strategy: str = Field(
|
|
description="The placement strategy for the Train run."
|
|
)
|
|
accelerator_type: Optional[str] = Field(
|
|
None, description="The accelerator type for the Train run."
|
|
)
|
|
use_tpu: bool = Field(description="Whether to use TPUs for the Train run.")
|
|
topology: Optional[str] = Field(None, description="The topology for the Train run.")
|
|
bundle_label_selector: Optional[
|
|
Union[Dict[str, str], List[Dict[str, str]]]
|
|
] = Field(None, description="The bundle label selector for the Train run.")
|
|
|
|
|
|
@DeveloperAPI
|
|
class FailureConfig(BaseModel):
|
|
"""Failure config for a Train run."""
|
|
|
|
max_failures: int = Field(
|
|
description="The maximum number of failures for a Train run."
|
|
)
|
|
controller_failure_limit: int = Field(
|
|
description="The maximum number of controller failures to tolerate."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class CheckpointConfig(BaseModel):
|
|
"""Checkpoint config for a Train run."""
|
|
|
|
num_to_keep: Optional[int] = Field(
|
|
None,
|
|
description="The number of most recent checkpoints to keep. Older checkpoints may be deleted.",
|
|
)
|
|
checkpoint_score_attribute: Optional[str] = Field(
|
|
None,
|
|
description="Attribute used to score and rank checkpoints; can be a metric key or attribute.",
|
|
)
|
|
checkpoint_score_order: Literal["max", "min"] = Field(
|
|
description="Order to rank checkpoint scores, 'max' for higher-is-better, 'min' for lower-is-better.",
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class RunConfig(BaseModel):
|
|
"""Run configuration parameters for a Train run, encompassing failure,
|
|
runtime environment, checkpoint settings, and storage path."""
|
|
|
|
name: str = Field(description="The name of the Train run.")
|
|
failure_config: FailureConfig = Field(
|
|
description="The failure config for a Train run."
|
|
)
|
|
worker_runtime_env: Dict[str, Any] = Field(
|
|
description="The worker runtime env for a Train run."
|
|
)
|
|
|
|
@field_validator("worker_runtime_env", mode="before")
|
|
@classmethod
|
|
def _sanitize_worker_runtime_env(cls, v):
|
|
return _to_json_serializable_value(v)
|
|
|
|
checkpoint_config: CheckpointConfig = Field(
|
|
description="The checkpoint config for a Train run."
|
|
)
|
|
storage_path: str = Field(description="The storage path for a Train run.")
|
|
storage_filesystem: Optional[str] = Field(
|
|
None, description="The storage filesystem for a Train run."
|
|
)
|
|
|
|
@field_validator("storage_filesystem", mode="before")
|
|
@classmethod
|
|
def _sanitize_storage_filesystem(cls, v):
|
|
return _to_json_serializable_value(v)
|
|
|
|
|
|
@DeveloperAPI
|
|
class BackendConfig(BaseModel):
|
|
"""Backend config for a Train run."""
|
|
|
|
framework: Optional[TrainingFramework] = Field(
|
|
None, description="The training framework for this backend config."
|
|
)
|
|
config: Dict[str, Any] = Field(
|
|
description="Training framework-specific configuration fields."
|
|
)
|
|
|
|
@field_validator("config", mode="before")
|
|
@classmethod
|
|
def _sanitize_config(cls, v):
|
|
return _to_json_serializable_value(v)
|
|
|
|
|
|
@DeveloperAPI
|
|
class RunSettings(BaseModel):
|
|
"""Settings for a Train run, primarily consisting of configs set before a train run starts.
|
|
|
|
This includes the train loop config, backend config, scaling config, dataset configs,
|
|
and runtime configuration.
|
|
"""
|
|
|
|
train_loop_config: Optional[Dict] = Field(
|
|
None, description="The user defined train loop config for a Train run."
|
|
)
|
|
|
|
@field_validator("train_loop_config", mode="before")
|
|
@classmethod
|
|
def _sanitize_train_loop_config(cls, v):
|
|
return _to_json_serializable_value(v)
|
|
|
|
backend_config: BackendConfig = Field(
|
|
description="The backend config for a Train run. Can vary with the framework (e.g. TorchConfig)"
|
|
)
|
|
scaling_config: ScalingConfig = Field(
|
|
description="The scaling config for this Train run."
|
|
)
|
|
datasets: List[str] = Field(
|
|
description="A list of dataset names for a Train run.",
|
|
)
|
|
data_config: DataConfig = Field(
|
|
description="The data config for a Train run.",
|
|
)
|
|
run_config: RunConfig = Field(
|
|
description="Run configuration for this Train run, including failure, runtime environment, checkpoint settings, and storage path."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class TrainRun(BaseModel):
|
|
"""Metadata for a Ray Train run, including its details and status."""
|
|
|
|
id: str = Field(description="Unique identifier for the Train run.")
|
|
name: str = Field(description="Human-readable name assigned to the Train run.")
|
|
job_id: str = Field(description="The Ray Job ID associated with this Train run.")
|
|
controller_actor_id: str = Field(
|
|
description="Unique ID of the actor managing the Train run."
|
|
)
|
|
status: RunStatus = Field(
|
|
description="The current execution status of the Train run."
|
|
)
|
|
status_detail: Optional[str] = Field(
|
|
None,
|
|
description="Additional details about the current status, "
|
|
"including error messages if applicable.",
|
|
)
|
|
start_time_ns: int = Field(
|
|
description="The UNIX timestamp (in nanoseconds) when the Train run started."
|
|
)
|
|
end_time_ns: Optional[int] = Field(
|
|
None,
|
|
description="The UNIX timestamp (in nanoseconds) when the Train run ended. "
|
|
"If null, the run is still in progress.",
|
|
)
|
|
controller_log_file_path: Optional[str] = Field(
|
|
None, description="The path to the log file for the Train run controller."
|
|
)
|
|
framework_versions: Dict[str, str] = Field(
|
|
description="The relevant framework versions for this Train run,"
|
|
"including the Ray version and training framework version."
|
|
)
|
|
run_settings: RunSettings = Field(
|
|
description="The run settings for this Train run, including train loop config, "
|
|
"backend config, scaling config, dataset details, and runtime configuration."
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class DecoratedTrainRun(TrainRun):
|
|
"""Detailed metadata for a Ray Train run, including attempts and job details."""
|
|
|
|
attempts: List[DecoratedTrainRunAttempt] = Field(
|
|
description="A list of attempts made to execute the Train run."
|
|
)
|
|
job_details: Optional[JobDetails] = Field(
|
|
None,
|
|
description="Detailed information about the job that initiated this Train run.",
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class TrainRunsResponse(BaseModel):
|
|
"""Response containing a list of decorated Train runs."""
|
|
|
|
train_runs: List[DecoratedTrainRun] = Field(
|
|
description="A list of Train runs with detailed metadata."
|
|
)
|