Update R related codes

qlib/workflow/exp.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 
 import mlflow
+from datetime import datetime
 from pathlib import Path
 from .recorder import MLflowRecorder
 from ..log import get_module_logger
@@ -11,12 +12,13 @@ logger = get_module_logger("workflow", "INFO")
 
 class Experiment:
     """
-    Thie is the `Experiment` class for each experiment being run. The API is designed
+    Thie is the `Experiment` class for each experiment being run. The API is designed similar to mlflow.
+    (The link: https://mlflow.org/docs/latest/python_api/mlflow.html)
     """
 
-    def __init__(self):
-        self.name = None
-        self.id = None
+    def __init__(self, id, name):
+        self.id = id
+        self.name = name
         self.active_recorder = None  # only one recorder can running each time
         self.recorders = dict()  # recorder id -> object
 
@@ -32,16 +34,14 @@ class Experiment:
         output["class"] = "Experiment"
         output["id"] = self.id
         output["name"] = self.name
-        output["active_recorder"] = self.active_recorder.id
+        output["active_recorder"] = self.active_recorder.id if self.active_recorder is not None else None
         output["recorders"] = list(self.recorders.keys())
+        return output
 
     def start(self):
         """
         Start the experiment.
 
-        Parameters
-        ----------
-
         Returns
         -------
         A running recorder instance.
@@ -63,9 +63,6 @@ class Experiment:
         """
         Create a recorder for each experiment.
 
-        Parameters
-        ----------
-
         Returns
         -------
         A recorder object.
@@ -124,13 +121,31 @@ class Experiment:
         """
         raise NotImplementedError(f"Please implement the `get_recorder` method.")
 
+    def list_recorders(self):
+        """
+        List all the existing recorders of this experiment.
+
+        Returns
+        -------
+        A dictionary (id -> recorder) of recorder information that being stored.
+        """
+        raise NotImplementedError(f"Please implement the `list_recorders` method.")
+
 
 class MLflowExperiment(Experiment):
     """
     Use mlflow to implement Experiment.
     """
 
+    def __init__(self, id, name, uri):
+        super(MLflowExperiment, self).__init__(id, name)
+        self._uri = uri
+        self._total_recorders = 0
+        self._default_name = None
+
     def start(self):
+        # get all the recorders of the experiment
+        self.recorders = self.list_recorders()
         # set up recorder
         recorder = self.create_recorder()
         self.active_recorder = recorder
@@ -138,17 +153,22 @@ class MLflowExperiment(Experiment):
         run = self.active_recorder.start_run()
         # store the recorder
         self.recorders[self.active_recorder.id] = recorder
+        self._total_recorders += 1  # update recorder num
+        logger.info(f"Experiment {self.id} starts running ...")
+
         return self.active_recorder
 
     def end(self, status):
         if self.active_recorder is not None:
             self.active_recorder.end_run(status)
             self.active_recorder = None
+            self._total_recorders -= 1
 
     def create_recorder(self):
         num = len(self.recorders)
         name = "Recorder_{}".format(num + 1)
-        recorder = MLflowRecorder(name, self.id)
+        recorder = MLflowRecorder(name, self.id, self._uri)
+
         return recorder
 
     def search_records(self, **kwargs):
@@ -156,21 +176,92 @@ class MLflowExperiment(Experiment):
         run_view_type = 1 if kwargs.get("run_view_type") is None else kwargs.get("run_view_type")
         max_results = 100000 if kwargs.get("max_results") is None else kwargs.get("max_results")
         order_by = kwargs.get("order_by")
+
         return mlflow.search_runs([self.id], filter_string, run_view_type, max_results, order_by)
 
-    def delete_recorder(self, recorder_id):
-        mlflow.delete_run(recorder_id)
-        self.recorders = [r for r in self.recorders if r.id == recorder_id]
+    def delete_recorder(self, recorder_id=None, recorder_name=None):
+        assert (
+            recorder_id is not None or recorder_name is not None
+        ), "Please input a valid recorder id or name before deleting."
+        try:
+            if recorder_id is not None:
+                mlflow.delete_run(recorder_id)
+                self.recorders = [r for r in self.recorders if r == recorder_id]
+            else:
+                for r in self.recorders:
+                    if self.recorders[r].name == recorder_name:
+                        recorder_id = r
+                        break
+                mlflow.delete_run(recorder_id)
+        except:
+            raise Exception(
+                "Something went wrong when deleting recorder. Please check if the name/id of the recorder is correct."
+            )
 
-    def get_recorder(self, recorder_id=None, recorder_name=None):
-        if recorder_id is not None:
-            return self.recorders[recorder_id]
-        elif recorder_name is not None:
-            for rid in self.recorders:
-                if self.recorders[rid].name == recorder_name:
-                    return self.recorders[rid]
-        elif self.active_recorder is None:
-            raise Exception("No valid active recorder exists. Please make sure the experiment is running.")
+    def get_recorder(self, recorder_id=None, recorder_name=None, create=True):
+        if recorder_id is None and recorder_name is None:
+            if self.active_recorder:
+                return self.active_recorder
+            else:
+                if create:
+                    self.start()
+                    logger.warning(
+                        f"Recorder {self.active_recorder.id} is running under the experiment with name {self.name}..."
+                    )
+                    return self.active_recorder
+                else:
+                    raise Exception(
+                        "Something went wrong when retrieving recorders. Please check if QlibRecorder is running or the name/id of the recorder is correct."
+                    )
         else:
-            logger.info("No experiment id or name is given. Return the current active experiment.")
-            return self.active_recorder
+            if recorder_id is not None:
+                if recorder_id in self.recorders:
+                    return self.recorders[recorder_id]
+                else:
+                    # mlflow does not support create a run with given id
+                    raise Exception(
+                        "Something went wrong when retrieving recorders. Please check if QlibRecorder is running or the name/id of the recorder is correct."
+                    )
+            else:
+                for rid in self.recorders:
+                    if self.recorders[rid].name == recorder_name:
+                        return self.recorders[rid]
+                if create:
+                    self.recorders = self.list_recorders()
+                    logger.warning(f"No valid recorder found. Create a new recorder with name {recorder_name}.")
+                    recorder = self.create_recorder()
+                    recorder.name = recorder_name
+                    recorder.start_run()
+                    return recorder
+                else:
+                    raise Exception(
+                        "Something went wrong when retrieving experiments. Please check if QlibRecorder is running or the name/id of the experiment is correct."
+                    )
+
+    def list_recorders(self):
+        client = mlflow.tracking.MlflowClient(tracking_uri=self._uri)
+        runs = client.list_run_infos(self.id)[::-1]
+        recorders = dict()
+        self._total_recorders = len(runs)
+        for i in range(len(runs)):
+            rid = runs[i].run_id
+            status = runs[i].status
+            start_time = runs[i].start_time
+            end_time = runs[i].end_time
+            recorder = MLflowRecorder(f"Recorder_{i+1}", self.id, self._uri)
+            recorder.id = rid
+            recorder.status = status
+            recorder.start_time = (
+                datetime.fromtimestamp(float(start_time) / 1000.0).strftime("%Y-%m-%d %H:%M:%S")
+                if start_time is not None
+                else None
+            )
+            recorder.end_time = (
+                datetime.fromtimestamp(float(end_time) / 1000.0).strftime("%Y-%m-%d %H:%M:%S")
+                if end_time is not None
+                else None
+            )
+            recorder._uri = self._uri
+            recorders[rid] = recorder
+
+        return recorders