wehub-resource-sync
166 lines
5.0 KiB
Python
166 lines
5.0 KiB
Python
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] = None
|
|
pageins: Optional[int] = None
|
|
|
|
|
|
@DeveloperAPI
|
|
class ProcessStats(BaseModel):
|
|
cpuPercent: float
|
|
# total memory, free memory, memory used ratio
|
|
mem: Optional[List[int]] = None
|
|
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] = None
|
|
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(None, 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(
|
|
None,
|
|
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]
|