genesis-flow 1.0.0__py3-none-any.whl

mlflow/store/artifact/unity_catalog_models_artifact_repo.py

@@ -0,0 +1,168 @@
+import base64
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    MODEL_VERSION_OPERATION_READ,
+    GenerateTemporaryModelVersionCredentialsRequest,
+    GenerateTemporaryModelVersionCredentialsResponse,
+    ModelVersionLineageDirection,
+    StorageMode,
+)
+from mlflow.protos.databricks_uc_registry_service_pb2 import UcModelRegistryService
+from mlflow.store._unity_catalog.lineage.constants import _DATABRICKS_LINEAGE_ID_HEADER
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.databricks_sdk_models_artifact_repo import (
+    DatabricksSDKModelsArtifactRepository,
+)
+from mlflow.store.artifact.presigned_url_artifact_repo import PresignedUrlArtifactRepository
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+)
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils._unity_catalog_utils import (
+    emit_model_version_lineage,
+    get_artifact_repo_from_storage_info,
+    get_full_name_from_sc,
+    is_databricks_sdk_models_artifact_repository_enabled,
+)
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.uri import (
+    _DATABRICKS_UNITY_CATALOG_SCHEME,
+    get_databricks_profile_uri_from_artifact_uri,
+    get_db_info_from_uri,
+    is_databricks_unity_catalog_uri,
+)
+
+_METHOD_TO_INFO = extract_api_info_for_service(UcModelRegistryService, _REST_API_PATH_PREFIX)
+
+
+class UnityCatalogModelsArtifactRepository(ArtifactRepository):
+    """
+    Performs storage operations on artifacts controlled by a Unity Catalog model registry
+
+    Temporary scoped tokens for the appropriate cloud storage locations are fetched from the
+    remote backend and used to download model artifacts.
+
+    The artifact_uri is expected to be of the form `models:/<model_name>/<model_version>`
+
+    Note : This artifact repository is meant is to be instantiated by the ModelsArtifactRepository
+    when the client is pointing to a Unity Catalog model registry.
+    """
+
+    def __init__(self, artifact_uri, registry_uri):
+        if not is_databricks_unity_catalog_uri(registry_uri):
+            raise MlflowException(
+                message="Attempted to instantiate an artifact repo to access models in the "
+                f"Unity Catalog with non-Unity Catalog registry URI '{registry_uri}'. "
+                f"Please specify a Unity Catalog registry URI of the "
+                f"form '{_DATABRICKS_UNITY_CATALOG_SCHEME}[://profile]', e.g. by calling "
+                f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}') if using the "
+                f"MLflow Python client",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        super().__init__(artifact_uri)
+        from mlflow.tracking.client import MlflowClient
+
+        registry_uri_from_artifact_uri = get_databricks_profile_uri_from_artifact_uri(
+            artifact_uri, result_scheme=_DATABRICKS_UNITY_CATALOG_SCHEME
+        )
+        if registry_uri_from_artifact_uri is not None:
+            registry_uri = registry_uri_from_artifact_uri
+        _, key_prefix = get_db_info_from_uri(registry_uri)
+        if key_prefix is not None:
+            raise MlflowException(
+                "Remote model registry access via model URIs of the form "
+                "'models://<scope>@<prefix>/<model_name>/<version_or_stage>' is unsupported for "
+                "models in the Unity Catalog. We recommend that you access the Unity Catalog "
+                "from the current Databricks workspace instead."
+            )
+        self.registry_uri = registry_uri
+        self.client = MlflowClient(registry_uri=self.registry_uri)
+        try:
+            spark = _get_active_spark_session()
+        except Exception:
+            pass
+        model_name, self.model_version = get_model_name_and_version(self.client, artifact_uri)
+        self.model_name = get_full_name_from_sc(model_name, spark)
+
+    def _get_blob_storage_path(self):
+        return self.client.get_model_version_download_uri(self.model_name, self.model_version)
+
+    def _get_scoped_token(self, lineage_header_info=None):
+        extra_headers = {}
+        if lineage_header_info:
+            header_json = message_to_json(lineage_header_info)
+            header_base64 = base64.b64encode(header_json.encode())
+            extra_headers[_DATABRICKS_LINEAGE_ID_HEADER] = header_base64
+
+        db_creds = get_databricks_host_creds(self.registry_uri)
+        endpoint, method = _METHOD_TO_INFO[GenerateTemporaryModelVersionCredentialsRequest]
+        req_body = message_to_json(
+            GenerateTemporaryModelVersionCredentialsRequest(
+                name=self.model_name,
+                version=self.model_version,
+                operation=MODEL_VERSION_OPERATION_READ,
+            )
+        )
+        response_proto = GenerateTemporaryModelVersionCredentialsResponse()
+        return call_endpoint(
+            host_creds=db_creds,
+            endpoint=endpoint,
+            method=method,
+            json_body=req_body,
+            response_proto=response_proto,
+            extra_headers=extra_headers,
+        ).credentials
+
+    def _get_artifact_repo(self, lineage_header_info=None):
+        """
+        Get underlying ArtifactRepository instance for model version blob
+        storage
+        """
+        host_creds = get_databricks_host_creds(self.registry_uri)
+        if is_databricks_sdk_models_artifact_repository_enabled(host_creds):
+            entities = lineage_header_info.entities if lineage_header_info else []
+            emit_model_version_lineage(
+                host_creds,
+                self.model_name,
+                self.model_version,
+                entities,
+                ModelVersionLineageDirection.DOWNSTREAM,
+            )
+            return DatabricksSDKModelsArtifactRepository(self.model_name, self.model_version)
+        scoped_token = self._get_scoped_token(lineage_header_info=lineage_header_info)
+        if scoped_token.storage_mode == StorageMode.DEFAULT_STORAGE:
+            return PresignedUrlArtifactRepository(
+                get_databricks_host_creds(self.registry_uri), self.model_name, self.model_version
+            )
+
+        blob_storage_path = self._get_blob_storage_path()
+        return get_artifact_repo_from_storage_info(
+            storage_location=blob_storage_path,
+            scoped_token=scoped_token,
+            base_credential_refresh_def=self._get_scoped_token,
+        )
+
+    def list_artifacts(self, path=None):
+        return self._get_artifact_repo().list_artifacts(path=path)
+
+    def download_artifacts(self, artifact_path, dst_path=None, lineage_header_info=None):
+        return self._get_artifact_repo(lineage_header_info=lineage_header_info).download_artifacts(
+            artifact_path, dst_path
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")

mlflow/store/artifact/unity_catalog_oss_models_artifact_repo.py

@@ -0,0 +1,168 @@
+import base64
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_uc_registry_service_pb2 import UcModelRegistryService
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    READ_MODEL_VERSION as MODEL_VERSION_OPERATION_READ_OSS,
+)
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    GenerateTemporaryModelVersionCredential as GenerateTemporaryModelVersionCredentialsOSS,
+)
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    TemporaryCredentials,
+)
+from mlflow.protos.unity_catalog_oss_service_pb2 import UnityCatalogService
+from mlflow.store._unity_catalog.lineage.constants import _DATABRICKS_LINEAGE_ID_HEADER
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+)
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils._unity_catalog_utils import (
+    get_artifact_repo_from_storage_info,
+    get_full_name_from_sc,
+)
+from mlflow.utils.oss_registry_utils import get_oss_host_creds
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    _UC_OSS_REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.uri import (
+    _OSS_UNITY_CATALOG_SCHEME,
+    get_databricks_profile_uri_from_artifact_uri,
+    get_db_info_from_uri,
+    is_oss_unity_catalog_uri,
+)
+
+_METHOD_TO_INFO = extract_api_info_for_service(UcModelRegistryService, _REST_API_PATH_PREFIX)
+_METHOD_TO_INFO_OSS = extract_api_info_for_service(
+    UnityCatalogService, _UC_OSS_REST_API_PATH_PREFIX
+)
+
+import urllib.parse
+
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.utils.uri import is_file_uri
+
+
+class UnityCatalogOSSModelsArtifactRepository(ArtifactRepository):
+    """
+    Performs storage operations on artifacts controlled by a Unity Catalog model registry
+
+    Temporary scoped tokens for the appropriate cloud storage locations are fetched from the
+    remote backend and used to download model artifacts.
+
+    The artifact_uri is expected to be of the form `models:/<model_name>/<model_version>`
+
+    Note : This artifact repository is meant is to be instantiated by the ModelsArtifactRepository
+    when the client is pointing to a Unity Catalog model registry.
+    """
+
+    def __init__(self, artifact_uri, registry_uri):
+        if not is_oss_unity_catalog_uri(registry_uri):
+            raise MlflowException(
+                message="Attempted to instantiate an artifact repo to access models in the "
+                f"OSS Unity Catalog with non-Unity Catalog registry URI '{registry_uri}'. "
+                f"Please specify a Unity Catalog registry URI of the "
+                f"form '{_OSS_UNITY_CATALOG_SCHEME}[://profile]', e.g. by calling "
+                f"mlflow.set_registry_uri('{_OSS_UNITY_CATALOG_SCHEME}') if using the "
+                f"MLflow Python client",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        super().__init__(artifact_uri)
+        from mlflow.tracking.client import MlflowClient
+
+        registry_uri_from_artifact_uri = get_databricks_profile_uri_from_artifact_uri(
+            artifact_uri, result_scheme=_OSS_UNITY_CATALOG_SCHEME
+        )
+        if registry_uri_from_artifact_uri is not None:
+            registry_uri = registry_uri_from_artifact_uri
+
+        _, key_prefix = get_db_info_from_uri(urllib.parse.urlparse(registry_uri).path)
+        if key_prefix is not None:
+            raise MlflowException(
+                "Remote model registry access via model URIs of the form "
+                "'models://<scope>@<prefix>/<model_name>/<version_or_stage>' is unsupported for "
+                "models in the Unity Catalog. We recommend that you access the Unity Catalog "
+                "from the current Databricks workspace instead."
+            )
+        self.registry_uri = registry_uri
+        self.client = MlflowClient(registry_uri=self.registry_uri)
+        try:
+            spark = _get_active_spark_session()
+        except Exception:
+            pass
+        model_name, self.model_version = get_model_name_and_version(self.client, artifact_uri)
+        self.model_name = get_full_name_from_sc(model_name, spark)
+
+    def _get_blob_storage_path(self):
+        return self.client.get_model_version_download_uri(self.model_name, self.model_version)
+
+    def _get_scoped_token(self, lineage_header_info=None):
+        extra_headers = {}
+        if lineage_header_info:
+            header_json = message_to_json(lineage_header_info)
+            header_base64 = base64.b64encode(header_json.encode())
+            extra_headers[_DATABRICKS_LINEAGE_ID_HEADER] = header_base64
+        oss_creds = get_oss_host_creds(
+            self.registry_uri
+        )  # Implement ENV variable the same way the databricks user/token is specified
+        oss_endpoint, oss_method = _METHOD_TO_INFO_OSS[GenerateTemporaryModelVersionCredentialsOSS]
+        [catalog_name, schema_name, model_name] = self.model_name.split(
+            "."
+        )  # self.model_name is actually the full name
+        oss_req_body = message_to_json(
+            GenerateTemporaryModelVersionCredentialsOSS(
+                catalog_name=catalog_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version=int(self.model_version),
+                operation=MODEL_VERSION_OPERATION_READ_OSS,
+            )
+        )
+        oss_response_proto = TemporaryCredentials()
+        return call_endpoint(
+            host_creds=oss_creds,
+            endpoint=oss_endpoint,
+            method=oss_method,
+            json_body=oss_req_body,
+            response_proto=oss_response_proto,
+            extra_headers=extra_headers,
+        )
+
+    def _get_artifact_repo(self, lineage_header_info=None):
+        """
+        Get underlying ArtifactRepository instance for model version blob
+        storage
+        """
+        blob_storage_path = self._get_blob_storage_path()
+        if is_file_uri(blob_storage_path):
+            return LocalArtifactRepository(artifact_uri=blob_storage_path)
+        scoped_token = self._get_scoped_token(lineage_header_info=lineage_header_info)
+        return get_artifact_repo_from_storage_info(
+            storage_location=blob_storage_path,
+            scoped_token=scoped_token,
+            base_credential_refresh_def=self._get_scoped_token,
+            is_oss=True,
+        )
+
+    def list_artifacts(self, path=None):
+        return self._get_artifact_repo().list_artifacts(path=path)
+
+    def download_artifacts(self, artifact_path, dst_path=None, lineage_header_info=None):
+        return self._get_artifact_repo(lineage_header_info=lineage_header_info).download_artifacts(
+            artifact_path, dst_path
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")

mlflow/store/artifact/utils/models.py

@@ -0,0 +1,148 @@
+import urllib.parse
+from pathlib import Path
+from typing import NamedTuple, Optional, Union
+
+import mlflow.tracking
+from mlflow.exceptions import MlflowException
+from mlflow.utils.uri import (
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_uri,
+    is_models_uri,
+)
+
+_MODELS_URI_SUFFIX_LATEST = "latest"
+
+
+def is_using_databricks_registry(uri):
+    profile_uri = get_databricks_profile_uri_from_artifact_uri(uri) or mlflow.get_registry_uri()
+    return is_databricks_uri(profile_uri)
+
+
+def _improper_model_uri_msg(uri, scheme: str = "models"):
+    if scheme not in ("models", "prompts"):
+        raise ValueError(f"Unsupported scheme for model/prompt URI: {scheme!r}")
+    entity_type = "Models" if scheme == "models" else "Prompts"
+    return (
+        f"Not a proper {scheme}:/ URI: {uri}. "
+        + f"{entity_type} URIs must be of the form '{scheme}:/name/suffix' "
+        + f"or '{scheme}:/name@alias' where suffix is a version"
+        + (f", stage, or the string {_MODELS_URI_SUFFIX_LATEST!r}" if scheme == "models" else "")
+        + f" and where alias is a registered {scheme[:-1]} alias. "
+        + "Only one of suffix or alias can be defined at a time."
+    )
+
+
+def _get_latest_model_version(client, name, stage):
+    """
+    Returns the latest version of the stage if stage is not None. Otherwise return the latest of all
+    versions.
+    """
+    latest = client.get_latest_versions(name, None if stage is None else [stage])
+    if len(latest) == 0:
+        stage_str = "" if stage is None else f" and stage '{stage}'"
+        raise MlflowException(f"No versions of model with name '{name}'{stage_str} found")
+    return max(int(x.version) for x in latest)
+
+
+class ParsedModelUri(NamedTuple):
+    model_id: Optional[str] = None
+    name: Optional[str] = None
+    version: Optional[str] = None
+    stage: Optional[str] = None
+    alias: Optional[str] = None
+
+
+def _parse_model_uri(uri, scheme: str = "models") -> ParsedModelUri:
+    """
+    Returns a ParsedModelUri tuple. Since a models:/ or prompts:/ URI can only have one of
+    {version, stage, 'latest', alias}, it will return
+        - (id, None, None, None) to look for a specific model by ID,
+        - (name, version, None, None) to look for a specific version,
+        - (name, None, stage, None) to look for the latest version of a stage,
+        - (name, None, None, None) to look for the latest of all versions.
+        - (name, None, None, alias) to look for a registered model alias.
+
+    Args:
+        uri: The URI to parse (e.g., "models:/name/version" or "prompts:/name@alias")
+        scheme: The expected URI scheme (default: "models", can be "prompts")
+    """
+    parsed = urllib.parse.urlparse(uri, allow_fragments=False)
+    if parsed.scheme != scheme:
+        raise MlflowException(_improper_model_uri_msg(uri, scheme))
+    path = parsed.path
+    if not path.startswith("/") or len(path) <= 1:
+        raise MlflowException(_improper_model_uri_msg(uri, scheme))
+
+    parts = path.lstrip("/").split("/")
+    if len(parts) > 2 or parts[0].strip() == "":
+        raise MlflowException(_improper_model_uri_msg(uri, scheme))
+
+    if len(parts) == 2:
+        name, suffix = parts
+        if suffix.strip() == "":
+            raise MlflowException(_improper_model_uri_msg(uri, scheme))
+        # The URI is in the suffix format
+        if suffix.isdigit():
+            # The suffix is a specific version, e.g. "models:/AdsModel1/123"
+            return ParsedModelUri(name=name, version=suffix)
+        elif suffix.lower() == _MODELS_URI_SUFFIX_LATEST.lower() and scheme == "models":
+            # The suffix is the 'latest' string (case insensitive), e.g. "models:/AdsModel1/latest"
+            # Only supported for models, not prompts
+            return ParsedModelUri(name=name)
+        elif scheme == "models":
+            # The suffix is a specific stage (case insensitive), e.g. "models:/AdsModel1/Production"
+            # Only supported for models, not prompts
+            return ParsedModelUri(name=name, stage=suffix)
+        else:
+            # For prompts, only version numbers are supported, not stages or 'latest'
+            raise MlflowException(_improper_model_uri_msg(uri, scheme))
+    elif "@" in path:
+        # The URI is an alias URI, e.g. "models:/AdsModel1@Champion"
+        alias_parts = parts[0].rsplit("@", 1)
+        if len(alias_parts) != 2 or alias_parts[1].strip() == "":
+            raise MlflowException(_improper_model_uri_msg(uri, scheme))
+        return ParsedModelUri(name=alias_parts[0], alias=alias_parts[1])
+    else:
+        # The URI is of the form "models:/<model_id>"
+        return ParsedModelUri(parts[0])
+
+
+def _parse_model_id_if_present(possible_model_uri: Union[str, Path]) -> Optional[str]:
+    """
+    Parses the model ID from the given string. If the string represents a UC model URI, we get the
+    model version to extract the model ID. If the string is not a models:/ URI, returns None.
+
+    Args:
+        possible_model_uri: The string that may be a models:/ URI.
+
+    Returns:
+        The model ID if the string is a models:/ URI, otherwise None.
+    """
+    uri = str(possible_model_uri)
+    if is_models_uri(uri):
+        parsed_model_uri = _parse_model_uri(uri)
+        if parsed_model_uri.model_id is not None:
+            return parsed_model_uri.model_id
+        elif parsed_model_uri.name is not None and parsed_model_uri.version is not None:
+            client = mlflow.tracking.MlflowClient()
+            return client.get_model_version(
+                parsed_model_uri.name, parsed_model_uri.version
+            ).model_id
+    return None
+
+
+def get_model_name_and_version(client, models_uri):
+    (model_id, model_name, model_version, model_stage, model_alias) = _parse_model_uri(models_uri)
+    if model_id is not None:
+        return (model_id,)
+    if model_version is not None:
+        return model_name, model_version
+
+    # NB: Call get_model_version_by_alias of registry client directly to bypass prompt check
+    if isinstance(client, mlflow.MlflowClient):
+        client = client._get_registry_client()
+
+    if model_alias is not None:
+        mv = client.get_model_version_by_alias(model_name, model_alias)
+        return model_name, mv.version
+    return model_name, str(_get_latest_model_version(client, model_name, model_stage))

mlflow/store/db/base_sql_model.py

@@ -0,0 +1,3 @@
+from sqlalchemy.orm import declarative_base
+
+Base = declarative_base()

mlflow/store/db/db_types.py

@@ -0,0 +1,10 @@
+"""
+Set of SQLAlchemy database schemas supported in MLflow for tracking server backends.
+"""
+
+POSTGRES = "postgresql"
+MYSQL = "mysql"
+SQLITE = "sqlite"
+MSSQL = "mssql"
+
+DATABASE_ENGINES = [POSTGRES, MYSQL, SQLITE, MSSQL]