chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,559 @@
|
||||
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
|
||||
Reference in New Issue
Block a user