Simplify TSDataset and async recorder

qlib/utils/paral.py

@@ -1,9 +1,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import pandas as pd
+from functools import partial
+from threading import Thread
+from typing import Callable
+
 from joblib import Parallel, delayed
 from joblib._parallel_backends import MultiprocessingBackend
+import pandas as pd
+
+from queue import Queue
 
 
 class ParallelExt(Parallel):
@@ -46,3 +52,54 @@ def datetime_groupby_apply(df, apply_func, axis=0, level="datetime", resample_ru
         return pd.concat(dfs, axis=axis).sort_index()
     else:
         return _naive_group_apply(df)
+
+
+class AsyncCaller:
+    """
+    This AsyncCaller tries to make it easier to async call
+
+    Currently, it is used in MLflowRecorder to make functions like `log_params` async
+
+    NOTE:
+    - This caller didn't consider the return value
+    """
+
+    STOP_MARK = "__STOP"
+
+    def __init__(self) -> None:
+        self._q = Queue()
+        self._stop = False
+        self._t = Thread(target=self.run)
+        self._t.start()
+
+    def close(self):
+        self._q.put(self.STOP_MARK)
+
+    def run(self):
+        while True:
+            data = self._q.get()
+            if data == self.STOP_MARK:
+                break
+            else:
+                data()
+
+    def __call__(self, func, *args, **kwargs):
+        self._q.put(partial(func, *args, **kwargs))
+
+    def wait(self, close=True):
+        if close:
+            self.close()
+        self._t.join()
+
+    @staticmethod
+    def async_dec(ac_attr):
+        def decorator_func(func):
+            def wrapper(self, *args, **kwargs):
+                if isinstance(getattr(self, ac_attr, None), Callable):
+                    return getattr(self, ac_attr)(func, self, *args, **kwargs)
+                else:
+                    return func(self, *args, **kwargs)
+
+            return wrapper
+
+        return decorator_func