[tune/train] Consolidate checkpoint manager 3: Ray Tune

doc/source/tune/api_docs/trainable.rst

@@ -142,7 +142,7 @@ You can restore a single trial checkpoint by using ``tune.run(restore=<checkpoin
                 "max_iter": 5
             },
         ).trials
-        last_ckpt = trial.checkpoint.value
+        last_ckpt = trial.checkpoint.dir_or_data
         analysis = tune.run(train, config={"max_iter": 10}, restore=last_ckpt)
 
 Tune also may copy or move checkpoints during the course of tuning. For this purpose,

doc/source/tune/examples/tune-pytorch-cifar.ipynb

@@ -298,7 +298,7 @@
     "    device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
     "    best_trained_model.to(device)\n",
     "\n",
-    "    checkpoint_path = os.path.join(best_trial.checkpoint.value, \"checkpoint\")\n",
+    "    checkpoint_path = os.path.join(best_trial.checkpoint.dir_or_data, \"checkpoint\")\n",
     "\n",
     "    model_state, optimizer_state = torch.load(checkpoint_path)\n",
     "    best_trained_model.load_state_dict(model_state)\n",

doc/source/tune/examples/tune-serve-integration-mnist.ipynb

@@ -439,7 +439,7 @@
     "    best_trial = analysis.get_best_trial(\"mean_accuracy\", \"max\", \"last\")\n",
     "    best_accuracy = best_trial.metric_analysis[\"mean_accuracy\"][\"last\"]\n",
     "    best_trial_config = best_trial.config\n",
-    "    best_checkpoint = best_trial.checkpoint.value\n",
+    "    best_checkpoint = best_trial.checkpoint.dir_or_data\n",
     "\n",
     "    return best_accuracy, best_trial_config, best_checkpoint, num_examples"
    ]
@@ -517,7 +517,7 @@
     "    best_trial = analysis.get_best_trial(\"mean_accuracy\", \"max\", \"last\")\n",
     "    best_accuracy = best_trial.metric_analysis[\"mean_accuracy\"][\"last\"]\n",
     "    best_trial_config = best_trial.config\n",
-    "    best_checkpoint = best_trial.checkpoint.value\n",
+    "    best_checkpoint = best_trial.checkpoint.dir_or_data\n",
     "\n",
     "    return best_accuracy, best_trial_config, best_checkpoint, num_examples"
    ]

python/ray/train/tests/test_tune.py

@@ -117,7 +117,7 @@ def train_func():
     TestTrainable = trainer.to_tune_trainable(train_func)
 
     [trial] = tune.run(TestTrainable).trials
-    checkpoint_path = trial.checkpoint.value
+    checkpoint_path = trial.checkpoint.dir_or_data
     assert os.path.exists(checkpoint_path)
     checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict()
     assert checkpoint["hello"] == "world"
@@ -138,7 +138,7 @@ def train_func(config):
     TestTrainable = trainer.to_tune_trainable(train_func)
 
     [trial] = tune.run(TestTrainable, config={"max_iter": 5}).trials
-    checkpoint_path = trial.checkpoint.value
+    checkpoint_path = trial.checkpoint.dir_or_data
     checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict()
     assert checkpoint["iter"] == 4
     analysis = tune.run(TestTrainable, config={"max_iter": 10}, restore=checkpoint_path)
@@ -164,7 +164,7 @@ def train_func():
     TestTrainable = trainer.to_tune_trainable(train_func)
 
     analysis = tune.run(TestTrainable, max_failures=3)
-    checkpoint_path = analysis.trials[0].checkpoint.value
+    checkpoint_path = analysis.trials[0].checkpoint.dir_or_data
     checkpoint = Checkpoint.from_directory(checkpoint_path).to_dict()
     assert checkpoint["iter"] == 3
 

python/ray/tune/analysis/experiment_analysis.py

@@ -440,7 +440,8 @@ def get_trial_checkpoints_paths(
             # Support metrics given as paths, e.g.
             # "info/learner/default_policy/policy_loss".
             return [
-                (c.value, unflattened_lookup(metric, c.result)) for c in checkpoints
+                (c.dir_or_data, unflattened_lookup(metric, c.metrics))
+                for c in checkpoints
             ]
         else:
             raise ValueError("trial should be a string or a Trial instance.")

python/ray/tune/callback.py

@@ -2,8 +2,8 @@
 from abc import ABCMeta
 import warnings
 
-from ray.tune.checkpoint_manager import _TuneCheckpoint
 from ray.util.annotations import PublicAPI, DeveloperAPI
+from ray.util.ml_utils.checkpoint_manager import _TrackedCheckpoint
 
 if TYPE_CHECKING:
     from ray.tune.trial import Trial
@@ -245,7 +245,7 @@ def on_checkpoint(
         iteration: int,
         trials: List["Trial"],
         trial: "Trial",
-        checkpoint: _TuneCheckpoint,
+        checkpoint: _TrackedCheckpoint,
         **info,
     ):
         """Called after a trial saved a checkpoint with Tune.