[train] support memory per worker

python/ray/train/_internal/worker_group.py

@@ -174,6 +174,13 @@ def __init__(
         self.num_cpus_per_worker = num_cpus_per_worker
         self.num_gpus_per_worker = num_gpus_per_worker
         self.additional_resources_per_worker = additional_resources_per_worker
+
+        self.memory_per_worker = None
+        if self.additional_resources_per_worker is not None:
+            self.memory_per_worker = self.additional_resources_per_worker.pop(
+                "memory", None
+            )
+
         self.workers = []
         self._base_cls = create_executable_class(actor_cls)
         assert issubclass(self._base_cls, RayTrainWorker)
@@ -188,6 +195,7 @@ def __init__(
         self._remote_cls = ray.remote(
             num_cpus=self.num_cpus_per_worker,
             num_gpus=self.num_gpus_per_worker,
+            memory=self.memory_per_worker,
             resources=self.additional_resources_per_worker,
         )(self._base_cls)
         self.start()

python/ray/train/tests/test_worker_group.py

@@ -32,6 +32,14 @@ def ray_start_2_cpus_and_neuron_core_accelerator():
     ray.shutdown()
 
 
+@pytest.fixture
+def ray_start_2_cpus_and_10kb_memory():
+    address_info = ray.init(num_cpus=2, _memory=10_000)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
 def test_worker_creation(ray_start_2_cpus):
     assert ray.available_resources()["CPU"] == 2
     wg = WorkerGroup(num_workers=2)
@@ -52,6 +60,11 @@ def test_worker_creation_num_cpus(ray_start_2_cpus):
     wg.shutdown()
 
 
+def test_worker_creation_with_memory(ray_start_2_cpus_and_10kb_memory):
+    wg = WorkerGroup(num_workers=2, additional_resources_per_worker={"memory": 1_000})
+    assert len(wg.workers) == 2
+
+
 def test_worker_shutdown(ray_start_2_cpus):
     assert ray.available_resources()["CPU"] == 2
     wg = WorkerGroup(num_workers=2)