[tune] get checkpoints paths for a trial after tuning (#6643)

python/ray/tune/trainable.py

@@ -3,8 +3,10 @@ from datetime import datetime
 import copy
 import io
 import logging
+import glob
 import os
 import pickle
+import pandas as pd
 from six import string_types
 import shutil
 import tempfile
@@ -73,6 +75,36 @@ class TrainableUtil:
         # Drop marker in directory to identify it as a checkpoint dir.
         open(os.path.join(checkpoint_dir, ".is_checkpoint"), "a").close()
 
+    @staticmethod
+    def get_checkpoints_paths(logdir):
+        """ Finds the checkpoints within a specific folder.
+
+        Returns a pandas DataFrame of training iterations and checkpoint
+        paths within a specific folder.
+
+        Raises:
+            FileNotFoundError if the directory is not found.
+        """
+        marker_paths = glob.glob(
+            os.path.join(logdir, "checkpoint_*/.is_checkpoint"))
+        iter_chkpt_pairs = []
+        for marker_path in marker_paths:
+            chkpt_dir = os.path.dirname(marker_path)
+            metadata_file = glob.glob(
+                os.path.join(chkpt_dir, "*.tune_metadata"))
+            if len(metadata_file) != 1:
+                raise ValueError(
+                    "{} has zero or more than one tune_metadata.".format(
+                        chkpt_dir))
+
+            chkpt_path = metadata_file[0][:-len(".tune_metadata")]
+            chkpt_iter = int(chkpt_dir[chkpt_dir.rfind("_") + 1:])
+            iter_chkpt_pairs.append([chkpt_iter, chkpt_path])
+
+        chkpt_df = pd.DataFrame(
+            iter_chkpt_pairs, columns=["training_iteration", "chkpt_path"])
+        return chkpt_df
+
 
 class Trainable:
     """Abstract class for trainable models, functions, etc.