diff --git a/python/ray/serve/examples/doc/tutorial_rllib.py b/python/ray/serve/examples/doc/tutorial_rllib.py index d9a794a06..e2ed0626a 100644 --- a/python/ray/serve/examples/doc/tutorial_rllib.py +++ b/python/ray/serve/examples/doc/tutorial_rllib.py @@ -45,7 +45,7 @@ def train_ppo_model(): # Train for one iteration trainer.train() trainer.save("/tmp/rllib_checkpoint") - return "/tmp/rllib_checkpoint/checkpoint_1/checkpoint-1" + return "/tmp/rllib_checkpoint/checkpoint_000001/checkpoint-1" checkpoint_path = train_ppo_model() diff --git a/python/ray/tune/tests/test_experiment_analysis.py b/python/ray/tune/tests/test_experiment_analysis.py index 911c5c32a..059f7d296 100644 --- a/python/ray/tune/tests/test_experiment_analysis.py +++ b/python/ray/tune/tests/test_experiment_analysis.py @@ -108,14 +108,15 @@ class ExperimentAnalysisSuite(unittest.TestCase): best_trial = self.ea.get_best_trial(self.metric, mode="max") checkpoints_metrics = self.ea.get_trial_checkpoints_paths(best_trial) logdir = self.ea.get_best_logdir(self.metric, mode="max") - expected_path = os.path.join(logdir, "checkpoint_1", "checkpoint") + expected_path = os.path.join(logdir, "checkpoint_000001", "checkpoint") assert checkpoints_metrics[0][0] == expected_path assert checkpoints_metrics[0][1] == 1 def testGetTrialCheckpointsPathsByPath(self): logdir = self.ea.get_best_logdir(self.metric, mode="max") checkpoints_metrics = self.ea.get_trial_checkpoints_paths(logdir) - expected_path = os.path.join(logdir, "checkpoint_1/", "checkpoint") + expected_path = os.path.join(logdir, "checkpoint_000001/", + "checkpoint") assert checkpoints_metrics[0][0] == expected_path assert checkpoints_metrics[0][1] == 1 @@ -123,7 +124,7 @@ class ExperimentAnalysisSuite(unittest.TestCase): best_trial = self.ea.get_best_trial(self.metric, mode="max") paths = self.ea.get_trial_checkpoints_paths(best_trial, self.metric) logdir = self.ea.get_best_logdir(self.metric, mode="max") - expected_path = os.path.join(logdir, "checkpoint_1", "checkpoint") + expected_path = os.path.join(logdir, "checkpoint_000001", "checkpoint") assert paths[0][0] == expected_path assert paths[0][1] == best_trial.metric_analysis[self.metric]["last"] @@ -131,7 +132,7 @@ class ExperimentAnalysisSuite(unittest.TestCase): best_trial = self.ea.get_best_trial(self.metric, mode="max") logdir = self.ea.get_best_logdir(self.metric, mode="max") paths = self.ea.get_trial_checkpoints_paths(best_trial, self.metric) - expected_path = os.path.join(logdir, "checkpoint_1", "checkpoint") + expected_path = os.path.join(logdir, "checkpoint_000001", "checkpoint") assert paths[0][0] == expected_path assert paths[0][1] == best_trial.metric_analysis[self.metric]["last"] diff --git a/python/ray/tune/tests/test_tune_save_restore.py b/python/ray/tune/tests/test_tune_save_restore.py index 63122e859..5d25080c4 100644 --- a/python/ray/tune/tests/test_tune_save_restore.py +++ b/python/ray/tune/tests/test_tune_save_restore.py @@ -87,7 +87,7 @@ class SerialTuneRelativeLocalDirTest(unittest.TestCase): self.assertTrue(os.path.isdir(abs_trial_dir)) self.assertTrue( os.path.isfile( - os.path.join(abs_trial_dir, "checkpoint_1/checkpoint-1"))) + os.path.join(abs_trial_dir, "checkpoint_000001/checkpoint-1"))) def _restore(self, exp_name, local_dir, absolute_local_dir): trial_name, abs_trial_dir = self._get_trial_dir( @@ -95,7 +95,7 @@ class SerialTuneRelativeLocalDirTest(unittest.TestCase): checkpoint_path = os.path.join( local_dir, exp_name, trial_name, - "checkpoint_1/checkpoint-1") # Relative checkpoint path + "checkpoint_000001/checkpoint-1") # Relative checkpoint path # The file tune would find. The absolute checkpoint path. tune_find_file = os.path.abspath(os.path.expanduser(checkpoint_path)) diff --git a/python/ray/tune/utils/trainable.py b/python/ray/tune/utils/trainable.py index cc61e3a83..9872788bf 100644 --- a/python/ray/tune/utils/trainable.py +++ b/python/ray/tune/utils/trainable.py @@ -101,14 +101,15 @@ class TrainableUtil: Args: checkpoint_dir (str): Path to checkpoint directory. - index (str): A subdirectory will be created + index (int|str): A subdirectory will be created at the checkpoint directory named 'checkpoint_{index}'. override (bool): Deletes checkpoint_dir before creating a new one. """ suffix = "checkpoint" if index is not None: - suffix += "_{}".format(index) + suffix += f"_{index:06d}" if isinstance(index, + int) else f"_{index}" checkpoint_dir = os.path.join(checkpoint_dir, suffix) if override and os.path.exists(checkpoint_dir): diff --git a/rllib/scripts.py b/rllib/scripts.py index 58e9ad447..5f7b5811e 100644 --- a/rllib/scripts.py +++ b/rllib/scripts.py @@ -10,7 +10,7 @@ Example usage for training: rllib train --run DQN --env CartPole-v0 Example usage for rollout: - rllib rollout /trial_dir/checkpoint_1/checkpoint-1 --run DQN + rllib rollout /trial_dir/checkpoint_000001/checkpoint-1 --run DQN """ diff --git a/rllib/tests/test_rollout.py b/rllib/tests/test_rollout.py index 62e3c2fd7..f16027b66 100644 --- a/rllib/tests/test_rollout.py +++ b/rllib/tests/test_rollout.py @@ -40,7 +40,7 @@ def rollout_test(algo, env="CartPole-v0", test_episode_rollout=False): "}' --stop='{\"training_iteration\": 1}'" + " --env={}".format(env)) - checkpoint_path = os.popen("ls {}/default/*/checkpoint_1/" + checkpoint_path = os.popen("ls {}/default/*/checkpoint_000001/" "checkpoint-1".format(tmp_dir)).read()[:-1] if not os.path.exists(checkpoint_path): sys.exit(1)