genesis-flow 1.0.0__py3-none-any.whl

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.
+        """

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

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/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

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

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

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()),
+        }

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}

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

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())

mlflow/tracking/default_experiment/__init__.py

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

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.
+
+        """

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

mlflow/tracking/default_experiment/registry.py

@@ -0,0 +1,75 @@
+import logging
+import warnings
+from typing import Optional
+
+from mlflow.tracking import get_tracking_uri
+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
+from mlflow.utils.uri import is_databricks_uri
+
+_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() -> Optional[str]:
+    """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.
+    """
+    for provider in _default_experiment_provider_registry:
+        try:
+            if provider.in_context():
+                return provider.get_experiment_id()
+        except Exception as e:
+            _logger.warning("Encountered unexpected error while getting experiment_id: %s", e)
+
+    return DEFAULT_EXPERIMENT_ID if not is_databricks_uri(get_tracking_uri()) else None