wehub-resource-sync
560 lines
22 KiB
Python
560 lines
22 KiB
Python
import logging
|
|
import sys
|
|
import threading
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from queue import Queue
|
|
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
|
|
|
|
import ray
|
|
from ray._common.retry import retry
|
|
from ray._common.utils import env_float
|
|
from ray.actor import ActorHandle
|
|
from ray.train.v2._internal.constants import (
|
|
AWS_RETRYABLE_TOKENS,
|
|
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
|
|
DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S,
|
|
)
|
|
from ray.train.v2._internal.execution.checkpoint.sync_actor import (
|
|
SynchronizationActor,
|
|
SynchronizationBarrierResetError,
|
|
)
|
|
from ray.train.v2._internal.execution.preemption import PreemptionContext
|
|
from ray.train.v2._internal.execution.storage import StorageContext, delete_fs_path
|
|
from ray.train.v2._internal.execution.training_report import (
|
|
_TrainingReport,
|
|
)
|
|
from ray.train.v2._internal.util import (
|
|
construct_user_exception_with_traceback,
|
|
context_watchdog,
|
|
invoke_context_managers,
|
|
)
|
|
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
|
from ray.train.v2.api.report_config import (
|
|
CheckpointConsistencyMode,
|
|
CheckpointUploadMode,
|
|
)
|
|
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data import DataIterator
|
|
from ray.train import BackendConfig, Checkpoint, DataConfig
|
|
from ray.train.v2._internal.data_integration.interfaces import (
|
|
DatasetShardMetadata,
|
|
DatasetShardProvider,
|
|
)
|
|
from ray.train.v2._internal.execution.callback import TrainContextCallback
|
|
from ray.train.v2._internal.execution.worker_group.thread_runner import ThreadRunner
|
|
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
|
|
|
|
|
logger = logging.getLogger(__file__)
|
|
|
|
|
|
# TODO: make this value manually or automatically configurable.
|
|
MAX_CHECKPOINT_UPLOAD_THREADS = 1
|
|
|
|
DEFAULT_CHECKPOINT_UPLOAD_WARN_MESSAGE = "Checkpoint upload for {checkpoint_dir_name} has been running for {elapsed}s (warning interval: {interval}s). This may indicate a network issue or slow storage backend. Consider specifying a different filesystem via RunConfig(storage_filesystem=...)."
|
|
CUSTOM_CHECKPOINT_UPLOAD_WARN_MESSAGE = "Custom checkpoint upload for {checkpoint_dir_name} has been running for {elapsed}s (warning interval: {interval}s). This may indicate an issue in your custom upload function passed to `ray.train.report(custom_upload_fn)`."
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainRunContext:
|
|
"""Holds the metadata and context for the current training run."""
|
|
|
|
# The unique ID of the training run.
|
|
run_id: str = field(init=False, default_factory=lambda: uuid.uuid4().hex)
|
|
|
|
# The run configuration for the current training run.
|
|
run_config: RunConfig
|
|
|
|
# The configuration passed to the training function.
|
|
train_loop_config: Optional[Dict]
|
|
|
|
# The scaling configuration for the current training run.
|
|
scaling_config: ScalingConfig
|
|
|
|
# The configuration for the training backend (e.g., PyTorch, XGBoost).
|
|
backend_config: "BackendConfig"
|
|
|
|
# The configuration for dataset ingestion and sharding.
|
|
dataset_config: "DataConfig"
|
|
|
|
def get_run_config(self) -> RunConfig:
|
|
"""Returns the run config of the current training run."""
|
|
return self.run_config
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DistributedContext:
|
|
world_rank: int
|
|
world_size: int
|
|
local_rank: int
|
|
local_world_size: int
|
|
node_rank: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExecutionContext:
|
|
"""Holds the execution context for the current worker process.
|
|
|
|
Every worker process has a single execution context accessed via the
|
|
`TrainContext`, which includes the training thread that is actually
|
|
running the user code.
|
|
"""
|
|
|
|
# A shared synchronization actor that helps broadcast data across ranks.
|
|
synchronization_actor: SynchronizationActor
|
|
|
|
# A queue that receives training results from the user training code.
|
|
# `ray.train.report` in user code populates this queue.
|
|
result_queue: Queue
|
|
|
|
# The thread launcher that runs the user training loop.
|
|
training_thread_runner: "ThreadRunner"
|
|
|
|
# The callbacks that are run in the worker train context.
|
|
train_context_callbacks: List["TrainContextCallback"]
|
|
|
|
|
|
@dataclass
|
|
class TrainContext:
|
|
train_run_context: TrainRunContext
|
|
distributed_context: DistributedContext
|
|
execution_context: ExecutionContext
|
|
storage_context: StorageContext
|
|
preemption_context: PreemptionContext
|
|
controller_actor: ActorHandle
|
|
|
|
dataset_shard_provider: "DatasetShardProvider"
|
|
has_validation_fn: Optional[bool] = None
|
|
|
|
# TODO: consolidate into CheckpointContext
|
|
checkpoint: Optional["Checkpoint"] = None
|
|
current_report_index: int = 0
|
|
report_call_index: int = 0
|
|
report_order_condition: threading.Condition = threading.Condition()
|
|
checkpoint_upload_threadpool: ThreadPoolExecutor = ThreadPoolExecutor(
|
|
max_workers=MAX_CHECKPOINT_UPLOAD_THREADS
|
|
)
|
|
|
|
def __post_init__(self):
|
|
# Ray train initializes worker with current report index
|
|
# report_call_index should start at the current report index
|
|
self.report_call_index = self.current_report_index
|
|
|
|
def get_experiment_name(self) -> str:
|
|
return self.train_run_context.run_config.name
|
|
|
|
def get_world_size(self) -> int:
|
|
return self.distributed_context.world_size
|
|
|
|
def get_world_rank(self) -> int:
|
|
return self.distributed_context.world_rank
|
|
|
|
def get_local_rank(self) -> int:
|
|
return self.distributed_context.local_rank
|
|
|
|
def get_local_world_size(self) -> int:
|
|
return self.distributed_context.local_world_size
|
|
|
|
def get_node_rank(self) -> int:
|
|
return self.distributed_context.node_rank
|
|
|
|
def get_storage(self):
|
|
return self.storage_context
|
|
|
|
# TODO: Don't allow these private methods to be called from user code.
|
|
def get_result_queue(self):
|
|
return self.execution_context.result_queue
|
|
|
|
def get_synchronization_actor(self):
|
|
return self.execution_context.synchronization_actor
|
|
|
|
def get_checkpoint(self):
|
|
with self.report_order_condition:
|
|
return self.checkpoint
|
|
|
|
def get_all_reported_checkpoints(
|
|
self,
|
|
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
|
timeout_s: Optional[float] = None,
|
|
) -> List["ReportedCheckpoint"]:
|
|
return ray.get(
|
|
self.controller_actor.get_all_reported_checkpoints.remote(
|
|
self.report_call_index,
|
|
consistency_mode,
|
|
timeout_s,
|
|
)
|
|
)
|
|
|
|
def get_dataset_shard(self, dataset_info: "DatasetShardMetadata") -> "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.
|
|
|
|
Args:
|
|
dataset_info: The shard metadata, including the dataset name and worker rank.
|
|
Returns:
|
|
The ``DataIterator`` shard with the given name for this worker.
|
|
Raises:
|
|
KeyError: If the dataset shard with the given name is not found.
|
|
"""
|
|
return self.dataset_shard_provider.get_dataset_shard(dataset_info)
|
|
|
|
def get_context_callbacks(self) -> List["TrainContextCallback"]:
|
|
return self.execution_context.train_context_callbacks
|
|
|
|
def _sync_checkpoint_dir_name_across_ranks(
|
|
self, checkpoint_dir_name: Optional[str] = None
|
|
) -> str:
|
|
"""Sync the checkpoint dir name across ranks.
|
|
|
|
Args:
|
|
checkpoint_dir_name: The checkpoint dir name to sync.
|
|
|
|
Returns:
|
|
The synced checkpoint dir name.
|
|
"""
|
|
# If checkpoint_dir_name is not set, use default checkpoint_dir_name
|
|
# created by the storage context.
|
|
checkpoint_dir_name = (
|
|
checkpoint_dir_name
|
|
or self.storage_context.make_default_checkpoint_dir_name()
|
|
)
|
|
# Get a consensus across ranks on the remote storage path, so distributed
|
|
# checkpoints will be stored to the same place.
|
|
sync_actor = self.get_synchronization_actor()
|
|
with invoke_context_managers(
|
|
[
|
|
callback.on_checkpoint_sync
|
|
for callback in self.execution_context.train_context_callbacks
|
|
]
|
|
):
|
|
return ray.get(
|
|
sync_actor.broadcast_from_rank_zero.remote(
|
|
world_rank=self.distributed_context.world_rank,
|
|
world_size=self.distributed_context.world_size,
|
|
data=checkpoint_dir_name,
|
|
caller_method_name="ray.train.report",
|
|
)
|
|
)
|
|
|
|
# TODO: make retry configurable
|
|
@retry(description="upload checkpoint", max_attempts=3, match=AWS_RETRYABLE_TOKENS)
|
|
def _upload_checkpoint(
|
|
self,
|
|
checkpoint_dir_name: str,
|
|
metrics: Dict[str, Any],
|
|
checkpoint: Optional["Checkpoint"] = None,
|
|
delete_local_checkpoint_after_upload: bool = False,
|
|
checkpoint_upload_fn: Optional[
|
|
Callable[["Checkpoint", str], "Checkpoint"]
|
|
] = None,
|
|
validation: Union[bool, ValidationTaskConfig] = False,
|
|
) -> _TrainingReport:
|
|
"""Save the checkpoint to remote storage.
|
|
|
|
Args:
|
|
checkpoint_dir_name: The checkpoint dir to persist to.
|
|
metrics: The metrics to report.
|
|
checkpoint: The checkpoint to report.
|
|
delete_local_checkpoint_after_upload: Whether to delete the checkpoint after it is uploaded.
|
|
checkpoint_upload_fn: A user defined function that will be called with the
|
|
checkpoint to upload it. If not provided, defaults to using the `pyarrow.fs.copy_files`
|
|
utility for copying to the destination `storage_path`.
|
|
validation: The validation configuration.
|
|
|
|
Returns:
|
|
The training result object containing the persisted checkpoint.
|
|
"""
|
|
|
|
if not checkpoint:
|
|
return _TrainingReport(checkpoint=None, metrics=metrics, validation=False)
|
|
|
|
def slow_upload_warning(stop_event: threading.Event, message: str):
|
|
# Log a warning for the checkpoint upload every `CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR`
|
|
# seconds until `stop_event` is set.
|
|
elapsed = 0.0
|
|
interval = env_float(
|
|
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
|
|
DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S,
|
|
)
|
|
while not stop_event.wait(interval):
|
|
elapsed += interval
|
|
logger.warning(
|
|
message.format(
|
|
checkpoint_dir_name=checkpoint_dir_name,
|
|
elapsed=elapsed,
|
|
interval=interval,
|
|
)
|
|
)
|
|
|
|
# Records how long the checkpoint transfer took
|
|
warn_message = (
|
|
CUSTOM_CHECKPOINT_UPLOAD_WARN_MESSAGE
|
|
if checkpoint_upload_fn
|
|
else DEFAULT_CHECKPOINT_UPLOAD_WARN_MESSAGE
|
|
)
|
|
with invoke_context_managers(
|
|
[
|
|
callback.on_checkpoint_transfer
|
|
for callback in self.execution_context.train_context_callbacks
|
|
]
|
|
):
|
|
try:
|
|
with context_watchdog(slow_upload_warning, warn_message):
|
|
if checkpoint_upload_fn:
|
|
# Upload the checkpoint using the custom checkpoint_upload_fn
|
|
persisted_checkpoint = checkpoint_upload_fn(
|
|
checkpoint, checkpoint_dir_name
|
|
)
|
|
else:
|
|
# Upload the checkpoint using PyArrow
|
|
persisted_checkpoint = (
|
|
self.storage_context.persist_current_checkpoint(
|
|
checkpoint, checkpoint_dir_name
|
|
)
|
|
)
|
|
except FileNotFoundError:
|
|
logger.exception(
|
|
f"Failed to find local checkpoint ({checkpoint}) when attempting to upload it. "
|
|
"This could be caused by multiple workers on a node attempting to upload the "
|
|
"same directory, and then one of the workers deletes the directory before the "
|
|
"others finish."
|
|
)
|
|
raise
|
|
|
|
# Check that the checkpoint generated is a `ray.train.Checkpoint` instance
|
|
if checkpoint_upload_fn and not isinstance(
|
|
persisted_checkpoint, ray.train.Checkpoint
|
|
):
|
|
raise ValueError(
|
|
f"checkpoint_upload_fn must return a `ray.train.Checkpoint`. Actual type is {type(persisted_checkpoint)}"
|
|
)
|
|
|
|
# TODO: consider deleting local checkpoint as async callback instead
|
|
if delete_local_checkpoint_after_upload:
|
|
try:
|
|
delete_fs_path(checkpoint.filesystem, checkpoint.path)
|
|
except Exception:
|
|
logger.exception(
|
|
f"Failed to delete the local checkpoint after a successful upload: {checkpoint}"
|
|
)
|
|
|
|
return _TrainingReport(
|
|
checkpoint=persisted_checkpoint,
|
|
metrics=metrics,
|
|
validation=validation,
|
|
)
|
|
|
|
def _wait_then_report(
|
|
self, training_report: _TrainingReport, report_call_index: int
|
|
):
|
|
"""Thread waits for its turn before reporting training result to result queue.
|
|
|
|
It does this in order to guarantee the FIFO processing of checkpoints.
|
|
|
|
The queue size is set to 1 to avoid accumulating unprocessed results.
|
|
If the queue is full, the put operation blocks until a result is consumed.
|
|
|
|
TODO: Add a metric to track the blocking time waiting for the
|
|
training result to be consumed by the controller.
|
|
"""
|
|
with self.report_order_condition:
|
|
self.report_order_condition.wait_for(
|
|
lambda: self.current_report_index == report_call_index - 1
|
|
)
|
|
logger.info(
|
|
f"Reporting training result {report_call_index}: {training_report} "
|
|
f"from rank {self.get_world_rank()}"
|
|
)
|
|
# Update latest checkpoint as the persisted checkpoint.
|
|
if training_report.checkpoint:
|
|
self.checkpoint = training_report.checkpoint
|
|
self.get_result_queue().put(training_report)
|
|
self.current_report_index += 1
|
|
self.report_order_condition.notify_all()
|
|
|
|
def report(
|
|
self,
|
|
metrics: Dict[str, Any],
|
|
checkpoint: Optional["Checkpoint"] = None,
|
|
checkpoint_dir_name: Optional[str] = None,
|
|
checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC,
|
|
delete_local_checkpoint_after_upload: Optional[bool] = None,
|
|
checkpoint_upload_fn: Optional[
|
|
Callable[["Checkpoint", str], "Checkpoint"]
|
|
] = None,
|
|
validation: Union[bool, ValidationTaskConfig] = False,
|
|
) -> None:
|
|
"""
|
|
Upload checkpoint to remote storage and put a training
|
|
result on the result queue of this worker process.
|
|
|
|
TODO: the report function should be implemented in the worker instead
|
|
of in the train context. The train context should only keep the train
|
|
related information and not the worker related actions. This refactor
|
|
would also require the `TrainContextCallback` to be updated as well.
|
|
"""
|
|
if "torch" in sys.modules:
|
|
from ray.air._internal.torch_utils import contains_tensor
|
|
|
|
if contains_tensor(metrics):
|
|
raise ValueError(
|
|
"Passing objects containing Torch tensors as metrics "
|
|
"is not supported as it will throw an exception on "
|
|
"deserialization. You can either convert the tensors "
|
|
"to Python objects (ex: `.numpy()`, `.item()`, etc.) "
|
|
"or save tensors as part of the checkpoint files instead."
|
|
)
|
|
|
|
if validation and not self.has_validation_fn:
|
|
raise ValueError(
|
|
"`validation_config` was not set on the trainer, but a validation was requested."
|
|
)
|
|
|
|
if delete_local_checkpoint_after_upload and checkpoint is not None:
|
|
experiment_path = Path(self.storage_context.experiment_fs_path)
|
|
checkpoint_path = Path(checkpoint.path)
|
|
|
|
# Resolve symlinks only for local (absolute) paths.
|
|
# Remote paths (S3, GCS, etc.) are relative after URI and resolve()
|
|
# would prepend CWD, producing a meaningless local path.
|
|
# Mixed absolute/relative paths return False
|
|
if experiment_path.is_absolute():
|
|
experiment_path = experiment_path.resolve()
|
|
if checkpoint_path.is_absolute():
|
|
checkpoint_path = checkpoint_path.resolve()
|
|
|
|
if experiment_path.is_relative_to(checkpoint_path):
|
|
raise ValueError(
|
|
f"Ray Train's experiment directory ({self.storage_context.experiment_fs_path}) "
|
|
f"is contained within the checkpoint path ({checkpoint.path}) "
|
|
f"and `ray.train.report(delete_local_checkpoint_after_upload=True)`. "
|
|
"As a result, this would delete the experiment directory. "
|
|
"Please write the checkpoint to a temporary directory, "
|
|
"a subdirectory of the experiment directory, "
|
|
"or use `delete_local_checkpoint_after_upload=False`."
|
|
)
|
|
|
|
with invoke_context_managers(
|
|
[
|
|
callback.on_report
|
|
for callback in self.execution_context.train_context_callbacks
|
|
]
|
|
):
|
|
self.report_call_index += 1
|
|
report_call_index = self.report_call_index
|
|
|
|
# Sync the checkpoint dir name across ranks.
|
|
try:
|
|
checkpoint_dir_name = self._sync_checkpoint_dir_name_across_ranks(
|
|
checkpoint_dir_name
|
|
)
|
|
except ray.exceptions.RayTaskError as e:
|
|
if not isinstance(e.cause, SynchronizationBarrierResetError):
|
|
raise e
|
|
logger.warning(
|
|
"Synchronization barrier was reset (likely due to a "
|
|
"worker failure). Skipping this report."
|
|
)
|
|
# Keep report indexes aligned across workers.
|
|
self.report_call_index -= 1
|
|
return
|
|
|
|
# Upload checkpoint, wait for turn, and report.
|
|
if checkpoint_upload_mode == CheckpointUploadMode.SYNC:
|
|
training_report = self._upload_checkpoint(
|
|
checkpoint_dir_name,
|
|
metrics,
|
|
checkpoint,
|
|
delete_local_checkpoint_after_upload,
|
|
checkpoint_upload_fn,
|
|
validation,
|
|
)
|
|
self._wait_then_report(training_report, report_call_index)
|
|
|
|
elif checkpoint_upload_mode == CheckpointUploadMode.NO_UPLOAD:
|
|
training_report = _TrainingReport(
|
|
checkpoint=checkpoint,
|
|
metrics=metrics,
|
|
validation=validation,
|
|
)
|
|
self._wait_then_report(training_report, report_call_index)
|
|
|
|
elif checkpoint_upload_mode == CheckpointUploadMode.ASYNC:
|
|
|
|
def _upload_checkpoint_and_report(
|
|
checkpoint_dir_name: str,
|
|
metrics: Dict[str, Any],
|
|
checkpoint: Optional["Checkpoint"],
|
|
report_call_index: int,
|
|
) -> None:
|
|
try:
|
|
training_report = self._upload_checkpoint(
|
|
checkpoint_dir_name,
|
|
metrics,
|
|
checkpoint,
|
|
delete_local_checkpoint_after_upload,
|
|
checkpoint_upload_fn,
|
|
validation,
|
|
)
|
|
self._wait_then_report(training_report, report_call_index)
|
|
except Exception as e:
|
|
# TODO: env var to disable eager raising
|
|
logger.exception(
|
|
"Checkpoint upload failed in the background thread. Raising eagerly "
|
|
"to avoid training in a corrupted state with more potential progress "
|
|
"lost due to checkpointing failures."
|
|
)
|
|
self.execution_context.training_thread_runner.get_exception_queue().put(
|
|
construct_user_exception_with_traceback(e)
|
|
)
|
|
|
|
self.checkpoint_upload_threadpool.submit(
|
|
_upload_checkpoint_and_report,
|
|
checkpoint_dir_name,
|
|
metrics,
|
|
checkpoint,
|
|
report_call_index,
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid checkpoint upload mode: {checkpoint_upload_mode}"
|
|
)
|
|
|
|
|
|
# The global variable holding the current TrainContext
|
|
_train_context: Optional[TrainContext] = None
|
|
|
|
# Thread lock to protect the global TrainContext
|
|
_context_lock = threading.Lock()
|
|
|
|
|
|
def get_train_context() -> TrainContext:
|
|
"""Get the internal train context.
|
|
|
|
Note:
|
|
This should not be used directly by user-facing APIs. User-facing APIs should
|
|
call :class:`~ray.train.v2._internal.execution.train_fn_utils.TrainFnUtils`
|
|
or use :class:`~ray.train.v2.api.context.TrainContext` instead.
|
|
|
|
Returns:
|
|
The internal TrainContext for this worker.
|
|
"""
|
|
with _context_lock:
|
|
if _train_context is None:
|
|
raise RuntimeError("TrainContext has not been initialized.")
|
|
return _train_context
|
|
|
|
|
|
def set_train_context(context) -> None:
|
|
global _train_context
|
|
with _context_lock:
|
|
_train_context = context
|