Databricks artifact repo

generate-protos.sh

@@ -14,16 +14,18 @@ protoc -I="$PROTOS" \
     "$PROTOS"/databricks.proto \
     "$PROTOS"/service.proto \
     "$PROTOS"/model_registry.proto \
+    "$PROTOS"/databricks_artifacts.proto \
     "$PROTOS"/scalapb/scalapb.proto
 
 OLD_SCALAPB="from scalapb import scalapb_pb2 as scalapb_dot_scalapb__pb2"
 NEW_SCALAPB="from .scalapb import scalapb_pb2 as scalapb_dot_scalapb__pb2"
-sed -i'.old' -e "s/$OLD_SCALAPB/$NEW_SCALAPB/g" "$PROTOS/databricks_pb2.py" "$PROTOS/service_pb2.py" "$PROTOS/model_registry_pb2.py"
+sed -i'.old' -e "s/$OLD_SCALAPB/$NEW_SCALAPB/g" "$PROTOS/databricks_pb2.py" "$PROTOS/service_pb2.py" "$PROTOS/model_registry_pb2.py" "$PROTOS/databricks_artifacts_pb2.py"
 
 OLD_DATABRICKS="import databricks_pb2 as databricks__pb2"
 NEW_DATABRICKS="from . import databricks_pb2 as databricks__pb2"
-sed -i'.old' -e "s/$OLD_DATABRICKS/$NEW_DATABRICKS/g" "$PROTOS/service_pb2.py" "$PROTOS/model_registry_pb2.py"
+sed -i'.old' -e "s/$OLD_DATABRICKS/$NEW_DATABRICKS/g" "$PROTOS/service_pb2.py" "$PROTOS/model_registry_pb2.py" "$PROTOS/databricks_artifacts_pb2.py"
 
 rm "$PROTOS/databricks_pb2.py.old"
 rm "$PROTOS/service_pb2.py.old"
 rm "$PROTOS/model_registry_pb2.py.old"
+rm "$PROTOS/databricks_artifacts__pb2.py.old"

mlflow/protos/databricks_artifacts.proto

@@ -0,0 +1,113 @@
+syntax = "proto2";
+
+package mlflow;
+
+import "scalapb/scalapb.proto";
+import "databricks.proto";
+
+option java_package = "com.databricks.api.proto.mlflow";
+option java_generate_equals_and_hash = true;
+option py_generic_services = true;
+option (scalapb.options) = {
+  flat_package: true,
+};
+
+service DatabricksMlflowArtifactsService {
+
+  // Fetch credentials to read from the specified MLflow artifact location
+  //
+  // Note: Even if no artifacts exist at the specified artifact location, this API will
+  // still provide read credentials as long as the format of the location is valid.
+  // Callers must subsequently check for the existence of the artifacts using the appropriate
+  // cloud storage APIs (as determined by the `ArtifactCredentialType` property of the response)
+  rpc getCredentialsForRead (GetCredentialsForRead) returns (GetCredentialsForRead.Response) {
+    option (rpc) = {
+      endpoints: [{
+        method: "GET",
+        path: "/mlflow/artifacts/credentials-for-read"
+        since { major: 2, minor: 0 },
+      }],
+      visibility: PUBLIC_UNDOCUMENTED,
+    };
+  }
+
+  // Fetch credentials to write to the specified MLflow artifact location
+  rpc getCredentialsForWrite (GetCredentialsForWrite) returns (GetCredentialsForWrite.Response) {
+    option (rpc) = {
+      endpoints: [{
+        method: "GET",
+        path: "/mlflow/artifacts/credentials-for-write"
+        since { major: 2, minor: 0 },
+      }],
+      visibility: PUBLIC_UNDOCUMENTED,
+    };
+  }
+}
+
+// The type of a given artifact access credential
+enum ArtifactCredentialType {
+
+  // The credential is an Azure Shared Access Signature URI. For more information, see
+  // https://docs.microsoft.com/en-us/azure/storage/common/storage-sas-overview
+  AZURE_SAS_URI = 1;
+
+  // The credential is an AWS Presigned URL. For more information, see
+  // https://docs.aws.amazon.com/AmazonS3/latest/dev/ShareObjectPreSignedURL.html
+  AWS_PRESIGNED_URL = 2;
+
+}
+
+message ArtifactCredentialInfo {
+
+  // The ID of the MLflow Run containing the artifact that can be accessed
+  // with the credential
+  optional string run_id = 1;
+
+  // The path, relative to the Run's artifact root location, of the artifact
+  // that can be accessed with the credential
+  optional string path = 2;
+
+  // The signed URI credential that provides access to the artifact
+  optional string signed_uri = 3;
+
+  // The type of the signed credential URI (e.g., an AWS presigned URL
+  // or an Azure Shared Access Signature URI)
+  optional ArtifactCredentialType type = 4;
+
+}
+
+message GetCredentialsForRead {
+  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";
+
+  // The ID of the MLflow Run for which to fetch artifact read credentials
+  optional string run_id = 1;
+
+  // The artifact path, relative to the Run's artifact root location, for which to
+  // fetch artifact read credentials
+  optional string path = 2;
+
+  message Response {
+
+    // Credentials for reading from the specified artifact location
+    optional ArtifactCredentialInfo credentials = 1;
+
+  }
+}
+
+message GetCredentialsForWrite {
+  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";
+
+  // The ID of the MLflow Run for which to fetch artifact write credentials
+  optional string run_id = 1;
+
+  // The artifact path, relative to the Run's artifact root location, for which to
+  // fetch artifact write credentials
+  optional string path = 2;
+
+  message Response {
+
+    // Credentials for writing to the specified artifacts location
+    optional ArtifactCredentialInfo credentials = 1;
+
+  }
+}

mlflow/store/artifact/databricks_artifact_repo.py

@@ -0,0 +1,94 @@
+from azure.storage.blob import BlobClient
+
+import os
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils.string_utils import strip_suffix
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+from mlflow.utils.rest_utils import call_endpoint, extract_api_info_for_service
+from mlflow.protos.databricks_artifacts_pb2 import DatabricksMlflowArtifactsService
+from mlflow.protos.databricks_artifacts_pb2 import GetCredentialsForWrite, GetCredentialsForRead
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.protos.service_pb2 import MlflowService, ListArtifacts
+
+_PATH_PREFIX = "/api/2.0"
+
+
+class DatabricksArtifactRepository(ArtifactRepository):
+    """
+    SOMETHING : TYPING TILL IT WORKS LOL
+    """
+
+    def __init__(self, artifact_uri):
+        super(DatabricksArtifactRepository, self).__init__(artifact_uri)
+
+    def _extract_run_id(self, artifact_uri):
+        return artifact_uri.lstrip('/').split('/')[4]
+
+    def _call_endpoint(self, service, api, json_body):
+        _METHOD_TO_INFO = extract_api_info_for_service(service, _PATH_PREFIX)
+        endpoint, method = _METHOD_TO_INFO[api]
+        response_proto = api.Response()
+        return call_endpoint(get_databricks_host_creds(), endpoint, method, json_body, response_proto)
+
+    def _create_json_body(self, run_id, path=None):
+        path = path or '.'
+        return {
+            "run_id": run_id,
+            "path": path
+        }
+
+    def _get_azure_write_credentials(self, run_id, path=None):
+        return self._call_endpoint(DatabricksMlflowArtifactsService, GetCredentialsForWrite,
+                                   self._create_json_body(run_id, path))
+
+    def _get_azure_read_credentials(self, run_id, path=None):
+        return self._call_endpoint(DatabricksMlflowArtifactsService, GetCredentialsForRead,
+                                   self._create_json_body(run_id, path))
+
+    def _upload_file(self, local_file, artifact_path):
+        run_id = self._extract_run_id(self.artifact_uri)
+        write_credentials = self._get_azure_write_credentials(run_id, artifact_path)
+        signed_write_uri = write_credentials.credentials.signed_uri
+        service = BlobClient.from_blob_url(blob_url=signed_write_uri, credential=None)
+        try:
+            with open(local_file, "rb") as data:
+                service.upload_blob(data, overwrite=True)
+        except Exception as err:
+            raise MlflowException(err)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        self._upload_file(local_file, artifact_path)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        artifact_path = artifact_path or ''
+        basename = os.path.basename(strip_suffix(local_dir, '/'))
+        for (dirpath, _, filenames) in os.walk(local_dir):
+            artifact_subdir = basename
+            if dirpath != local_dir:
+                rel_path = os.path.relpath(dirpath, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                artifact_subdir = os.path.join(artifact_subdir, rel_path)
+            for name in filenames:
+                local_file = os.path.join(dirpath, name)
+                artifact_location = os.path.join(artifact_path, artifact_subdir)
+                self._upload_file(local_file, artifact_location)
+
+    def list_artifacts(self, path=None):
+        run_id = self._extract_run_id(self.artifact_uri)
+        return self._call_endpoint(MlflowService, ListArtifacts, self._create_json_body(run_id, path))
+
+    def _download_file(self, remote_file_path, local_path):
+        run_id = self._extract_run_id(self.artifact_uri)
+        read_credentials = self._get_azure_read_credentials(run_id, remote_file_path)
+        signed_read_uri = read_credentials.credentials.signed_uri
+        service = BlobClient.from_blob_url(blob_url=signed_read_uri, credential=None)
+        try:
+            with open(local_path, "wb") as output_file:
+                blob = service.download_blob()
+                output_file.write(blob.readall())
+        except Exception as err:
+            raise MlflowException(err)
+
+    def delete_artifacts(self, artifact_path=None):
+        raise MlflowException('Not implemented yet')

mlflow/store/artifact/dbfs_artifact_repo.py

@@ -6,12 +6,14 @@
 from mlflow.exceptions import MlflowException
 from mlflow.store.tracking.rest_store import RestStore
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.databricks_artifact_repo import DatabricksArtifactRepository
 from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
 from mlflow.tracking._tracking_service import utils
 from mlflow.utils.file_utils import relative_path_to_artifact_path
 from mlflow.utils.rest_utils import http_request, http_request_safe, RESOURCE_DOES_NOT_EXIST
 from mlflow.utils.string_utils import strip_prefix
 import mlflow.utils.databricks_utils
+from mlflow.utils.uri import is_artifact_acled_uri
 
 LIST_API_ENDPOINT = '/api/2.0/dbfs/list'
 GET_STATUS_ENDPOINT = '/api/2.0/dbfs/get-status'
@@ -163,7 +165,9 @@ def dbfs_artifact_repo_factory(artifact_uri):
     :return: Subclass of ArtifactRepository capable of storing artifacts on DBFS.
     """
     cleaned_artifact_uri = artifact_uri.rstrip('/')
-    if mlflow.utils.databricks_utils.is_dbfs_fuse_available() \
+    if is_artifact_acled_uri(artifact_uri):
+        return DatabricksArtifactRepository(artifact_uri)
+    elif mlflow.utils.databricks_utils.is_dbfs_fuse_available() \
             and os.environ.get(USE_FUSE_ENV_VAR, "").lower() != "false" \
             and not artifact_uri.startswith("dbfs:/databricks/mlflow-registry"):
         # If the DBFS FUSE mount is available, write artifacts directly to /dbfs/... using

mlflow/utils/databricks_utils.py

@@ -6,7 +6,6 @@
 from mlflow.utils.rest_utils import MlflowHostCreds
 from databricks_cli.configure import provider
 
-
 _logger = logging.getLogger(__name__)
 
 

mlflow/utils/uri.py

@@ -8,7 +8,7 @@
 
 _INVALID_DB_URI_MSG = "Please refer to https://mlflow.org/docs/latest/tracking.html#storage for " \
                       "format specifications."
-
+_ACLED_ARTIFACT_URI = "dbfs:/databricks/mlflow-tracking/"
 
 def is_local_uri(uri):
     """Returns true if this is a local file path (/foo or file:/foo)."""
@@ -129,3 +129,7 @@ def _join_posixpaths_and_append_absolute_suffixes(prefix_path, suffix_path):
     # joined path
     suffix_path = suffix_path.lstrip(posixpath.sep)
     return posixpath.join(prefix_path, suffix_path)
+
+
+def is_artifact_acled_uri(artifact_uri):
+    return artifact_uri.startswith(_ACLED_ARTIFACT_URI.lstrip('/'))