Created command context provider to set job group ID as part of run t…

…ags (mlflow#4521)

mlflow/tracking/context/databricks_command_context.py

@@ -0,0 +1,15 @@
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID
+
+
+class DatabricksCommandRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.get_job_group_id() is not None
+
+    def tags(self):
+        job_group_id = databricks_utils.get_job_group_id()
+        tags = {}
+        if job_group_id is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID] = job_group_id
+        return tags

mlflow/tracking/context/registry.py

@@ -7,6 +7,7 @@
 from mlflow.tracking.context.databricks_notebook_context import DatabricksNotebookRunContext
 from mlflow.tracking.context.databricks_job_context import DatabricksJobRunContext
 from mlflow.tracking.context.databricks_cluster_context import DatabricksClusterRunContext
+from mlflow.tracking.context.databricks_command_context import DatabricksCommandRunContext
 
 
 _logger = logging.getLogger(__name__)
@@ -53,6 +54,7 @@ def __iter__(self):
 _run_context_provider_registry.register(DatabricksNotebookRunContext)
 _run_context_provider_registry.register(DatabricksJobRunContext)
 _run_context_provider_registry.register(DatabricksClusterRunContext)
+_run_context_provider_registry.register(DatabricksCommandRunContext)
 
 _run_context_provider_registry.register_entrypoints()
 

mlflow/utils/databricks_utils.py

@@ -151,6 +151,16 @@ def get_cluster_id():
     return spark_session.conf.get("spark.databricks.clusterUsageTags.clusterId")
 
 
+def get_job_group_id():
+    try:
+        dbutils = _get_dbutils()
+        job_group_id = dbutils.entry_point.getJobGroupId()
+        if job_group_id is not None:
+            return job_group_id
+    except Exception:
+        return None
+
+
 def get_job_id():
     try:
         return _get_command_context().jobId().get()

mlflow/utils/mlflow_tags.py

@@ -27,6 +27,8 @@
 MLFLOW_DATABRICKS_WEBAPP_URL = "mlflow.databricks.webappURL"
 MLFLOW_DATABRICKS_RUN_URL = "mlflow.databricks.runURL"
 MLFLOW_DATABRICKS_CLUSTER_ID = "mlflow.databricks.cluster.id"
+# The unique ID of a command execution in a Databricks notebook
+MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID = "mlflow.databricks.notebook.commandID"
 # The SHELL_JOB_ID and SHELL_JOB_RUN_ID tags are used for tracking the
 # Databricks Job ID and Databricks Job Run ID associated with an MLflow Project run
 MLFLOW_DATABRICKS_SHELL_JOB_ID = "mlflow.databricks.shellJobID"

tests/tracking/context/test_databricks_command_context.py

@@ -0,0 +1,21 @@
+from unittest import mock
+
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID
+from mlflow.tracking.context.databricks_command_context import DatabricksCommandRunContext
+
+
+def test_databricks_command_run_context_in_context():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id", return_value="1"):
+        assert DatabricksCommandRunContext().in_context()
+
+
+def test_databricks_command_run_context_tags():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id") as job_group_id_mock:
+        assert DatabricksCommandRunContext().tags() == {
+            MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID: job_group_id_mock.return_value
+        }
+
+
+def test_databricks_command_run_context_tags_nones():
+    with mock.patch("mlflow.utils.databricks_utils.get_job_group_id", return_value=None):
+        assert DatabricksCommandRunContext().tags() == {}