chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_serialization.py

@@ -0,0 +1,103 @@
+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__]))