[tune] Fix memory resources for head bundle (#23861)

Fixes memory and object_store_memory actor options not being set properly for the Tune trainable.
ray-project · Apr 12, 2022 · ff60ebd · ff60ebd
1 parent 71e142b
commit ff60ebd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/ray/tune/utils/placement_groups.py b/python/ray/tune/utils/placement_groups.py
@@ -508,6 +508,8 @@ def get_full_actor_cls(
             head_bundle = pg.bundle_specs[0].copy()
             num_cpus = head_bundle.pop("CPU", 0)
             num_gpus = head_bundle.pop("GPU", 0)
+            memory = head_bundle.pop("memory", None)
+            object_store_memory = head_bundle.pop("object_store_memory", None)
 
             # Only custom resources remain in `head_bundle`
             resources = head_bundle
@@ -517,6 +519,8 @@ def get_full_actor_cls(
                 placement_group_capture_child_tasks=True,
                 num_cpus=num_cpus,
                 num_gpus=num_gpus,
+                memory=memory,
+                object_store_memory=object_store_memory,
                 resources=resources,
             )
         else: