chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,656 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import ray
|
||||
from ray._common.pydantic_compat import BaseModel
|
||||
from ray._private.ray_constants import env_float
|
||||
from ray.air.config import CheckpointConfig
|
||||
from ray.train._checkpoint import Checkpoint
|
||||
from ray.train._internal.checkpoint_manager import (
|
||||
_CheckpointManager,
|
||||
_insert_into_sorted_list,
|
||||
)
|
||||
from ray.train._internal.session import _TrainingResult
|
||||
from ray.train.v2._internal.constants import (
|
||||
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
|
||||
)
|
||||
from ray.train.v2._internal.exceptions import CheckpointManagerInitializationError
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReportCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import StorageContext
|
||||
from ray.train.v2._internal.execution.storage import _exists_at_fs_path, delete_fs_path
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2._internal.execution.worker_group import Worker
|
||||
from ray.train.v2._internal.util import wait_with_logging
|
||||
from ray.train.v2.api.report_config import CheckpointConsistencyMode
|
||||
from ray.train.v2.api.reported_checkpoint import (
|
||||
ReportedCheckpoint,
|
||||
ReportedCheckpointStatus,
|
||||
)
|
||||
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
GET_ALL_REPORTED_CHECKPOINTS_PERIODIC_WARNING = """
|
||||
`get_all_reported_checkpoints` has been waiting for all checkpoints to get to the {consistency_mode} state for {time_elapsed_s:.2f} s.
|
||||
You can set the {warn_interval_env_var} environment variable to change the frequency of this warning (current value: {warn_interval_s} s).
|
||||
"""
|
||||
|
||||
|
||||
class _TrainingResultState(BaseModel):
|
||||
# Increment version if the schema changes
|
||||
version: int = 0
|
||||
checkpoint_dir_name: str
|
||||
metrics: dict
|
||||
|
||||
|
||||
class _CheckpointManagerState(BaseModel):
|
||||
ray_version: str = ray.__version__
|
||||
checkpoint_results: List[_TrainingResultState]
|
||||
checkpoint_report_indices: List[int]
|
||||
latest_checkpoint_result: Optional[_TrainingResultState] = None
|
||||
pending_training_results: List[_TrainingResultState]
|
||||
pending_validation_specs: List[Union[bool, ValidationTaskConfig]]
|
||||
current_report_index: int
|
||||
|
||||
# List of processed checkpoints based on if successfully validated,
|
||||
# timed out or failed due to an error or canceled for some reason.
|
||||
validated_checkpoint_dir_names: List[str]
|
||||
timed_out_validation_checkpoint_dir_names: List[str]
|
||||
failed_validation_checkpoint_dir_names: List[str]
|
||||
|
||||
|
||||
def _get_training_result_from_state(
|
||||
state: _TrainingResultState,
|
||||
storage_context: StorageContext,
|
||||
) -> _TrainingResult:
|
||||
"""Get a TrainingResult object from a Pydantic state object."""
|
||||
return _TrainingResult(
|
||||
checkpoint=Checkpoint(
|
||||
path=storage_context.build_checkpoint_path_from_name(
|
||||
state.checkpoint_dir_name
|
||||
),
|
||||
filesystem=storage_context.storage_filesystem,
|
||||
),
|
||||
metrics=state.metrics,
|
||||
)
|
||||
|
||||
|
||||
def _get_state_from_training_result(
|
||||
training_result: _TrainingResult,
|
||||
storage_context: StorageContext,
|
||||
) -> _TrainingResultState:
|
||||
"""Get a Pydantic state object from a TrainingResult object."""
|
||||
return _TrainingResultState(
|
||||
checkpoint_dir_name=storage_context.extract_checkpoint_dir_name_from_path(
|
||||
training_result.checkpoint.path
|
||||
),
|
||||
metrics=training_result.metrics,
|
||||
)
|
||||
|
||||
|
||||
class CheckpointManager(_CheckpointManager, ReportCallback, WorkerGroupCallback):
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint_config: CheckpointConfig,
|
||||
storage_context: StorageContext,
|
||||
):
|
||||
self._storage_context = storage_context
|
||||
self._checkpoint_config = checkpoint_config
|
||||
|
||||
# This tracks the number of report calls that have been processed
|
||||
# for the current worker group.
|
||||
self._current_report_index = 0
|
||||
|
||||
# Map from pending checkpoint to validation.
|
||||
self._pending_training_results: Dict[
|
||||
Checkpoint, Tuple[_TrainingResult, Union[bool, ValidationTaskConfig]]
|
||||
] = {}
|
||||
|
||||
# Set of checkpoints that have successfully completed, been timed out
|
||||
# or failed validation.
|
||||
self._validated_checkpoints: set = set()
|
||||
self._timed_out_validation_checkpoints: set = set()
|
||||
self._failed_validation_checkpoints: set = set()
|
||||
|
||||
# Map from checkpoint to report index. Used to order checkpoints.
|
||||
self._checkpoint_to_report_index = {}
|
||||
|
||||
self._condition = asyncio.Condition()
|
||||
|
||||
# Strong references to background tasks created via
|
||||
# ``asyncio.create_task`` to prevent them from being garbage
|
||||
# collected mid-execution. The event loop only keeps weak refs.
|
||||
self._background_tasks: set = set()
|
||||
|
||||
self._collective_warn_interval_s = env_float(
|
||||
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
|
||||
)
|
||||
|
||||
super().__init__(checkpoint_config)
|
||||
# If the snapshot is found, the checkpoint manager will restore its state.
|
||||
# TODO(xgui): CheckpointManager is used to save or restore the checkpoint manager state.
|
||||
# We should sanity check if we should see old state in the storage folder.
|
||||
self._maybe_load_state_from_storage()
|
||||
|
||||
def register_checkpoint(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
):
|
||||
"""Register new checkpoint and add to bookkeeping.
|
||||
|
||||
This method will register a new checkpoint and add it to the internal
|
||||
bookkeeping logic. This means the checkpoint manager will decide if
|
||||
this checkpoint should be kept, and if older or worse performing
|
||||
checkpoints should be deleted.
|
||||
|
||||
Args:
|
||||
training_report: Training report to register.
|
||||
"""
|
||||
checkpoint_result = _TrainingResult(
|
||||
checkpoint=training_report.checkpoint,
|
||||
metrics=training_report.metrics,
|
||||
)
|
||||
self._latest_checkpoint_result = checkpoint_result
|
||||
self._checkpoint_to_report_index[
|
||||
checkpoint_result.checkpoint
|
||||
] = self._current_report_index
|
||||
|
||||
if self._checkpoint_config.checkpoint_score_attribute is not None:
|
||||
# If we're ordering by a score, insert the checkpoint
|
||||
# so that the list remains sorted.
|
||||
_insert_into_sorted_list(
|
||||
self._checkpoint_results,
|
||||
checkpoint_result,
|
||||
key=self._get_checkpoint_score,
|
||||
checkpoint_to_report_index=self._checkpoint_to_report_index,
|
||||
)
|
||||
else:
|
||||
# If no metric is provided, just append (ordering by time of registration).
|
||||
self._checkpoint_results.append(checkpoint_result)
|
||||
|
||||
if training_report.validation:
|
||||
self._pending_training_results[checkpoint_result.checkpoint] = (
|
||||
checkpoint_result,
|
||||
training_report.validation,
|
||||
)
|
||||
|
||||
self._current_report_index += 1
|
||||
|
||||
self._save_state_and_delete_old_checkpoints()
|
||||
|
||||
self._notify()
|
||||
|
||||
def update_checkpoints_with_validation_result(
|
||||
self,
|
||||
checkpoint_updates: Dict[
|
||||
Checkpoint, Tuple[Dict[str, Any], ReportedCheckpointStatus]
|
||||
],
|
||||
):
|
||||
"""Finalize pending validations by recording terminal status and metrics.
|
||||
|
||||
* For VALIDATED checkpoints, metrics are merged into the checkpoint's
|
||||
existing metrics and the checkpoint is re-sorted.
|
||||
* For VALIDATION_TIMEOUT and VALIDATION_FAILED checkpoints, metrics are
|
||||
left untouched and the checkpoint retains its original training-time
|
||||
score position.
|
||||
"""
|
||||
for checkpoint, (metrics, status) in checkpoint_updates.items():
|
||||
if checkpoint not in self._pending_training_results:
|
||||
logger.warning(
|
||||
f"Checkpoint {checkpoint} not found in pending training results. "
|
||||
)
|
||||
continue
|
||||
checkpoint_result, _ = self._pending_training_results[checkpoint]
|
||||
if checkpoint_result not in self._checkpoint_results:
|
||||
raise ValueError(
|
||||
f"Checkpoint {checkpoint} was in pending training results but not "
|
||||
"checkpoint results. "
|
||||
)
|
||||
self._pending_training_results.pop(checkpoint)
|
||||
|
||||
if status == ReportedCheckpointStatus.VALIDATED:
|
||||
# Update the metrics and sort into checkpoint_results
|
||||
checkpoint_result.metrics.update(metrics)
|
||||
self._checkpoint_results.remove(checkpoint_result)
|
||||
_insert_into_sorted_list(
|
||||
self._checkpoint_results,
|
||||
checkpoint_result,
|
||||
key=self._get_checkpoint_score,
|
||||
checkpoint_to_report_index=self._checkpoint_to_report_index,
|
||||
)
|
||||
self._validated_checkpoints.add(checkpoint)
|
||||
elif status == ReportedCheckpointStatus.VALIDATION_TIMEOUT:
|
||||
self._timed_out_validation_checkpoints.add(checkpoint)
|
||||
elif status == ReportedCheckpointStatus.VALIDATION_FAILED:
|
||||
self._failed_validation_checkpoints.add(checkpoint)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unexpected terminal validation status {status} for "
|
||||
f"checkpoint {checkpoint}."
|
||||
)
|
||||
|
||||
self._save_state_and_delete_old_checkpoints()
|
||||
self._notify()
|
||||
|
||||
def get_pending_training_results(
|
||||
self,
|
||||
) -> Dict[Checkpoint, Tuple[_TrainingResult, Union[bool, ValidationTaskConfig]]]:
|
||||
"""Get the pending training results which includes their validation specs."""
|
||||
return self._pending_training_results
|
||||
|
||||
def _notify(self):
|
||||
"""Notify condition so all listeners know state has changed."""
|
||||
|
||||
async def async_notify():
|
||||
async with self._condition:
|
||||
self._condition.notify_all()
|
||||
|
||||
# Keep a strong reference to the task so it isn't garbage
|
||||
# collected before completing, which would silently drop
|
||||
# the notification and could leave listeners waiting forever.
|
||||
task = asyncio.create_task(async_notify())
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
|
||||
def _save_state_and_delete_old_checkpoints(self):
|
||||
"""Delete the old checkpoints."""
|
||||
# Get checkpoints to delete
|
||||
results_to_delete = set()
|
||||
if self._checkpoint_config.num_to_keep is not None:
|
||||
# Delete the bottom (N - K) checkpoints
|
||||
worst_results = set(
|
||||
self._checkpoint_results[: -self._checkpoint_config.num_to_keep]
|
||||
)
|
||||
# Except for the latest checkpoint and pending checkpoints
|
||||
results_to_delete = worst_results - {self._latest_checkpoint_result}
|
||||
results_to_delete = results_to_delete - {
|
||||
v for v, _ in self._pending_training_results.values()
|
||||
}
|
||||
|
||||
# Update internal state before actually deleting them.
|
||||
self._checkpoint_results = [
|
||||
checkpoint_result
|
||||
for checkpoint_result in self._checkpoint_results
|
||||
if checkpoint_result not in results_to_delete
|
||||
]
|
||||
for checkpoint_result in results_to_delete:
|
||||
del self._checkpoint_to_report_index[checkpoint_result.checkpoint]
|
||||
|
||||
# discard doesn't raise an error if the element isn't found
|
||||
self._validated_checkpoints.discard(checkpoint_result.checkpoint)
|
||||
self._timed_out_validation_checkpoints.discard(
|
||||
checkpoint_result.checkpoint
|
||||
)
|
||||
self._failed_validation_checkpoints.discard(
|
||||
checkpoint_result.checkpoint
|
||||
)
|
||||
|
||||
# Save the checkpoint manager state to storage.
|
||||
# Note: We save the state before deleting the old checkpoints.
|
||||
# If deletion happens first and the process crashes, our snapshot
|
||||
# may point to some stale checkpoints that are already deleted.
|
||||
# TODO: Make this writing operation non-blocking.
|
||||
self._write_state_to_storage()
|
||||
|
||||
# Delete the old checkpoints.
|
||||
for checkpoint_result in results_to_delete:
|
||||
checkpoint = checkpoint_result.checkpoint
|
||||
logger.debug("Deleting checkpoint: %s", checkpoint)
|
||||
delete_fs_path(fs=checkpoint.filesystem, fs_path=checkpoint.path)
|
||||
|
||||
# --------------------------
|
||||
# CheckpointManager state
|
||||
# --------------------------
|
||||
|
||||
def _save_state(self) -> str:
|
||||
"""Save the checkpoint manager state to a JSON str."""
|
||||
|
||||
checkpoint_results = [
|
||||
_get_state_from_training_result(checkpoint_result, self._storage_context)
|
||||
for checkpoint_result in self._checkpoint_results
|
||||
]
|
||||
|
||||
checkpoint_report_indices = [
|
||||
self._checkpoint_to_report_index[checkpoint_result.checkpoint]
|
||||
for checkpoint_result in self._checkpoint_results
|
||||
]
|
||||
|
||||
latest_checkpoint_result = (
|
||||
_get_state_from_training_result(
|
||||
self._latest_checkpoint_result, self._storage_context
|
||||
)
|
||||
if self._latest_checkpoint_result is not None
|
||||
else None
|
||||
)
|
||||
|
||||
pending_training_results = [
|
||||
_get_state_from_training_result(v, self._storage_context)
|
||||
for v, _ in self._pending_training_results.values()
|
||||
]
|
||||
pending_validation_specs = [
|
||||
v for _, v in self._pending_training_results.values()
|
||||
]
|
||||
|
||||
validated_ckpt_dir_names = [
|
||||
self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
|
||||
for checkpoint in self._validated_checkpoints
|
||||
]
|
||||
timed_out_validation_ckpt_dir_names = [
|
||||
self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
|
||||
for checkpoint in self._timed_out_validation_checkpoints
|
||||
]
|
||||
failed_validation_ckpt_dir_names = [
|
||||
self._storage_context.extract_checkpoint_dir_name_from_path(checkpoint.path)
|
||||
for checkpoint in self._failed_validation_checkpoints
|
||||
]
|
||||
|
||||
manager_snapshot = _CheckpointManagerState(
|
||||
checkpoint_results=checkpoint_results,
|
||||
checkpoint_report_indices=checkpoint_report_indices,
|
||||
latest_checkpoint_result=latest_checkpoint_result,
|
||||
pending_training_results=pending_training_results,
|
||||
pending_validation_specs=pending_validation_specs,
|
||||
current_report_index=self._current_report_index,
|
||||
validated_checkpoint_dir_names=validated_ckpt_dir_names,
|
||||
timed_out_validation_checkpoint_dir_names=timed_out_validation_ckpt_dir_names,
|
||||
failed_validation_checkpoint_dir_names=failed_validation_ckpt_dir_names,
|
||||
)
|
||||
return manager_snapshot.json()
|
||||
|
||||
def _load_state(self, json_state: str):
|
||||
"""Load the checkpoint manager state from a JSON str."""
|
||||
json_dict = None
|
||||
try:
|
||||
json_dict = json.loads(json_state)
|
||||
manager_snapshot = _CheckpointManagerState.parse_obj(json_dict)
|
||||
except Exception as e:
|
||||
if not json_dict:
|
||||
error = e
|
||||
elif "ray_version" not in json_dict:
|
||||
error = (
|
||||
"You are loading a checkpoint manager snapshot saved with an unknown Ray version "
|
||||
f"but you are running Ray version {ray.__version__}. Please use the same Ray version "
|
||||
"the checkpoint manager snapshot was saved with."
|
||||
)
|
||||
elif json_dict["ray_version"] != ray.__version__:
|
||||
error = (
|
||||
f"You are loading a checkpoint manager snapshot saved with Ray version "
|
||||
f"{json_dict['ray_version']} but you are running Ray version "
|
||||
f"{ray.__version__}. Please use the same Ray version the checkpoint "
|
||||
"manager snapshot was saved with."
|
||||
)
|
||||
else:
|
||||
error = e
|
||||
raise CheckpointManagerInitializationError(error) from e
|
||||
|
||||
# Do this so we are using the same checkpoint and trainingresult objects.
|
||||
# TODO: consider asserting that every checkpoint has a unique dir name
|
||||
checkpoint_dir_name_to_checkpoint_result = {}
|
||||
|
||||
for training_result_state in manager_snapshot.checkpoint_results:
|
||||
training_result = _get_training_result_from_state(
|
||||
training_result_state, self._storage_context
|
||||
)
|
||||
checkpoint_dir_name_to_checkpoint_result[
|
||||
training_result_state.checkpoint_dir_name
|
||||
] = training_result
|
||||
self._checkpoint_results.append(training_result)
|
||||
self._assert_checkpoints_exist()
|
||||
|
||||
assert len(self._checkpoint_results) == len(
|
||||
manager_snapshot.checkpoint_report_indices
|
||||
)
|
||||
self._checkpoint_to_report_index = {
|
||||
checkpoint_result.checkpoint: report_index
|
||||
for checkpoint_result, report_index in zip(
|
||||
self._checkpoint_results, manager_snapshot.checkpoint_report_indices
|
||||
)
|
||||
}
|
||||
|
||||
self._latest_checkpoint_result = (
|
||||
checkpoint_dir_name_to_checkpoint_result[
|
||||
manager_snapshot.latest_checkpoint_result.checkpoint_dir_name
|
||||
]
|
||||
if manager_snapshot.latest_checkpoint_result is not None
|
||||
else None
|
||||
)
|
||||
|
||||
assert len(manager_snapshot.pending_training_results) == len(
|
||||
manager_snapshot.pending_validation_specs
|
||||
)
|
||||
for training_result_state, validation_spec in zip(
|
||||
manager_snapshot.pending_training_results,
|
||||
manager_snapshot.pending_validation_specs,
|
||||
):
|
||||
training_result = checkpoint_dir_name_to_checkpoint_result[
|
||||
training_result_state.checkpoint_dir_name
|
||||
]
|
||||
self._pending_training_results[training_result.checkpoint] = (
|
||||
training_result,
|
||||
validation_spec,
|
||||
)
|
||||
|
||||
# Restore terminal validation statuses. Only checkpoints still in
|
||||
# _checkpoint_results can be looked up; evicted checkpoints are irrelevant.
|
||||
for dir_names, target_set in (
|
||||
(
|
||||
manager_snapshot.validated_checkpoint_dir_names,
|
||||
self._validated_checkpoints,
|
||||
),
|
||||
(
|
||||
manager_snapshot.timed_out_validation_checkpoint_dir_names,
|
||||
self._timed_out_validation_checkpoints,
|
||||
),
|
||||
(
|
||||
manager_snapshot.failed_validation_checkpoint_dir_names,
|
||||
self._failed_validation_checkpoints,
|
||||
),
|
||||
):
|
||||
for dir_name in dir_names:
|
||||
if dir_name in checkpoint_dir_name_to_checkpoint_result:
|
||||
target_set.add(
|
||||
checkpoint_dir_name_to_checkpoint_result[dir_name].checkpoint
|
||||
)
|
||||
|
||||
self._current_report_index = manager_snapshot.current_report_index
|
||||
|
||||
def _maybe_load_state_from_storage(self):
|
||||
"""Load the checkpoint manager state from storage.
|
||||
If no snapshot is found, start with a clean state.
|
||||
"""
|
||||
if not _exists_at_fs_path(
|
||||
fs=self._storage_context.storage_filesystem,
|
||||
fs_path=self._storage_context.checkpoint_manager_snapshot_path,
|
||||
):
|
||||
logger.debug(
|
||||
"No checkpoint manager snapshot found. "
|
||||
"No checkpoint will be available via `ray.train.get_checkpoint`, "
|
||||
"so training will start from scratch."
|
||||
)
|
||||
return
|
||||
with self._storage_context.storage_filesystem.open_input_stream(
|
||||
self._storage_context.checkpoint_manager_snapshot_path
|
||||
) as f:
|
||||
logger.info(
|
||||
"A run snapshot was found in storage folder at: "
|
||||
f"'{self._storage_context.experiment_fs_path}'\n"
|
||||
"This snapshot contains a list of checkpoints reported via "
|
||||
"`ray.train.report` and will be loaded. "
|
||||
"This allows the latest checkpoint found in the snapshot to be "
|
||||
"accessible within your training function via "
|
||||
"`ray.train.get_checkpoint`.\n"
|
||||
"If you meant to start a brand new training job without any "
|
||||
"information about previous checkpoints found in this directory, "
|
||||
"please configure a new, unique `RunConfig(name)` or delete the "
|
||||
f"existing folder at '{self._storage_context.experiment_fs_path}'."
|
||||
)
|
||||
json_state = f.read().decode("utf-8")
|
||||
self._load_state(json_state)
|
||||
|
||||
def _write_state_to_storage(self):
|
||||
"""Write the checkpoint manager state to storage."""
|
||||
checkpoint_manager_snapshot = self._save_state()
|
||||
with self._storage_context.storage_filesystem.open_output_stream(
|
||||
self._storage_context.checkpoint_manager_snapshot_path
|
||||
) as f:
|
||||
f.write(checkpoint_manager_snapshot.encode("utf-8"))
|
||||
|
||||
def _assert_checkpoints_exist(self):
|
||||
"""Validate the checkpoint manager state.
|
||||
|
||||
This method will validate the checkpoint manager state by checking if
|
||||
the checkpoints specified in manager snapshot is compatible with the
|
||||
checkpoint folders of the experiment storage filesystem.
|
||||
|
||||
Raises:
|
||||
CheckpointManagerInitializationError: If the checkpoint manager snapshot
|
||||
is not consistent with the stored checkpoints.
|
||||
"""
|
||||
for checkpoint_result in self._checkpoint_results:
|
||||
checkpoint = checkpoint_result.checkpoint
|
||||
assert checkpoint is not None
|
||||
if not _exists_at_fs_path(
|
||||
fs=checkpoint.filesystem, fs_path=checkpoint.path
|
||||
):
|
||||
raise CheckpointManagerInitializationError(
|
||||
"The run snapshot contains a reference to a checkpoint "
|
||||
f"that does not exist anymore ({checkpoint}). You are "
|
||||
"running in a corrupted run directory `experiment_fs_path`. "
|
||||
"Please configure a new, unique `RunConfig(name)` "
|
||||
"or delete the existing folder at "
|
||||
f"`{self._storage_context.experiment_fs_path}`."
|
||||
)
|
||||
|
||||
# --------------------------
|
||||
# ReportCallback
|
||||
# --------------------------
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
metrics: List[Dict[str, Any]],
|
||||
):
|
||||
if not training_report.checkpoint:
|
||||
self._current_report_index += 1
|
||||
self._notify()
|
||||
return
|
||||
|
||||
self.register_checkpoint(training_report)
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def before_init_train_context(self, workers: List[Worker]) -> Dict[str, List[Any]]:
|
||||
latest_checkpoint = (
|
||||
self.latest_checkpoint_result.checkpoint
|
||||
if self.latest_checkpoint_result
|
||||
else None
|
||||
)
|
||||
train_context_args = {
|
||||
"checkpoint": [latest_checkpoint] * len(workers),
|
||||
"current_report_index": [self._current_report_index] * len(workers),
|
||||
}
|
||||
return train_context_args
|
||||
|
||||
# --------------------------------
|
||||
# Get all reported checkpoints API
|
||||
# --------------------------------
|
||||
|
||||
def _get_checkpoint_status(
|
||||
self, checkpoint: Checkpoint
|
||||
) -> ReportedCheckpointStatus:
|
||||
"""Get ReportedCheckpoint's status."""
|
||||
if checkpoint in self._pending_training_results:
|
||||
return ReportedCheckpointStatus.PENDING_VALIDATION
|
||||
elif checkpoint in self._timed_out_validation_checkpoints:
|
||||
return ReportedCheckpointStatus.VALIDATION_TIMEOUT
|
||||
elif checkpoint in self._failed_validation_checkpoints:
|
||||
return ReportedCheckpointStatus.VALIDATION_FAILED
|
||||
elif checkpoint in self._validated_checkpoints:
|
||||
return ReportedCheckpointStatus.VALIDATED
|
||||
else:
|
||||
return ReportedCheckpointStatus.COMMITTED
|
||||
|
||||
def _generate_get_all_reported_checkpoints_periodic_warning(
|
||||
self, start_time: float, consistency_mode: CheckpointConsistencyMode
|
||||
) -> str:
|
||||
"""Generates the warning message for the get_all_reported_checkpoints periodic warning."""
|
||||
return GET_ALL_REPORTED_CHECKPOINTS_PERIODIC_WARNING.format(
|
||||
consistency_mode=consistency_mode,
|
||||
time_elapsed_s=asyncio.get_event_loop().time() - start_time,
|
||||
warn_interval_env_var=COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
warn_interval_s=self._collective_warn_interval_s,
|
||||
)
|
||||
|
||||
async def get_all_reported_checkpoints(
|
||||
self,
|
||||
current_report_index: int,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List[ReportedCheckpoint]:
|
||||
"""Get all the reported checkpoints so far.
|
||||
|
||||
Args:
|
||||
current_report_index: The current report index.
|
||||
consistency_mode: Read semantics for checkpoint retrieval. Defaults to VALIDATED.
|
||||
timeout_s: Timeout in seconds. Defaults to None to run forever.
|
||||
|
||||
Returns:
|
||||
A list of ReportedCheckpoint objects that represent the checkpoints and
|
||||
corresponding metrics reported by the workers.
|
||||
"""
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
if consistency_mode == CheckpointConsistencyMode.COMMITTED:
|
||||
|
||||
def predicate() -> bool:
|
||||
return self._current_report_index == current_report_index
|
||||
|
||||
elif consistency_mode == CheckpointConsistencyMode.VALIDATED:
|
||||
|
||||
def predicate() -> bool:
|
||||
return (
|
||||
self._current_report_index == current_report_index
|
||||
and not self._pending_training_results
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unexpected CheckpointConsistencyMode: {consistency_mode}"
|
||||
)
|
||||
|
||||
async with self._condition:
|
||||
try:
|
||||
await wait_with_logging(
|
||||
self._condition,
|
||||
predicate=predicate,
|
||||
generate_warning_message=lambda: self._generate_get_all_reported_checkpoints_periodic_warning(
|
||||
start_time, consistency_mode
|
||||
),
|
||||
warn_interval_s=self._collective_warn_interval_s,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
# Time out due to checkpoint upload or validation in progress
|
||||
logger.debug(
|
||||
"Timed out waiting for reported_checkpoint to become available."
|
||||
)
|
||||
|
||||
# TODO: might be nice for CheckpointManager to manage ReportedCheckpoint
|
||||
# instead of _TrainingResult but that is a large refactor.
|
||||
return [
|
||||
ReportedCheckpoint(
|
||||
checkpoint=tr.checkpoint,
|
||||
metrics=tr.metrics,
|
||||
status=self._get_checkpoint_status(tr.checkpoint),
|
||||
)
|
||||
for tr in self._checkpoint_results
|
||||
]
|
||||
@@ -0,0 +1,129 @@
|
||||
from collections import deque
|
||||
from typing import Deque, List, Optional
|
||||
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReplicaGroupCallback,
|
||||
ReportCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupPollStatus,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group.execution_group import ReplicaGroup
|
||||
|
||||
|
||||
class ReportCallbackHandler(ReplicaGroupCallback, WorkerGroupCallback):
|
||||
"""Consolidate training results from multiple workers and call
|
||||
subscribers implementing the `ReportCallback` interface sequentially.
|
||||
"""
|
||||
|
||||
def __init__(self, report_callbacks: List[ReportCallback]):
|
||||
# We set the worker group after it has been started and remove it after it
|
||||
# has been shut down.
|
||||
self._worker_group: Optional[WorkerGroup] = None
|
||||
# A list of queues holding training reports from workers.
|
||||
self._training_report_queues: Optional[List[Deque[_TrainingReport]]] = None
|
||||
|
||||
self._report_callbacks = report_callbacks
|
||||
|
||||
def _assert_initialized(self):
|
||||
assert (
|
||||
self._worker_group and self._training_report_queues
|
||||
), "Need to call initialize state with `after_worker_group_start` first."
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def after_worker_group_poll_status(
|
||||
self, worker_group_status: WorkerGroupPollStatus
|
||||
) -> None:
|
||||
"""Handle training results as they roll in from worker status polls.
|
||||
|
||||
Wait for all workers to report training results to collect
|
||||
a consolidated training result.
|
||||
"""
|
||||
# Step 1: Assert that the worker group has been started and not shut down.
|
||||
self._assert_initialized()
|
||||
|
||||
assert len(self._worker_group) == len(worker_group_status.worker_statuses), (
|
||||
f"The number of workers in the worker group has changed unexpectedly. "
|
||||
f"Expected: {len(self._worker_group)}, got: {len(worker_group_status.worker_statuses)}"
|
||||
)
|
||||
|
||||
# Step 2: Update training_reports_queues with poll_results.
|
||||
for i in range(len(self._worker_group)):
|
||||
training_report = worker_group_status.worker_statuses[i].training_report
|
||||
if training_report:
|
||||
self._training_report_queues[i].append(training_report)
|
||||
|
||||
# Directly return if any of the worker result queues are empty.
|
||||
if not all(self._training_report_queues):
|
||||
return
|
||||
|
||||
training_reports = [q.popleft() for q in self._training_report_queues]
|
||||
|
||||
# Step 3: Consolidate a list of checkpoints to single checkpoint.
|
||||
# Use the first checkpoint as the consolidated checkpoint.
|
||||
checkpoint_results = [
|
||||
tr for tr in training_reports if tr.checkpoint is not None
|
||||
]
|
||||
|
||||
consolidated_checkpoint = None
|
||||
validation = False
|
||||
if checkpoint_results:
|
||||
# Double check the storage path of the checkpoints in the training results.
|
||||
unique_checkpoint_paths = {tr.checkpoint.path for tr in checkpoint_results}
|
||||
if len(unique_checkpoint_paths) > 1:
|
||||
# TODO: Support for inconsistent checkpoints path from workers
|
||||
# instead of hard raising error. Maybe drop this iteration of
|
||||
# training results and continue with the next iteration.
|
||||
raise RuntimeError(
|
||||
"The storage path of the checkpoints in the training results "
|
||||
"is not the same. This means the checkpoints are not consistent."
|
||||
"Got a mix of the following checkpoint paths: "
|
||||
f"{unique_checkpoint_paths}\n"
|
||||
"This is unexpected -- please file a Github issue."
|
||||
)
|
||||
consolidated_checkpoint = checkpoint_results[0].checkpoint
|
||||
validation = checkpoint_results[0].validation
|
||||
|
||||
# Step 4: Invoke all dependent `ReportCallback`s.
|
||||
metrics_per_worker = [
|
||||
training_report.metrics for training_report in training_reports
|
||||
]
|
||||
for callback in self._report_callbacks:
|
||||
callback.after_report(
|
||||
training_report=_TrainingReport(
|
||||
checkpoint=consolidated_checkpoint,
|
||||
metrics=metrics_per_worker[0],
|
||||
validation=validation,
|
||||
),
|
||||
metrics=metrics_per_worker,
|
||||
)
|
||||
|
||||
def after_worker_group_start(self, worker_group: WorkerGroup) -> None:
|
||||
"""Handle worker group start. Initialize internal states."""
|
||||
self._worker_group = worker_group
|
||||
self._training_report_queues = [deque() for _ in range(len(self._worker_group))]
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: WorkerGroup) -> None:
|
||||
"""Handle worker group shutdown. Clear internal states.
|
||||
|
||||
None of the partial reported results are valid at this point.
|
||||
"""
|
||||
self._worker_group = None
|
||||
self._training_report_queues = None
|
||||
|
||||
# --------------------------
|
||||
# ReplicaGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def after_replica_group_start(self, replica_group: ReplicaGroup) -> None:
|
||||
"""Handle replica group start. Initialize internal states."""
|
||||
self._assert_initialized()
|
||||
# TODO: it might be possible to reuse existing queues.
|
||||
# For example, if 3/4 ddp workers reported a checkpoint, that checkpoint is usable.
|
||||
self._training_report_queues = [deque() for _ in range(len(self._worker_group))]
|
||||
@@ -0,0 +1,226 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import List, Optional, TypeVar
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.constants import (
|
||||
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
DEFAULT_COLLECTIVE_TIMEOUT_S,
|
||||
DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
|
||||
)
|
||||
from ray.train.v2._internal.exceptions import BroadcastCollectiveTimeoutError
|
||||
from ray.train.v2._internal.util import wait_with_logging
|
||||
|
||||
T = TypeVar("T", bound=Optional[object])
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SynchronizationBarrierResetError(Exception):
|
||||
"""Raised when the synchronization barrier is reset, e.g. due to a worker failure."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
BROADCAST_PERIODIC_WARNING = """
|
||||
`{caller_method_name}` has not been called by all {world_size} workers in the group.
|
||||
The workers have been waiting for {max_time_elapsed_s:.2f} s for the following ranks to join the `{caller_method_name}` call: {missing_ranks}.
|
||||
Also ensure that workers are not hanging on other operations, causing them to miss this synchronization barrier.
|
||||
You can set the {warn_interval_env_var} environment variable to change the frequency of this warning (current value: {warn_interval_s} s).
|
||||
"""
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0) # type: ignore
|
||||
class SynchronizationActor:
|
||||
"""A Ray actor that synchronizes the workers in a distributed training job.
|
||||
|
||||
This actor forms a synchronization barrier on a group of processes.
|
||||
Every time a worker calls the broadcast_from_rank_zero method,
|
||||
the counter is incremented. When the counter equals to the world size,
|
||||
the actor notifies all the workers to continue.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
timeout_s: Optional[float] = DEFAULT_COLLECTIVE_TIMEOUT_S,
|
||||
warn_interval_s: float = DEFAULT_COLLECTIVE_WARN_INTERVAL_S,
|
||||
):
|
||||
self._counter: int = 0
|
||||
self._world_size: int = 0
|
||||
self._condition = asyncio.Condition()
|
||||
self._reduced_data = None
|
||||
self._reset = False
|
||||
# The time when workers from different ranks
|
||||
# enters the synchronization barrier.
|
||||
self._sync_start_times: List[Optional[float]] = []
|
||||
# The timeout in seconds for the synchronization barrier.
|
||||
self._timeout_s: Optional[float] = timeout_s
|
||||
# The interval in seconds to log a warning when waiting for the barrier.
|
||||
self._warn_interval_s: float = warn_interval_s
|
||||
|
||||
def get_counter(self):
|
||||
"""Returns the current value of the counter."""
|
||||
return self._counter
|
||||
|
||||
def get_world_size(self):
|
||||
"""Returns the current value of the world_size."""
|
||||
return self._world_size
|
||||
|
||||
def get_reduced_data(self):
|
||||
"""Returns the current value of the reduced_data."""
|
||||
return self._reduced_data
|
||||
|
||||
def _clear_states(self):
|
||||
"""Clears the states of the actor. When the last worker has
|
||||
called the _clear_states method, the actor clears its states
|
||||
"""
|
||||
self._counter -= 1
|
||||
if self._counter == 0:
|
||||
self._reduced_data = None
|
||||
self._world_size = 0
|
||||
self._reset = False
|
||||
self._condition.notify_all()
|
||||
|
||||
async def _setup_or_validate_collective_op(self, world_size: int):
|
||||
"""The setup method for the synchronization actor if it is not setup yet.
|
||||
It initializes the world size and the start times for the
|
||||
synchronization barrier.
|
||||
"""
|
||||
# Wait for previous collective reset to finish.
|
||||
await self._condition.wait_for(lambda: not self._reset)
|
||||
if self._world_size == 0:
|
||||
self._world_size = world_size
|
||||
self._sync_start_times = [None] * world_size
|
||||
elif world_size != self._world_size:
|
||||
raise ValueError(
|
||||
f"Expected all callers to provide the same world size. \
|
||||
Got {world_size} and expected {self._world_size}."
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _broadcast_collective_context_manager(
|
||||
self, world_rank: int, world_size: int, data: T
|
||||
):
|
||||
"""A context manager that ensures the synchronization barrier is lifted
|
||||
after the block of code is executed.
|
||||
"""
|
||||
try:
|
||||
await self._setup_or_validate_collective_op(world_size)
|
||||
if world_rank == 0:
|
||||
self._reduced_data = data
|
||||
if self._counter < self._world_size:
|
||||
self._counter += 1
|
||||
yield
|
||||
finally:
|
||||
self._clear_states()
|
||||
|
||||
def _get_time_elapsed(self) -> Optional[float]:
|
||||
"""Return the time elapsed since the first worker entered the barrier.
|
||||
If no workers have entered the barrier, returns None.
|
||||
"""
|
||||
start_times = [t for t in self._sync_start_times if t is not None]
|
||||
if not start_times:
|
||||
return None
|
||||
|
||||
return asyncio.get_event_loop().time() - min(start_times)
|
||||
|
||||
def _get_missing_ranks(self) -> List[int]:
|
||||
"""Returns the ranks that have not entered the synchronization barrier."""
|
||||
return [i for i, t in enumerate(self._sync_start_times) if t is None]
|
||||
|
||||
def _generate_broadcast_periodic_warning(self, caller_method_name: str) -> str:
|
||||
"""Generates the warning message for the broadcast periodic warning."""
|
||||
|
||||
return BROADCAST_PERIODIC_WARNING.format(
|
||||
caller_method_name=caller_method_name,
|
||||
world_size=self._world_size,
|
||||
max_time_elapsed_s=self._get_time_elapsed(),
|
||||
missing_ranks=self._get_missing_ranks(),
|
||||
warn_interval_env_var=COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
warn_interval_s=self._warn_interval_s,
|
||||
)
|
||||
|
||||
async def reset(self):
|
||||
"""Reset the synchronization barrier, unblocking any waiting workers.
|
||||
|
||||
If no workers are currently at the barrier, this is a no-op.
|
||||
Waiting workers will raise SynchronizationBarrierResetError.
|
||||
The actor remains alive and usable for subsequent barriers.
|
||||
"""
|
||||
async with self._condition:
|
||||
if self._counter == 0:
|
||||
return
|
||||
self._reset = True
|
||||
self._condition.notify_all()
|
||||
|
||||
async def broadcast_from_rank_zero(
|
||||
self,
|
||||
world_rank: int,
|
||||
world_size: int,
|
||||
data: T,
|
||||
caller_method_name: str,
|
||||
) -> T:
|
||||
"""Broadcasts a data from the worker with rank 0 to all other workers.
|
||||
|
||||
This method is a coroutine that blocks until all workers have called this
|
||||
method with the their data. The data from the worker with rank 0 will
|
||||
be returned.
|
||||
|
||||
Args:
|
||||
world_rank: The rank of the worker that calls this method.
|
||||
world_size: The total number of workers in the group.
|
||||
data: The data to broadcast.
|
||||
caller_method_name: The name of the method that calls this method.
|
||||
|
||||
Returns:
|
||||
The data broadcasted from the worker with rank 0.
|
||||
"""
|
||||
# TODO: resolve https://github.com/ray-project/ray/pull/54066#discussion_r2180657435
|
||||
# We couldn't reproduce the issue but the asyncio docs don't say it can't happen.
|
||||
|
||||
# Ensures that all global states manipulation is done within the async context
|
||||
# manager which makes the condition variable awaiting and the counter
|
||||
# incrementing an atomic operation.
|
||||
async with self._condition:
|
||||
async with self._broadcast_collective_context_manager(
|
||||
world_rank, world_size, data
|
||||
):
|
||||
# If the counter is equal to the world size, it means the last worker
|
||||
# has called the broadcast_from_rank_zero method. The actor notifies
|
||||
# all the workers to continue.
|
||||
if self._counter == self._world_size:
|
||||
self._condition.notify_all()
|
||||
return self._reduced_data
|
||||
# If the counter is less than the world size, the actor waits for the
|
||||
# other workers to call the broadcast_from_rank_zero method.
|
||||
try:
|
||||
current_time = asyncio.get_event_loop().time()
|
||||
self._sync_start_times[world_rank] = current_time
|
||||
await wait_with_logging(
|
||||
self._condition,
|
||||
predicate=None,
|
||||
generate_warning_message=(
|
||||
lambda: self._generate_broadcast_periodic_warning(
|
||||
caller_method_name
|
||||
)
|
||||
)
|
||||
if world_rank == 0
|
||||
else None,
|
||||
warn_interval_s=self._warn_interval_s,
|
||||
timeout_s=self._timeout_s,
|
||||
)
|
||||
if self._reset:
|
||||
raise SynchronizationBarrierResetError(
|
||||
"Synchronization barrier was reset, likely due "
|
||||
"to a worker failure and replica group replacement."
|
||||
)
|
||||
return self._reduced_data
|
||||
except (asyncio.TimeoutError, TimeoutError) as e:
|
||||
raise BroadcastCollectiveTimeoutError(
|
||||
time_elapsed=self._get_time_elapsed(),
|
||||
missing_ranks=self._get_missing_ranks(),
|
||||
timeout_s=self._timeout_s,
|
||||
) from e
|
||||
|
||||
# TODO: Implement a general consensus_from_votes method that takes a callable
|
||||
# reduce_fn and a list of votes from each worker. The method returns the consensus
|
||||
@@ -0,0 +1,289 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections import OrderedDict, deque
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import ray
|
||||
from ray.train._checkpoint import Checkpoint
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
ReportCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.training_report import (
|
||||
_TrainingReport,
|
||||
)
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
|
||||
from ray.train.v2.api.validation_config import ValidationConfig, ValidationTaskConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.controller import TrainControllerState
|
||||
from ray.train.v2._internal.execution.worker_group.worker import Worker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
VALIDATION_TASK_POLL_INTERVAL_S = 1
|
||||
MAX_IN_FLIGHT_VALIDATIONS = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PendingValidation:
|
||||
checkpoint: Checkpoint
|
||||
start_time: float
|
||||
# None when no timeout applies.
|
||||
timeout_s: Optional[float]
|
||||
|
||||
def __post_init__(self):
|
||||
assert (
|
||||
self.timeout_s is None or self.timeout_s > 0
|
||||
), f"timeout_s needs to be None (for no timeout) or a positive value in seconds. Actual value: {self.timeout_s}"
|
||||
|
||||
|
||||
@ray.remote
|
||||
def run_validation_fn(
|
||||
validation_config: ValidationConfig,
|
||||
validation_task_config: Union[bool, ValidationTaskConfig],
|
||||
checkpoint: Checkpoint,
|
||||
) -> Dict:
|
||||
"""Run the user-defined validation function.
|
||||
|
||||
Merges fn_kwargs from validation_config.task_config (defaults) with
|
||||
fn_kwargs from validation_task_config (per-report overrides).
|
||||
"""
|
||||
# Merge kwargs: defaults from validation_config, overrides from validation_task_config
|
||||
if validation_task_config is True:
|
||||
merged_kwargs = validation_config.task_config.fn_kwargs
|
||||
else:
|
||||
merged_kwargs = {
|
||||
**validation_config.task_config.fn_kwargs,
|
||||
**validation_task_config.fn_kwargs,
|
||||
}
|
||||
metrics_dict = validation_config.fn(
|
||||
checkpoint,
|
||||
**merged_kwargs,
|
||||
)
|
||||
if not isinstance(metrics_dict, dict):
|
||||
raise ValueError(
|
||||
"The validation function must return a dictionary of metrics. "
|
||||
f"Got {type(metrics_dict)} instead."
|
||||
)
|
||||
return metrics_dict
|
||||
|
||||
|
||||
class ValidationManager(ControllerCallback, ReportCallback, WorkerGroupCallback):
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint_manager: CheckpointManager,
|
||||
validation_config: ValidationConfig,
|
||||
):
|
||||
self._checkpoint_manager = checkpoint_manager
|
||||
self._validation_config = validation_config
|
||||
|
||||
# _TrainingReports that we will validate
|
||||
self._training_report_queue = deque()
|
||||
|
||||
# Map from in flight validation task to its pending-validation record
|
||||
# (checkpoint + start_time + resolved timeout).
|
||||
self._pending_validations: "OrderedDict[ray.ObjectRef, _PendingValidation]" = (
|
||||
OrderedDict()
|
||||
)
|
||||
|
||||
# Tasks that this manager proactively cancelled due to timeout. Used to
|
||||
# distinguish timeout-cancels from controller-abort-cancels (both raise
|
||||
# TaskCancelledError on ray.get).
|
||||
self._timed_out_tasks: set = set()
|
||||
|
||||
# Map from validation task to checkpoint
|
||||
# Finished validations that have yet to be processed
|
||||
self._finished_validations: "OrderedDict[ray.ObjectRef, Checkpoint]" = (
|
||||
OrderedDict()
|
||||
)
|
||||
|
||||
self._requeue_incomplete_validations()
|
||||
|
||||
def _requeue_incomplete_validations(self):
|
||||
"""Add _TrainingReports for incomplete validations to the queue."""
|
||||
for checkpoint, (
|
||||
training_result,
|
||||
validation,
|
||||
) in self._checkpoint_manager.get_pending_training_results().items():
|
||||
if validation:
|
||||
self._training_report_queue.append(
|
||||
_TrainingReport(
|
||||
metrics=training_result.metrics,
|
||||
checkpoint=checkpoint,
|
||||
validation=validation,
|
||||
)
|
||||
)
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
metrics: List[Dict[str, Any]],
|
||||
):
|
||||
if training_report.validation:
|
||||
self._training_report_queue.append(training_report)
|
||||
|
||||
def _cancel_timed_out_validations(self):
|
||||
"""Cancel any in-flight validation that has exceeded its timeout_s.
|
||||
|
||||
Cancelled tasks are moved directly from `_pending_validations` to
|
||||
`_finished_validations` so the MAX_IN_FLIGHT slot is freed immediately
|
||||
and the task flows through the normal finished-processing pipeline
|
||||
without waiting for `ray.wait` to echo the cancellation.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
for task, pending in list(self._pending_validations.items()):
|
||||
if (
|
||||
pending.timeout_s is None
|
||||
or now - pending.start_time < pending.timeout_s
|
||||
):
|
||||
continue
|
||||
self._pending_validations.pop(task)
|
||||
logger.warning(
|
||||
f"Validation for checkpoint {pending.checkpoint} exceeded "
|
||||
f"timeout_s={pending.timeout_s}s. Cancelling."
|
||||
)
|
||||
self._timed_out_tasks.add(task)
|
||||
ray.cancel(task, force=True)
|
||||
self._finished_validations[task] = pending.checkpoint
|
||||
|
||||
def _poll_validations(self) -> int:
|
||||
"""Poll/process validations, update checkpoint manager, return num pending validations."""
|
||||
self._cancel_timed_out_validations()
|
||||
|
||||
# Move pending validations to finished validations
|
||||
validation_tasks = list(self._pending_validations.keys())
|
||||
done, _ = ray.wait(
|
||||
validation_tasks, timeout=0, num_returns=len(validation_tasks)
|
||||
)
|
||||
done_checkpoints = []
|
||||
for task in done:
|
||||
pending = self._pending_validations.pop(task)
|
||||
done_checkpoints.append(pending.checkpoint)
|
||||
self._finished_validations[task] = pending.checkpoint
|
||||
if done_checkpoints:
|
||||
logger.info(
|
||||
f"Finished async validation task(s) for checkpoint(s): {done_checkpoints}.\n"
|
||||
f"Running validations for checkpoint(s): {[p.checkpoint for p in self._pending_validations.values()]}.\n"
|
||||
f"Staged validations for checkpoint(s): {[tr.checkpoint for tr in self._training_report_queue]}."
|
||||
)
|
||||
|
||||
# Process finished validations (one at a time)
|
||||
if self._finished_validations:
|
||||
task, checkpoint = self._finished_validations.popitem(last=False)
|
||||
update = self._process_finished_validation(task, checkpoint)
|
||||
if update is not None:
|
||||
self._checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{checkpoint: update}
|
||||
)
|
||||
|
||||
return len(self._pending_validations)
|
||||
|
||||
def _kick_off_validations(self) -> int:
|
||||
"""Kick off validations and return the number of pending validations."""
|
||||
# TODO: figure out where to place run_validation_fn task:
|
||||
# TODO: provide option to run this on gpu?
|
||||
num_validations_to_start = max(
|
||||
MAX_IN_FLIGHT_VALIDATIONS - len(self._pending_validations), 0
|
||||
)
|
||||
num_validations_to_start = min(
|
||||
num_validations_to_start, len(self._training_report_queue)
|
||||
)
|
||||
for _ in range(num_validations_to_start):
|
||||
training_report = self._training_report_queue.popleft()
|
||||
run_validation_fn_with_options = run_validation_fn.options(
|
||||
**self._validation_config.ray_remote_kwargs,
|
||||
)
|
||||
validate_task = run_validation_fn_with_options.remote(
|
||||
self._validation_config,
|
||||
training_report.validation,
|
||||
training_report.checkpoint,
|
||||
)
|
||||
if isinstance(training_report.validation, ValidationTaskConfig):
|
||||
timeout_s = training_report.validation.timeout_s
|
||||
else:
|
||||
timeout_s = self._validation_config.task_config.timeout_s
|
||||
self._pending_validations[validate_task] = _PendingValidation(
|
||||
checkpoint=training_report.checkpoint,
|
||||
start_time=time.monotonic(),
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
logger.info(
|
||||
f"Launched async validation task for checkpoint {training_report.checkpoint}"
|
||||
)
|
||||
return len(self._pending_validations)
|
||||
|
||||
def _process_finished_validation(
|
||||
self, task: ray.ObjectRef, checkpoint: Checkpoint
|
||||
) -> Optional[Tuple[Dict[str, Any], ReportedCheckpointStatus]]:
|
||||
"""Process finished validation. Returns (metrics, status) or None.
|
||||
|
||||
Returns None when the task was cancelled by a controller abort (not a
|
||||
timeout), leaving it pending so it re-queues on resumption.
|
||||
"""
|
||||
was_timed_out = task in self._timed_out_tasks
|
||||
self._timed_out_tasks.discard(task)
|
||||
if was_timed_out:
|
||||
logger.info(
|
||||
f"Validation for checkpoint {checkpoint} was cancelled due to timeout."
|
||||
)
|
||||
return {}, ReportedCheckpointStatus.VALIDATION_TIMEOUT
|
||||
|
||||
try:
|
||||
metrics = ray.get(task)
|
||||
return metrics, ReportedCheckpointStatus.VALIDATED
|
||||
except ray.exceptions.TaskCancelledError:
|
||||
logger.info(
|
||||
f"Validation was cancelled for checkpoint {checkpoint}, likely because the train run was aborted. "
|
||||
"It will be retried in the next train run with the same storage path if there is one."
|
||||
)
|
||||
return None
|
||||
except ray.exceptions.RayTaskError:
|
||||
logger.exception(f"Validation failed for checkpoint {checkpoint}")
|
||||
return {}, ReportedCheckpointStatus.VALIDATION_FAILED
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
while self._poll_validations() != 0 or self._kick_off_validations() != 0:
|
||||
await asyncio.sleep(VALIDATION_TASK_POLL_INTERVAL_S)
|
||||
checkpoint_updates: Dict[
|
||||
Checkpoint, Tuple[Dict[str, Any], ReportedCheckpointStatus]
|
||||
] = {}
|
||||
tasks = list(self._finished_validations.keys())
|
||||
for task in tasks:
|
||||
checkpoint = self._finished_validations[task]
|
||||
self._finished_validations.pop(task)
|
||||
update = self._process_finished_validation(task, checkpoint)
|
||||
if update is not None:
|
||||
checkpoint_updates[checkpoint] = update
|
||||
self._checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
checkpoint_updates
|
||||
)
|
||||
|
||||
def before_controller_abort(self):
|
||||
for task in self._pending_validations.keys():
|
||||
ray.cancel(task)
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: "TrainControllerState",
|
||||
current_state: "TrainControllerState",
|
||||
):
|
||||
# TODO: figure out if there's a better place to poll validations
|
||||
if current_state.is_terminal():
|
||||
return
|
||||
self._poll_validations()
|
||||
self._kick_off_validations()
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List["Worker"]
|
||||
) -> Dict[str, List[bool]]:
|
||||
return {
|
||||
"has_validation_fn": [True] * len(workers),
|
||||
}
|
||||
Reference in New Issue
Block a user