chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,889 @@
|
||||
load("@rules_python//python:defs.bzl", "py_library", "py_test")
|
||||
load("//bazel:python.bzl", "doctest")
|
||||
|
||||
doctest(
|
||||
name = "py_doctest[train_v2]",
|
||||
size = "large",
|
||||
env = {
|
||||
"RAY_TRAIN_V2_ENABLED": "1",
|
||||
"TF_USE_LEGACY_KERAS": "1",
|
||||
},
|
||||
files = glob(
|
||||
["**/*.py"],
|
||||
exclude = [
|
||||
"tests/**",
|
||||
"horovod/**",
|
||||
"jax/**",
|
||||
"examples/pytorch/torchft_linear_example.py",
|
||||
],
|
||||
),
|
||||
tags = ["team:ml"],
|
||||
)
|
||||
|
||||
py_library(
|
||||
name = "conftest",
|
||||
srcs = ["tests/conftest.py"],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_accelerator_utils",
|
||||
size = "small",
|
||||
srcs = ["tests/test_accelerator_utils.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_async_checkpointing_validation",
|
||||
size = "large",
|
||||
srcs = ["tests/test_async_checkpointing_validation.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_checkpoint_manager",
|
||||
size = "small",
|
||||
srcs = ["tests/test_checkpoint_manager.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_config",
|
||||
size = "small",
|
||||
srcs = ["tests/test_config.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_circular_imports",
|
||||
size = "small",
|
||||
srcs = ["tests/test_circular_imports.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_circular_import_linter",
|
||||
size = "small",
|
||||
srcs = ["tests/test_circular_import_linter.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_validation_manager",
|
||||
size = "small",
|
||||
srcs = ["tests/test_validation_manager.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_collective",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_collective.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_callback_manager",
|
||||
size = "small",
|
||||
srcs = ["tests/test_callback_manager.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_controller",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_controller.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_elastic_scaling_policy",
|
||||
size = "small",
|
||||
srcs = ["tests/test_elastic_scaling_policy.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_elastic_e2e",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_elastic_e2e.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_controller_callback_behaviour",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_controller_callback_behaviour.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_data_integration",
|
||||
size = "large",
|
||||
srcs = ["tests/test_data_integration.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"data_integration",
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_data_parallel_trainer",
|
||||
size = "large",
|
||||
srcs = ["tests/test_data_parallel_trainer.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_report_fault_tolerance",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_report_fault_tolerance.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_data_resource_cleanup",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_data_resource_cleanup.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"data_integration",
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_dataset_manager",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_dataset_manager.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"data_integration",
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_env_callbacks",
|
||||
size = "small",
|
||||
srcs = ["tests/test_env_callbacks.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_failure_policy",
|
||||
size = "small",
|
||||
srcs = ["tests/test_failure_policy.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_jax_elastic_e2e",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_jax_elastic_e2e.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_jax_trainer",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_jax_trainer.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_jax_gpu",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_jax_gpu.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
# Temporarily disabled in CI: Ray GPU CI is on CUDA 12.1, but JAX wheels
|
||||
# do not support CUDA 12.1 (and older JAX versions were removed from PyPI).
|
||||
# Tagging as 'manual' excludes it from Ray's default Bazel CI runs.
|
||||
"manual",
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2_gpu",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_lightgbm_trainer",
|
||||
size = "small",
|
||||
srcs = ["tests/test_lightgbm_trainer.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_lightning_integration",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_lightning_integration.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_logging",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_logging.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_metrics",
|
||||
size = "small",
|
||||
srcs = ["tests/test_metrics.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_persistence",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_persistence.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_report_handler",
|
||||
size = "small",
|
||||
srcs = ["tests/test_report_handler.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_placement_group_cleaner",
|
||||
size = "small",
|
||||
srcs = ["tests/test_placement_group_cleaner.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_placement_group_handle",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_placement_group_handle.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_preemption_watcher",
|
||||
size = "small",
|
||||
srcs = ["tests/test_preemption_watcher.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_result",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_result.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_scheduling",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_scheduling.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_serialization",
|
||||
size = "small",
|
||||
srcs = ["tests/test_serialization.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_state",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_state.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_state_export",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_state_export.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_storage",
|
||||
size = "small",
|
||||
srcs = ["tests/test_storage.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_sync_actor",
|
||||
size = "small",
|
||||
srcs = ["tests/test_sync_actor.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_telemetry",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_telemetry.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_tensorflow_trainer",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_tensorflow_trainer.py"],
|
||||
env = {
|
||||
"RAY_TRAIN_V2_ENABLED": "1",
|
||||
"TF_USE_LEGACY_KERAS": "1",
|
||||
},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_thread_runner",
|
||||
size = "small",
|
||||
srcs = ["tests/test_thread_runner.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_torch_gpu",
|
||||
size = "large",
|
||||
srcs = ["tests/test_torch_gpu.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2_gpu",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_torch_trainer",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_torch_trainer.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"torchft",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_torch_transformers_train",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_torch_transformers_train.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2_gpu",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_util",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_util.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_v2_api",
|
||||
size = "small",
|
||||
srcs = ["tests/test_v2_api.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_worker",
|
||||
size = "small",
|
||||
srcs = ["tests/test_worker.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_worker_group",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_worker_group.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_worker_group_poll_status",
|
||||
size = "small",
|
||||
srcs = ["tests/test_worker_group_poll_status.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_xgboost_trainer",
|
||||
size = "small",
|
||||
srcs = ["tests/test_xgboost_trainer.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_local_mode",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_local_mode.py"],
|
||||
env = {
|
||||
"RAY_TRAIN_V2_ENABLED": "1",
|
||||
"TF_USE_LEGACY_KERAS": "1",
|
||||
},
|
||||
tags = [
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
|
||||
py_test(
|
||||
name = "test_data_config",
|
||||
size = "medium",
|
||||
srcs = ["tests/test_data_config.py"],
|
||||
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
||||
tags = [
|
||||
"data_integration",
|
||||
"exclusive",
|
||||
"team:ml",
|
||||
"train_v2",
|
||||
],
|
||||
deps = [
|
||||
":conftest",
|
||||
"//:ray_lib",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,16 @@
|
||||
from .accelerators import AcceleratorSetupCallback
|
||||
from .backend_setup import BackendSetupCallback
|
||||
from .datasets import DatasetsCallback
|
||||
from .state_manager import StateManagerCallback
|
||||
from .working_dir_setup import WorkingDirectorySetupCallback
|
||||
|
||||
__all__ = [
|
||||
"AcceleratorSetupCallback",
|
||||
"BackendSetupCallback",
|
||||
"DatasetsCallback",
|
||||
"StateManagerCallback",
|
||||
"WorkingDirectorySetupCallback",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,160 @@
|
||||
import logging
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING, Any, Dict, List
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.train import BackendConfig
|
||||
from ray.train.constants import ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.worker_group import ActorMetadata
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.worker_group.worker import Worker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AcceleratorSetupCallback(WorkerGroupCallback):
|
||||
"""Perform accelerator setup for workers.
|
||||
|
||||
For example, this callback can be used to share CUDA_VISIBLE_DEVICES
|
||||
among workers on the same node.
|
||||
"""
|
||||
|
||||
def __init__(self, backend_config: BackendConfig, scaling_config: ScalingConfig):
|
||||
self._backend = backend_config.backend_cls()
|
||||
self._scaling_config = scaling_config
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List["Worker"]
|
||||
) -> Dict[str, List[Any]]:
|
||||
self._maybe_share_cuda_visible_devices(workers)
|
||||
# TODO: Add support for sharing other accelerator resources.
|
||||
|
||||
return {}
|
||||
|
||||
def _maybe_share_cuda_visible_devices(self, workers: List["Worker"]):
|
||||
"""Set CUDA visible devices environment variables on workers."""
|
||||
share_cuda_visible_devices_enabled = env_bool(
|
||||
ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
|
||||
self._backend.share_cuda_visible_devices,
|
||||
)
|
||||
|
||||
if (
|
||||
self._scaling_config._resources_per_worker_not_none.get("GPU", 0) > 0
|
||||
and share_cuda_visible_devices_enabled
|
||||
):
|
||||
_share_cuda_visible_devices(workers)
|
||||
|
||||
|
||||
def _share_cuda_visible_devices(workers: List["Worker"]):
|
||||
"""Sets CUDA_VISIBLE_DEVICES on all workers.
|
||||
For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs
|
||||
visible to all workers on that worker's node.
|
||||
This allows GPU workers on the same node to communicate with one
|
||||
another.
|
||||
|
||||
Example:
|
||||
Setup:
|
||||
- Node1:
|
||||
- Worker1: {0, 1}
|
||||
- Worker2: {2, 3}
|
||||
- Node2:
|
||||
- Worker3: {0, 1}
|
||||
CUDA_VISIBLE_DEVICES:
|
||||
- Worker1: "0,1,2,3"
|
||||
- Worker2: "0,1,2,3"
|
||||
- Worker3: "0,1"
|
||||
|
||||
Args:
|
||||
workers: List of worker objects.
|
||||
"""
|
||||
_share_accelerator_ids(workers, ray_constants.GPU, CUDA_VISIBLE_DEVICES_ENV_VAR)
|
||||
|
||||
|
||||
def _share_accelerator_ids(
|
||||
workers: List["Worker"], accelerator_name: str, env_var: str
|
||||
):
|
||||
"""Sets the given env_var on all workers.
|
||||
For each worker, the cores/devices are visible to all the
|
||||
workers on that worker's node. This allows workers on the
|
||||
same node to communicate with one another.
|
||||
|
||||
Example:
|
||||
Setup:
|
||||
- Node1:
|
||||
- Worker1: {0, 1}
|
||||
- Worker2: {2, 3}
|
||||
- Node2:
|
||||
- Worker3: {0, 1}
|
||||
NEURON_RT_VISIBLE_CORES/TPU_VISIBLE_CHIPS/...:
|
||||
- Worker1: "0,1,2,3"
|
||||
- Worker2: "0,1,2,3"
|
||||
- Worker3: "0,1"
|
||||
|
||||
Args:
|
||||
workers: List of worker objects.
|
||||
accelerator_name: The name of the accelerator.
|
||||
env_var: The name of the environment variable to set.
|
||||
"""
|
||||
worker_metadatas = [worker.metadata for worker in workers]
|
||||
visible_accelerator_ids_per_worker = _get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas=worker_metadatas, accelerator_name=accelerator_name
|
||||
)
|
||||
|
||||
def set_accelerator_ids(accelerator_ids):
|
||||
os.environ[env_var] = accelerator_ids
|
||||
|
||||
futures = []
|
||||
for rank, visible_accelerator_ids in enumerate(visible_accelerator_ids_per_worker):
|
||||
futures.append(
|
||||
workers[rank].execute_async(
|
||||
set_accelerator_ids, accelerator_ids=visible_accelerator_ids
|
||||
)
|
||||
)
|
||||
ray.get(futures)
|
||||
|
||||
|
||||
def _get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas: List[ActorMetadata], accelerator_name: str
|
||||
) -> List[str]:
|
||||
"""Returns a list of comma-separated accelerator IDs visible to each worker.
|
||||
|
||||
All workers on a node should have the same set of visible accelerators,
|
||||
which is the union of accelerator ids of the workers.
|
||||
|
||||
Args:
|
||||
worker_metadatas: The actor metadata for each worker.
|
||||
accelerator_name: The name of the accelerator resource to inspect.
|
||||
|
||||
Returns:
|
||||
A list of comma-separated accelerator ID strings. This list is the
|
||||
same length as the number of workers.
|
||||
"""
|
||||
for metadata in worker_metadatas:
|
||||
if accelerator_name not in metadata.accelerator_ids:
|
||||
raise ValueError(
|
||||
f"Accelerator '{accelerator_name}' is not available on all workers. "
|
||||
f"Got these available accelerators instead: {metadata.accelerator_ids}"
|
||||
)
|
||||
|
||||
node_id_to_accelerator_ids = defaultdict(set)
|
||||
|
||||
for metadata in worker_metadatas:
|
||||
node_id_to_accelerator_ids[metadata.node_id].update(
|
||||
metadata.accelerator_ids[accelerator_name]
|
||||
)
|
||||
|
||||
visible_accelerator_ids_per_worker = []
|
||||
for worker_id in range(len(worker_metadatas)):
|
||||
node_id = worker_metadatas[worker_id].node_id
|
||||
accelerator_ids = sorted(node_id_to_accelerator_ids[node_id])
|
||||
all_resource_ids = ",".join([str(id) for id in accelerator_ids])
|
||||
visible_accelerator_ids_per_worker.append(all_resource_ids)
|
||||
|
||||
return visible_accelerator_ids_per_worker
|
||||
@@ -0,0 +1,33 @@
|
||||
import logging
|
||||
|
||||
from ray.exceptions import RayActorError, RayTaskError
|
||||
from ray.train.backend import BackendConfig
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReplicaGroupCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
ExecutionGroup,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BackendSetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
|
||||
def __init__(self, backend_config: BackendConfig):
|
||||
self._backend_config = backend_config
|
||||
self._backend = backend_config.backend_cls()
|
||||
|
||||
def after_execution_group_start(self, execution_group: ExecutionGroup):
|
||||
self._backend.on_start(execution_group, self._backend_config)
|
||||
self._backend.on_training_start(execution_group, self._backend_config)
|
||||
|
||||
def before_execution_group_shutdown(self, execution_group: ExecutionGroup):
|
||||
try:
|
||||
self._backend.on_shutdown(execution_group, self._backend_config)
|
||||
except (RayActorError, RayTaskError):
|
||||
logger.warning(
|
||||
"Graceful shutdown of backend failed. This is "
|
||||
"expected if one of the workers has crashed.",
|
||||
exc_info=True,
|
||||
)
|
||||
@@ -0,0 +1,175 @@
|
||||
import copy
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
import ray.train
|
||||
from ray.train.v2._internal.data_integration.interfaces import (
|
||||
DatasetShardMetadata,
|
||||
DatasetShardProvider,
|
||||
GenDataset,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.worker_group.worker_group import (
|
||||
Worker,
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
)
|
||||
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator, Dataset, NodeIdStr
|
||||
from ray.data.context import DataContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RayDatasetShardProvider:
|
||||
def __init__(
|
||||
self,
|
||||
datasets: Dict[str, GenDataset],
|
||||
data_config: ray.train.DataConfig,
|
||||
data_context: "DataContext",
|
||||
world_size: int,
|
||||
worker_node_ids: List["NodeIdStr"],
|
||||
):
|
||||
from ray.train.v2._internal.data_integration.dataset_manager import (
|
||||
DatasetManager,
|
||||
)
|
||||
|
||||
self._dataset_names = set(datasets)
|
||||
self._dataset_manager = (
|
||||
ray.remote(DatasetManager)
|
||||
.options(
|
||||
num_cpus=0,
|
||||
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
||||
ray.get_runtime_context().get_node_id(), soft=False
|
||||
),
|
||||
)
|
||||
.remote(
|
||||
datasets=datasets,
|
||||
data_config=data_config,
|
||||
data_context=data_context,
|
||||
world_size=world_size,
|
||||
worker_node_ids=worker_node_ids,
|
||||
)
|
||||
)
|
||||
self._cached_dataset_shards: Dict[str, "DataIterator"] = {}
|
||||
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
dataset_name = dataset_info.dataset_name
|
||||
if dataset_name not in self._dataset_names:
|
||||
raise KeyError(
|
||||
f"Dataset shard for '{dataset_name}' not found. "
|
||||
"Please ensure that the dataset is passed through the Trainer `datasets` "
|
||||
"argument."
|
||||
)
|
||||
|
||||
if dataset_name not in self._cached_dataset_shards:
|
||||
self._cached_dataset_shards[dataset_name] = ray.get(
|
||||
self._dataset_manager.get_dataset_shard.remote(dataset_info)
|
||||
)
|
||||
|
||||
return self._cached_dataset_shards[dataset_name]
|
||||
|
||||
def shutdown_data_executors(self) -> None:
|
||||
"""
|
||||
Attempts to eagerly shutdown the data executors for datasets, freeing resources allocated to data execution.
|
||||
"""
|
||||
try:
|
||||
self._dataset_manager.shutdown_data_executors.remote()
|
||||
except Exception:
|
||||
logger.debug("Failed to invoke remote cleanup of Dataset Manager.")
|
||||
|
||||
|
||||
class DatasetsCallback(WorkerGroupCallback):
|
||||
"""A callback for managing Ray Datasets for the worker group."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_run_context: TrainRunContext,
|
||||
datasets: Dict[str, "Dataset"],
|
||||
):
|
||||
self._datasets = datasets
|
||||
self._data_config = copy.deepcopy(train_run_context.dataset_config)
|
||||
self._scaling_config = train_run_context.scaling_config
|
||||
self._dataset_shard_provider: Optional[RayDatasetShardProvider] = None
|
||||
|
||||
# Capture the current DataContext to propagate it to
|
||||
# the Train workers later.
|
||||
# The propagation works in the following way:
|
||||
# 1. This callback is created when user create the Trainer.
|
||||
# 2. Then this callback will be passed to the Controller actor.
|
||||
# 3. Lastly, when the worker group is initialized, the Controller
|
||||
# will call the `after_worker_group_start` callback to propagate
|
||||
# the DataContext to Train workers.
|
||||
from ray.data.context import DataContext
|
||||
|
||||
self._data_context = copy.deepcopy(DataContext.get_current())
|
||||
|
||||
def get_train_total_resources(
|
||||
self, scaling_config: ray.train.ScalingConfig
|
||||
) -> Dict[str, float]:
|
||||
"""Return the resources reserved for training, so that Data can exclude
|
||||
these resources logically from its available pool."""
|
||||
if scaling_config.elasticity_enabled:
|
||||
# If Train is running with a variable number of workers,
|
||||
# we can't provide a fixed number of resources to exclude.
|
||||
# Instead, Train and Data should coordinate via the autoscaling
|
||||
# coordinator to allocate resources dynamically.
|
||||
return {}
|
||||
return scaling_config.total_resources
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List[Worker]
|
||||
) -> Dict[str, List[DatasetShardProvider]]:
|
||||
world_size = len(workers)
|
||||
worker_node_ids = [worker.metadata.node_id for worker in workers]
|
||||
datasets = {k: v() if callable(v) else v for k, v in self._datasets.items()}
|
||||
|
||||
# TODO: Move this to the constructor.
|
||||
# Notify the DataConfig about the total resources reserved for training.
|
||||
total_train_resources = self.get_train_total_resources(self._scaling_config)
|
||||
self._data_config.set_train_total_resources(
|
||||
total_train_resources.get("CPU", 0), total_train_resources.get("GPU", 0)
|
||||
)
|
||||
|
||||
self._dataset_shard_provider = RayDatasetShardProvider(
|
||||
datasets=datasets,
|
||||
data_config=self._data_config,
|
||||
data_context=self._data_context,
|
||||
world_size=world_size,
|
||||
worker_node_ids=worker_node_ids,
|
||||
)
|
||||
return {"dataset_shard_provider": [self._dataset_shard_provider] * world_size}
|
||||
|
||||
def after_worker_group_start(self, worker_group: WorkerGroup):
|
||||
# Propagate DataContext
|
||||
from ray.data.context import DataContext
|
||||
|
||||
def _propagate_data_context(ctx: "DataContext"):
|
||||
DataContext._set_current(ctx)
|
||||
|
||||
worker_group.execute(
|
||||
_propagate_data_context,
|
||||
self._data_context,
|
||||
)
|
||||
|
||||
def after_worker_group_shutdown(
|
||||
self, worker_group_context: WorkerGroupContext
|
||||
) -> None:
|
||||
shard_provider = self._dataset_shard_provider
|
||||
if shard_provider:
|
||||
shard_provider.shutdown_data_executors()
|
||||
|
||||
def after_worker_group_abort(
|
||||
self, worker_group_context: WorkerGroupContext
|
||||
) -> None:
|
||||
shard_provider = self._dataset_shard_provider
|
||||
if shard_provider:
|
||||
shard_provider.shutdown_data_executors()
|
||||
@@ -0,0 +1,39 @@
|
||||
import importlib
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
from ray.train.v2._internal.constants import RAY_TRAIN_CALLBACKS_ENV_VAR
|
||||
from ray.train.v2._internal.execution.callback import RayTrainCallback
|
||||
|
||||
|
||||
def _initialize_env_callbacks() -> List[RayTrainCallback]:
|
||||
"""Initialize callbacks from environment variable.
|
||||
|
||||
Returns:
|
||||
List of callbacks initialized from environment variable.
|
||||
"""
|
||||
callbacks = []
|
||||
callbacks_str = os.environ.get(RAY_TRAIN_CALLBACKS_ENV_VAR, "")
|
||||
if not callbacks_str:
|
||||
return callbacks
|
||||
|
||||
for callback_path in callbacks_str.split(","):
|
||||
callback_path = callback_path.strip()
|
||||
if not callback_path:
|
||||
continue
|
||||
|
||||
try:
|
||||
module_path, class_name = callback_path.rsplit(".", 1)
|
||||
module = importlib.import_module(module_path)
|
||||
callback_cls = getattr(module, class_name)
|
||||
if not issubclass(callback_cls, RayTrainCallback):
|
||||
raise TypeError(
|
||||
f"Callback class '{callback_path}' must be a subclass of "
|
||||
f"RayTrainCallback, got {type(callback_cls).__name__}"
|
||||
)
|
||||
callback = callback_cls()
|
||||
callbacks.append(callback)
|
||||
except (ImportError, AttributeError, ValueError, TypeError) as e:
|
||||
raise ValueError(f"Failed to import callback from '{callback_path}'") from e
|
||||
|
||||
return callbacks
|
||||
@@ -0,0 +1,124 @@
|
||||
from contextlib import contextmanager
|
||||
from typing import Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
TrainContextCallback,
|
||||
WorkerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext, get_train_context
|
||||
from ray.train.v2._internal.execution.controller.state import (
|
||||
TrainControllerState,
|
||||
TrainControllerStateType,
|
||||
)
|
||||
from ray.train.v2._internal.metrics.base import Metric
|
||||
from ray.train.v2._internal.metrics.controller import ControllerMetrics
|
||||
from ray.train.v2._internal.metrics.worker import WorkerMetrics
|
||||
from ray.train.v2._internal.util import time_monotonic
|
||||
|
||||
|
||||
class ControllerMetricsCallback(ControllerCallback, WorkerGroupCallback):
|
||||
"""Callback that records controller-specific metrics."""
|
||||
|
||||
def after_controller_start(self, train_run_context: TrainRunContext):
|
||||
"""Initialize metrics after controller starts."""
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
self._metrics: Dict[str, Metric] = ControllerMetrics.get_controller_metrics(
|
||||
self._run_name, self._run_id
|
||||
)
|
||||
# Record initial state
|
||||
self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
|
||||
TrainControllerStateType.INITIALIZING
|
||||
)
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
"""Shutdown metrics before controller shuts down."""
|
||||
for metric in self._metrics.values():
|
||||
metric.reset()
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: TrainControllerState,
|
||||
current_state: TrainControllerState,
|
||||
):
|
||||
"""Record state transitions after controller state updates."""
|
||||
self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
|
||||
current_state._state_type
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def on_worker_group_start(self):
|
||||
"""Measure time taken to start worker group."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def on_worker_group_shutdown(self):
|
||||
"""Measure time taken to shutdown worker group."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
|
||||
|
||||
class WorkerMetricsCallback(WorkerCallback, TrainContextCallback):
|
||||
"""Callback that records worker-specific metrics."""
|
||||
|
||||
def __init__(self, train_run_context: TrainRunContext):
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
self._metrics: Optional[Dict[str, Metric]] = None
|
||||
|
||||
def after_init_train_context(self):
|
||||
"""Initialize metrics after train context is initialized."""
|
||||
train_context = get_train_context()
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
world_rank = train_context.get_world_rank()
|
||||
worker_actor_id = core_context.get_actor_id()
|
||||
self._metrics = WorkerMetrics.get_worker_metrics(
|
||||
self._run_name, self._run_id, world_rank, worker_actor_id
|
||||
)
|
||||
|
||||
def before_worker_shutdown(self):
|
||||
"""Shutdown metrics before shutdown."""
|
||||
if self._metrics:
|
||||
for metric in self._metrics.values():
|
||||
metric.reset()
|
||||
|
||||
@contextmanager
|
||||
def on_report(self):
|
||||
"""
|
||||
Context manager to measure the time taken to report a checkpoint to the storage.
|
||||
"""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].record(elapsed_time_s)
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_sync(self):
|
||||
"""Measure time spent in the cross-rank barrier that synchronizes the
|
||||
checkpoint directory name across all workers."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.CHECKPOINT_SYNC_TOTAL_TIME_S].record(elapsed_time_s)
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_transfer(self):
|
||||
"""Measure time spent transferring checkpoint files to storage."""
|
||||
start_time_s = time_monotonic()
|
||||
yield
|
||||
elapsed_time_s = time_monotonic() - start_time_s
|
||||
self._metrics[WorkerMetrics.CHECKPOINT_TRANSFER_TOTAL_TIME_S].record(
|
||||
elapsed_time_s
|
||||
)
|
||||
@@ -0,0 +1,131 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import ray
|
||||
from ray.exceptions import RayActorError
|
||||
from ray.train.v2._internal.constants import GET_ACTOR_TIMEOUT_S
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.controller.placement_group_cleaner import (
|
||||
PlacementGroupCleaner,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.worker_group import WorkerGroup
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PlacementGroupCleanerCallback(ControllerCallback, WorkerGroupCallback):
|
||||
"""Callback that manages a PlacementGroupCleaner for the training controller.
|
||||
|
||||
This callback ensures that placement groups are cleaned up even if the controller
|
||||
dies ungracefully.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
check_interval_s: float = 1.0,
|
||||
get_actor_timeout_s: float = GET_ACTOR_TIMEOUT_S,
|
||||
stop_timeout: Optional[float] = None,
|
||||
):
|
||||
"""Initialize the callback.
|
||||
|
||||
Args:
|
||||
check_interval_s: How often (in seconds) the cleaner should check
|
||||
if the controller is still alive.
|
||||
get_actor_timeout_s: How long to wait when calling the get actor state api.
|
||||
stop_timeout: How long to wait for the cleaner to stop.
|
||||
"""
|
||||
self._check_interval_s = check_interval_s
|
||||
self._get_actor_timeout_s = get_actor_timeout_s
|
||||
self._stop_timeout = stop_timeout
|
||||
if self._stop_timeout is None:
|
||||
self._stop_timeout = max(
|
||||
2.0, self._check_interval_s * 2 + self._get_actor_timeout_s
|
||||
)
|
||||
self._cleaner: Optional[PlacementGroupCleaner] = None
|
||||
self._controller_actor_id: Optional[str] = None
|
||||
|
||||
def after_controller_start(self, train_run_context: "TrainRunContext"):
|
||||
"""Launch the detached PlacementGroupCleaner actor and start monitoring."""
|
||||
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
self._controller_actor_id = core_context.get_actor_id()
|
||||
try:
|
||||
# Launch the cleaner as a detached actor so it survives controller death
|
||||
cleaner_actor_cls = ray.remote(num_cpus=0)(PlacementGroupCleaner)
|
||||
self._cleaner = cleaner_actor_cls.options(
|
||||
lifetime="detached",
|
||||
get_if_exists=False,
|
||||
).remote(
|
||||
controller_actor_id=self._controller_actor_id,
|
||||
check_interval_s=self._check_interval_s,
|
||||
get_actor_timeout_s=self._get_actor_timeout_s,
|
||||
stop_timeout=self._stop_timeout,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"PlacementGroupCleaner launched for run_id={train_run_context.run_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to launch PlacementGroupCleaner: {e}. "
|
||||
"Placement groups may not be cleaned up if controller exits ungracefully."
|
||||
)
|
||||
self._cleaner = None
|
||||
return
|
||||
|
||||
self._cleaner.start_monitoring.remote()
|
||||
|
||||
def after_worker_group_start(self, worker_group: "WorkerGroup"):
|
||||
"""Register the worker group's placement group with the cleaner.
|
||||
|
||||
This is called after a worker group is successfully started.
|
||||
"""
|
||||
if not self._cleaner or not self._controller_actor_id:
|
||||
logger.warning(
|
||||
"PlacementGroupCleaner not available. "
|
||||
"Placement groups may not be cleaned up if controller exits ungracefully."
|
||||
)
|
||||
return
|
||||
worker_group_state = worker_group.get_worker_group_state()
|
||||
placement_group = worker_group_state.placement_group_handle.placement_group
|
||||
|
||||
try:
|
||||
ray.get(self._cleaner.register_placement_group.remote(placement_group))
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to register placement group with cleaner: {e}. "
|
||||
"Placement group may not be cleaned up if controller dies ungracefully."
|
||||
)
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Registered placement group {placement_group.id} with PlacementGroupCleaner."
|
||||
)
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
self._stop_cleaner()
|
||||
|
||||
def before_controller_abort(self):
|
||||
self._stop_cleaner()
|
||||
|
||||
def _stop_cleaner(self):
|
||||
if not self._cleaner:
|
||||
return
|
||||
|
||||
try:
|
||||
# Stop the cleaner gracefully (it won't clean up the PG)
|
||||
ray.get(self._cleaner.stop.remote(), timeout=self._stop_timeout)
|
||||
except RayActorError:
|
||||
logger.debug(
|
||||
"PlacementGroupCleaner exited before stop completed; ignoring."
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to stop PlacementGroupCleaner gracefully.")
|
||||
finally:
|
||||
self._cleaner = None
|
||||
@@ -0,0 +1,93 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_PREEMPTION_POLL_INTERVAL_S,
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import WorkerGroupCallback
|
||||
from ray.train.v2._internal.execution.preemption import PreemptionWatcher
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PreemptionCallback(WorkerGroupCallback):
|
||||
"""Manages a :class:`PreemptionWatcher` across worker-group lifecycles.
|
||||
|
||||
Spawns a fresh watcher in :meth:`after_worker_group_start` and stops it on
|
||||
every teardown path (shutdown and abort). Each worker group gets its own
|
||||
watcher and failure-domain map, so elastic resizes and restarts never
|
||||
leak stale state.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._poll_interval_s: float = float(
|
||||
os.getenv(
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
||||
str(DEFAULT_PREEMPTION_POLL_INTERVAL_S),
|
||||
)
|
||||
)
|
||||
self._watcher: Optional[ActorHandle] = None
|
||||
|
||||
def after_worker_group_start(self, worker_group: "WorkerGroup") -> None:
|
||||
# Tear down any watcher from a previous worker group first. Worker-group
|
||||
# startup can fail after this hook without running the shutdown hook, so
|
||||
# this also prevents leaking an orphaned watcher across a reschedule.
|
||||
self._stop_watcher()
|
||||
|
||||
# These handles are captured once per worker-group start. With the
|
||||
# standard backend, any worker replacement goes through a full worker
|
||||
# group restart (this hook runs again with fresh handles), so they
|
||||
# never go stale.
|
||||
# TODO(lehui): refresh worker handles on in-place replica replacement
|
||||
# when adding preemption support for replica groups (TorchFT).
|
||||
node_to_ranks: Dict[str, List[int]] = {}
|
||||
worker_actors_by_rank: Dict[int, ActorHandle] = {}
|
||||
for w in worker_group.get_workers():
|
||||
rank = w.distributed_context.world_rank
|
||||
node_to_ranks.setdefault(w.metadata.node_id, []).append(rank)
|
||||
worker_actors_by_rank[rank] = w.actor
|
||||
|
||||
watcher_cls = ray.remote(num_cpus=0, max_restarts=-1)(PreemptionWatcher)
|
||||
self._watcher = watcher_cls.remote(
|
||||
node_to_ranks=node_to_ranks,
|
||||
poll_interval_s=self._poll_interval_s,
|
||||
worker_actors_by_rank=worker_actors_by_rank,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"PreemptionCallback: started watcher for %d node(s).",
|
||||
len(node_to_ranks),
|
||||
)
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: "WorkerGroup") -> None:
|
||||
self._stop_watcher()
|
||||
|
||||
def after_worker_group_abort(
|
||||
self, worker_group_context: "WorkerGroupContext"
|
||||
) -> None:
|
||||
# abort() doesn't run the shutdown hook, so tear the watcher down here
|
||||
# too — otherwise it keeps polling GCS until the cluster reaps it.
|
||||
self._stop_watcher()
|
||||
|
||||
def _stop_watcher(self) -> None:
|
||||
if self._watcher is None:
|
||||
return
|
||||
watcher = self._watcher
|
||||
self._watcher = None
|
||||
# Force-kill (non-blocking) rather than a synchronous graceful stop, so
|
||||
# we never block the controller's event loop. The watcher's daemon poll
|
||||
# thread dies with the actor process and holds no external resources.
|
||||
try:
|
||||
ray.kill(watcher)
|
||||
except Exception:
|
||||
logger.warning("Failed to kill PreemptionWatcher actor.", exc_info=True)
|
||||
@@ -0,0 +1,221 @@
|
||||
import importlib
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import Dataset
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller.state import (
|
||||
AbortedState,
|
||||
ErroredState,
|
||||
FinishedState,
|
||||
ReschedulingState,
|
||||
ResizingState,
|
||||
RestartingState,
|
||||
RunningState,
|
||||
SchedulingState,
|
||||
ShuttingDownState,
|
||||
TrainControllerState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
ResizeDecision,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
WorkerGroupState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group.poll import WorkerGroupPollStatus
|
||||
from ray.train.v2._internal.logging.logging import (
|
||||
get_train_application_controller_log_path,
|
||||
)
|
||||
from ray.train.v2._internal.state.state_manager import TrainStateManager
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_framework_version(framework: Optional[TrainingFramework]):
|
||||
versions = {}
|
||||
|
||||
try:
|
||||
import ray
|
||||
|
||||
versions["ray"] = ray.__version__
|
||||
except ImportError:
|
||||
logger.warning("Failed to collect ray version on worker.")
|
||||
|
||||
if framework is None:
|
||||
return versions
|
||||
|
||||
for module_name in framework.module_names():
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
versions[module_name] = module.__version__
|
||||
except ModuleNotFoundError:
|
||||
# Module is not installed, skip without recording a version.
|
||||
continue
|
||||
except Exception:
|
||||
logger.warning(f"Failed to collect {module_name} version on worker.")
|
||||
continue
|
||||
|
||||
return versions
|
||||
|
||||
|
||||
class StateManagerCallback(ControllerCallback, WorkerGroupCallback):
|
||||
def __init__(self, datasets: Dict[str, "Dataset"]):
|
||||
self._datasets = datasets
|
||||
|
||||
def after_controller_start(self, train_run_context: TrainRunContext):
|
||||
self._state_manager = TrainStateManager()
|
||||
self._run_name = train_run_context.get_run_config().name
|
||||
self._run_id = train_run_context.run_id
|
||||
|
||||
# TODO: Should this be generated by the caller?
|
||||
# NOTE: These must be called on the Controller.
|
||||
# The Callback is first initialized on the Driver.
|
||||
core_context = ray.runtime_context.get_runtime_context()
|
||||
self._job_id = core_context.get_job_id()
|
||||
self._controller_actor_id = core_context.get_actor_id()
|
||||
controller_log_file_path = get_train_application_controller_log_path()
|
||||
self._state_manager.create_train_run(
|
||||
id=self._run_id,
|
||||
name=self._run_name,
|
||||
job_id=self._job_id,
|
||||
controller_actor_id=self._controller_actor_id,
|
||||
controller_log_file_path=controller_log_file_path,
|
||||
run_config=train_run_context.run_config,
|
||||
train_loop_config=train_run_context.train_loop_config,
|
||||
scaling_config=train_run_context.scaling_config,
|
||||
backend_config=train_run_context.backend_config,
|
||||
datasets=self._datasets,
|
||||
dataset_config=train_run_context.dataset_config,
|
||||
)
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: TrainControllerState,
|
||||
current_state: TrainControllerState,
|
||||
):
|
||||
if previous_state._state_type == current_state._state_type:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"[State Transition] {previous_state._state_type.state_name} -> "
|
||||
f"{current_state._state_type.state_name}."
|
||||
)
|
||||
|
||||
if isinstance(current_state, SchedulingState):
|
||||
# TODO: This should probably always be ResizeDecision.
|
||||
if isinstance(current_state.scaling_decision, ResizeDecision):
|
||||
resize_decision = current_state.scaling_decision
|
||||
else:
|
||||
resize_decision = None
|
||||
|
||||
self._state_manager.update_train_run_scheduling(
|
||||
run_id=self._run_id,
|
||||
resize_decision=resize_decision,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, RunningState):
|
||||
self._state_manager.update_train_run_running(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, RestartingState):
|
||||
self._state_manager.update_train_run_restarting(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ResizingState):
|
||||
self._state_manager.update_train_run_resizing(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ErroredState):
|
||||
self._state_manager.update_train_run_errored(
|
||||
run_id=self._run_id,
|
||||
status_detail=str(current_state.training_failed_error),
|
||||
)
|
||||
|
||||
elif isinstance(current_state, FinishedState):
|
||||
self._state_manager.update_train_run_finished(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, AbortedState):
|
||||
self._state_manager.update_train_run_aborted(
|
||||
run_id=self._run_id,
|
||||
)
|
||||
|
||||
elif isinstance(current_state, ReschedulingState):
|
||||
# substate of SchedulingState
|
||||
pass
|
||||
|
||||
elif isinstance(current_state, ShuttingDownState):
|
||||
# substate of RunningState
|
||||
pass
|
||||
|
||||
def before_worker_group_start(self, worker_group_context: WorkerGroupContext):
|
||||
self._state_manager.create_train_run_attempt(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
num_workers=worker_group_context.num_workers,
|
||||
resources_per_worker=worker_group_context.resources_per_worker,
|
||||
)
|
||||
|
||||
def after_worker_group_start(self, worker_group: WorkerGroup):
|
||||
worker_group_context: WorkerGroupContext = (
|
||||
worker_group.get_worker_group_context()
|
||||
)
|
||||
worker_group_state: WorkerGroupState = worker_group.get_worker_group_state()
|
||||
self._state_manager.update_train_run_attempt_running(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
workers=worker_group_state.workers,
|
||||
)
|
||||
|
||||
# Update train run framework version
|
||||
framework = self._state_manager.get_train_run_framework(self._run_id)
|
||||
framework_versions = worker_group.execute_single(
|
||||
0, _get_framework_version, framework
|
||||
)
|
||||
self._state_manager.update_train_run_framework_versions(
|
||||
run_id=self._run_id,
|
||||
framework_versions=framework_versions,
|
||||
)
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: WorkerGroup):
|
||||
worker_group_context: WorkerGroupContext = (
|
||||
worker_group.get_worker_group_context()
|
||||
)
|
||||
# TODO: Consider passing error reason directly to the callback.
|
||||
# Something along the lines of:
|
||||
# WorkerGroup.shutdown(reason)
|
||||
# -> WorkerGroupCallback.before_worker_group_shutdown(reason)
|
||||
worker_group_poll_status: Optional[
|
||||
WorkerGroupPollStatus
|
||||
] = worker_group.get_latest_poll_status()
|
||||
if worker_group_poll_status and worker_group_poll_status.errors:
|
||||
self._state_manager.update_train_run_attempt_errored(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
status_detail=worker_group_poll_status.get_error_string(),
|
||||
)
|
||||
else:
|
||||
self._state_manager.update_train_run_attempt_finished(
|
||||
run_id=self._run_id,
|
||||
attempt_id=worker_group_context.run_attempt_id,
|
||||
)
|
||||
|
||||
def before_worker_group_abort(self, worker_group_context: WorkerGroupContext):
|
||||
self._state_manager.update_train_run_attempt_aborted(
|
||||
self._run_id,
|
||||
worker_group_context.run_attempt_id,
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReportCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2._internal.execution.worker_group import WorkerGroupPollStatus
|
||||
from ray.train.v2.api.callback import UserCallback
|
||||
|
||||
|
||||
class UserCallbackHandler(WorkerGroupCallback, ReportCallback):
|
||||
"""Responsible for calling methods of subscribers implementing
|
||||
the `UserCallback` interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, user_callbacks: List[UserCallback], train_run_context: TrainRunContext
|
||||
):
|
||||
self._user_callbacks = user_callbacks
|
||||
self._train_run_context = train_run_context
|
||||
|
||||
# --------------------------
|
||||
# ReportCallback
|
||||
# --------------------------
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
metrics: List[Dict[str, Any]],
|
||||
):
|
||||
for user_callback in self._user_callbacks:
|
||||
user_callback.after_report(
|
||||
run_context=self._train_run_context,
|
||||
metrics=metrics,
|
||||
checkpoint=training_report.checkpoint,
|
||||
)
|
||||
|
||||
# --------------------------
|
||||
# WorkerGroupCallback
|
||||
# --------------------------
|
||||
|
||||
def after_worker_group_poll_status(
|
||||
self, worker_group_status: WorkerGroupPollStatus
|
||||
):
|
||||
if not worker_group_status.errors:
|
||||
return
|
||||
|
||||
for user_callback in self._user_callbacks:
|
||||
user_callback.after_exception(
|
||||
run_context=self._train_run_context,
|
||||
worker_exceptions=worker_group_status.errors,
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ReplicaGroupCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import get_train_context
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
ExecutionGroup,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WorkingDirectorySetupCallback(ReplicaGroupCallback, WorkerGroupCallback):
|
||||
def after_execution_group_start(self, execution_group: ExecutionGroup):
|
||||
"""Shared logic for setting up the working directory on an execution group."""
|
||||
|
||||
def chdir_to_working_dir() -> None:
|
||||
"""Create the local working directory for the experiment."""
|
||||
local_working_directory = (
|
||||
get_train_context().get_storage().local_working_directory
|
||||
)
|
||||
os.makedirs(local_working_directory, exist_ok=True)
|
||||
logger.debug(
|
||||
f"Changing the working directory to: {local_working_directory}"
|
||||
)
|
||||
os.chdir(local_working_directory)
|
||||
|
||||
execution_group.execute(chdir_to_working_dir)
|
||||
@@ -0,0 +1,159 @@
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
|
||||
from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
|
||||
from ray._private.ray_constants import env_bool, env_set_by_user
|
||||
|
||||
# Unsupported configs can use this value to detect if the user has set it.
|
||||
_UNSUPPORTED = "UNSUPPORTED"
|
||||
_DEPRECATED = "DEPRECATED"
|
||||
|
||||
# The name of the file that is used to validate the storage.
|
||||
VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker"
|
||||
# The name of the file that is used to store the checkpoint manager snapshot.
|
||||
CHECKPOINT_MANAGER_SNAPSHOT_FILENAME = "checkpoint_manager_snapshot.json"
|
||||
|
||||
AWS_RETRYABLE_TOKENS = (
|
||||
"AWS Error SLOW_DOWN",
|
||||
"AWS Error INTERNAL_FAILURE",
|
||||
"AWS Error SERVICE_UNAVAILABLE",
|
||||
"AWS Error NETWORK_CONNECTION",
|
||||
"AWS Error UNKNOWN",
|
||||
)
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Environment variables used in the controller, workers, and state actor.
|
||||
#
|
||||
# Be sure to update ENV_VARS_TO_PROPAGATE when adding new
|
||||
# environment variables in this section.
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
# Polling interval for the Train controller.
|
||||
# This determines how many seconds the controller will wait between
|
||||
# polling the worker group for its status.
|
||||
HEALTH_CHECK_INTERVAL_S_ENV_VAR = "RAY_TRAIN_HEALTH_CHECK_INTERVAL_S"
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL_S: float = 2.0
|
||||
|
||||
# The time in seconds a worker health check must be hanging for
|
||||
# before the controller marks the worker as dead and handles the failure.
|
||||
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_HEALTH_CHECK_TIMEOUT_S"
|
||||
DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S: float = 10 * 60
|
||||
|
||||
# Timeout in seconds for the worker group to start.
|
||||
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_GROUP_START_TIMEOUT_S"
|
||||
DEFAULT_WORKER_GROUP_START_TIMEOUT_S: float = 60.0
|
||||
|
||||
# Time in seconds for collective operations before raising a timeout error.
|
||||
COLLECTIVE_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_TIMEOUT_S"
|
||||
# NOTE: Default to no timeout to avoid introducing more timeouts for users to configure.
|
||||
# For example, users can already configure timeouts in torch distributed.
|
||||
DEFAULT_COLLECTIVE_TIMEOUT_S: Optional[float] = None
|
||||
# Interval in seconds to log a warning when waiting for a collective operation to complete.
|
||||
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_WARN_INTERVAL_S"
|
||||
DEFAULT_COLLECTIVE_WARN_INTERVAL_S: float = 60
|
||||
|
||||
# Interval in seconds to log a warning when waiting for a checkpoint upload fn operation to complete.
|
||||
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR = (
|
||||
"RAY_TRAIN_CHECKPOINT_UPLOAD_WARN_INTERVAL_S"
|
||||
)
|
||||
DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S: float = 60
|
||||
|
||||
# Feature flag for the preemption watcher. Default-on; provides a quick
|
||||
# rollback path if the watcher actor misbehaves in a cluster.
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR = "RAY_TRAIN_ENABLE_PREEMPTION_WATCHER"
|
||||
DEFAULT_ENABLE_PREEMPTION_WATCHER: bool = True
|
||||
|
||||
# How often the preemption watcher polls Ray Core's drain state.
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR = "RAY_TRAIN_PREEMPTION_POLL_INTERVAL_S"
|
||||
DEFAULT_PREEMPTION_POLL_INTERVAL_S: float = 5.0
|
||||
|
||||
# Environment variable to enable the print function patching.
|
||||
ENABLE_PRINT_PATCH_ENV_VAR = "RAY_TRAIN_ENABLE_PRINT_PATCH"
|
||||
DEFAULT_ENABLE_PRINT_PATCH = True
|
||||
|
||||
# V2 feature flag.
|
||||
V2_ENABLED_ENV_VAR = "RAY_TRAIN_V2_ENABLED"
|
||||
|
||||
# Environment variables to enable/disable controller/worker structured logging.
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR = (
|
||||
"RAY_TRAIN_ENABLE_CONTROLLER_STRUCTURED_LOGGING"
|
||||
)
|
||||
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR = "RAY_TRAIN_ENABLE_WORKER_STRUCTURED_LOGGING"
|
||||
DEFAULT_ENABLE_CONTROLLER_LOGGING = True
|
||||
DEFAULT_ENABLE_WORKER_LOGGING = True
|
||||
|
||||
# Environment variables to configure reconciliation interval for Train state actor.
|
||||
# This determines how many seconds the state actor will wait between
|
||||
# polling the controller for its status.
|
||||
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR = (
|
||||
"RAY_TRAIN_ENABLE_STATE_ACTOR_RECONCILIATION"
|
||||
)
|
||||
DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION = True
|
||||
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR = (
|
||||
"RAY_TRAIN_STATE_ACTOR_RECONCILIATION_INTERVAL_S"
|
||||
)
|
||||
DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S: float = 30.0
|
||||
# TODO: `ray.util.state.api.get_actor` typically takes 10-50ms but can take longer
|
||||
# when there is high load on the cluster.
|
||||
GET_ACTOR_TIMEOUT_S: int = 10
|
||||
# GET_ACTOR_TIMEOUT_S * CONTROLLERS_TO_POLL_PER_ITERATION should be
|
||||
# way less than STATE_ACTOR_RECONCILIATION_INTERVAL_S to give the state actor
|
||||
# time to update live train run state.
|
||||
CONTROLLERS_TO_POLL_PER_ITERATION: int = 1
|
||||
|
||||
# Environment variable for Train execution callbacks
|
||||
RAY_TRAIN_CALLBACKS_ENV_VAR = "RAY_TRAIN_CALLBACKS"
|
||||
|
||||
# Ray Train does not warn by default when using blocking ray.get inside async actor.
|
||||
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE = "0"
|
||||
|
||||
# torchft lighthouse address
|
||||
TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR = "TORCHFT_LIGHTHOUSE"
|
||||
|
||||
# Environment variables to propagate from the driver to the controller,
|
||||
# and then from the controller to the workers.
|
||||
ENV_VARS_TO_PROPAGATE = {
|
||||
V2_ENABLED_ENV_VAR,
|
||||
HEALTH_CHECK_INTERVAL_S_ENV_VAR,
|
||||
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR,
|
||||
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
|
||||
COLLECTIVE_TIMEOUT_S_ENV_VAR,
|
||||
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
||||
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
|
||||
ENABLE_PRINT_PATCH_ENV_VAR,
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
|
||||
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
|
||||
RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
|
||||
TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR,
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
||||
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
||||
}
|
||||
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# Environment variables used in the driver script only.
|
||||
# ------------------------------------------------------------
|
||||
|
||||
# The environment variable to enable the Ray Train Metrics.
|
||||
METRICS_ENABLED_ENV_VAR = "RAY_TRAIN_METRICS_ENABLED"
|
||||
|
||||
|
||||
def is_v2_enabled() -> bool:
|
||||
return env_bool(V2_ENABLED_ENV_VAR, True)
|
||||
|
||||
|
||||
def get_env_vars_to_propagate() -> Dict[str, str]:
|
||||
"""Returns a dictionary of environment variables that should be propagated
|
||||
from the driver to the controller, and then from the controller
|
||||
to each training worker.
|
||||
|
||||
This way, users only need to set environment variables in one place
|
||||
when launching the script instead of needing to manually set a runtime environment.
|
||||
"""
|
||||
env_vars = {}
|
||||
for env_var in ENV_VARS_TO_PROPAGATE:
|
||||
if env_set_by_user(env_var):
|
||||
env_vars[env_var] = os.environ[env_var]
|
||||
return env_vars
|
||||
@@ -0,0 +1,174 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, List
|
||||
|
||||
import ray
|
||||
from ray.exceptions import GetTimeoutError
|
||||
from ray.train.v2._internal.data_integration.interfaces import (
|
||||
DatasetShardMetadata,
|
||||
GenDataset,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataContext, DataIterator, Dataset, NodeIdStr
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatasetManager:
|
||||
"""Manages the dataset shards for datasets configured in the trainer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
datasets: Dict[str, GenDataset],
|
||||
data_config: ray.train.DataConfig,
|
||||
data_context: "DataContext",
|
||||
world_size: int,
|
||||
worker_node_ids: List["NodeIdStr"],
|
||||
):
|
||||
self._datasets = datasets
|
||||
self._data_config = data_config
|
||||
self._datasets_to_split = (
|
||||
set(self._datasets.keys())
|
||||
if data_config._datasets_to_split == "all"
|
||||
else set(data_config._datasets_to_split)
|
||||
)
|
||||
self._world_size = world_size
|
||||
self._worker_node_ids = worker_node_ids
|
||||
self._coordinator_actors: List[ray.actor.ActorHandle] = []
|
||||
|
||||
# Maps dataset name to a list of cached `DataIterator`s corresponding to
|
||||
# Train worker ranks.
|
||||
self._dataset_iterators: Dict[str, List["DataIterator"]] = {}
|
||||
|
||||
# A condition variable to synchronize the calls to the async `get_dataset_shard` method.
|
||||
self._condition = asyncio.Condition()
|
||||
|
||||
from ray.data import DataContext
|
||||
|
||||
DataContext._set_current(data_context)
|
||||
|
||||
def _create_dataset_iterators(
|
||||
self, dataset_info: DatasetShardMetadata, base_dataset: "Dataset"
|
||||
) -> List["DataIterator"]:
|
||||
dataset_name = dataset_info.dataset_name
|
||||
|
||||
iterators_per_rank = self._data_config.configure(
|
||||
datasets={dataset_name: base_dataset},
|
||||
world_size=self._world_size,
|
||||
worker_handles=None,
|
||||
worker_node_ids=self._worker_node_ids,
|
||||
)
|
||||
assert len(iterators_per_rank) == self._world_size
|
||||
# Convert the List[Dict[str, DataIterator]] to a List[DataIterator],
|
||||
# since we only configured one dataset.
|
||||
return [iterators_per_rank[i][dataset_name] for i in range(self._world_size)]
|
||||
|
||||
def _get_unsharded_dataset_iterator(
|
||||
self, dataset_info: DatasetShardMetadata
|
||||
) -> "DataIterator":
|
||||
"""Returns the dataset iterator for a dataset that is excluded
|
||||
from `DataConfig.datasets_to_split`.
|
||||
Note that this method is NOT a barrier across workers and can be called
|
||||
by any subset of workers and will return immediately.
|
||||
"""
|
||||
dataset_name = dataset_info.dataset_name
|
||||
world_rank = dataset_info.world_rank
|
||||
|
||||
if dataset_name not in self._dataset_iterators:
|
||||
self._dataset_iterators[dataset_name] = self._create_dataset_iterators(
|
||||
dataset_info, self._datasets[dataset_name]
|
||||
)
|
||||
|
||||
return self._dataset_iterators[dataset_name][world_rank]
|
||||
|
||||
async def _get_sharded_dataset_iterator(
|
||||
self, dataset_info: DatasetShardMetadata
|
||||
) -> "DataIterator":
|
||||
"""Returns the dataset iterator for a dataset that is included
|
||||
in `DataConfig.datasets_to_split`.
|
||||
Note that this method is a barrier across workers,
|
||||
and all workers must call this method before training.
|
||||
"""
|
||||
dataset_name = dataset_info.dataset_name
|
||||
world_rank = dataset_info.world_rank
|
||||
|
||||
async with self._condition:
|
||||
if dataset_name in self._dataset_iterators:
|
||||
# If the dataset iterators have already been created, return the
|
||||
# existing one.
|
||||
iterator = self._dataset_iterators[dataset_name][world_rank]
|
||||
elif world_rank == 0:
|
||||
# In this case, the dataset iterators have not been created yet.
|
||||
# The dataset only needs to be configured once globally for all workers.
|
||||
# Do it only when the rank 0 worker calls this method.
|
||||
iterators = self._create_dataset_iterators(
|
||||
dataset_info, self._datasets[dataset_name]
|
||||
)
|
||||
iterator = iterators[world_rank]
|
||||
|
||||
# Cache the split coordinators for resource cleanup.
|
||||
from ray.data._internal.iterator.stream_split_iterator import (
|
||||
StreamSplitDataIterator,
|
||||
)
|
||||
|
||||
if isinstance(iterator, StreamSplitDataIterator):
|
||||
self._coordinator_actors.append(iterator._coord_actor)
|
||||
|
||||
# Cache the dataset iterators for future use.
|
||||
self._dataset_iterators[dataset_name] = iterators
|
||||
self._condition.notify_all()
|
||||
else:
|
||||
# Wait for the dataset iterators to be created by the rank 0 worker.
|
||||
await self._condition.wait_for(
|
||||
lambda: dataset_name in self._dataset_iterators
|
||||
)
|
||||
iterator = self._dataset_iterators[dataset_name][world_rank]
|
||||
return iterator
|
||||
|
||||
async def get_dataset_shard(
|
||||
self,
|
||||
dataset_info: DatasetShardMetadata,
|
||||
) -> "DataIterator":
|
||||
"""Create and return the dataset shard iterator for a Ray Train worker's
|
||||
call to `ray.train.get_dataset_shard`.
|
||||
|
||||
This method is a barrier that should be called by all Ray Train workers at once.
|
||||
If the dataset iterators have already been created, return the existing ones.
|
||||
|
||||
Otherwise, create the dataset iterators and cache them.
|
||||
Here's an example of how this method is used with 4 workers:
|
||||
Rank 2 calls get_dataset_shard, waits on the condition variable.
|
||||
Rank 1 calls get_dataset_shard, waits on the condition variable.
|
||||
Rank 0 calls get_dataset_shard, creates the dataset iterators, caches them,
|
||||
and notifies all workers hanging on the condition variable.
|
||||
Rank 3 calls get_dataset_shard, returns the cached iterator.
|
||||
"""
|
||||
dataset_name = dataset_info.dataset_name
|
||||
|
||||
if dataset_name in self._datasets_to_split:
|
||||
return await self._get_sharded_dataset_iterator(dataset_info)
|
||||
else:
|
||||
return self._get_unsharded_dataset_iterator(dataset_info)
|
||||
|
||||
def shutdown_data_executors(self) -> None:
|
||||
"""
|
||||
Attempts to shut down the data executors of each sharded dataset,
|
||||
freeing resources allocated to data execution.
|
||||
|
||||
Note: The data executors for unsharded datasets are not managed by
|
||||
SplitCoordinator actors and hence, are not accessible via the DatasetManager
|
||||
so their cleanup is not handled by this method.
|
||||
"""
|
||||
try:
|
||||
shutdown_refs = [
|
||||
coord.shutdown_executor.remote() for coord in self._coordinator_actors
|
||||
]
|
||||
ray.get(shutdown_refs, timeout=5)
|
||||
except GetTimeoutError:
|
||||
logger.error("Ray Data executor shutdown task timed out after 5 seconds.")
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to gracefully terminate the Ray Data executor for each running dataset."
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Callable, Protocol, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator, Dataset
|
||||
|
||||
# A type representing either a ray.data.Dataset or a function that returns a
|
||||
# ray.data.Dataset and accepts no arguments.
|
||||
GenDataset = Union["Dataset", Callable[[], "Dataset"]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatasetShardMetadata:
|
||||
"""Metadata about a dataset shard used for lookup and configuration."""
|
||||
|
||||
dataset_name: str
|
||||
world_rank: int
|
||||
|
||||
|
||||
class DatasetShardProvider(Protocol):
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
"""Get the dataset shard for the given dataset info.
|
||||
Args:
|
||||
dataset_info: The metadata of the shard to retrieve,
|
||||
including the dataset name.
|
||||
Returns:
|
||||
The :class:`~ray.data.DataIterator` shard for the given dataset info.
|
||||
Raises:
|
||||
KeyError: If the dataset shard for the given dataset info is not found.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,165 @@
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
from ray.train.v2._internal.constants import (
|
||||
COLLECTIVE_TIMEOUT_S_ENV_VAR,
|
||||
DEFAULT_WORKER_GROUP_START_TIMEOUT_S,
|
||||
DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S,
|
||||
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
|
||||
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR,
|
||||
)
|
||||
|
||||
|
||||
# TODO: Distinguish between user and system exceptions.
|
||||
class RayTrainError(Exception):
|
||||
"""Base class for all Ray Train exceptions."""
|
||||
|
||||
|
||||
class WorkerHealthCheckTimeoutError(RayTrainError):
|
||||
"""Exception raised when a worker health check hangs for long enough."""
|
||||
|
||||
def __init__(self, message):
|
||||
timeout = os.getenv(
|
||||
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR, DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S
|
||||
)
|
||||
message += (
|
||||
f"\nSet the {WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR} "
|
||||
"environment variable to increase the timeout "
|
||||
f"(current value: {timeout} seconds)."
|
||||
)
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class WorkerHealthCheckFailedError(RayTrainError):
|
||||
"""Exception raised when a worker health check fails."""
|
||||
|
||||
def __init__(self, message, failure: Exception):
|
||||
super().__init__(message)
|
||||
self._message = message
|
||||
self.health_check_failure = failure
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self._message, self.health_check_failure))
|
||||
|
||||
def __str__(self):
|
||||
return self._message + "\n" + str(self.health_check_failure)
|
||||
|
||||
|
||||
class WorkerGroupStartupTimeoutError(RayTrainError):
|
||||
"""Exception raised when the worker group startup times out.
|
||||
|
||||
Example scenario: 4 GPUs are detected in the cluster, but when the worker
|
||||
are actually scheduled, one of the nodes goes down and only 3 GPUs are
|
||||
available. One of the worker tasks may be stuck pending, until a timeout is reached.
|
||||
"""
|
||||
|
||||
def __init__(self, num_workers: int):
|
||||
timeout = float(
|
||||
os.environ.get(
|
||||
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
|
||||
DEFAULT_WORKER_GROUP_START_TIMEOUT_S,
|
||||
)
|
||||
)
|
||||
self.num_workers = num_workers
|
||||
super().__init__(
|
||||
f"The worker group startup timed out after {timeout} seconds waiting "
|
||||
f"for {num_workers} workers. "
|
||||
"Potential causes include: "
|
||||
"(1) temporary insufficient cluster resources while waiting for "
|
||||
"autoscaling (ignore this warning in this case), "
|
||||
"(2) infeasible resource request where the provided `ScalingConfig` "
|
||||
"cannot be satisfied), "
|
||||
"and (3) transient network issues. "
|
||||
f"Set the {WORKER_GROUP_START_TIMEOUT_S_ENV_VAR} "
|
||||
"environment variable to increase the timeout."
|
||||
)
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self.num_workers,))
|
||||
|
||||
|
||||
class WorkerGroupStartupFailedError(RayTrainError):
|
||||
"""Exception raised when the worker group fails to start.
|
||||
|
||||
Example scenario: A worker is scheduled onto a node that dies while
|
||||
the worker actor is initializing.
|
||||
"""
|
||||
|
||||
|
||||
class InsufficientClusterResourcesError(RayTrainError):
|
||||
"""Exception raised when the cluster has insufficient resources.
|
||||
|
||||
Example scenario: A worker that requires 1 GPU is scheduled onto a cluster
|
||||
that only has CPU worker node types.
|
||||
"""
|
||||
|
||||
|
||||
class CheckpointManagerInitializationError(RayTrainError):
|
||||
"""Exception raised when the checkpoint manager fails to initialize from a snapshot.
|
||||
|
||||
Example scenarios:
|
||||
1. The checkpoint manager snapshot version is old and
|
||||
incompatible with the current version of Ray Train.
|
||||
2. The checkpoint manager snapshot JSON file is corrupted.
|
||||
3. The checkpoint manager snapshot references checkpoints that cannot be found
|
||||
in the run storage path.
|
||||
"""
|
||||
|
||||
|
||||
class CollectiveTimeoutError(RayTrainError):
|
||||
"""Exception raised when an internal Ray Train collective operation of
|
||||
the worker group times out.
|
||||
"""
|
||||
|
||||
|
||||
class BroadcastCollectiveTimeoutError(CollectiveTimeoutError):
|
||||
"""Exception raised when the broadcast operation times out.
|
||||
|
||||
There are two main timeout examples:
|
||||
1. If not all workers call `ray.train.report`, the entire worker group will
|
||||
hang until the timeout before raising. This prevents indefinite worker
|
||||
group hangs.
|
||||
2. If a worker is slow in the training loop and fails to reach the broadcast
|
||||
time, the collective will time out.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, time_elapsed: Optional[float], missing_ranks: List[int], timeout_s: float
|
||||
):
|
||||
self._time_elapsed = time_elapsed
|
||||
self._missing_ranks = missing_ranks
|
||||
self._timeout_s = timeout_s
|
||||
|
||||
message = (
|
||||
f"The collective operation timed out after {time_elapsed:.2f} seconds. "
|
||||
f"The following ranks have not joined the collective operation: {missing_ranks}\n"
|
||||
f"You can set the timeout with the {COLLECTIVE_TIMEOUT_S_ENV_VAR} "
|
||||
f"environment variable (current value: {timeout_s:.2f} seconds). "
|
||||
"Disable the timeout by setting the environment variable to `None`."
|
||||
)
|
||||
super().__init__(message)
|
||||
|
||||
def __reduce__(self):
|
||||
return (
|
||||
self.__class__,
|
||||
(self._time_elapsed, self._missing_ranks, self._timeout_s),
|
||||
)
|
||||
|
||||
|
||||
class UserExceptionWithTraceback(RayTrainError):
|
||||
"""This class wraps a user code exception raised on the worker
|
||||
with its original traceback string, for logging and debugging purposes.
|
||||
|
||||
This is needed because the original exception traceback is not serialized
|
||||
with the exception when it is *returned* back to the main process.
|
||||
"""
|
||||
|
||||
def __init__(self, exc: BaseException, traceback_str: str):
|
||||
self._base_exc = exc
|
||||
self._traceback_str = traceback_str
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self._base_exc, self._traceback_str))
|
||||
|
||||
def __str__(self):
|
||||
return self._traceback_str
|
||||
@@ -0,0 +1,240 @@
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2.api.callback import RayTrainCallback
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller import (
|
||||
TrainControllerState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.failure_handling import FailureDecision
|
||||
from ray.train.v2._internal.execution.scaling_policy import ResizeDecision
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
ExecutionGroup,
|
||||
ReplicaGroup,
|
||||
Worker,
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
WorkerGroupPollStatus,
|
||||
)
|
||||
from ray.train.v2.api.result import Result
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ExecutionGroupCallback(RayTrainCallback):
|
||||
"""Base callback for execution groups (worker groups and replica groups)."""
|
||||
|
||||
def before_init_train_context(
|
||||
self, workers: List["Worker"]
|
||||
) -> Dict[str, List[Any]]:
|
||||
"""Called before initializing the TrainContext for an execution group.
|
||||
|
||||
Return:
|
||||
A dictionary of additional arguments for TrainContext.
|
||||
The key is the argument name and the value is a list of argument values
|
||||
to pass to the TrainContext constructor of each worker in the group.
|
||||
"""
|
||||
return {}
|
||||
|
||||
def after_execution_group_start(self, execution_group: "ExecutionGroup"):
|
||||
"""Called after an execution group is started or replaced.
|
||||
All workers in the execution group should be ready to execute tasks."""
|
||||
pass
|
||||
|
||||
def before_execution_group_shutdown(self, execution_group: "ExecutionGroup"):
|
||||
"""Called before an execution group is shut down.
|
||||
Workers may be dead at this point due to actor failures."""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class WorkerGroupCallback(ExecutionGroupCallback):
|
||||
@contextmanager
|
||||
def on_worker_group_start(self):
|
||||
yield
|
||||
|
||||
def before_worker_group_start(self, worker_group_context: "WorkerGroupContext"):
|
||||
"""Called before the worker group actors are initialized."""
|
||||
pass
|
||||
|
||||
def after_worker_group_start(self, worker_group: "WorkerGroup"):
|
||||
"""Called after the worker group actors are initialized.
|
||||
All workers should be ready to execute tasks."""
|
||||
return self.after_execution_group_start(worker_group)
|
||||
|
||||
def after_worker_group_training_start(self, worker_group: "WorkerGroup"):
|
||||
pass
|
||||
|
||||
@contextmanager
|
||||
def on_worker_group_shutdown(self):
|
||||
yield
|
||||
|
||||
def before_worker_group_shutdown(self, worker_group: "WorkerGroup"):
|
||||
"""Called before the worker group is shut down.
|
||||
Workers may be dead at this point due to actor failures, so this method
|
||||
should catch and handle exceptions if attempting to execute tasks."""
|
||||
return self.before_execution_group_shutdown(worker_group)
|
||||
|
||||
def after_worker_group_shutdown(self, worker_group_context: "WorkerGroupContext"):
|
||||
"""Called after the worker group is shut down."""
|
||||
pass
|
||||
|
||||
def after_worker_group_poll_status(
|
||||
self, worker_group_status: "WorkerGroupPollStatus"
|
||||
):
|
||||
pass
|
||||
|
||||
def before_worker_group_abort(self, worker_group_context: "WorkerGroupContext"):
|
||||
"""Called before the worker group is aborted."""
|
||||
pass
|
||||
|
||||
def after_worker_group_abort(self, worker_group_context: "WorkerGroupContext"):
|
||||
"""Called after the worker group is aborted."""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ReplicaGroupCallback(ExecutionGroupCallback):
|
||||
"""Callback for replica group lifecycle events."""
|
||||
|
||||
def after_replica_group_start(self, replica_group: "ReplicaGroup"):
|
||||
"""Called after a replica group is started or replaced.
|
||||
All workers in the replica group should be ready to execute tasks."""
|
||||
return self.after_execution_group_start(replica_group)
|
||||
|
||||
def before_replica_group_shutdown(self, replica_group: "ReplicaGroup"):
|
||||
"""Called before a replica group is shut down.
|
||||
Workers may be dead at this point due to actor failures."""
|
||||
return self.before_execution_group_shutdown(replica_group)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ControllerCallback(RayTrainCallback):
|
||||
def after_controller_start(self, train_run_context: "TrainRunContext"):
|
||||
"""Called immediately after `TrainController.run` is called,
|
||||
before the control loop starts executing."""
|
||||
pass
|
||||
|
||||
# TODO(matthewdeng): Revisit this callback interface for better extensibility.
|
||||
# This hook was added for the specific use case of setting a `label_selector`
|
||||
# for new worker groups (e.g., for TPU reservations). The current interface is
|
||||
# tightly coupled to this purpose and limits its reuse for other use-cases.
|
||||
def on_controller_start_worker_group(
|
||||
self, *, scaling_config: ScalingConfig, num_workers: int
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Called by the TrainController before the worker group is started.
|
||||
|
||||
This hook can be used to perform setup that modifies the worker group's
|
||||
placement, such as reserving an accelerator slice.
|
||||
|
||||
Args:
|
||||
scaling_config: The scaling configuration for the run.
|
||||
num_workers: The number of workers to be started.
|
||||
|
||||
Returns:
|
||||
An optional dictionary defining a `label_selector`
|
||||
to gang schedule the worker group on the reserved TPU slice.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
"""Called before `TrainController.run` exits,
|
||||
after the control loop has exited."""
|
||||
pass
|
||||
|
||||
def after_controller_state_update(
|
||||
self,
|
||||
previous_state: "TrainControllerState",
|
||||
current_state: "TrainControllerState",
|
||||
):
|
||||
"""Called whenever the controller state is updated."""
|
||||
pass
|
||||
|
||||
def before_controller_execute_failure_decision(
|
||||
self,
|
||||
failure_decision: "FailureDecision",
|
||||
):
|
||||
"""Called before the controller executes a failure decision."""
|
||||
pass
|
||||
|
||||
def before_controller_execute_resize_decision(
|
||||
self,
|
||||
resize_decision: "ResizeDecision",
|
||||
):
|
||||
"""Called before the controller executes a resize decision."""
|
||||
pass
|
||||
|
||||
def after_controller_finish(self, result: "Result"):
|
||||
"""Called after the training run completes, providing access to the final result.
|
||||
|
||||
Args:
|
||||
result: The final training result containing metrics and checkpoint.
|
||||
"""
|
||||
pass
|
||||
|
||||
def before_controller_abort(self):
|
||||
"""Called during `TrainController.abort` before the actor process exits."""
|
||||
pass
|
||||
|
||||
|
||||
# TODO: consider consolidating all metrics into one dict, possibly with UDF
|
||||
@DeveloperAPI
|
||||
class ReportCallback(RayTrainCallback):
|
||||
def after_report(
|
||||
self,
|
||||
training_report: _TrainingReport,
|
||||
metrics: List[Dict[str, Any]],
|
||||
):
|
||||
"""Called after all workers have reported a training result.
|
||||
|
||||
Note that this differs from `after_worker_group_poll_status`,
|
||||
which may only contain a subset of workers that have reported.
|
||||
For example, if only rank 0 is performing checkpointing, then
|
||||
rank 0 would report a training result the slowest.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class WorkerCallback(RayTrainCallback):
|
||||
"""
|
||||
Callbacks that are hooked to the worker event.
|
||||
|
||||
These callbacks are created on the train driver process and then
|
||||
copied and passed to all the workers.
|
||||
The execution of these callbacks happens on each of the workers,
|
||||
not on the train driver process.
|
||||
"""
|
||||
|
||||
def after_init_train_context(self):
|
||||
pass
|
||||
|
||||
def before_worker_shutdown(self):
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainContextCallback(RayTrainCallback):
|
||||
"""
|
||||
Callbacks that are hooked to the train context event.
|
||||
|
||||
These callbacks are created on the train driver process and then
|
||||
copied and passed to all the workers.
|
||||
The execution of these callbacks happens on the train context of the workers.
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def on_report(self):
|
||||
yield
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_sync(self):
|
||||
yield
|
||||
|
||||
@contextmanager
|
||||
def on_checkpoint_transfer(self):
|
||||
yield
|
||||
@@ -0,0 +1,67 @@
|
||||
import logging
|
||||
|
||||
from ray.train.v2.api.exceptions import ControllerError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CallbackManager:
|
||||
def __init__(self, callbacks):
|
||||
self._callbacks = callbacks
|
||||
|
||||
def _get_method(self, callback, hook_name: str):
|
||||
"""Look up a hook method on a callback, raising if missing."""
|
||||
callback_name = type(callback).__name__
|
||||
method = getattr(callback, hook_name, None)
|
||||
if method is None or not callable(method):
|
||||
raise ControllerError(
|
||||
AttributeError(
|
||||
f"Callback '{callback_name}' hook '{hook_name}' is missing "
|
||||
"or not callable."
|
||||
)
|
||||
)
|
||||
return method, callback_name
|
||||
|
||||
def invoke(self, hook_name: str, *args, **context) -> None:
|
||||
for callback in self._callbacks:
|
||||
method, callback_name = self._get_method(callback, hook_name)
|
||||
try:
|
||||
method(*args, **context)
|
||||
except Exception as e:
|
||||
# TODO: Enable configuration to suppress exceptions.
|
||||
logger.exception(
|
||||
f"Exception raised in callback hook '{hook_name}' from callback "
|
||||
f"'{callback_name}'."
|
||||
)
|
||||
raise ControllerError(e) from e
|
||||
|
||||
async def async_invoke(self, hook_name: str, *args, **context) -> None:
|
||||
for callback in self._callbacks:
|
||||
method, callback_name = self._get_method(callback, hook_name)
|
||||
try:
|
||||
await method(*args, **context)
|
||||
except Exception as e:
|
||||
# TODO: Enable configuration to suppress exceptions.
|
||||
logger.exception(
|
||||
f"Exception raised in callback hook '{hook_name}' from callback "
|
||||
f"'{callback_name}'."
|
||||
)
|
||||
raise ControllerError(e) from e
|
||||
|
||||
def invoke_best_effort(self, hook_name: str, *args, **context) -> None:
|
||||
"""Invoke a hook on every callback, logging and suppressing errors.
|
||||
|
||||
Unlike ``invoke``, this does not fail fast — every callback is
|
||||
attempted even if earlier ones raise. Used for cleanup hooks
|
||||
(e.g. ``before_controller_abort``) where partial execution is
|
||||
better than skipping remaining callbacks.
|
||||
"""
|
||||
for callback in self._callbacks:
|
||||
method, callback_name = self._get_method(callback, hook_name)
|
||||
try:
|
||||
method(*args, **context)
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in callback hook '{hook_name}' from callback "
|
||||
f"'{callback_name}': {e}"
|
||||
)
|
||||
@@ -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),
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
import ray.cloudpickle as pickle
|
||||
from ray.train.v2._internal.execution.context import get_train_context
|
||||
|
||||
# For reference, {1:1} is 19 bytes, {"1":"1"} is 21 bytes,
|
||||
# and {"12345": "12345"} is 25 bytes.
|
||||
_MAX_BROADCAST_SIZE_BYTES = 1000
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def barrier() -> None:
|
||||
"""
|
||||
Create a barrier across all training workers.
|
||||
"""
|
||||
train_context = get_train_context()
|
||||
sync_actor = train_context.get_synchronization_actor()
|
||||
return ray.get(
|
||||
sync_actor.broadcast_from_rank_zero.remote(
|
||||
world_rank=train_context.get_world_rank(),
|
||||
world_size=train_context.get_world_size(),
|
||||
data=None,
|
||||
caller_method_name="ray.train.collective.barrier",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def broadcast_from_rank_zero(data: Any) -> Any:
|
||||
"""Broadcast data from the rank 0 worker to all other workers.
|
||||
|
||||
This method is used by the public API function :func:`ray.train.collective.broadcast_from_rank_zero`.
|
||||
Users should typically call ``ray.train.collective.broadcast_from_rank_zero()`` instead of calling this method directly.
|
||||
"""
|
||||
# Validate data.
|
||||
if data is not None:
|
||||
data_bytes = len(pickle.dumps(data))
|
||||
if data_bytes > _MAX_BROADCAST_SIZE_BYTES:
|
||||
logger.warning(
|
||||
f"Data size {data_bytes} bytes exceeds the maximum broadcast "
|
||||
f"size of {_MAX_BROADCAST_SIZE_BYTES} bytes"
|
||||
)
|
||||
|
||||
train_context = get_train_context()
|
||||
sync_actor = train_context.get_synchronization_actor()
|
||||
return ray.get(
|
||||
sync_actor.broadcast_from_rank_zero.remote(
|
||||
world_rank=train_context.get_world_rank(),
|
||||
world_size=train_context.get_world_size(),
|
||||
data=data,
|
||||
caller_method_name="ray.train.collective.broadcast_from_rank_zero",
|
||||
)
|
||||
)
|
||||
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
from .controller import TrainController
|
||||
|
||||
__all__ = ["TrainController"]
|
||||
@@ -0,0 +1,865 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Callable, List, Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray.exceptions import AsyncioActorExit
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_ENABLE_CONTROLLER_LOGGING,
|
||||
DEFAULT_ENABLE_PREEMPTION_WATCHER,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL_S,
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
||||
HEALTH_CHECK_INTERVAL_S_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
ReportCallback,
|
||||
TrainContextCallback,
|
||||
WorkerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback_manager import CallbackManager
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.checkpoint.report_handler import (
|
||||
ReportCallbackHandler,
|
||||
)
|
||||
from ray.train.v2._internal.execution.checkpoint.validation_manager import (
|
||||
ValidationManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller.state import (
|
||||
AbortedState,
|
||||
ErroredState,
|
||||
FinishedState,
|
||||
InitializingState,
|
||||
ReschedulingState,
|
||||
ResizingState,
|
||||
RestartingState,
|
||||
RunningState,
|
||||
SchedulingState,
|
||||
ShuttingDownState,
|
||||
TrainControllerState,
|
||||
)
|
||||
from ray.train.v2._internal.execution.failure_handling import (
|
||||
FailureDecision,
|
||||
FailurePolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.scaling_policy import (
|
||||
NoopDecision,
|
||||
ResizeDecision,
|
||||
ScalingPolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
WorkerGroupPollStatus,
|
||||
)
|
||||
from ray.train.v2._internal.logging import LoggingManager
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
|
||||
from ray.train.v2.api.callback import RayTrainCallback
|
||||
from ray.train.v2.api.exceptions import (
|
||||
ControllerError,
|
||||
TrainingFailedError,
|
||||
)
|
||||
from ray.train.v2.api.report_config import CheckpointConsistencyMode
|
||||
from ray.train.v2.api.result import Result
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
||||
|
||||
from ray.util.tpu import get_tpu_num_slices_for_workers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainControllerLoopIterationResult:
|
||||
"""The result of a single iteration of the control loop."""
|
||||
|
||||
run_attempt_id: str
|
||||
previous_state: TrainControllerState
|
||||
next_state: TrainControllerState
|
||||
training_failed_error: Optional[TrainingFailedError] = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"TrainControllerLoopIterationResult(\n"
|
||||
f" run_attempt_id={self.run_attempt_id},\n"
|
||||
f" previous_state={self.previous_state._state_type.state_name},\n"
|
||||
f" next_state={self.next_state._state_type.state_name}\n"
|
||||
f" training_failed_error={self.training_failed_error}\n"
|
||||
f")"
|
||||
)
|
||||
|
||||
|
||||
class TrainController:
|
||||
"""Manages the execution of a distributed training job.
|
||||
|
||||
Responsibilities include:
|
||||
* Triggering the training function to run on the worker group.
|
||||
* Monitoring the status of the worker group.
|
||||
* Handling scaling decisions by restarting the worker group.
|
||||
* Handling failure decisions by restarting the worker group or terminating training.
|
||||
* Running callback logic on different hooks in the control loop.
|
||||
"""
|
||||
|
||||
worker_group_cls = WorkerGroup
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_fn_ref: ObjectRefWrapper[Callable[[], None]],
|
||||
train_run_context: TrainRunContext,
|
||||
scaling_policy: ScalingPolicy,
|
||||
failure_policy: FailurePolicy,
|
||||
callbacks: Optional[List[RayTrainCallback]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
self._train_run_context = train_run_context
|
||||
if ray_constants.env_bool(
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
DEFAULT_ENABLE_CONTROLLER_LOGGING,
|
||||
):
|
||||
LoggingManager.configure_controller_logger(self._train_run_context)
|
||||
self._train_fn_ref = train_fn_ref
|
||||
self._scaling_policy = scaling_policy
|
||||
self._failure_policy = failure_policy
|
||||
self._run_config = self._train_run_context.run_config
|
||||
self._callbacks = callbacks or []
|
||||
self._storage_context = self._train_run_context.run_config.storage_context
|
||||
|
||||
self._checkpoint_manager = CheckpointManager(
|
||||
checkpoint_config=self._run_config.checkpoint_config,
|
||||
storage_context=self._storage_context,
|
||||
)
|
||||
if validation_config:
|
||||
validation_manager = ValidationManager(
|
||||
checkpoint_manager=self._checkpoint_manager,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
else:
|
||||
validation_manager = None
|
||||
report_handler = ReportCallbackHandler(
|
||||
report_callbacks=(
|
||||
[self._checkpoint_manager]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [c for c in self._callbacks if isinstance(c, ReportCallback)]
|
||||
)
|
||||
)
|
||||
|
||||
# Group callbacks by the hooks they're subscribed to.
|
||||
self._controller_callbacks = (
|
||||
[
|
||||
self._scaling_policy,
|
||||
]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [c for c in self._callbacks if isinstance(c, ControllerCallback)]
|
||||
)
|
||||
self._controller_callback_manager = CallbackManager(self._controller_callbacks)
|
||||
# Group callbacks that will be propagated to the worker group,
|
||||
# train worker and the train context.
|
||||
self._worker_group_callbacks_to_propagate = (
|
||||
[report_handler]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [
|
||||
c
|
||||
for c in self._callbacks
|
||||
if isinstance(
|
||||
c, (WorkerGroupCallback, WorkerCallback, TrainContextCallback)
|
||||
)
|
||||
]
|
||||
+ [self._checkpoint_manager]
|
||||
)
|
||||
|
||||
self._health_check_interval_s = float(
|
||||
os.getenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, DEFAULT_HEALTH_CHECK_INTERVAL_S)
|
||||
)
|
||||
|
||||
self._manages_replica_groups = (
|
||||
train_run_context.backend_config.backend_cls.has_replica_groups
|
||||
if train_run_context.backend_config
|
||||
else False
|
||||
)
|
||||
|
||||
# Register the preemption-observability callback when not in TorchFT
|
||||
# mode (replica groups handle peer loss via their own quorum).
|
||||
enable_preemption_watcher = ray_constants.env_bool(
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
||||
DEFAULT_ENABLE_PREEMPTION_WATCHER,
|
||||
)
|
||||
if self._manages_replica_groups:
|
||||
if enable_preemption_watcher and ray_constants.env_set_by_user(
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR
|
||||
):
|
||||
logger.info(
|
||||
"The preemption watcher is not compatible with replica "
|
||||
"groups (e.g. TorchFT), which handle peer loss via their "
|
||||
"own quorum; skipping it."
|
||||
)
|
||||
elif enable_preemption_watcher:
|
||||
from ray.train.v2._internal.callbacks.preemption_callback import (
|
||||
PreemptionCallback,
|
||||
)
|
||||
|
||||
self._worker_group_callbacks_to_propagate.append(PreemptionCallback())
|
||||
|
||||
self._worker_group: Optional[WorkerGroup] = None
|
||||
self._state = InitializingState()
|
||||
self._return_value: Optional[Any] = None
|
||||
|
||||
# TODO: These can be attributes of a RunAttempt?
|
||||
self._latest_poll_time = float("-inf")
|
||||
|
||||
# Generate an initial run attempt ID so that `_run_controller_hook`
|
||||
# can reference it if a callback fails during `_start`.
|
||||
self._generate_run_attempt_id()
|
||||
self._start()
|
||||
|
||||
def _run_controller_hook(
|
||||
self,
|
||||
hook_name: str,
|
||||
*args,
|
||||
invoke_failure_decision_callbacks: bool = True,
|
||||
**context,
|
||||
) -> Optional["TrainControllerLoopIterationResult"]:
|
||||
"""Invoke a named controller hook and catch any exceptions.
|
||||
|
||||
This method invokes all callbacks registered for the given controller hook.
|
||||
If a callback raises an error, the error is routed through the failure policy
|
||||
and may produce a ``TrainControllerLoopIterationResult``, indicating that the
|
||||
current controller step should exit early with this failure result.
|
||||
|
||||
Args:
|
||||
hook_name: The controller hook name to invoke.
|
||||
*args: Positional arguments to pass to the hook.
|
||||
invoke_failure_decision_callbacks: Whether to invoke failure-decision hooks
|
||||
when handling a callback failure.
|
||||
**context: Keyword arguments to pass to the hook.
|
||||
|
||||
Returns:
|
||||
failure_result: A``TrainControllerLoopIterationResult`` if the hook execution results
|
||||
in an early exit from the controller loop to raise the callback error,
|
||||
or ``None`` if hook execution completes successfully.
|
||||
"""
|
||||
try:
|
||||
self._controller_callback_manager.invoke(hook_name, *args, **context)
|
||||
except ControllerError as error:
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=error,
|
||||
)
|
||||
# Avoid re-entering controller callback hooks while handling a callback failure.
|
||||
return self._execute_failure_decision(
|
||||
failure_decision,
|
||||
training_failed_error=error,
|
||||
invoke_failure_decision_callbacks=invoke_failure_decision_callbacks,
|
||||
)
|
||||
return None
|
||||
|
||||
def _execute_resize_decision(
|
||||
self, decision: ResizeDecision
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Executes resize decisions.
|
||||
|
||||
Errors from worker group shutdown, callbacks, or worker group startup
|
||||
are allowed to propagate to the catch-all in ``run()``.
|
||||
"""
|
||||
|
||||
failure_result = self._run_controller_hook(
|
||||
"before_controller_execute_resize_decision", decision
|
||||
)
|
||||
if failure_result:
|
||||
return failure_result
|
||||
current_num_workers = (
|
||||
len(self._worker_group.get_workers()) if self._worker_group else 0
|
||||
)
|
||||
poll_status = (
|
||||
self._worker_group.get_latest_poll_status() if self._worker_group else None
|
||||
)
|
||||
failing_rgs = (
|
||||
poll_status.failing_replica_group_indices if poll_status else set()
|
||||
)
|
||||
all_rgs = poll_status.all_replica_group_indices if poll_status else set()
|
||||
if (
|
||||
self._manages_replica_groups
|
||||
and bool(failing_rgs)
|
||||
and failing_rgs != all_rgs
|
||||
and self._worker_group
|
||||
# TODO: relax this after integrating replica groups with elastic training.
|
||||
and decision.num_workers == current_num_workers
|
||||
):
|
||||
# Torchft: replace only failing replica groups.
|
||||
self._replace_bad_workers(poll_status)
|
||||
else:
|
||||
# Standard: full restart.
|
||||
if self._worker_group:
|
||||
self._shutdown_worker_group()
|
||||
self._start_worker_group(
|
||||
num_workers=decision.num_workers,
|
||||
resources_per_worker=decision.resources_per_worker,
|
||||
)
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=self._state,
|
||||
next_state=RunningState(),
|
||||
)
|
||||
|
||||
def _replace_bad_workers(self, poll_status: WorkerGroupPollStatus):
|
||||
"""Replace failing replica groups in the worker group.
|
||||
|
||||
Args:
|
||||
poll_status: The poll status containing error information.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
failing_rg_indices = poll_status.failing_replica_group_indices
|
||||
|
||||
if not failing_rg_indices:
|
||||
logger.warning("No failing replica groups found in poll status.")
|
||||
return
|
||||
|
||||
logger.info(f"Replacing failing replica groups: {failing_rg_indices}")
|
||||
|
||||
for rg_index in failing_rg_indices:
|
||||
# TODO: parallelize this.
|
||||
# TODO: also ensure that if earlier replacements succeed and later replacements fail,
|
||||
# we don't redo the earlier replacements.
|
||||
# See https://github.com/ray-project/ray/pull/61475#discussion_r3055217289
|
||||
self._worker_group.replace_replica_group(rg_index)
|
||||
|
||||
def _get_retry_state(
|
||||
self,
|
||||
controller_state: TrainControllerState,
|
||||
training_failed_error: TrainingFailedError,
|
||||
) -> TrainControllerState:
|
||||
if isinstance(controller_state, RunningState):
|
||||
return RestartingState(training_failed_error=training_failed_error)
|
||||
elif isinstance(controller_state, SchedulingState):
|
||||
return ReschedulingState(training_failed_error=training_failed_error)
|
||||
else:
|
||||
# Cannot retry from this state (e.g. InitializingState,
|
||||
# ShuttingDownState); force shutdown with error.
|
||||
logger.warning(
|
||||
"Cannot retry from state %s; forcing shutdown.",
|
||||
type(controller_state).__name__,
|
||||
)
|
||||
return ShuttingDownState(
|
||||
next_state=ErroredState(training_failed_error=training_failed_error)
|
||||
)
|
||||
|
||||
def _execute_failure_decision(
|
||||
self,
|
||||
failure_decision: FailureDecision,
|
||||
training_failed_error: TrainingFailedError,
|
||||
invoke_failure_decision_callbacks: bool = True,
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Executes failure handling decisions for a scheduling or poll error."""
|
||||
|
||||
controller_state = self.get_state()
|
||||
|
||||
if invoke_failure_decision_callbacks:
|
||||
failure_result = self._run_controller_hook(
|
||||
"before_controller_execute_failure_decision",
|
||||
failure_decision,
|
||||
invoke_failure_decision_callbacks=False,
|
||||
)
|
||||
if failure_result:
|
||||
return failure_result
|
||||
|
||||
# TODO: What should we do here?
|
||||
# This currently never happens because there must be errors.
|
||||
if failure_decision == FailureDecision.NOOP:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=controller_state,
|
||||
training_failed_error=training_failed_error,
|
||||
)
|
||||
|
||||
if failure_decision == FailureDecision.RETRY:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=self._get_retry_state(
|
||||
controller_state, training_failed_error
|
||||
),
|
||||
)
|
||||
elif failure_decision == FailureDecision.RAISE:
|
||||
next_state = ShuttingDownState(
|
||||
next_state=ErroredState(
|
||||
training_failed_error=training_failed_error,
|
||||
),
|
||||
)
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
training_failed_error=training_failed_error,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unexpected failure decision: {failure_decision}")
|
||||
|
||||
async def _poll_workers(self) -> WorkerGroupPollStatus:
|
||||
# Ensure that the time between polls is at least HEALTH_CHECK_INTERVAL_S.
|
||||
time_since_last_poll = time_monotonic() - self._latest_poll_time
|
||||
if time_since_last_poll < self._health_check_interval_s:
|
||||
remaining_time = max(
|
||||
self._health_check_interval_s - time_since_last_poll, 0
|
||||
)
|
||||
await asyncio.sleep(remaining_time)
|
||||
if self.get_state().is_terminal():
|
||||
logger.debug(
|
||||
f"Controller is unexpectedly in terminal state {self.get_state()} after "
|
||||
"sleeping and before polling workers. Exiting actor."
|
||||
)
|
||||
ray.actor.exit_actor()
|
||||
|
||||
status = self._worker_group.poll_status(timeout=self._health_check_interval_s)
|
||||
self._latest_poll_time = time_monotonic()
|
||||
return status
|
||||
|
||||
def _start_worker_group(self, num_workers: int, resources_per_worker: dict) -> None:
|
||||
"""Start the worker group and launch the train function.
|
||||
|
||||
Args:
|
||||
num_workers: The number of workers to start.
|
||||
resources_per_worker: The resources per worker to start.
|
||||
|
||||
Raises:
|
||||
Exception: If the worker group failed to start.
|
||||
"""
|
||||
placement_strategy = self._scaling_policy.scaling_config.placement_strategy
|
||||
scaling_config = self._train_run_context.scaling_config
|
||||
|
||||
# Check for `label_selector` to influence WorkerGroup scheduling.
|
||||
label_selector = scaling_config._label_selector_per_worker(num_workers)
|
||||
for callback in self._controller_callbacks:
|
||||
selector = callback.on_controller_start_worker_group(
|
||||
scaling_config=scaling_config, num_workers=num_workers
|
||||
)
|
||||
if selector:
|
||||
if label_selector:
|
||||
logger.warning(
|
||||
f"Overriding `ScalingConfig.label_selector` {label_selector} "
|
||||
f"with label_selector returned by user-specified callback {selector}"
|
||||
)
|
||||
label_selector = [selector.copy() for _ in range(num_workers)]
|
||||
|
||||
# Calculate num_slices for the worker group if using TPU.
|
||||
num_slices = 1
|
||||
if scaling_config.use_tpu:
|
||||
num_slices = get_tpu_num_slices_for_workers(
|
||||
topology=scaling_config.topology,
|
||||
accelerator_type=scaling_config.accelerator_type,
|
||||
num_workers=num_workers,
|
||||
resources_per_worker=resources_per_worker,
|
||||
)
|
||||
|
||||
worker_group_context = WorkerGroupContext(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
train_fn_ref=self._train_fn_ref,
|
||||
num_workers=num_workers,
|
||||
resources_per_worker=resources_per_worker,
|
||||
placement_strategy=placement_strategy,
|
||||
label_selector=label_selector,
|
||||
num_slices=num_slices,
|
||||
)
|
||||
self._worker_group = self.worker_group_cls.create(
|
||||
train_run_context=self._train_run_context,
|
||||
worker_group_context=worker_group_context,
|
||||
callbacks=self._worker_group_callbacks_to_propagate,
|
||||
)
|
||||
|
||||
def _start(self):
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_start", self._train_run_context
|
||||
)
|
||||
if failure_result:
|
||||
self._set_state(failure_result.next_state)
|
||||
|
||||
async def _shutdown(self) -> "TrainControllerLoopIterationResult":
|
||||
"""Execute shutdown and return the final state transition.
|
||||
|
||||
Shutdown errors are never retried. If an error occurs during shutdown:
|
||||
- If we're already shutting down after a training error
|
||||
(next_state is ErroredState), the original error is preserved.
|
||||
- Otherwise the shutdown error becomes the training failure.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
assert isinstance(controller_state, ShuttingDownState)
|
||||
|
||||
shutdown_error = None
|
||||
|
||||
# TODO: move to __del__ after https://github.com/ray-project/ray/issues/53169
|
||||
if self._worker_group:
|
||||
try:
|
||||
self._shutdown_worker_group()
|
||||
except Exception as e:
|
||||
logger.exception("Error shutting down worker group.")
|
||||
shutdown_error = ControllerError(e)
|
||||
|
||||
try:
|
||||
await self._controller_callback_manager.async_invoke(
|
||||
"before_controller_shutdown"
|
||||
)
|
||||
except ControllerError as e:
|
||||
if shutdown_error:
|
||||
logger.warning(
|
||||
"An additional error occurred in the before_controller_shutdown "
|
||||
"callback after a worker group shutdown error. "
|
||||
"This error is being ignored to preserve the original "
|
||||
"shutdown error. Error: %s",
|
||||
e,
|
||||
)
|
||||
else:
|
||||
shutdown_error = e
|
||||
|
||||
if shutdown_error:
|
||||
if isinstance(controller_state.next_state, ErroredState):
|
||||
logger.warning(
|
||||
"Another error occurred during shutdown after a training error. "
|
||||
"This error is being ignored to preserve the original "
|
||||
"training error. Error: %s",
|
||||
shutdown_error,
|
||||
)
|
||||
else:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=ErroredState(training_failed_error=shutdown_error),
|
||||
training_failed_error=shutdown_error,
|
||||
)
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=controller_state.next_state,
|
||||
)
|
||||
|
||||
def _shutdown_worker_group(self):
|
||||
"""Shutdown the worker group and set the worker group to None."""
|
||||
self._worker_group.shutdown()
|
||||
self._worker_group = None
|
||||
|
||||
def get_worker_group(self) -> Optional[WorkerGroup]:
|
||||
return self._worker_group
|
||||
|
||||
def get_state(self) -> TrainControllerState:
|
||||
return self._state
|
||||
|
||||
def _set_state(self, state: TrainControllerState):
|
||||
previous_state = self._state
|
||||
self._state = state
|
||||
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_state_update", previous_state, state
|
||||
)
|
||||
if failure_result:
|
||||
# If we're transitioning into a terminal state, or if we're already in the shutdown path to an errored terminal state
|
||||
# (ShuttingDownState -> ErroredState), preserve the original failure as the
|
||||
# surfaced error. A failure in a state-update callback should not overwrite
|
||||
# the underlying root-cause error.
|
||||
if state.is_terminal() or (
|
||||
isinstance(state, ShuttingDownState)
|
||||
and isinstance(state.next_state, ErroredState)
|
||||
):
|
||||
logger.warning(
|
||||
"A callback failed during a terminal state transition. "
|
||||
"This failure is being ignored to preserve the original "
|
||||
"training result. Error: %s",
|
||||
failure_result.training_failed_error,
|
||||
)
|
||||
return
|
||||
|
||||
# NOTE: We intentionally do *not* re-invoke `after_controller_state_update`
|
||||
# for this transition to avoid re-entering callback hooks while handling
|
||||
# a callback failure.
|
||||
self._state = failure_result.next_state
|
||||
|
||||
def _make_and_handle_scaling_decision_for_non_running_worker_group(
|
||||
self,
|
||||
controller_state: TrainControllerState,
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Make a scaling decision for a non-running worker group and return the appropriate next state.
|
||||
|
||||
This method should be called when entering a state that requires a scaling decision
|
||||
for a non-running worker group.
|
||||
|
||||
This method handles the complete flow of:
|
||||
1. Shutting down the non-running worker group if it still exists.
|
||||
2. Getting a scaling decision for a non-running worker group
|
||||
3. Determining the next state based on the decision type
|
||||
4. Creating and returning the iteration result
|
||||
|
||||
Args:
|
||||
controller_state: The current controller state
|
||||
|
||||
Returns:
|
||||
TrainControllerLoopIterationResult with the appropriate next state
|
||||
"""
|
||||
scaling_decision = (
|
||||
self._scaling_policy.make_decision_for_non_running_worker_group()
|
||||
)
|
||||
|
||||
if isinstance(scaling_decision, NoopDecision):
|
||||
next_state = controller_state
|
||||
elif isinstance(scaling_decision, ResizeDecision):
|
||||
next_state = SchedulingState(scaling_decision)
|
||||
else:
|
||||
raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
)
|
||||
|
||||
async def _step(self) -> TrainControllerLoopIterationResult:
|
||||
"""Run a single iteration of the control loop.
|
||||
|
||||
Returns:
|
||||
The result of the iteration.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
|
||||
if isinstance(
|
||||
controller_state, (InitializingState, RestartingState, ReschedulingState)
|
||||
):
|
||||
return self._make_and_handle_scaling_decision_for_non_running_worker_group(
|
||||
controller_state
|
||||
)
|
||||
elif isinstance(controller_state, SchedulingState):
|
||||
assert isinstance(controller_state.scaling_decision, ResizeDecision)
|
||||
return self._execute_resize_decision(controller_state.scaling_decision)
|
||||
elif isinstance(controller_state, RunningState):
|
||||
worker_group_status: WorkerGroupPollStatus = await self._poll_workers()
|
||||
|
||||
if worker_group_status.finished and not worker_group_status.errors:
|
||||
self._return_value = worker_group_status.worker_statuses[0].return_value
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=ShuttingDownState(
|
||||
next_state=FinishedState(),
|
||||
),
|
||||
)
|
||||
if worker_group_status.errors:
|
||||
worker_group_error = worker_group_status.get_worker_group_error()
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=worker_group_error,
|
||||
)
|
||||
return self._execute_failure_decision(
|
||||
failure_decision, training_failed_error=worker_group_error
|
||||
)
|
||||
|
||||
scaling_decision = (
|
||||
self._scaling_policy.make_decision_for_running_worker_group(
|
||||
worker_group_state=self.get_worker_group().get_worker_group_state(),
|
||||
worker_group_status=worker_group_status,
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(scaling_decision, NoopDecision):
|
||||
next_state = RunningState()
|
||||
elif isinstance(scaling_decision, ResizeDecision):
|
||||
next_state = ResizingState(
|
||||
scaling_decision=scaling_decision,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
)
|
||||
elif isinstance(controller_state, ResizingState):
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=SchedulingState(
|
||||
scaling_decision=controller_state.scaling_decision
|
||||
),
|
||||
)
|
||||
elif isinstance(controller_state, ShuttingDownState):
|
||||
return await self._shutdown()
|
||||
else:
|
||||
raise ValueError(f"Unexpected controller state: {controller_state}")
|
||||
|
||||
def _generate_run_attempt_id(self):
|
||||
self._run_attempt_id = uuid.uuid4().hex
|
||||
return self._run_attempt_id
|
||||
|
||||
def _get_run_attempt_id(self):
|
||||
return self._run_attempt_id
|
||||
|
||||
async def _run_control_loop_iteration(self):
|
||||
"""Run a single iteration of the control loop.
|
||||
|
||||
Steps:
|
||||
1. Poll the worker group for status.
|
||||
2. If the worker group is initializing or recovering from an error,
|
||||
make a scaling decision and execute it.
|
||||
3. If the worker group has finished, set the controller state to FINISHED.
|
||||
4. If the worker group has errors, make a failure decision and execute it.
|
||||
5. Otherwise, the worker group is running healthily.
|
||||
Query the scaling policy for a scaling decision and execute it.
|
||||
|
||||
Errors raised by ``_step`` are caught and routed through the failure
|
||||
policy (retry / raise). If the failure policy itself fails, the
|
||||
controller is forced into ``ErroredState`` as a last resort.
|
||||
|
||||
``AsyncioActorExit`` is always re-raised so that the actor can shut
|
||||
down cleanly.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
assert not controller_state.is_terminal()
|
||||
|
||||
if controller_state.needs_new_run_attempt():
|
||||
self._generate_run_attempt_id()
|
||||
|
||||
try:
|
||||
result = await self._step()
|
||||
except AsyncioActorExit:
|
||||
raise
|
||||
except Exception as e:
|
||||
# Preserve the original error type if it is already a
|
||||
# TrainingFailedError (e.g. WorkerGroupError); otherwise
|
||||
# wrap it in a ControllerError.
|
||||
if isinstance(e, TrainingFailedError):
|
||||
training_error = e
|
||||
else:
|
||||
# Log the full traceback only for unexpected errors.
|
||||
logger.exception("Error in control loop iteration: %s", e)
|
||||
training_error = ControllerError(e)
|
||||
try:
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=training_error,
|
||||
)
|
||||
result = self._execute_failure_decision(
|
||||
failure_decision,
|
||||
training_failed_error=training_error,
|
||||
)
|
||||
except Exception:
|
||||
# Last resort: force into errored state, bypassing callbacks.
|
||||
logger.exception(
|
||||
"Failed to execute failure decision, forcing error state."
|
||||
)
|
||||
self._state = ErroredState(training_failed_error=training_error)
|
||||
return
|
||||
|
||||
self._set_state(result.next_state)
|
||||
|
||||
async def run(self):
|
||||
"""Run the main control loop. Exits when training is finished or errored."""
|
||||
while not self.get_state().is_terminal():
|
||||
await self._run_control_loop_iteration()
|
||||
|
||||
# Call after_controller_finish with the final result.
|
||||
result = self._build_result()
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_finish", result, invoke_failure_decision_callbacks=False
|
||||
)
|
||||
# Since we are already in a terminal state, a callback failure should
|
||||
# not overwrite the training outcome — log and preserve the result.
|
||||
if failure_result:
|
||||
logger.warning(
|
||||
"A callback failed after training finished. "
|
||||
"This failure is being ignored to preserve the original "
|
||||
"training result. Error: %s",
|
||||
failure_result.training_failed_error,
|
||||
)
|
||||
|
||||
async def abort(self):
|
||||
"""Trigger callback abort hooks and terminate the controller process."""
|
||||
# Do not abort run if it's already finished.
|
||||
if self.get_state().is_terminal():
|
||||
return
|
||||
|
||||
self._controller_callback_manager.invoke_best_effort("before_controller_abort")
|
||||
|
||||
# Intentionally abort worker group before setting train run state because
|
||||
# we only reconcile the states of live train runs.
|
||||
try:
|
||||
if self._worker_group:
|
||||
self._worker_group.abort()
|
||||
self._set_state(AbortedState())
|
||||
except Exception as e:
|
||||
logger.exception("Error aborting worker group: %s", e)
|
||||
|
||||
ray.actor.exit_actor()
|
||||
|
||||
def _build_result(self) -> Result:
|
||||
storage = self._checkpoint_manager._storage_context
|
||||
|
||||
latest_checkpoint_result = self._checkpoint_manager.latest_checkpoint_result
|
||||
latest_metrics = (
|
||||
latest_checkpoint_result.metrics if latest_checkpoint_result else None
|
||||
)
|
||||
latest_checkpoint = (
|
||||
latest_checkpoint_result.checkpoint if latest_checkpoint_result else None
|
||||
)
|
||||
best_checkpoints = [
|
||||
(r.checkpoint, r.metrics)
|
||||
for r in self._checkpoint_manager.best_checkpoint_results
|
||||
]
|
||||
|
||||
# Provide the history of metrics attached to checkpoints as a dataframe.
|
||||
metrics_dataframe = None
|
||||
if best_checkpoints:
|
||||
metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
|
||||
|
||||
return Result(
|
||||
metrics=latest_metrics,
|
||||
checkpoint=latest_checkpoint,
|
||||
error=self.get_training_failed_error(),
|
||||
path=storage.experiment_fs_path,
|
||||
best_checkpoints=best_checkpoints,
|
||||
metrics_dataframe=metrics_dataframe,
|
||||
_storage_filesystem=storage.storage_filesystem,
|
||||
return_value=self._return_value,
|
||||
)
|
||||
|
||||
def get_result(self) -> Result:
|
||||
"""Get the final training result from the TrainController."""
|
||||
|
||||
controller_state = self.get_state()
|
||||
if not controller_state.is_terminal():
|
||||
raise ValueError(
|
||||
f"Cannot get result when controller is in state {controller_state}"
|
||||
)
|
||||
return self._build_result()
|
||||
|
||||
def get_training_failed_error(self) -> Optional[TrainingFailedError]:
|
||||
"""Get the training failed error from the controller state.
|
||||
|
||||
Returns:
|
||||
The training failed error if the controller is in an errored state,
|
||||
None otherwise.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
|
||||
if isinstance(controller_state, ErroredState):
|
||||
return controller_state.training_failed_error
|
||||
|
||||
return None
|
||||
|
||||
async def get_all_reported_checkpoints(
|
||||
self,
|
||||
current_report_index: int,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
return await self._checkpoint_manager.get_all_reported_checkpoints(
|
||||
current_report_index, consistency_mode, timeout_s
|
||||
)
|
||||
@@ -0,0 +1,183 @@
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.state.util import is_actor_alive
|
||||
from ray.util.placement_group import PlacementGroup, remove_placement_group
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PlacementGroupCleaner:
|
||||
"""Detached helper that ensures PG cleanup if Ray Train Controller exits ungracefully.
|
||||
|
||||
This actor should be created with lifetime='detached' to avoid being
|
||||
fate-shared with the Train controller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
controller_actor_id: str,
|
||||
check_interval_s: float,
|
||||
get_actor_timeout_s: float,
|
||||
stop_timeout: Optional[float],
|
||||
):
|
||||
self._controller_actor_id = controller_actor_id
|
||||
self._check_interval_s = check_interval_s
|
||||
self._get_actor_timeout_s = get_actor_timeout_s
|
||||
self._stop_timeout = stop_timeout
|
||||
self._pg_queue: queue.Queue = queue.Queue()
|
||||
self._stop_event = threading.Event()
|
||||
self._monitor_thread: Optional[threading.Thread] = None
|
||||
self._exiting: bool = False
|
||||
|
||||
def register_placement_group(self, placement_group: PlacementGroup):
|
||||
logger.debug(
|
||||
"PlacementGroupCleaner registered placement group %s for controller %s",
|
||||
placement_group.id,
|
||||
self._controller_actor_id,
|
||||
)
|
||||
# Send placement group update to the monitor thread via queue
|
||||
self._pg_queue.put(placement_group)
|
||||
|
||||
def start_monitoring(self):
|
||||
"""Start monitoring the controller and placement group."""
|
||||
if self._monitor_thread is not None and self._monitor_thread.is_alive():
|
||||
# Thread already running, just return True
|
||||
logger.debug("Monitor thread already running")
|
||||
return True
|
||||
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._monitor_loop,
|
||||
name="PlacementGroupCleanerMonitor",
|
||||
daemon=True,
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
logger.debug("PlacementGroupCleaner started monitoring in background thread")
|
||||
return True
|
||||
|
||||
def _monitor_loop(self):
|
||||
"""Monitor controller; remove PG when controller is gone.
|
||||
|
||||
This runs continuously until controller dies or stop() is called.
|
||||
Uses a queue to receive placement group updates.
|
||||
"""
|
||||
curr_placement_group: Optional[PlacementGroup] = None
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
# Check for new placement group updates from queue
|
||||
try:
|
||||
pg = self._pg_queue.get(timeout=self._check_interval_s)
|
||||
curr_placement_group = pg
|
||||
logger.debug(f"Updated current placement group to {pg.id}")
|
||||
except queue.Empty:
|
||||
pass # continue to monitor current placement group
|
||||
|
||||
# Check if controller is still alive
|
||||
try:
|
||||
alive = is_actor_alive(
|
||||
actor_id=self._controller_actor_id,
|
||||
timeout=self._get_actor_timeout_s,
|
||||
)
|
||||
except ray.util.state.exception.RayStateApiException:
|
||||
logger.warning(
|
||||
"Failed to query Ray Train Controller actor state. "
|
||||
"State API may be temporarily unavailable. Continuing to monitor."
|
||||
)
|
||||
continue
|
||||
|
||||
# Cleanup if controller is dead
|
||||
if not alive:
|
||||
# Drain any queued placement groups
|
||||
while True:
|
||||
try:
|
||||
pg = self._pg_queue.get_nowait()
|
||||
curr_placement_group = pg
|
||||
except queue.Empty:
|
||||
break
|
||||
self._cleanup_placement_group(curr_placement_group)
|
||||
break
|
||||
|
||||
# Exit the actor after cleanup since controller is dead
|
||||
self._exit()
|
||||
self._monitor_thread = None
|
||||
|
||||
def _cleanup_placement_group(self, placement_group: Optional[PlacementGroup]):
|
||||
"""Clean up the current placement group if it hasn't been removed."""
|
||||
if placement_group is None:
|
||||
logger.debug("No placement group registered; skipping cleanup.")
|
||||
return
|
||||
|
||||
if self._is_placement_group_removed(placement_group):
|
||||
logger.debug(
|
||||
"Controller actor died but placement group already removed; "
|
||||
"skipping cleanup."
|
||||
)
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
f"Detected that the Ray Train controller actor ({self._controller_actor_id}) is dead. "
|
||||
f"Cleaning up placement group = [{placement_group.id}] created by this run."
|
||||
)
|
||||
try:
|
||||
remove_placement_group(placement_group)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to clean up placement group: {e}")
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Placement group = [{placement_group.id}] cleaned up successfully"
|
||||
)
|
||||
|
||||
def _stop_monitor_thread(self):
|
||||
"""Stop the monitor thread and wait for it to exit.
|
||||
|
||||
Returns:
|
||||
bool: True if the thread was stopped, False if there was no active thread.
|
||||
"""
|
||||
if self._monitor_thread is None or not self._monitor_thread.is_alive():
|
||||
return False
|
||||
|
||||
# Signal stop and wait for thread to exit
|
||||
self._stop_event.set()
|
||||
self._monitor_thread.join(timeout=self._stop_timeout)
|
||||
if self._monitor_thread.is_alive():
|
||||
logger.warning(
|
||||
"Monitor thread did not exit within %.2f seconds", self._stop_timeout
|
||||
)
|
||||
return False
|
||||
|
||||
self._monitor_thread = None
|
||||
return True
|
||||
|
||||
def stop(self):
|
||||
"""Request the cleaner to stop monitoring and exit."""
|
||||
self._stop_monitor_thread()
|
||||
self._exit()
|
||||
|
||||
def _is_placement_group_removed(self, placement_group: PlacementGroup) -> bool:
|
||||
"""Check if a placement group has been removed."""
|
||||
try:
|
||||
table = ray.util.placement_group_table(placement_group)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to query placement group table: {e}. "
|
||||
"Assuming placement group is not removed."
|
||||
)
|
||||
return False
|
||||
if "state" not in table:
|
||||
return True
|
||||
return table["state"] == "REMOVED"
|
||||
|
||||
def _exit(self):
|
||||
"""Exit the actor."""
|
||||
if self._exiting:
|
||||
return
|
||||
self._exiting = True
|
||||
try:
|
||||
ray.actor.exit_actor()
|
||||
except Exception as e:
|
||||
# If exit fails for any reason, just log it.
|
||||
logger.warning(f"Failed to exit actor: {e}")
|
||||
@@ -0,0 +1,156 @@
|
||||
from enum import Enum
|
||||
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
ScalingDecision,
|
||||
)
|
||||
from ray.train.v2.api.exceptions import TrainingFailedError
|
||||
|
||||
|
||||
class TrainControllerStateType(Enum):
|
||||
"""Enum representing different states of the train controller.
|
||||
|
||||
States:
|
||||
INITIALIZING: The train controller is starting up. This is always the initial
|
||||
state of the controller.
|
||||
SCHEDULING: The train controller is in the process of scheduling a new worker
|
||||
group.
|
||||
RESCHEDULING: The train controller is in the process of rescheduling the worker
|
||||
group.
|
||||
RUNNING: The train controller is actively running training tasks.
|
||||
RESTARTING: The train controller is in the process of recovering from an error.
|
||||
RESIZING: The train controller is in the process of resizing a running worker
|
||||
group.
|
||||
SHUTTING_DOWN: The train controller has already shut down the worker group and
|
||||
and is in the process of shutting itself down.
|
||||
ERRORED: A terminal state indicating that training has encountered an error and
|
||||
cannot continue.
|
||||
FINISHED: A terminal state indicating that training has completed.
|
||||
ABORTED: A terminal state indicating that training has been aborted.
|
||||
|
||||
Args:
|
||||
state_name: The name of the state.
|
||||
is_terminal: Whether this is a terminal state that should not be further processed.
|
||||
needs_new_run_attempt: Whether this state requires starting a new run attempt, where
|
||||
a run attempt is a logical unit that encompasses both scheduling workers and
|
||||
executing training on those workers.
|
||||
"""
|
||||
|
||||
INITIALIZING = ("INITIALIZING", False, True)
|
||||
SCHEDULING = ("SCHEDULING", False, False)
|
||||
RESCHEDULING = ("RESCHEDULING", False, False)
|
||||
RUNNING = ("RUNNING", False, False)
|
||||
RESTARTING = ("RESTARTING", False, True)
|
||||
RESIZING = ("RESIZING", False, True)
|
||||
SHUTTING_DOWN = ("SHUTTING_DOWN", False, False)
|
||||
ERRORED = ("ERRORED", True, False)
|
||||
FINISHED = ("FINISHED", True, False)
|
||||
ABORTED = ("ABORTED", True, False)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
state_name: str,
|
||||
is_terminal: bool,
|
||||
needs_new_run_attempt: bool,
|
||||
):
|
||||
self.state_name = state_name
|
||||
self.is_terminal = is_terminal
|
||||
self.needs_new_run_attempt = needs_new_run_attempt
|
||||
|
||||
|
||||
class TrainControllerState:
|
||||
"""Base class for all train controller states.
|
||||
|
||||
Methods:
|
||||
get_type() -> TrainControllerStateType: Returns the type of the state.
|
||||
is_terminal() -> bool: Returns whether the state is terminal.
|
||||
needs_new_run_attempt() -> bool: Returns whether a new run attempt is needed.
|
||||
"""
|
||||
|
||||
def __init__(self, state_type: TrainControllerStateType):
|
||||
self._state_type = state_type
|
||||
|
||||
def __repr__(self) -> str:
|
||||
attrs = {
|
||||
"type": self._state_type.name,
|
||||
"is_terminal": self._state_type.is_terminal,
|
||||
"needs_new_run_attempt": self._state_type.needs_new_run_attempt,
|
||||
**{k: v for k, v in vars(self).items() if not k.startswith("_")},
|
||||
}
|
||||
attrs_str = "\n ".join(f"{k}={v}" for k, v in attrs.items())
|
||||
return f"{self.__class__.__name__}(\n {attrs_str}\n)"
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
return self._state_type.is_terminal
|
||||
|
||||
def needs_new_run_attempt(self) -> bool:
|
||||
return self._state_type.needs_new_run_attempt
|
||||
|
||||
|
||||
class InitializingState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.INITIALIZING)
|
||||
|
||||
|
||||
class SchedulingState(TrainControllerState):
|
||||
def __init__(self, scaling_decision: ScalingDecision):
|
||||
super().__init__(state_type=TrainControllerStateType.SCHEDULING)
|
||||
self.scaling_decision = scaling_decision
|
||||
|
||||
|
||||
class ReschedulingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESCHEDULING)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class RunningState(TrainControllerState):
|
||||
# TODO: Split into multiple more granular states, or add more fields.
|
||||
# For example, we may want to indicate if any health checks failed.
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.RUNNING)
|
||||
|
||||
|
||||
class RestartingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESTARTING)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class ResizingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
scaling_decision: ScalingDecision,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESIZING)
|
||||
self.scaling_decision = scaling_decision
|
||||
|
||||
|
||||
class ShuttingDownState(TrainControllerState):
|
||||
def __init__(self, next_state: "TrainControllerState"):
|
||||
super().__init__(state_type=TrainControllerStateType.SHUTTING_DOWN)
|
||||
self.next_state = next_state
|
||||
|
||||
|
||||
class ErroredState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.ERRORED)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class FinishedState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.FINISHED)
|
||||
|
||||
|
||||
class AbortedState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.ABORTED)
|
||||
@@ -0,0 +1,16 @@
|
||||
# isort: off
|
||||
from .failure_policy import FailureDecision, FailurePolicy
|
||||
from .default import DefaultFailurePolicy
|
||||
from .factory import create_failure_policy
|
||||
|
||||
# isort: on
|
||||
|
||||
__all__ = [
|
||||
"DefaultFailurePolicy",
|
||||
"FailureDecision",
|
||||
"FailurePolicy",
|
||||
"create_failure_policy",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,100 @@
|
||||
import logging
|
||||
|
||||
from .failure_policy import FailureDecision, FailurePolicy
|
||||
from ray.train.v2._internal.exceptions import (
|
||||
WorkerGroupStartupFailedError,
|
||||
WorkerGroupStartupTimeoutError,
|
||||
)
|
||||
from ray.train.v2.api.config import FailureConfig
|
||||
from ray.train.v2.api.exceptions import (
|
||||
ControllerError,
|
||||
TrainingFailedError,
|
||||
WorkerGroupError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
RETRYABLE_CONTROLLER_ERRORS = (
|
||||
WorkerGroupStartupFailedError,
|
||||
WorkerGroupStartupTimeoutError,
|
||||
)
|
||||
|
||||
|
||||
class DefaultFailurePolicy(FailurePolicy):
|
||||
def __init__(self, failure_config: FailureConfig):
|
||||
super().__init__(failure_config)
|
||||
self._worker_group_failures = 0
|
||||
self._controller_failures = 0
|
||||
|
||||
def _log_decision(
|
||||
self,
|
||||
decision: FailureDecision,
|
||||
training_failed_error: TrainingFailedError,
|
||||
error_count: int,
|
||||
retry_limit: int,
|
||||
):
|
||||
if isinstance(training_failed_error, ControllerError):
|
||||
error_source = "controller"
|
||||
elif isinstance(training_failed_error, WorkerGroupError):
|
||||
error_source = "worker group"
|
||||
else:
|
||||
raise ValueError(f"Unknown error type: {type(training_failed_error)}")
|
||||
|
||||
logger.info(
|
||||
f"[FailurePolicy] {decision.value}\n"
|
||||
f" Source: {error_source}\n"
|
||||
f" Error count: {error_count} (max allowed: {retry_limit})\n"
|
||||
f"Error: {training_failed_error}",
|
||||
exc_info=(
|
||||
type(training_failed_error),
|
||||
training_failed_error,
|
||||
training_failed_error.__traceback__,
|
||||
),
|
||||
)
|
||||
|
||||
def _is_retryable_error(self, training_failed_error: TrainingFailedError) -> bool:
|
||||
if isinstance(training_failed_error, WorkerGroupError):
|
||||
return True
|
||||
elif isinstance(training_failed_error, ControllerError):
|
||||
return isinstance(
|
||||
training_failed_error.controller_failure, RETRYABLE_CONTROLLER_ERRORS
|
||||
)
|
||||
return False
|
||||
|
||||
def make_decision(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
) -> FailureDecision:
|
||||
|
||||
if not self._is_retryable_error(training_failed_error):
|
||||
decision = FailureDecision.RAISE
|
||||
error_count = 1
|
||||
retry_limit = 0
|
||||
else:
|
||||
if isinstance(training_failed_error, ControllerError):
|
||||
self._controller_failures += 1
|
||||
error_count = self._controller_failures
|
||||
retry_limit = (
|
||||
self.failure_config.controller_failure_limit
|
||||
if self.failure_config.controller_failure_limit != -1
|
||||
else float("inf")
|
||||
)
|
||||
elif isinstance(training_failed_error, WorkerGroupError):
|
||||
self._worker_group_failures += 1
|
||||
error_count = self._worker_group_failures
|
||||
retry_limit = (
|
||||
self.failure_config.max_failures
|
||||
if self.failure_config.max_failures != -1
|
||||
else float("inf")
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown error type: {type(training_failed_error)}")
|
||||
|
||||
if error_count > retry_limit:
|
||||
decision = FailureDecision.RAISE
|
||||
else:
|
||||
decision = FailureDecision.RETRY
|
||||
|
||||
self._log_decision(decision, training_failed_error, error_count, retry_limit)
|
||||
return decision
|
||||
@@ -0,0 +1,13 @@
|
||||
from ray.train import FailureConfig
|
||||
from ray.train.v2._internal.execution.failure_handling import (
|
||||
DefaultFailurePolicy,
|
||||
FailurePolicy,
|
||||
)
|
||||
|
||||
|
||||
def create_failure_policy(failure_config: FailureConfig) -> FailurePolicy:
|
||||
"""Create a failure policy from the given failure config.
|
||||
|
||||
Defaults to the `DefaultFailurePolicy` implementation.
|
||||
"""
|
||||
return DefaultFailurePolicy(failure_config=failure_config)
|
||||
@@ -0,0 +1,29 @@
|
||||
import abc
|
||||
from enum import Enum
|
||||
|
||||
from ray.train.v2.api.config import FailureConfig
|
||||
from ray.train.v2.api.exceptions import TrainingFailedError
|
||||
|
||||
|
||||
class FailureDecision(Enum):
|
||||
RETRY = "RETRY"
|
||||
RAISE = "RAISE"
|
||||
NOOP = "NOOP"
|
||||
|
||||
|
||||
class FailurePolicy(abc.ABC):
|
||||
"""A policy that determines how to handle user and system failures.
|
||||
FailurePolicy will handle the controller failure and worker errors during training.
|
||||
|
||||
This can be used to implement fault tolerance and error recovery.
|
||||
"""
|
||||
|
||||
def __init__(self, failure_config: FailureConfig):
|
||||
self.failure_config = failure_config
|
||||
|
||||
@abc.abstractmethod
|
||||
def make_decision(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
) -> FailureDecision:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,93 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ray.train import Result
|
||||
from ray.train.v2._internal.execution.local_mode.utils import LocalController
|
||||
from ray.train.v2._internal.execution.train_fn_utils import (
|
||||
LocalTrainFnUtils,
|
||||
get_train_fn_utils,
|
||||
set_train_fn_utils,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def has_torchrun_env() -> bool:
|
||||
"""Return True if this process has torch.distributed env vars set.
|
||||
|
||||
For torch.distributed.init_process_group with init_method="env://", these variables are required:
|
||||
- RANK: The rank of the current process
|
||||
- LOCAL_RANK: The local rank of the current process
|
||||
- WORLD_SIZE: Total number of processes participating in the job
|
||||
- LOCAL_WORLD_SIZE: Total number of processes participating in the job on the current node
|
||||
- MASTER_ADDR: The IP address or hostname of the master node (rank 0)
|
||||
- MASTER_PORT: A free port on the master node for communication
|
||||
|
||||
"""
|
||||
torch_dist_required_vars = {
|
||||
"RANK",
|
||||
"LOCAL_RANK",
|
||||
"WORLD_SIZE",
|
||||
"LOCAL_WORLD_SIZE",
|
||||
"MASTER_ADDR",
|
||||
"MASTER_PORT",
|
||||
}
|
||||
|
||||
return torch_dist_required_vars.issubset(os.environ.keys())
|
||||
|
||||
|
||||
class LocalTorchController(LocalController):
|
||||
def _set_train_fn_utils(self) -> None:
|
||||
world_size = 1
|
||||
global_rank = 0
|
||||
local_rank = 0
|
||||
nproc_per_node = 1
|
||||
node_rank = 0
|
||||
if has_torchrun_env():
|
||||
assert not dist.is_initialized(), "torch.distributed is already initialized"
|
||||
torch.distributed.init_process_group(
|
||||
backend="nccl" if torch.cuda.is_available() else "gloo"
|
||||
)
|
||||
world_size = torch.distributed.get_world_size()
|
||||
global_rank = torch.distributed.get_rank()
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(local_rank)
|
||||
nproc_per_node = int(os.environ.get("LOCAL_WORLD_SIZE"))
|
||||
node_rank = global_rank // nproc_per_node
|
||||
|
||||
if world_size != 1:
|
||||
assert (
|
||||
self.datasets is None or len(self.datasets) == 0
|
||||
), "Ray Data is not supported in local mode with multiple workers."
|
||||
set_train_fn_utils(
|
||||
LocalTrainFnUtils(
|
||||
experiment_name=self.experiment_name,
|
||||
world_size=world_size,
|
||||
world_rank=global_rank,
|
||||
local_rank=local_rank,
|
||||
local_world_size=nproc_per_node,
|
||||
node_rank=node_rank,
|
||||
dataset_shards=self.datasets,
|
||||
)
|
||||
)
|
||||
|
||||
def run(self, train_func: Callable[[], Any]) -> Result:
|
||||
self._set_train_fn_utils()
|
||||
train_result = train_func()
|
||||
train_fn_utils = get_train_fn_utils()
|
||||
assert isinstance(train_fn_utils, LocalTrainFnUtils)
|
||||
result = Result(
|
||||
metrics=train_fn_utils._get_last_metrics(),
|
||||
checkpoint=train_fn_utils.get_checkpoint(),
|
||||
path=None,
|
||||
error=None,
|
||||
return_value=train_result,
|
||||
)
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
return result
|
||||
@@ -0,0 +1,41 @@
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
from ray.train import Result
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.v2._internal.execution.train_fn_utils import (
|
||||
LocalTrainFnUtils,
|
||||
get_train_fn_utils,
|
||||
set_train_fn_utils,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LocalController:
|
||||
def __init__(
|
||||
self, experiment_name: str, datasets: Optional[Dict[str, GenDataset]] = None
|
||||
):
|
||||
if datasets is not None:
|
||||
datasets = {k: v() if callable(v) else v for k, v in datasets.items()}
|
||||
|
||||
self.datasets = datasets
|
||||
self.experiment_name = experiment_name
|
||||
|
||||
def run(self, train_func: Callable[[], Any]) -> Result:
|
||||
set_train_fn_utils(
|
||||
LocalTrainFnUtils(
|
||||
experiment_name=self.experiment_name,
|
||||
dataset_shards=self.datasets,
|
||||
)
|
||||
)
|
||||
result = train_func()
|
||||
train_fn_utils = get_train_fn_utils()
|
||||
assert isinstance(train_fn_utils, LocalTrainFnUtils)
|
||||
return Result(
|
||||
metrics=train_fn_utils._get_last_metrics(),
|
||||
checkpoint=train_fn_utils.get_checkpoint(),
|
||||
path=None,
|
||||
error=None,
|
||||
return_value=result,
|
||||
)
|
||||
@@ -0,0 +1,243 @@
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Set
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.constants import DEFAULT_PREEMPTION_POLL_INTERVAL_S
|
||||
from ray.util.tpu import get_tpu_slice_name_from_node
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.worker import RayTrainWorker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PreemptionInfo:
|
||||
"""Information about an imminent preemption event.
|
||||
|
||||
Attributes:
|
||||
deadline_ms: Earliest preemption deadline (UNIX time in milliseconds)
|
||||
across all preempted nodes. ``None`` if no deadline was reported.
|
||||
preempted_node_to_ranks: Map of preempted ``node_id`` to the worker ``world_rank``s affected when that node
|
||||
is preempted.
|
||||
"""
|
||||
|
||||
deadline_ms: Optional[int]
|
||||
preempted_node_to_ranks: Dict[str, List[int]]
|
||||
|
||||
@property
|
||||
def preempted_node_ids(self) -> List[str]:
|
||||
"""Preempted node IDs, sorted lexicographically."""
|
||||
return sorted(self.preempted_node_to_ranks)
|
||||
|
||||
@property
|
||||
def preempted_ranks(self) -> List[int]:
|
||||
"""All affected ranks across the preempted nodes, sorted ascending."""
|
||||
return sorted(
|
||||
{r for ranks in self.preempted_node_to_ranks.values() for r in ranks}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PreemptionContext:
|
||||
"""Thread-shared preemption signal for one worker actor."""
|
||||
|
||||
_preemption_info: Optional[PreemptionInfo] = field(default=None, init=False)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock, init=False)
|
||||
|
||||
def set(self, info: PreemptionInfo) -> None:
|
||||
with self._lock:
|
||||
self._preemption_info = info
|
||||
|
||||
def get(self) -> Optional[PreemptionInfo]:
|
||||
"""Return the current preemption signal, or ``None`` if none received."""
|
||||
with self._lock:
|
||||
return self._preemption_info
|
||||
|
||||
|
||||
def _get_draining_nodes() -> Dict[str, int]:
|
||||
"""Ray Core's draining nodes as ``{node_id_hex: deadline_ms}`` (0 = no deadline)."""
|
||||
return ray._private.state.state.get_draining_nodes()
|
||||
|
||||
|
||||
class PreemptionWatcher:
|
||||
"""Polls Ray Core for node drains and logs detected preemption events.
|
||||
|
||||
One watcher per worker group, spawned as a ``num_cpus=0`` actor by
|
||||
``PreemptionCallback``. The poll loop runs in a background thread. The
|
||||
failure-domain map is built once on construction and is immutable for the
|
||||
watcher's lifetime.
|
||||
|
||||
The failure-domain map records which of our ranks are affected if a node is
|
||||
preempted: for a GPU node, the ranks on that node; for a TPU node, every
|
||||
rank in the node's slice, since a TPU slice is preempted atomically.
|
||||
|
||||
Args:
|
||||
node_to_ranks: Map ``node_id_hex -> [ranks on that node]``. Used both
|
||||
as the set of nodes we care about (drains elsewhere are ignored)
|
||||
and as the seed for failure-domain expansion.
|
||||
poll_interval_s: Seconds between drain-state polls.
|
||||
worker_actors_by_rank: Map ``world_rank -> worker actor handle``. On a
|
||||
detected preemption, ``mark_preempt`` is called on every worker.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
node_to_ranks: Dict[str, List[int]],
|
||||
poll_interval_s: float = DEFAULT_PREEMPTION_POLL_INTERVAL_S,
|
||||
worker_actors_by_rank: Optional[
|
||||
Dict[int, ActorHandle["RayTrainWorker"]]
|
||||
] = None,
|
||||
):
|
||||
self._node_to_ranks: Dict[str, List[int]] = {
|
||||
nid: sorted(ranks) for nid, ranks in node_to_ranks.items()
|
||||
}
|
||||
self._poll_interval_s = poll_interval_s
|
||||
self._worker_actors_by_rank: Dict[int, ActorHandle["RayTrainWorker"]] = (
|
||||
worker_actors_by_rank or {}
|
||||
)
|
||||
self._failure_domain_map: Dict[str, List[int]] = self._build_failure_domain_map(
|
||||
self._node_to_ranks
|
||||
)
|
||||
|
||||
self._stop_event = threading.Event()
|
||||
self._last_drained: Dict[str, int] = {}
|
||||
self._latest_info: Optional[PreemptionInfo] = None
|
||||
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._watch_loop,
|
||||
name="PreemptionWatcher",
|
||||
daemon=True,
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
|
||||
@staticmethod
|
||||
def _build_failure_domain_map(
|
||||
node_to_ranks: Dict[str, List[int]],
|
||||
) -> Dict[str, List[int]]:
|
||||
"""Map each node we host to all ranks in its failure domain.
|
||||
|
||||
- Non-TPU (e.g. GPU) clusters: the failure domain is the node itself,
|
||||
so a drain on a node flags only the ranks this job runs there.
|
||||
- TPU multislice: every host in a slice is reclaimed atomically, so a
|
||||
drain on any host is fate-shared with the rest.
|
||||
"""
|
||||
per_node = {nid: sorted(set(ranks)) for nid, ranks in node_to_ranks.items()}
|
||||
|
||||
try:
|
||||
all_nodes = ray.nodes()
|
||||
|
||||
# Slice label for each node we host (None for non-TPU nodes).
|
||||
node_to_slice: Dict[str, Optional[str]] = {
|
||||
node["NodeID"]: get_tpu_slice_name_from_node(node)
|
||||
for node in all_nodes
|
||||
if node["NodeID"] in node_to_ranks
|
||||
}
|
||||
|
||||
# Union our ranks per slice.
|
||||
slice_to_ranks: Dict[str, Set[int]] = {}
|
||||
for node_id, ranks in node_to_ranks.items():
|
||||
slice_label = node_to_slice.get(node_id)
|
||||
if slice_label:
|
||||
slice_to_ranks.setdefault(slice_label, set()).update(ranks)
|
||||
|
||||
# Non-TPU cluster (or none of our nodes are on a slice): per-node.
|
||||
if not slice_to_ranks:
|
||||
return per_node
|
||||
|
||||
result: Dict[str, List[int]] = {}
|
||||
for node_id, ranks in node_to_ranks.items():
|
||||
slice_label = node_to_slice.get(node_id)
|
||||
if slice_label:
|
||||
result[node_id] = sorted(slice_to_ranks[slice_label])
|
||||
else:
|
||||
result[node_id] = sorted(set(ranks))
|
||||
return result
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Could not build failure-domain map; falling back to per-node "
|
||||
"domains (no TPU-slice expansion).",
|
||||
exc_info=True,
|
||||
)
|
||||
return per_node
|
||||
|
||||
def get_latest_preemption_info(self) -> Optional[PreemptionInfo]:
|
||||
"""Most recent :class:`PreemptionInfo` observed, or ``None``."""
|
||||
return self._latest_info
|
||||
|
||||
def _watch_loop(self) -> None:
|
||||
logger.debug(
|
||||
"PreemptionWatcher polling %d node(s) every %.1fs.",
|
||||
len(self._node_to_ranks),
|
||||
self._poll_interval_s,
|
||||
)
|
||||
while not self._stop_event.is_set():
|
||||
self._poll_once()
|
||||
self._stop_event.wait(timeout=self._poll_interval_s)
|
||||
logger.debug("PreemptionWatcher stopped.")
|
||||
|
||||
def _poll_once(self) -> None:
|
||||
"""Poll the drain source once and dispatch on change.
|
||||
|
||||
Per-poll exceptions are caught and logged so a transient GCS hiccup
|
||||
doesn't kill the watcher loop.
|
||||
"""
|
||||
try:
|
||||
drained = _get_draining_nodes() or {}
|
||||
# Keep only drains on this job's own nodes (others are ignored).
|
||||
# That's complete for TPU — an SPMD job fully occupies its slice, so
|
||||
# every fate-shared host is one of our nodes and a drain on any slice
|
||||
# host appears here. For GPU, a drain on a host we don't run on is
|
||||
# correctly irrelevant.
|
||||
relevant = {
|
||||
n: d for n, d in drained.items() if n in self._failure_domain_map
|
||||
}
|
||||
if relevant != self._last_drained:
|
||||
self._on_drain_change(relevant)
|
||||
self._last_drained = relevant
|
||||
except Exception:
|
||||
# TODO(lehui): consider exponential backoff when the drain API keeps
|
||||
# failing, instead of retrying at the fixed poll interval.
|
||||
logger.warning("PreemptionWatcher poll failed", exc_info=True)
|
||||
|
||||
def _on_drain_change(self, drained: Dict[str, int]) -> None:
|
||||
"""Handle a change in the drained-node set.
|
||||
|
||||
``drained`` has already been narrowed to this job's nodes by the
|
||||
caller (``_poll_once``).
|
||||
"""
|
||||
if not drained:
|
||||
return
|
||||
|
||||
affected_node_ids = sorted(drained.keys())
|
||||
preempted_node_to_ranks = {
|
||||
node_id: self._failure_domain_map[node_id] for node_id in affected_node_ids
|
||||
}
|
||||
|
||||
# Earliest deadline across the preempted nodes; None if none reported one
|
||||
# (Ray Core uses 0 for "no deadline", which is falsy and filtered out).
|
||||
reported_deadlines = [drained[n] for n in affected_node_ids if drained[n]]
|
||||
deadline_ms = min(reported_deadlines) if reported_deadlines else None
|
||||
|
||||
info = PreemptionInfo(
|
||||
deadline_ms=deadline_ms,
|
||||
preempted_node_to_ranks=preempted_node_to_ranks,
|
||||
)
|
||||
self._latest_info = info
|
||||
|
||||
logger.warning(
|
||||
"PreemptionWatcher: preemption detected — "
|
||||
"preempted_node_ids=%s, preempted_ranks=%s, deadline_ms=%s",
|
||||
info.preempted_node_ids,
|
||||
info.preempted_ranks,
|
||||
deadline_ms,
|
||||
)
|
||||
|
||||
for rank, actor in self._worker_actors_by_rank.items():
|
||||
actor.mark_preempt.remote(info)
|
||||
# TODO(lehui): coalesce preemptions seen within one window into a single
|
||||
# worker-group restart, so a staggered drain (node A at t, node B at
|
||||
# t+60s) doesn't cause back-to-back restarts.
|
||||
@@ -0,0 +1,29 @@
|
||||
# isort: off
|
||||
from .scaling_policy import ScalingDecision, ScalingPolicy, NoopDecision, ResizeDecision
|
||||
from .scaling_policy import (
|
||||
AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
|
||||
AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
|
||||
AUTOSCALING_REQUESTS_INTERVAL_S,
|
||||
)
|
||||
from .elastic import ElasticScalingPolicy
|
||||
from .fixed import FixedScalingPolicy
|
||||
from .factory import create_scaling_policy
|
||||
|
||||
# isort: on
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AUTOSCALING_REQUESTS_EXPIRE_TIME_S",
|
||||
"AUTOSCALING_REQUESTS_GET_TIMEOUT_S",
|
||||
"AUTOSCALING_REQUESTS_INTERVAL_S",
|
||||
"ScalingPolicy",
|
||||
"ElasticScalingPolicy",
|
||||
"FixedScalingPolicy",
|
||||
"ScalingDecision",
|
||||
"NoopDecision",
|
||||
"ResizeDecision",
|
||||
"create_scaling_policy",
|
||||
]
|
||||
|
||||
|
||||
# DO NOT ADD ANYTHING AFTER THIS LINE.
|
||||
@@ -0,0 +1,291 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.scaling_policy import (
|
||||
NoopDecision,
|
||||
ResizeDecision,
|
||||
ScalingDecision,
|
||||
ScalingPolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroupPollStatus,
|
||||
WorkerGroupState,
|
||||
)
|
||||
from ray.train.v2._internal.util import time_monotonic
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
|
||||
ResourceDict,
|
||||
)
|
||||
|
||||
|
||||
class ElasticScalingPolicy(ScalingPolicy):
|
||||
|
||||
# Minimum interval in seconds between querying the AutoscalingCoordinator for allocated resources.
|
||||
GET_ALLOCATED_RESOURCES_INTERVAL_S = 1
|
||||
# Minimum interval in seconds between logging warnings about insufficient workers.
|
||||
INSUFFICIENT_WORKERS_WARNING_INTERVAL_S = 30
|
||||
|
||||
def __init__(self, scaling_config: ScalingConfig):
|
||||
super().__init__(scaling_config)
|
||||
|
||||
self._latest_monitor_time = float("-inf")
|
||||
self._latest_insufficient_workers_warning_time = float("-inf")
|
||||
self._latest_allocated_resources_query_time = float("-inf")
|
||||
self._latest_allocated_resources: Optional[List["ResourceDict"]] = None
|
||||
|
||||
def _get_num_workers_for_resource_request(self) -> int:
|
||||
return self.scaling_config.max_workers
|
||||
|
||||
def _count_possible_workers(
|
||||
self, allocated_resources: List[Dict[str, float]]
|
||||
) -> int:
|
||||
"""Count the number of workers that can be started/restarted with the given
|
||||
the list of node resources. The returned number is capped at the maximum
|
||||
number of workers.
|
||||
|
||||
For GPUs, this divides raw allocated resources by per-worker requirements.
|
||||
For TPUs, an additional check ensures workers align with physically intact
|
||||
TPU slices (see ``_get_strict_tpu_worker_count``).
|
||||
|
||||
Args:
|
||||
allocated_resources: The resources currently allocated by the AutoscalingCoordinator.
|
||||
|
||||
Returns:
|
||||
The number of workers that can be started/restarted with the current resources.
|
||||
"""
|
||||
# TODO: Fractional resources do not work well here.
|
||||
single_worker_resources = self.scaling_config._resources_per_worker_not_none
|
||||
total_num_workers = 0
|
||||
|
||||
# If workers require no resources, we can run as many as we want.
|
||||
if sum(single_worker_resources.values()) == 0:
|
||||
return self.scaling_config.max_workers
|
||||
|
||||
for resources in allocated_resources:
|
||||
num_workers = min(
|
||||
[
|
||||
resources.get(resource, 0.0) // single_worker_resources[resource]
|
||||
for resource in single_worker_resources
|
||||
if single_worker_resources[resource] > 0
|
||||
]
|
||||
)
|
||||
total_num_workers += num_workers
|
||||
|
||||
total_num_workers = min(int(total_num_workers), self.scaling_config.max_workers)
|
||||
|
||||
# Multi-host TPUs are scheduled atomically in interconnected slices defined by a topology.
|
||||
if (
|
||||
self.scaling_config.use_tpu
|
||||
and self.scaling_config.topology
|
||||
and self.scaling_config.accelerator_type
|
||||
):
|
||||
total_num_workers = self._get_strict_tpu_worker_count(
|
||||
total_num_workers=total_num_workers,
|
||||
)
|
||||
|
||||
return total_num_workers
|
||||
|
||||
def _get_strict_tpu_worker_count(self, total_num_workers: int) -> int:
|
||||
"""Calculate the number of workers that can run on intact TPU slices.
|
||||
|
||||
The Autoscaler's allocated resources might overestimate the number of
|
||||
schedulable TPU workers because it counts raw resources. TPUs require
|
||||
atomic, interconnected slices. This function checks the cluster for
|
||||
physically intact slices to prevent scaling onto fractional/broken
|
||||
topologies.
|
||||
|
||||
The worker count is: min(resource_based_slices, intact_slices) *
|
||||
workers_per_slice, where resource_based_slices =
|
||||
total_num_workers // workers_per_slice.
|
||||
|
||||
Args:
|
||||
total_num_workers: The initial estimate of workers based on raw
|
||||
allocated resources.
|
||||
|
||||
Returns:
|
||||
The number of workers aligned to fully intact TPU slices.
|
||||
"""
|
||||
from ray.util.tpu import get_num_tpu_slices, get_tpu_worker_resources
|
||||
|
||||
single_worker_resources = self.scaling_config._resources_per_worker_not_none
|
||||
|
||||
try:
|
||||
workers_per_slice, _ = get_tpu_worker_resources(
|
||||
topology=self.scaling_config.topology,
|
||||
accelerator_type=self.scaling_config.accelerator_type,
|
||||
resources_per_unit=single_worker_resources,
|
||||
num_slices=1,
|
||||
)
|
||||
|
||||
if workers_per_slice == 0:
|
||||
# A single worker requires more resources than exist in a
|
||||
# full slice — impossible scheduling configuration for TPU.
|
||||
return 0
|
||||
|
||||
num_slices_from_resources = total_num_workers // workers_per_slice
|
||||
|
||||
if num_slices_from_resources > 0:
|
||||
try:
|
||||
num_intact_slices = get_num_tpu_slices(
|
||||
topology=self.scaling_config.topology,
|
||||
accelerator_type=self.scaling_config.accelerator_type,
|
||||
)
|
||||
num_slices_from_resources = min(
|
||||
num_slices_from_resources, num_intact_slices
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to check cluster state for intact TPU slices: {e}"
|
||||
)
|
||||
|
||||
return num_slices_from_resources * workers_per_slice
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Could not calculate TPU slice boundaries for elastic scaling: {e}. "
|
||||
"Worker counts may not align with TPU topology."
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
def _get_resize_decision(self, num_workers: int) -> ResizeDecision:
|
||||
return ResizeDecision(
|
||||
num_workers=num_workers,
|
||||
resources_per_worker=self.scaling_config._resources_per_worker_not_none,
|
||||
)
|
||||
|
||||
def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
|
||||
self._maybe_send_resource_request()
|
||||
|
||||
allocated_resources = self._get_allocated_resources()
|
||||
if allocated_resources is None:
|
||||
return NoopDecision()
|
||||
|
||||
num_workers = self._count_possible_workers(allocated_resources)
|
||||
|
||||
if num_workers < self.scaling_config.min_workers:
|
||||
now = time_monotonic()
|
||||
# Only log this warning periodically to avoid spamming logs
|
||||
if (
|
||||
now - self._latest_insufficient_workers_warning_time
|
||||
>= self.INSUFFICIENT_WORKERS_WARNING_INTERVAL_S
|
||||
):
|
||||
logger.info(
|
||||
f"Detected ready resources for {num_workers} workers "
|
||||
"in the cluster. "
|
||||
"Deciding NOT to start/restart training due to the number of workers "
|
||||
"falling below the minimum "
|
||||
f"(min_workers={self.scaling_config.min_workers})."
|
||||
)
|
||||
self._latest_insufficient_workers_warning_time = now
|
||||
return NoopDecision()
|
||||
|
||||
logger.info(
|
||||
f"Detected ready resources for {num_workers} workers "
|
||||
"in the cluster. "
|
||||
"Deciding to start/restart training with this worker group size."
|
||||
)
|
||||
return self._get_resize_decision(num_workers)
|
||||
|
||||
def make_decision_for_running_worker_group(
|
||||
self,
|
||||
worker_group_state: WorkerGroupState,
|
||||
worker_group_status: WorkerGroupPollStatus,
|
||||
) -> ScalingDecision:
|
||||
self._maybe_send_resource_request()
|
||||
|
||||
# Ensure that we don't make resizing decisions too frequently.
|
||||
# The latest restart time and the latest monitor time (whichever is later)
|
||||
# determine the time of the next resize consideration.
|
||||
latest_consideration_time = max(
|
||||
worker_group_state.start_time, self._latest_monitor_time
|
||||
)
|
||||
|
||||
now = time_monotonic()
|
||||
time_since_latest_consideration = now - latest_consideration_time
|
||||
if (
|
||||
time_since_latest_consideration
|
||||
< self.scaling_config.elastic_resize_monitor_interval_s
|
||||
):
|
||||
logger.debug(
|
||||
"Skipping resize decision due to the latest resizing consideration "
|
||||
"happening too recently: "
|
||||
"%.2f seconds < ScalingConfig(elastic_resize_monitor_interval_s=%.2f seconds).",
|
||||
time_since_latest_consideration,
|
||||
self.scaling_config.elastic_resize_monitor_interval_s,
|
||||
)
|
||||
return NoopDecision()
|
||||
|
||||
self._latest_monitor_time = now
|
||||
|
||||
allocated_resources = self._get_allocated_resources()
|
||||
if allocated_resources is None:
|
||||
return NoopDecision()
|
||||
|
||||
num_workers = self._count_possible_workers(allocated_resources)
|
||||
|
||||
if num_workers == worker_group_state.num_workers:
|
||||
logger.info(
|
||||
"Did not detect any changes in the cluster resources. "
|
||||
"Training will continue with the same worker group size "
|
||||
f"({num_workers})."
|
||||
)
|
||||
return NoopDecision()
|
||||
elif num_workers < self.scaling_config.min_workers:
|
||||
# This covers an edge case where allocated resources decrease to less
|
||||
# than the minimum number of workers.
|
||||
# This situation is rare, since cluster downsizing typically involves
|
||||
# worker failures. However, this check is still useful to fully
|
||||
# avoid entering an invalid state with fewer workers than the minimum.
|
||||
return NoopDecision()
|
||||
|
||||
logger.info(
|
||||
"Detected changes in the cluster resources. "
|
||||
"Deciding to resize the worker group from "
|
||||
f"{worker_group_state.num_workers} -> {num_workers} workers."
|
||||
)
|
||||
return self._get_resize_decision(num_workers)
|
||||
|
||||
# ---------------------------------------------------
|
||||
# Methods for interacting with AutoscalingCoordinator
|
||||
# ---------------------------------------------------
|
||||
|
||||
def _get_allocated_resources(self) -> Optional[List["ResourceDict"]]:
|
||||
"""Get allocated resources from AutoscalingCoordinator.
|
||||
Return None if there is an error."""
|
||||
now = time_monotonic()
|
||||
time_since_last_call = now - self._latest_allocated_resources_query_time
|
||||
|
||||
if time_since_last_call < self.GET_ALLOCATED_RESOURCES_INTERVAL_S:
|
||||
return self._latest_allocated_resources
|
||||
|
||||
allocated_resources = None
|
||||
try:
|
||||
allocated_resources = ray.get(
|
||||
self._autoscaling_coordinator.get_allocated_resources.remote(
|
||||
self._requester_id
|
||||
),
|
||||
timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
|
||||
)
|
||||
except Exception:
|
||||
msg = (
|
||||
f"Failed to get allocated resources for {self._requester_id}."
|
||||
" Will not resize the worker group."
|
||||
" If this only happens transiently during network partition or"
|
||||
" CPU being overloaded, it's safe to ignore this error."
|
||||
" If this error persists, file a GitHub issue."
|
||||
)
|
||||
logger.warning(msg, exc_info=True)
|
||||
finally:
|
||||
self._latest_allocated_resources_query_time = time_monotonic()
|
||||
self._latest_allocated_resources = allocated_resources
|
||||
|
||||
return self._latest_allocated_resources
|
||||
@@ -0,0 +1,16 @@
|
||||
from ray.train.v2._internal.execution.scaling_policy import (
|
||||
ElasticScalingPolicy,
|
||||
FixedScalingPolicy,
|
||||
ScalingPolicy,
|
||||
)
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
|
||||
|
||||
def create_scaling_policy(scaling_config: ScalingConfig) -> ScalingPolicy:
|
||||
"""Create a scaling policy from the given scaling config.
|
||||
|
||||
Defaults to the `FixedScalingPolicy` implementation.
|
||||
"""
|
||||
if scaling_config.elasticity_enabled:
|
||||
return ElasticScalingPolicy(scaling_config=scaling_config)
|
||||
return FixedScalingPolicy(scaling_config=scaling_config)
|
||||
@@ -0,0 +1,30 @@
|
||||
from ray.train.v2._internal.execution.scaling_policy import (
|
||||
NoopDecision,
|
||||
ResizeDecision,
|
||||
ScalingDecision,
|
||||
ScalingPolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroupPollStatus,
|
||||
WorkerGroupState,
|
||||
)
|
||||
|
||||
|
||||
class FixedScalingPolicy(ScalingPolicy):
|
||||
def _get_num_workers_for_resource_request(self) -> int:
|
||||
return self.scaling_config.num_workers
|
||||
|
||||
def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
|
||||
self._maybe_send_resource_request()
|
||||
return ResizeDecision(
|
||||
num_workers=self.scaling_config.num_workers,
|
||||
resources_per_worker=self.scaling_config._resources_per_worker_not_none,
|
||||
)
|
||||
|
||||
def make_decision_for_running_worker_group(
|
||||
self,
|
||||
worker_group_state: WorkerGroupState,
|
||||
worker_group_status: WorkerGroupPollStatus,
|
||||
) -> ScalingDecision:
|
||||
self._maybe_send_resource_request()
|
||||
return NoopDecision()
|
||||
@@ -0,0 +1,183 @@
|
||||
import abc
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import Dict
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.execution.callback import ControllerCallback
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroupPollStatus,
|
||||
WorkerGroupState,
|
||||
)
|
||||
from ray.train.v2._internal.util import time_monotonic
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The time in seconds after which an autoscaling request will expire.
|
||||
AUTOSCALING_REQUESTS_EXPIRE_TIME_S = 180
|
||||
# Timeout in seconds for getting the result of a call to the AutoscalingCoordinator.
|
||||
AUTOSCALING_REQUESTS_GET_TIMEOUT_S = 5
|
||||
# Interval in seconds between resource requests to the AutoscalingCoordinator.
|
||||
AUTOSCALING_REQUESTS_INTERVAL_S = 20
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScalingDecision:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class NoopDecision(ScalingDecision):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResizeDecision(ScalingDecision):
|
||||
num_workers: int
|
||||
resources_per_worker: Dict[str, float]
|
||||
|
||||
|
||||
class ScalingPolicy(abc.ABC, ControllerCallback):
|
||||
"""A policy that determines when and how to scale a worker group.
|
||||
|
||||
This can be used to implement elasticity and fault tolerance.
|
||||
|
||||
Recovery decisions are made when workers are in an inactive or unhealthy state.
|
||||
Upscale decisions are optional and are made when workers are healthy.
|
||||
|
||||
Note: When adding new scaling policies, revisit the shared defaults- particularly if:
|
||||
- AutoscalingCoordinator integration is not needed or a different interface
|
||||
becomes available
|
||||
- Timeout/expiry constants need to diverge between policies
|
||||
- _get_num_workers_for_resource_request() needs variable worker counts
|
||||
- Controller lifecycle behavior diverges
|
||||
"""
|
||||
|
||||
# TODO: Restructure these APIs to consider different TrainControllerStates
|
||||
# instead of just running and non-running worker groups.
|
||||
|
||||
def __init__(self, scaling_config: ScalingConfig):
|
||||
self.scaling_config = scaling_config
|
||||
self._requester_id = "train-" + uuid.uuid4().hex
|
||||
self._latest_autoscaling_request_time = float("-inf")
|
||||
|
||||
@abc.abstractmethod
|
||||
def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
|
||||
"""Makes a scaling decision when the worker group is initializing
|
||||
or recovering from an error."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def make_decision_for_running_worker_group(
|
||||
self,
|
||||
worker_group_state: WorkerGroupState,
|
||||
worker_group_status: WorkerGroupPollStatus,
|
||||
) -> ScalingDecision:
|
||||
"""Makes a scaling decision when monitoring healthy, running workers."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def _get_num_workers_for_resource_request(self) -> int:
|
||||
"""Return the number of workers to request resources for."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ---------------------------------------------------
|
||||
# Methods for interacting with AutoscalingCoordinator
|
||||
# ---------------------------------------------------
|
||||
|
||||
@cached_property
|
||||
def _autoscaling_coordinator(self):
|
||||
from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
|
||||
get_or_create_autoscaling_coordinator,
|
||||
)
|
||||
|
||||
return get_or_create_autoscaling_coordinator()
|
||||
|
||||
def _maybe_send_resource_request(self):
|
||||
"""Send a resource request to AutoscalingCoordinator,
|
||||
if AUTOSCALING_REQUESTS_INTERVAL_S has passed since the last send."""
|
||||
now = time_monotonic()
|
||||
if (
|
||||
now - self._latest_autoscaling_request_time
|
||||
< AUTOSCALING_REQUESTS_INTERVAL_S
|
||||
):
|
||||
return
|
||||
self._send_resource_request()
|
||||
|
||||
def _send_resource_request(self):
|
||||
"""Register training resources with the AutoscalingCoordinator."""
|
||||
resources_per_worker = self.scaling_config._resources_per_worker_not_none
|
||||
num_workers = self._get_num_workers_for_resource_request()
|
||||
label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
|
||||
try:
|
||||
from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
|
||||
ResourceRequestPriority,
|
||||
)
|
||||
|
||||
ray.get(
|
||||
self._autoscaling_coordinator.request_resources.remote(
|
||||
requester_id=self._requester_id,
|
||||
resources=[resources_per_worker] * num_workers,
|
||||
label_selectors=label_selectors,
|
||||
expire_after_s=AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
|
||||
priority=ResourceRequestPriority.HIGH,
|
||||
),
|
||||
timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
|
||||
)
|
||||
self._latest_autoscaling_request_time = time_monotonic()
|
||||
except Exception:
|
||||
msg = (
|
||||
f"Failed to send resource request for {self._requester_id}."
|
||||
" If this only happens transiently during network partition or"
|
||||
" CPU being overloaded, it's safe to ignore this error."
|
||||
" If this error persists, file a GitHub issue."
|
||||
)
|
||||
logger.warning(msg, exc_info=True)
|
||||
|
||||
def _cancel_resource_request(self):
|
||||
"""Cancel the resource request to AutoscalingCoordinator."""
|
||||
try:
|
||||
ray.get(
|
||||
self._autoscaling_coordinator.cancel_request.remote(
|
||||
requester_id=self._requester_id,
|
||||
),
|
||||
timeout=AUTOSCALING_REQUESTS_GET_TIMEOUT_S,
|
||||
)
|
||||
except Exception:
|
||||
msg = (
|
||||
f"Failed to cancel resource request for {self._requester_id}."
|
||||
" The request will still expire after the timeout of"
|
||||
f" {AUTOSCALING_REQUESTS_EXPIRE_TIME_S} seconds."
|
||||
)
|
||||
logger.warning(msg, exc_info=True)
|
||||
|
||||
# --------------------------
|
||||
# ControllerCallback
|
||||
# --------------------------
|
||||
|
||||
def after_controller_start(self, train_run_context: TrainRunContext):
|
||||
"""Register training resources with the AutoscalingCoordinator."""
|
||||
self._requester_id = f"train-{train_run_context.run_id}"
|
||||
resources_per_worker = self.scaling_config._resources_per_worker_not_none
|
||||
num_workers = self._get_num_workers_for_resource_request()
|
||||
label_selectors = self.scaling_config._label_selector_per_worker(num_workers)
|
||||
if label_selectors:
|
||||
logger.info(
|
||||
f"Requesting resources: {resources_per_worker} * {num_workers} "
|
||||
f"with label_selectors={label_selectors}"
|
||||
)
|
||||
else:
|
||||
logger.info(f"Requesting resources: {resources_per_worker} * {num_workers}")
|
||||
self._send_resource_request()
|
||||
|
||||
async def before_controller_shutdown(self):
|
||||
"""Cancel the resource request when the controller shuts down."""
|
||||
self._cancel_resource_request()
|
||||
|
||||
def before_controller_abort(self):
|
||||
"""Cancel the resource request when the controller is aborted."""
|
||||
self._cancel_resource_request()
|
||||
@@ -0,0 +1,573 @@
|
||||
# Try import ray[train] core requirements (defined in setup.py)
|
||||
# isort: off
|
||||
try:
|
||||
import fsspec # noqa
|
||||
from fsspec.implementations.local import LocalFileSystem
|
||||
|
||||
except (ImportError, ModuleNotFoundError) as e:
|
||||
raise RuntimeError(
|
||||
"fsspec is a required dependency of Ray Train and Ray Tune. "
|
||||
"Please install with: `pip install fsspec`"
|
||||
) from e
|
||||
|
||||
try:
|
||||
import pyarrow
|
||||
import pyarrow.fs
|
||||
|
||||
except (ImportError, ModuleNotFoundError) as e:
|
||||
raise RuntimeError(
|
||||
"pyarrow is a required dependency of Ray Train and Ray Tune. "
|
||||
"Please install with: `pip install pyarrow`"
|
||||
) from e
|
||||
# isort: on
|
||||
|
||||
import fnmatch
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
|
||||
|
||||
from ray.air._internal.filelock import TempFileLock
|
||||
from ray.train.constants import _get_ray_train_session_dir
|
||||
from ray.train.v2._internal.constants import (
|
||||
CHECKPOINT_MANAGER_SNAPSHOT_FILENAME,
|
||||
VALIDATE_STORAGE_MARKER_FILENAME,
|
||||
)
|
||||
from ray.train.v2._internal.util import date_str
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _ExcludingLocalFilesystem(LocalFileSystem):
|
||||
"""LocalFileSystem wrapper to exclude files according to patterns.
|
||||
|
||||
Args:
|
||||
root_path: Root path to strip when matching with the exclude pattern.
|
||||
Ex: root_path="/tmp/a/b/c", exclude=["*a*"], will exclude
|
||||
/tmp/a/b/c/_a_.txt but not ALL of /tmp/a/*.
|
||||
exclude: List of patterns that are applied to files returned by
|
||||
``self.find()``. If a file path matches this pattern, it will
|
||||
be excluded.
|
||||
**kwargs: Additional keyword arguments forwarded to
|
||||
``pyarrow.fs.LocalFileSystem``.
|
||||
"""
|
||||
|
||||
def __init__(self, root_path: Path, exclude: List[str], **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._exclude = exclude
|
||||
self._root_path = root_path
|
||||
|
||||
@property
|
||||
def fsid(self):
|
||||
return "_excluding_local"
|
||||
|
||||
def _should_exclude(self, path: str) -> bool:
|
||||
"""Return True if `path` (relative to `root_path`) matches any of the
|
||||
`self._exclude` patterns."""
|
||||
path = Path(path)
|
||||
relative_path = path.relative_to(self._root_path).as_posix()
|
||||
match_candidates = [relative_path]
|
||||
if path.is_dir():
|
||||
# Everything is in posix path format ('/')
|
||||
match_candidates.append(relative_path + "/")
|
||||
|
||||
for excl in self._exclude:
|
||||
if any(fnmatch.fnmatch(candidate, excl) for candidate in match_candidates):
|
||||
return True
|
||||
return False
|
||||
|
||||
def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs):
|
||||
"""Call parent find() and exclude from result."""
|
||||
paths = super().find(
|
||||
path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, **kwargs
|
||||
)
|
||||
if detail:
|
||||
return {
|
||||
path: out
|
||||
for path, out in paths.items()
|
||||
if not self._should_exclude(path)
|
||||
}
|
||||
else:
|
||||
return [path for path in paths if not self._should_exclude(path)]
|
||||
|
||||
|
||||
def _pyarrow_fs_copy_files(
|
||||
source, destination, source_filesystem=None, destination_filesystem=None, **kwargs
|
||||
):
|
||||
if isinstance(destination_filesystem, pyarrow.fs.S3FileSystem):
|
||||
# Workaround multi-threading issue with pyarrow. Note that use_threads=True
|
||||
# is safe for download, just not for uploads, see:
|
||||
# https://github.com/apache/arrow/issues/32372
|
||||
kwargs.setdefault("use_threads", False)
|
||||
|
||||
# Use a large chunk size to speed up large checkpoint transfers.
|
||||
kwargs.setdefault("chunk_size", 64 * 1024 * 1024)
|
||||
|
||||
return pyarrow.fs.copy_files(
|
||||
source,
|
||||
destination,
|
||||
source_filesystem=source_filesystem,
|
||||
destination_filesystem=destination_filesystem,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
# TODO(justinvyu): Add unit tests for all these utils.
|
||||
|
||||
|
||||
def delete_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str):
|
||||
"""Deletes (fs, fs_path) or raises FileNotFoundError if it doesn't exist."""
|
||||
is_dir = _is_directory(fs, fs_path)
|
||||
|
||||
try:
|
||||
if is_dir:
|
||||
fs.delete_dir(fs_path)
|
||||
else:
|
||||
fs.delete_file(fs_path)
|
||||
except Exception:
|
||||
logger.exception(f"Caught exception when deleting path at ({fs}, {fs_path}):")
|
||||
|
||||
|
||||
def _download_from_fs_path(
|
||||
fs: pyarrow.fs.FileSystem,
|
||||
fs_path: str,
|
||||
local_path: str,
|
||||
filelock: bool = True,
|
||||
):
|
||||
"""Downloads a directory or file from (fs, fs_path) to a local path.
|
||||
|
||||
If fs_path points to a directory:
|
||||
- The full directory contents are downloaded directly into `local_path`,
|
||||
rather than to a subdirectory of `local_path`.
|
||||
|
||||
If fs_path points to a file:
|
||||
- The file is downloaded to `local_path`, which is expected to be a file path.
|
||||
|
||||
If the download fails, the `local_path` contents are
|
||||
cleaned up before raising, if the directory did not previously exist.
|
||||
|
||||
NOTE: This method creates `local_path`'s parent directories if they do not
|
||||
already exist. If the download fails, this does NOT clean up all the parent
|
||||
directories that were created.
|
||||
|
||||
Args:
|
||||
fs: The filesystem to download from.
|
||||
fs_path: The filesystem path (either a directory or a file) to download.
|
||||
local_path: The local path to download to.
|
||||
filelock: Whether to require a file lock before downloading, useful for
|
||||
multiple downloads to the same directory that may be happening in parallel.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: if (fs, fs_path) doesn't exist.
|
||||
"""
|
||||
|
||||
_local_path = Path(local_path).resolve()
|
||||
exists_before = _local_path.exists()
|
||||
if _is_directory(fs=fs, fs_path=fs_path):
|
||||
_local_path.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
_local_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
if filelock:
|
||||
with TempFileLock(f"{os.path.normpath(local_path)}.lock"):
|
||||
_pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
|
||||
else:
|
||||
_pyarrow_fs_copy_files(fs_path, local_path, source_filesystem=fs)
|
||||
except Exception as e:
|
||||
# Clean up the directory if downloading was unsuccessful
|
||||
if not exists_before:
|
||||
shutil.rmtree(local_path, ignore_errors=True)
|
||||
raise e
|
||||
|
||||
|
||||
def _upload_to_fs_path(
|
||||
local_path: str,
|
||||
fs: pyarrow.fs.FileSystem,
|
||||
fs_path: str,
|
||||
exclude: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Uploads a local directory or file to (fs, fs_path).
|
||||
|
||||
NOTE: This will create all necessary parent directories at the destination.
|
||||
|
||||
Args:
|
||||
local_path: The local path to upload.
|
||||
fs: The filesystem to upload to.
|
||||
fs_path: The filesystem path where the dir/file will be uploaded to.
|
||||
exclude: A list of filename matches to exclude from upload. This includes
|
||||
all files under subdirectories as well.
|
||||
This pattern will match with the relative paths of all files under
|
||||
`local_path`.
|
||||
Ex: ["*.png"] to exclude all .png images.
|
||||
"""
|
||||
|
||||
if not exclude:
|
||||
# TODO(justinvyu): uploading a single file doesn't work
|
||||
# (since we always create a directory at fs_path)
|
||||
_create_directory(fs=fs, fs_path=fs_path)
|
||||
_pyarrow_fs_copy_files(local_path, fs_path, destination_filesystem=fs)
|
||||
return
|
||||
|
||||
_upload_to_uri_with_exclude_fsspec(
|
||||
local_path=local_path, fs=fs, fs_path=fs_path, exclude=exclude
|
||||
)
|
||||
|
||||
|
||||
def _upload_to_uri_with_exclude_fsspec(
|
||||
local_path: str, fs: "pyarrow.fs", fs_path: str, exclude: Optional[List[str]]
|
||||
) -> None:
|
||||
local_fs = _ExcludingLocalFilesystem(root_path=local_path, exclude=exclude)
|
||||
handler = pyarrow.fs.FSSpecHandler(local_fs)
|
||||
source_fs = pyarrow.fs.PyFileSystem(handler)
|
||||
|
||||
_create_directory(fs=fs, fs_path=fs_path)
|
||||
_pyarrow_fs_copy_files(
|
||||
local_path, fs_path, source_filesystem=source_fs, destination_filesystem=fs
|
||||
)
|
||||
|
||||
|
||||
def _list_at_fs_path(
|
||||
fs: pyarrow.fs.FileSystem,
|
||||
fs_path: str,
|
||||
file_filter: Callable[[pyarrow.fs.FileInfo], bool] = lambda x: True,
|
||||
) -> List[str]:
|
||||
"""Returns the list of filenames at (fs, fs_path), similar to os.listdir.
|
||||
|
||||
If the path doesn't exist, returns an empty list.
|
||||
"""
|
||||
selector = pyarrow.fs.FileSelector(fs_path, allow_not_found=True, recursive=False)
|
||||
return [
|
||||
os.path.relpath(file_info.path.lstrip("/"), start=fs_path.lstrip("/"))
|
||||
for file_info in fs.get_file_info(selector)
|
||||
if file_filter(file_info)
|
||||
]
|
||||
|
||||
|
||||
def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
|
||||
"""Returns True if (fs, fs_path) exists."""
|
||||
|
||||
valid = fs.get_file_info(fs_path)
|
||||
return valid.type != pyarrow.fs.FileType.NotFound
|
||||
|
||||
|
||||
def _is_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> bool:
|
||||
"""Checks if (fs, fs_path) is a directory or a file.
|
||||
|
||||
Args:
|
||||
fs: The filesystem to query.
|
||||
fs_path: The path on the filesystem.
|
||||
|
||||
Returns:
|
||||
``True`` if the path is a directory, ``False`` if it is a file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: if (fs, fs_path) doesn't exist.
|
||||
"""
|
||||
|
||||
file_info = fs.get_file_info(fs_path)
|
||||
if file_info.type == pyarrow.fs.FileType.NotFound:
|
||||
raise FileNotFoundError(f"Path not found: ({fs}, {fs_path})")
|
||||
|
||||
return not file_info.is_file
|
||||
|
||||
|
||||
def _create_directory(fs: pyarrow.fs.FileSystem, fs_path: str) -> None:
|
||||
"""Create directory at (fs, fs_path).
|
||||
|
||||
Some external filesystems require directories to already exist, or at least
|
||||
the `netloc` to be created (e.g. PyArrows ``mock://`` filesystem).
|
||||
|
||||
Generally this should be done before and outside of Ray applications. This
|
||||
utility is thus primarily used in testing, e.g. of ``mock://` URIs.
|
||||
"""
|
||||
try:
|
||||
fs.create_dir(fs_path)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
f"Caught exception when creating directory at ({fs}, {fs_path}):"
|
||||
)
|
||||
|
||||
|
||||
def get_fs_and_path(
|
||||
storage_path: Union[str, os.PathLike],
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
) -> Tuple[pyarrow.fs.FileSystem, str]:
|
||||
"""Returns the fs and path from a storage path and an optional custom fs.
|
||||
|
||||
Args:
|
||||
storage_path: A storage path or URI. (ex: s3://bucket/path or /tmp/ray_results)
|
||||
storage_filesystem: A custom filesystem to use. If not provided,
|
||||
this will be auto-resolved by pyarrow. If provided, the storage_path
|
||||
is assumed to be prefix-stripped already, and must be a valid path
|
||||
on the filesystem.
|
||||
|
||||
Returns:
|
||||
A ``(filesystem, path)`` tuple.
|
||||
"""
|
||||
storage_path = str(storage_path)
|
||||
|
||||
if storage_filesystem:
|
||||
return storage_filesystem, storage_path
|
||||
|
||||
return pyarrow.fs.FileSystem.from_uri(storage_path)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class StorageContext:
|
||||
"""Shared context that holds the source of truth for all paths and
|
||||
storage utilities, passed along from the driver to workers.
|
||||
|
||||
This object defines a few types of paths:
|
||||
1. *_fs_path: A path on the `storage_filesystem`. This is a regular path
|
||||
which has been prefix-stripped by pyarrow.fs.FileSystem.from_uri and
|
||||
can be joined with `Path(...).as_posix()`.
|
||||
2. *_driver_staging_path: The temporary staging directory on the local filesystem
|
||||
where driver artifacts are saved to before persisting them to storage.
|
||||
3. trial_working_directory: The local filesystem path that the remote
|
||||
actors' working directories are moved to by default.
|
||||
This is separated from the driver staging path so that driver syncing
|
||||
does not implicitly upload the trial working directory, for trials on the
|
||||
driver node.
|
||||
|
||||
Example with storage_path="mock:///bucket/path?param=1":
|
||||
|
||||
>>> import ray
|
||||
>>> from ray.train._internal.storage import StorageContext
|
||||
>>> import os
|
||||
>>> _ = ray.init()
|
||||
>>> storage = StorageContext(
|
||||
... storage_path="mock://netloc/bucket/path?param=1",
|
||||
... experiment_dir_name="exp_name",
|
||||
... )
|
||||
>>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS
|
||||
<pyarrow._fs._MockFileSystem object...
|
||||
>>> storage.experiment_fs_path
|
||||
'bucket/path/exp_name'
|
||||
>>> storage.experiment_driver_staging_path # doctest: +ELLIPSIS
|
||||
'/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts'
|
||||
>>> storage.trial_dir_name = "trial_dir"
|
||||
>>> storage.trial_fs_path
|
||||
'bucket/path/exp_name/trial_dir'
|
||||
>>> storage.trial_driver_staging_path # doctest: +ELLIPSIS
|
||||
'/tmp/ray/session_.../artifacts/.../exp_name/driver_artifacts/trial_dir'
|
||||
>>> storage.trial_working_directory # doctest: +ELLIPSIS
|
||||
'/tmp/ray/session_.../artifacts/.../exp_name/working_dirs/trial_dir'
|
||||
>>> ray.shutdown()
|
||||
|
||||
Example with storage_path="/tmp/ray_results":
|
||||
|
||||
>>> from ray.train._internal.storage import StorageContext
|
||||
>>> storage = StorageContext(
|
||||
... storage_path="/tmp/ray_results",
|
||||
... experiment_dir_name="exp_name",
|
||||
... )
|
||||
>>> storage.storage_fs_path
|
||||
'/tmp/ray_results'
|
||||
>>> storage.experiment_fs_path
|
||||
'/tmp/ray_results/exp_name'
|
||||
>>> storage.storage_filesystem # Auto-resolved # doctest: +ELLIPSIS
|
||||
<pyarrow._fs.LocalFileSystem object...
|
||||
|
||||
Internal Usage Examples:
|
||||
- To copy files to the trial directory on the storage filesystem:
|
||||
|
||||
pyarrow.fs.copy_files(
|
||||
local_dir,
|
||||
Path(storage.trial_fs_path, "subdir").as_posix(),
|
||||
destination_filesystem=storage.filesystem
|
||||
)
|
||||
|
||||
.. warning::
|
||||
This is an experimental developer API and is subject to change
|
||||
without notice between versions.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage_path: Union[str, os.PathLike],
|
||||
experiment_dir_name: str,
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
read_only: bool = False,
|
||||
):
|
||||
self.custom_fs_provided = storage_filesystem is not None
|
||||
|
||||
# Invariant: (`storage_filesystem`, `storage_path`) is the location where
|
||||
# *all* results can be accessed.
|
||||
self.experiment_dir_name = experiment_dir_name
|
||||
|
||||
self.storage_filesystem, self.storage_fs_path = get_fs_and_path(
|
||||
storage_path, storage_filesystem
|
||||
)
|
||||
self.storage_fs_path = Path(self.storage_fs_path).as_posix()
|
||||
|
||||
self.read_only = read_only
|
||||
if not self.read_only:
|
||||
self._create_validation_file()
|
||||
self._check_validation_file()
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
"StorageContext<\n"
|
||||
f" storage_filesystem='{self.storage_filesystem.type_name}',\n"
|
||||
f" storage_fs_path='{self.storage_fs_path}',\n"
|
||||
f" experiment_dir_name='{self.experiment_dir_name}',\n"
|
||||
">"
|
||||
)
|
||||
|
||||
def _create_validation_file(self):
|
||||
"""On the creation of a storage context, create a validation file at the
|
||||
storage path to verify that the storage path can be written to.
|
||||
This validation file is also used to check whether the storage path is
|
||||
accessible by all nodes in the cluster."""
|
||||
valid_file = Path(
|
||||
self.experiment_fs_path, VALIDATE_STORAGE_MARKER_FILENAME
|
||||
).as_posix()
|
||||
self.storage_filesystem.create_dir(self.experiment_fs_path)
|
||||
with self.storage_filesystem.open_output_stream(valid_file):
|
||||
pass
|
||||
|
||||
def _check_validation_file(self):
|
||||
"""Checks that the validation file exists at the storage path."""
|
||||
valid_file = Path(
|
||||
self.experiment_fs_path, VALIDATE_STORAGE_MARKER_FILENAME
|
||||
).as_posix()
|
||||
if not _exists_at_fs_path(fs=self.storage_filesystem, fs_path=valid_file):
|
||||
raise RuntimeError(
|
||||
f"Unable to set up cluster storage with the following settings:\n{self}"
|
||||
"\nCheck that all nodes in the cluster have read/write access "
|
||||
"to the configured storage path. `RunConfig(storage_path)` should be "
|
||||
"set to a cloud storage URI or a shared filesystem path accessible "
|
||||
"by all nodes in your cluster ('s3://bucket' or '/mnt/nfs'). "
|
||||
"A local path on the head node is not accessible by worker nodes. "
|
||||
"See: https://docs.ray.io/en/latest/train/user-guides/persistent-storage.html" # noqa: E501
|
||||
)
|
||||
|
||||
def persist_current_checkpoint(
|
||||
self, checkpoint: "Checkpoint", checkpoint_dir_name: str
|
||||
) -> "Checkpoint":
|
||||
"""Persists a given checkpoint to the current checkpoint path on the filesystem.
|
||||
|
||||
This method copies the checkpoint files to the storage location.
|
||||
It's up to the user to delete the original checkpoint files if desired.
|
||||
|
||||
For example, the original directory is typically a local temp directory.
|
||||
|
||||
Args:
|
||||
checkpoint: The checkpoint to persist to
|
||||
(fs, experiment_fs_path / checkpoint_dir_name).
|
||||
checkpoint_dir_name: Name of the destination directory for the
|
||||
checkpoint, relative to ``experiment_fs_path``.
|
||||
|
||||
Returns:
|
||||
Checkpoint: A Checkpoint pointing to the persisted checkpoint location.
|
||||
"""
|
||||
if self.read_only:
|
||||
raise RuntimeError(
|
||||
"Cannot perform write/validation operations as the StorageContext is read-only."
|
||||
)
|
||||
|
||||
# TODO(justinvyu): Fix this cyclical import.
|
||||
from ray.train import Checkpoint
|
||||
|
||||
checkpoint_fs_path = self.build_checkpoint_path_from_name(checkpoint_dir_name)
|
||||
|
||||
logger.debug(
|
||||
"Copying checkpoint files to storage path:\n"
|
||||
"({source_fs}, {source}) -> ({dest_fs}, {destination})".format(
|
||||
source=checkpoint.path,
|
||||
destination=checkpoint_fs_path,
|
||||
source_fs=checkpoint.filesystem,
|
||||
dest_fs=self.storage_filesystem,
|
||||
)
|
||||
)
|
||||
|
||||
# Raise an error if the storage path is not accessible when
|
||||
# attempting to upload a checkpoint from a remote worker.
|
||||
# Ex: If storage_path is a local path, then a validation marker
|
||||
# will only exist on the head node but not the worker nodes.
|
||||
self._check_validation_file()
|
||||
|
||||
self.storage_filesystem.create_dir(checkpoint_fs_path)
|
||||
_pyarrow_fs_copy_files(
|
||||
source=checkpoint.path,
|
||||
destination=checkpoint_fs_path,
|
||||
source_filesystem=checkpoint.filesystem,
|
||||
destination_filesystem=self.storage_filesystem,
|
||||
)
|
||||
|
||||
persisted_checkpoint = Checkpoint(
|
||||
filesystem=self.storage_filesystem,
|
||||
path=checkpoint_fs_path,
|
||||
)
|
||||
logger.info(f"Checkpoint successfully created at: {persisted_checkpoint}")
|
||||
return persisted_checkpoint
|
||||
|
||||
@property
|
||||
def experiment_fs_path(self) -> str:
|
||||
"""The path on the `storage_filesystem` to the experiment directory.
|
||||
|
||||
NOTE: This does not have a URI prefix anymore, since it has been stripped
|
||||
by pyarrow.fs.FileSystem.from_uri already. The URI scheme information is
|
||||
kept in `storage_filesystem` instead.
|
||||
"""
|
||||
return Path(self.storage_fs_path, self.experiment_dir_name).as_posix()
|
||||
|
||||
@property
|
||||
def local_working_directory(self) -> str:
|
||||
"""Every ray train worker will set this directory as its working directory."""
|
||||
if self.experiment_dir_name is None:
|
||||
raise RuntimeError(
|
||||
"Cannot access `local_working_directory` without "
|
||||
"setting `experiment_dir_name`"
|
||||
)
|
||||
return Path(_get_ray_train_session_dir(), self.experiment_dir_name).as_posix()
|
||||
|
||||
@property
|
||||
def checkpoint_manager_snapshot_path(self) -> str:
|
||||
"""The path to the checkpoint manager snapshot file."""
|
||||
return Path(
|
||||
self.experiment_fs_path, CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
|
||||
).as_posix()
|
||||
|
||||
@staticmethod
|
||||
def get_experiment_dir_name(run_obj: Union[str, Callable, Type]) -> str:
|
||||
from ray.tune.experiment import Experiment
|
||||
|
||||
run_identifier = Experiment.get_trainable_name(run_obj)
|
||||
|
||||
if bool(int(os.environ.get("TUNE_DISABLE_DATED_SUBDIR", 0))):
|
||||
dir_name = run_identifier
|
||||
else:
|
||||
dir_name = "{}_{}".format(run_identifier, date_str())
|
||||
return dir_name
|
||||
|
||||
@staticmethod
|
||||
def make_default_checkpoint_dir_name():
|
||||
"""Get the name of the checkpoint directory by timestamp."""
|
||||
return f"checkpoint_{date_str(include_ms=True)}"
|
||||
|
||||
def extract_checkpoint_dir_name_from_path(self, checkpoint_path: str) -> str:
|
||||
"""Get the checkpoint name from the checkpoint path.
|
||||
The parent directory of the checkpoint path should be the experiment directory.
|
||||
"""
|
||||
# TODO: Use Pathlib to extract the name when supports at least Python 3.9
|
||||
experiment_fs_path = self.experiment_fs_path + "/"
|
||||
if not checkpoint_path.startswith(experiment_fs_path):
|
||||
raise ValueError(
|
||||
f"Checkpoint path {checkpoint_path} is not under the experiment "
|
||||
f"directory {self.experiment_fs_path}."
|
||||
)
|
||||
return checkpoint_path[len(experiment_fs_path) :]
|
||||
|
||||
def build_checkpoint_path_from_name(self, checkpoint_name: str) -> str:
|
||||
"""Get the checkpoint path from the checkpoint name.
|
||||
The parent directory of the checkpoint path should be the experiment directory.
|
||||
"""
|
||||
return Path(self.experiment_fs_path, checkpoint_name).as_posix()
|
||||
@@ -0,0 +1,297 @@
|
||||
import logging
|
||||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
|
||||
from ray.train.v2._internal.execution import collective_impl
|
||||
from ray.train.v2._internal.execution.context import (
|
||||
get_train_context as get_internal_train_context,
|
||||
)
|
||||
from ray.train.v2.api.context import (
|
||||
DistributedTrainContext,
|
||||
LocalTrainContext,
|
||||
TrainContext as ExternalTrainContext,
|
||||
)
|
||||
from ray.train.v2.api.report_config import (
|
||||
CheckpointConsistencyMode,
|
||||
CheckpointUploadMode,
|
||||
)
|
||||
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
||||
|
||||
|
||||
class TrainFnUtils(ABC):
|
||||
"""Utility class providing an abstraction layer between user-facing APIs
|
||||
and :class:`~ray.train.v2.api.context.TrainContext`.
|
||||
|
||||
It should be set before the users' training function is called.
|
||||
This class can be patched if new user APIs behaviors is wanted.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
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.
|
||||
|
||||
Args:
|
||||
metrics: The metrics to report.
|
||||
checkpoint: The checkpoint to report.
|
||||
checkpoint_dir_name: The name of the checkpoint dir
|
||||
in this iteration. Note: If not set, the checkpoint will
|
||||
be stored in the default storage path. If set, make sure
|
||||
this value is unique for each iteration.
|
||||
checkpoint_upload_mode: The manner in which we want to upload the checkpoint.
|
||||
Defaults to uploading the checkpoint synchronously.
|
||||
This works when no checkpoint is provided but is not useful in that case.
|
||||
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: [Alpha] If True, triggers validation with default kwargs from validation_config.
|
||||
If a ValidationTaskConfig, validation is run using fn_kwargs merged with validation_config
|
||||
defaults, with fn_kwargs taking precedence on conflicts. If False, no validation.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_checkpoint(self) -> Optional["Checkpoint"]:
|
||||
"""Get the latest checkpoint to resume training from.
|
||||
|
||||
Returns:
|
||||
The latest checkpoint if available, None otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_all_reported_checkpoints(
|
||||
self,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
"""Get all the checkpoints reported by the workers.
|
||||
|
||||
Args:
|
||||
consistency_mode: Read semantics for checkpoint retrieval. Defaults to VALIDATED.
|
||||
timeout_s: Timeout in seconds for reading checkpoints and validation data.
|
||||
Defaults to ``None`` to not time out.
|
||||
|
||||
Returns:
|
||||
A list of ReportedCheckpoint objects that represent the checkpoints and
|
||||
corresponding metrics reported by the workers.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
"""Get the dataset shard for this training process.
|
||||
|
||||
Args:
|
||||
dataset_info: The metadata of the dataset to get the shard for.
|
||||
|
||||
Returns:
|
||||
The DataIterator shard for this worker.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_context(self) -> ExternalTrainContext:
|
||||
"""Get the TrainContext for this training process.
|
||||
The specific type of TrainContext returned depends on the implementation of TrainFnUtils.
|
||||
|
||||
Returns:
|
||||
The train context for this training process.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def is_distributed(self) -> bool:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def barrier(self) -> None:
|
||||
"""Create a barrier across all workers.
|
||||
|
||||
All workers must call this method before the training function can continue.
|
||||
|
||||
This method is used by the public API function :func:`ray.train.collective.barrier`.
|
||||
Users should typically call ``ray.train.collective.barrier()`` instead of calling this method directly.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def broadcast_from_rank_zero(self, data: Any) -> Any:
|
||||
"""Broadcast data from the rank 0 worker to all other workers.
|
||||
|
||||
This method is used by the public API function :func:`ray.train.collective.broadcast_from_rank_zero`.
|
||||
Users should typically call ``ray.train.collective.broadcast_from_rank_zero()`` instead of calling this method directly.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class DistributedTrainFnUtils(TrainFnUtils):
|
||||
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:
|
||||
return get_internal_train_context().report(
|
||||
metrics,
|
||||
checkpoint,
|
||||
checkpoint_dir_name,
|
||||
checkpoint_upload_mode,
|
||||
delete_local_checkpoint_after_upload,
|
||||
checkpoint_upload_fn,
|
||||
validation,
|
||||
)
|
||||
|
||||
def get_checkpoint(self):
|
||||
return get_internal_train_context().get_checkpoint()
|
||||
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
return get_internal_train_context().get_dataset_shard(dataset_info)
|
||||
|
||||
def get_context(self) -> DistributedTrainContext:
|
||||
return DistributedTrainContext()
|
||||
|
||||
def is_distributed(self) -> bool:
|
||||
return True
|
||||
|
||||
def barrier(self) -> None:
|
||||
return collective_impl.barrier()
|
||||
|
||||
def broadcast_from_rank_zero(self, data: Any) -> Any:
|
||||
return collective_impl.broadcast_from_rank_zero(data)
|
||||
|
||||
def get_all_reported_checkpoints(
|
||||
self,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
return get_internal_train_context().get_all_reported_checkpoints(
|
||||
consistency_mode=consistency_mode, timeout_s=timeout_s
|
||||
)
|
||||
|
||||
|
||||
class LocalTrainFnUtils(TrainFnUtils):
|
||||
def __init__(
|
||||
self,
|
||||
experiment_name: str,
|
||||
dataset_shards: Optional[Dict[str, "DataIterator"]] = None,
|
||||
world_size: int = 1,
|
||||
world_rank: int = 0,
|
||||
local_rank: int = 0,
|
||||
local_world_size: int = 1,
|
||||
node_rank: int = 0,
|
||||
):
|
||||
self._context = LocalTrainContext(
|
||||
experiment_name=experiment_name,
|
||||
world_size=world_size,
|
||||
world_rank=world_rank,
|
||||
local_rank=local_rank,
|
||||
local_world_size=local_world_size,
|
||||
node_rank=node_rank,
|
||||
)
|
||||
self._dataset_shards = dataset_shards
|
||||
self._last_metrics = None
|
||||
self._last_checkpoint = None
|
||||
|
||||
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:
|
||||
self._last_metrics = metrics
|
||||
self._last_checkpoint = checkpoint
|
||||
logger.info(f"Reported metrics: {metrics}")
|
||||
|
||||
def get_checkpoint(self) -> Optional["Checkpoint"]:
|
||||
return self._last_checkpoint
|
||||
|
||||
def get_dataset_shard(self, dataset_info: DatasetShardMetadata) -> "DataIterator":
|
||||
dataset_name = dataset_info.dataset_name
|
||||
assert (
|
||||
self._dataset_shards is not None and dataset_name in self._dataset_shards
|
||||
), f"Dataset shard {dataset_name} not found."
|
||||
return self._dataset_shards[dataset_name]
|
||||
|
||||
def get_context(self) -> LocalTrainContext:
|
||||
return self._context
|
||||
|
||||
def is_distributed(self) -> bool:
|
||||
return False
|
||||
|
||||
def barrier(self) -> None:
|
||||
pass
|
||||
|
||||
def broadcast_from_rank_zero(self, data: Any) -> Any:
|
||||
return data
|
||||
|
||||
def _get_last_metrics(self) -> Optional[Dict[str, Any]]:
|
||||
"""Return the last metrics reported by the training function.
|
||||
This function should only be called by LocalController
|
||||
"""
|
||||
return self._last_metrics
|
||||
|
||||
def get_all_reported_checkpoints(
|
||||
self,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
return []
|
||||
|
||||
|
||||
_train_fn_utils: Optional[TrainFnUtils] = None
|
||||
_train_fn_utils_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_train_fn_utils() -> TrainFnUtils:
|
||||
"""Return the Ray Train function utilities.
|
||||
|
||||
Returns:
|
||||
The TrainFnUtils instance for the current worker.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the Ray Train function utilities are not initialized.
|
||||
"""
|
||||
global _train_fn_utils
|
||||
with _train_fn_utils_lock:
|
||||
if _train_fn_utils is None:
|
||||
raise RuntimeError("Ray Train function utilities not initialized.")
|
||||
return _train_fn_utils
|
||||
|
||||
|
||||
def set_train_fn_utils(train_fn_utils) -> None:
|
||||
global _train_fn_utils
|
||||
with _train_fn_utils_lock:
|
||||
_train_fn_utils = train_fn_utils
|
||||
@@ -0,0 +1,22 @@
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
||||
|
||||
|
||||
class _TrainingReport:
|
||||
"""Checkpoint and metrics reported by user, as well as optional validation configuration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint: Optional["Checkpoint"],
|
||||
metrics: Dict[str, Any],
|
||||
validation: Union[bool, "ValidationTaskConfig"],
|
||||
):
|
||||
self.checkpoint = checkpoint
|
||||
self.metrics = metrics
|
||||
self.validation = validation
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"TrainingReport(checkpoint={self.checkpoint}, metrics={self.metrics}, validation={self.validation})"
|
||||
@@ -0,0 +1,31 @@
|
||||
from .execution_group import ExecutionGroup, ReplicaGroup
|
||||
from .placement_group_handle import (
|
||||
DefaultPlacementGroupHandle,
|
||||
PlacementGroupHandle,
|
||||
SlicePlacementGroupHandle,
|
||||
)
|
||||
from .poll import WorkerGroupPollStatus, WorkerStatus
|
||||
from .state import (
|
||||
WorkerGroupContext,
|
||||
WorkerGroupState,
|
||||
WorkerGroupStateBuilder,
|
||||
)
|
||||
from .worker import ActorMetadata, RayTrainWorker, Worker
|
||||
from .worker_group import WorkerGroup
|
||||
|
||||
__all__ = [
|
||||
"ActorMetadata",
|
||||
"DefaultPlacementGroupHandle",
|
||||
"ExecutionGroup",
|
||||
"PlacementGroupHandle",
|
||||
"RayTrainWorker",
|
||||
"ReplicaGroup",
|
||||
"SlicePlacementGroupHandle",
|
||||
"Worker",
|
||||
"WorkerGroup",
|
||||
"WorkerGroupContext",
|
||||
"WorkerGroupPollStatus",
|
||||
"WorkerGroupState",
|
||||
"WorkerGroupStateBuilder",
|
||||
"WorkerStatus",
|
||||
]
|
||||
@@ -0,0 +1,117 @@
|
||||
import abc
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, TypeVar
|
||||
|
||||
import ray
|
||||
from ray.train._internal.base_worker_group import BaseWorkerGroup
|
||||
from ray.train.v2._internal.execution.worker_group.state import _shutdown_workers
|
||||
from ray.train.v2._internal.execution.worker_group.worker import Worker
|
||||
from ray.types import ObjectRef
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.callback import ReplicaGroupCallback
|
||||
from ray.train.v2._internal.execution.worker_group.state import WorkerGroupContext
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class ExecutionGroup(BaseWorkerGroup):
|
||||
"""Base class for groups that can execute functions on workers.
|
||||
|
||||
Provides concrete implementations of the 4 execution methods and __len__
|
||||
based on two abstract primitives: _assert_active() and get_workers().
|
||||
"""
|
||||
|
||||
@abc.abstractmethod
|
||||
def _assert_active(self):
|
||||
"""Assert that this execution group is active."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_workers(self) -> List[Worker]:
|
||||
"""Return the list of workers in this group."""
|
||||
pass
|
||||
|
||||
def execute_async(self, fn: Callable, *fn_args, **fn_kwargs) -> List[ObjectRef]:
|
||||
self._assert_active()
|
||||
workers = self.get_workers()
|
||||
|
||||
return [worker.execute_async(fn, *fn_args, **fn_kwargs) for worker in workers]
|
||||
|
||||
def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> List[T]:
|
||||
return ray.get(self.execute_async(fn, *fn_args, **fn_kwargs))
|
||||
|
||||
def execute_single_async(
|
||||
self, rank: int, fn: Callable[..., T], *fn_args, **fn_kwargs
|
||||
) -> ObjectRef:
|
||||
self._assert_active()
|
||||
workers = self.get_workers()
|
||||
|
||||
if rank >= len(workers):
|
||||
raise ValueError(
|
||||
f"The provided {rank=} is " f"not valid for {len(workers)} workers."
|
||||
)
|
||||
|
||||
return workers[rank].execute_async(fn, *fn_args, **fn_kwargs)
|
||||
|
||||
def execute_single(
|
||||
self, rank: int, fn: Callable[..., T], *fn_args, **fn_kwargs
|
||||
) -> T:
|
||||
return ray.get(self.execute_single_async(rank, fn, *fn_args, **fn_kwargs))
|
||||
|
||||
def __len__(self) -> int:
|
||||
self._assert_active()
|
||||
return len(self.get_workers())
|
||||
|
||||
|
||||
class ReplicaGroup(ExecutionGroup):
|
||||
"""A group representing a subset of workers from a WorkerGroup.
|
||||
|
||||
Used to pass a replica's workers to backend methods (on_start, etc.)
|
||||
as if they were a standalone worker group.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workers: List[Worker],
|
||||
resources_per_worker: dict,
|
||||
callbacks: Optional[List["ReplicaGroupCallback"]] = None,
|
||||
):
|
||||
self._workers = workers
|
||||
self._resources_per_worker = resources_per_worker
|
||||
self._callbacks = callbacks or []
|
||||
# An inactive ReplicaGroup still needs to keep track of workers
|
||||
# so we can replace them later.
|
||||
self._active = True
|
||||
|
||||
def _assert_active(self):
|
||||
if not self.is_active():
|
||||
raise ValueError("ReplicaGroup has been shut down.")
|
||||
|
||||
def is_active(self) -> bool:
|
||||
return self._active
|
||||
|
||||
def get_workers(self) -> List[Worker]:
|
||||
return self._workers
|
||||
|
||||
def get_resources_per_worker(self) -> dict:
|
||||
return self._resources_per_worker
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown all workers in this replica group and clear state."""
|
||||
if self.is_active():
|
||||
for cb in self._callbacks:
|
||||
cb.before_replica_group_shutdown(self)
|
||||
|
||||
_shutdown_workers(self._workers)
|
||||
self._active = False
|
||||
|
||||
def start_training(self, worker_group_context: "WorkerGroupContext"):
|
||||
"""Start training on all workers in this replica group."""
|
||||
for cb in self._callbacks:
|
||||
cb.after_replica_group_start(self)
|
||||
ray.get(
|
||||
[
|
||||
worker.actor.run_train_fn.remote(worker_group_context.train_fn_ref)
|
||||
for worker in self._workers
|
||||
]
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
from ray.types import ObjectRef
|
||||
from ray.util.placement_group import PlacementGroup, remove_placement_group
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.util.tpu import SlicePlacementGroup
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PlacementGroupHandle(ABC):
|
||||
"""Unified interface for placement groups in Ray Train.
|
||||
|
||||
This abstract base class provides a common interface for both standard
|
||||
PlacementGroup and SlicePlacementGroup, allowing WorkerGroup to handle
|
||||
them uniformly without conditional logic.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def placement_group(self) -> PlacementGroup:
|
||||
"""The underlying PlacementGroup for worker scheduling."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def ready(self) -> ObjectRef:
|
||||
"""Returns an ObjectRef to check if the placement group is ready.
|
||||
|
||||
Compatible with ray.get() and ray.wait().
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
|
||||
"""Wait for the placement group to be ready within the specified time.
|
||||
Args:
|
||||
timeout_seconds: Timeout in seconds.
|
||||
Returns:
|
||||
True if the placement group is created. False otherwise.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def shutdown(self) -> None:
|
||||
"""Release all resources associated with this placement group.
|
||||
|
||||
After calling this method, the placement group should no longer be used.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DefaultPlacementGroupHandle(PlacementGroupHandle):
|
||||
"""Wrapper for standard PlacementGroup."""
|
||||
|
||||
def __init__(self, pg: PlacementGroup):
|
||||
self._pg = pg
|
||||
|
||||
@property
|
||||
def placement_group(self) -> PlacementGroup:
|
||||
return self._pg
|
||||
|
||||
def ready(self) -> ObjectRef:
|
||||
return self._pg.ready()
|
||||
|
||||
def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
|
||||
try:
|
||||
return self._pg.wait(timeout_seconds)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Placement group wait failed; treating as not ready.",
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
def shutdown(self) -> None:
|
||||
remove_placement_group(self._pg)
|
||||
|
||||
|
||||
class SlicePlacementGroupHandle(PlacementGroupHandle):
|
||||
"""Wrapper for SlicePlacementGroup that delegates to its underlying PlacementGroup."""
|
||||
|
||||
def __init__(self, spg: "SlicePlacementGroup"):
|
||||
self._spg = spg
|
||||
|
||||
@property
|
||||
def placement_group(self) -> PlacementGroup:
|
||||
return self._spg.placement_group
|
||||
|
||||
def ready(self) -> ObjectRef:
|
||||
return self._spg.placement_group.ready()
|
||||
|
||||
def wait(self, timeout_seconds: Union[float, int] = 30) -> bool:
|
||||
try:
|
||||
return self._spg.placement_group.wait(timeout_seconds)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Slice placement group wait failed; treating as not ready.",
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._spg.shutdown()
|
||||
@@ -0,0 +1,152 @@
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional, Set
|
||||
|
||||
from ray._private.ray_logging import NUMBERS
|
||||
from ray.train.v2._internal.exceptions import (
|
||||
UserExceptionWithTraceback,
|
||||
WorkerHealthCheckFailedError,
|
||||
)
|
||||
from ray.train.v2._internal.execution.preemption import PreemptionInfo
|
||||
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
||||
from ray.train.v2.api.exceptions import WorkerGroupError
|
||||
from ray.types import ObjectRef
|
||||
|
||||
ERR_CHAR_LIMIT = 1000
|
||||
|
||||
|
||||
def _normalize_error_string(error_str: str) -> str:
|
||||
# Replace numbers with <NUM> based on NUMBERS regex
|
||||
normalized = re.sub(NUMBERS, "<NUM>", error_str)
|
||||
return normalized
|
||||
|
||||
|
||||
def _truncate_error_string(error_str: str) -> str:
|
||||
"""
|
||||
Truncates error strings to include the first ERR_CHAR_LIMIT // 2
|
||||
characters and the last ERR_CHAR_LIMIT // 2 characters.
|
||||
"""
|
||||
if len(error_str) >= ERR_CHAR_LIMIT:
|
||||
return (
|
||||
error_str[: ERR_CHAR_LIMIT // 2]
|
||||
+ "...\n... (Output truncated. See individual worker logs for full details) ...\n"
|
||||
+ error_str[len(error_str) - ERR_CHAR_LIMIT // 2 :]
|
||||
)
|
||||
return error_str
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerStatus:
|
||||
running: bool
|
||||
error: Optional[Exception] = None
|
||||
training_report: Optional[_TrainingReport] = None
|
||||
return_value: Any = field(default=None)
|
||||
preemption_info: Optional[PreemptionInfo] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerGroupPollStatus:
|
||||
worker_statuses: Dict[int, WorkerStatus]
|
||||
worker_rank_to_replica_group_rank: Optional[Dict[int, int]] = None
|
||||
|
||||
@property
|
||||
def all_replica_group_indices(self) -> Set[int]:
|
||||
"""Return the set of all replica group indices."""
|
||||
if self.worker_rank_to_replica_group_rank is None:
|
||||
return set()
|
||||
return set(self.worker_rank_to_replica_group_rank.values())
|
||||
|
||||
@property
|
||||
def failing_replica_group_indices(self) -> Set[int]:
|
||||
"""Return the set of replica group indices that have failing workers."""
|
||||
if self.worker_rank_to_replica_group_rank is None:
|
||||
return set()
|
||||
return {
|
||||
self.worker_rank_to_replica_group_rank[rank]
|
||||
for rank in self.errors
|
||||
if rank in self.worker_rank_to_replica_group_rank
|
||||
}
|
||||
|
||||
@property
|
||||
def errors(self) -> Dict[int, Exception]:
|
||||
errors = {}
|
||||
for world_rank, status in self.worker_statuses.items():
|
||||
if status.error is not None:
|
||||
error = status.error
|
||||
if isinstance(error, UserExceptionWithTraceback):
|
||||
error = error._base_exc
|
||||
errors[world_rank] = error
|
||||
return errors
|
||||
|
||||
def get_worker_group_error(self) -> WorkerGroupError:
|
||||
return WorkerGroupError(
|
||||
error_message=self.get_error_string(),
|
||||
worker_failures=self.errors,
|
||||
)
|
||||
|
||||
@property
|
||||
def finished(self) -> bool:
|
||||
return self.worker_statuses and all(
|
||||
not status.running for status in self.worker_statuses.values()
|
||||
)
|
||||
|
||||
def get_error_string(self) -> str:
|
||||
"""
|
||||
Returns a string representation of worker group errors.
|
||||
Groups similar errors (ignoring numbers) and shows original error examples.
|
||||
"""
|
||||
# Group errors by normalized strings (ignoring numbers)
|
||||
normalized_error_to_ranks = defaultdict(list)
|
||||
normalized_error_to_original = {}
|
||||
show_full_error = set()
|
||||
|
||||
for world_rank, status in self.worker_statuses.items():
|
||||
if status.error:
|
||||
error_str = str(status.error)
|
||||
normalized_error = _normalize_error_string(error_str)
|
||||
|
||||
normalized_error_to_ranks[normalized_error].append(str(world_rank))
|
||||
|
||||
# Store the first original error for this normalized group
|
||||
if normalized_error not in normalized_error_to_original:
|
||||
normalized_error_to_original[normalized_error] = error_str
|
||||
|
||||
# Fully show errors for non-graceful worker failures or running workers
|
||||
if (
|
||||
isinstance(status.error, WorkerHealthCheckFailedError)
|
||||
or status.running
|
||||
):
|
||||
show_full_error.add(normalized_error)
|
||||
|
||||
errors = []
|
||||
for normalized_error, ranks in normalized_error_to_ranks.items():
|
||||
# Show the original error
|
||||
orig_error = normalized_error_to_original[normalized_error]
|
||||
|
||||
# Convert rank list to comma-separated strings
|
||||
ranks_str = ",".join(ranks)
|
||||
|
||||
if normalized_error in show_full_error:
|
||||
errors.append(f"[Rank {ranks_str} Error Snippet]:\n{orig_error}")
|
||||
else:
|
||||
errors.append(
|
||||
f"[Rank {ranks_str} Error Snippet]:\n{_truncate_error_string(orig_error)}"
|
||||
)
|
||||
|
||||
error_str = "\n".join(errors)
|
||||
|
||||
return error_str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PollTask:
|
||||
"""Represents a poll task for a worker.
|
||||
|
||||
Attributes:
|
||||
start_time: The time when the poll task was started.
|
||||
task: The ObjectRef representing the poll task.
|
||||
"""
|
||||
|
||||
start_time: float
|
||||
task: ObjectRef
|
||||
@@ -0,0 +1,170 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Callable, Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.execution.checkpoint.sync_actor import SynchronizationActor
|
||||
from ray.train.v2._internal.execution.worker_group.worker import Worker
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.execution.worker_group.placement_group_handle import (
|
||||
PlacementGroupHandle,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerGroupContext:
|
||||
"""Context for a worker group.
|
||||
|
||||
This stores the context that is shared when starting a worker group.
|
||||
|
||||
Attributes:
|
||||
run_attempt_id: The ID of the run attempt.
|
||||
train_fn_ref: An object store reference to the training function to execute.
|
||||
num_workers: The number of workers in the worker group.
|
||||
resources_per_worker: The resources per worker.
|
||||
placement_strategy: Strategy for placing workers.
|
||||
label_selector: Optional label selectors to apply per-bundle for workers.
|
||||
num_slices: The number of TPU slices (if using TPU). Defaults to 1.
|
||||
"""
|
||||
|
||||
run_attempt_id: str
|
||||
train_fn_ref: ObjectRefWrapper[Callable[[], None]]
|
||||
num_workers: int
|
||||
resources_per_worker: Dict[str, float]
|
||||
placement_strategy: str = "PACK"
|
||||
label_selector: Optional[List[Dict[str, str]]] = None
|
||||
num_slices: int = 1
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerGroupState:
|
||||
"""Ongoing state of an active worker group.
|
||||
|
||||
Attributes:
|
||||
start_time: The time when the worker group was started.
|
||||
workers: The workers in the worker group.
|
||||
These should always be in sorted order by world rank.
|
||||
placement_group_handle: The placement group handle for the worker group.
|
||||
sync_actor: The synchronization actor for the worker group.
|
||||
"""
|
||||
|
||||
start_time: float
|
||||
placement_group_handle: "PlacementGroupHandle"
|
||||
workers: List[Worker]
|
||||
sync_actor: ActorHandle
|
||||
|
||||
@property
|
||||
def num_workers(self) -> int:
|
||||
return len(self.workers)
|
||||
|
||||
def replace_workers(
|
||||
self, old_workers: List[Worker], new_workers: List[Worker]
|
||||
) -> "WorkerGroupState":
|
||||
"""Return a new WorkerGroupState with old_workers replaced by new_workers."""
|
||||
current_workers = list(self.workers)
|
||||
for old_w, new_w in zip(old_workers, new_workers):
|
||||
idx = current_workers.index(old_w)
|
||||
current_workers[idx] = new_w
|
||||
return WorkerGroupState(
|
||||
start_time=self.start_time,
|
||||
placement_group_handle=self.placement_group_handle,
|
||||
workers=current_workers,
|
||||
sync_actor=self.sync_actor,
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
_shutdown_workers(self.workers)
|
||||
_shutdown_sync_actor(self.sync_actor)
|
||||
self.placement_group_handle.shutdown()
|
||||
|
||||
|
||||
class WorkerGroupStateBuilder:
|
||||
"""Builder for WorkerGroupState.
|
||||
|
||||
Example usage:
|
||||
```python
|
||||
builder = WorkerGroupStateBuilder()
|
||||
builder.with_placement_group_handle(placement_group_handle)
|
||||
builder.with_workers(workers)
|
||||
builder.with_sync_actor(sync_actor)
|
||||
state = builder.build()
|
||||
|
||||
builder.shutdown(patience_s=10)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.placement_group_handle = None
|
||||
self.workers = None
|
||||
self.sync_actor = None
|
||||
|
||||
def with_placement_group_handle(
|
||||
self, placement_group_handle: "PlacementGroupHandle"
|
||||
) -> "WorkerGroupStateBuilder":
|
||||
self.placement_group_handle = placement_group_handle
|
||||
return self
|
||||
|
||||
def with_workers(self, workers: List[Worker]) -> "WorkerGroupStateBuilder":
|
||||
self.workers = workers
|
||||
return self
|
||||
|
||||
def with_sync_actor(
|
||||
self, sync_actor: SynchronizationActor
|
||||
) -> "WorkerGroupStateBuilder":
|
||||
self.sync_actor = sync_actor
|
||||
return self
|
||||
|
||||
def build(self) -> WorkerGroupState:
|
||||
required_attrs = {
|
||||
"placement_group_handle": self.placement_group_handle,
|
||||
"workers": self.workers,
|
||||
"sync_actor": self.sync_actor,
|
||||
}
|
||||
missing = [name for name, attr in required_attrs.items() if attr is None]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Cannot build incomplete state. Missing: {', '.join(missing)}"
|
||||
)
|
||||
return WorkerGroupState(
|
||||
start_time=time_monotonic(),
|
||||
placement_group_handle=self.placement_group_handle,
|
||||
workers=self.workers,
|
||||
sync_actor=self.sync_actor,
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
if self.workers:
|
||||
_shutdown_workers(self.workers)
|
||||
self.workers = None
|
||||
|
||||
if self.sync_actor:
|
||||
_shutdown_sync_actor(self.sync_actor)
|
||||
self.sync_actor = None
|
||||
|
||||
if self.placement_group_handle:
|
||||
self.placement_group_handle.shutdown()
|
||||
self.placement_group_handle = None
|
||||
|
||||
|
||||
def _shutdown_workers(workers: List[Worker], patience_s: float = 5):
|
||||
"""Shuts down workers after allowing a maximum of patience_s seconds for shutdown hooks to run."""
|
||||
if patience_s < 0:
|
||||
raise ValueError("Invalid patience_s: must be non-negative")
|
||||
|
||||
done_refs = [w.actor.shutdown.remote() for w in workers]
|
||||
|
||||
logger.debug(f"Shutting down {len(workers)} workers.")
|
||||
|
||||
ray.wait(done_refs, num_returns=len(done_refs), timeout=patience_s)
|
||||
|
||||
for worker in workers:
|
||||
ray.kill(worker.actor)
|
||||
|
||||
|
||||
def _shutdown_sync_actor(sync_actor: SynchronizationActor):
|
||||
ray.kill(sync_actor)
|
||||
@@ -0,0 +1,93 @@
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
from typing import Callable, Optional, TypeVar
|
||||
|
||||
from ray.train.v2._internal.exceptions import UserExceptionWithTraceback
|
||||
from ray.train.v2._internal.util import (
|
||||
construct_user_exception_with_traceback,
|
||||
get_callable_name,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ThreadRunner:
|
||||
"""Utility to run a user function as a thread and capture its return value
|
||||
or exception.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._ret: Optional[T] = None
|
||||
self._exc: Optional[UserExceptionWithTraceback] = None
|
||||
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._monitor_thread: Optional[threading.Thread] = None
|
||||
self._lock = threading.Lock()
|
||||
self._exc_queue: queue.SimpleQueue[Optional[Exception]] = queue.SimpleQueue()
|
||||
|
||||
def run(self, target: Callable[[], T]) -> None:
|
||||
if self._thread is not None:
|
||||
raise RuntimeError("Thread is already running.")
|
||||
|
||||
def _run_target():
|
||||
try:
|
||||
result = target()
|
||||
with self._lock:
|
||||
self._ret = result
|
||||
self._exc_queue.put(None)
|
||||
except BaseException as e:
|
||||
# Exclude the first 3 frames from the traceback, which are
|
||||
# the `ThreadRunner._run_target`, `construct_train_func`, and
|
||||
# train_fn_with_final_checkpoint_flush calls.
|
||||
self._exc_queue.put(
|
||||
construct_user_exception_with_traceback(e, exclude_frames=3)
|
||||
)
|
||||
|
||||
# Join the monitor thread. This ensures that a queued exception
|
||||
# is processed before the target function is considered done.
|
||||
self._monitor_thread.join()
|
||||
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._monitor_target,
|
||||
daemon=True,
|
||||
name=f"MonitoringThread({get_callable_name(target)})",
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=_run_target,
|
||||
daemon=True,
|
||||
name=f"TrainingThread({get_callable_name(target)})",
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
def _monitor_target(self):
|
||||
"""Monitor the exception queue and set the exception if an exception is found.
|
||||
|
||||
This should run as a daemon thread and exit when None is put into the exception queue.
|
||||
"""
|
||||
exc: Optional[UserExceptionWithTraceback] = self._exc_queue.get()
|
||||
if exc is None:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
self._exc = exc
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Returns whether the target function is still running."""
|
||||
return self._thread is not None and self._thread.is_alive()
|
||||
|
||||
def get_error(self) -> Optional[BaseException]:
|
||||
with self._lock:
|
||||
return self._exc
|
||||
|
||||
def get_return_value(self) -> Optional[T]:
|
||||
with self._lock:
|
||||
return self._ret
|
||||
|
||||
def get_exception_queue(self) -> queue.SimpleQueue:
|
||||
"""Returns a queue that nested threads can add exceptions to."""
|
||||
return self._exc_queue
|
||||
@@ -0,0 +1,319 @@
|
||||
import logging
|
||||
import os
|
||||
import queue
|
||||
import socket
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypeVar, Union
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from .thread_runner import ThreadRunner
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_ENABLE_WORKER_LOGGING,
|
||||
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
TrainContextCallback,
|
||||
WorkerCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import (
|
||||
DistributedContext,
|
||||
ExecutionContext,
|
||||
TrainContext,
|
||||
TrainRunContext,
|
||||
get_train_context,
|
||||
set_train_context,
|
||||
)
|
||||
from ray.train.v2._internal.execution.preemption import (
|
||||
PreemptionContext,
|
||||
PreemptionInfo,
|
||||
)
|
||||
from ray.train.v2._internal.execution.storage import StorageContext
|
||||
from ray.train.v2._internal.execution.train_fn_utils import (
|
||||
DistributedTrainFnUtils,
|
||||
set_train_fn_utils,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group.poll import WorkerStatus
|
||||
from ray.train.v2._internal.logging.logging import LoggingManager
|
||||
from ray.train.v2._internal.logging.patch_print import patch_print_function
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper
|
||||
from ray.types import ObjectRef
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2._internal.data_integration.interfaces import DatasetShardProvider
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActorMetadata:
|
||||
hostname: str
|
||||
node_id: str
|
||||
node_ip: str
|
||||
pid: int
|
||||
accelerator_ids: Dict[str, List[Union[int, str]]]
|
||||
|
||||
@property
|
||||
def gpu_ids(self) -> List[Union[int, str]]:
|
||||
return self.accelerator_ids.get("GPU", [])
|
||||
|
||||
@cached_property
|
||||
def _repr(self) -> str:
|
||||
indent = " "
|
||||
repr_lines = [
|
||||
"ActorMetadata(",
|
||||
f"{indent}hostname={repr(self.hostname)},",
|
||||
f"{indent}node_id={repr(self.node_id)},",
|
||||
f"{indent}node_ip={repr(self.node_ip)},",
|
||||
f"{indent}pid={repr(self.pid)},",
|
||||
]
|
||||
non_empty_accelerator_ids = {k: v for k, v in self.accelerator_ids.items() if v}
|
||||
if non_empty_accelerator_ids:
|
||||
repr_lines.append(f"{indent}accelerator_ids={non_empty_accelerator_ids},")
|
||||
|
||||
repr_lines.append(")")
|
||||
return "\n".join(repr_lines)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self._repr
|
||||
|
||||
|
||||
@dataclass
|
||||
class Worker:
|
||||
actor: ActorHandle
|
||||
metadata: ActorMetadata
|
||||
resources: Dict[str, float]
|
||||
distributed_context: Optional[DistributedContext] = None
|
||||
log_file_path: Optional[str] = None
|
||||
placement_group_bundle_index: Optional[int] = None
|
||||
|
||||
@cached_property
|
||||
def _repr(self) -> str:
|
||||
indent = " "
|
||||
metadata_repr = repr(self.metadata).replace("\n", f"\n{indent}")
|
||||
context_repr = repr(self.distributed_context).replace("\n", f"\n{indent}")
|
||||
|
||||
repr_lines = [
|
||||
"Worker(",
|
||||
f"{indent}actor={repr(self.actor)},",
|
||||
f"{indent}metadata={metadata_repr},",
|
||||
f"{indent}distributed_context={context_repr},",
|
||||
f"{indent}log_file_path={repr(self.log_file_path)},",
|
||||
")",
|
||||
]
|
||||
return "\n".join(repr_lines)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self._repr
|
||||
|
||||
def execute_async(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> ObjectRef:
|
||||
"""Execute ``func`` on worker.
|
||||
|
||||
Args:
|
||||
fn: The function to execute on the worker.
|
||||
*fn_args: Positional arguments to forward to ``fn``.
|
||||
**fn_kwargs: Keyword arguments to forward to ``fn``.
|
||||
|
||||
Returns:
|
||||
(ObjectRef) An ObjectRef representing the output of func.
|
||||
|
||||
"""
|
||||
return self.actor.execute.options(name=f"execute.{fn.__name__}").remote(
|
||||
fn, *fn_args, **fn_kwargs
|
||||
)
|
||||
|
||||
|
||||
class RayTrainWorker:
|
||||
def __init__(self):
|
||||
self._callbacks: List[WorkerCallback] = []
|
||||
|
||||
def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:
|
||||
return fn(*fn_args, **fn_kwargs)
|
||||
|
||||
def run_train_fn(self, train_fn_ref: ObjectRefWrapper[Callable[[], None]]):
|
||||
"""Run the training function in a separate thread.
|
||||
|
||||
This function should return immediately, freeing up the main actor thread
|
||||
to perform other tasks such as polling the status.
|
||||
"""
|
||||
try:
|
||||
train_fn = train_fn_ref.get()
|
||||
except Exception as e:
|
||||
logger.error(f"Error deserializing the training function: {e}")
|
||||
raise
|
||||
|
||||
def train_fn_with_final_checkpoint_flush():
|
||||
result = train_fn()
|
||||
get_train_context().checkpoint_upload_threadpool.shutdown()
|
||||
|
||||
if "torch" in sys.modules:
|
||||
from ray.air._internal.torch_utils import contains_tensor
|
||||
|
||||
if contains_tensor(result):
|
||||
raise ValueError(
|
||||
"Returning objects containing Torch tensors from the "
|
||||
"training function 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."
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
# Create and start the training thread.
|
||||
logger.debug(
|
||||
f"Rank {get_train_context().get_world_rank()}: Launching training function."
|
||||
)
|
||||
get_train_context().execution_context.training_thread_runner.run(
|
||||
train_fn_with_final_checkpoint_flush
|
||||
)
|
||||
|
||||
def get_metadata(self) -> ActorMetadata:
|
||||
return ActorMetadata(
|
||||
hostname=socket.gethostname(),
|
||||
node_id=ray.get_runtime_context().get_node_id(),
|
||||
node_ip=ray.util.get_node_ip_address(),
|
||||
pid=os.getpid(),
|
||||
accelerator_ids=ray.get_runtime_context().get_accelerator_ids(),
|
||||
)
|
||||
|
||||
def mark_preempt(self, info: PreemptionInfo) -> None:
|
||||
"""Store an incoming preemption signal for the UDF to read.
|
||||
|
||||
Called by the PreemptionWatcher on every worker when a preemption
|
||||
affecting the worker group is detected.
|
||||
"""
|
||||
train_context = get_train_context()
|
||||
rank = train_context.get_world_rank()
|
||||
train_context.preemption_context.set(info)
|
||||
logger.info(
|
||||
"Rank %d received preemption signal "
|
||||
"(this_worker_preempted=%s, preempted_ranks=%s, deadline_ms=%s).",
|
||||
rank,
|
||||
rank in info.preempted_ranks,
|
||||
info.preempted_ranks,
|
||||
info.deadline_ms,
|
||||
)
|
||||
|
||||
def poll_status(self) -> WorkerStatus:
|
||||
train_context = get_train_context()
|
||||
execution_context = train_context.execution_context
|
||||
|
||||
# TODO: We can implement two phase commit here.
|
||||
# Only mark the task done when the result has been processed by the controller.
|
||||
try:
|
||||
training_report = execution_context.result_queue.get_nowait()
|
||||
execution_context.result_queue.task_done()
|
||||
except queue.Empty:
|
||||
training_report = None
|
||||
|
||||
error = execution_context.training_thread_runner.get_error()
|
||||
|
||||
# TODO: The running state should not be conflated with queue flushing.
|
||||
# Running should only be true if the user code is still running.
|
||||
# This relies on `worker_group_status.finished` returning False
|
||||
# until all training results have been flushed.
|
||||
running = execution_context.training_thread_runner.is_running() or bool(
|
||||
training_report
|
||||
)
|
||||
|
||||
return_value = (
|
||||
execution_context.training_thread_runner.get_return_value()
|
||||
if not running
|
||||
else None
|
||||
)
|
||||
|
||||
return WorkerStatus(
|
||||
running=running,
|
||||
error=error,
|
||||
training_report=training_report,
|
||||
return_value=return_value,
|
||||
preemption_info=train_context.preemption_context.get(),
|
||||
)
|
||||
|
||||
def clear_result_queue(self) -> bool:
|
||||
"""Drain the result queue, discarding any pending training reports.
|
||||
|
||||
Returns:
|
||||
True if the queue had at least one result, False if it was empty.
|
||||
"""
|
||||
execution_context = get_train_context().execution_context
|
||||
had_result = False
|
||||
while True:
|
||||
try:
|
||||
execution_context.result_queue.get_nowait()
|
||||
execution_context.result_queue.task_done()
|
||||
had_result = True
|
||||
except queue.Empty:
|
||||
break
|
||||
return had_result
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown the worker.
|
||||
|
||||
This method is not doing the real shutdown, but it is used by the worker
|
||||
group to signal the worker to stop running the training function.
|
||||
Any shutdown worker callbacks can hook on this method to implement the
|
||||
corresponding shutdown logic. Note that the shutdown logic needs to be
|
||||
thread-safe if it is running in a separate thread.
|
||||
"""
|
||||
for callback in self._callbacks:
|
||||
callback.before_worker_shutdown()
|
||||
|
||||
def init_train_context(
|
||||
self,
|
||||
train_run_context: TrainRunContext,
|
||||
distributed_context: DistributedContext,
|
||||
synchronization_actor: ActorHandle,
|
||||
storage_context: StorageContext,
|
||||
worker_callbacks: List[Union[WorkerCallback, TrainContextCallback]],
|
||||
controller_actor: ActorHandle,
|
||||
dataset_shard_provider: Optional["DatasetShardProvider"] = None,
|
||||
checkpoint: Optional[Checkpoint] = None,
|
||||
has_validation_fn: Optional[bool] = None,
|
||||
current_report_index: int = 0,
|
||||
):
|
||||
self._callbacks = [c for c in worker_callbacks if isinstance(c, WorkerCallback)]
|
||||
context_callbacks_to_propagate = [
|
||||
c for c in worker_callbacks if isinstance(c, TrainContextCallback)
|
||||
]
|
||||
context = TrainContext(
|
||||
train_run_context=train_run_context,
|
||||
distributed_context=distributed_context,
|
||||
execution_context=ExecutionContext(
|
||||
synchronization_actor=synchronization_actor,
|
||||
# Make the queue size 1 to avoid building up too
|
||||
# many unprocessed results.
|
||||
result_queue=queue.Queue(maxsize=1),
|
||||
training_thread_runner=ThreadRunner(),
|
||||
train_context_callbacks=context_callbacks_to_propagate,
|
||||
),
|
||||
storage_context=storage_context,
|
||||
preemption_context=PreemptionContext(),
|
||||
controller_actor=controller_actor,
|
||||
checkpoint=checkpoint,
|
||||
dataset_shard_provider=dataset_shard_provider,
|
||||
has_validation_fn=has_validation_fn,
|
||||
current_report_index=current_report_index,
|
||||
)
|
||||
# Configure the train and root logger for the worker processes.
|
||||
if ray_constants.env_bool(
|
||||
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR, DEFAULT_ENABLE_WORKER_LOGGING
|
||||
):
|
||||
LoggingManager.configure_worker_logger(context)
|
||||
patch_print_function()
|
||||
# Set the train context global variable for the worker.
|
||||
set_train_context(context)
|
||||
|
||||
# user facing train fn utils
|
||||
set_train_fn_utils(DistributedTrainFnUtils())
|
||||
|
||||
for callback in self._callbacks:
|
||||
callback.after_init_train_context()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
from .logging import LoggingManager
|
||||
|
||||
__all__ = ["LoggingManager"]
|
||||
@@ -0,0 +1,340 @@
|
||||
import logging.config
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
|
||||
import ray
|
||||
from ray._common.filters import CoreContextFilter
|
||||
from ray._common.formatters import JSONFormatter
|
||||
from ray._private.log import PlainRayHandler
|
||||
from ray.train.v2._internal.execution.context import TrainContext, TrainRunContext
|
||||
from ray.train.v2._internal.util import get_module_name
|
||||
|
||||
|
||||
class TrainContextFilter(logging.Filter):
|
||||
"""Add Ray Train metadata to the log records.
|
||||
|
||||
This filter is applied to Ray Train controller and worker processes.
|
||||
"""
|
||||
|
||||
# Log keys for Ray Train controller and worker processes.
|
||||
class LogKey(str, Enum):
|
||||
RUN_NAME = "run_name"
|
||||
COMPONENT = "component"
|
||||
WORLD_RANK = "world_rank"
|
||||
LOCAL_RANK = "local_rank"
|
||||
NODE_RANK = "node_rank"
|
||||
|
||||
# Ray Train Component by process types
|
||||
class TrainComponent(str, Enum):
|
||||
CONTROLLER = "controller"
|
||||
WORKER = "worker"
|
||||
|
||||
def __init__(self, context: Union[TrainRunContext, TrainContext]):
|
||||
self._is_worker: bool = isinstance(context, TrainContext)
|
||||
if self._is_worker:
|
||||
self._run_name: str = context.train_run_context.get_run_config().name
|
||||
self._world_rank: int = context.get_world_rank()
|
||||
self._local_rank: int = context.get_local_rank()
|
||||
self._node_rank: int = context.get_node_rank()
|
||||
self._component: str = TrainContextFilter.TrainComponent.WORKER
|
||||
else:
|
||||
self._run_name: str = context.get_run_config().name
|
||||
self._component: str = TrainContextFilter.TrainComponent.CONTROLLER
|
||||
|
||||
def controller_filter(self, record):
|
||||
# Add the run_id and component to Ray Train controller processes.
|
||||
setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
|
||||
setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
|
||||
return True
|
||||
|
||||
def worker_filter(self, record):
|
||||
# Add the run_id and component to Ray Train worker processes.
|
||||
setattr(record, TrainContextFilter.LogKey.RUN_NAME, self._run_name)
|
||||
setattr(record, TrainContextFilter.LogKey.COMPONENT, self._component)
|
||||
# Add all the rank related information to the log record for worker processes.
|
||||
setattr(record, TrainContextFilter.LogKey.WORLD_RANK, self._world_rank)
|
||||
setattr(record, TrainContextFilter.LogKey.LOCAL_RANK, self._local_rank)
|
||||
setattr(record, TrainContextFilter.LogKey.NODE_RANK, self._node_rank)
|
||||
return True
|
||||
|
||||
def filter(self, record):
|
||||
if self._is_worker:
|
||||
return self.worker_filter(record)
|
||||
else:
|
||||
return self.controller_filter(record)
|
||||
|
||||
|
||||
class TrainLogLevelFilter(logging.Filter):
|
||||
"""Filter that applies log level filtering only to ray.train log records."""
|
||||
|
||||
def __init__(self, log_level: str = "INFO"):
|
||||
super().__init__()
|
||||
self._log_level = getattr(logging, log_level)
|
||||
|
||||
def filter(self, record):
|
||||
if record.name == "ray.train" or record.name.startswith("ray.train."):
|
||||
return record.levelno >= self._log_level
|
||||
return True
|
||||
|
||||
|
||||
class SessionFileHandler(logging.Handler):
|
||||
"""A handler that writes to a log file in the Ray session directory.
|
||||
|
||||
The Ray session directory isn't available until Ray is initialized, so any logs
|
||||
emitted before Ray is initialized will be lost.
|
||||
This handler will not create the file handler until you emit a log record.
|
||||
|
||||
Args:
|
||||
filename: The name of the log file. The file is created in the 'logs/train'
|
||||
directory of the Ray session directory.
|
||||
"""
|
||||
|
||||
# TODO (hpguo): This handler class is shared by both Ray Train and ray data. We
|
||||
# should move this to ray core and make it available to both libraries.
|
||||
|
||||
def __init__(self, filename: str):
|
||||
super().__init__()
|
||||
self._filename = filename
|
||||
self._handler = None
|
||||
self._formatter = None
|
||||
self._path = None
|
||||
|
||||
def emit(self, record):
|
||||
if self._handler is None:
|
||||
self._try_create_handler()
|
||||
if self._handler is not None:
|
||||
self._handler.emit(record)
|
||||
|
||||
def setFormatter(self, fmt: logging.Formatter) -> None:
|
||||
if self._handler is not None:
|
||||
self._handler.setFormatter(fmt)
|
||||
self._formatter = fmt
|
||||
|
||||
def get_log_file_path(self) -> Optional[str]:
|
||||
if self._handler is None:
|
||||
self._try_create_handler()
|
||||
return self._path
|
||||
|
||||
def _try_create_handler(self):
|
||||
assert self._handler is None
|
||||
|
||||
# Get the Ray Train log directory. If not in a Ray session, return.
|
||||
# This handler will only be created within a Ray session.
|
||||
log_directory = LoggingManager.get_log_directory()
|
||||
if log_directory is None:
|
||||
return
|
||||
|
||||
os.makedirs(log_directory, exist_ok=True)
|
||||
|
||||
# Create the log file.
|
||||
self._path = os.path.join(log_directory, self._filename)
|
||||
self._handler = logging.FileHandler(self._path)
|
||||
if self._formatter is not None:
|
||||
self._handler.setFormatter(self._formatter)
|
||||
|
||||
|
||||
class LoggingManager:
|
||||
"""
|
||||
A utility class for managing the logging configuration of Ray Train.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _get_base_logger_config_dict(
|
||||
context: Union[TrainRunContext, TrainContext],
|
||||
) -> dict:
|
||||
"""Return the base logging configuration dictionary."""
|
||||
log_level = LoggingManager._resolve_log_level(context)
|
||||
# Using Ray worker ID as the file identifier where logs are written to.
|
||||
file_identifier = ray.get_runtime_context().get_worker_id()
|
||||
# Return the base logging configuration as a Python dictionary.
|
||||
return {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
"formatters": {
|
||||
"ray_json": {"class": get_module_name(JSONFormatter)},
|
||||
},
|
||||
"filters": {
|
||||
"core_context_filter": {"()": CoreContextFilter},
|
||||
"train_context_filter": {"()": TrainContextFilter, "context": context},
|
||||
"train_log_level_filter": {
|
||||
"()": TrainLogLevelFilter,
|
||||
"log_level": log_level,
|
||||
},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"class": get_module_name(PlainRayHandler),
|
||||
"filters": ["train_log_level_filter"],
|
||||
},
|
||||
"file_train_sys_controller": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-sys-controller-{file_identifier}.log",
|
||||
"filters": ["core_context_filter", "train_context_filter"],
|
||||
},
|
||||
"file_train_app_controller": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-app-controller-{file_identifier}.log",
|
||||
"filters": [
|
||||
"core_context_filter",
|
||||
"train_context_filter",
|
||||
"train_log_level_filter",
|
||||
],
|
||||
},
|
||||
"file_train_sys_worker": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-sys-worker-{file_identifier}.log",
|
||||
"filters": ["core_context_filter", "train_context_filter"],
|
||||
},
|
||||
"file_train_app_worker": {
|
||||
"class": get_module_name(SessionFileHandler),
|
||||
"formatter": "ray_json",
|
||||
"filename": f"ray-train-app-worker-{file_identifier}.log",
|
||||
"filters": [
|
||||
"core_context_filter",
|
||||
"train_context_filter",
|
||||
"train_log_level_filter",
|
||||
],
|
||||
},
|
||||
},
|
||||
"loggers": {},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_log_level(
|
||||
context: Union[TrainRunContext, TrainContext],
|
||||
) -> str:
|
||||
"""Returns the log level from RunConfig's LoggingConfig."""
|
||||
if isinstance(context, TrainContext):
|
||||
run_config = context.train_run_context.get_run_config()
|
||||
else:
|
||||
run_config = context.get_run_config()
|
||||
|
||||
return run_config.logging_config.log_level
|
||||
|
||||
@staticmethod
|
||||
def _get_controller_logger_config_dict(context: TrainRunContext) -> dict:
|
||||
"""Return the controller logger configuration dictionary.
|
||||
|
||||
On the controller process, only the `ray.train` logger is configured.
|
||||
It is broadly set to level DEBUG, with downstream processing by log handlers.
|
||||
This logger emits logs to the following three locations:
|
||||
- `file_train_sys_controller`: Ray Train system logs.
|
||||
- `file_train_app_controller`: Ray Train application logs.
|
||||
- `console`: Logs to the console.
|
||||
"""
|
||||
|
||||
config_dict = LoggingManager._get_base_logger_config_dict(context)
|
||||
config_dict["loggers"]["ray.train"] = {
|
||||
"level": "DEBUG",
|
||||
"handlers": [
|
||||
"file_train_sys_controller",
|
||||
"file_train_app_controller",
|
||||
"console",
|
||||
],
|
||||
"propagate": False,
|
||||
}
|
||||
return config_dict
|
||||
|
||||
@staticmethod
|
||||
def _get_worker_logger_config_dict(context: TrainContext) -> dict:
|
||||
"""Return the worker loggers configuration dictionary.
|
||||
|
||||
On the worker process, there are two loggers being configured:
|
||||
|
||||
First, the `ray.train` logger is configured and emits logs to the
|
||||
following three locations:
|
||||
- `file_train_sys_worker`: Ray Train system logs.
|
||||
- `file_train_app_worker`: Ray Train application logs.
|
||||
- `console`: Logs to the console.
|
||||
It is broadly set to level DEBUG, with downstream processing by log handlers.
|
||||
|
||||
Second, the root logger is configured and emits logs to the following
|
||||
two locations:
|
||||
- `console`: Logs to the console.
|
||||
- `file_train_app_worker`: Ray Train application logs.
|
||||
The root logger will not emit Ray Train system logs and thus not writing to
|
||||
`file_train_sys_worker` file handler.
|
||||
"""
|
||||
|
||||
config_dict = LoggingManager._get_base_logger_config_dict(context)
|
||||
config_dict["loggers"]["ray.train"] = {
|
||||
"level": "DEBUG",
|
||||
"handlers": ["file_train_sys_worker", "file_train_app_worker", "console"],
|
||||
"propagate": False,
|
||||
}
|
||||
config_dict["root"] = {
|
||||
"level": "INFO",
|
||||
"handlers": ["file_train_app_worker", "console"],
|
||||
}
|
||||
return config_dict
|
||||
|
||||
@staticmethod
|
||||
def configure_controller_logger(context: TrainRunContext) -> None:
|
||||
"""
|
||||
Configure the logger on the controller process, which is the `ray.train` logger.
|
||||
"""
|
||||
config = LoggingManager._get_controller_logger_config_dict(context)
|
||||
logging.config.dictConfig(config)
|
||||
# TODO: Return the controller log file path.
|
||||
|
||||
@staticmethod
|
||||
def configure_worker_logger(context: TrainContext) -> None:
|
||||
"""
|
||||
Configure the loggers on the worker process, which contains the
|
||||
`ray.train` logger and the root logger.
|
||||
"""
|
||||
config = LoggingManager._get_worker_logger_config_dict(context)
|
||||
logging.config.dictConfig(config)
|
||||
# TODO: Return the worker log file path.
|
||||
|
||||
@staticmethod
|
||||
def get_log_directory() -> Optional[str]:
|
||||
"""Return the directory where Ray Train writes log files.
|
||||
|
||||
If not in a Ray session, return None.
|
||||
|
||||
This path looks like: "/tmp/ray/session_xxx/logs/train/"
|
||||
"""
|
||||
global_node = ray._private.worker._global_node
|
||||
|
||||
if global_node is None:
|
||||
return None
|
||||
|
||||
root_dir = global_node.get_session_dir_path()
|
||||
return os.path.join(root_dir, "logs", "train")
|
||||
|
||||
|
||||
def get_train_application_controller_log_path() -> Optional[str]:
|
||||
"""
|
||||
Return the path to the file train application controller log file.
|
||||
"""
|
||||
# TODO: This is a temporary solution. We should return the log file path in
|
||||
# the `configure_controller_logger` function.
|
||||
logger = logging.getLogger("ray.train")
|
||||
for handler in logger.handlers:
|
||||
if (
|
||||
isinstance(handler, SessionFileHandler)
|
||||
and "ray-train-app-controller" in handler._filename
|
||||
):
|
||||
return handler.get_log_file_path()
|
||||
return None
|
||||
|
||||
|
||||
def get_train_application_worker_log_path() -> Optional[str]:
|
||||
"""
|
||||
Return the path to the file train application worker log file.
|
||||
"""
|
||||
# TODO: This is a temporary solution. We should return the log file path in
|
||||
# the `configure_worker_logger` function.
|
||||
logger = logging.getLogger("ray.train")
|
||||
for handler in logger.handlers:
|
||||
if (
|
||||
isinstance(handler, SessionFileHandler)
|
||||
and "ray-train-app-worker" in handler._filename
|
||||
):
|
||||
return handler.get_log_file_path()
|
||||
return None
|
||||
@@ -0,0 +1,76 @@
|
||||
import builtins
|
||||
import contextlib
|
||||
import logging
|
||||
import sys
|
||||
from typing import Callable
|
||||
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_ENABLE_PRINT_PATCH,
|
||||
ENABLE_PRINT_PATCH_ENV_VAR,
|
||||
)
|
||||
|
||||
# Save the original print function
|
||||
_original_print = builtins.print
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def print_context_manager(print_fn: Callable):
|
||||
"""Context manager to set the builtin print function as print_fn."""
|
||||
current_print = builtins.print
|
||||
builtins.print = print_fn
|
||||
yield
|
||||
builtins.print = current_print
|
||||
|
||||
|
||||
def redirected_print(*objects, sep=" ", end="\n", file=None, flush=False):
|
||||
"""Implement python's print function to redirect logs to Train's logger.
|
||||
|
||||
If the file is set to anything other than stdout, stderr, or None, call the
|
||||
builtin print. Else, construct the message and redirect to Train's logger.
|
||||
|
||||
This makes sure that print to customized file in user defined function will not
|
||||
be overwritten by the redirected print function.
|
||||
|
||||
See https://docs.python.org/3/library/functions.html#print
|
||||
"""
|
||||
# TODO (hpguo): This handler class is shared by both ray train and ray serve. We
|
||||
# should move this to ray core and make it available to both libraries.
|
||||
|
||||
if file not in [sys.stdout, sys.stderr, None]:
|
||||
_original_print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
return
|
||||
|
||||
# If sys.stdout/stderr has been redirected (e.g. contextlib.redirect_stdout(),
|
||||
# or wrapping by libraries like wandb / MLflow / colorama / IPython), tee to
|
||||
# the original print so the redirect target also receives the output. The
|
||||
# logger still gets the message below, so structured logs aren't silently
|
||||
# dropped when a third-party library wraps the stream.
|
||||
if (file in (sys.stdout, None) and sys.stdout is not sys.__stdout__) or (
|
||||
file is sys.stderr and sys.stderr is not sys.__stderr__
|
||||
):
|
||||
_original_print(*objects, sep=sep, end=end, file=file, flush=flush)
|
||||
|
||||
root_logger = logging.getLogger()
|
||||
message = sep.join(map(str, objects))
|
||||
# Use the original `print` method for the scope of the logger call, in order to
|
||||
# avoid infinite recursion errors if any exceptions get raised (since exception
|
||||
# handling involves another `print(..., file=sys.stderr)`.
|
||||
# Note that an exception being raised here is not expected (e.g. it would be a
|
||||
# bug in our own logging code), so this is just to keep the error logs sane
|
||||
# during development.
|
||||
with print_context_manager(_original_print):
|
||||
# We want this log to be associated with the line of code where user calls
|
||||
# `print`, which is stacklevel 2.
|
||||
# Frame [stacklevel]:
|
||||
# User's call to print [2] -> `redirected_print` [1] -> root_logger.log [0]
|
||||
root_logger.log(logging.INFO, message, stacklevel=2)
|
||||
|
||||
|
||||
def patch_print_function() -> None:
|
||||
"""
|
||||
Patch the print function to redirect logs to Train's logger.
|
||||
Only patch the print function if the environment variable is set to "1"
|
||||
"""
|
||||
if env_bool(ENABLE_PRINT_PATCH_ENV_VAR, DEFAULT_ENABLE_PRINT_PATCH):
|
||||
builtins.print = redirected_print
|
||||
@@ -0,0 +1,166 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Dict, Generic, Optional, Tuple, TypeVar
|
||||
|
||||
from ray.util.metrics import Gauge
|
||||
|
||||
RUN_NAME_TAG_KEY = "ray_train_run_name"
|
||||
RUN_ID_TAG_KEY = "ray_train_run_id"
|
||||
|
||||
T = TypeVar("T")
|
||||
E = TypeVar("E", bound=Enum)
|
||||
|
||||
|
||||
class Metric(ABC):
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
default: T,
|
||||
description: str,
|
||||
base_tags: Dict[str, str],
|
||||
):
|
||||
"""
|
||||
Initialize a new metric.
|
||||
|
||||
Args:
|
||||
name: The name of the metric.
|
||||
default: The default value of the metric.
|
||||
description: The description of the metric.
|
||||
base_tags: The base tags for the metric.
|
||||
"""
|
||||
self._default = default
|
||||
self._base_tags = base_tags
|
||||
self._gauge = Gauge(
|
||||
name,
|
||||
description=description,
|
||||
tag_keys=self._get_tag_keys(),
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def record(self, value: T):
|
||||
"""Update the metric value.
|
||||
|
||||
Args:
|
||||
value: The value to update the metric with.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_value(self) -> T:
|
||||
"""Get the value of the metric.
|
||||
|
||||
Returns:
|
||||
The value of the metric. If the metric has not been recorded,
|
||||
the default value is returned.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def reset(self):
|
||||
"""Reset values and clean up resources."""
|
||||
pass
|
||||
|
||||
def _get_tag_keys(self) -> Tuple[str, ...]:
|
||||
return tuple(self._base_tags.keys())
|
||||
|
||||
|
||||
class TimeMetric(Metric):
|
||||
"""A metric for tracking elapsed time."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
base_tags: Dict[str, str],
|
||||
):
|
||||
self._current_value = 0.0
|
||||
super().__init__(
|
||||
name=name,
|
||||
default=0.0,
|
||||
description=description,
|
||||
base_tags=base_tags,
|
||||
)
|
||||
|
||||
def record(self, value: float):
|
||||
"""Update the time metric value by accumulating the time.
|
||||
|
||||
Args:
|
||||
value: The time value to increment the metric by.
|
||||
"""
|
||||
self._current_value += value
|
||||
self._gauge.set(self._current_value, self._base_tags)
|
||||
|
||||
def get_value(self) -> float:
|
||||
return self._current_value
|
||||
|
||||
def reset(self):
|
||||
self._current_value = self._default
|
||||
self._gauge.set(self._default, self._base_tags)
|
||||
|
||||
|
||||
class EnumMetric(Metric, Generic[E]):
|
||||
"""A metric for tracking enum values."""
|
||||
|
||||
DEFAULT_VALUE = 0
|
||||
RECORDED_VALUE = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
base_tags: Dict[str, str],
|
||||
enum_tag_key: str,
|
||||
):
|
||||
self._enum_tag_key = enum_tag_key
|
||||
self._current_value: Optional[E] = None
|
||||
super().__init__(
|
||||
name=name,
|
||||
default=self.DEFAULT_VALUE,
|
||||
description=description,
|
||||
base_tags=base_tags,
|
||||
)
|
||||
|
||||
def record(self, enum_value: E) -> None:
|
||||
"""Record a specific enum value.
|
||||
|
||||
The metric will be reset to 0 for the previous value and set to 1 for the new value.
|
||||
|
||||
Args:
|
||||
enum_value: The enum value to record for.
|
||||
"""
|
||||
if enum_value == self._current_value:
|
||||
return
|
||||
|
||||
if self._current_value is not None:
|
||||
previous_tags = self._get_tags(self._current_value)
|
||||
self._gauge.set(self._default, previous_tags)
|
||||
|
||||
current_tags = self._get_tags(enum_value)
|
||||
self._gauge.set(self.RECORDED_VALUE, current_tags)
|
||||
|
||||
self._current_value = enum_value
|
||||
|
||||
def get_value(self, enum_value: E) -> int:
|
||||
"""Get the value for a specific enum value.
|
||||
|
||||
Args:
|
||||
enum_value: The enum value to get the value for
|
||||
|
||||
Returns:
|
||||
The value for the enum value
|
||||
"""
|
||||
return int(enum_value == self._current_value)
|
||||
|
||||
def reset(self):
|
||||
if self._current_value is not None:
|
||||
tags = self._get_tags(self._current_value)
|
||||
self._gauge.set(self._default, tags)
|
||||
self._current_value = None
|
||||
|
||||
def _get_tag_keys(self) -> Tuple[str, ...]:
|
||||
return tuple(self._base_tags.keys()) + (self._enum_tag_key,)
|
||||
|
||||
def _get_tags(self, enum_value: E) -> Dict[str, str]:
|
||||
tags = self._base_tags.copy()
|
||||
tags[self._enum_tag_key] = enum_value.name
|
||||
return tags
|
||||
@@ -0,0 +1,66 @@
|
||||
from typing import Dict, Union
|
||||
|
||||
from ray.train.v2._internal.execution.controller.state import TrainControllerStateType
|
||||
from ray.train.v2._internal.metrics.base import (
|
||||
RUN_ID_TAG_KEY,
|
||||
RUN_NAME_TAG_KEY,
|
||||
EnumMetric,
|
||||
TimeMetric,
|
||||
)
|
||||
|
||||
|
||||
class ControllerMetrics:
|
||||
"""Factory for creating controller-specific metrics.
|
||||
|
||||
This class defines all metrics used to track the state and performance of the
|
||||
training controller. Each metric is defined with its name, type, default value,
|
||||
description, and required tags.
|
||||
"""
|
||||
|
||||
# ===== Metric Names =====
|
||||
CONTROLLER_STATE = "train_controller_state"
|
||||
WORKER_GROUP_START_TOTAL_TIME_S = "train_worker_group_start_total_time_s"
|
||||
WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S = "train_worker_group_shutdown_total_time_s"
|
||||
|
||||
# ===== Tag Keys =====
|
||||
CONTROLLER_STATE_TAG_KEY = "ray_train_controller_state"
|
||||
|
||||
@classmethod
|
||||
def _create_time_metric(
|
||||
cls, name: str, description: str, base_tags: Dict[str, str]
|
||||
) -> TimeMetric:
|
||||
return TimeMetric(
|
||||
name=name,
|
||||
description=description,
|
||||
base_tags=base_tags,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _create_controller_state_metric(
|
||||
cls, base_tags: Dict[str, str]
|
||||
) -> EnumMetric[TrainControllerStateType]:
|
||||
return EnumMetric[TrainControllerStateType](
|
||||
name=cls.CONTROLLER_STATE,
|
||||
description="Current state of the Ray Train controller",
|
||||
base_tags=base_tags,
|
||||
enum_tag_key=cls.CONTROLLER_STATE_TAG_KEY,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_controller_metrics(
|
||||
cls, run_name: str, run_id: str
|
||||
) -> Dict[str, Union[TimeMetric, EnumMetric[TrainControllerStateType]]]:
|
||||
base_tags = {RUN_NAME_TAG_KEY: run_name, RUN_ID_TAG_KEY: run_id}
|
||||
return {
|
||||
cls.WORKER_GROUP_START_TOTAL_TIME_S: cls._create_time_metric(
|
||||
cls.WORKER_GROUP_START_TOTAL_TIME_S,
|
||||
"Total time taken to start the worker group",
|
||||
base_tags,
|
||||
),
|
||||
cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S: cls._create_time_metric(
|
||||
cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S,
|
||||
"Total time taken to shutdown the worker group",
|
||||
base_tags,
|
||||
),
|
||||
cls.CONTROLLER_STATE: cls._create_controller_state_metric(base_tags),
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
from typing import Dict
|
||||
|
||||
from ray.train.v2._internal.metrics.base import (
|
||||
RUN_ID_TAG_KEY,
|
||||
RUN_NAME_TAG_KEY,
|
||||
TimeMetric,
|
||||
)
|
||||
|
||||
WORKER_WORLD_RANK_TAG_KEY = "ray_train_worker_world_rank"
|
||||
WORKER_ACTOR_ID_TAG_KEY = "ray_train_worker_actor_id"
|
||||
|
||||
|
||||
class WorkerMetrics:
|
||||
"""Factory for creating worker-specific metrics.
|
||||
|
||||
This class defines all metrics used to track the state and performance of the
|
||||
training workers. Each metric is defined with its name, type, default value,
|
||||
description, and required tags.
|
||||
"""
|
||||
|
||||
# ===== Metric Names =====
|
||||
REPORT_TOTAL_BLOCKED_TIME_S = "train_report_total_blocked_time_s"
|
||||
CHECKPOINT_SYNC_TOTAL_TIME_S = "train_checkpoint_sync_total_time_s"
|
||||
CHECKPOINT_TRANSFER_TOTAL_TIME_S = "train_checkpoint_transfer_total_time_s"
|
||||
|
||||
@classmethod
|
||||
def _create_time_metric(
|
||||
cls, name: str, description: str, base_tags: Dict[str, str]
|
||||
) -> TimeMetric:
|
||||
"""Create a time-based metric."""
|
||||
return TimeMetric(
|
||||
name=name,
|
||||
description=description,
|
||||
base_tags=base_tags,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_worker_metrics(
|
||||
cls, run_name: str, run_id: str, world_rank: int, worker_actor_id: str
|
||||
) -> Dict[str, TimeMetric]:
|
||||
"""Get all worker metrics."""
|
||||
base_tags = {
|
||||
RUN_NAME_TAG_KEY: run_name,
|
||||
RUN_ID_TAG_KEY: run_id,
|
||||
WORKER_WORLD_RANK_TAG_KEY: str(world_rank),
|
||||
WORKER_ACTOR_ID_TAG_KEY: worker_actor_id,
|
||||
}
|
||||
return {
|
||||
cls.REPORT_TOTAL_BLOCKED_TIME_S: cls._create_time_metric(
|
||||
cls.REPORT_TOTAL_BLOCKED_TIME_S,
|
||||
"Cumulative time in seconds to report a checkpoint to the storage.",
|
||||
base_tags,
|
||||
),
|
||||
cls.CHECKPOINT_SYNC_TOTAL_TIME_S: cls._create_time_metric(
|
||||
cls.CHECKPOINT_SYNC_TOTAL_TIME_S,
|
||||
"Cumulative time in seconds spent synchronizing the checkpoint directory name across all ranks.",
|
||||
base_tags,
|
||||
),
|
||||
cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S: cls._create_time_metric(
|
||||
cls.CHECKPOINT_TRANSFER_TOTAL_TIME_S,
|
||||
"Cumulative time in seconds spent transferring checkpoint files to storage.",
|
||||
base_tags,
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
from ray.train.constants import V2_MIGRATION_GUIDE_MESSAGE
|
||||
|
||||
FAIL_FAST_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.FailureConfig(fail_fast)` is deprecated since it is "
|
||||
"only relevant in the context of multiple trials running in Ray Tune. "
|
||||
"This parameter is still available in `ray.tune.FailureConfig` "
|
||||
"for passing into a `ray.tune.Tuner`. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
TRAINER_RESOURCES_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.ScalingConfig(trainer_resources)` is deprecated. "
|
||||
"This parameter was an advanced configuration that specified "
|
||||
"resources for the Ray Train driver actor, which doesn't "
|
||||
"need to reserve logical resources because it doesn't perform "
|
||||
"any heavy computation. "
|
||||
"Only the `resources_per_worker` parameter should be used "
|
||||
"to specify resources for the training workers. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
VERBOSE_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(verbose)` is deprecated. "
|
||||
"This parameter controls Ray Tune logging verbosity, "
|
||||
"and is only relevant when using Ray Tune. "
|
||||
"This parameter is still available in `ray.tune.RunConfig` "
|
||||
"for passing into a `ray.tune.Tuner`. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
LOG_TO_FILE_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(log_to_file)` is deprecated. "
|
||||
"The Ray Train driver actor and the training worker actors "
|
||||
"already log stdout/stderr as part of Ray's logging system. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
STOP_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(stop)` is deprecated. "
|
||||
"This parameter is only relevant when using Ray Tune "
|
||||
"and is still available in `ray.tune.RunConfig` "
|
||||
"for passing into a `ray.tune.Tuner`. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
CALLBACKS_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(callbacks: List[ray.tune.Callback])` is deprecated. "
|
||||
"Ray Train no longer accepts Ray Tune callbacks, since the Ray Train "
|
||||
"execution backend is being separated from Ray Tune. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
PROGRESS_REPORTER_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(progress_reporter)` is deprecated. "
|
||||
"This parameter controls the Ray Tune console output reporter, "
|
||||
"and is only relevant when using Ray Tune. "
|
||||
"This parameter is still available in `ray.tune.RunConfig` "
|
||||
"for passing into a `ray.tune.Tuner`. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
|
||||
SYNC_CONFIG_DEPRECATION_MESSAGE = (
|
||||
"`ray.train.RunConfig(sync_config)` is deprecated. "
|
||||
"This configuration controls advanced syncing behavior, "
|
||||
"which is either not supported or not relevant in the reworked Ray Train. "
|
||||
"This parameter is still available in `ray.tune.RunConfig` "
|
||||
"for passing into a `ray.tune.Tuner`. "
|
||||
"The `SyncConfig` class has been moved to `ray.tune.SyncConfig`. "
|
||||
f"{V2_MIGRATION_GUIDE_MESSAGE}"
|
||||
)
|
||||
@@ -0,0 +1,313 @@
|
||||
import logging
|
||||
|
||||
from google.protobuf.struct_pb2 import Struct
|
||||
|
||||
from ray.core.generated.export_train_state_pb2 import (
|
||||
ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
|
||||
ExportTrainRunEventData as ProtoTrainRun,
|
||||
)
|
||||
from ray.dashboard.modules.metrics.dashboards.common import Panel
|
||||
from ray.dashboard.modules.metrics.dashboards.train_dashboard_panels import (
|
||||
TRAIN_RUN_PANELS,
|
||||
TRAIN_WORKER_PANELS,
|
||||
)
|
||||
from ray.train.v2._internal.state.schema import (
|
||||
ActorStatus,
|
||||
BackendConfig,
|
||||
DataConfig,
|
||||
ExecutionOptions,
|
||||
RunAttemptStatus,
|
||||
RunConfig,
|
||||
RunSettings,
|
||||
RunStatus,
|
||||
ScalingConfig,
|
||||
TrainRun,
|
||||
TrainRunAttempt,
|
||||
TrainWorker,
|
||||
)
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
|
||||
# Increment each time the exported Train schema changes (proto, pydantic, or
|
||||
# exported json) so downstream consumers can distinguish schema versions.
|
||||
TRAIN_SCHEMA_VERSION = 4
|
||||
RAY_TRAIN_VERSION = 2
|
||||
|
||||
# Status mapping dictionaries
|
||||
_ACTOR_STATUS_MAP = {
|
||||
ActorStatus.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
|
||||
ActorStatus.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
|
||||
}
|
||||
|
||||
_RUN_ATTEMPT_STATUS_MAP = {
|
||||
RunAttemptStatus.PENDING: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
|
||||
RunAttemptStatus.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
|
||||
RunAttemptStatus.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
|
||||
RunAttemptStatus.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
|
||||
RunAttemptStatus.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
|
||||
}
|
||||
|
||||
_RUN_STATUS_MAP = {
|
||||
RunStatus.INITIALIZING: ProtoTrainRun.RunStatus.INITIALIZING,
|
||||
RunStatus.SCHEDULING: ProtoTrainRun.RunStatus.SCHEDULING,
|
||||
RunStatus.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
|
||||
RunStatus.RESTARTING: ProtoTrainRun.RunStatus.RESTARTING,
|
||||
RunStatus.RESIZING: ProtoTrainRun.RunStatus.RESIZING,
|
||||
RunStatus.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
|
||||
RunStatus.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
|
||||
RunStatus.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
|
||||
}
|
||||
|
||||
_TRAINING_FRAMEWORK_MAP = {
|
||||
None: ProtoTrainRun.BackendConfig.TrainingFramework.TRAINING_FRAMEWORK_UNSPECIFIED,
|
||||
TrainingFramework.TORCH: ProtoTrainRun.BackendConfig.TrainingFramework.TORCH,
|
||||
TrainingFramework.JAX: ProtoTrainRun.BackendConfig.TrainingFramework.JAX,
|
||||
TrainingFramework.TENSORFLOW: ProtoTrainRun.BackendConfig.TrainingFramework.TENSORFLOW,
|
||||
TrainingFramework.XGBOOST: ProtoTrainRun.BackendConfig.TrainingFramework.XGBOOST,
|
||||
TrainingFramework.LIGHTGBM: ProtoTrainRun.BackendConfig.TrainingFramework.LIGHTGBM,
|
||||
}
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _dict_to_struct(d: dict) -> Struct:
|
||||
"""Returns a protobuf Struct from a dictionary."""
|
||||
s = Struct()
|
||||
s.update(d)
|
||||
return s
|
||||
|
||||
|
||||
def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
|
||||
"""Convert resources dictionary to protobuf TrainResources."""
|
||||
return ProtoTrainRunAttempt.TrainResources(resources=resources)
|
||||
|
||||
|
||||
def _to_proto_worker(worker: TrainWorker) -> ProtoTrainRunAttempt.TrainWorker:
|
||||
"""Convert TrainWorker to protobuf format."""
|
||||
status = None
|
||||
if worker.status is not None:
|
||||
status = _ACTOR_STATUS_MAP[worker.status]
|
||||
|
||||
return ProtoTrainRunAttempt.TrainWorker(
|
||||
world_rank=worker.world_rank,
|
||||
local_rank=worker.local_rank,
|
||||
node_rank=worker.node_rank,
|
||||
actor_id=bytes.fromhex(worker.actor_id),
|
||||
node_id=bytes.fromhex(worker.node_id),
|
||||
node_ip=worker.node_ip,
|
||||
pid=worker.pid,
|
||||
gpu_ids=worker.gpu_ids,
|
||||
status=status,
|
||||
resources=_to_proto_resources(worker.resources.resources),
|
||||
log_file_path=worker.log_file_path,
|
||||
)
|
||||
|
||||
|
||||
# Main conversion functions
|
||||
def train_run_attempt_to_proto(attempt: TrainRunAttempt) -> ProtoTrainRunAttempt:
|
||||
"""Convert TrainRunAttempt to protobuf format."""
|
||||
proto_attempt = ProtoTrainRunAttempt(
|
||||
schema_version=TRAIN_SCHEMA_VERSION,
|
||||
ray_train_version=RAY_TRAIN_VERSION,
|
||||
run_id=attempt.run_id,
|
||||
attempt_id=attempt.attempt_id,
|
||||
status=_RUN_ATTEMPT_STATUS_MAP[attempt.status],
|
||||
status_detail=attempt.status_detail,
|
||||
start_time_ns=attempt.start_time_ns,
|
||||
end_time_ns=attempt.end_time_ns,
|
||||
resources=[_to_proto_resources(r.resources) for r in attempt.resources],
|
||||
workers=[_to_proto_worker(w) for w in attempt.workers],
|
||||
)
|
||||
|
||||
return proto_attempt
|
||||
|
||||
|
||||
def _to_proto_dashboard_panel(panel: Panel) -> ProtoTrainRun.DashboardPanelMetadata:
|
||||
"""Convert Dashboard Panel to protobuf format."""
|
||||
proto_panel = ProtoTrainRun.DashboardPanelMetadata(
|
||||
id=str(panel.id),
|
||||
title=panel.title,
|
||||
)
|
||||
|
||||
return proto_panel
|
||||
|
||||
|
||||
def to_proto_backend_config(
|
||||
backend_config: BackendConfig,
|
||||
) -> ProtoTrainRun.BackendConfig:
|
||||
"""Convert BackendConfig to protobuf format."""
|
||||
proto_backend_config = ProtoTrainRun.BackendConfig(
|
||||
framework=_TRAINING_FRAMEWORK_MAP[backend_config.framework],
|
||||
)
|
||||
|
||||
proto_backend_config.config.CopyFrom(_dict_to_struct(backend_config.config))
|
||||
|
||||
return proto_backend_config
|
||||
|
||||
|
||||
def to_proto_scaling_config(
|
||||
scaling_config: ScalingConfig,
|
||||
) -> ProtoTrainRun.ScalingConfig:
|
||||
"""Convert ScalingConfig to protobuf format."""
|
||||
proto_scaling_config = ProtoTrainRun.ScalingConfig(
|
||||
use_gpu=scaling_config.use_gpu,
|
||||
placement_strategy=scaling_config.placement_strategy,
|
||||
use_tpu=scaling_config.use_tpu,
|
||||
)
|
||||
|
||||
if isinstance(scaling_config.num_workers, tuple):
|
||||
proto_scaling_config.num_workers_range.CopyFrom(
|
||||
ProtoTrainRun.ScalingConfig.IntRange(
|
||||
min=scaling_config.num_workers[0],
|
||||
max=scaling_config.num_workers[1],
|
||||
)
|
||||
)
|
||||
else:
|
||||
proto_scaling_config.num_workers_fixed = scaling_config.num_workers
|
||||
|
||||
if scaling_config.resources_per_worker is not None:
|
||||
proto_scaling_config.resources_per_worker.values.update(
|
||||
scaling_config.resources_per_worker
|
||||
)
|
||||
|
||||
if scaling_config.accelerator_type is not None:
|
||||
proto_scaling_config.accelerator_type = scaling_config.accelerator_type
|
||||
if scaling_config.topology is not None:
|
||||
proto_scaling_config.topology = scaling_config.topology
|
||||
|
||||
if scaling_config.bundle_label_selector is not None:
|
||||
selectors = scaling_config.bundle_label_selector
|
||||
if isinstance(selectors, dict):
|
||||
proto_scaling_config.label_selector_single.values.update(selectors)
|
||||
else:
|
||||
proto_scaling_config.label_selector_list.values.extend(
|
||||
[ProtoTrainRun.ScalingConfig.StringMap(values=s) for s in selectors]
|
||||
)
|
||||
|
||||
return proto_scaling_config
|
||||
|
||||
|
||||
def _to_proto_execution_options(
|
||||
execution_options: ExecutionOptions,
|
||||
) -> ProtoTrainRun.ExecutionOptions:
|
||||
"""Convert a single ExecutionOptions schema model to protobuf."""
|
||||
return ProtoTrainRun.ExecutionOptions(
|
||||
resource_limits=_dict_to_struct(execution_options.resource_limits),
|
||||
exclude_resources=_dict_to_struct(execution_options.exclude_resources),
|
||||
preserve_order=execution_options.preserve_order,
|
||||
actor_locality_enabled=execution_options.actor_locality_enabled,
|
||||
verbose_progress=execution_options.verbose_progress,
|
||||
)
|
||||
|
||||
|
||||
def to_proto_data_config(data_config: DataConfig) -> ProtoTrainRun.DataConfig:
|
||||
"""Convert DataConfig to protobuf format."""
|
||||
data_execution_options = data_config.data_execution_options
|
||||
proto_data_config = ProtoTrainRun.DataConfig(
|
||||
enable_shard_locality=data_config.enable_shard_locality,
|
||||
data_execution_options=ProtoTrainRun.DataExecutionOptions(
|
||||
default=_to_proto_execution_options(data_execution_options.default),
|
||||
per_dataset_execution_options={
|
||||
name: _to_proto_execution_options(opts)
|
||||
for name, opts in data_execution_options.per_dataset_execution_options.items()
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
if data_config.datasets_to_split == "all":
|
||||
proto_data_config.all.SetInParent()
|
||||
else:
|
||||
proto_data_config.datasets.values.extend(data_config.datasets_to_split)
|
||||
|
||||
return proto_data_config
|
||||
|
||||
|
||||
def _to_proto_failure_config(run_config: RunConfig) -> ProtoTrainRun.FailureConfig:
|
||||
"""Convert RunConfig.failure_config to protobuf format."""
|
||||
return ProtoTrainRun.FailureConfig(
|
||||
max_failures=run_config.failure_config.max_failures,
|
||||
controller_failure_limit=run_config.failure_config.controller_failure_limit,
|
||||
)
|
||||
|
||||
|
||||
def _to_proto_checkpoint_config(
|
||||
run_config: RunConfig,
|
||||
) -> ProtoTrainRun.CheckpointConfig:
|
||||
"""Convert RunConfig.checkpoint_config to protobuf format."""
|
||||
checkpoint_score_order = ProtoTrainRun.CheckpointConfig.CheckpointScoreOrder.Value(
|
||||
run_config.checkpoint_config.checkpoint_score_order.upper()
|
||||
)
|
||||
|
||||
proto_checkpoint_config = ProtoTrainRun.CheckpointConfig(
|
||||
checkpoint_score_order=checkpoint_score_order
|
||||
)
|
||||
if run_config.checkpoint_config.num_to_keep is not None:
|
||||
proto_checkpoint_config.num_to_keep = run_config.checkpoint_config.num_to_keep
|
||||
if run_config.checkpoint_config.checkpoint_score_attribute is not None:
|
||||
proto_checkpoint_config.checkpoint_score_attribute = (
|
||||
run_config.checkpoint_config.checkpoint_score_attribute
|
||||
)
|
||||
return proto_checkpoint_config
|
||||
|
||||
|
||||
def to_proto_run_config(run_config: RunConfig) -> ProtoTrainRun.RunConfig:
|
||||
"""Convert RunConfig to protobuf format."""
|
||||
proto_run_config = ProtoTrainRun.RunConfig(
|
||||
name=run_config.name,
|
||||
failure_config=_to_proto_failure_config(run_config),
|
||||
worker_runtime_env=_dict_to_struct(run_config.worker_runtime_env),
|
||||
checkpoint_config=_to_proto_checkpoint_config(run_config),
|
||||
storage_path=run_config.storage_path,
|
||||
)
|
||||
|
||||
if run_config.storage_filesystem is not None:
|
||||
proto_run_config.storage_filesystem = run_config.storage_filesystem
|
||||
|
||||
return proto_run_config
|
||||
|
||||
|
||||
def _to_proto_run_settings(run_settings: RunSettings) -> ProtoTrainRun.RunSettings:
|
||||
"""Convert RunSettings to protobuf format."""
|
||||
|
||||
proto_run_settings = ProtoTrainRun.RunSettings(
|
||||
backend_config=to_proto_backend_config(run_settings.backend_config),
|
||||
scaling_config=to_proto_scaling_config(run_settings.scaling_config),
|
||||
datasets=run_settings.datasets,
|
||||
data_config=to_proto_data_config(run_settings.data_config),
|
||||
run_config=to_proto_run_config(run_settings.run_config),
|
||||
)
|
||||
|
||||
if run_settings.train_loop_config is not None:
|
||||
proto_run_settings.train_loop_config.CopyFrom(
|
||||
_dict_to_struct(run_settings.train_loop_config)
|
||||
)
|
||||
|
||||
return proto_run_settings
|
||||
|
||||
|
||||
def train_run_to_proto(run: TrainRun) -> ProtoTrainRun:
|
||||
"""Convert TrainRun to protobuf format."""
|
||||
|
||||
proto_train_run_panels = [_to_proto_dashboard_panel(p) for p in TRAIN_RUN_PANELS]
|
||||
proto_train_worker_panels = [
|
||||
_to_proto_dashboard_panel(p) for p in TRAIN_WORKER_PANELS
|
||||
]
|
||||
|
||||
proto_train_run = ProtoTrainRun(
|
||||
schema_version=TRAIN_SCHEMA_VERSION,
|
||||
ray_train_version=RAY_TRAIN_VERSION,
|
||||
id=run.id,
|
||||
name=run.name,
|
||||
job_id=bytes.fromhex(run.job_id),
|
||||
controller_actor_id=bytes.fromhex(run.controller_actor_id),
|
||||
status=_RUN_STATUS_MAP[run.status],
|
||||
status_detail=run.status_detail,
|
||||
start_time_ns=run.start_time_ns,
|
||||
end_time_ns=run.end_time_ns,
|
||||
controller_log_file_path=run.controller_log_file_path,
|
||||
train_run_panels=proto_train_run_panels,
|
||||
train_worker_panels=proto_train_worker_panels,
|
||||
framework_versions=run.framework_versions,
|
||||
run_settings=_to_proto_run_settings(run.run_settings),
|
||||
)
|
||||
|
||||
return proto_train_run
|
||||
@@ -0,0 +1,532 @@
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
from pydantic import field_validator
|
||||
|
||||
from ray._common.pydantic_compat import BaseModel, Field
|
||||
from ray.dashboard.modules.job.pydantic_models import JobDetails
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
MAX_ERROR_STACK_TRACE_LENGTH = 50000
|
||||
|
||||
|
||||
def _to_json_serializable_value(value: Any, *, max_depth: int = 3) -> Any:
|
||||
"""Recursively coerce a value into a human-readable, JSON serializable representation.
|
||||
|
||||
If ``value`` is a list or dict, this function walks through it and replaces non-JSON
|
||||
serializable fields (e.g. custom objects, modules, tensors, callables, etc.) with a
|
||||
human-readable string representation.
|
||||
|
||||
Args:
|
||||
value: Any Python value. Primitives pass through; collections recurse;
|
||||
other types are stringified.
|
||||
max_depth: Truncates dicts nested beyond ``max_depth`` to ``"..."``.
|
||||
Lists do not consume depth.
|
||||
|
||||
Returns:
|
||||
The JSON serializable representation of the value.
|
||||
"""
|
||||
if max_depth <= 0:
|
||||
raise ValueError("max_depth must be greater than 0")
|
||||
|
||||
def _safe_str(v):
|
||||
try:
|
||||
return str(v)
|
||||
except Exception:
|
||||
return type(v).__name__
|
||||
|
||||
def _walk(value, depth):
|
||||
if value is None or isinstance(value, (bool, int, str)):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return str(value) if not math.isfinite(value) else value
|
||||
if isinstance(value, Mapping):
|
||||
if depth <= 0:
|
||||
return "..."
|
||||
try:
|
||||
items = list(value.items())
|
||||
except Exception:
|
||||
# Custom Mapping subclass with a broken `.items()`.
|
||||
return type(value).__name__
|
||||
return {_safe_str(k): _walk(v, depth - 1) for k, v in items}
|
||||
|
||||
# Tuples, sets, and frozensets all become lists in JSON.
|
||||
if isinstance(value, (list, tuple, set, frozenset)):
|
||||
return [_walk(v, depth) for v in value]
|
||||
|
||||
cls = type(value)
|
||||
# Use class name if no custom string representation is defined.
|
||||
if cls.__str__ is object.__str__ and cls.__repr__ is object.__repr__:
|
||||
return cls.__name__
|
||||
return _safe_str(value)
|
||||
|
||||
return _walk(value, max_depth)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RunStatus(str, Enum):
|
||||
"""Enumeration of the possible statuses for a Train run."""
|
||||
|
||||
# ====== Active States ======
|
||||
# The Train run is currently in the process of initializing.
|
||||
INITIALIZING = "INITIALIZING"
|
||||
# The Train run is waiting to be scheduled.
|
||||
SCHEDULING = "SCHEDULING"
|
||||
# The Train run is currently in progress.
|
||||
RUNNING = "RUNNING"
|
||||
# The Train run is recovering from a failure or restart.
|
||||
RESTARTING = "RESTARTING"
|
||||
# The Train run is resizing.
|
||||
RESIZING = "RESIZING"
|
||||
|
||||
# ===== Terminal States ======
|
||||
# The Train run completed successfully.
|
||||
FINISHED = "FINISHED"
|
||||
# The Train run failed due to an error in the training workers.
|
||||
ERRORED = "ERRORED"
|
||||
# The Train run was terminated due to system or controller errors.
|
||||
ABORTED = "ABORTED"
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
return self in [RunStatus.FINISHED, RunStatus.ERRORED, RunStatus.ABORTED]
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RunAttemptStatus(str, Enum):
|
||||
"""Enumeration of the possible statuses for a Train run attempt."""
|
||||
|
||||
# ====== Active States ======
|
||||
# The run attempt is waiting to be scheduled.
|
||||
PENDING = "PENDING"
|
||||
# The run attempt is currently in progress.
|
||||
RUNNING = "RUNNING"
|
||||
|
||||
# ===== Terminal States =====
|
||||
# The run attempt completed successfully.
|
||||
FINISHED = "FINISHED"
|
||||
# The run attempt failed due to an error in the training workers.
|
||||
ERRORED = "ERRORED"
|
||||
# The run attempt was terminated due to system or controller errors.
|
||||
ABORTED = "ABORTED"
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
return self in [
|
||||
RunAttemptStatus.FINISHED,
|
||||
RunAttemptStatus.ERRORED,
|
||||
RunAttemptStatus.ABORTED,
|
||||
]
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ActorStatus(str, Enum):
|
||||
"""Enumeration of the statuses for a Train worker actor."""
|
||||
|
||||
# The actor is currently active.
|
||||
ALIVE = "ALIVE"
|
||||
# The actor is no longer active.
|
||||
DEAD = "DEAD"
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainResources(BaseModel):
|
||||
"""Resources allocated for a Train worker or run."""
|
||||
|
||||
resources: Dict[str, float] = Field(
|
||||
description="A dictionary specifying the types and amounts of resources "
|
||||
"allocated (e.g., CPU, GPU)."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainWorker(BaseModel):
|
||||
"""Metadata about a Ray Train worker."""
|
||||
|
||||
world_rank: int = Field(
|
||||
description="The global rank of the worker in the training cluster."
|
||||
)
|
||||
local_rank: int = Field(description="The local rank of the worker on its node.")
|
||||
node_rank: int = Field(description="The rank of the worker's node in the cluster.")
|
||||
actor_id: str = Field(description="The unique ID of the worker's actor.")
|
||||
node_id: str = Field(
|
||||
description="The unique ID of the node where the worker is running."
|
||||
)
|
||||
node_ip: str = Field(
|
||||
description="The IP address of the node where the worker is running."
|
||||
)
|
||||
pid: int = Field(description="The process ID of the worker.")
|
||||
gpu_ids: List[int] = Field(description="A list of GPU IDs allocated to the worker.")
|
||||
status: Optional[ActorStatus] = Field(
|
||||
None, description="The current status of the worker actor."
|
||||
)
|
||||
resources: TrainResources = Field(
|
||||
description="The resources allocated to this Train worker."
|
||||
)
|
||||
log_file_path: Optional[str] = Field(
|
||||
None, description="The path to the log file for the Train worker."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class MemoryInfo(BaseModel):
|
||||
"""Memory usage information for a process."""
|
||||
|
||||
rss: int = Field(description="The resident set size (RSS) memory usage in bytes.")
|
||||
vms: int = Field(description="The virtual memory size (VMS) usage in bytes.")
|
||||
pfaults: Optional[int] = Field(None, description="The number of page faults.")
|
||||
pageins: Optional[int] = Field(None, description="The number of page-ins.")
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ProcessStats(BaseModel):
|
||||
"""CPU and memory statistics for a process."""
|
||||
|
||||
cpuPercent: float = Field(description="The percentage of CPU usage.")
|
||||
mem: Optional[List[int]] = Field(
|
||||
None,
|
||||
description="Memory statistics, including total memory, free memory, "
|
||||
"and memory usage ratio.",
|
||||
)
|
||||
memoryInfo: MemoryInfo = Field(description="Detailed memory usage information.")
|
||||
|
||||
|
||||
class ProcessGPUUsage(BaseModel):
|
||||
"""GPU usage statistics for a process."""
|
||||
|
||||
pid: int = Field(description="The process ID.")
|
||||
gpuMemoryUsage: int = Field(description="The GPU memory usage in bytes.")
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class GPUStats(BaseModel):
|
||||
"""Statistics for a GPU."""
|
||||
|
||||
uuid: str = Field(description="The unique identifier of the GPU.")
|
||||
index: int = Field(description="The index of the GPU.")
|
||||
name: str = Field(description="The name of the GPU.")
|
||||
utilizationGpu: Optional[float] = Field(
|
||||
None, description="The percentage utilization of the GPU."
|
||||
)
|
||||
memoryUsed: float = Field(description="The amount of GPU memory used in bytes.")
|
||||
memoryTotal: float = Field(description="The total amount of GPU memory in bytes.")
|
||||
processInfo: ProcessGPUUsage = Field(
|
||||
description="GPU usage statistics for the associated process."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DecoratedTrainWorker(TrainWorker):
|
||||
"""Detailed metadata for a Ray Train worker, including process and GPU stats."""
|
||||
|
||||
processStats: Optional[ProcessStats] = Field(
|
||||
None, description="CPU and memory statistics for the worker process."
|
||||
)
|
||||
gpus: List[GPUStats] = Field(
|
||||
default_factory=list,
|
||||
description="A list of GPUs used by the worker process,"
|
||||
" with detailed statistics.",
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainRunAttempt(BaseModel):
|
||||
"""Metadata for an individual attempt to execute a Train run."""
|
||||
|
||||
run_id: str = Field(description="Unique identifier for the parent Train run.")
|
||||
attempt_id: str = Field(
|
||||
description="Unique identifier for this specific Train run attempt."
|
||||
)
|
||||
status: RunAttemptStatus = Field(
|
||||
description="The current execution status of the Train run attempt."
|
||||
)
|
||||
status_detail: Optional[str] = Field(
|
||||
None,
|
||||
description="Additional details about the status,"
|
||||
" including error messages if applicable.",
|
||||
)
|
||||
start_time_ns: int = Field(
|
||||
description="The UNIX timestamp (in nanoseconds)"
|
||||
" when the Train run attempt started."
|
||||
)
|
||||
end_time_ns: Optional[int] = Field(
|
||||
None,
|
||||
description="The UNIX timestamp (in nanoseconds)"
|
||||
" when the Train run attempt ended. "
|
||||
"If null, the attempt is still ongoing.",
|
||||
)
|
||||
resources: List[TrainResources] = Field(
|
||||
description="The resources (e.g., CPU, GPU) allocated to the Train run attempt."
|
||||
)
|
||||
workers: List[TrainWorker] = Field(
|
||||
description="List of Train workers participating in this attempt, "
|
||||
"sorted by global ranks."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DecoratedTrainRunAttempt(TrainRunAttempt):
|
||||
"""Detailed metadata for a Train run attempt, including decorated worker data."""
|
||||
|
||||
workers: List[DecoratedTrainWorker] = Field(
|
||||
description="A list of Train workers with detailed statistics, "
|
||||
"sorted by global ranks."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ExecutionOptions(BaseModel):
|
||||
"""ExecutionOptions for a single Ray Data ingest pipeline."""
|
||||
|
||||
resource_limits: Dict[str, Any] = Field(
|
||||
description="The resource limits applied to the Ray Data execution plan."
|
||||
)
|
||||
exclude_resources: Dict[str, Any] = Field(
|
||||
description="The resources excluded from the Ray Data execution plan "
|
||||
"(e.g. resources reserved by Ray Train workers)."
|
||||
)
|
||||
|
||||
@field_validator("resource_limits", "exclude_resources", mode="before")
|
||||
@classmethod
|
||||
def _sanitize_dict(cls, v):
|
||||
return _to_json_serializable_value(v)
|
||||
|
||||
preserve_order: bool = Field(
|
||||
description="Whether to preserve the order of outputs across operators."
|
||||
)
|
||||
actor_locality_enabled: bool = Field(
|
||||
description="Whether actor-based locality optimizations are enabled."
|
||||
)
|
||||
verbose_progress: bool = Field(
|
||||
description="Whether verbose progress reporting is enabled."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DataExecutionOptions(BaseModel):
|
||||
"""ExecutionOptions for a Ray Train run, split into defaults and per-dataset overrides."""
|
||||
|
||||
default: ExecutionOptions = Field(
|
||||
description="Execution options applied to any dataset without a per-dataset override."
|
||||
)
|
||||
per_dataset_execution_options: Dict[str, ExecutionOptions] = Field(
|
||||
default_factory=dict,
|
||||
description="Per-dataset execution option overrides, keyed by dataset name.",
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DataConfig(BaseModel):
|
||||
"""Configuration for dataset splitting and execution options within Ray Train."""
|
||||
|
||||
datasets_to_split: Union[Literal["all"], List[str]] = Field(
|
||||
description="Which datasets to split; either 'all' or a list of dataset names."
|
||||
)
|
||||
execution_options: Optional[Dict] = Field(
|
||||
default=None,
|
||||
deprecated="DEPRECATED: Use data_execution_options instead.",
|
||||
)
|
||||
data_execution_options: DataExecutionOptions = Field(
|
||||
description="Data execution options"
|
||||
)
|
||||
enable_shard_locality: bool = Field(
|
||||
description="Whether to enable shard locality optimization."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ScalingConfig(BaseModel):
|
||||
"""Scaling config for a Train run."""
|
||||
|
||||
num_workers: Union[int, Tuple[int, int]] = Field(
|
||||
description="The number of workers for the Train run."
|
||||
)
|
||||
use_gpu: bool = Field(description="Whether to use GPUs for the Train run.")
|
||||
resources_per_worker: Optional[Dict[str, float]] = Field(
|
||||
None, description="The resources per worker for a Train run."
|
||||
)
|
||||
placement_strategy: str = Field(
|
||||
description="The placement strategy for the Train run."
|
||||
)
|
||||
accelerator_type: Optional[str] = Field(
|
||||
None, description="The accelerator type for the Train run."
|
||||
)
|
||||
use_tpu: bool = Field(description="Whether to use TPUs for the Train run.")
|
||||
topology: Optional[str] = Field(None, description="The topology for the Train run.")
|
||||
bundle_label_selector: Optional[
|
||||
Union[Dict[str, str], List[Dict[str, str]]]
|
||||
] = Field(None, description="The bundle label selector for the Train run.")
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class FailureConfig(BaseModel):
|
||||
"""Failure config for a Train run."""
|
||||
|
||||
max_failures: int = Field(
|
||||
description="The maximum number of failures for a Train run."
|
||||
)
|
||||
controller_failure_limit: int = Field(
|
||||
description="The maximum number of controller failures to tolerate."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class CheckpointConfig(BaseModel):
|
||||
"""Checkpoint config for a Train run."""
|
||||
|
||||
num_to_keep: Optional[int] = Field(
|
||||
None,
|
||||
description="The number of most recent checkpoints to keep. Older checkpoints may be deleted.",
|
||||
)
|
||||
checkpoint_score_attribute: Optional[str] = Field(
|
||||
None,
|
||||
description="Attribute used to score and rank checkpoints; can be a metric key or attribute.",
|
||||
)
|
||||
checkpoint_score_order: Literal["max", "min"] = Field(
|
||||
description="Order to rank checkpoint scores, 'max' for higher-is-better, 'min' for lower-is-better.",
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RunConfig(BaseModel):
|
||||
"""Run configuration parameters for a Train run, encompassing failure,
|
||||
runtime environment, checkpoint settings, and storage path."""
|
||||
|
||||
name: str = Field(description="The name of the Train run.")
|
||||
failure_config: FailureConfig = Field(
|
||||
description="The failure config for a Train run."
|
||||
)
|
||||
worker_runtime_env: Dict[str, Any] = Field(
|
||||
description="The worker runtime env for a Train run."
|
||||
)
|
||||
|
||||
@field_validator("worker_runtime_env", mode="before")
|
||||
@classmethod
|
||||
def _sanitize_worker_runtime_env(cls, v):
|
||||
return _to_json_serializable_value(v)
|
||||
|
||||
checkpoint_config: CheckpointConfig = Field(
|
||||
description="The checkpoint config for a Train run."
|
||||
)
|
||||
storage_path: str = Field(description="The storage path for a Train run.")
|
||||
storage_filesystem: Optional[str] = Field(
|
||||
None, description="The storage filesystem for a Train run."
|
||||
)
|
||||
|
||||
@field_validator("storage_filesystem", mode="before")
|
||||
@classmethod
|
||||
def _sanitize_storage_filesystem(cls, v):
|
||||
return _to_json_serializable_value(v)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class BackendConfig(BaseModel):
|
||||
"""Backend config for a Train run."""
|
||||
|
||||
framework: Optional[TrainingFramework] = Field(
|
||||
None, description="The training framework for this backend config."
|
||||
)
|
||||
config: Dict[str, Any] = Field(
|
||||
description="Training framework-specific configuration fields."
|
||||
)
|
||||
|
||||
@field_validator("config", mode="before")
|
||||
@classmethod
|
||||
def _sanitize_config(cls, v):
|
||||
return _to_json_serializable_value(v)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RunSettings(BaseModel):
|
||||
"""Settings for a Train run, primarily consisting of configs set before a train run starts.
|
||||
|
||||
This includes the train loop config, backend config, scaling config, dataset configs,
|
||||
and runtime configuration.
|
||||
"""
|
||||
|
||||
train_loop_config: Optional[Dict] = Field(
|
||||
None, description="The user defined train loop config for a Train run."
|
||||
)
|
||||
|
||||
@field_validator("train_loop_config", mode="before")
|
||||
@classmethod
|
||||
def _sanitize_train_loop_config(cls, v):
|
||||
return _to_json_serializable_value(v)
|
||||
|
||||
backend_config: BackendConfig = Field(
|
||||
description="The backend config for a Train run. Can vary with the framework (e.g. TorchConfig)"
|
||||
)
|
||||
scaling_config: ScalingConfig = Field(
|
||||
description="The scaling config for this Train run."
|
||||
)
|
||||
datasets: List[str] = Field(
|
||||
description="A list of dataset names for a Train run.",
|
||||
)
|
||||
data_config: DataConfig = Field(
|
||||
description="The data config for a Train run.",
|
||||
)
|
||||
run_config: RunConfig = Field(
|
||||
description="Run configuration for this Train run, including failure, runtime environment, checkpoint settings, and storage path."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainRun(BaseModel):
|
||||
"""Metadata for a Ray Train run, including its details and status."""
|
||||
|
||||
id: str = Field(description="Unique identifier for the Train run.")
|
||||
name: str = Field(description="Human-readable name assigned to the Train run.")
|
||||
job_id: str = Field(description="The Ray Job ID associated with this Train run.")
|
||||
controller_actor_id: str = Field(
|
||||
description="Unique ID of the actor managing the Train run."
|
||||
)
|
||||
status: RunStatus = Field(
|
||||
description="The current execution status of the Train run."
|
||||
)
|
||||
status_detail: Optional[str] = Field(
|
||||
None,
|
||||
description="Additional details about the current status, "
|
||||
"including error messages if applicable.",
|
||||
)
|
||||
start_time_ns: int = Field(
|
||||
description="The UNIX timestamp (in nanoseconds) when the Train run started."
|
||||
)
|
||||
end_time_ns: Optional[int] = Field(
|
||||
None,
|
||||
description="The UNIX timestamp (in nanoseconds) when the Train run ended. "
|
||||
"If null, the run is still in progress.",
|
||||
)
|
||||
controller_log_file_path: Optional[str] = Field(
|
||||
None, description="The path to the log file for the Train run controller."
|
||||
)
|
||||
framework_versions: Dict[str, str] = Field(
|
||||
description="The relevant framework versions for this Train run,"
|
||||
"including the Ray version and training framework version."
|
||||
)
|
||||
run_settings: RunSettings = Field(
|
||||
description="The run settings for this Train run, including train loop config, "
|
||||
"backend config, scaling config, dataset details, and runtime configuration."
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DecoratedTrainRun(TrainRun):
|
||||
"""Detailed metadata for a Ray Train run, including attempts and job details."""
|
||||
|
||||
attempts: List[DecoratedTrainRunAttempt] = Field(
|
||||
description="A list of attempts made to execute the Train run."
|
||||
)
|
||||
job_details: Optional[JobDetails] = Field(
|
||||
None,
|
||||
description="Detailed information about the job that initiated this Train run.",
|
||||
)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class TrainRunsResponse(BaseModel):
|
||||
"""Response containing a list of decorated Train runs."""
|
||||
|
||||
train_runs: List[DecoratedTrainRun] = Field(
|
||||
description="A list of Train runs with detailed metadata."
|
||||
)
|
||||
@@ -0,0 +1,301 @@
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict, defaultdict
|
||||
from typing import Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray._private import ray_constants
|
||||
from ray._private.event.export_event_logger import (
|
||||
EventLogType,
|
||||
check_export_api_enabled,
|
||||
get_export_event_logger,
|
||||
)
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train.v2._internal.constants import (
|
||||
CONTROLLERS_TO_POLL_PER_ITERATION,
|
||||
DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
|
||||
DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
|
||||
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
|
||||
GET_ACTOR_TIMEOUT_S,
|
||||
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.state.schema import (
|
||||
TrainRun,
|
||||
TrainRunAttempt,
|
||||
)
|
||||
from ray.train.v2._internal.state.util import (
|
||||
is_actor_alive,
|
||||
update_train_run_aborted,
|
||||
update_train_run_attempt_aborted,
|
||||
)
|
||||
from ray.train.v2._internal.util import time_monotonic
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TrainStateActor:
|
||||
def __init__(
|
||||
self,
|
||||
# TODO: group into single config if we need to do similar polling elsewhere
|
||||
enable_state_actor_reconciliation: bool = False,
|
||||
reconciliation_interval_s: float = 30,
|
||||
get_actor_timeout_s: int = GET_ACTOR_TIMEOUT_S,
|
||||
controllers_to_poll_per_iteration: int = CONTROLLERS_TO_POLL_PER_ITERATION,
|
||||
):
|
||||
# NOTE: All runs and attempts are stored in memory.
|
||||
# This may be a memory issue for large runs.
|
||||
# TODO: consider cleaning up runs over time.
|
||||
self._runs: Dict[str, TrainRun] = OrderedDict()
|
||||
# {run_id: {attempt_id: TrainRunAttempt}}
|
||||
self._run_attempts: Dict[str, OrderedDict[str, TrainRunAttempt]] = defaultdict(
|
||||
OrderedDict
|
||||
)
|
||||
(
|
||||
self._export_logger,
|
||||
self._is_train_run_export_api_enabled,
|
||||
self._is_train_run_attempt_export_api_enabled,
|
||||
) = self._init_export_logger()
|
||||
|
||||
# TODO: consider row level locking if loop takes too long.
|
||||
self._runs_lock = threading.RLock()
|
||||
self._run_attempts_lock = threading.RLock()
|
||||
|
||||
# Set env vars related to reconciling train run/attempt state.
|
||||
if enable_state_actor_reconciliation:
|
||||
self._reconciliation_interval_s = reconciliation_interval_s
|
||||
self._controllers_to_poll_per_iteration = controllers_to_poll_per_iteration
|
||||
self._get_actor_timeout_s = get_actor_timeout_s
|
||||
self._start_run_state_reconciliation_thread()
|
||||
|
||||
def _abort_live_runs_with_dead_controllers(
|
||||
self, last_poll_run_id: Optional[str]
|
||||
) -> str:
|
||||
aborted_run_ids = []
|
||||
with self._runs_lock:
|
||||
runs = list(self._runs.values())
|
||||
|
||||
# Start iterating from poll index.
|
||||
starting_poll_index = 0
|
||||
if last_poll_run_id is not None:
|
||||
for poll_index, run in enumerate(runs):
|
||||
if run.id == last_poll_run_id:
|
||||
starting_poll_index = (poll_index + 1) % len(runs)
|
||||
break
|
||||
|
||||
# Abort runs.
|
||||
num_polled_runs = 0
|
||||
poll_index = starting_poll_index
|
||||
while (
|
||||
poll_index < starting_poll_index + len(runs)
|
||||
and num_polled_runs < self._controllers_to_poll_per_iteration
|
||||
):
|
||||
run = runs[poll_index % len(runs)]
|
||||
poll_index += 1
|
||||
last_poll_run_id = run.id
|
||||
if run.status.is_terminal():
|
||||
continue
|
||||
try:
|
||||
if not is_actor_alive(
|
||||
run.controller_actor_id, self._get_actor_timeout_s
|
||||
):
|
||||
update_train_run_aborted(run, False)
|
||||
self.create_or_update_train_run(run)
|
||||
aborted_run_ids.append(run.id)
|
||||
except ray.util.state.exception.RayStateApiException:
|
||||
logger.exception(
|
||||
"State API unavailable when checking if actor is alive. "
|
||||
"Will check again on next poll."
|
||||
)
|
||||
num_polled_runs += 1
|
||||
|
||||
# Abort run attempts.
|
||||
with self._run_attempts_lock:
|
||||
for run_id in aborted_run_ids:
|
||||
latest_run_attempt = self._get_latest_run_attempt(run_id)
|
||||
if latest_run_attempt and not latest_run_attempt.status.is_terminal():
|
||||
update_train_run_attempt_aborted(latest_run_attempt, False)
|
||||
self.create_or_update_train_run_attempt(latest_run_attempt)
|
||||
|
||||
return last_poll_run_id
|
||||
|
||||
def _start_run_state_reconciliation_thread(self) -> None:
|
||||
def _reconciliation_loop():
|
||||
last_poll_run_id = None
|
||||
latest_poll_time = float("-inf")
|
||||
while True:
|
||||
# Wait for the poll interval to elapse.
|
||||
time_since_last_poll = time_monotonic() - latest_poll_time
|
||||
if time_since_last_poll < self._reconciliation_interval_s:
|
||||
remaining_time = (
|
||||
self._reconciliation_interval_s - time_since_last_poll
|
||||
)
|
||||
time.sleep(remaining_time)
|
||||
|
||||
last_poll_run_id = self._abort_live_runs_with_dead_controllers(
|
||||
last_poll_run_id
|
||||
)
|
||||
latest_poll_time = time_monotonic()
|
||||
|
||||
threading.Thread(target=_reconciliation_loop, daemon=True).start()
|
||||
|
||||
def _get_latest_run_attempt(self, run_id: str) -> Optional[TrainRunAttempt]:
|
||||
with self._run_attempts_lock:
|
||||
# NOTE: run_attempts is OrderedDict from attempt_id to TrainRunAttempt.
|
||||
run_attempts = self._run_attempts.get(run_id, {})
|
||||
if not run_attempts:
|
||||
return None
|
||||
return next(reversed(run_attempts.values()))
|
||||
|
||||
def create_or_update_train_run(self, run: TrainRun) -> None:
|
||||
with self._runs_lock:
|
||||
self._runs[run.id] = run
|
||||
run_copy = copy.deepcopy(run)
|
||||
self._maybe_export_train_run(run_copy)
|
||||
|
||||
def create_or_update_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
|
||||
with self._run_attempts_lock:
|
||||
self._run_attempts[run_attempt.run_id][run_attempt.attempt_id] = run_attempt
|
||||
run_attempt_copy = copy.deepcopy(run_attempt)
|
||||
self._maybe_export_train_run_attempt(run_attempt_copy)
|
||||
|
||||
def get_train_runs(self) -> Dict[str, TrainRun]:
|
||||
with self._runs_lock:
|
||||
return self._runs
|
||||
|
||||
def get_train_run_attempts(self) -> Dict[str, Dict[str, TrainRunAttempt]]:
|
||||
with self._run_attempts_lock:
|
||||
return self._run_attempts
|
||||
|
||||
# ============================
|
||||
# Export API
|
||||
# ============================
|
||||
|
||||
def is_export_api_enabled(self) -> bool:
|
||||
return self._export_logger is not None
|
||||
|
||||
def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]:
|
||||
"""Initialize the export logger and check if the export API is enabled.
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- The export logger (or None if export API is not enabled).
|
||||
- A boolean indicating if the export API is enabled for train runs.
|
||||
- A boolean indicating if the export API is enabled for train run attempts.
|
||||
"""
|
||||
# Proto schemas should be imported within the scope of TrainStateActor to
|
||||
# prevent serialization errors.
|
||||
from ray.core.generated.export_event_pb2 import ExportEvent
|
||||
|
||||
is_train_run_export_api_enabled = check_export_api_enabled(
|
||||
ExportEvent.SourceType.EXPORT_TRAIN_RUN
|
||||
)
|
||||
is_train_run_attempt_export_api_enabled = check_export_api_enabled(
|
||||
ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
|
||||
)
|
||||
export_api_enabled = (
|
||||
is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled
|
||||
)
|
||||
|
||||
if not export_api_enabled:
|
||||
return None, False, False
|
||||
|
||||
log_directory = os.path.join(
|
||||
ray._private.worker._global_node.get_session_dir_path(), "logs"
|
||||
)
|
||||
logger = None
|
||||
try:
|
||||
logger = get_export_event_logger(
|
||||
EventLogType.TRAIN_STATE,
|
||||
log_directory,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Unable to initialize the export event logger, so no Train export "
|
||||
"events will be written."
|
||||
)
|
||||
|
||||
if logger is None:
|
||||
return None, False, False
|
||||
|
||||
return (
|
||||
logger,
|
||||
is_train_run_export_api_enabled,
|
||||
is_train_run_attempt_export_api_enabled,
|
||||
)
|
||||
|
||||
def _maybe_export_train_run(self, run: TrainRun) -> None:
|
||||
if not self._is_train_run_export_api_enabled:
|
||||
return
|
||||
|
||||
from ray.train.v2._internal.state.export import train_run_to_proto
|
||||
|
||||
run_proto = train_run_to_proto(run)
|
||||
self._export_logger.send_event(run_proto)
|
||||
|
||||
def _maybe_export_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
|
||||
if not self._is_train_run_attempt_export_api_enabled:
|
||||
return
|
||||
|
||||
from ray.train.v2._internal.state.export import train_run_attempt_to_proto
|
||||
|
||||
run_attempt_proto = train_run_attempt_to_proto(run_attempt)
|
||||
self._export_logger.send_event(run_attempt_proto)
|
||||
|
||||
|
||||
TRAIN_STATE_ACTOR_NAME = "train_v2_state_actor"
|
||||
TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor"
|
||||
|
||||
_state_actor_lock: threading.RLock = threading.RLock()
|
||||
|
||||
|
||||
def get_or_create_state_actor() -> ActorHandle:
|
||||
"""Get or create the Ray Train state actor singleton.
|
||||
|
||||
This is a long-living, detached actor living on the head node
|
||||
that gets initialized when the first Train run happens on the
|
||||
Ray cluster.
|
||||
"""
|
||||
with _state_actor_lock:
|
||||
state_actor = (
|
||||
ray.remote(TrainStateActor)
|
||||
.options(
|
||||
num_cpus=0,
|
||||
name=TRAIN_STATE_ACTOR_NAME,
|
||||
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
|
||||
get_if_exists=True,
|
||||
lifetime="detached",
|
||||
resources={"node:__internal_head__": 0.001},
|
||||
# Escape from the parent's placement group
|
||||
scheduling_strategy="DEFAULT",
|
||||
max_restarts=-1,
|
||||
max_task_retries=-1,
|
||||
)
|
||||
.remote(
|
||||
enable_state_actor_reconciliation=ray_constants.env_bool(
|
||||
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
|
||||
DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
|
||||
),
|
||||
reconciliation_interval_s=float(
|
||||
os.getenv(
|
||||
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
|
||||
DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return state_actor
|
||||
|
||||
|
||||
def get_state_actor() -> Optional[ActorHandle]:
|
||||
"""Get the `TrainStateActor` if exists, otherwise return None."""
|
||||
try:
|
||||
return ray.get_actor(
|
||||
name=TRAIN_STATE_ACTOR_NAME,
|
||||
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
|
||||
)
|
||||
except ValueError:
|
||||
return None
|
||||
@@ -0,0 +1,326 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.actor import ActorHandle
|
||||
from ray.train import BackendConfig
|
||||
from ray.train._internal.data_config import DataConfig
|
||||
from ray.train.v2._internal.execution.context import DistributedContext
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
ResizeDecision,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import ActorMetadata, Worker
|
||||
from ray.train.v2._internal.state.schema import (
|
||||
ActorStatus,
|
||||
BackendConfig as BackendConfigSchema,
|
||||
CheckpointConfig as CheckpointConfigSchema,
|
||||
FailureConfig as FailureConfigSchema,
|
||||
RunAttemptStatus,
|
||||
RunConfig as RunConfigSchema,
|
||||
RunSettings,
|
||||
RunStatus,
|
||||
ScalingConfig as ScalingConfigSchema,
|
||||
TrainResources,
|
||||
TrainRun,
|
||||
TrainRunAttempt,
|
||||
TrainWorker,
|
||||
)
|
||||
from ray.train.v2._internal.state.state_actor import get_or_create_state_actor
|
||||
from ray.train.v2._internal.state.util import (
|
||||
construct_data_config,
|
||||
current_time_ns,
|
||||
mark_workers_dead,
|
||||
update_train_run_aborted,
|
||||
update_train_run_attempt_aborted,
|
||||
)
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TrainStateManager:
|
||||
"""Manages the state of a train run and run attempts."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._state_actor = get_or_create_state_actor()
|
||||
# NOTE: All runs and attempts are stored in memory.
|
||||
# This may be a memory issue for large runs.
|
||||
self._runs: Dict[str, TrainRun] = {}
|
||||
# {run_id: {attempt_id: TrainRunAttempt}}
|
||||
self._run_attempts: Dict[str, Dict[str, TrainRunAttempt]] = defaultdict(dict)
|
||||
|
||||
def create_train_run(
|
||||
self,
|
||||
id: str,
|
||||
name: str,
|
||||
job_id: str,
|
||||
controller_actor_id: str,
|
||||
controller_log_file_path: str,
|
||||
run_config: RunConfig,
|
||||
train_loop_config: Optional[Dict],
|
||||
scaling_config: ScalingConfig,
|
||||
backend_config: BackendConfig,
|
||||
datasets: Dict[str, ray.data.Dataset],
|
||||
dataset_config: DataConfig,
|
||||
) -> None:
|
||||
run_config_schema = RunConfigSchema(
|
||||
name=run_config.name,
|
||||
failure_config=FailureConfigSchema(
|
||||
max_failures=run_config.failure_config.max_failures,
|
||||
controller_failure_limit=run_config.failure_config.controller_failure_limit,
|
||||
),
|
||||
worker_runtime_env=run_config.worker_runtime_env,
|
||||
checkpoint_config=CheckpointConfigSchema(
|
||||
num_to_keep=run_config.checkpoint_config.num_to_keep,
|
||||
checkpoint_score_attribute=run_config.checkpoint_config.checkpoint_score_attribute,
|
||||
checkpoint_score_order=run_config.checkpoint_config.checkpoint_score_order,
|
||||
),
|
||||
storage_path=run_config.storage_path,
|
||||
storage_filesystem=run_config.storage_filesystem,
|
||||
)
|
||||
|
||||
scaling_config_schema = ScalingConfigSchema(
|
||||
num_workers=scaling_config.num_workers,
|
||||
use_gpu=scaling_config.use_gpu,
|
||||
resources_per_worker=scaling_config.resources_per_worker,
|
||||
placement_strategy=scaling_config.placement_strategy,
|
||||
accelerator_type=scaling_config.accelerator_type,
|
||||
use_tpu=scaling_config.use_tpu,
|
||||
topology=scaling_config.topology,
|
||||
bundle_label_selector=scaling_config.label_selector,
|
||||
)
|
||||
|
||||
backend_config_schema = BackendConfigSchema(
|
||||
framework=backend_config.framework,
|
||||
config=backend_config.to_dict(),
|
||||
)
|
||||
|
||||
run_settings = RunSettings(
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=backend_config_schema,
|
||||
scaling_config=scaling_config_schema,
|
||||
datasets=list(datasets.keys()),
|
||||
data_config=construct_data_config(dataset_config),
|
||||
run_config=run_config_schema,
|
||||
)
|
||||
|
||||
run = TrainRun(
|
||||
id=id,
|
||||
name=name,
|
||||
job_id=job_id,
|
||||
status=RunStatus.INITIALIZING,
|
||||
status_detail=None,
|
||||
controller_actor_id=controller_actor_id,
|
||||
start_time_ns=current_time_ns(),
|
||||
end_time_ns=None,
|
||||
controller_log_file_path=controller_log_file_path,
|
||||
framework_versions={"ray": ray.__version__},
|
||||
run_settings=run_settings,
|
||||
)
|
||||
self._runs[run.id] = run
|
||||
# Block so the initial run state isn't lost if the controller exits
|
||||
# right after. Without this, the .remote() task could still be in the
|
||||
# caller's outbound queue when the controller dies, leaving the state
|
||||
# actor with no record of the run.
|
||||
self._create_or_update_train_run(run, block=True)
|
||||
|
||||
def update_train_run_scheduling(
|
||||
self,
|
||||
run_id: str,
|
||||
resize_decision: Optional[ResizeDecision] = None,
|
||||
) -> None:
|
||||
if resize_decision is not None:
|
||||
status_detail = _get_scheduling_status_detail(
|
||||
resize_decision.num_workers, resize_decision.resources_per_worker
|
||||
)
|
||||
else:
|
||||
status_detail = None
|
||||
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.SCHEDULING
|
||||
run.status_detail = status_detail
|
||||
self._create_or_update_train_run(run)
|
||||
|
||||
def update_train_run_running(
|
||||
self,
|
||||
run_id: str,
|
||||
) -> None:
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.RUNNING
|
||||
run.status_detail = None
|
||||
self._create_or_update_train_run(run)
|
||||
|
||||
def update_train_run_restarting(
|
||||
self,
|
||||
run_id: str,
|
||||
) -> None:
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.RESTARTING
|
||||
run.status_detail = None
|
||||
self._create_or_update_train_run(run)
|
||||
|
||||
def update_train_run_resizing(
|
||||
self,
|
||||
run_id: str,
|
||||
) -> None:
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.RESIZING
|
||||
run.status_detail = None
|
||||
self._create_or_update_train_run(run)
|
||||
|
||||
def update_train_run_finished(
|
||||
self,
|
||||
run_id: str,
|
||||
):
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.FINISHED
|
||||
run.status_detail = None
|
||||
run.end_time_ns = current_time_ns()
|
||||
# Block on terminal status so the final state isn't lost if the controller exits right after.
|
||||
self._create_or_update_train_run(run, block=True)
|
||||
|
||||
def update_train_run_errored(
|
||||
self,
|
||||
run_id: str,
|
||||
status_detail: str,
|
||||
):
|
||||
run = self._runs[run_id]
|
||||
run.status = RunStatus.ERRORED
|
||||
run.status_detail = status_detail
|
||||
run.end_time_ns = current_time_ns()
|
||||
# Block on terminal status so the final state isn't lost if the controller exits right after.
|
||||
self._create_or_update_train_run(run, block=True)
|
||||
|
||||
def update_train_run_aborted(
|
||||
self,
|
||||
run_id: str,
|
||||
):
|
||||
run = self._runs[run_id]
|
||||
update_train_run_aborted(run=run, graceful=True)
|
||||
# Block on terminal status so the final state isn't lost if the controller exits right after.
|
||||
self._create_or_update_train_run(run, block=True)
|
||||
|
||||
def update_train_run_framework_versions(
|
||||
self, run_id: str, framework_versions: Dict[str, str]
|
||||
):
|
||||
run = self._runs[run_id]
|
||||
run.framework_versions = framework_versions
|
||||
self._create_or_update_train_run(run)
|
||||
|
||||
def create_train_run_attempt(
|
||||
self,
|
||||
run_id: str,
|
||||
attempt_id: str,
|
||||
num_workers: int,
|
||||
resources_per_worker: Dict[str, float],
|
||||
) -> None:
|
||||
status_detail = _get_scheduling_status_detail(num_workers, resources_per_worker)
|
||||
resources = [
|
||||
TrainResources(resources=resources_per_worker) for _ in range(num_workers)
|
||||
]
|
||||
run_attempt = TrainRunAttempt(
|
||||
run_id=run_id,
|
||||
attempt_id=attempt_id,
|
||||
start_time_ns=current_time_ns(),
|
||||
status=RunAttemptStatus.PENDING,
|
||||
status_detail=status_detail,
|
||||
resources=resources,
|
||||
workers=[], # Not started yet.
|
||||
)
|
||||
|
||||
self._run_attempts[run_id][attempt_id] = run_attempt
|
||||
self._create_or_update_train_run_attempt(run_attempt)
|
||||
|
||||
def update_train_run_attempt_running(
|
||||
self, run_id: str, attempt_id: str, workers: List[Worker]
|
||||
) -> None:
|
||||
def _convert_worker(worker: Worker) -> TrainWorker:
|
||||
|
||||
actor: ActorHandle = worker.actor
|
||||
distributed_context: DistributedContext = worker.distributed_context
|
||||
actor_metadata: ActorMetadata = worker.metadata
|
||||
|
||||
return TrainWorker(
|
||||
world_rank=distributed_context.world_rank,
|
||||
local_rank=distributed_context.local_rank,
|
||||
node_rank=distributed_context.node_rank,
|
||||
actor_id=actor._actor_id.hex(),
|
||||
node_id=actor_metadata.node_id,
|
||||
node_ip=actor_metadata.node_ip,
|
||||
pid=actor_metadata.pid,
|
||||
gpu_ids=actor_metadata.gpu_ids,
|
||||
status=ActorStatus.ALIVE,
|
||||
resources=TrainResources(resources=worker.resources),
|
||||
log_file_path=worker.log_file_path,
|
||||
)
|
||||
|
||||
workers: List[TrainWorker] = [_convert_worker(worker) for worker in workers]
|
||||
|
||||
run_attempt = self._run_attempts[run_id][attempt_id]
|
||||
run_attempt.status = RunAttemptStatus.RUNNING
|
||||
run_attempt.status_detail = None
|
||||
run_attempt.workers = workers
|
||||
self._create_or_update_train_run_attempt(run_attempt)
|
||||
|
||||
def update_train_run_attempt_finished(
|
||||
self,
|
||||
run_id: str,
|
||||
attempt_id: str,
|
||||
):
|
||||
run_attempt = self._run_attempts[run_id][attempt_id]
|
||||
run_attempt.status = RunAttemptStatus.FINISHED
|
||||
run_attempt.status_detail = None
|
||||
run_attempt.end_time_ns = current_time_ns()
|
||||
mark_workers_dead(run_attempt)
|
||||
# Block to avoid case where controller is dead but attempt is not terminal.
|
||||
self._create_or_update_train_run_attempt(run_attempt, block=True)
|
||||
|
||||
def update_train_run_attempt_errored(
|
||||
self,
|
||||
run_id: str,
|
||||
attempt_id: str,
|
||||
status_detail: str,
|
||||
):
|
||||
run_attempt = self._run_attempts[run_id][attempt_id]
|
||||
run_attempt.status = RunAttemptStatus.ERRORED
|
||||
run_attempt.status_detail = status_detail
|
||||
run_attempt.end_time_ns = current_time_ns()
|
||||
mark_workers_dead(run_attempt)
|
||||
# Block to avoid case where controller is dead but attempt is not terminal.
|
||||
self._create_or_update_train_run_attempt(run_attempt, block=True)
|
||||
|
||||
def update_train_run_attempt_aborted(
|
||||
self,
|
||||
run_id: str,
|
||||
attempt_id: str,
|
||||
):
|
||||
run_attempt = self._run_attempts[run_id][attempt_id]
|
||||
update_train_run_attempt_aborted(run_attempt=run_attempt, graceful=True)
|
||||
# Block to avoid case where controller is dead but attempt is not terminal.
|
||||
self._create_or_update_train_run_attempt(run_attempt, block=True)
|
||||
|
||||
def get_train_run_framework(self, run_id: str) -> Optional[TrainingFramework]:
|
||||
run = self._runs[run_id]
|
||||
return run.run_settings.backend_config.framework
|
||||
|
||||
def _create_or_update_train_run(
|
||||
self, run: TrainRun, *, block: bool = False
|
||||
) -> None:
|
||||
ref = self._state_actor.create_or_update_train_run.remote(run)
|
||||
if block:
|
||||
ray.get(ref)
|
||||
|
||||
def _create_or_update_train_run_attempt(
|
||||
self, run_attempt: TrainRunAttempt, *, block: bool = False
|
||||
) -> None:
|
||||
ref = self._state_actor.create_or_update_train_run_attempt.remote(run_attempt)
|
||||
if block:
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def _get_scheduling_status_detail(
|
||||
num_workers: int, resources_per_worker: Dict[str, float]
|
||||
) -> str:
|
||||
return f"Scheduling {num_workers} workers, each requiring: {resources_per_worker}."
|
||||
@@ -0,0 +1,97 @@
|
||||
import time
|
||||
|
||||
from ray.data._internal.execution.interfaces.execution_options import ExecutionOptions
|
||||
from ray.train._internal.data_config import DataConfig
|
||||
from ray.train.v2._internal.state.schema import (
|
||||
ActorStatus,
|
||||
DataConfig as DataConfigSchema,
|
||||
DataExecutionOptions,
|
||||
ExecutionOptions as ExecutionOptionsSchema,
|
||||
RunAttemptStatus,
|
||||
RunStatus,
|
||||
TrainRun,
|
||||
TrainRunAttempt,
|
||||
)
|
||||
from ray.util.state import get_actor
|
||||
|
||||
_GRACEFUL_ABORT_STATUS_DETAIL = "Run aborted due to user interrupt (SIGINT)."
|
||||
_DEAD_CONTROLLER_ABORT_STATUS_DETAIL = (
|
||||
"Run aborted because the driver process exited unexpectedly."
|
||||
)
|
||||
|
||||
|
||||
def update_train_run_aborted(run: TrainRun, graceful: bool) -> None:
|
||||
run.status = RunStatus.ABORTED
|
||||
if graceful:
|
||||
run.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
|
||||
else:
|
||||
run.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
|
||||
run.end_time_ns = current_time_ns()
|
||||
|
||||
|
||||
def update_train_run_attempt_aborted(
|
||||
run_attempt: TrainRunAttempt, graceful: bool
|
||||
) -> None:
|
||||
if graceful:
|
||||
run_attempt.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
|
||||
else:
|
||||
run_attempt.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
|
||||
run_attempt.status = RunAttemptStatus.ABORTED
|
||||
run_attempt.end_time_ns = current_time_ns()
|
||||
mark_workers_dead(run_attempt)
|
||||
|
||||
|
||||
def mark_workers_dead(run_attempt: TrainRunAttempt) -> None:
|
||||
for worker in run_attempt.workers:
|
||||
worker.status = ActorStatus.DEAD
|
||||
|
||||
|
||||
def current_time_ns() -> int:
|
||||
return time.time_ns()
|
||||
|
||||
|
||||
def is_actor_alive(actor_id: str, timeout: int) -> bool:
|
||||
"""Returns whether actor is alive."""
|
||||
actor_state = get_actor(actor_id, timeout=timeout)
|
||||
return actor_state and actor_state.state != "DEAD"
|
||||
|
||||
|
||||
def construct_data_config(data_config: DataConfig) -> DataConfigSchema:
|
||||
"""Serialize a user-facing DataConfig into the exportable schema.
|
||||
|
||||
Note: This function assumes data_config._execution_options (a defaultdict)
|
||||
hasn't been read between initialization of the field and this function call.
|
||||
Any read materializes a dataset key and affects the data config shape,
|
||||
wrongly capturing a per dataset execution options even if the user only
|
||||
provided a default.
|
||||
"""
|
||||
exec_options = data_config._execution_options
|
||||
|
||||
per_dataset_execution_options = {}
|
||||
if exec_options:
|
||||
per_dataset_execution_options = {
|
||||
ds_name: execution_options_to_model(opts)
|
||||
for ds_name, opts in exec_options.items()
|
||||
}
|
||||
|
||||
return DataConfigSchema(
|
||||
datasets_to_split=data_config._datasets_to_split,
|
||||
data_execution_options=DataExecutionOptions(
|
||||
default=execution_options_to_model(exec_options.default_factory()),
|
||||
per_dataset_execution_options=per_dataset_execution_options,
|
||||
),
|
||||
enable_shard_locality=data_config._enable_shard_locality,
|
||||
)
|
||||
|
||||
|
||||
def execution_options_to_model(
|
||||
execution_options: ExecutionOptions,
|
||||
) -> ExecutionOptionsSchema:
|
||||
"""Convert a ray.data ExecutionOptions object into the export schema model."""
|
||||
return ExecutionOptionsSchema(
|
||||
resource_limits=execution_options.resource_limits.to_resource_dict(),
|
||||
exclude_resources=execution_options.exclude_resources.to_resource_dict(),
|
||||
preserve_order=execution_options.preserve_order,
|
||||
actor_locality_enabled=execution_options.actor_locality_enabled,
|
||||
verbose_progress=execution_options.verbose_progress,
|
||||
)
|
||||
@@ -0,0 +1,369 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import functools
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
ContextManager,
|
||||
Dict,
|
||||
Generator,
|
||||
Generic,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray.train._internal.utils import count_required_parameters
|
||||
from ray.train.v2._internal.exceptions import UserExceptionWithTraceback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def bundle_to_remote_args(bundle: dict) -> dict:
|
||||
"""Convert a bundle of resources to Ray actor/task arguments.
|
||||
|
||||
>>> bundle_to_remote_args({"GPU": 1, "memory": 1, "custom": 0.1})
|
||||
{'num_cpus': 0, 'num_gpus': 1, 'memory': 1, 'resources': {'custom': 0.1}}
|
||||
"""
|
||||
bundle = bundle.copy()
|
||||
args = {
|
||||
"num_cpus": bundle.pop("CPU", 0),
|
||||
"num_gpus": bundle.pop("GPU", 0),
|
||||
"memory": bundle.pop("memory", 0),
|
||||
}
|
||||
if bundle:
|
||||
args["resources"] = bundle
|
||||
return args
|
||||
|
||||
|
||||
def construct_train_func(
|
||||
train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
|
||||
config: Optional[Dict[str, Any]],
|
||||
train_func_context: ContextManager,
|
||||
fn_arg_name: Optional[str] = "train_loop_per_worker",
|
||||
) -> Callable[[], T]:
|
||||
"""Validates and constructs the training function to execute.
|
||||
|
||||
Args:
|
||||
train_func: The training function to execute.
|
||||
This can either take in no arguments or a ``config`` dict.
|
||||
config: Configurations to pass into ``train_func``. If None then an
|
||||
empty Dict will be created.
|
||||
train_func_context: Context manager for user's `train_func`, which executes
|
||||
backend-specific logic before and after the training function.
|
||||
fn_arg_name: The name of training function to use for error messages.
|
||||
|
||||
Returns:
|
||||
A valid training function.
|
||||
|
||||
Raises:
|
||||
ValueError: if the input ``train_func`` is invalid.
|
||||
"""
|
||||
num_required_params = count_required_parameters(train_func)
|
||||
|
||||
if num_required_params > 1:
|
||||
err_msg = (
|
||||
f"{fn_arg_name} should take in 0 or 1 required arguments, but it accepts "
|
||||
f"{num_required_params} required arguments instead."
|
||||
)
|
||||
raise ValueError(err_msg)
|
||||
|
||||
if num_required_params == 1:
|
||||
config = config or {}
|
||||
|
||||
@functools.wraps(train_func)
|
||||
def train_fn():
|
||||
with train_func_context():
|
||||
return train_func(config)
|
||||
|
||||
else: # num_params == 0
|
||||
|
||||
@functools.wraps(train_func)
|
||||
def train_fn():
|
||||
with train_func_context():
|
||||
return train_func()
|
||||
|
||||
return train_fn
|
||||
|
||||
|
||||
class TrainingFramework(Enum):
|
||||
TORCH = "torch"
|
||||
JAX = "jax"
|
||||
TENSORFLOW = "tensorflow"
|
||||
XGBOOST = "xgboost"
|
||||
LIGHTGBM = "lightgbm"
|
||||
|
||||
def module_names(self) -> tuple[str, ...]:
|
||||
"""Returns the relevant module names for the training framework.
|
||||
|
||||
These module names are used by Train state version collection (see
|
||||
`_get_framework_version`) to gather versions of key framework-related packages.
|
||||
|
||||
Note: If adding a new module, make sure to use the module name rather than
|
||||
the distribution name. (e.g. sklearn instead of scikit-learn)
|
||||
"""
|
||||
if self is TrainingFramework.TORCH:
|
||||
return ("torch",)
|
||||
if self is TrainingFramework.JAX:
|
||||
return ("jax", "jaxlib")
|
||||
if self is TrainingFramework.TENSORFLOW:
|
||||
return ("tensorflow", "keras")
|
||||
if self is TrainingFramework.XGBOOST:
|
||||
return ("xgboost",)
|
||||
if self is TrainingFramework.LIGHTGBM:
|
||||
return ("lightgbm",)
|
||||
|
||||
return (self.value,)
|
||||
|
||||
|
||||
class ObjectRefWrapper(Generic[T]):
|
||||
"""Thin wrapper around ray.put to manually control dereferencing."""
|
||||
|
||||
def __init__(self, obj: T):
|
||||
self._ref = ray.put(obj)
|
||||
|
||||
def get(self) -> T:
|
||||
return ray.get(self._ref)
|
||||
|
||||
|
||||
def date_str(include_ms: bool = False):
|
||||
pattern = "%Y-%m-%d_%H-%M-%S"
|
||||
if include_ms:
|
||||
pattern += ".%f"
|
||||
return datetime.today().strftime(pattern)
|
||||
|
||||
|
||||
def time_monotonic():
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _copy_doc(copy_func):
|
||||
def wrapped(func):
|
||||
func.__doc__ = copy_func.__doc__
|
||||
return func
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def invoke_context_managers(
|
||||
context_managers: List[ContextManager],
|
||||
) -> Generator[None, None, None]:
|
||||
"""
|
||||
Utility to invoke a list of context managers and yield sequentially.
|
||||
|
||||
Args:
|
||||
context_managers: List of context managers to invoke.
|
||||
"""
|
||||
with contextlib.ExitStack() as stack:
|
||||
for context_manager in context_managers:
|
||||
stack.enter_context(context_manager())
|
||||
yield
|
||||
|
||||
|
||||
def get_module_name(obj: object) -> str:
|
||||
"""Returns the full module name of the given object, including its qualified name.
|
||||
|
||||
Args:
|
||||
obj: The object (class, function, etc.) whose module name is required.
|
||||
|
||||
Returns:
|
||||
Full module and qualified name as a string.
|
||||
"""
|
||||
return f"{obj.__module__}.{obj.__qualname__}"
|
||||
|
||||
|
||||
def get_callable_name(fn: Callable) -> str:
|
||||
"""Returns a readable name for any callable.
|
||||
|
||||
Args:
|
||||
fn: The callable to extract a name from.
|
||||
|
||||
Returns:
|
||||
A human-readable name for the callable.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> get_callable_name(lambda x: x)
|
||||
'<lambda>'
|
||||
>>> def foo(a, b): pass
|
||||
>>> get_callable_name(foo)
|
||||
'foo'
|
||||
>>> from functools import partial
|
||||
>>> bar = partial(partial(foo, a=1), b=2)
|
||||
>>> get_callable_name(bar)
|
||||
'foo'
|
||||
>>> class Dummy:
|
||||
... def __call__(self, a, b): pass
|
||||
>>> get_callable_name(Dummy())
|
||||
'Dummy'
|
||||
"""
|
||||
if isinstance(fn, functools.partial):
|
||||
return get_callable_name(fn.func)
|
||||
|
||||
# Use __name__ for regular functions and lambdas
|
||||
if hasattr(fn, "__name__"):
|
||||
return fn.__name__
|
||||
|
||||
# Fallback to the class name for objects that implement __call__
|
||||
return fn.__class__.__name__
|
||||
|
||||
|
||||
def construct_user_exception_with_traceback(
|
||||
e: BaseException, exclude_frames: int = 0
|
||||
) -> UserExceptionWithTraceback:
|
||||
"""Construct a UserExceptionWithTraceback from a base exception.
|
||||
|
||||
Args:
|
||||
e: The base exception to construct a UserExceptionWithTraceback from.
|
||||
exclude_frames: The number of frames to exclude from the beginnning of
|
||||
the traceback.
|
||||
|
||||
Returns:
|
||||
A UserExceptionWithTraceback object.
|
||||
"""
|
||||
# TODO(justinvyu): This is brittle and may break if the call stack
|
||||
# changes. Figure out a more robust way to exclude these frames.
|
||||
exc_traceback_str = traceback.format_exc(
|
||||
limit=-(len(traceback.extract_tb(e.__traceback__)) - exclude_frames)
|
||||
)
|
||||
logger.error(f"Error in training function:\n{exc_traceback_str}")
|
||||
return UserExceptionWithTraceback(e, traceback_str=exc_traceback_str)
|
||||
|
||||
|
||||
def _in_ray_train_worker() -> bool:
|
||||
"""Check if the current process is a Ray Train V2 worker."""
|
||||
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
||||
|
||||
try:
|
||||
get_train_fn_utils()
|
||||
return True
|
||||
except RuntimeError:
|
||||
return False
|
||||
|
||||
|
||||
def requires_train_worker(raise_in_tune_session: bool = False) -> Callable:
|
||||
"""Check that the caller is a Ray Train worker spawned by Ray Train,
|
||||
with access to training function utilities.
|
||||
|
||||
Args:
|
||||
raise_in_tune_session: Whether to raise a specific error message if the caller
|
||||
is in a Tune session. If True, will raise a DeprecationWarning.
|
||||
|
||||
Returns:
|
||||
A decorator that performs this check, which raises an error if the caller
|
||||
is not a Ray Train worker.
|
||||
"""
|
||||
|
||||
def _wrap(fn: Callable) -> Callable:
|
||||
@functools.wraps(fn)
|
||||
def _wrapped_fn(*args, **kwargs):
|
||||
from ray.tune.trainable.trainable_fn_utils import _in_tune_session
|
||||
|
||||
if raise_in_tune_session and _in_tune_session():
|
||||
raise DeprecationWarning(
|
||||
f"`ray.train.{fn.__name__}` is deprecated when running in a function "
|
||||
"passed to Ray Tune. Please use the equivalent `ray.tune` API instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not _in_ray_train_worker():
|
||||
raise RuntimeError(
|
||||
f"`{fn.__name__}` cannot be used outside of a Ray Train training function. "
|
||||
"You are calling this API from the driver or another non-training process. "
|
||||
"These utilities are only available within a function launched by `trainer.fit()`."
|
||||
)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return _wrapped_fn
|
||||
|
||||
return _wrap
|
||||
|
||||
|
||||
async def wait_with_logging(
|
||||
condition: asyncio.Condition,
|
||||
predicate: Optional[Callable[[], bool]] = None,
|
||||
generate_warning_message: Optional[Callable[[], str]] = None,
|
||||
warn_interval_s: float = 60,
|
||||
timeout_s: Optional[float] = None,
|
||||
):
|
||||
"""Waits for condition to be notified, logging warnings and eventually timing out.
|
||||
|
||||
You must acquire the condition before calling this function.
|
||||
|
||||
Args:
|
||||
condition: The condition to wait for.
|
||||
predicate: Wait until this predicate is True. If None, wait until the condition
|
||||
is notified.
|
||||
generate_warning_message: A function that generates the warning message to log.
|
||||
If None, no warning is logged.
|
||||
warn_interval_s: The interval in seconds to log a warning.
|
||||
timeout_s: The timeout in seconds. Defaults to``None`` to not time out.
|
||||
"""
|
||||
|
||||
async def _wait_loop():
|
||||
while True:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
condition.wait()
|
||||
if predicate is None
|
||||
else condition.wait_for(predicate),
|
||||
timeout=warn_interval_s,
|
||||
)
|
||||
return
|
||||
# asyncio.wait_for() raises `asyncio.TimeoutError` for asyncio<=3.10
|
||||
# and raises `TimeoutError` for asyncio>=3.11
|
||||
# https://docs.python.org/3/library/asyncio-task.html#asyncio.wait_for
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
if generate_warning_message is not None:
|
||||
warning_message = generate_warning_message()
|
||||
logger.warning(warning_message)
|
||||
|
||||
await asyncio.wait_for(
|
||||
_wait_loop(),
|
||||
timeout=timeout_s,
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def context_watchdog(fn: Callable, *args: Any) -> Iterator[None]:
|
||||
"""Run a function in a background thread for the duration of the context.
|
||||
|
||||
The function is started in a daemon thread on entry. On exit, a
|
||||
threading.Event is set to signal the thread to stop. The function is
|
||||
responsible for checking the event and returning promptly once it is set.
|
||||
|
||||
Args:
|
||||
fn: A function whose first argument is a threading.Event stop signal.
|
||||
The function should return when stop_event.is_set() or
|
||||
stop_event.wait(...) returns True.
|
||||
*args: Additional arguments forwarded to fn after the stop event.
|
||||
|
||||
Yields:
|
||||
None: Control is yielded to the caller while the watchdog thread runs.
|
||||
"""
|
||||
stop_event = threading.Event()
|
||||
thread = threading.Thread(
|
||||
target=fn,
|
||||
args=(stop_event, *args),
|
||||
daemon=True, # thread will end even if the finally is bypassed by an abnormal exit
|
||||
)
|
||||
thread.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
stop_event.set()
|
||||
thread.join()
|
||||
@@ -0,0 +1,53 @@
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class RayTrainCallback:
|
||||
"""Base Ray Train callback interface."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class UserCallback(RayTrainCallback):
|
||||
"""Callback interface for custom user-defined callbacks to handling events
|
||||
during training.
|
||||
|
||||
This callback is called on the Ray Train controller process, not on the
|
||||
worker processes.
|
||||
"""
|
||||
|
||||
def after_report(
|
||||
self,
|
||||
run_context: TrainRunContext,
|
||||
metrics: List[Dict[str, Any]],
|
||||
checkpoint: Optional[Checkpoint],
|
||||
):
|
||||
"""Called after all workers have reported a metric + checkpoint
|
||||
via `ray.train.report`.
|
||||
|
||||
Args:
|
||||
run_context: The `TrainRunContext` for the current training run.
|
||||
metrics: A list of metric dictionaries reported by workers,
|
||||
where metrics[i] is the metrics dict reported by worker i.
|
||||
checkpoint: A Checkpoint object that has been persisted to
|
||||
storage. This is None if no workers reported a checkpoint
|
||||
(e.g. `ray.train.report(metrics, checkpoint=None)`).
|
||||
"""
|
||||
pass
|
||||
|
||||
def after_exception(
|
||||
self, run_context: TrainRunContext, worker_exceptions: Dict[int, Exception]
|
||||
):
|
||||
"""Called after one or more workers have raised an exception.
|
||||
|
||||
Args:
|
||||
run_context: The `TrainRunContext` for the current training run.
|
||||
worker_exceptions: A dict from worker world rank to the exception
|
||||
raised by that worker.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,527 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
import pyarrow.fs
|
||||
|
||||
from ray.air.config import (
|
||||
FailureConfig as FailureConfigV1,
|
||||
ScalingConfig as ScalingConfigV1,
|
||||
)
|
||||
from ray.runtime_env import RuntimeEnv
|
||||
from ray.train.v2._internal.constants import _DEPRECATED
|
||||
from ray.train.v2._internal.execution.storage import StorageContext
|
||||
from ray.train.v2._internal.migration_utils import (
|
||||
FAIL_FAST_DEPRECATION_MESSAGE,
|
||||
TRAINER_RESOURCES_DEPRECATION_MESSAGE,
|
||||
)
|
||||
from ray.train.v2._internal.util import date_str
|
||||
from ray.util.annotations import PublicAPI
|
||||
from ray.util.tpu import get_tpu_worker_resources
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import UserCallback
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScalingConfig(ScalingConfigV1):
|
||||
"""Configuration for scaling training.
|
||||
|
||||
Args:
|
||||
num_workers: The number of workers (Ray actors) to launch.
|
||||
Each worker will reserve 1 CPU by default. The number of CPUs
|
||||
reserved by each worker can be overridden with the
|
||||
``resources_per_worker`` argument. If the number of workers is 0,
|
||||
the training function will run in local mode, meaning the training
|
||||
function runs in the same process. To enable elasticity, provide a
|
||||
``(min_workers, max_workers)`` tuple of ints.
|
||||
elastic_resize_monitor_interval_s: While the worker group is healthy,
|
||||
consider resizing the worker group every
|
||||
``elastic_resize_monitor_interval_s`` seconds.
|
||||
use_gpu: If True, training will be done on GPUs (1 per worker).
|
||||
Defaults to False. The number of GPUs reserved by each
|
||||
worker can be overridden with the ``resources_per_worker``
|
||||
argument.
|
||||
resources_per_worker: If specified, the resources
|
||||
defined in this Dict is reserved for each worker.
|
||||
Define the ``"CPU"`` and ``"GPU"`` keys (case-sensitive) to
|
||||
override the number of CPU or GPUs used by each worker.
|
||||
|
||||
Accepts the same resource keys that Ray uses for scheduling tasks
|
||||
and actors (see :ref:`Resources <core-resources>`):
|
||||
|
||||
- ``"CPU"``: number of logical CPUs per worker.
|
||||
- ``"GPU"``: number of logical GPUs per worker. Prefer setting
|
||||
``use_gpu=True`` (which reserves 1 GPU per worker) and only
|
||||
override this key when you need a different per-worker count.
|
||||
- ``"TPU"``: number of logical TPUs per worker, when ``use_tpu=True``.
|
||||
- ``"memory"``: heap memory reserved per worker, in bytes
|
||||
(for example, ``"memory": 1e9`` reserves 1 GB per worker).
|
||||
- Any :ref:`custom resource <custom-resources>` name configured on
|
||||
your cluster (for example, ``"special_hardware": 1``).
|
||||
|
||||
Keys are case-sensitive: use ``"CPU"``, ``"GPU"``, and ``"TPU"``
|
||||
(uppercase), and ``"memory"`` (lowercase).
|
||||
placement_strategy: The placement strategy to use for the
|
||||
placement group of the Ray actors. See :ref:`Placement Group
|
||||
Strategies <pgroup-strategy>` for the possible options.
|
||||
label_selector: A list of label selectors for Ray Train worker placement.
|
||||
If a single label selector is provided, it will be applied to all Ray Train workers.
|
||||
If a list is provided, it must be the same length as the max number of Ray Train workers.
|
||||
accelerator_type: [Experimental] If specified, Ray Train will launch the
|
||||
training coordinator and workers on the nodes with the specified type
|
||||
of accelerators.
|
||||
See :ref:`the available accelerator types <accelerator_types>`.
|
||||
Ensure that your cluster has instances with the specified accelerator type
|
||||
or is able to autoscale to fulfill the request. This field is required
|
||||
when `use_tpu` is True and `num_workers` is greater than 1.
|
||||
use_tpu: [Experimental] If True, training will be done on TPUs (1 TPU VM
|
||||
per worker). Defaults to False. The number of TPUs reserved by each
|
||||
worker can be overridden with the ``resources_per_worker``
|
||||
argument. This arg enables SPMD execution of the training workload.
|
||||
topology: [Experimental] If specified, Ray Train will launch the training
|
||||
coordinator and workers on nodes with the specified topology. Topology is
|
||||
auto-detected for TPUs and added as Ray node labels. This arg enables
|
||||
SPMD execution of the training workload. This field is required
|
||||
when `use_tpu` is True and `num_workers` is greater than 1.
|
||||
"""
|
||||
|
||||
num_workers: Union[int, Tuple[int, int]] = 1
|
||||
trainer_resources: Optional[dict] = None
|
||||
label_selector: Optional[Union[Dict[str, str], List[Dict[str, str]]]] = None
|
||||
|
||||
# Accelerator specific fields.
|
||||
use_tpu: Union[bool] = False
|
||||
topology: Optional[str] = None
|
||||
|
||||
# Elasticity specific fields.
|
||||
elastic_resize_monitor_interval_s: float = 60.0
|
||||
|
||||
def __post_init__(self):
|
||||
if self.trainer_resources is not None:
|
||||
raise DeprecationWarning(TRAINER_RESOURCES_DEPRECATION_MESSAGE)
|
||||
|
||||
is_fixed = isinstance(self.num_workers, int)
|
||||
is_elastic = (
|
||||
isinstance(self.num_workers, tuple)
|
||||
and len(self.num_workers) == 2
|
||||
and all(isinstance(x, int) for x in self.num_workers)
|
||||
)
|
||||
if not (is_fixed or is_elastic):
|
||||
raise ValueError(
|
||||
"ScalingConfig(num_workers) must be an int or a tuple of two ints."
|
||||
)
|
||||
if self.elastic_resize_monitor_interval_s < 0:
|
||||
raise ValueError(
|
||||
"ScalingConfig(elastic_resize_monitor_interval_s) must be non-negative."
|
||||
)
|
||||
if self.min_workers < 0:
|
||||
raise ValueError(
|
||||
f"Invalid ScalingConfig(num_workers={self.num_workers}): "
|
||||
"Number of workers cannot be negative."
|
||||
)
|
||||
if self.min_workers > self.max_workers:
|
||||
raise ValueError(
|
||||
f"Invalid ScalingConfig(num_workers={self.num_workers}): "
|
||||
f"min_workers={self.min_workers} must be <= max_workers={self.max_workers}."
|
||||
)
|
||||
|
||||
self._validate_tpu_config()
|
||||
|
||||
if (
|
||||
isinstance(self.label_selector, list)
|
||||
and len(self.label_selector) != self.max_workers
|
||||
):
|
||||
raise ValueError(
|
||||
"If `label_selector` is a list, it must be the same length as "
|
||||
"`max_workers` (or `num_workers` when fixed)."
|
||||
)
|
||||
|
||||
if self.num_workers == 0:
|
||||
logger.info(
|
||||
"Running in local mode. The training function will run in the same process. "
|
||||
"If you are using it and running into issues please file a report at "
|
||||
"https://github.com/ray-project/ray/issues."
|
||||
)
|
||||
|
||||
super().__post_init__()
|
||||
|
||||
@property
|
||||
def elasticity_enabled(self) -> bool:
|
||||
return isinstance(self.num_workers, tuple)
|
||||
|
||||
@property
|
||||
def min_workers(self) -> int:
|
||||
return (
|
||||
self.num_workers
|
||||
if isinstance(self.num_workers, int)
|
||||
else self.num_workers[0]
|
||||
)
|
||||
|
||||
@property
|
||||
def max_workers(self) -> int:
|
||||
return (
|
||||
self.num_workers
|
||||
if isinstance(self.num_workers, int)
|
||||
else self.num_workers[1]
|
||||
)
|
||||
|
||||
def _label_selector_per_worker(
|
||||
self, num_workers: int
|
||||
) -> Optional[List[Dict[str, str]]]:
|
||||
"""Normalize ``label_selector`` into a per-worker list of length ``num_workers``.
|
||||
|
||||
- ``None`` -> ``None`` (no constraint; downstream consumers — the
|
||||
placement-group path and the autoscaling coordinator — both
|
||||
accept ``None`` and treat it as "no label requirement").
|
||||
- ``Dict`` -> the same dict replicated for each worker
|
||||
- ``List`` -> the first ``num_workers`` entries (validated to be
|
||||
``max_workers`` long in ``__post_init__``)
|
||||
"""
|
||||
if isinstance(self.label_selector, list):
|
||||
return [s.copy() for s in self.label_selector[:num_workers]]
|
||||
if isinstance(self.label_selector, dict):
|
||||
return [self.label_selector.copy() for _ in range(num_workers)]
|
||||
return None
|
||||
|
||||
@property
|
||||
def total_resources(self):
|
||||
"""Map of total resources required for training.
|
||||
|
||||
For elastic configs, this returns an upper bound based on max_workers.
|
||||
"""
|
||||
total_resource_map = dict(self._trainer_resources_not_none)
|
||||
for k, value in self._resources_per_worker_not_none.items():
|
||||
total_resource_map[k] = total_resource_map.get(k, 0.0) + (
|
||||
value * self.max_workers
|
||||
)
|
||||
return total_resource_map
|
||||
|
||||
def _validate_tpu_config(self):
|
||||
"""Validates configuration specifically for TPU usage."""
|
||||
max_workers = self.max_workers
|
||||
|
||||
if self.use_gpu and self.use_tpu:
|
||||
raise ValueError("Cannot specify both `use_gpu=True` and `use_tpu=True`.")
|
||||
|
||||
if not self.use_tpu:
|
||||
if self.num_tpus_per_worker > 0:
|
||||
raise ValueError(
|
||||
"`use_tpu` is False but `TPU` was found in "
|
||||
"`resources_per_worker`. Either set `use_tpu` to True or "
|
||||
"remove `TPU` from `resources_per_worker."
|
||||
)
|
||||
# If not using TPU, we are done validating TPU-specific logic.
|
||||
return
|
||||
|
||||
if self.num_tpus_per_worker == 0:
|
||||
raise ValueError(
|
||||
"`use_tpu` is True but `TPU` is set to 0 in "
|
||||
"`resources_per_worker`. Either set `use_tpu` to False or "
|
||||
"request a positive number of `TPU` in "
|
||||
"`resources_per_worker."
|
||||
)
|
||||
|
||||
if max_workers > 1:
|
||||
if not self.topology:
|
||||
raise ValueError(
|
||||
"`topology` must be specified in ScalingConfig when `use_tpu=True` "
|
||||
" and `num_workers` > 1."
|
||||
)
|
||||
if not self.accelerator_type:
|
||||
raise ValueError(
|
||||
"`accelerator_type` must be specified in ScalingConfig when "
|
||||
"`use_tpu=True` and `num_workers` > 1."
|
||||
)
|
||||
if self.label_selector:
|
||||
raise ValueError(
|
||||
"Cannot set `label_selector` when `use_tpu=True` because "
|
||||
"Ray Train automatically reserves a TPU slice with a predefined label."
|
||||
)
|
||||
|
||||
# Validate TPU resources when both topology and accelerator type are specified.
|
||||
if self.topology and self.accelerator_type:
|
||||
try:
|
||||
workers_per_slice, tpu_resources = get_tpu_worker_resources(
|
||||
topology=self.topology,
|
||||
accelerator_type=self.accelerator_type,
|
||||
resources_per_unit=self.resources_per_worker,
|
||||
num_slices=1,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Could not parse TPU topology details for "
|
||||
f"type={self.accelerator_type}, "
|
||||
f"topology={self.topology}. Error: {e}"
|
||||
)
|
||||
|
||||
if workers_per_slice > 0 and max_workers % workers_per_slice != 0:
|
||||
raise ValueError(
|
||||
f"The configured `num_workers` ({self.num_workers}) must be a "
|
||||
f"multiple of {workers_per_slice} for the specified topology ({self.topology}). "
|
||||
"TPU workloads typically require symmetric resource distribution "
|
||||
"across all slices to function correctly."
|
||||
)
|
||||
if workers_per_slice > 0 and self.min_workers % workers_per_slice != 0:
|
||||
raise ValueError(
|
||||
f"The configured `min_workers` ({self.min_workers}) must be a "
|
||||
f"multiple of {workers_per_slice} for the specified topology ({self.topology}). "
|
||||
"TPU workloads typically require symmetric resource distribution "
|
||||
"across all slices to function correctly."
|
||||
)
|
||||
|
||||
if self.resources_per_worker is None:
|
||||
self.resources_per_worker = tpu_resources
|
||||
|
||||
@property
|
||||
def _resources_per_worker_not_none(self):
|
||||
if self.resources_per_worker is None:
|
||||
if self.use_tpu:
|
||||
return {"TPU": 1}
|
||||
|
||||
return super()._resources_per_worker_not_none
|
||||
|
||||
@property
|
||||
def _trainer_resources_not_none(self):
|
||||
return {}
|
||||
|
||||
@property
|
||||
def num_tpus_per_worker(self):
|
||||
"""The number of TPUs to set per worker."""
|
||||
return self._resources_per_worker_not_none.get("TPU", 0)
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="stable")
|
||||
class CheckpointConfig:
|
||||
"""Configuration for checkpointing.
|
||||
|
||||
Default behavior is to persist all checkpoints reported with
|
||||
:meth:`ray.train.report` to disk. If ``num_to_keep`` is set,
|
||||
the default retention policy is to keep the most recent checkpoints.
|
||||
|
||||
Args:
|
||||
num_to_keep: The maximum number of checkpoints to keep.
|
||||
If you report more checkpoints than this, the oldest
|
||||
(or lowest-scoring, if ``checkpoint_score_attribute`` is set)
|
||||
checkpoint will be deleted.
|
||||
If this is ``None`` then all checkpoints will be kept. Must be >= 1.
|
||||
checkpoint_score_attribute: The attribute that will be used to
|
||||
score checkpoints to determine which checkpoints should be kept.
|
||||
This attribute must be a key from the metrics dictionary
|
||||
attached to the checkpoint. This attribute must have a numerical value.
|
||||
checkpoint_score_order: Either "max" or "min".
|
||||
If "max"/"min", then checkpoints with highest/lowest values of
|
||||
the ``checkpoint_score_attribute`` will be kept. Defaults to "max".
|
||||
checkpoint_frequency: [Deprecated]
|
||||
checkpoint_at_end: [Deprecated]
|
||||
"""
|
||||
|
||||
num_to_keep: Optional[int] = None
|
||||
checkpoint_score_attribute: Optional[str] = None
|
||||
checkpoint_score_order: Literal["max", "min"] = "max"
|
||||
checkpoint_frequency: Union[Optional[int], Literal[_DEPRECATED]] = _DEPRECATED
|
||||
checkpoint_at_end: Union[Optional[bool], Literal[_DEPRECATED]] = _DEPRECATED
|
||||
|
||||
def __post_init__(self):
|
||||
if self.checkpoint_frequency != _DEPRECATED:
|
||||
raise DeprecationWarning(
|
||||
"`checkpoint_frequency` is deprecated since it does not "
|
||||
"apply to user-defined training functions. "
|
||||
"Please remove this argument from your CheckpointConfig."
|
||||
)
|
||||
|
||||
if self.checkpoint_at_end != _DEPRECATED:
|
||||
raise DeprecationWarning(
|
||||
"`checkpoint_at_end` is deprecated since it does not "
|
||||
"apply to user-defined training functions. "
|
||||
"Please remove this argument from your CheckpointConfig."
|
||||
)
|
||||
|
||||
if self.num_to_keep is not None and self.num_to_keep <= 0:
|
||||
raise ValueError(
|
||||
f"Received invalid num_to_keep: {self.num_to_keep}. "
|
||||
"Must be None or an integer >= 1."
|
||||
)
|
||||
|
||||
if self.checkpoint_score_order not in ("max", "min"):
|
||||
raise ValueError(
|
||||
f"Received invalid checkpoint_score_order: {self.checkpoint_score_order}. "
|
||||
"Must be 'max' or 'min'."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FailureConfig(FailureConfigV1):
|
||||
"""Configuration related to failure handling of each training run.
|
||||
|
||||
Args:
|
||||
max_failures: Tries to recover a run from training worker errors at least this many times.
|
||||
Will recover from the latest checkpoint if present.
|
||||
Setting to -1 will lead to infinite recovery retries.
|
||||
Setting to 0 will disable retries. Defaults to 0.
|
||||
controller_failure_limit: [DeveloperAPI] The maximum number of controller failures to tolerate.
|
||||
Setting to -1 will lead to infinite controller retries.
|
||||
Setting to 0 will disable controller retries. Defaults to -1.
|
||||
"""
|
||||
|
||||
fail_fast: Union[bool, str] = _DEPRECATED
|
||||
controller_failure_limit: int = -1
|
||||
|
||||
def __post_init__(self):
|
||||
if self.fail_fast != _DEPRECATED:
|
||||
raise DeprecationWarning(FAIL_FAST_DEPRECATION_MESSAGE)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@dataclass
|
||||
class LoggingConfig:
|
||||
"""Configuration for Ray Train's logging behavior.
|
||||
|
||||
Args:
|
||||
log_level: The log level for Ray Train's internal ``ray.train`` logs
|
||||
on console output and application-level log files. Accepts standard
|
||||
Python logging level names. Defaults to ``"INFO"``.
|
||||
System-level log files always capture all levels (DEBUG and above),
|
||||
and the ``ray`` logger (set by ``ray.init()``) and root logger
|
||||
are unaffected.
|
||||
"""
|
||||
|
||||
log_level: str = "INFO"
|
||||
|
||||
def __post_init__(self):
|
||||
valid_levels = set(logging._nameToLevel)
|
||||
if (
|
||||
not isinstance(self.log_level, str)
|
||||
or self.log_level.upper() not in valid_levels
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid log_level: {self.log_level!r}. "
|
||||
f"Must be one of: {', '.join(repr(x) for x in sorted(valid_levels))}."
|
||||
)
|
||||
self.log_level = self.log_level.upper()
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="stable")
|
||||
class RunConfig:
|
||||
"""Runtime configuration for training runs.
|
||||
|
||||
Args:
|
||||
name: Name of the trial or experiment. If not provided, will be deduced
|
||||
from the Trainable.
|
||||
storage_path: Path where all results and checkpoints are persisted.
|
||||
Can be a local directory or a destination on cloud storage.
|
||||
For multi-node training/tuning runs, this must be set to a
|
||||
shared storage location (e.g., S3, NFS).
|
||||
This defaults to the local ``~/ray_results`` directory.
|
||||
storage_filesystem: A custom filesystem to use for storage.
|
||||
If this is provided, `storage_path` should be a path with its
|
||||
prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`).
|
||||
failure_config: Failure mode configuration.
|
||||
checkpoint_config: Checkpointing configuration.
|
||||
callbacks: [DeveloperAPI] A list of callbacks that the Ray Train controller
|
||||
will invoke during training.
|
||||
worker_runtime_env: [DeveloperAPI] Runtime environment configuration
|
||||
for all Ray Train worker actors.
|
||||
logging_config: Configuration for Ray Train's logging behavior.
|
||||
See :class:`LoggingConfig` for details.
|
||||
"""
|
||||
|
||||
name: Optional[str] = None
|
||||
storage_path: Optional[str] = None
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None
|
||||
failure_config: Optional[FailureConfig] = None
|
||||
checkpoint_config: Optional[CheckpointConfig] = None
|
||||
callbacks: Optional[List["UserCallback"]] = None
|
||||
worker_runtime_env: Optional[Union[dict, RuntimeEnv]] = None
|
||||
logging_config: Optional[LoggingConfig] = None
|
||||
|
||||
sync_config: str = _DEPRECATED
|
||||
verbose: str = _DEPRECATED
|
||||
stop: str = _DEPRECATED
|
||||
progress_reporter: str = _DEPRECATED
|
||||
log_to_file: str = _DEPRECATED
|
||||
|
||||
def __post_init__(self):
|
||||
from ray.train.constants import DEFAULT_STORAGE_PATH
|
||||
|
||||
if self.storage_path is None:
|
||||
self.storage_path = DEFAULT_STORAGE_PATH
|
||||
|
||||
if not self.failure_config:
|
||||
self.failure_config = FailureConfig()
|
||||
|
||||
if not self.checkpoint_config:
|
||||
self.checkpoint_config = CheckpointConfig()
|
||||
|
||||
if not self.logging_config:
|
||||
self.logging_config = LoggingConfig()
|
||||
|
||||
if isinstance(self.storage_path, Path):
|
||||
self.storage_path = self.storage_path.as_posix()
|
||||
|
||||
run_config_deprecation_message = (
|
||||
"`RunConfig({})` is deprecated. This configuration was a "
|
||||
"Ray Tune API that did not support Ray Train usage well, "
|
||||
"so we are dropping support going forward. "
|
||||
"If you heavily rely on these configurations, "
|
||||
"you can run Ray Train as a single Ray Tune trial. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
unsupported_params = [
|
||||
"sync_config",
|
||||
"verbose",
|
||||
"stop",
|
||||
"progress_reporter",
|
||||
"log_to_file",
|
||||
]
|
||||
for param in unsupported_params:
|
||||
if getattr(self, param) != _DEPRECATED:
|
||||
raise DeprecationWarning(run_config_deprecation_message.format(param))
|
||||
|
||||
if not self.name:
|
||||
self.name = f"ray_train_run-{date_str()}"
|
||||
|
||||
self.callbacks = self.callbacks or []
|
||||
self.worker_runtime_env = self.worker_runtime_env or {}
|
||||
|
||||
from ray.train.v2.api.callback import RayTrainCallback
|
||||
|
||||
if not all(isinstance(cb, RayTrainCallback) for cb in self.callbacks):
|
||||
raise ValueError(
|
||||
"All callbacks must be instances of `ray.train.UserCallback`. "
|
||||
"Passing in a Ray Tune callback is no longer supported. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.checkpoint_config, CheckpointConfig):
|
||||
raise ValueError(
|
||||
f"Invalid `CheckpointConfig` type: {self.checkpoint_config.__class__}. "
|
||||
"Use `ray.train.CheckpointConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.failure_config, FailureConfig):
|
||||
raise ValueError(
|
||||
f"Invalid `FailureConfig` type: {self.failure_config.__class__}. "
|
||||
"Use `ray.train.FailureConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def storage_context(self) -> StorageContext:
|
||||
return StorageContext(
|
||||
storage_path=self.storage_path,
|
||||
experiment_dir_name=self.name,
|
||||
storage_filesystem=self.storage_filesystem,
|
||||
)
|
||||
@@ -0,0 +1,261 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict
|
||||
|
||||
from ray.train.v2._internal.execution.context import (
|
||||
get_train_context as get_internal_train_context,
|
||||
)
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class TrainContext(ABC):
|
||||
"""Abstract interface for training context."""
|
||||
|
||||
@Deprecated
|
||||
def get_metadata(self) -> Dict[str, Any]:
|
||||
"""[Deprecated] User metadata dict passed to the Trainer constructor."""
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_name(self) -> str:
|
||||
"""[Deprecated] Trial name for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_name")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_id(self) -> str:
|
||||
"""[Deprecated] Trial id for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_id")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_resources(self):
|
||||
"""[Deprecated] Trial resources for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_resources")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_dir(self) -> str:
|
||||
"""[Deprecated] Log directory corresponding to the trial directory for a Tune session.
|
||||
This is deprecated for Ray Train and should no longer be called in Ray Train workers.
|
||||
|
||||
If this directory is needed, please pass it into the `train_loop_config` directly.
|
||||
"""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_dir")
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_experiment_name(self) -> str:
|
||||
"""Experiment name for the corresponding trial."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_size(self) -> int:
|
||||
"""Get the current world size (i.e. total number of workers) for this run.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
NUM_WORKERS = 2
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
assert ray.train.get_context().get_world_size() == NUM_WORKERS
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=NUM_WORKERS),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_rank(self) -> int:
|
||||
"""Get the world rank of this worker.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
print("Worker 0")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_rank(self) -> int:
|
||||
"""Get the local rank of this worker (rank of the worker on its node).
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_local_rank() == 0:
|
||||
print("Local rank 0 worker")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_world_size(self) -> int:
|
||||
"""Get the local world size of this node (i.e. number of workers on this node).
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_local_world_size())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The number of workers running on this node.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_node_rank(self) -> int:
|
||||
"""Get the rank of this node.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_node_rank())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=1),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The rank of this node among the nodes participating in training.
|
||||
"""
|
||||
pass
|
||||
|
||||
@DeveloperAPI
|
||||
@abstractmethod
|
||||
def get_storage(self):
|
||||
"""Returns the :class:`~ray.train._internal.storage.StorageContext` storage
|
||||
context which gives advanced access to the filesystem and paths
|
||||
configured through `RunConfig`.
|
||||
|
||||
NOTE: This is a DeveloperAPI, and the `StorageContext` interface may change
|
||||
without notice between minor versions.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DistributedTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for distributed mode."""
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return get_internal_train_context().get_experiment_name()
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return get_internal_train_context().get_world_size()
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return get_internal_train_context().get_world_rank()
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return get_internal_train_context().get_local_rank()
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return get_internal_train_context().get_local_world_size()
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return get_internal_train_context().get_node_rank()
|
||||
|
||||
def get_storage(self):
|
||||
return get_internal_train_context().get_storage()
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class LocalTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for local mode."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
experiment_name: str,
|
||||
world_size: int = 1,
|
||||
world_rank: int = 0,
|
||||
local_rank: int = 0,
|
||||
local_world_size: int = 1,
|
||||
node_rank: int = 0,
|
||||
):
|
||||
self.experiment_name = experiment_name
|
||||
self.world_size = world_size
|
||||
self.world_rank = world_rank
|
||||
self.local_rank = local_rank
|
||||
self.local_world_size = local_world_size
|
||||
self.node_rank = node_rank
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return self.experiment_name
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return self.world_size
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return self.world_rank
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return self.local_rank
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return self.local_world_size
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return self.node_rank
|
||||
|
||||
def get_storage(self):
|
||||
raise NotImplementedError("Local storage context not yet implemented. ")
|
||||
@@ -0,0 +1,335 @@
|
||||
import logging
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
|
||||
from ray._common.usage import usage_lib
|
||||
from ray._private.ray_constants import env_bool
|
||||
from ray.actor import ActorHandle
|
||||
from ray.air._internal.usage import tag_train_v2_trainer
|
||||
from ray.train import (
|
||||
BackendConfig,
|
||||
Checkpoint,
|
||||
DataConfig,
|
||||
Result,
|
||||
RunConfig,
|
||||
ScalingConfig,
|
||||
)
|
||||
from ray.train.base_trainer import (
|
||||
_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING,
|
||||
_TRAINER_RESTORE_DEPRECATION_WARNING,
|
||||
)
|
||||
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR, RAY_TRAIN_ENABLE_STATE_TRACKING
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
from ray.train.v2._internal.callbacks import (
|
||||
AcceleratorSetupCallback,
|
||||
BackendSetupCallback,
|
||||
DatasetsCallback,
|
||||
WorkingDirectorySetupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.env_callback import _initialize_env_callbacks
|
||||
from ray.train.v2._internal.callbacks.metrics import (
|
||||
ControllerMetricsCallback,
|
||||
WorkerMetricsCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.placement_group_callback import (
|
||||
PlacementGroupCleanerCallback,
|
||||
)
|
||||
from ray.train.v2._internal.callbacks.state_manager import StateManagerCallback
|
||||
from ray.train.v2._internal.callbacks.user_callback import UserCallbackHandler
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
|
||||
METRICS_ENABLED_ENV_VAR,
|
||||
V2_ENABLED_ENV_VAR,
|
||||
get_env_vars_to_propagate,
|
||||
is_v2_enabled,
|
||||
)
|
||||
from ray.train.v2._internal.data_integration.interfaces import GenDataset
|
||||
from ray.train.v2._internal.execution.callback import RayTrainCallback
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
from ray.train.v2._internal.execution.controller import TrainController
|
||||
from ray.train.v2._internal.execution.failure_handling import create_failure_policy
|
||||
from ray.train.v2._internal.execution.local_mode.utils import LocalController
|
||||
from ray.train.v2._internal.execution.scaling_policy import create_scaling_policy
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper, construct_train_func
|
||||
from ray.train.v2.api.callback import UserCallback
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DataParallelTrainer:
|
||||
"""Base class for distributed data parallel training on Ray.
|
||||
|
||||
This class supports the SPMD parallelization pattern, where a single
|
||||
training function is executed in parallel across multiple workers,
|
||||
and different shards of data are processed by each worker.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
backend_config: Optional[BackendConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[DataConfig] = None,
|
||||
# TODO: [Deprecated] Remove in future release
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
self.run_config = run_config or RunConfig()
|
||||
self.train_loop_per_worker = train_loop_per_worker
|
||||
self.validation_config = validation_config
|
||||
self.train_loop_config = train_loop_config
|
||||
self.scaling_config = scaling_config or ScalingConfig()
|
||||
self.backend_config = backend_config or BackendConfig()
|
||||
self.datasets = datasets or {}
|
||||
self.data_config = dataset_config or DataConfig()
|
||||
|
||||
self.running_in_local_mode = self.scaling_config.num_workers == 0
|
||||
|
||||
self.train_run_context = TrainRunContext(
|
||||
run_config=self.run_config,
|
||||
train_loop_config=self.train_loop_config,
|
||||
scaling_config=self.scaling_config,
|
||||
backend_config=self.backend_config,
|
||||
dataset_config=self.data_config,
|
||||
)
|
||||
|
||||
if resume_from_checkpoint is not None:
|
||||
raise DeprecationWarning(_RESUME_FROM_CHECKPOINT_DEPRECATION_WARNING)
|
||||
|
||||
if metadata is not None:
|
||||
raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
|
||||
self._validate_configs()
|
||||
|
||||
usage_lib.record_library_usage("train")
|
||||
tag_train_v2_trainer(self)
|
||||
if self.scaling_config.elasticity_enabled:
|
||||
usage_lib.record_extra_usage_tag(
|
||||
usage_lib.TagKey.TRAIN_ELASTICITY_ENABLED, "1"
|
||||
)
|
||||
|
||||
def _validate_configs(self):
|
||||
if not is_v2_enabled():
|
||||
raise ValueError(
|
||||
f"Ray Train V2 must be enabled with `{V2_ENABLED_ENV_VAR}=1` "
|
||||
"when using this V2 Trainer API."
|
||||
)
|
||||
|
||||
from ray.train.v2.api.config import (
|
||||
RunConfig as RunConfigV2,
|
||||
ScalingConfig as ScalingConfigV2,
|
||||
)
|
||||
|
||||
if not isinstance(self.run_config, RunConfigV2):
|
||||
raise ValueError(
|
||||
f"Invalid `RunConfig` type: {self.run_config.__class__}. "
|
||||
"Use `ray.train.RunConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
if not isinstance(self.scaling_config, ScalingConfigV2):
|
||||
raise ValueError(
|
||||
f"Invalid `ScalingConfig` type: {self.scaling_config.__class__}. "
|
||||
"Use `ray.train.ScalingConfig` instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/49454"
|
||||
)
|
||||
|
||||
def _get_train_func(self) -> Callable[[], Any]:
|
||||
return construct_train_func(
|
||||
self.train_loop_per_worker,
|
||||
config=self.train_loop_config,
|
||||
train_func_context=self.backend_config.train_func_context,
|
||||
fn_arg_name="train_loop_per_worker",
|
||||
)
|
||||
|
||||
def fit(self) -> Result:
|
||||
"""Launches the Ray Train controller to run training on workers.
|
||||
|
||||
Returns:
|
||||
A Result object containing the training result.
|
||||
|
||||
Raises:
|
||||
ray.train.TrainingFailedError: This is a union of the ControllerError and WorkerGroupError.
|
||||
This returns a :class:`ray.train.ControllerError` if internal Ray Train controller logic
|
||||
encounters a non-retryable error or reaches the controller failure limit configured in `FailureConfig`.
|
||||
This returns a :class:`ray.train.WorkerGroupError` if one or more workers fail during
|
||||
training and reaches the worker group failure limit configured in `FailureConfig(max_failures)`.
|
||||
"""
|
||||
train_fn = self._get_train_func()
|
||||
if self.running_in_local_mode:
|
||||
return self._initialize_and_run_local_controller(train_fn)
|
||||
else:
|
||||
train_fn_ref = ObjectRefWrapper(train_fn)
|
||||
|
||||
result = self._initialize_and_run_controller(
|
||||
train_fn_ref=train_fn_ref,
|
||||
scaling_policy=create_scaling_policy(self.scaling_config),
|
||||
failure_policy=create_failure_policy(self.run_config.failure_config),
|
||||
train_run_context=self.train_run_context,
|
||||
callbacks=self._create_default_callbacks(),
|
||||
validation_config=self.validation_config,
|
||||
)
|
||||
|
||||
if result.error:
|
||||
# NOTE: If the training run errored out, raise an error back to the
|
||||
# user's driver script.
|
||||
# For example, if the Train `FailurePolicy` runs out of retries,
|
||||
# and one of the workers errors. The controller will exit, and
|
||||
# the error will be raised here.
|
||||
raise result.error
|
||||
|
||||
return result
|
||||
|
||||
def _get_local_controller(self) -> LocalController:
|
||||
return LocalController(
|
||||
experiment_name=self.run_config.name,
|
||||
datasets=self.datasets,
|
||||
)
|
||||
|
||||
def _create_default_callbacks(self) -> List[RayTrainCallback]:
|
||||
# Initialize callbacks from environment variable
|
||||
callbacks = _initialize_env_callbacks()
|
||||
|
||||
accelerator_setup_callback = AcceleratorSetupCallback(
|
||||
self.backend_config, self.scaling_config
|
||||
)
|
||||
backend_setup_callback = BackendSetupCallback(self.backend_config)
|
||||
datasets_callback = DatasetsCallback(
|
||||
train_run_context=self.train_run_context,
|
||||
datasets=self.datasets,
|
||||
)
|
||||
placement_group_cleaner_callback = PlacementGroupCleanerCallback()
|
||||
callbacks.extend(
|
||||
[
|
||||
accelerator_setup_callback,
|
||||
backend_setup_callback,
|
||||
placement_group_cleaner_callback,
|
||||
datasets_callback,
|
||||
]
|
||||
)
|
||||
if env_bool(RAY_CHDIR_TO_TRIAL_DIR, True):
|
||||
working_directory_setup_callback = WorkingDirectorySetupCallback()
|
||||
callbacks.append(working_directory_setup_callback)
|
||||
|
||||
if env_bool(METRICS_ENABLED_ENV_VAR, True):
|
||||
callbacks.append(ControllerMetricsCallback())
|
||||
callbacks.append(WorkerMetricsCallback(self.train_run_context))
|
||||
|
||||
if env_bool(RAY_TRAIN_ENABLE_STATE_TRACKING, False):
|
||||
callbacks.append(StateManagerCallback(datasets=self.datasets))
|
||||
|
||||
run_config_callbacks = (
|
||||
self.run_config.callbacks if self.run_config.callbacks is not None else []
|
||||
)
|
||||
|
||||
# Add internal callback that invokes all user-defined callbacks.
|
||||
user_callbacks = [
|
||||
cb for cb in run_config_callbacks if isinstance(cb, UserCallback)
|
||||
]
|
||||
callbacks.append(
|
||||
UserCallbackHandler(
|
||||
user_callbacks=user_callbacks, train_run_context=self.train_run_context
|
||||
)
|
||||
)
|
||||
|
||||
# Append all other callbacks to the full list. This allows custom workarounds
|
||||
# built on top of internal callbacks to work.
|
||||
callbacks.extend(
|
||||
[cb for cb in run_config_callbacks if not isinstance(cb, UserCallback)]
|
||||
)
|
||||
return callbacks
|
||||
|
||||
def _initialize_and_run_local_controller(
|
||||
self, train_func: Callable[[], Any]
|
||||
) -> Result:
|
||||
return self._get_local_controller().run(train_func)
|
||||
|
||||
def _initialize_and_run_controller(self, **controller_init_kwargs) -> Result:
|
||||
env_vars = get_env_vars_to_propagate()
|
||||
env_vars.setdefault(
|
||||
RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
|
||||
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE,
|
||||
)
|
||||
|
||||
# Attach the controller to the node running the driver script.
|
||||
controller_actor_cls = ray.remote(
|
||||
num_cpus=0,
|
||||
label_selector={
|
||||
ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
|
||||
},
|
||||
# TODO: Extract env variables that affect controller behavior
|
||||
# and pass them as explicit args
|
||||
runtime_env={"env_vars": env_vars},
|
||||
)(TrainController)
|
||||
|
||||
controller = controller_actor_cls.remote(**controller_init_kwargs)
|
||||
|
||||
# If this is not the main thread - as is the case when running in Tune -
|
||||
# registering the SIGINT handler raises an exception.
|
||||
if threading.current_thread() is threading.main_thread():
|
||||
self._register_sigint_handler(controller)
|
||||
|
||||
ray.get(controller.run.remote())
|
||||
return ray.get(controller.get_result.remote())
|
||||
|
||||
def _register_sigint_handler(self, controller: ActorHandle[TrainController]):
|
||||
"""Register SIGINT handler so user Ctrl C gracefully aborts run."""
|
||||
sigint_count = 0
|
||||
|
||||
def sigint_handler(signum, frame):
|
||||
logger.info(
|
||||
"Received SIGINT. Gracefully aborting the training run — this "
|
||||
"may take a few seconds. To forcefully abort immediately, you "
|
||||
"can send a different signal, such as SIGKILL."
|
||||
)
|
||||
nonlocal sigint_count
|
||||
sigint_count += 1
|
||||
if sigint_count >= 3:
|
||||
logger.info(
|
||||
"Received SIGINT at least 3 times. "
|
||||
"Forcefully aborting the training run."
|
||||
)
|
||||
sys.exit(0)
|
||||
if sigint_count <= 1:
|
||||
try:
|
||||
ray.get(controller.abort.remote())
|
||||
except ray.exceptions.RayActorError:
|
||||
# We catch the error and exit 0 to indicate graceful termination.
|
||||
# However, for some reason the process still exits with 1.
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, sigint_handler)
|
||||
|
||||
@classmethod
|
||||
@Deprecated
|
||||
def restore(cls, *args, **kwargs):
|
||||
"""[Deprecated] Restores a Train experiment from a previously
|
||||
interrupted/failed run.
|
||||
|
||||
This method is deprecated and will be removed in a future release.
|
||||
"""
|
||||
raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
|
||||
@classmethod
|
||||
@Deprecated
|
||||
def can_restore(cls, *args, **kwargs):
|
||||
"""[Deprecated] Checks if a Train experiment can be restored from
|
||||
a previously interrupted/failed run.
|
||||
|
||||
This method is deprecated and will be removed in a future release.
|
||||
"""
|
||||
raise DeprecationWarning(_TRAINER_RESTORE_DEPRECATION_WARNING)
|
||||
@@ -0,0 +1,49 @@
|
||||
from typing import Dict
|
||||
|
||||
from ray.train.v2._internal.exceptions import RayTrainError
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class TrainingFailedError(RayTrainError):
|
||||
"""Exception raised when training fails from a `trainer.fit()` call.
|
||||
This is either :class:`ray.train.WorkerGroupError` or :class:`ray.train.ControllerError`.
|
||||
"""
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class WorkerGroupError(TrainingFailedError):
|
||||
"""Exception raised from the worker group during training.
|
||||
|
||||
Args:
|
||||
error_message: A human-readable error message describing the training worker failures.
|
||||
worker_failures: A mapping from worker rank to the exception that
|
||||
occurred on that worker during training.
|
||||
"""
|
||||
|
||||
def __init__(self, error_message: str, worker_failures: Dict[int, Exception]):
|
||||
super().__init__("Training failed due to worker errors:\n" + error_message)
|
||||
self._error_message = error_message
|
||||
self.worker_failures = worker_failures
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self._error_message, self.worker_failures))
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ControllerError(TrainingFailedError):
|
||||
"""Exception raised when training fails due to a controller error.
|
||||
|
||||
Args:
|
||||
controller_failure: The exception that occurred on the controller.
|
||||
"""
|
||||
|
||||
def __init__(self, controller_failure: Exception):
|
||||
super().__init__(
|
||||
"Training failed due to controller error:\n" + str(controller_failure)
|
||||
)
|
||||
self.controller_failure = controller_failure
|
||||
self.with_traceback(controller_failure.__traceback__)
|
||||
|
||||
def __reduce__(self):
|
||||
return (self.__class__, (self.controller_failure,))
|
||||
@@ -0,0 +1,36 @@
|
||||
from enum import Enum
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class CheckpointUploadMode(Enum):
|
||||
"""The manner in which we want to upload the checkpoint.
|
||||
|
||||
Members:
|
||||
ASYNC: Upload checkpoint asynchronously.
|
||||
SYNC: Upload checkpoint synchronously.
|
||||
NO_UPLOAD: Do not upload checkpoint.
|
||||
"""
|
||||
|
||||
ASYNC = "ASYNC"
|
||||
SYNC = "SYNC"
|
||||
NO_UPLOAD = "NO_UPLOAD"
|
||||
|
||||
def default_delete_local_checkpoint_after_upload(self) -> bool:
|
||||
return self == CheckpointUploadMode.ASYNC
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class CheckpointConsistencyMode(Enum):
|
||||
"""Read semantics for checkpoint retrieval during an ongoing run.
|
||||
|
||||
Members:
|
||||
COMMITTED: Block until the checkpoint from the latest ray.train.report
|
||||
has been uploaded and committed.
|
||||
VALIDATED: Block until the checkpoint from the latest ray.train.report
|
||||
has been uploaded and validated.
|
||||
"""
|
||||
|
||||
COMMITTED = "COMMITTED"
|
||||
VALIDATED = "VALIDATED"
|
||||
@@ -0,0 +1,43 @@
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, Dict
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ReportedCheckpointStatus(Enum):
|
||||
"""ReportedCheckpoint status enum.
|
||||
|
||||
* COMMITTED: The checkpoint is saved, and no validation was requested.
|
||||
* PENDING_VALIDATION: The checkpoint is saved, and validation is in progress.
|
||||
* VALIDATED: The checkpoint is saved, and validation is complete.
|
||||
* VALIDATION_TIMEOUT: The checkpoint is saved, and validation is timed out according to
|
||||
`ValidationTaskConfig(..., timeout_s=N)`.
|
||||
* VALIDATION_FAILED: The checkpoint is saved, and validation failed.
|
||||
"""
|
||||
|
||||
COMMITTED = "COMMITTED"
|
||||
PENDING_VALIDATION = "PENDING_VALIDATION"
|
||||
VALIDATED = "VALIDATED"
|
||||
VALIDATION_TIMEOUT = "VALIDATION_TIMEOUT"
|
||||
VALIDATION_FAILED = "VALIDATION_FAILED"
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="alpha")
|
||||
class ReportedCheckpoint:
|
||||
"""A user-reported checkpoint and its associated metrics.
|
||||
|
||||
Attributes:
|
||||
checkpoint: The checkpoint reported by the user.
|
||||
metrics: The metrics associated with that checkpoint.
|
||||
status: The status of the checkpoint.
|
||||
"""
|
||||
|
||||
checkpoint: "Checkpoint"
|
||||
metrics: Dict[str, Any]
|
||||
status: ReportedCheckpointStatus
|
||||
@@ -0,0 +1,162 @@
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow
|
||||
|
||||
import ray
|
||||
from ray.air.result import Result as ResultV1
|
||||
from ray.train import Checkpoint, CheckpointConfig
|
||||
from ray.train.v2._internal.constants import CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.storage import (
|
||||
StorageContext,
|
||||
_exists_at_fs_path,
|
||||
get_fs_and_path,
|
||||
)
|
||||
from ray.train.v2.api.exceptions import TrainingFailedError
|
||||
from ray.util.annotations import Deprecated, PublicAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Result(ResultV1):
|
||||
"""The output of ``trainer.fit()``.
|
||||
|
||||
Attributes:
|
||||
metrics: The latest set of metrics reported by the training function
|
||||
via :func:`ray.train.report`.
|
||||
checkpoint: The latest checkpoint saved by the training function
|
||||
via :func:`ray.train.report`.
|
||||
return_value: The value returned by the user-defined training function on the
|
||||
rank 0 worker, or ``None`` if no value was returned or if training did
|
||||
not complete successfully. The return value must be serializable.
|
||||
metrics_dataframe: A DataFrame of metrics from all checkpoints saved
|
||||
during the run. Each row corresponds to a checkpoint.
|
||||
best_checkpoints: A list of ``(checkpoint, metrics)`` tuples for the
|
||||
best checkpoints saved during the run. The checkpoints retained
|
||||
are determined by :class:`~ray.train.CheckpointConfig`
|
||||
(by default, all checkpoints are kept).
|
||||
path: Path pointing to the run output directory on persistent storage.
|
||||
This can point to a remote storage location (e.g. S3) or to a
|
||||
local location on the head node.
|
||||
error: The execution error of the training run, if the run finished
|
||||
in error. This is a :class:`~ray.train.v2.api.exceptions.TrainingFailedError`
|
||||
wrapping the original exception.
|
||||
"""
|
||||
|
||||
checkpoint: Optional[Checkpoint]
|
||||
error: Optional[TrainingFailedError]
|
||||
best_checkpoints: Optional[List[Tuple[Checkpoint, Dict[str, Any]]]] = None
|
||||
return_value: Optional[Any] = None
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
def get_best_checkpoint(
|
||||
self, metric: str, mode: str
|
||||
) -> Optional["ray.train.Checkpoint"]:
|
||||
return super().get_best_checkpoint(metric, mode)
|
||||
|
||||
@classmethod
|
||||
def from_path(
|
||||
cls,
|
||||
path: Union[str, os.PathLike],
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
) -> "Result":
|
||||
"""Restore a training result from a previously saved training run path.
|
||||
|
||||
Args:
|
||||
path: Path to the run output directory
|
||||
storage_filesystem: Optional filesystem to use for accessing the path
|
||||
|
||||
Returns:
|
||||
Result object with restored checkpoints and metrics
|
||||
"""
|
||||
fs, fs_path = get_fs_and_path(str(path), storage_filesystem)
|
||||
|
||||
# Validate that the experiment directory exists
|
||||
if not _exists_at_fs_path(fs, fs_path):
|
||||
raise RuntimeError(f"Experiment folder {fs_path} doesn't exist.")
|
||||
|
||||
# Remove trailing slashes to handle paths correctly
|
||||
# os.path.basename() returns empty string for paths with trailing slashes
|
||||
fs_path = fs_path.rstrip("/")
|
||||
storage_path, experiment_dir_name = os.path.dirname(fs_path), os.path.basename(
|
||||
fs_path
|
||||
)
|
||||
|
||||
storage_context = StorageContext(
|
||||
storage_path=storage_path,
|
||||
experiment_dir_name=experiment_dir_name,
|
||||
storage_filesystem=fs,
|
||||
read_only=True,
|
||||
)
|
||||
|
||||
# Validate that the checkpoint manager snapshot file exists
|
||||
if not _exists_at_fs_path(
|
||||
storage_context.storage_filesystem,
|
||||
storage_context.checkpoint_manager_snapshot_path,
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Failed to restore the Result object: "
|
||||
f"{CHECKPOINT_MANAGER_SNAPSHOT_FILENAME} doesn't exist in the "
|
||||
f"experiment folder. Make sure that this is an output directory created by a Ray Train run."
|
||||
)
|
||||
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
|
||||
# When we build a Result object from checkpoints, the error is not loaded.
|
||||
return cls._from_checkpoint_manager(
|
||||
checkpoint_manager=checkpoint_manager,
|
||||
storage_context=storage_context,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _from_checkpoint_manager(
|
||||
cls,
|
||||
checkpoint_manager: CheckpointManager,
|
||||
storage_context: StorageContext,
|
||||
error: Optional[TrainingFailedError] = None,
|
||||
) -> "Result":
|
||||
"""Create a Result object from a CheckpointManager."""
|
||||
latest_checkpoint_result = checkpoint_manager.latest_checkpoint_result
|
||||
if latest_checkpoint_result:
|
||||
latest_metrics = latest_checkpoint_result.metrics
|
||||
latest_checkpoint = latest_checkpoint_result.checkpoint
|
||||
else:
|
||||
latest_metrics = None
|
||||
latest_checkpoint = None
|
||||
best_checkpoints = [
|
||||
(r.checkpoint, r.metrics)
|
||||
for r in checkpoint_manager.best_checkpoint_results
|
||||
]
|
||||
|
||||
# Provide the history of metrics attached to checkpoints as a dataframe.
|
||||
metrics_dataframe = None
|
||||
if best_checkpoints:
|
||||
metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
|
||||
|
||||
return Result(
|
||||
metrics=latest_metrics,
|
||||
checkpoint=latest_checkpoint,
|
||||
error=error,
|
||||
path=storage_context.experiment_fs_path,
|
||||
best_checkpoints=best_checkpoints,
|
||||
metrics_dataframe=metrics_dataframe,
|
||||
_storage_filesystem=storage_context.storage_filesystem,
|
||||
)
|
||||
|
||||
@property
|
||||
@Deprecated
|
||||
def config(self) -> Optional[Dict[str, Any]]:
|
||||
raise DeprecationWarning(
|
||||
"The `config` property for a `ray.train.Result` is deprecated, "
|
||||
"since it is only relevant in the context of Ray Tune."
|
||||
)
|
||||
@@ -0,0 +1,297 @@
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
|
||||
from ray.train.v2._internal.data_integration.interfaces import DatasetShardMetadata
|
||||
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
||||
from ray.train.v2._internal.util import requires_train_worker
|
||||
from ray.train.v2.api.context import TrainContext
|
||||
from ray.train.v2.api.report_config import (
|
||||
CheckpointConsistencyMode,
|
||||
CheckpointUploadMode,
|
||||
)
|
||||
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data import DataIterator
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def report(
|
||||
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,
|
||||
):
|
||||
"""Report metrics and optionally save a checkpoint.
|
||||
|
||||
If a checkpoint is provided, it will be
|
||||
:ref:`persisted to storage <persistent-storage-guide>`.
|
||||
|
||||
If this is called in multiple distributed training workers:
|
||||
|
||||
- Only the metrics reported by the rank 0 worker will be attached to the checkpoint.
|
||||
- A checkpoint will be registered as long as one or more workers reports
|
||||
checkpoint that is not None.
|
||||
See the :ref:`checkpointing guide <train-dl-saving-checkpoints>`.
|
||||
- Checkpoints from multiple workers will be merged into one directory
|
||||
in persistent storage.
|
||||
See :ref:`the distributed checkpointing guide <train-distributed-checkpointing>`.
|
||||
|
||||
|
||||
.. warning::
|
||||
|
||||
All workers must call `ray.train.report` the same number of times
|
||||
so that Ray Train can properly synchronize the training state across
|
||||
workers. This method acts as a barrier across all workers, so be sure
|
||||
that every worker reaches this method.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 10)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": ...}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
# torch.save(...)
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
|
||||
# Example: Only the rank 0 worker uploads the checkpoint.
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
else:
|
||||
ray.train.report(metrics, checkpoint=None)
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
|
||||
Args:
|
||||
metrics: The metrics you want to report.
|
||||
checkpoint: The optional checkpoint you want to report.
|
||||
checkpoint_dir_name: Custom name for the checkpoint directory.
|
||||
If not provided, a unique directory name will be automatically generated.
|
||||
If provided, it must be unique across all checkpoints per worker to avoid
|
||||
naming collisions. Consider including identifiers such as the epoch or batch
|
||||
index in the name.
|
||||
checkpoint_upload_mode: The manner in which we want to upload the checkpoint.
|
||||
Defaults to uploading the checkpoint synchronously.
|
||||
This works when no checkpoint is provided but is not useful in that case.
|
||||
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: [Alpha] If True, triggers validation with default kwargs from validation_config.
|
||||
If a ValidationTaskConfig, validation is run using fn_kwargs merged with validation_config
|
||||
defaults, with fn_kwargs taking precedence on conflicts. If False, no validation.
|
||||
"""
|
||||
if validation and not checkpoint:
|
||||
raise ValueError("Validation requires a checkpoint to be provided.")
|
||||
|
||||
if delete_local_checkpoint_after_upload is None:
|
||||
delete_local_checkpoint_after_upload = (
|
||||
checkpoint_upload_mode.default_delete_local_checkpoint_after_upload()
|
||||
)
|
||||
|
||||
if checkpoint:
|
||||
record_extra_usage_tag(
|
||||
TagKey.TRAIN_CHECKPOINT_MODE, checkpoint_upload_mode.value
|
||||
)
|
||||
if validation:
|
||||
record_extra_usage_tag(TagKey.TRAIN_ASYNCHRONOUS_VALIDATION, "1")
|
||||
|
||||
get_train_fn_utils().report(
|
||||
metrics=metrics,
|
||||
checkpoint=checkpoint,
|
||||
checkpoint_dir_name=checkpoint_dir_name,
|
||||
checkpoint_upload_mode=checkpoint_upload_mode,
|
||||
delete_local_checkpoint_after_upload=delete_local_checkpoint_after_upload,
|
||||
checkpoint_upload_fn=checkpoint_upload_fn,
|
||||
validation=validation,
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def get_context() -> TrainContext:
|
||||
"""Get or create a singleton training context.
|
||||
|
||||
The context is only available within a function passed to Ray Train.
|
||||
|
||||
See the :class:`~ray.train.TrainContext` API reference to see available methods.
|
||||
"""
|
||||
return get_train_fn_utils().get_context()
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker(raise_in_tune_session=True)
|
||||
def get_checkpoint() -> Optional["Checkpoint"]:
|
||||
"""Access the latest reported checkpoint to resume from if one exists.
|
||||
|
||||
See :ref:`the checkpoint loading guide <train-dl-loading-checkpoints>` for more details.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
checkpoint = ray.train.get_checkpoint()
|
||||
if checkpoint:
|
||||
with checkpoint.as_directory() as checkpoint_dir:
|
||||
# Load back training state
|
||||
...
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 10)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": ...}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
|
||||
Returns:
|
||||
Checkpoint object if the session is currently being resumed.
|
||||
Otherwise, return None.
|
||||
"""
|
||||
return get_train_fn_utils().get_checkpoint()
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@requires_train_worker()
|
||||
def get_all_reported_checkpoints(
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
"""Get all the reported checkpoints so far.
|
||||
|
||||
Blocks until Ray Train has finished processing every in-flight `ray.train.report` call.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import tempfile
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def train_func(config):
|
||||
start_epoch = 0
|
||||
|
||||
for epoch in range(start_epoch, config.get("num_epochs", 2)):
|
||||
# Do training...
|
||||
|
||||
metrics = {"loss": 0.1}
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
# Save the checkpoint...
|
||||
|
||||
checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
|
||||
ray.train.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
reported_checkpoints = ray.train.get_all_reported_checkpoints()
|
||||
# Report artifacts/metrics to experiment tracking framework...
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Args:
|
||||
consistency_mode: Read semantics for checkpoint retrieval during an ongoing run.
|
||||
Defaults to CheckpointConsistencyMode.VALIDATED.
|
||||
See :class:`~ray.train.CheckpointConsistencyMode` for more details.
|
||||
timeout_s: Timeout in seconds to collecting checkpoint and validation information.
|
||||
Defaults to None to wait indefinitely.
|
||||
|
||||
Returns:
|
||||
List of ReportedCheckpoint objects that represent the checkpoints and
|
||||
corresponding metrics reported by the workers.
|
||||
"""
|
||||
return get_train_fn_utils().get_all_reported_checkpoints(
|
||||
consistency_mode=consistency_mode, timeout_s=timeout_s
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
@requires_train_worker()
|
||||
def get_dataset_shard(dataset_name: Optional[str] = None) -> Optional["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.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
...
|
||||
for epoch in range(2):
|
||||
# Trainer will automatically handle sharding.
|
||||
data_shard = ray.train.get_dataset_shard("train")
|
||||
for batch in data_shard.iter_torch_batches():
|
||||
...
|
||||
|
||||
train_dataset = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
datasets={"train": train_dataset}
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Args:
|
||||
dataset_name: If a Dictionary of Datasets was passed to ``Trainer``, then
|
||||
specifies which dataset shard to return.
|
||||
|
||||
Returns:
|
||||
The ``DataIterator`` shard to use for this worker.
|
||||
If no dataset is passed into Trainer, then return None.
|
||||
"""
|
||||
train_fn_utils = get_train_fn_utils()
|
||||
return train_fn_utils.get_dataset_shard(
|
||||
DatasetShardMetadata(
|
||||
dataset_name=dataset_name,
|
||||
world_rank=train_fn_utils.get_context().get_world_rank(),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,74 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol
|
||||
|
||||
from ray.util.annotations import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train import Checkpoint
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationFn(Protocol):
|
||||
"""Protocol for a function that validates a checkpoint."""
|
||||
|
||||
def __call__(self, checkpoint: "Checkpoint", **kwargs: Any) -> Dict:
|
||||
...
|
||||
|
||||
|
||||
@dataclass
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationTaskConfig:
|
||||
"""Configuration for a specific validation task, passed to report().
|
||||
|
||||
Args:
|
||||
fn_kwargs: json-serializable keyword arguments to pass to the validation function.
|
||||
Note that we always pass `checkpoint` as the first argument to the
|
||||
validation function.
|
||||
timeout_s: Timeout in seconds for this validation task.
|
||||
``None`` is no timeout.
|
||||
"""
|
||||
|
||||
fn_kwargs: Optional[Dict[str, Any]] = None
|
||||
timeout_s: Optional[float] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.fn_kwargs is None:
|
||||
self.fn_kwargs = {}
|
||||
assert (
|
||||
self.timeout_s is None or self.timeout_s > 0
|
||||
), f"The `timeout_s` should be None or greater than zero, actual value: {self.timeout_s}"
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class ValidationConfig:
|
||||
"""Configuration for validation, passed to the trainer.
|
||||
|
||||
Args:
|
||||
fn: The validation function to run on checkpoints.
|
||||
This function should accept a checkpoint as the first argument
|
||||
and return a dictionary of metrics.
|
||||
task_config: Default configuration for validation tasks.
|
||||
The fn_kwargs in this config can be overridden by
|
||||
ValidationTaskConfig passed to report().
|
||||
ray_remote_kwargs: Keyword arguments to pass to `ray.remote()` for the validation task.
|
||||
This can be used to specify resource requirements, number of retries, etc.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fn: ValidationFn,
|
||||
task_config: Optional[ValidationTaskConfig] = None,
|
||||
ray_remote_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.fn = fn
|
||||
if task_config is None:
|
||||
self.task_config = ValidationTaskConfig()
|
||||
else:
|
||||
self.task_config = task_config
|
||||
# TODO: ray_remote_kwargs is not json-serializable because retry_exceptions
|
||||
# can be a list of exception types. If ray core makes ray_remote_kwargs json-serializable
|
||||
# we can move this to ValidationTaskConfig.
|
||||
if ray_remote_kwargs is None:
|
||||
self.ray_remote_kwargs = {}
|
||||
else:
|
||||
self.ray_remote_kwargs = ray_remote_kwargs
|
||||
@@ -0,0 +1,209 @@
|
||||
import argparse
|
||||
import tempfile
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchft import (
|
||||
DistributedDataParallel,
|
||||
DistributedSampler,
|
||||
Manager,
|
||||
Optimizer,
|
||||
ProcessGroupGloo,
|
||||
)
|
||||
from torchft.checkpointing.pg_transport import PGTransport
|
||||
|
||||
import ray.train
|
||||
from ray.train import RunConfig, ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
from ray.train.v2.torch.torchft_config import TorchftConfig
|
||||
|
||||
|
||||
class LinearDataset(torch.utils.data.Dataset):
|
||||
"""y = a * x + b"""
|
||||
|
||||
def __init__(self, a, b, size=1000):
|
||||
x = np.arange(0, 10, 10 / size, dtype=np.float32)
|
||||
self.x = torch.from_numpy(x)
|
||||
self.y = torch.from_numpy(a * x + b)
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.x[index, None], self.y[index, None]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.x)
|
||||
|
||||
|
||||
def train_func(config):
|
||||
data_size = config.get("data_size", 1000)
|
||||
batch_size = config.get("batch_size", 4)
|
||||
hidden_size = config.get("hidden_size", 1)
|
||||
lr = config.get("lr", 1e-2)
|
||||
num_steps = config.get("num_steps", 100)
|
||||
num_replicas = config.get("num_replicas", 1)
|
||||
report_interval = config.get("report_interval", 10)
|
||||
error_step = config.get("error_step")
|
||||
error_rank = config.get("error_rank", 0)
|
||||
|
||||
context = ray.train.get_context()
|
||||
world_rank = context.get_world_rank()
|
||||
world_size = context.get_world_size()
|
||||
# Each worker is its own replica group with rank 0.
|
||||
group_rank = 0
|
||||
replica_group_id = world_rank
|
||||
|
||||
# Model and optimizer
|
||||
model = nn.Linear(1, hidden_size)
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
|
||||
loss_fn = nn.MSELoss()
|
||||
|
||||
# torchft process group and checkpoint transport.
|
||||
# Timeouts must be generous enough to re-form the gloo process group after a
|
||||
# replica fails. On loaded CI machines a 5s gloo store wait is too short, which
|
||||
# makes the post-failure reconfigure time out (DistStoreError) and breaks
|
||||
# recovery. Keep these <= the Manager timeout so the PG wait isn't cancelled
|
||||
# by the outer quorum timeout first.
|
||||
pg = ProcessGroupGloo(timeout=timedelta(seconds=30))
|
||||
transport = PGTransport(
|
||||
pg,
|
||||
timeout=timedelta(seconds=30),
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
# State dict callbacks for torchft recovery
|
||||
def load_state_dict(state_dict):
|
||||
model.load_state_dict(state_dict["model"])
|
||||
optimizer.load_state_dict(state_dict["optim"])
|
||||
|
||||
def state_dict():
|
||||
return {
|
||||
"model": model.state_dict(),
|
||||
"optim": optimizer.state_dict(),
|
||||
}
|
||||
|
||||
manager = Manager(
|
||||
pg=pg,
|
||||
min_replica_size=num_replicas,
|
||||
load_state_dict=load_state_dict,
|
||||
state_dict=state_dict,
|
||||
world_size=1,
|
||||
rank=0,
|
||||
replica_id=f"train_ddp_{world_rank}",
|
||||
timeout=timedelta(seconds=60),
|
||||
checkpoint_transport=transport,
|
||||
)
|
||||
|
||||
# Wrap model and optimizer with torchft primitives
|
||||
model = DistributedDataParallel(manager, model)
|
||||
optimizer = Optimizer(manager, optimizer)
|
||||
|
||||
# Data
|
||||
train_dataset = LinearDataset(2, 5, size=data_size)
|
||||
sampler = DistributedSampler(
|
||||
train_dataset,
|
||||
replica_rank=replica_group_id,
|
||||
num_replica_groups=world_size,
|
||||
group_rank=group_rank,
|
||||
num_replicas=1,
|
||||
shuffle=False,
|
||||
)
|
||||
train_loader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=batch_size, sampler=sampler
|
||||
)
|
||||
|
||||
# Training
|
||||
results = []
|
||||
train_iter = iter(train_loader)
|
||||
running_loss = 0.0
|
||||
num_batches = 0
|
||||
|
||||
while manager.current_step() < num_steps:
|
||||
try:
|
||||
X, y = next(train_iter)
|
||||
except StopIteration:
|
||||
train_iter = iter(train_loader)
|
||||
X, y = next(train_iter)
|
||||
|
||||
optimizer.zero_grad()
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
running_loss += loss.item()
|
||||
num_batches += 1
|
||||
|
||||
step = manager.current_step()
|
||||
if error_step is not None and step >= error_step and world_rank == error_rank:
|
||||
marker = Path(
|
||||
ray.train.get_context()
|
||||
.get_storage()
|
||||
.build_checkpoint_path_from_name("error_marker")
|
||||
)
|
||||
if not marker.exists():
|
||||
marker.parent.mkdir(parents=True, exist_ok=True)
|
||||
marker.touch()
|
||||
raise RuntimeError(
|
||||
f"Simulated replica failure at step {step} on rank {world_rank}"
|
||||
)
|
||||
if step % report_interval == 0 or step >= num_steps:
|
||||
avg_loss = running_loss / max(num_batches, 1)
|
||||
weight = model.module.weight.detach().flatten().tolist()
|
||||
bias = model.module.bias.detach().flatten().tolist()
|
||||
result = {"loss": avg_loss, "weight": weight, "bias": bias, "step": step}
|
||||
# TODO(tseah): remove this check once we support reporting with 1/2 workers.
|
||||
if config.get("training_requires_all_workers", True):
|
||||
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
||||
ray.train.report(
|
||||
result,
|
||||
checkpoint=ray.train.Checkpoint.from_directory(
|
||||
temp_checkpoint_dir
|
||||
),
|
||||
)
|
||||
results.append(result)
|
||||
running_loss = 0.0
|
||||
num_batches = 0
|
||||
|
||||
# Needed to avoid "split brain" where worker X dies, worker Y finishes, worker X resumes,
|
||||
# and worker X gets stuck in loss.backward()
|
||||
print(f"Shutting down manager on rank {world_rank}")
|
||||
manager.shutdown()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def train_torchft(num_workers=2, num_steps=100, storage_path=None):
|
||||
config = {
|
||||
"num_steps": num_steps,
|
||||
}
|
||||
trainer = TorchTrainer(
|
||||
train_loop_per_worker=train_func,
|
||||
train_loop_config=config,
|
||||
scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=False),
|
||||
torch_config=TorchftConfig(
|
||||
lighthouse_kwargs={"min_replicas": 1}, backend="gloo"
|
||||
),
|
||||
run_config=RunConfig(storage_path=storage_path),
|
||||
)
|
||||
result = trainer.fit()
|
||||
|
||||
print(result.metrics)
|
||||
return result.metrics
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--num-workers",
|
||||
"-n",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Sets number of workers for training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-steps", type=int, default=100, help="Number of training steps."
|
||||
)
|
||||
|
||||
args, _ = parser.parse_known_args()
|
||||
train_torchft(num_workers=args.num_workers, num_steps=args.num_steps)
|
||||
@@ -0,0 +1 @@
|
||||
import ray.train.horovod # noqa: F401
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.air.config import RunConfig, ScalingConfig
|
||||
from ray.train import Checkpoint, DataConfig
|
||||
from ray.train.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.horovod.config import HorovodConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.util.annotations import Deprecated
|
||||
|
||||
|
||||
@Deprecated
|
||||
class HorovodTrainer(DataParallelTrainer):
|
||||
"""A Trainer for data parallel Horovod training. HorovodTrainer is deprecated."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
horovod_config: Optional[HorovodConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
dataset_config: Optional[DataConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
raise DeprecationWarning(
|
||||
"`HorovodTrainer` is not supported and is scheduled to be removed "
|
||||
"in the future. "
|
||||
"Please consider using `TorchTrainer` or `TensorflowTrainer`, "
|
||||
"fall back to the old implementation with `RAY_TRAIN_V2_ENABLED=0`, "
|
||||
"or file an issue on Github describing your use case."
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
try:
|
||||
import jax # noqa: F401
|
||||
except ModuleNotFoundError as exception:
|
||||
raise ModuleNotFoundError(
|
||||
"Jax isn't installed. To install Jax, please check"
|
||||
" `https://github.com/google/jax#installation` for the instructions."
|
||||
) from exception
|
||||
|
||||
from ray.train.v2.jax.config import JaxConfig
|
||||
from ray.train.v2.jax.jax_trainer import JaxTrainer
|
||||
|
||||
__all__ = ["JaxConfig", "JaxTrainer"]
|
||||
@@ -0,0 +1,227 @@
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import ray
|
||||
from ray._private import ray_constants
|
||||
from ray.train._internal.utils import get_address_and_port
|
||||
from ray.train._internal.worker_group import WorkerGroup
|
||||
from ray.train.backend import Backend, BackendConfig
|
||||
from ray.train.constants import (
|
||||
DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
from ray.train.v2._internal.util import TrainingFramework
|
||||
from ray.util import PublicAPI
|
||||
from ray.util.tpu import get_tpu_coordinator_env_vars, get_tpu_worker_resources
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# JAX multi-slice (megascale) coordination env vars whose values are
|
||||
# authoritatively computed by the controller for the current worker group.
|
||||
# These must always be overwritten on each worker (even when the underlying
|
||||
# TPU node provider has already set them in the pod environment), because
|
||||
# stale values from a previous worker group configuration -- e.g. after a
|
||||
# slice was preempted and replaced -- would otherwise cause libtpu's
|
||||
# megascale topology coordinator to wait for slices that no longer exist
|
||||
# and hang TPU initialization.
|
||||
_JAX_MULTISLICE_OVERRIDE_KEYS = frozenset(
|
||||
{
|
||||
"MEGASCALE_COORDINATOR_ADDRESS",
|
||||
"MEGASCALE_NUM_SLICES",
|
||||
"MEGASCALE_SLICE_ID",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
@dataclass
|
||||
class JaxConfig(BackendConfig):
|
||||
use_tpu: bool = False
|
||||
use_gpu: bool = False
|
||||
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _JaxBackend
|
||||
|
||||
@property
|
||||
def framework(self):
|
||||
return TrainingFramework.JAX
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
config_dict = {
|
||||
"use_tpu": self.use_tpu,
|
||||
"use_gpu": self.use_gpu,
|
||||
}
|
||||
return config_dict
|
||||
|
||||
|
||||
def _setup_jax_distributed_environment(
|
||||
master_addr_with_port: str,
|
||||
num_workers: int,
|
||||
index: int,
|
||||
use_tpu: bool,
|
||||
use_gpu: bool,
|
||||
resources_per_worker: dict,
|
||||
jax_env_vars: Optional[dict] = None,
|
||||
):
|
||||
"""Set up distributed Jax training information.
|
||||
|
||||
This function should be called on each worker. It sets JAX environment
|
||||
variables and initializes JAX distributed training.
|
||||
|
||||
Args:
|
||||
master_addr_with_port: The master address with port for coordination.
|
||||
num_workers: Total number of workers.
|
||||
index: Index of this worker.
|
||||
use_tpu: Whether to configure for TPU. If True and JAX_PLATFORMS is not
|
||||
already set, it will be set to "tpu".
|
||||
use_gpu: Whether to configure for GPU. If True and JAX_PLATFORMS is not
|
||||
already set, it will be set to "cuda".
|
||||
resources_per_worker: The resources per worker.
|
||||
jax_env_vars: The JAX coordinator env vars to inject for multi-slice.
|
||||
Multi-slice coordination keys (``MEGASCALE_*``) always overwrite
|
||||
any pre-existing value in the environment; all other keys are
|
||||
only set if they are not already present.
|
||||
"""
|
||||
# Get JAX_PLATFORMS from environment if already set
|
||||
jax_platforms = os.environ.get("JAX_PLATFORMS", "").lower()
|
||||
|
||||
if not jax_platforms and use_tpu:
|
||||
os.environ["JAX_PLATFORMS"] = "tpu"
|
||||
jax_platforms = "tpu"
|
||||
|
||||
if jax_env_vars:
|
||||
for k, v in jax_env_vars.items():
|
||||
# For multi-slice coordination keys, always override -- the
|
||||
# controller's freshly computed value is the source of truth and
|
||||
# may differ from what the TPU node provider baked into the pod
|
||||
# environment (e.g. after a slice replacement following preemption).
|
||||
# For all other keys, respect any pre-existing value.
|
||||
if k in _JAX_MULTISLICE_OVERRIDE_KEYS or k not in os.environ:
|
||||
os.environ[k] = v
|
||||
|
||||
if not jax_platforms and use_gpu:
|
||||
os.environ["JAX_PLATFORMS"] = "cuda"
|
||||
jax_platforms = "cuda"
|
||||
|
||||
if "cuda" in jax_platforms.split(","):
|
||||
num_gpus_per_worker = resources_per_worker.get("GPU", 0)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
|
||||
str(i) for i in range(num_gpus_per_worker)
|
||||
)
|
||||
|
||||
import jax
|
||||
|
||||
jax_platforms_list = jax_platforms.split(",")
|
||||
|
||||
if "tpu" in jax_platforms_list:
|
||||
jax.distributed.initialize(master_addr_with_port, num_workers, index)
|
||||
logger.info("Initialized JAX distributed on TPU.")
|
||||
elif "cuda" in jax_platforms_list:
|
||||
if num_gpus_per_worker > 0:
|
||||
local_device_ids = list(range(num_gpus_per_worker))
|
||||
else:
|
||||
local_device_ids = 0
|
||||
jax.distributed.initialize(
|
||||
master_addr_with_port, num_workers, index, local_device_ids
|
||||
)
|
||||
logger.info("Initialized JAX distributed on CUDA.")
|
||||
elif "cpu" in jax_platforms_list:
|
||||
jax.distributed.initialize(master_addr_with_port, num_workers, index)
|
||||
logger.info("Initialized JAX distributed on CPU.")
|
||||
|
||||
|
||||
def _shutdown_jax_distributed():
|
||||
"""Shutdown JAX distributed environment.
|
||||
|
||||
This function should be called on each worker during cleanup.
|
||||
If JAX distributed was not initialized, this is a no-op.
|
||||
"""
|
||||
try:
|
||||
import jax
|
||||
|
||||
jax.distributed.shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during JAX distributed shutdown: {e}")
|
||||
|
||||
|
||||
class _JaxBackend(Backend):
|
||||
def on_start(self, worker_group: WorkerGroup, backend_config: JaxConfig):
|
||||
if not backend_config.use_tpu and not backend_config.use_gpu:
|
||||
return
|
||||
|
||||
master_addr, master_port = worker_group.execute_single(0, get_address_and_port)
|
||||
master_addr_with_port = f"{master_addr}:{master_port}"
|
||||
|
||||
if backend_config.use_tpu and hasattr(worker_group, "get_worker_group_context"):
|
||||
num_slices = worker_group.get_worker_group_context().num_slices
|
||||
else:
|
||||
num_slices = 1
|
||||
|
||||
# Calculate the number of workers per slice for multi-slice env setup.
|
||||
if backend_config.use_tpu and num_slices > 1:
|
||||
# Handle the case where a user requests less workers than the total
|
||||
# capacity of the TPU slice.
|
||||
scaling_config = worker_group._train_run_context.scaling_config
|
||||
workers_per_slice, _ = get_tpu_worker_resources(
|
||||
topology=scaling_config.topology,
|
||||
accelerator_type=scaling_config.accelerator_type,
|
||||
resources_per_unit=scaling_config.resources_per_worker,
|
||||
num_slices=1,
|
||||
)
|
||||
else:
|
||||
# Assume even distribution based on the requested number of workers.
|
||||
workers_per_slice = max(1, len(worker_group) // num_slices)
|
||||
|
||||
# Set up JAX distributed environment on all workers
|
||||
num_workers_total = len(worker_group)
|
||||
setup_futures = []
|
||||
for i in range(num_workers_total):
|
||||
env_vars = {}
|
||||
if num_slices > 1:
|
||||
slice_id = min(i // workers_per_slice, num_slices - 1)
|
||||
env_vars = get_tpu_coordinator_env_vars(
|
||||
coordinator_address=master_addr,
|
||||
num_slices=num_slices,
|
||||
slice_id=slice_id,
|
||||
)
|
||||
|
||||
setup_futures.append(
|
||||
worker_group.execute_single_async(
|
||||
i,
|
||||
_setup_jax_distributed_environment,
|
||||
master_addr_with_port=master_addr_with_port,
|
||||
num_workers=len(worker_group),
|
||||
index=i,
|
||||
use_tpu=backend_config.use_tpu,
|
||||
use_gpu=backend_config.use_gpu,
|
||||
resources_per_worker=worker_group.get_resources_per_worker(),
|
||||
jax_env_vars=env_vars,
|
||||
)
|
||||
)
|
||||
ray.get(setup_futures)
|
||||
|
||||
def on_shutdown(self, worker_group: WorkerGroup, backend_config: JaxConfig):
|
||||
"""Cleanup JAX distributed resources when shutting down worker group."""
|
||||
if not backend_config.use_tpu and not backend_config.use_gpu:
|
||||
return
|
||||
|
||||
# Shutdown JAX distributed on all workers
|
||||
shutdown_futures = worker_group.execute_async(_shutdown_jax_distributed)
|
||||
|
||||
timeout_s = ray_constants.env_integer(
|
||||
JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
DEFAULT_JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
|
||||
)
|
||||
try:
|
||||
ray.get(shutdown_futures, timeout=timeout_s)
|
||||
logger.debug("JAX distributed shutdown completed")
|
||||
except ray.exceptions.GetTimeoutError:
|
||||
logger.warning(
|
||||
f"JAX distributed shutdown timed out after {timeout_s} seconds. "
|
||||
"This may indicate workers are hung or unresponsive."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during JAX distributed shutdown: {e}")
|
||||
@@ -0,0 +1,169 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
|
||||
from ray.train import DataConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
||||
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.train.v2.jax.config import JaxConfig
|
||||
from ray.util import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class JaxTrainer(DataParallelTrainer):
|
||||
"""A Trainer for Single-Program Multi-Data (SPMD) JAX training.
|
||||
|
||||
At a high level, this Trainer does the following:
|
||||
|
||||
1. Launches multiple workers as defined by the ``scaling_config``.
|
||||
2. Sets up a distributed JAX environment for TPUs or GPUs
|
||||
on these workers as defined by the ``jax_config``.
|
||||
3. Ingests the input ``datasets`` based on the ``dataset_config``.
|
||||
4. Runs the input ``train_loop_per_worker(train_loop_config)``
|
||||
on all workers.
|
||||
|
||||
For more details, see:
|
||||
|
||||
* :ref:`Jax Guide <train-jax>`
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import os
|
||||
from absl import app
|
||||
import logging
|
||||
from typing import Sequence
|
||||
|
||||
import ray
|
||||
from ray.train import ScalingConfig, RunConfig
|
||||
from ray.train.v2.jax import JaxTrainer
|
||||
from MaxText.train import main as maxtext_main
|
||||
|
||||
def train_loop_per_worker(config):
|
||||
argv = config["argv"]
|
||||
maxtext_main(argv)
|
||||
|
||||
def main(argv: Sequence[str]):
|
||||
ray.init()
|
||||
|
||||
# If you want to use TPUs, specify the TPU topology and accelerator type.
|
||||
tpu_scaling_config = ScalingConfig(
|
||||
use_tpu=True,
|
||||
num_workers=4,
|
||||
topology="4x4",
|
||||
accelerator_type="TPU-V6E",
|
||||
placement_strategy="SPREAD",
|
||||
resources_per_worker={"TPU": 4},
|
||||
)
|
||||
|
||||
# If you want to use GPUs, specify the GPU scaling config like below.
|
||||
# gpu_scaling_config = ScalingConfig(
|
||||
# use_gpu=True,
|
||||
# num_workers=4,
|
||||
# resources_per_worker={"GPU": 1},
|
||||
# )
|
||||
|
||||
|
||||
trainer = JaxTrainer(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config={"argv": absolute_argv},
|
||||
scaling_config=tpu_scaling_config,
|
||||
run_config=RunConfig(
|
||||
name="maxtext_jaxtrainer",
|
||||
worker_runtime_env={
|
||||
"env_vars": {
|
||||
"JAX_PLATFORMS": "tpu",
|
||||
# If you want to use GPUs, set the JAX_PLATFORMS to "cuda".
|
||||
# "JAX_PLATFORMS": "cuda",
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
result = trainer.fit()
|
||||
|
||||
If the ``datasets`` dict contains datasets (e.g. "train" and "val"), then it will be split into multiple dataset
|
||||
shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` and ``ray.train.get_dataset_shard("val")``.
|
||||
|
||||
Note:
|
||||
* If you are using TPUs, importing `jax` should occur within `train_loop_per_worker` to
|
||||
avoid driver-side TPU lock issues.
|
||||
|
||||
Args:
|
||||
train_loop_per_worker: The training function to execute on each worker.
|
||||
This function can either take in zero arguments or a single ``Dict``
|
||||
argument which is set by defining ``train_loop_config``.
|
||||
Within this function you can use any of the
|
||||
:ref:`Ray Train Loop utilities <train-loop-api>`.
|
||||
train_loop_config: A configuration ``Dict`` to pass in as an argument to
|
||||
``train_loop_per_worker``.
|
||||
This is typically used for specifying hyperparameters. Passing large
|
||||
datasets via `train_loop_config` is not recommended and may introduce
|
||||
large overhead and unknown issues with serialization and deserialization.
|
||||
jax_config: The configuration for setting up the JAX backend.
|
||||
If set to None, a default configuration will be used based on the ``scaling_config`` and ``JAX_PLATFORMS`` environment variable.
|
||||
scaling_config: Configuration for how to scale data parallel training
|
||||
with SPMD. ``num_workers`` should be set to the number of TPU hosts or GPU workers.
|
||||
If using TPUs, ``topology`` should be set to the TPU topology.
|
||||
See :class:`~ray.train.ScalingConfig` for more info.
|
||||
dataset_config: The configuration for ingesting the input ``datasets``.
|
||||
By default, all the Ray Dataset are split equally across workers.
|
||||
See :class:`~ray.train.DataConfig` for more details.
|
||||
run_config: The configuration for the execution of the training run.
|
||||
See :class:`~ray.train.RunConfig` for more info.
|
||||
datasets: The Ray Datasets to ingest for training.
|
||||
Datasets are keyed by name (``{name: dataset}``).
|
||||
Each dataset can be accessed from within the ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_dataset_shard(name)``.
|
||||
Sharding and additional configuration can be done by
|
||||
passing in a ``dataset_config``.
|
||||
validation_config: [Alpha] Configuration for checkpoint validation.
|
||||
If provided and ``ray.train.report`` is called with the ``validation``
|
||||
argument, Ray Train will validate the reported checkpoint using
|
||||
the validation function specified in this config.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
jax_config: Optional[JaxConfig] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
dataset_config: Optional[Dict[str, DataConfig]] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
if not jax_config:
|
||||
jax_config = JaxConfig(
|
||||
use_tpu=scaling_config.use_tpu,
|
||||
use_gpu=scaling_config.use_gpu,
|
||||
)
|
||||
super(JaxTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=jax_config,
|
||||
scaling_config=scaling_config,
|
||||
dataset_config=dataset_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _validate_scaling_config(cls, scaling_config: ScalingConfig) -> ScalingConfig:
|
||||
"""Return scaling config dataclass after validating updated keys."""
|
||||
ensure_only_allowed_dataclass_keys_updated(
|
||||
dataclass=scaling_config,
|
||||
allowed_keys=cls._scaling_config_allowed_keys,
|
||||
)
|
||||
|
||||
return scaling_config
|
||||
@@ -0,0 +1,184 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
|
||||
|
||||
import ray.train
|
||||
from ray.train import Checkpoint
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
||||
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.util.annotations import Deprecated
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.lightgbm import LightGBMConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LightGBMTrainer(DataParallelTrainer):
|
||||
"""A Trainer for distributed data-parallel LightGBM training.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
import lightgbm as lgb
|
||||
|
||||
import ray.data
|
||||
import ray.train
|
||||
from ray.train.lightgbm import (
|
||||
LightGBMTrainer,
|
||||
RayTrainReportCallback,
|
||||
normalize_pandas_for_lightgbm,
|
||||
)
|
||||
|
||||
|
||||
def train_fn_per_worker(config: dict):
|
||||
# (Optional) Add logic to resume training state from a checkpoint.
|
||||
# ray.train.get_checkpoint()
|
||||
|
||||
# 1. Get the dataset shard for the worker and convert to a `lgb.Dataset`
|
||||
train_ds_iter, eval_ds_iter = (
|
||||
ray.train.get_dataset_shard("train"),
|
||||
ray.train.get_dataset_shard("validation"),
|
||||
)
|
||||
train_ds, eval_ds = train_ds_iter.materialize(), eval_ds_iter.materialize()
|
||||
train_df = normalize_pandas_for_lightgbm(train_ds.to_pandas())
|
||||
eval_df = normalize_pandas_for_lightgbm(eval_ds.to_pandas())
|
||||
train_X, train_y = train_df.drop("y", axis=1), train_df["y"]
|
||||
eval_X, eval_y = eval_df.drop("y", axis=1), eval_df["y"]
|
||||
|
||||
train_set = lgb.Dataset(train_X, label=train_y)
|
||||
eval_set = lgb.Dataset(eval_X, label=eval_y)
|
||||
|
||||
# 2. Run distributed data-parallel training.
|
||||
# `get_network_params` sets up the necessary configurations for LightGBM
|
||||
# to set up the data parallel training worker group on your Ray cluster.
|
||||
params = {
|
||||
"objective": "regression",
|
||||
# Adding the lines below are the only changes needed
|
||||
# for your `lgb.train` call!
|
||||
"tree_learner": "data_parallel",
|
||||
"pre_partition": True,
|
||||
**ray.train.lightgbm.get_network_params(),
|
||||
}
|
||||
lgb.train(
|
||||
params,
|
||||
train_set,
|
||||
valid_sets=[eval_set],
|
||||
valid_names=["eval"],
|
||||
num_boost_round=1,
|
||||
# To access the checkpoint from trainer, you need this callback.
|
||||
callbacks=[RayTrainReportCallback()],
|
||||
)
|
||||
|
||||
train_ds = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
|
||||
eval_ds = ray.data.from_items(
|
||||
[{"x": x, "y": x + 1} for x in range(32, 32 + 16)]
|
||||
)
|
||||
trainer = LightGBMTrainer(
|
||||
train_fn_per_worker,
|
||||
datasets={"train": train_ds, "validation": eval_ds},
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
result = trainer.fit()
|
||||
booster = RayTrainReportCallback.get_model(result.checkpoint)
|
||||
|
||||
Args:
|
||||
train_loop_per_worker: The training function to execute on each worker.
|
||||
This function can either take in zero arguments or a single ``Dict``
|
||||
argument which is set by defining ``train_loop_config``.
|
||||
Within this function you can use any of the
|
||||
:ref:`Ray Train Loop utilities <train-loop-api>`.
|
||||
train_loop_config: A configuration ``Dict`` to pass in as an argument to
|
||||
``train_loop_per_worker``.
|
||||
This is typically used for specifying hyperparameters.
|
||||
lightgbm_config: The configuration for setting up the distributed lightgbm
|
||||
backend. See :class:`~ray.train.lightgbm.LightGBMConfig` for more info.
|
||||
scaling_config: The configuration for how to scale data parallel training.
|
||||
``num_workers`` determines how many Python processes are used for training,
|
||||
and ``use_gpu`` determines whether or not each process should use GPUs.
|
||||
See :class:`~ray.train.ScalingConfig` for more info.
|
||||
run_config: The configuration for the execution of the training run.
|
||||
See :class:`~ray.train.RunConfig` for more info.
|
||||
datasets: The Ray Datasets to ingest for training.
|
||||
Datasets are keyed by name (``{name: dataset}``).
|
||||
Each dataset can be accessed from within the ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_dataset_shard(name)``.
|
||||
Sharding and additional configuration can be done by
|
||||
passing in a ``dataset_config``.
|
||||
dataset_config: The configuration for ingesting the input ``datasets``.
|
||||
By default, all the Ray Dataset are split equally across workers.
|
||||
See :class:`~ray.train.DataConfig` for more details.
|
||||
validation_config: [Alpha] Configuration for checkpoint validation.
|
||||
If provided and ``ray.train.report`` is called with the ``validation``
|
||||
argument, Ray Train will validate the reported checkpoint using
|
||||
the validation function specified in this config.
|
||||
metadata: [Deprecated]
|
||||
resume_from_checkpoint: [Deprecated]
|
||||
label_column: [Deprecated] Legacy LightGBMTrainer API.
|
||||
params: [Deprecated] Legacy LightGBMTrainer API.
|
||||
num_boost_round: [Deprecated] Legacy LightGBMTrainer API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
lightgbm_config: Optional["LightGBMConfig"] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
dataset_config: Optional[ray.train.DataConfig] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
# TODO: [Deprecated]
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
# TODO: [Deprecated] Legacy LightGBMTrainer API
|
||||
label_column: Optional[str] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
num_boost_round: Optional[int] = None,
|
||||
):
|
||||
if (
|
||||
label_column is not None
|
||||
or params is not None
|
||||
or num_boost_round is not None
|
||||
):
|
||||
raise DeprecationWarning(
|
||||
"The legacy LightGBMTrainer API is deprecated. "
|
||||
"Please switch to passing in a custom `train_loop_per_worker` "
|
||||
"function instead. "
|
||||
"See this issue for more context: "
|
||||
"https://github.com/ray-project/ray/issues/50042"
|
||||
)
|
||||
from ray.train.lightgbm import LightGBMConfig
|
||||
|
||||
super(LightGBMTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=lightgbm_config or LightGBMConfig(),
|
||||
scaling_config=scaling_config,
|
||||
dataset_config=dataset_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
resume_from_checkpoint=resume_from_checkpoint,
|
||||
metadata=metadata,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@Deprecated
|
||||
def get_model(
|
||||
cls,
|
||||
checkpoint: Checkpoint,
|
||||
):
|
||||
"""Retrieve the LightGBM model stored in this checkpoint.
|
||||
|
||||
This API is deprecated. Use `RayTrainReportCallback.get_model` instead.
|
||||
"""
|
||||
raise DeprecationWarning(
|
||||
"`LightGBMTrainer.get_model` is deprecated. "
|
||||
"Use `RayTrainReportCallback.get_model` instead."
|
||||
)
|
||||
@@ -0,0 +1,199 @@
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
|
||||
|
||||
from ray.train import Checkpoint, DataConfig
|
||||
from ray.train.trainer import GenDataset
|
||||
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
||||
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
from ray.util import PublicAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.tensorflow import TensorflowConfig
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
class TensorflowTrainer(DataParallelTrainer):
|
||||
"""A Trainer for data parallel Tensorflow training.
|
||||
|
||||
At a high level, this Trainer does the following:
|
||||
|
||||
1. Launches multiple workers as defined by the ``scaling_config``.
|
||||
2. Sets up a distributed Tensorflow environment
|
||||
on these workers as defined by the ``tensorflow_config``.
|
||||
3. Ingests the input ``datasets`` based on the ``dataset_config``.
|
||||
4. Runs the input ``train_loop_per_worker(train_loop_config)``
|
||||
on all workers.
|
||||
|
||||
For more details, see:
|
||||
|
||||
* :ref:`Tensorflow Guide <train-tensorflow-overview>`
|
||||
|
||||
Inside the ``train_loop_per_worker`` function, you can use any of the
|
||||
:ref:`Ray Train loop methods <train-loop-api>`.
|
||||
|
||||
.. warning::
|
||||
Ray will not automatically set any environment variables or configuration
|
||||
related to local parallelism / threading
|
||||
:ref:`aside from "OMP_NUM_THREADS" <omp-num-thread-note>`.
|
||||
If you desire greater control over TensorFlow threading, use
|
||||
the ``tf.config.threading`` module (eg.
|
||||
``tf.config.threading.set_inter_op_parallelism_threads(num_cpus)``)
|
||||
at the beginning of your ``train_loop_per_worker`` function.
|
||||
|
||||
|
||||
.. testcode::
|
||||
|
||||
from ray import train
|
||||
|
||||
def train_loop_per_worker():
|
||||
# Report intermediate results for callbacks or logging and
|
||||
# checkpoint data.
|
||||
train.report(...)
|
||||
|
||||
# Returns dict of last saved checkpoint.
|
||||
train.get_checkpoint()
|
||||
|
||||
# Returns the Dataset shard for the given key.
|
||||
train.get_dataset_shard("my_dataset")
|
||||
|
||||
# Returns the total number of workers executing training.
|
||||
train.get_context().get_world_size()
|
||||
|
||||
# Returns the rank of this worker.
|
||||
train.get_context().get_world_rank()
|
||||
|
||||
# Returns the rank of the worker on the current node.
|
||||
train.get_context().get_local_rank()
|
||||
|
||||
Any returns from the ``train_loop_per_worker`` will be discarded and not
|
||||
used or persisted anywhere.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import tensorflow as tf
|
||||
|
||||
import ray
|
||||
from ray import train
|
||||
from ray.train import Checkpoint, ScalingConfig
|
||||
from ray.train.tensorflow import TensorflowTrainer
|
||||
|
||||
def build_model():
|
||||
# toy neural network : 1-layer
|
||||
return tf.keras.Sequential(
|
||||
[tf.keras.layers.Dense(
|
||||
1, activation="linear", input_shape=(1,))]
|
||||
)
|
||||
|
||||
def train_loop_per_worker(config):
|
||||
dataset_shard = train.get_dataset_shard("train")
|
||||
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
|
||||
with strategy.scope():
|
||||
model = build_model()
|
||||
model.compile(
|
||||
optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
|
||||
|
||||
tf_dataset = dataset_shard.to_tf(
|
||||
feature_columns="x",
|
||||
label_columns="y",
|
||||
batch_size=1,
|
||||
)
|
||||
for epoch in range(config["num_epochs"]):
|
||||
model.fit(tf_dataset)
|
||||
|
||||
# Create checkpoint.
|
||||
checkpoint_dir = tempfile.mkdtemp()
|
||||
model.save_weights(
|
||||
os.path.join(checkpoint_dir, "my_checkpoint")
|
||||
)
|
||||
checkpoint = Checkpoint.from_directory(checkpoint_dir)
|
||||
|
||||
train.report(
|
||||
{},
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
|
||||
trainer = TensorflowTrainer(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
scaling_config=ScalingConfig(num_workers=3, use_gpu=False),
|
||||
datasets={"train": train_dataset},
|
||||
train_loop_config={"num_epochs": 2},
|
||||
)
|
||||
result = trainer.fit()
|
||||
|
||||
.. testoutput::
|
||||
:options:+ELLIPSIS
|
||||
:hide:
|
||||
|
||||
...
|
||||
|
||||
Args:
|
||||
train_loop_per_worker: The training function to execute on each worker.
|
||||
This function can either take in zero arguments or a single ``Dict``
|
||||
argument which is set by defining ``train_loop_config``.
|
||||
Within this function you can use any of the
|
||||
:ref:`Ray Train Loop utilities <train-loop-api>`.
|
||||
train_loop_config: A configuration ``Dict`` to pass in as an argument to
|
||||
``train_loop_per_worker``.
|
||||
This is typically used for specifying hyperparameters. Passing large
|
||||
datasets via `train_loop_config` is not recommended and may introduce
|
||||
large overhead and unknown issues with serialization and deserialization.
|
||||
tensorflow_config: The configuration for setting up the Tensorflow
|
||||
Distributed backend. If set to None, a default configuration will be
|
||||
used in which GPU training uses NCCL and CPU training uses Gloo.
|
||||
scaling_config: The configuration for how to scale data parallel training.
|
||||
``num_workers`` determines how many Python processes are used for training,
|
||||
and ``use_gpu`` determines whether or not each process should use GPUs.
|
||||
See :class:`~ray.train.ScalingConfig` for more info.
|
||||
dataset_config: The configuration for ingesting the input ``datasets``.
|
||||
By default, all the Ray Datasets are split equally across workers.
|
||||
See :class:`~ray.train.DataConfig` for more details.
|
||||
run_config: The configuration for the execution of the training run.
|
||||
See :class:`~ray.train.RunConfig` for more info.
|
||||
datasets: The Ray Datasets to ingest for training.
|
||||
Datasets are keyed by name (``{name: dataset}``).
|
||||
Each dataset can be accessed from within the ``train_loop_per_worker``
|
||||
by calling ``ray.train.get_dataset_shard(name)``.
|
||||
Sharding and additional configuration can be done by
|
||||
passing in a ``dataset_config``.
|
||||
validation_config: [Alpha] Configuration for checkpoint validation.
|
||||
If provided and ``ray.train.report`` is called with the ``validation``
|
||||
argument, Ray Train will validate the reported checkpoint using
|
||||
the validation function specified in this config.
|
||||
metadata: [Deprecated]
|
||||
resume_from_checkpoint: [Deprecated]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_loop_per_worker: Union[Callable[[], Any], Callable[[Dict], Any]],
|
||||
*,
|
||||
train_loop_config: Optional[Dict] = None,
|
||||
tensorflow_config: Optional["TensorflowConfig"] = None,
|
||||
scaling_config: Optional[ScalingConfig] = None,
|
||||
dataset_config: Optional[DataConfig] = None,
|
||||
run_config: Optional[RunConfig] = None,
|
||||
datasets: Optional[Dict[str, GenDataset]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
# TODO: [Deprecated]
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
resume_from_checkpoint: Optional[Checkpoint] = None,
|
||||
):
|
||||
from ray.train.tensorflow import TensorflowConfig
|
||||
|
||||
super(TensorflowTrainer, self).__init__(
|
||||
train_loop_per_worker=train_loop_per_worker,
|
||||
train_loop_config=train_loop_config,
|
||||
backend_config=tensorflow_config or TensorflowConfig(),
|
||||
scaling_config=scaling_config,
|
||||
dataset_config=dataset_config,
|
||||
run_config=run_config,
|
||||
datasets=datasets,
|
||||
resume_from_checkpoint=resume_from_checkpoint,
|
||||
metadata=metadata,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
@@ -0,0 +1,90 @@
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import runtime_context
|
||||
from ray._common import utils as ray_utils
|
||||
from ray.cluster_utils import Cluster
|
||||
from ray.train.v2._internal.constants import (
|
||||
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def ray_start_4_cpus():
|
||||
ray.init(num_cpus=4)
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def ray_start_4_cpus_2_gpus():
|
||||
ray.init(num_cpus=4, num_gpus=2)
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_2x2_gpu_cluster():
|
||||
cluster = Cluster()
|
||||
for _ in range(2):
|
||||
cluster.add_node(num_cpus=4, num_gpus=2)
|
||||
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
yield
|
||||
|
||||
ray.shutdown()
|
||||
cluster.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_logging():
|
||||
logger = logging.getLogger("ray.train")
|
||||
orig_level = logger.getEffectiveLevel()
|
||||
logger.setLevel(logging.INFO)
|
||||
yield
|
||||
logger.setLevel(orig_level)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_ray_address(monkeypatch):
|
||||
ray_utils.reset_ray_address()
|
||||
yield
|
||||
ray_utils.reset_ray_address()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shutdown_only():
|
||||
yield None
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def disable_state_actor_polling(monkeypatch):
|
||||
monkeypatch.setenv(ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR, "0")
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_runtime_context(monkeypatch):
|
||||
@ray.remote(num_cpus=0)
|
||||
class DummyActor:
|
||||
pass
|
||||
|
||||
# Must return real actor handle so it can get passed to other actors
|
||||
# Cannot create actor here since ray has not been initialized yet
|
||||
def mock_current_actor(self):
|
||||
return DummyActor.remote()
|
||||
|
||||
# In unit tests where the controller is not an actor, current_actor is
|
||||
# a DummyActor, which is ok because it won't be called in those tests.
|
||||
# In unit tests where the controller is an actor, current_actor is the
|
||||
# controller actor because monkeypatch doesn't propagate to the actor
|
||||
# process. Those tests can successfully test methods on that actor.
|
||||
monkeypatch.setattr(
|
||||
runtime_context.RuntimeContext, "current_actor", property(mock_current_actor)
|
||||
)
|
||||
|
||||
yield
|
||||
@@ -0,0 +1,143 @@
|
||||
import collections
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.cluster_utils import Cluster
|
||||
from ray.train import BackendConfig
|
||||
from ray.train.backend import Backend
|
||||
from ray.train.v2._internal.callbacks.accelerators import (
|
||||
AcceleratorSetupCallback,
|
||||
_get_visible_accelerator_ids_per_worker,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import ActorMetadata, WorkerGroup
|
||||
from ray.train.v2._internal.execution.worker_group.worker_group import (
|
||||
WorkerGroupContext,
|
||||
)
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper
|
||||
from ray.train.v2.api.config import ScalingConfig
|
||||
from ray.train.v2.tests.util import create_dummy_run_context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_gpu_cluster():
|
||||
"""Yields a GPU cluster with 3 nodes (4 GPU, 1 GPU, 1 GPU)."""
|
||||
cluster = Cluster()
|
||||
cluster.add_node(num_gpus=4)
|
||||
cluster.add_node(num_gpus=1)
|
||||
cluster.add_node(num_gpus=1)
|
||||
cluster.wait_for_nodes()
|
||||
cluster.connect()
|
||||
yield cluster
|
||||
ray.shutdown()
|
||||
cluster.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"node_ids, accelerator_ids_per_worker, expected",
|
||||
[
|
||||
(["0"], [[0]], ["0"]),
|
||||
(
|
||||
["0", "0", "1"],
|
||||
[[0, 1], [2, 3], [0, 1]],
|
||||
["0,1,2,3", "0,1,2,3", "0,1"],
|
||||
),
|
||||
(
|
||||
["0", "0", "1", "1", "1", "1"],
|
||||
[["1"], ["3"], ["3"], ["0"], ["1"], ["2"]],
|
||||
["1,3", "1,3", "0,1,2,3", "0,1,2,3", "0,1,2,3", "0,1,2,3"],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_visible_accelerator_ids_per_worker(
|
||||
node_ids, accelerator_ids_per_worker, expected
|
||||
):
|
||||
worker_metadatas = [
|
||||
ActorMetadata(
|
||||
hostname=node_id,
|
||||
node_id=node_id,
|
||||
node_ip=node_id,
|
||||
pid=0,
|
||||
accelerator_ids={"GPU": accelerator_ids},
|
||||
)
|
||||
for node_id, accelerator_ids in zip(node_ids, accelerator_ids_per_worker)
|
||||
]
|
||||
|
||||
assert (
|
||||
_get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas=worker_metadatas, accelerator_name="GPU"
|
||||
)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
def test_missing_accelerator():
|
||||
"""Trying to share accelerator ids on a heterogeneous worker group
|
||||
(where some workers do not have access to certain accelerators)
|
||||
should raise an error."""
|
||||
with pytest.raises(ValueError):
|
||||
_get_visible_accelerator_ids_per_worker(
|
||||
worker_metadatas=[
|
||||
ActorMetadata(
|
||||
hostname="0",
|
||||
node_id="0",
|
||||
node_ip="0",
|
||||
pid=0,
|
||||
accelerator_ids={"GPU": [0]},
|
||||
),
|
||||
ActorMetadata(
|
||||
hostname="0",
|
||||
node_id="0",
|
||||
node_ip="0",
|
||||
pid=0,
|
||||
accelerator_ids={},
|
||||
),
|
||||
],
|
||||
accelerator_name="GPU",
|
||||
)
|
||||
|
||||
|
||||
def test_accelerator_setup_callback(mock_gpu_cluster, mock_runtime_context):
|
||||
"""The accelerator setup callback should set the CUDA_VISIBLE_DEVICES
|
||||
on each worker properly."""
|
||||
|
||||
class DummyBackendConfig(BackendConfig):
|
||||
def backend_cls(self):
|
||||
return DummyBackend
|
||||
|
||||
class DummyBackend(Backend):
|
||||
share_cuda_visible_devices = True
|
||||
|
||||
scaling_config = ScalingConfig(num_workers=6, use_gpu=True)
|
||||
setup_callback = AcceleratorSetupCallback(
|
||||
backend_config=DummyBackendConfig(),
|
||||
scaling_config=scaling_config,
|
||||
)
|
||||
|
||||
worker_group_context = WorkerGroupContext(
|
||||
run_attempt_id="attempt_1",
|
||||
train_fn_ref=ObjectRefWrapper(lambda: None),
|
||||
num_workers=scaling_config.num_workers,
|
||||
resources_per_worker=scaling_config._resources_per_worker_not_none,
|
||||
)
|
||||
|
||||
worker_group = WorkerGroup(
|
||||
train_run_context=create_dummy_run_context(),
|
||||
worker_group_context=worker_group_context,
|
||||
)
|
||||
|
||||
worker_group._start()
|
||||
|
||||
setup_callback.before_init_train_context(worker_group.get_workers())
|
||||
|
||||
visible_devices_per_worker = worker_group.execute(
|
||||
lambda: os.environ["CUDA_VISIBLE_DEVICES"]
|
||||
)
|
||||
assert collections.Counter(visible_devices_per_worker) == {"0,1,2,3": 4, "0": 2}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,86 @@
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.train.v2._internal.execution.callback import ControllerCallback
|
||||
from ray.train.v2._internal.execution.callback_manager import CallbackManager
|
||||
from ray.train.v2.api.exceptions import ControllerError
|
||||
|
||||
|
||||
def test_invoke_callback_without_hook():
|
||||
"""Test that callbacks without the hook method raise an error."""
|
||||
|
||||
class CallbackWithoutHook:
|
||||
pass
|
||||
|
||||
manager = CallbackManager([CallbackWithoutHook()])
|
||||
with pytest.raises(ControllerError) as exc_info:
|
||||
manager.invoke("test_hook", "arg1")
|
||||
assert isinstance(exc_info.value.controller_failure, AttributeError)
|
||||
|
||||
|
||||
def test_invoke_multiple_callbacks_all_succeed():
|
||||
"""Test that multiple callbacks are invoked in sequence."""
|
||||
callback1 = MagicMock()
|
||||
callback1.test_hook = MagicMock()
|
||||
callback2 = MagicMock()
|
||||
callback2.test_hook = MagicMock()
|
||||
|
||||
manager = CallbackManager([callback1, callback2])
|
||||
result = manager.invoke("test_hook", "arg")
|
||||
|
||||
callback1.test_hook.assert_called_once_with("arg")
|
||||
callback2.test_hook.assert_called_once_with("arg")
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_invoke_with_real_controller_callback_error_returned():
|
||||
"""Test with a real ControllerCallback implementation that raises an error."""
|
||||
|
||||
class MockControllerCallback(ControllerCallback):
|
||||
def __init__(self):
|
||||
self.called = False
|
||||
|
||||
def after_controller_start(self, train_run_context):
|
||||
self.called = True
|
||||
raise ValueError("Intentional error")
|
||||
|
||||
callback = MockControllerCallback()
|
||||
|
||||
manager = CallbackManager([callback])
|
||||
train_run_context = MagicMock()
|
||||
with pytest.raises(ControllerError) as exc_info:
|
||||
manager.invoke("after_controller_start", train_run_context)
|
||||
|
||||
assert callback.called is True
|
||||
assert isinstance(exc_info.value.controller_failure, ValueError)
|
||||
|
||||
|
||||
def test_invoke_callback_error_returns_controller_error():
|
||||
callback = MagicMock()
|
||||
callback.test_hook = MagicMock(side_effect=ValueError("Original hook error"))
|
||||
|
||||
manager = CallbackManager([callback])
|
||||
with pytest.raises(ControllerError) as exc_info:
|
||||
manager.invoke("test_hook", "arg1", key1="value1")
|
||||
|
||||
assert isinstance(exc_info.value.controller_failure, ValueError)
|
||||
|
||||
|
||||
def test_invoke_best_effort_calls_all_callbacks_even_on_failure():
|
||||
"""invoke_best_effort continues calling remaining callbacks after failure."""
|
||||
callback1 = MagicMock()
|
||||
callback1.test_hook = MagicMock(side_effect=ValueError("fail"))
|
||||
callback2 = MagicMock()
|
||||
callback2.test_hook = MagicMock()
|
||||
|
||||
manager = CallbackManager([callback1, callback2])
|
||||
manager.invoke_best_effort("test_hook", "arg")
|
||||
|
||||
callback1.test_hook.assert_called_once_with("arg")
|
||||
callback2.test_hook.assert_called_once_with("arg")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
@@ -0,0 +1,359 @@
|
||||
import uuid
|
||||
from typing import Optional
|
||||
from unittest.mock import create_autospec
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.train import CheckpointConfig
|
||||
from ray.train._internal.session import _TrainingResult
|
||||
from ray.train.v2._internal.exceptions import CheckpointManagerInitializationError
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.storage import StorageContext
|
||||
from ray.train.v2._internal.execution.worker_group import Worker
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
|
||||
from ray.train.v2.tests.util import create_dummy_training_reports
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def ray_start_4_cpus():
|
||||
ray.init(num_cpus=4)
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def _checkpoint_managers_equal(cm1: CheckpointManager, cm2: CheckpointManager) -> bool:
|
||||
"""
|
||||
Compare two checkpoint managers for equality.
|
||||
Ignore uuid differences of all the checkpoints recorded.
|
||||
"""
|
||||
|
||||
def _training_results_equal(
|
||||
tr1: Optional[_TrainingResult], tr2: Optional[_TrainingResult]
|
||||
) -> bool:
|
||||
if not tr1 and not tr2:
|
||||
return True
|
||||
if not tr1 or not tr2:
|
||||
return False
|
||||
return (
|
||||
tr1.metrics == tr2.metrics
|
||||
and tr1.checkpoint.path == tr2.checkpoint.path
|
||||
and tr1.checkpoint.filesystem == tr2.checkpoint.filesystem
|
||||
)
|
||||
|
||||
def _checkpoint_to_report_indices_equal(
|
||||
cm1: CheckpointManager, cm2: CheckpointManager
|
||||
) -> bool:
|
||||
# Do this because Checkpoint and Filesystem are not hashable.
|
||||
cm1_dict = {
|
||||
checkpoint.path: report_index
|
||||
for checkpoint, report_index in cm1._checkpoint_to_report_index.items()
|
||||
}
|
||||
cm2_dict = {
|
||||
checkpoint.path: report_index
|
||||
for checkpoint, report_index in cm2._checkpoint_to_report_index.items()
|
||||
}
|
||||
return cm1_dict == cm2_dict
|
||||
|
||||
if cm1._checkpoint_config != cm2._checkpoint_config:
|
||||
return False
|
||||
if not _training_results_equal(
|
||||
cm1.latest_checkpoint_result, cm2.latest_checkpoint_result
|
||||
):
|
||||
return False
|
||||
if not _training_results_equal(
|
||||
cm1.best_checkpoint_result, cm2.best_checkpoint_result
|
||||
):
|
||||
return False
|
||||
if len(cm1.best_checkpoint_results) != len(cm2.best_checkpoint_results):
|
||||
return False
|
||||
for tr1, tr2 in zip(cm1.best_checkpoint_results, cm2.best_checkpoint_results):
|
||||
if not _training_results_equal(tr1, tr2):
|
||||
return False
|
||||
if cm1._current_report_index != cm2._current_report_index:
|
||||
return False
|
||||
if not _checkpoint_to_report_indices_equal(cm1, cm2):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"checkpoint_config",
|
||||
[
|
||||
CheckpointConfig(),
|
||||
CheckpointConfig(
|
||||
num_to_keep=1,
|
||||
checkpoint_score_attribute="score",
|
||||
checkpoint_score_order="max",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_load_state_equivalence(
|
||||
monkeypatch, tmp_path, checkpoint_config: CheckpointConfig
|
||||
):
|
||||
# Use async here because register_checkpoint creates an async task
|
||||
|
||||
# Mock the delete function as we don't want report checkpoints to be deleted.
|
||||
monkeypatch.setattr(
|
||||
ray.train.v2._internal.execution.checkpoint.checkpoint_manager,
|
||||
"delete_fs_path",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
exp_name = f"checkpoint_manager_test-{uuid.uuid4().hex}"
|
||||
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name=exp_name,
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=checkpoint_config,
|
||||
)
|
||||
training_reports = create_dummy_training_reports(
|
||||
num_results=2, storage_context=storage_context
|
||||
) + create_dummy_training_reports(
|
||||
num_results=1,
|
||||
storage_context=storage_context,
|
||||
include_validation=True,
|
||||
starting_checkpoint_index=2,
|
||||
)
|
||||
|
||||
# Register the training results into checkpoint manager
|
||||
for i, tr in enumerate(training_reports):
|
||||
checkpoint_manager.register_checkpoint(tr)
|
||||
assert checkpoint_manager._current_report_index == i + 1
|
||||
loaded_checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=checkpoint_config,
|
||||
)
|
||||
assert _checkpoint_managers_equal(checkpoint_manager, loaded_checkpoint_manager)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"json_state,match",
|
||||
[
|
||||
(
|
||||
'{"dummy": "1", "dummy_dict": {"key": "value"}}',
|
||||
"You are loading a checkpoint manager snapshot saved with an unknown Ray version but",
|
||||
),
|
||||
('{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"', None),
|
||||
(
|
||||
'{"ray_version": "2.0.0", "dummy": "1", "dummy_dict": {"key": "value"}}',
|
||||
"You are loading a checkpoint manager snapshot saved with Ray version 2.0.0 but",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_load_state_error(tmp_path, json_state, match):
|
||||
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="load_state_error_experiment",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
with pytest.raises(
|
||||
CheckpointManagerInitializationError,
|
||||
match=match,
|
||||
):
|
||||
checkpoint_manager._load_state(json_state)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_before_init_train_context(tmp_path):
|
||||
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="my_experiment_name",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
workers = [create_autospec(Worker, instance=True) for _ in range(4)]
|
||||
|
||||
# Assert without a checkpoint.
|
||||
assert checkpoint_manager.before_init_train_context(workers) == {
|
||||
"checkpoint": [None] * 4,
|
||||
"current_report_index": [0] * 4,
|
||||
}
|
||||
|
||||
# Assert with a checkpoint
|
||||
latest_checkpoint_report = create_dummy_training_reports(1, storage_context)[0]
|
||||
checkpoint_manager.register_checkpoint(latest_checkpoint_report)
|
||||
assert checkpoint_manager.before_init_train_context(workers) == {
|
||||
"checkpoint": [latest_checkpoint_report.checkpoint] * 4,
|
||||
"current_report_index": [1] * 4,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_checkpoint_management(tmp_path):
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="pending_checkpoint_management_experiment",
|
||||
)
|
||||
checkpoint_config = CheckpointConfig(
|
||||
num_to_keep=1,
|
||||
checkpoint_score_attribute="score",
|
||||
checkpoint_score_order="max",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=checkpoint_config,
|
||||
)
|
||||
(
|
||||
low_initial_high_final_training_report,
|
||||
high_initial_low_final_training_report,
|
||||
final_training_report,
|
||||
) = create_dummy_training_reports(
|
||||
num_results=3, storage_context=storage_context, include_validation=True
|
||||
)
|
||||
final_training_report.validation = False
|
||||
scoreless_training_report = create_dummy_training_reports(
|
||||
num_results=1,
|
||||
storage_context=storage_context,
|
||||
include_metrics=False,
|
||||
starting_checkpoint_index=3,
|
||||
)[0]
|
||||
|
||||
# Register pending/final/unknown checkpoints and verify their storage
|
||||
checkpoint_manager.register_checkpoint(low_initial_high_final_training_report)
|
||||
checkpoint_manager.register_checkpoint(final_training_report)
|
||||
checkpoint_manager.register_checkpoint(scoreless_training_report)
|
||||
checkpoint_manager.register_checkpoint(high_initial_low_final_training_report)
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
low_initial_high_final_training_report.checkpoint, # keep pending
|
||||
high_initial_low_final_training_report.checkpoint, # keep pending/latest
|
||||
final_training_report.checkpoint, # keep highest final score so far
|
||||
]
|
||||
|
||||
# Assert checkpoint state after all tasks are done
|
||||
checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{
|
||||
low_initial_high_final_training_report.checkpoint: (
|
||||
{"score": 200},
|
||||
ReportedCheckpointStatus.VALIDATED,
|
||||
),
|
||||
high_initial_low_final_training_report.checkpoint: (
|
||||
{"score": 100},
|
||||
ReportedCheckpointStatus.VALIDATED,
|
||||
),
|
||||
}
|
||||
)
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
high_initial_low_final_training_report.checkpoint, # keep latest checkpoint
|
||||
low_initial_high_final_training_report.checkpoint, # keep highest score checkpoint
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_checkpoint_management_break_ties_by_report_index(tmp_path):
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="pending_checkpoint_management_break_ties_by_report_index_experiment",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
training_reports = create_dummy_training_reports(
|
||||
num_results=2, storage_context=storage_context, include_validation=True
|
||||
)
|
||||
checkpoint_manager.register_checkpoint(training_reports[0])
|
||||
checkpoint_manager.register_checkpoint(training_reports[1])
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
training_reports[0].checkpoint,
|
||||
training_reports[1].checkpoint,
|
||||
]
|
||||
checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{
|
||||
training_reports[1].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
|
||||
}
|
||||
)
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
training_reports[0].checkpoint,
|
||||
training_reports[1].checkpoint,
|
||||
]
|
||||
checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{
|
||||
training_reports[0].checkpoint: ({}, ReportedCheckpointStatus.VALIDATED),
|
||||
}
|
||||
)
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
training_reports[0].checkpoint,
|
||||
training_reports[1].checkpoint,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_checkpoint_management_finalized_checkpoint(tmp_path):
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="pending_checkpoint_management_experiment",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(
|
||||
checkpoint_score_attribute="score",
|
||||
checkpoint_score_order="max",
|
||||
),
|
||||
)
|
||||
training_reports = create_dummy_training_reports(
|
||||
num_results=2, storage_context=storage_context
|
||||
)
|
||||
checkpoint_manager.register_checkpoint(training_reports[0])
|
||||
checkpoint_manager.register_checkpoint(training_reports[1])
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
training_reports[0].checkpoint,
|
||||
training_reports[1].checkpoint,
|
||||
]
|
||||
checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{
|
||||
training_reports[0].checkpoint: (
|
||||
{"score": 100},
|
||||
ReportedCheckpointStatus.VALIDATED,
|
||||
),
|
||||
}
|
||||
)
|
||||
assert [tr.checkpoint for tr in checkpoint_manager._checkpoint_results] == [
|
||||
training_reports[0].checkpoint,
|
||||
training_reports[1].checkpoint,
|
||||
]
|
||||
|
||||
|
||||
def test_update_checkpoints_with_metrics_not_in_checkpoint_results(tmp_path):
|
||||
storage_context = StorageContext(
|
||||
storage_path=tmp_path,
|
||||
experiment_dir_name="update_checkpoints_with_metrics_error_experiment",
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
storage_context=storage_context,
|
||||
checkpoint_config=CheckpointConfig(),
|
||||
)
|
||||
training_reports = create_dummy_training_reports(
|
||||
num_results=1, storage_context=storage_context
|
||||
)
|
||||
checkpoint_manager._pending_training_results[training_reports[0].checkpoint] = (
|
||||
_TrainingResult(training_reports[0].checkpoint, training_reports[0].metrics),
|
||||
training_reports[0].validation,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
checkpoint_manager.update_checkpoints_with_validation_result(
|
||||
{
|
||||
training_reports[0].checkpoint: (
|
||||
{"score": 100},
|
||||
ReportedCheckpointStatus.VALIDATED,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", "-x", __file__]))
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user