[Tune] Enable experiment restore from moved cloud uri

python/ray/tune/analysis/experiment_analysis.py

@@ -30,10 +30,7 @@
     TRAINING_ITERATION,
 )
 from ray.tune.experiment import Trial
-from ray.tune.execution.trial_runner import (
-    _find_newest_experiment_checkpoint,
-    _load_trial_from_checkpoint,
-)
+from ray.tune.execution.trial_runner import _find_newest_experiment_checkpoint
 from ray.tune.trainable.util import TrainableUtil
 from ray.tune.utils.util import unflattened_lookup
 
@@ -143,13 +140,10 @@ def _load_checkpoints_from_latest(self, latest_checkpoint: List[str]) -> None:
 
             if "checkpoints" not in experiment_state:
                 raise TuneError("Experiment state invalid; no checkpoints found.")
+
             self._checkpoints_and_paths += [
-                (_decode_checkpoint_from_experiment_state(cp), Path(path).parent)
-                for cp in experiment_state["checkpoints"]
+                (cp, Path(path).parent) for cp in experiment_state["checkpoints"]
             ]
-            self._checkpoints_and_paths = sorted(
-                self._checkpoints_and_paths, key=lambda tup: tup[0]["trial_id"]
-            )
 
     def _get_latest_checkpoint(self, experiment_checkpoint_path: Path) -> List[str]:
         # Case 1: Dir specified, find latest checkpoint.
@@ -798,12 +792,10 @@ def _get_trial_paths(self) -> List[str]:
                 "out of sync, as checkpointing is periodic."
             )
             self.trials = []
-            _trial_paths = []
-            for checkpoint, path in self._checkpoints_and_paths:
+            for trial_json_state, path in self._checkpoints_and_paths:
                 try:
-                    trial = _load_trial_from_checkpoint(
-                        checkpoint, stub=True, new_local_dir=str(path)
-                    )
+                    trial = Trial.from_json_state(trial_json_state, stub=True)
+                    trial.local_dir = str(path)
                 except Exception:
                     logger.warning(
                         f"Could not load trials from experiment checkpoint. "
@@ -814,7 +806,9 @@ def _get_trial_paths(self) -> List[str]:
                     )
                     continue
                 self.trials.append(trial)
-                _trial_paths.append(str(trial.logdir))
+
+            self.trials.sort(key=lambda trial: trial.trial_id)
+            _trial_paths = [str(trial.logdir) for trial in self.trials]
 
         if not _trial_paths:
             raise TuneError("No trials found.")

python/ray/tune/execution/trial_runner.py

@@ -1,6 +1,6 @@
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, DefaultDict, List, Mapping, Optional, Union, Tuple, Set
+from typing import DefaultDict, List, Optional, Union, Tuple, Set
 
 import click
 from datetime import datetime
@@ -67,65 +67,6 @@ def _find_newest_experiment_checkpoint(ckpt_dir) -> Optional[str]:
     return max(full_paths)
 
 
-def _load_trial_from_checkpoint(
-    trial_cp: dict, stub: bool = False, new_local_dir: Optional[str] = None
-) -> Trial:
-    """Create a Trial from the state stored in the experiment checkpoint.
-
-    Args:
-        trial_cp: Trial state from the experiment checkpoint, which is loaded
-            from the trial's `Trial.get_json_state`.
-        stub: Whether or not to validate the trainable name when creating the Trial.
-            Used for testing purposes for creating mocks.
-        new_local_dir: If set, this `local_dir` will overwrite what's saved in the
-            `trial_cp` state. Used in the case that the trial directory has moved.
-            The Trial `logdir` and the persistent trial checkpoints will have their
-            paths updated relative to this new directory.
-
-    Returns:
-        new_trial: New trial with state loaded from experiment checkpoint
-    """
-    new_trial = Trial(
-        trial_cp["trainable_name"],
-        stub=stub,
-        _setup_default_resource=False,
-    )
-    if new_local_dir:
-        trial_cp["local_dir"] = new_local_dir
-    new_trial.__setstate__(trial_cp)
-    new_trial.refresh_default_resource_request()
-    return new_trial
-
-
-def _load_trials_from_experiment_checkpoint(
-    experiment_checkpoint: Mapping[str, Any],
-    stub: bool = False,
-    new_local_dir: Optional[str] = None,
-) -> List[Trial]:
-    """Create trial objects from experiment checkpoint.
-
-    Given an experiment checkpoint (TrialRunner state dict), return
-    list of trials. See `_ExperimentCheckpointManager.checkpoint` for
-    what's saved in the TrialRunner state dict.
-    """
-    checkpoints = [
-        json.loads(cp, cls=TuneFunctionDecoder) if isinstance(cp, str) else cp
-        for cp in experiment_checkpoint["checkpoints"]
-    ]
-
-    trials = []
-    for trial_cp in checkpoints:
-        trials.append(
-            _load_trial_from_checkpoint(
-                trial_cp,
-                stub=stub,
-                new_local_dir=new_local_dir,
-            )
-        )
-
-    return trials
-
-
 @dataclass
 class _ResumeConfig:
     resume_unfinished: bool = True
@@ -154,17 +95,16 @@ class _ExperimentCheckpointManager:
 
     def __init__(
         self,
-        checkpoint_dir: str,
+        local_checkpoint_dir: str,
         checkpoint_period: Union[int, float, str],
         start_time: float,
         session_str: str,
         syncer: Syncer,
         sync_trial_checkpoints: bool,
-        local_dir: str,
-        remote_dir: str,
+        remote_checkpoint_dir: str,
         sync_every_n_trial_checkpoints: Optional[int] = None,
     ):
-        self._checkpoint_dir = checkpoint_dir
+        self._local_checkpoint_dir = local_checkpoint_dir
         self._auto_checkpoint_enabled = checkpoint_period == "auto"
         if self._auto_checkpoint_enabled:
             self._checkpoint_period = 10.0  # Initial value
@@ -176,8 +116,7 @@ def __init__(
 
         self._syncer = syncer
         self._sync_trial_checkpoints = sync_trial_checkpoints
-        self._local_dir = local_dir
-        self._remote_dir = remote_dir
+        self._remote_checkpoint_dir = remote_checkpoint_dir
 
         self._last_checkpoint_time = 0.0
         self._last_sync_time = 0.0
@@ -225,7 +164,7 @@ def checkpoint(
         Args:
             force: Forces a checkpoint despite checkpoint_period.
         """
-        if not self._checkpoint_dir:
+        if not self._local_checkpoint_dir:
             return
 
         force = force or self._should_force_cloud_sync
@@ -243,12 +182,14 @@ def _serialize_and_write():
                     "timestamp": self._last_checkpoint_time,
                 },
             }
-            tmp_file_name = os.path.join(self._checkpoint_dir, ".tmp_checkpoint")
+            tmp_file_name = os.path.join(self._local_checkpoint_dir, ".tmp_checkpoint")
             with open(tmp_file_name, "w") as f:
                 json.dump(runner_state, f, indent=2, cls=TuneFunctionEncoder)
 
             os.replace(tmp_file_name, checkpoint_file)
-            search_alg.save_to_dir(self._checkpoint_dir, session_str=self._session_str)
+            search_alg.save_to_dir(
+                self._local_checkpoint_dir, session_str=self._session_str
+            )
 
         checkpoint_time_start = time.monotonic()
         with out_of_band_serialize_dataset():
@@ -274,14 +215,14 @@ def _serialize_and_write():
                         "`sync_timeout` in `SyncConfig`."
                     )
                 synced = self._syncer.sync_up(
-                    local_dir=self._local_dir,
-                    remote_dir=self._remote_dir,
+                    local_dir=self._local_checkpoint_dir,
+                    remote_dir=self._remote_checkpoint_dir,
                     exclude=exclude,
                 )
             else:
                 synced = self._syncer.sync_up_if_needed(
-                    local_dir=self._local_dir,
-                    remote_dir=self._remote_dir,
+                    local_dir=self._local_checkpoint_dir,
+                    remote_dir=self._remote_checkpoint_dir,
                     exclude=exclude,
                 )
 
@@ -320,7 +261,7 @@ def _serialize_and_write():
             )
 
         self._last_checkpoint_time = time.time()
-        return self._checkpoint_dir
+        return self._local_checkpoint_dir
 
 
 @DeveloperAPI
@@ -350,14 +291,15 @@ class TrialRunner:
         search_alg: SearchAlgorithm for generating
             Trial objects.
         scheduler: Defaults to FIFOScheduler.
-        local_checkpoint_dir: Path where
-            global checkpoints are stored and restored from.
-        remote_checkpoint_dir: Remote path where
-            global checkpoints are stored and restored from. Used
-            if `resume` == REMOTE.
-        sync_config: See `tune.py:run`.
-        stopper: Custom class for stopping whole experiments. See
-            ``Stopper``.
+        local_checkpoint_dir: Path where global experiment state checkpoints
+            are saved and restored from.
+        sync_config: See :class:`~ray.tune.syncer.SyncConfig`.
+            Within sync config, the `upload_dir` specifies cloud storage, and
+            experiment state checkpoints will be synced to the `remote_checkpoint_dir`:
+            `{sync_config.upload_dir}/{experiment_name}`.
+        experiment_dir_name: Experiment directory name.
+            See :class:`~ray.tune.experiment.Experiment`.
+        stopper: Custom class for stopping whole experiments. See ``Stopper``.
         resume: see `tune.py:run`.
         server_port: Port number for launching TuneServer.
         fail_fast: Finishes as soon as a trial fails if True.
@@ -388,8 +330,8 @@ def __init__(
         search_alg: Optional[SearchAlgorithm] = None,
         scheduler: Optional[TrialScheduler] = None,
         local_checkpoint_dir: Optional[str] = None,
-        remote_checkpoint_dir: Optional[str] = None,
         sync_config: Optional[SyncConfig] = None,
+        experiment_dir_name: Optional[str] = None,
         stopper: Optional[Stopper] = None,
         resume: Union[str, bool] = False,
         server_port: Optional[int] = None,
@@ -436,11 +378,11 @@ def __init__(
             # Manual override
             self._max_pending_trials = int(max_pending_trials)
 
-        sync_config = sync_config or SyncConfig()
+        self._sync_config = sync_config or SyncConfig()
 
         self.trial_executor.setup(
             max_pending_trials=self._max_pending_trials,
-            trainable_kwargs={"sync_timeout": sync_config.sync_timeout},
+            trainable_kwargs={"sync_timeout": self._sync_config.sync_timeout},
         )
 
         self._metric = metric
@@ -485,9 +427,9 @@ def __init__(
         if self._local_checkpoint_dir:
             os.makedirs(self._local_checkpoint_dir, exist_ok=True)
 
-        self._remote_checkpoint_dir = remote_checkpoint_dir
+        self._experiment_dir_name = experiment_dir_name
 
-        self._syncer = get_node_to_storage_syncer(sync_config)
+        self._syncer = get_node_to_storage_syncer(self._sync_config)
         self._stopper = stopper or NoopStopper()
         self._resumed = False
 
@@ -562,14 +504,13 @@ def end_experiment_callbacks(self) -> None:
 
     def _create_checkpoint_manager(self, sync_trial_checkpoints: bool = True):
         return _ExperimentCheckpointManager(
-            checkpoint_dir=self._local_checkpoint_dir,
+            local_checkpoint_dir=self._local_checkpoint_dir,
             checkpoint_period=self._checkpoint_period,
             start_time=self._start_time,
             session_str=self._session_str,
             syncer=self._syncer,
             sync_trial_checkpoints=sync_trial_checkpoints,
-            local_dir=self._local_checkpoint_dir,
-            remote_dir=self._remote_checkpoint_dir,
+            remote_checkpoint_dir=self._remote_checkpoint_dir,
             sync_every_n_trial_checkpoints=self._trial_checkpoint_config.num_to_keep,
         )
 
@@ -585,6 +526,12 @@ def search_alg(self):
     def scheduler_alg(self):
         return self._scheduler_alg
 
+    @property
+    def _remote_checkpoint_dir(self):
+        if self._sync_config.upload_dir and self._experiment_dir_name:
+            return os.path.join(self._sync_config.upload_dir, self._experiment_dir_name)
+        return None
+
     def _validate_resume(
         self, resume_type: Union[str, bool], driver_sync_trial_checkpoints=True
     ) -> Tuple[bool, Optional[_ResumeConfig]]:
@@ -845,19 +792,34 @@ def resume(
             )
         )
 
-        trial_runner_data = runner_state["runner_data"]
-        # Don't overwrite the current `_local_checkpoint_dir`
-        # The current directory could be different from the checkpointed
-        # directory, if the experiment directory has changed.
-        trial_runner_data.pop("_local_checkpoint_dir", None)
+        # 1. Restore trial runner state
+        self.__setstate__(runner_state["runner_data"])
 
-        self.__setstate__(trial_runner_data)
+        # 2. Restore search algorithm state
         if self._search_alg.has_checkpoint(self._local_checkpoint_dir):
             self._search_alg.restore_from_dir(self._local_checkpoint_dir)
 
-        trials = _load_trials_from_experiment_checkpoint(
-            runner_state, new_local_dir=self._local_checkpoint_dir
-        )
+        # 3. Load trial table from experiment checkpoint
+        trials = []
+        for trial_json_state in runner_state["checkpoints"]:
+            trial = Trial.from_json_state(trial_json_state)
+
+            # The following properties may be updated on restoration
+            # Ex: moved local/cloud experiment directory
+            trial.local_dir = self._local_checkpoint_dir
+            trial.sync_config = self._sync_config
+            trial.experiment_dir_name = self._experiment_dir_name
+
+            # Avoid creating logdir in client mode for returned trial results,
+            # since the dir might not be creatable locally.
+            # TODO(ekl) this is kind of a hack.
+            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
         for trial in sorted(trials, key=lambda t: t.last_update_time, reverse=True):
             trial_to_add = trial
             if trial.status == Trial.ERROR:
@@ -1623,6 +1585,9 @@ def __getstate__(self):
             "_syncer",
             "_callbacks",
             "_checkpoint_manager",
+            "_local_checkpoint_dir",
+            "_sync_config",
+            "_experiment_dir_name",
         ]:
             del state[k]
         state["launch_web_server"] = bool(self._server)