From ff60ebd4b3b71665ca5ff7fa2473fda7266e8e71 Mon Sep 17 00:00:00 2001 From: Antoni Baum Date: Tue, 12 Apr 2022 21:56:05 +0200 Subject: [PATCH] [tune] Fix memory resources for head bundle (#23861) Fixes memory and object_store_memory actor options not being set properly for the Tune trainable. --- python/ray/tune/utils/placement_groups.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/ray/tune/utils/placement_groups.py b/python/ray/tune/utils/placement_groups.py index c0af68ef5080..87d1a482f750 100644 --- 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: