wehub-resource-sync
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray.cloudpickle as ray_pickle
|
|
from ray.train.v2._internal.execution.callback import (
|
|
ControllerCallback,
|
|
WorkerGroupCallback,
|
|
)
|
|
from ray.train.v2._internal.execution.context import TrainRunContext
|
|
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.v2.api.exceptions import ControllerError, WorkerGroupError
|
|
|
|
|
|
def block_import(import_name):
|
|
import sys
|
|
|
|
class BlockTorchImport:
|
|
def find_spec(self, fullname, path, target=None):
|
|
if fullname == import_name or fullname.startswith(import_name + "."):
|
|
raise ImportError(
|
|
f"Test error: {import_name} not installed on this node"
|
|
)
|
|
|
|
sys.meta_path.insert(0, BlockTorchImport())
|
|
|
|
|
|
def test_captured_imports(ray_start_4_cpus):
|
|
import torch
|
|
|
|
def capture_torch_import_fn():
|
|
# torch is captured in the closure of the train_fn
|
|
# and should be re-imported on each worker.
|
|
torch.ones(1)
|
|
|
|
class AssertImportsCallback(ControllerCallback):
|
|
def after_controller_start(self, train_run_context: TrainRunContext):
|
|
# Check that torch is not imported in the controller process.
|
|
# The train_fn should be deserialized directly on the workers.
|
|
assert "torch" not in sys.modules
|
|
|
|
trainer = DataParallelTrainer(
|
|
capture_torch_import_fn,
|
|
run_config=ray.train.RunConfig(callbacks=[AssertImportsCallback()]),
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_deserialization_error(ray_start_4_cpus):
|
|
"""Test that train_fn deserialization errors are propagated properly.
|
|
|
|
This test showcases a common deserialization error example, where
|
|
the driver script successfully imports torch, but torch is not
|
|
installed on the worker nodes.
|
|
"""
|
|
import torch
|
|
|
|
def capture_torch_import_fn():
|
|
torch.ones(1)
|
|
|
|
class BlockTorchImportCallback(WorkerGroupCallback):
|
|
def after_worker_group_start(self, worker_group):
|
|
# Make it so that the torch import that happens on
|
|
# train_fn deserialization will fail on workers.
|
|
worker_group.execute(block_import, "torch")
|
|
|
|
trainer = DataParallelTrainer(
|
|
capture_torch_import_fn,
|
|
run_config=ray.train.RunConfig(callbacks=[BlockTorchImportCallback()]),
|
|
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
|
)
|
|
with pytest.raises(ControllerError, match="torch not installed on this node"):
|
|
trainer.fit()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"error",
|
|
[
|
|
WorkerGroupError(
|
|
"Training failed on multiple workers",
|
|
{0: ValueError("worker 0 failed"), 1: RuntimeError("worker 1 failed")},
|
|
),
|
|
ControllerError(Exception("Controller crashed")),
|
|
],
|
|
)
|
|
def test_exceptions_are_picklable(error):
|
|
"""Test that WorkerGroupError and ControllerError are picklable."""
|
|
|
|
# Test pickle/unpickle for WorkerGroupError
|
|
pickled_error = ray_pickle.dumps(error)
|
|
unpickled_error = ray_pickle.loads(pickled_error)
|
|
|
|
# Verify attributes are preserved
|
|
assert str(unpickled_error) == str(error)
|
|
assert type(unpickled_error) is type(error)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|