[tune/placement group] dist. training placement group support (#11934)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

python/ray/tune/tests/test_tensorflow_trainable.py

@@ -1,5 +1,9 @@
+from typing import Dict, Optional
+
 import pytest
 import ray
+from ray import tune
+from ray.cluster_utils import Cluster
 from ray.tune.integration.tensorflow import DistributedTrainableCreator
 from ray.tune.examples.tf_distributed_keras_example import train_mnist
 
@@ -20,6 +24,34 @@ def ray_start_4_cpus():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_4_node():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=1)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture
+def ray_4_node_gpu():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=2, num_gpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
 @pytest.fixture
 def ray_connect_cluster():
     try:
@@ -31,8 +63,16 @@ def ray_connect_cluster():
     ray.shutdown()
 
 
+def _train_check_global(config: Dict, checkpoint_dir: Optional[str] = None):
+    """For testing only. Putting this here because Ray has problems
+    serializing within the test file."""
+    import time
+    time.sleep(0.1)
+    tune.report(is_distributed=True)
+
+
 def test_single_step(ray_start_2_cpus):  # noqa: F811
-    trainable_cls = DistributedTrainableCreator(train_mnist)
+    trainable_cls = DistributedTrainableCreator(train_mnist, num_workers=2)
     trainer = trainable_cls()
     trainer.train()
     trainer.stop()
@@ -50,9 +90,45 @@ def test_validation(ray_start_2_cpus):  # noqa: F811
     def bad_func(a, b, c):
         return 1
 
-    t_cls = DistributedTrainableCreator(bad_func)
     with pytest.raises(ValueError):
-        t_cls()
+        DistributedTrainableCreator(bad_func)
+
+
+def test_colocated(ray_4_node):  # noqa: F811
+    assert ray.available_resources()["CPU"] == 4
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global, num_workers=4, num_workers_per_host=1)
+    trainable = trainable_cls()
+    assert ray.available_resources().get("CPU", 0) == 0
+    trainable.train()
+    trainable.stop()
+
+
+def test_colocated_gpu(ray_4_node_gpu):  # noqa: F811
+    assert ray.available_resources()["GPU"] == 8
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global,
+        num_workers=4,
+        num_gpus_per_worker=2,
+        num_workers_per_host=1)
+    trainable = trainable_cls()
+    assert ray.available_resources().get("GPU", 0) == 0
+    trainable.train()
+    trainable.stop()
+
+
+def test_colocated_gpu_double(ray_4_node_gpu):  # noqa: F811
+    assert ray.available_resources()["GPU"] == 8
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global,
+        num_workers=8,
+        num_gpus_per_worker=1,
+        num_cpus_per_worker=1,
+        num_workers_per_host=2)
+    trainable = trainable_cls()
+    assert ray.available_resources().get("GPU", 0) == 0
+    trainable.train()
+    trainable.stop()
 
 
 if __name__ == "__main__":

python/ray/tune/tests/test_torch_trainable.py

@@ -6,6 +6,7 @@ import torch.distributed as dist
 
 import ray
 from ray import tune
+from ray.cluster_utils import Cluster
 from ray.tune.integration.torch import (DistributedTrainableCreator,
                                         distributed_checkpoint_dir,
                                         _train_simple, _train_check_global)
@@ -97,6 +98,79 @@ def test_checkpoint(ray_start_2_cpus, rank):  # noqa: F811
             assert not os.path.exists(path)
 
 
+@pytest.fixture
+def ray_4_node():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=1)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+    # Ensure that tests don't ALL fail
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+@pytest.fixture
+def ray_4_node_gpu():
+    cluster = Cluster()
+    for _ in range(4):
+        cluster.add_node(num_cpus=2, num_gpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+    # Ensure that tests don't ALL fail
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def test_colocated(ray_4_node):  # noqa: F811
+    assert ray.available_resources()["CPU"] == 4
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global, num_workers=4, num_workers_per_host=1)
+    trainable = trainable_cls()
+    assert ray.available_resources().get("CPU", 0) == 0
+    trainable.train()
+    trainable.stop()
+
+
+def test_colocated_gpu(ray_4_node_gpu):  # noqa: F811
+    assert ray.available_resources()["GPU"] == 8
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global,
+        num_workers=4,
+        num_gpus_per_worker=2,
+        num_workers_per_host=1)
+    trainable = trainable_cls()
+    assert ray.available_resources().get("GPU", 0) == 0
+    trainable.train()
+    trainable.stop()
+
+
+def test_colocated_gpu_double(ray_4_node_gpu):  # noqa: F811
+    assert ray.available_resources()["GPU"] == 8
+    trainable_cls = DistributedTrainableCreator(
+        _train_check_global,
+        num_workers=8,
+        num_gpus_per_worker=1,
+        num_workers_per_host=2,
+        timeout_s=30)
+    trainable = trainable_cls()
+    print("?????")
+    print(ray.available_resources().get("GPU"))
+    assert ray.available_resources().get("GPU", 0) == 0
+    trainable.train()
+    trainable.stop()
+
+
 if __name__ == "__main__":
     import pytest
     import sys