chore: import upstream snapshot with attribution

python/ray/train/tests/test_torch_device_manager.py

@@ -0,0 +1,95 @@
+import pytest
+import torch
+
+import ray
+from ray.air._internal.device_manager import (
+    CUDATorchDeviceManager,
+    NPUTorchDeviceManager,
+    get_torch_device_manager_by_context,
+)
+from ray.air._internal.device_manager.npu import NPU_TORCH_PACKAGE_AVAILABLE
+from ray.cluster_utils import Cluster
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+if NPU_TORCH_PACKAGE_AVAILABLE:
+    import torch_npu  # noqa: F401
+
+
+@pytest.fixture
+def ray_2_node_2_npus():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, resources={"NPU": 2})
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_1_node_1_gpu_1_npu():
+    cluster = Cluster()
+    cluster.add_node(num_cpus=4, num_gpus=1, resources={"NPU": 1})
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+def test_cuda_device_manager(ray_2_node_2_gpu):
+    def train_fn():
+        assert isinstance(get_torch_device_manager_by_context(), CUDATorchDeviceManager)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, resources_per_worker={"GPU": 1}
+        ),
+    )
+
+    trainer.fit()
+
+
+def test_npu_device_manager(ray_2_node_2_npus):
+    def train_fn():
+        assert isinstance(get_torch_device_manager_by_context(), NPUTorchDeviceManager)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn,
+        scaling_config=ScalingConfig(num_workers=1, resources_per_worker={"NPU": 1}),
+    )
+
+    if NPU_TORCH_PACKAGE_AVAILABLE and torch.npu.is_available():
+        # Except test run successfully when torch npu is available.
+        trainer.fit()
+    else:
+        # A RuntimeError will be triggered when NPU resources are declared
+        # but the torch npu is actually not available
+        with pytest.raises(RuntimeError):
+            trainer.fit()
+
+
+def test_device_manager_conflict(ray_1_node_1_gpu_1_npu):
+    trainer = TorchTrainer(
+        train_loop_per_worker=lambda: None,
+        scaling_config=ScalingConfig(
+            num_workers=1, use_gpu=True, resources_per_worker={"GPU": 1, "NPU": 1}
+        ),
+    )
+    # TODO: Do validation at the `ScalingConfig.__post_init__` level instead.
+    with pytest.raises(RuntimeError):
+        trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))