wehub-resource-sync
98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
import time
|
|
|
|
from ray.data._internal.execution.interfaces.execution_options import ExecutionOptions
|
|
from ray.train._internal.data_config import DataConfig
|
|
from ray.train.v2._internal.state.schema import (
|
|
ActorStatus,
|
|
DataConfig as DataConfigSchema,
|
|
DataExecutionOptions,
|
|
ExecutionOptions as ExecutionOptionsSchema,
|
|
RunAttemptStatus,
|
|
RunStatus,
|
|
TrainRun,
|
|
TrainRunAttempt,
|
|
)
|
|
from ray.util.state import get_actor
|
|
|
|
_GRACEFUL_ABORT_STATUS_DETAIL = "Run aborted due to user interrupt (SIGINT)."
|
|
_DEAD_CONTROLLER_ABORT_STATUS_DETAIL = (
|
|
"Run aborted because the driver process exited unexpectedly."
|
|
)
|
|
|
|
|
|
def update_train_run_aborted(run: TrainRun, graceful: bool) -> None:
|
|
run.status = RunStatus.ABORTED
|
|
if graceful:
|
|
run.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
|
|
else:
|
|
run.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
|
|
run.end_time_ns = current_time_ns()
|
|
|
|
|
|
def update_train_run_attempt_aborted(
|
|
run_attempt: TrainRunAttempt, graceful: bool
|
|
) -> None:
|
|
if graceful:
|
|
run_attempt.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
|
|
else:
|
|
run_attempt.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
|
|
run_attempt.status = RunAttemptStatus.ABORTED
|
|
run_attempt.end_time_ns = current_time_ns()
|
|
mark_workers_dead(run_attempt)
|
|
|
|
|
|
def mark_workers_dead(run_attempt: TrainRunAttempt) -> None:
|
|
for worker in run_attempt.workers:
|
|
worker.status = ActorStatus.DEAD
|
|
|
|
|
|
def current_time_ns() -> int:
|
|
return time.time_ns()
|
|
|
|
|
|
def is_actor_alive(actor_id: str, timeout: int) -> bool:
|
|
"""Returns whether actor is alive."""
|
|
actor_state = get_actor(actor_id, timeout=timeout)
|
|
return actor_state and actor_state.state != "DEAD"
|
|
|
|
|
|
def construct_data_config(data_config: DataConfig) -> DataConfigSchema:
|
|
"""Serialize a user-facing DataConfig into the exportable schema.
|
|
|
|
Note: This function assumes data_config._execution_options (a defaultdict)
|
|
hasn't been read between initialization of the field and this function call.
|
|
Any read materializes a dataset key and affects the data config shape,
|
|
wrongly capturing a per dataset execution options even if the user only
|
|
provided a default.
|
|
"""
|
|
exec_options = data_config._execution_options
|
|
|
|
per_dataset_execution_options = {}
|
|
if exec_options:
|
|
per_dataset_execution_options = {
|
|
ds_name: execution_options_to_model(opts)
|
|
for ds_name, opts in exec_options.items()
|
|
}
|
|
|
|
return DataConfigSchema(
|
|
datasets_to_split=data_config._datasets_to_split,
|
|
data_execution_options=DataExecutionOptions(
|
|
default=execution_options_to_model(exec_options.default_factory()),
|
|
per_dataset_execution_options=per_dataset_execution_options,
|
|
),
|
|
enable_shard_locality=data_config._enable_shard_locality,
|
|
)
|
|
|
|
|
|
def execution_options_to_model(
|
|
execution_options: ExecutionOptions,
|
|
) -> ExecutionOptionsSchema:
|
|
"""Convert a ray.data ExecutionOptions object into the export schema model."""
|
|
return ExecutionOptionsSchema(
|
|
resource_limits=execution_options.resource_limits.to_resource_dict(),
|
|
exclude_resources=execution_options.exclude_resources.to_resource_dict(),
|
|
preserve_order=execution_options.preserve_order,
|
|
actor_locality_enabled=execution_options.actor_locality_enabled,
|
|
verbose_progress=execution_options.verbose_progress,
|
|
)
|