[Tune] let categorical values return indices that get resolved in a separate step

python/ray/tune/BUILD

@@ -240,6 +240,14 @@ py_test(
     tags = ["team:ml", "exclusive", "medium_instance"],
 )
 
+py_test(
+    name = "test_placeholder",
+    size = "small",
+    srcs = ["tests/test_placeholder.py"],
+    deps = [":tune_lib"],
+    tags = ["team:ml", "exclusive"],
+)
+
 py_test(
     name = "test_searchers",
     size = "large",

python/ray/tune/execution/ray_trial_executor.py

@@ -13,6 +13,7 @@
 from typing import Callable, Dict, Iterable, List, Optional, Set, Union, Tuple
 
 import ray
+from ray.actor import ActorHandle
 from ray.air import Checkpoint, AcquiredResources, ResourceRequest
 from ray.air._internal.checkpoint_manager import CheckpointStorage, _TrackedCheckpoint
 from ray.air.constants import COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV
@@ -351,7 +352,7 @@ def get_ready_trial(self) -> Optional[Trial]:
 
         return None
 
-    def _maybe_use_cached_actor(self, trial, logger_creator) -> Optional:
+    def _maybe_use_cached_actor(self, trial, logger_creator) -> Optional[ActorHandle]:
         if not self._reuse_actors:
             return None
 
@@ -426,7 +427,6 @@ def _setup_remote_runner(self, trial):
         # configure the remote runner to use a noop-logger.
         trial_config = copy.deepcopy(trial.config)
         trial_config[TRIAL_INFO] = _TrialInfo(trial)
-
         stdout_file, stderr_file = trial.log_to_file
         trial_config[STDOUT_FILE] = stdout_file
         trial_config[STDERR_FILE] = stderr_file

python/ray/tune/execution/trial_runner.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import DefaultDict, List, Optional, Union, Tuple, Set
+from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple, Set
 
 import click
 from datetime import datetime
@@ -331,7 +331,9 @@ class TrialRunner:
 
     def __init__(
         self,
+        *,
         search_alg: Optional[SearchAlgorithm] = None,
+        placeholder_resolvers: Optional[Dict[Tuple, Any]] = None,
         scheduler: Optional[TrialScheduler] = None,
         local_checkpoint_dir: Optional[str] = None,
         sync_config: Optional[SyncConfig] = None,
@@ -349,6 +351,7 @@ def __init__(
         driver_sync_trial_checkpoints: bool = False,
     ):
         self._search_alg = search_alg or BasicVariantGenerator()
+        self._placeholder_resolvers = placeholder_resolvers
         self._scheduler_alg = scheduler or FIFOScheduler()
         self.trial_executor = trial_executor or RayTrialExecutor()
         self._callbacks = CallbackList(callbacks or [])
@@ -823,7 +826,6 @@ def resume(
             if not ray.util.client.ray.is_connected():
                 trial.init_logdir()  # Create logdir if it does not exist
 
-            trial.refresh_default_resource_request()
             trials.append(trial)
 
         # 4. Set trial statuses according to the resume configuration
@@ -1085,6 +1087,14 @@ def add_trial(self, trial: Trial):
         Args:
             trial: Trial to queue.
         """
+        # If the config map has had all the references replaced with placeholders,
+        # resolve them before adding the trial.
+        if self._placeholder_resolvers:
+            trial.resolve_config_placeholders(self._placeholder_resolvers)
+
+        # With trial.config resolved, create placement group factory if needed.
+        trial.create_placement_group_factory()
+
         self._trials.append(trial)
         if trial.status != Trial.TERMINATED:
             self._live_trials.add(trial)
@@ -1586,6 +1596,7 @@ def __getstate__(self):
             "_stop_queue",
             "_server",
             "_search_alg",
+            "_placeholder_resolvers",
             "_scheduler_alg",
             "_pending_trial_queue_times",
             "trial_executor",

python/ray/tune/experiment/trial.py

@@ -9,7 +9,7 @@
 import re
 import shutil
 import time
-from typing import Dict, Optional, Sequence, Union, Callable, List, Tuple
+from typing import Any, Dict, Optional, Sequence, Union, Callable, List, Tuple
 import uuid
 
 import ray
@@ -239,6 +239,8 @@ class Trial:
         "param_config",
         "extra_arg",
         "placement_group_factory",
+        "_resources",
+        "_default_placement_group_factory",
     ]
 
     PENDING = "PENDING"
@@ -292,42 +294,25 @@ def __init__(
         # Trial config
         self.trainable_name = trainable_name
         self.trial_id = Trial.generate_id() if trial_id is None else trial_id
-        self.config = config or {}
         self._local_dir = local_dir  # This remains unexpanded for syncing.
 
+        self.config = config or {}
+        # Save a copy of the original unresolved config so that we can swap
+        # out and update any reference config values after restoration.
+        self.__unresolved_config = self.config
+
         # Parameters that Tune varies across searches.
         self.evaluated_params = evaluated_params or {}
         self.experiment_tag = experiment_tag
         self.location = _Location()
-        trainable_cls = self.get_trainable_cls()
-        if trainable_cls and _setup_default_resource:
-            default_resources = trainable_cls.default_resource_request(self.config)
-
-            # If Trainable returns resources, do not allow manual override via
-            # `resources_per_trial` by the user.
-            if default_resources:
-                if resources or placement_group_factory:
-                    raise ValueError(
-                        "Resources for {} have been automatically set to {} "
-                        "by its `default_resource_request()` method. Please "
-                        "clear the `resources_per_trial` option.".format(
-                            trainable_cls, default_resources
-                        )
-                    )
-
-                if isinstance(default_resources, PlacementGroupFactory):
-                    placement_group_factory = default_resources
-                    resources = None
-                else:
-                    placement_group_factory = None
-                    resources = default_resources
-
-        self.placement_group_factory = _to_pg_factory(
-            resources, placement_group_factory
-        )
-
         self.stopping_criterion = stopping_criterion or {}
 
+        self._setup_default_resource = _setup_default_resource
+        self._resources = resources
+        self._default_placement_group_factory = placement_group_factory
+        # Will be created in create_placement_group_factory().
+        self.placement_group_factory = None
+
         self.log_to_file = log_to_file
         # Make sure `stdout_file, stderr_file = Trial.log_to_file` works
         if (
@@ -417,6 +402,48 @@ def __init__(
         self._state_json = None
         self._state_valid = False
 
+    def create_placement_group_factory(self):
+        """Compute placement group factor if needed.
+
+        Note: this must be called after all the placeholders in
+        self.config are resolved.
+        """
+        trainable_cls = self.get_trainable_cls()
+        if not trainable_cls or not self._setup_default_resource:
+            # Create placement group factory using default resources.
+            self.placement_group_factory = _to_pg_factory(
+                self._resources, self._default_placement_group_factory
+            )
+            return
+
+        default_resources = trainable_cls.default_resource_request(self.config)
+
+        # If Trainable returns resources, do not allow manual override via
+        # `resources_per_trial` by the user.
+        if default_resources:
+            if self._resources or self._default_placement_group_factory:
+                raise ValueError(
+                    "Resources for {} have been automatically set to {} "
+                    "by its `default_resource_request()` method. Please "
+                    "clear the `resources_per_trial` option.".format(
+                        trainable_cls, default_resources
+                    )
+                )
+
+            if isinstance(default_resources, PlacementGroupFactory):
+                default_placement_group_factory = default_resources
+                resources = None
+            else:
+                default_placement_group_factory = None
+                resources = default_resources
+        else:
+            default_placement_group_factory = self._default_placement_group_factory
+            resources = self._resources
+
+        self.placement_group_factory = _to_pg_factory(
+            resources, default_placement_group_factory
+        )
+
     def _get_default_result_or_future(self) -> Optional[dict]:
         """Calls ray.get on self._default_result_or_future and assigns back.
 
@@ -439,6 +466,13 @@ def _get_default_result_or_future(self) -> Optional[dict]:
             )
         return self._default_result_or_future
 
+    def resolve_config_placeholders(self, placeholder_resolvers: Dict[Tuple, Any]):
+        from ray.tune.impl.placeholder import resolve_placeholders
+
+        # Make a copy of the unresolved config before resolve it.
+        self.config = copy.deepcopy(self.__unresolved_config)
+        resolve_placeholders(self.config, placeholder_resolvers)
+
     @property
     def last_result(self) -> dict:
         # The logic in here is as follows:
@@ -654,15 +688,6 @@ def update_resources(
 
         self.invalidate_json_state()
 
-    def refresh_default_resource_request(self):
-        """Update trial resources according to the trainable's default resource
-        request, if it is provided."""
-        trainable_cls = self.get_trainable_cls()
-        if trainable_cls:
-            default_resources = trainable_cls.default_resource_request(self.config)
-            if default_resources:
-                self.update_resources(default_resources)
-
     def set_runner(self, runner):
         self.runner = runner
         if runner: