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