from mlflow.entities.model_registry import RegisteredModel, ModelVersion
from mlflow.protos.model_registry_pb2 import ModelRegistryService, CreateRegisteredModel, \
UpdateRegisteredModel, DeleteRegisteredModel, ListRegisteredModels, \
GetLatestVersions, CreateModelVersion, UpdateModelVersion, \
DeleteModelVersion, GetModelVersionDownloadUri, SearchModelVersions, \
RenameRegisteredModel, GetRegisteredModel, GetModelVersion, TransitionModelVersionStage, \
SearchRegisteredModels, SetRegisteredModelTag, SetModelVersionTag, \
DeleteRegisteredModelTag, DeleteModelVersionTag
from mlflow.store.entities.paged_list import PagedList
from mlflow.store.model_registry.abstract_store import AbstractStore
from mlflow.utils.proto_json_utils import message_to_json
from mlflow.utils.rest_utils import call_endpoint, extract_api_info_for_service
_PATH_PREFIX = "/api/2.0"
_METHOD_TO_INFO = extract_api_info_for_service(ModelRegistryService,
_PATH_PREFIX)
_logger = logging.getLogger(__name__)
class RestStore(AbstractStore):
"""
Note:: Experimental: This entity may change or be removed in a future release without warning.
Client for a remote model registry server accessed via REST API calls
:param get_host_creds: Method to be invoked prior to every REST request to get the
:py:class:`mlflow.rest_utils.MlflowHostCreds` for the request. Note that this
is a function so that we can obtain fresh credentials in the case of expiry.
"""
def __init__(self, get_host_creds):
super(RestStore, self).__init__()
UpdateExperiment,
LogBatch,
LogModel,
DeleteTag,
SetExperimentTag,
GetExperimentByName,
)
from mlflow.store.tracking.abstract_store import AbstractStore
from mlflow.utils.proto_json_utils import message_to_json
from mlflow.utils.rest_utils import (
call_endpoint,
extract_api_info_for_service,
_REST_API_PATH_PREFIX,
)
_METHOD_TO_INFO = extract_api_info_for_service(MlflowService, _REST_API_PATH_PREFIX)
class RestStore(AbstractStore):
"""
Client for a remote tracking server accessed via REST API calls
:param get_host_creds: Method to be invoked prior to every REST request to get the
:py:class:`mlflow.rest_utils.MlflowHostCreds` for the request. Note that this
is a function so that we can obtain fresh credentials in the case of expiry.
"""
def __init__(self, get_host_creds):
super(RestStore, self).__init__()
self.get_host_creds = get_host_creds
from mlflow.protos.databricks_artifacts_pb2 import DatabricksMlflowArtifactsService, \
GetCredentialsForWrite, GetCredentialsForRead, ArtifactCredentialType
from mlflow.protos.service_pb2 import MlflowService, GetRun, ListArtifacts
from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.utils.databricks_utils import get_databricks_host_creds
from mlflow.utils.file_utils import relative_path_to_artifact_path, yield_file_in_chunks
from mlflow.utils.proto_json_utils import message_to_json
from mlflow.utils.rest_utils import call_endpoint, extract_api_info_for_service
from mlflow.utils.uri import extract_and_normalize_path, is_databricks_acled_artifacts_uri
_logger = logging.getLogger(__name__)
_PATH_PREFIX = "/api/2.0"
_AZURE_MAX_BLOCK_CHUNK_SIZE = 100000000 # Max. size of each block allowed is 100 MB in stage_block
_DOWNLOAD_CHUNK_SIZE = 100000000
_SERVICE_AND_METHOD_TO_INFO = {
service: extract_api_info_for_service(service, _PATH_PREFIX)
for service in [MlflowService, DatabricksMlflowArtifactsService]
}
class DatabricksArtifactRepository(ArtifactRepository):
"""
Performs storage operations on artifacts in the access-controlled
`dbfs:/databricks/mlflow-tracking` location.
Signed access URIs for S3 / Azure Blob Storage are fetched from the MLflow service and used to
read and write files from/to this location.
The artifact_uri is expected to be of the form
dbfs:/databricks/mlflow-tracking/<EXP_ID>/<RUN_ID>/
"""
