From 45ea4bae4e9a8205f999aeb160940da17ccedc1d Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Sun, 26 Sep 2021 16:21:15 +0800
Subject: [PATCH] Add file lock for MLflowExpManager (#619)

---
 qlib/workflow/expm.py | 9 +++++++++
 setup.py              | 1 +
 2 files changed, 10 insertions(+)

diff --git a/qlib/workflow/expm.py b/qlib/workflow/expm.py
index 84cc6a13a..5d3596c4d 100644
--- a/qlib/workflow/expm.py
+++ b/qlib/workflow/expm.py
@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+from urllib.parse import urlparse
 import mlflow
+from filelock import FileLock
 from mlflow.exceptions import MlflowException
 from mlflow.entities import ViewType
 import os, logging
@@ -191,6 +193,13 @@ class ExpManager:
             if experiment_name is None:
                 experiment_name = self._default_exp_name
             logger.warning(f"No valid experiment found. Create a new experiment with name {experiment_name}.")
+
+            # NOTE: mlflow doesn't consider the lock for recording multiple runs
+            # So we supported it in the interface wrapper
+            pr = urlparse(self.uri)
+            if pr.scheme == "file":
+                with FileLock(os.path.join(pr.netloc, pr.path, "filelock")) as f:
+                    return self.create_exp(experiment_name), True
             return self.create_exp(experiment_name), True
 
     def _get_exp(self, experiment_id=None, experiment_name=None) -> Experiment:
diff --git a/setup.py b/setup.py
index 7553e7534..790ce0979 100644
--- a/setup.py
+++ b/setup.py
@@ -58,6 +58,7 @@ REQUIRED = [
     "pymongo==3.7.2",  # For task management
     "scikit-learn>=0.22",
     "dill",
+    "filelock",
 ]
 
 # Numpy include