genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracking/_tracking_service/registry.py

@@ -0,0 +1,56 @@
+import threading
+from functools import lru_cache
+
+from mlflow.tracking.registry import StoreRegistry
+
+_building_store_lock = threading.Lock()
+
+
+class TrackingStoreRegistry(StoreRegistry):
+    """Scheme-based registry for tracking store implementations
+
+    This class allows the registration of a function or class to provide an
+    implementation for a given scheme of `store_uri` through the `register`
+    methods. Implementations declared though the entrypoints
+    `mlflow.tracking_store` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    When instantiating a store through the `get_store` method, the scheme of
+    the store URI provided (or inferred from environment) will be used to
+    select which implementation to instantiate, which will be called with same
+    arguments passed to the `get_store` method.
+    """
+
+    def __init__(self):
+        super().__init__("mlflow.tracking_store")
+
+    def get_store(self, store_uri=None, artifact_uri=None):
+        """Get a store from the registry based on the scheme of store_uri
+
+        Args:
+            store_uri: The store URI. If None, it will be inferred from the environment. This URI
+                is used to select which tracking store implementation to instantiate and
+                is passed to the constructor of the implementation.
+            artifact_uri: Artifact repository URI. Passed through to the tracking store
+                implementation.
+
+        Returns:
+            An instance of `mlflow.store.tracking.AbstractStore` that fulfills the store URI
+            requirements.
+
+        """
+        from mlflow.tracking._tracking_service import utils
+
+        resolved_store_uri = utils._resolve_tracking_uri(store_uri)
+        return self._get_store_with_resolved_uri(resolved_store_uri, artifact_uri)
+
+    @lru_cache(maxsize=100)
+    def _get_store_with_resolved_uri(self, resolved_store_uri, artifact_uri=None):
+        """
+        Retrieve the store associated with a resolved (non-None) store URI and an artifact URI.
+        Caching is done on resolved URIs because the meaning of an unresolved (None) URI may change
+        depending on external configuration, such as environment variables
+        """
+        with _building_store_lock:
+            builder = self.get_store_builder(resolved_store_uri)
+            return builder(store_uri=resolved_store_uri, artifact_uri=artifact_uri)

mlflow/tracking/_tracking_service/utils.py

@@ -0,0 +1,275 @@
+import logging
+import os
+from collections import OrderedDict
+from contextlib import contextmanager
+from functools import partial
+from pathlib import Path
+from typing import Generator, Union
+
+from mlflow.environment_variables import MLFLOW_TRACKING_URI
+from mlflow.store.db.db_types import DATABASE_ENGINES
+from mlflow.store.tracking import DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+from mlflow.store.tracking.rest_store import RestStore
+from mlflow.tracing.provider import reset
+from mlflow.tracking._tracking_service.registry import TrackingStoreRegistry
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.uri import _DATABRICKS_UNITY_CATALOG_SCHEME, _OSS_UNITY_CATALOG_SCHEME
+
+_logger = logging.getLogger(__name__)
+_tracking_uri = None
+
+
+def is_tracking_uri_set():
+    """Returns True if the tracking URI has been set, False otherwise."""
+    if _tracking_uri or MLFLOW_TRACKING_URI.get():
+        return True
+    return False
+
+
+def set_tracking_uri(uri: Union[str, Path]) -> None:
+    """
+    Set the tracking server URI. This does not affect the
+    currently active run (if one exists), but takes effect for successive runs.
+
+    Args:
+        uri:
+
+            - An empty string, or a local file path, prefixed with ``file:/``. Data is stored
+              locally at the provided file (or ``./mlruns`` if empty).
+            - An HTTP URI like ``https://my-tracking-server:5000``.
+            - A Databricks workspace, provided as the string "databricks" or, to use a Databricks
+              CLI `profile <https://github.com/databricks/databricks-cli#installation>`_,
+              "databricks://<profileName>".
+            - A :py:class:`pathlib.Path` instance
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+
+        mlflow.set_tracking_uri("file:///tmp/my_tracking")
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+    .. code-block:: text
+        :caption: Output
+
+        Current tracking uri: file:///tmp/my_tracking
+    """
+    if isinstance(uri, Path):
+        # On Windows with Python3.8 (https://bugs.python.org/issue38671)
+        # .resolve() doesn't return the absolute path if the directory doesn't exist
+        # so we're calling .absolute() first to get the absolute path on Windows,
+        # then .resolve() to clean the path
+        uri = uri.absolute().resolve().as_uri()
+    global _tracking_uri
+
+    if _tracking_uri != uri:
+        _tracking_uri = uri
+        if _tracking_uri is not None:
+            # Set 'MLFLOW_TRACKING_URI' environment variable
+            # so that subprocess can inherit it.
+            MLFLOW_TRACKING_URI.set(_tracking_uri)
+        else:
+            MLFLOW_TRACKING_URI.unset()
+
+        # Tracer provider uses tracking URI to determine where to export traces.
+        # Tracer provider stores the URI as its state so we need to reset
+        # it explicitly when the global tracking URI changes.
+        reset()
+
+
+@contextmanager
+def _use_tracking_uri(uri: str) -> Generator[None, None, None]:
+    """Temporarily use the specified tracking URI.
+
+    Args:
+        uri: The tracking URI to use.
+
+    """
+    old_tracking_uri = _tracking_uri
+    try:
+        set_tracking_uri(uri)
+        yield
+    finally:
+        set_tracking_uri(old_tracking_uri)
+
+
+def _resolve_tracking_uri(tracking_uri=None):
+    return tracking_uri or get_tracking_uri()
+
+
+def get_tracking_uri() -> str:
+    """Get the current tracking URI. This may not correspond to the tracking URI of
+    the currently active run, since the tracking URI can be updated via ``set_tracking_uri``.
+
+    Returns:
+        The tracking URI.
+
+    .. code-block:: python
+
+        import mlflow
+
+        # Get the current tracking uri
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+    .. code-block:: text
+
+        Current tracking uri: file:///.../mlruns
+    """
+    if _tracking_uri is not None:
+        return _tracking_uri
+    elif uri := MLFLOW_TRACKING_URI.get():
+        return uri
+    else:
+        return path_to_local_file_uri(os.path.abspath(DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH))
+
+
+def _get_file_store(store_uri, **_):
+    from mlflow.store.tracking.file_store import FileStore
+
+    return FileStore(store_uri, store_uri)
+
+
+def _get_sqlalchemy_store(store_uri, artifact_uri):
+    from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+    
+    # Check if this is a PostgreSQL URI that should use Managed Identity
+    if store_uri.startswith("postgresql://") and (
+        "auth_method=managed_identity" in store_uri or
+        os.getenv("MLFLOW_POSTGRES_USE_MANAGED_IDENTITY", "").lower() == "true"
+    ):
+        from mlflow.store.tracking.postgres_managed_identity import get_postgres_store_with_managed_identity
+        if artifact_uri is None:
+            artifact_uri = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+        return get_postgres_store_with_managed_identity(store_uri, artifact_uri)
+
+    if artifact_uri is None:
+        artifact_uri = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+    return SqlAlchemyStore(store_uri, artifact_uri)
+
+
+def _get_rest_store(store_uri, **_):
+    return RestStore(partial(get_default_host_creds, store_uri))
+
+
+
+
+def _get_databricks_rest_store(store_uri, **_):
+    return RestStore(partial(get_databricks_host_creds, store_uri))
+
+
+def _get_databricks_uc_rest_store(store_uri, **_):
+    from mlflow.exceptions import MlflowException
+    from mlflow.version import VERSION
+
+    supported_schemes = [
+        scheme
+        for scheme in _tracking_store_registry._registry
+        if scheme not in {_DATABRICKS_UNITY_CATALOG_SCHEME, _OSS_UNITY_CATALOG_SCHEME}
+    ]
+    raise MlflowException(
+        f"Detected Unity Catalog tracking URI '{store_uri}'. "
+        "Setting the tracking URI to a Unity Catalog backend is not supported in the current "
+        f"version of the MLflow client ({VERSION}). "
+        "Please specify a different tracking URI via mlflow.set_tracking_uri, with "
+        "one of the supported schemes: "
+        f"{supported_schemes}. If you're trying to access models in the Unity "
+        "Catalog, please upgrade to the latest version of the MLflow Python "
+        "client, then specify a Unity Catalog model registry URI via "
+        f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}') or "
+        f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}://profile_name') where "
+        "'profile_name' is the name of the Databricks CLI profile to use for "
+        "authentication. A OSS Unity Catalog model registry URI can also be specified via "
+        f"mlflow.set_registry_uri('{_OSS_UNITY_CATALOG_SCHEME}:http://localhost:8080')."
+        "Be sure to leave the registry URI configured to use one of the supported"
+        "schemes listed above."
+    )
+
+
+_tracking_store_registry = TrackingStoreRegistry()
+
+
+def _register_tracking_stores():
+    _tracking_store_registry.register("", _get_file_store)
+    _tracking_store_registry.register("file", _get_file_store)
+    _tracking_store_registry.register("databricks", _get_databricks_rest_store)
+    _tracking_store_registry.register(
+        _DATABRICKS_UNITY_CATALOG_SCHEME, _get_databricks_uc_rest_store
+    )
+    _tracking_store_registry.register(_OSS_UNITY_CATALOG_SCHEME, _get_databricks_uc_rest_store)
+
+
+    for scheme in ["http", "https"]:
+        _tracking_store_registry.register(scheme, _get_rest_store)
+
+    for scheme in DATABASE_ENGINES:
+        _tracking_store_registry.register(scheme, _get_sqlalchemy_store)
+
+    _tracking_store_registry.register_entrypoints()
+
+
+def _register(scheme, builder):
+    _tracking_store_registry.register(scheme, builder)
+
+
+_register_tracking_stores()
+
+
+def _get_store(store_uri=None, artifact_uri=None):
+    return _tracking_store_registry.get_store(store_uri, artifact_uri)
+
+
+_artifact_repos_cache = OrderedDict()
+
+
+def _get_artifact_repo(run_id):
+    return _artifact_repos_cache.get(run_id)
+
+
+# TODO(sueann): move to a projects utils module
+def _get_git_url_if_present(uri):
+    """Return the path git_uri#sub_directory if the URI passed is a local path that's part of
+    a Git repo, or returns the original URI otherwise.
+
+    Args:
+        uri: The expanded uri.
+
+    Returns:
+        The git_uri#sub_directory if the uri is part of a Git repo, otherwise return the original
+        uri.
+    """
+    if "#" in uri:
+        # Already a URI in git repo format
+        return uri
+    try:
+        from git import GitCommandNotFound, InvalidGitRepositoryError, NoSuchPathError, Repo
+    except ImportError as e:
+        _logger.warning(
+            "Failed to import Git (the git executable is probably not on your PATH),"
+            " so Git SHA is not available. Error: %s",
+            e,
+        )
+        return uri
+    try:
+        # Check whether this is part of a git repo
+        repo = Repo(uri, search_parent_directories=True)
+
+        # Repo url
+        repo_url = f"file://{repo.working_tree_dir}"
+
+        # Sub directory
+        rlpath = uri.replace(repo.working_tree_dir, "")
+        if rlpath == "":
+            git_path = repo_url
+        elif rlpath[0] == "/":
+            git_path = repo_url + "#" + rlpath[1:]
+        else:
+            git_path = repo_url + "#" + rlpath
+        return git_path
+    except (InvalidGitRepositoryError, GitCommandNotFound, ValueError, NoSuchPathError):
+        return uri

mlflow/tracking/artifact_utils.py

@@ -0,0 +1,179 @@
+"""
+Utilities for dealing with artifacts in the context of a Run.
+"""
+
+import os
+import pathlib
+import posixpath
+import tempfile
+import urllib.parse
+import uuid
+from typing import Any, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.store.artifact.dbfs_artifact_repo import DbfsRestArtifactRepository
+from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
+from mlflow.tracking._tracking_service.utils import _get_store
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.os import is_windows
+from mlflow.utils.uri import add_databricks_profile_info_to_artifact_uri, append_to_uri_path
+
+
+def get_artifact_uri(run_id, artifact_path=None, tracking_uri=None):
+    """Get the absolute URI of the specified artifact in the specified run. If `path` is not
+    specified the artifact root URI of the specified run will be returned; calls to ``log_artifact``
+    and ``log_artifacts`` write artifact(s) to subdirectories of the artifact root URI.
+
+    Args:
+        run_id: The ID of the run for which to obtain an absolute artifact URI.
+        artifact_path: The run-relative artifact path. For example,
+            ``path/to/artifact``. If unspecified, the artifact root URI for the
+            specified run will be returned.
+        tracking_uri: The tracking URI from which to get the run and its artifact location. If
+            not given, the current default tracking URI is used.
+
+    Returns:
+        An *absolute* URI referring to the specified artifact or the specified run's artifact
+        root. For example, if an artifact path is provided and the specified run uses an
+        S3-backed  store, this may be a uri of the form
+        ``s3://<bucket_name>/path/to/artifact/root/path/to/artifact``. If an artifact path
+        is not provided and the specified run uses an S3-backed store, this may be a URI of
+        the form ``s3://<bucket_name>/path/to/artifact/root``.
+
+    """
+    if not run_id:
+        raise MlflowException(
+            message="A run_id must be specified in order to obtain an artifact uri!",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    store = _get_store(tracking_uri)
+    run = store.get_run(run_id)
+    # Maybe move this method to RunsArtifactRepository so the circular dependency is clearer.
+    assert urllib.parse.urlparse(run.info.artifact_uri).scheme != "runs"  # avoid an infinite loop
+    if artifact_path is None:
+        return run.info.artifact_uri
+    else:
+        return append_to_uri_path(run.info.artifact_uri, artifact_path)
+
+
+# TODO: This would be much simpler if artifact_repo.download_artifacts could take the absolute path
+# or no path.
+def _get_root_uri_and_artifact_path(artifact_uri):
+    """Parse the artifact_uri to get the root_uri and artifact_path.
+
+    Args:
+        artifact_uri: The *absolute* URI of the artifact.
+    """
+    if os.path.exists(artifact_uri):
+        if not is_windows():
+            # If we're dealing with local files, just reference the direct pathing.
+            # non-nt-based file systems can directly reference path information, while nt-based
+            # systems need to url-encode special characters in directory listings to be able to
+            # resolve them (i.e., spaces converted to %20 within a file name or path listing)
+            root_uri = os.path.dirname(artifact_uri)
+            artifact_path = os.path.basename(artifact_uri)
+            return root_uri, artifact_path
+        else:  # if we're dealing with nt-based systems, we need to utilize pathname2url to encode.
+            artifact_uri = path_to_local_file_uri(artifact_uri)
+
+    parsed_uri = urllib.parse.urlparse(str(artifact_uri))
+    prefix = ""
+    if parsed_uri.scheme and not parsed_uri.path.startswith("/"):
+        # relative path is a special case, urllib does not reconstruct it properly
+        prefix = parsed_uri.scheme + ":"
+        parsed_uri = parsed_uri._replace(scheme="")
+
+    # For models:/ URIs, it doesn't make sense to initialize a ModelsArtifactRepository with only
+    # the model name portion of the URI, then call download_artifacts with the version info.
+    if ModelsArtifactRepository.is_models_uri(artifact_uri):
+        root_uri, artifact_path = ModelsArtifactRepository.split_models_uri(artifact_uri)
+    else:
+        artifact_path = posixpath.basename(parsed_uri.path)
+        parsed_uri = parsed_uri._replace(path=posixpath.dirname(parsed_uri.path))
+        root_uri = prefix + urllib.parse.urlunparse(parsed_uri)
+    return root_uri, artifact_path
+
+
+def _download_artifact_from_uri(
+    artifact_uri: str,
+    output_path: Optional[str] = None,
+    lineage_header_info: Optional[dict[str, Any]] = None,
+    tracking_uri: Optional[str] = None,
+) -> str:
+    """
+    Args:
+        artifact_uri: The *absolute* URI of the artifact to download.
+        output_path: The local filesystem path to which to download the artifact. If unspecified,
+            a local output path will be created.
+        lineage_header_info: The model lineage header info to be consumed by lineage services.
+        tracking_uri: The tracking URI to be used when downloading artifacts.
+    """
+    root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
+    repo = get_artifact_repository(artifact_uri=root_uri, tracking_uri=tracking_uri)
+
+    try:
+        if isinstance(repo, ModelsArtifactRepository):
+            return repo.download_artifacts(
+                artifact_path=artifact_path,
+                dst_path=output_path,
+                lineage_header_info=lineage_header_info,
+            )
+        return repo.download_artifacts(artifact_path=artifact_path, dst_path=output_path)
+    except Exception as e:
+        if artifact_uri.startswith("m-"):
+            # When a Model ID like string is passed, suggest using 'models:/{artifact_uri}' instead.
+            raise MlflowException(
+                f"Invalid uri `{artifact_uri}` is passed. Maybe you meant 'models:/{artifact_uri}'?"
+            ) from e
+        raise
+
+
+def _upload_artifact_to_uri(local_path, artifact_uri):
+    """Uploads a local artifact (file) to a specified URI.
+
+    Args:
+        local_path: The local path of the file to upload.
+        artifact_uri: The *absolute* URI of the path to upload the artifact to.
+
+    """
+    root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
+    get_artifact_repository(artifact_uri=root_uri).log_artifact(local_path, artifact_path)
+
+
+def _upload_artifacts_to_databricks(
+    source, run_id, source_host_uri=None, target_databricks_profile_uri=None
+):
+    """Copy the artifacts from ``source`` to the destination Databricks workspace (DBFS) given by
+    ``databricks_profile_uri`` or the current tracking URI.
+
+    Args:
+        source: Source location for the artifacts to copy.
+        run_id: Run ID to associate the artifacts with.
+        source_host_uri: Specifies the source artifact's host URI (e.g. Databricks tracking URI)
+            if applicable. If not given, defaults to the current tracking URI.
+        target_databricks_profile_uri: Specifies the destination Databricks host. If not given,
+            defaults to the current tracking URI.
+
+    Returns:
+        The DBFS location in the target Databricks workspace the model files have been
+        uploaded to.
+    """
+
+    with tempfile.TemporaryDirectory() as local_dir:
+        source_with_profile = add_databricks_profile_info_to_artifact_uri(source, source_host_uri)
+        _download_artifact_from_uri(source_with_profile, local_dir)
+        dest_root = "dbfs:/databricks/mlflow/tmp-external-source/"
+        dest_root_with_profile = add_databricks_profile_info_to_artifact_uri(
+            dest_root, target_databricks_profile_uri
+        )
+        dest_repo = DbfsRestArtifactRepository(dest_root_with_profile)
+        dest_artifact_path = run_id if run_id else uuid.uuid4().hex
+        # Allow uploading from the same run id multiple times by randomizing a suffix
+        if len(dest_repo.list_artifacts(dest_artifact_path)) > 0:
+            dest_artifact_path = dest_artifact_path + "-" + uuid.uuid4().hex[0:4]
+        dest_repo.log_artifacts(local_dir, artifact_path=dest_artifact_path)
+        dirname = pathlib.PurePath(source).name  # innermost directory name
+        return posixpath.join(dest_root, dest_artifact_path, dirname)  # new source