Fix databricks runtime checking issue in databricks DCS cluster (mlfl…

…ow#11483)

Signed-off-by: Weichen Xu <[email protected]>

mlflow/utils/databricks_utils.py

@@ -183,8 +183,22 @@ def is_in_databricks_repo_notebook():
         return False
 
 
+_DATABRICKS_VERSION_FILE_PATH = "/databricks/DBR_VERSION"
+
+
+def get_databricks_runtime_version():
+    if ver := os.environ.get("DATABRICKS_RUNTIME_VERSION"):
+        return ver
+    if os.path.exists(_DATABRICKS_VERSION_FILE_PATH):
+        # In Databricks DCS cluster, it doesn't have DATABRICKS_RUNTIME_VERSION
+        # environment variable, we have to read version from the version file.
+        with open(_DATABRICKS_VERSION_FILE_PATH) as f:
+            return f.read().strip()
+    return None
+
+
 def is_in_databricks_runtime():
-    return "DATABRICKS_RUNTIME_VERSION" in os.environ
+    return get_databricks_runtime_version() is not None
 
 
 def is_dbfs_fuse_available():
@@ -805,7 +819,7 @@ class DatabricksRuntimeVersion(NamedTuple):
 
     @classmethod
     def parse(cls):
-        dbr_version = os.environ["DATABRICKS_RUNTIME_VERSION"]
+        dbr_version = get_databricks_runtime_version()
         try:
             dbr_version_splits = dbr_version.split(".", maxsplit=2)
             if dbr_version_splits[0] == "client":