star

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/__init__.py

@@ -0,0 +1 @@
+DEFAULT_EXPERIMENT_ID = "0"

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

@@ -0,0 +1,43 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class DefaultExperimentProvider:
+    """
+    Abstract base class for objects that provide the ID of an MLflow Experiment based on the
+    current client context. For example, when the MLflow client is running in a Databricks Job,
+    a provider is used to obtain the ID of the MLflow Experiment associated with the Job.
+
+    Usually the experiment_id is set explicitly by the user, but if the experiment is not set,
+    MLflow computes a default experiment id based on different contexts.
+    When an experiment is created via the fluent ``mlflow.start_run`` method, MLflow iterates
+    through the registered ``DefaultExperimentProvider``s until it finds one whose
+    ``in_context()`` method returns ``True``; MLflow then calls the provider's
+    ``get_experiment_id()`` method and uses the resulting experiment ID for Tracking operations.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if the MLflow client is running in a context where this provider can
+        identify an associated MLflow Experiment ID.
+
+        Returns:
+            True if the MLflow client is running in a context where the provider
+            can identify an associated MLflow Experiment ID. False otherwise.
+
+        """
+
+    @abstractmethod
+    def get_experiment_id(self):
+        """Provide the MLflow Experiment ID for the current MLflow client context.
+
+        Assumes that ``in_context()`` is ``True``.
+
+        Returns:
+            The ID of the MLflow Experiment associated with the current context.
+
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/databricks_notebook_experiment_provider.py

@@ -0,0 +1,44 @@
+from mlflow.exceptions import MlflowException
+from mlflow.protos import databricks_pb2
+from mlflow.tracking.client import MlflowClient
+from mlflow.tracking.default_experiment.abstract_context import DefaultExperimentProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import MLFLOW_EXPERIMENT_SOURCE_ID, MLFLOW_EXPERIMENT_SOURCE_TYPE
+
+
+class DatabricksNotebookExperimentProvider(DefaultExperimentProvider):
+    _resolved_notebook_experiment_id = None
+
+    def in_context(self):
+        return databricks_utils.is_in_databricks_notebook()
+
+    def get_experiment_id(self):
+        if DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id:
+            return DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id
+
+        source_notebook_id = databricks_utils.get_notebook_id()
+        source_notebook_name = databricks_utils.get_notebook_path()
+        tags = {
+            MLFLOW_EXPERIMENT_SOURCE_ID: source_notebook_id,
+        }
+
+        if databricks_utils.is_in_databricks_repo_notebook():
+            tags[MLFLOW_EXPERIMENT_SOURCE_TYPE] = "REPO_NOTEBOOK"
+
+        # With the presence of the source id, the following is a get or create in which it will
+        # return the corresponding experiment if one exists for the repo notebook.
+        # For non-repo notebooks, it will raise an exception and we will use source_notebook_id
+        try:
+            experiment_id = MlflowClient().create_experiment(source_notebook_name, None, tags)
+        except MlflowException as e:
+            if e.error_code == databricks_pb2.ErrorCode.Name(
+                databricks_pb2.INVALID_PARAMETER_VALUE
+            ):
+                # If determined that it is not a repo notebook
+                experiment_id = source_notebook_id
+            else:
+                raise e
+
+        DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id = experiment_id
+
+        return experiment_id

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

@@ -0,0 +1,75 @@
+import logging
+import warnings
+
+from mlflow.tracking.default_experiment import DEFAULT_EXPERIMENT_ID
+from mlflow.tracking.default_experiment.databricks_notebook_experiment_provider import (
+    DatabricksNotebookExperimentProvider,
+)
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+# Listed below are the list of providers, which are used to provide MLflow Experiment IDs based on
+# the current context where the MLflow client is running when the user has not explicitly set
+# an experiment. The order below is the order in which the these providers are registered.
+_EXPERIMENT_PROVIDERS = (DatabricksNotebookExperimentProvider,)
+
+
+class DefaultExperimentProviderRegistry:
+    """Registry for default experiment provider implementations
+
+    This class allows the registration of default experiment providers, which are used to provide
+    MLflow Experiment IDs based on the current context where the MLflow client is running when
+    the user has not explicitly set an experiment. Implementations declared though the entrypoints
+    `mlflow.default_experiment_provider` group can be automatically registered through the
+    `register_entrypoints` method.
+    """
+
+    def __init__(self):
+        self._registry = []
+
+    def register(self, default_experiment_provider_cls):
+        self._registry.append(default_experiment_provider_cls())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.default_experiment_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    "Failure attempting to register default experiment"
+                    + f'context provider "{entrypoint.name}": {exc}',
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_default_experiment_provider_registry = DefaultExperimentProviderRegistry()
+for exp_provider in _EXPERIMENT_PROVIDERS:
+    _default_experiment_provider_registry.register(exp_provider)
+
+_default_experiment_provider_registry.register_entrypoints()
+
+
+def get_experiment_id():
+    """Get an experiment ID for the current context.
+
+    The experiment ID is fetched by querying providers, in the order that they were registered.
+    This function iterates through all default experiment context providers in the registry.
+
+    Returns:
+        An experiment_id.
+    """
+
+    experiment_id = DEFAULT_EXPERIMENT_ID
+    for provider in _default_experiment_provider_registry:
+        try:
+            if provider.in_context():
+                experiment_id = provider.get_experiment_id()
+                break
+        except Exception as e:
+            _logger.warning("Encountered unexpected error while getting experiment_id: %s", e)
+
+    return experiment_id