star

venv/lib/python3.9/site-packages/mlflow/tracking/context/abstract_context.py

@@ -0,0 +1,35 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RunContextProvider:
+    """
+    Abstract base class for context provider objects specifying custom tags at run-creation time
+    (e.g. tags specifying the git repo with which the run is associated).
+
+    When a run is created via the fluent ``mlflow.start_run`` method, MLflow iterates through all
+    registered RunContextProviders. For each context provider where ``in_context`` returns ``True``,
+    MLflow calls the ``tags`` method on the context provider to compute context tags for the run.
+    All context tags are then merged together and set on the newly-created run.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if MLflow is running in this context.
+
+        Returns:
+            bool indicating if in this context.
+
+        """
+
+    @abstractmethod
+    def tags(self):
+        """Generate context-specific tags.
+
+        Returns:
+            dict of tags.
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_cluster_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_CLUSTER_ID
+
+
+class DatabricksClusterRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_cluster()
+
+    def tags(self):
+        cluster_id = databricks_utils.get_cluster_id()
+        tags = {}
+        if cluster_id is not None:
+            tags[MLFLOW_DATABRICKS_CLUSTER_ID] = cluster_id
+        return tags

venv/lib/python3.9/site-packages/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

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_job_context.py

@@ -0,0 +1,49 @@
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_JOB_ID,
+    MLFLOW_DATABRICKS_JOB_RUN_ID,
+    MLFLOW_DATABRICKS_JOB_TYPE,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+    MLFLOW_DATABRICKS_WORKSPACE_ID,
+    MLFLOW_DATABRICKS_WORKSPACE_URL,
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+)
+
+
+class DatabricksJobRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_job()
+
+    def tags(self):
+        job_id = databricks_utils.get_job_id()
+        job_run_id = databricks_utils.get_job_run_id()
+        job_type = databricks_utils.get_job_type()
+        webapp_url = databricks_utils.get_webapp_url()
+        workspace_url = databricks_utils.get_workspace_url()
+        workspace_url_fallback, workspace_id = databricks_utils.get_workspace_info_from_dbutils()
+        tags = {
+            MLFLOW_SOURCE_NAME: (
+                f"jobs/{job_id}/run/{job_run_id}"
+                if job_id is not None and job_run_id is not None
+                else None
+            ),
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB),
+        }
+        if job_id is not None:
+            tags[MLFLOW_DATABRICKS_JOB_ID] = job_id
+        if job_run_id is not None:
+            tags[MLFLOW_DATABRICKS_JOB_RUN_ID] = job_run_id
+        if job_type is not None:
+            tags[MLFLOW_DATABRICKS_JOB_TYPE] = job_type
+        if webapp_url is not None:
+            tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url
+        if workspace_url is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url
+        elif workspace_url_fallback is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url_fallback
+        if workspace_id is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_ID] = workspace_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_notebook_context.py

@@ -0,0 +1,41 @@
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_NOTEBOOK_ID,
+    MLFLOW_DATABRICKS_NOTEBOOK_PATH,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+    MLFLOW_DATABRICKS_WORKSPACE_ID,
+    MLFLOW_DATABRICKS_WORKSPACE_URL,
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+)
+
+
+class DatabricksNotebookRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_notebook()
+
+    def tags(self):
+        notebook_id = databricks_utils.get_notebook_id()
+        notebook_path = databricks_utils.get_notebook_path()
+        webapp_url = databricks_utils.get_webapp_url()
+        workspace_url = databricks_utils.get_workspace_url()
+        workspace_url_fallback, workspace_id = databricks_utils.get_workspace_info_from_dbutils()
+        tags = {
+            MLFLOW_SOURCE_NAME: notebook_path,
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK),
+        }
+        if notebook_id is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_ID] = notebook_id
+        if notebook_path is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_PATH] = notebook_path
+        if webapp_url is not None:
+            tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url
+        if workspace_url is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url
+        elif workspace_url_fallback is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url_fallback
+        if workspace_id is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_ID] = workspace_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_repo_context.py

@@ -0,0 +1,43 @@
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_GIT_REPO_COMMIT,
+    MLFLOW_DATABRICKS_GIT_REPO_PROVIDER,
+    MLFLOW_DATABRICKS_GIT_REPO_REFERENCE,
+    MLFLOW_DATABRICKS_GIT_REPO_REFERENCE_TYPE,
+    MLFLOW_DATABRICKS_GIT_REPO_RELATIVE_PATH,
+    MLFLOW_DATABRICKS_GIT_REPO_STATUS,
+    MLFLOW_DATABRICKS_GIT_REPO_URL,
+)
+
+
+class DatabricksRepoRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_repo()
+
+    def tags(self):
+        tags = {}
+        git_repo_url = databricks_utils.get_git_repo_url()
+        git_repo_provider = databricks_utils.get_git_repo_provider()
+        git_repo_commit = databricks_utils.get_git_repo_commit()
+        git_repo_relative_path = databricks_utils.get_git_repo_relative_path()
+        git_repo_reference = databricks_utils.get_git_repo_reference()
+        git_repo_reference_type = databricks_utils.get_git_repo_reference_type()
+        git_repo_status = databricks_utils.get_git_repo_status()
+
+        if git_repo_url is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_URL] = git_repo_url
+        if git_repo_provider is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_PROVIDER] = git_repo_provider
+        if git_repo_commit is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_COMMIT] = git_repo_commit
+        if git_repo_relative_path is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_RELATIVE_PATH] = git_repo_relative_path
+        if git_repo_reference is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_REFERENCE] = git_repo_reference
+        if git_repo_reference_type is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_REFERENCE_TYPE] = git_repo_reference_type
+        if git_repo_status is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_STATUS] = git_repo_status
+
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/default_context.py

@@ -0,0 +1,51 @@
+import getpass
+import sys
+
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils.credentials import read_mlflow_creds
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+    MLFLOW_USER,
+)
+
+_DEFAULT_USER = "unknown"
+
+
+def _get_user():
+    """Get the current computer username."""
+    try:
+        return getpass.getuser()
+    except ImportError:
+        return _DEFAULT_USER
+
+
+def _get_main_file():
+    if len(sys.argv) > 0:
+        return sys.argv[0]
+    return None
+
+
+def _get_source_name():
+    main_file = _get_main_file()
+    if main_file is not None:
+        return main_file
+    return "<console>"
+
+
+def _get_source_type():
+    return SourceType.LOCAL
+
+
+class DefaultRunContext(RunContextProvider):
+    def in_context(self):
+        return True
+
+    def tags(self):
+        creds = read_mlflow_creds()
+        return {
+            MLFLOW_USER: creds.username or _get_user(),
+            MLFLOW_SOURCE_NAME: _get_source_name(),
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(_get_source_type()),
+        }

venv/lib/python3.9/site-packages/mlflow/tracking/context/git_context.py

@@ -0,0 +1,32 @@
+import logging
+
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.tracking.context.default_context import _get_main_file
+from mlflow.utils.git_utils import get_git_commit
+from mlflow.utils.mlflow_tags import MLFLOW_GIT_COMMIT
+
+_logger = logging.getLogger(__name__)
+
+
+def _get_source_version():
+    main_file = _get_main_file()
+    if main_file is not None:
+        return get_git_commit(main_file)
+    return None
+
+
+class GitRunContext(RunContextProvider):
+    def __init__(self):
+        self._cache = {}
+
+    @property
+    def _source_version(self):
+        if "source_version" not in self._cache:
+            self._cache["source_version"] = _get_source_version()
+        return self._cache["source_version"]
+
+    def in_context(self):
+        return self._source_version is not None
+
+    def tags(self):
+        return {MLFLOW_GIT_COMMIT: self._source_version}

venv/lib/python3.9/site-packages/mlflow/tracking/context/registry.py

@@ -0,0 +1,98 @@
+import logging
+import warnings
+from typing import Optional
+
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.tracking.context.databricks_cluster_context import DatabricksClusterRunContext
+from mlflow.tracking.context.databricks_command_context import DatabricksCommandRunContext
+from mlflow.tracking.context.databricks_job_context import DatabricksJobRunContext
+from mlflow.tracking.context.databricks_notebook_context import DatabricksNotebookRunContext
+from mlflow.tracking.context.databricks_repo_context import DatabricksRepoRunContext
+from mlflow.tracking.context.default_context import DefaultRunContext
+from mlflow.tracking.context.git_context import GitRunContext
+from mlflow.tracking.context.system_environment_context import SystemEnvironmentContext
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+
+
+class RunContextProviderRegistry:
+    """Registry for run context provider implementations
+
+    This class allows the registration of a run context provider which can be used to infer meta
+    information about the context of an MLflow experiment run. Implementations declared though the
+    entrypoints `mlflow.run_context_provider` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    Registered run context providers can return tags that override those implemented in the core
+    library, however the order in which plugins are resolved is undefined.
+    """
+
+    def __init__(self):
+        self._registry = []
+
+    def register(self, run_context_provider_cls):
+        self._registry.append(run_context_provider_cls())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.run_context_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register context provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_run_context_provider_registry = RunContextProviderRegistry()
+_run_context_provider_registry.register(DefaultRunContext)
+_run_context_provider_registry.register(GitRunContext)
+_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(DatabricksRepoRunContext)
+_run_context_provider_registry.register(SystemEnvironmentContext)
+
+_run_context_provider_registry.register_entrypoints()
+
+
+def resolve_tags(tags=None, ignore: Optional[list[RunContextProvider]] = None):
+    """Generate a set of tags for the current run context. Tags are resolved in the order,
+    contexts are registered. Argument tags are applied last.
+
+    This function iterates through all run context providers in the registry. Additional context
+    providers can be registered as described in
+    :py:class:`mlflow.tracking.context.RunContextProvider`.
+
+    Args:
+        tags: A dictionary of tags to override. If specified, tags passed in this argument will
+            override those inferred from the context.
+        ignore: A list of RunContextProvider classes to exclude from the resolution.
+
+    Returns:
+        A dictionary of resolved tags.
+    """
+    ignore = ignore or []
+    all_tags = {}
+    for provider in _run_context_provider_registry:
+        if any(isinstance(provider, ig) for ig in ignore):
+            continue
+
+        try:
+            if provider.in_context():
+                all_tags.update(provider.tags())
+        except Exception as e:
+            _logger.warning("Encountered unexpected error during resolving tags: %s", e)
+
+    if tags is not None:
+        all_tags.update(tags)
+
+    return all_tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/system_environment_context.py

@@ -0,0 +1,15 @@
+import json
+
+from mlflow.environment_variables import MLFLOW_RUN_CONTEXT
+from mlflow.tracking.context.abstract_context import RunContextProvider
+
+# The constant MLFLOW_RUN_CONTEXT_ENV_VAR is marked as @developer_stable
+MLFLOW_RUN_CONTEXT_ENV_VAR = MLFLOW_RUN_CONTEXT.name
+
+
+class SystemEnvironmentContext(RunContextProvider):
+    def in_context(self):
+        return MLFLOW_RUN_CONTEXT.defined
+
+    def tags(self):
+        return json.loads(MLFLOW_RUN_CONTEXT.get())