genesis-flow 1.0.0__py3-none-any.whl

mlflow/store/artifact/s3_artifact_repo.py

@@ -0,0 +1,330 @@
+import json
+import os
+import posixpath
+import urllib.parse
+from datetime import datetime
+from functools import lru_cache
+from mimetypes import guess_type
+from typing import Optional
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+)
+from mlflow.environment_variables import (
+    MLFLOW_BOTO_CLIENT_ADDRESSING_STYLE,
+    MLFLOW_S3_ENDPOINT_URL,
+    MLFLOW_S3_IGNORE_TLS,
+    MLFLOW_S3_UPLOAD_EXTRA_ARGS,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    MultipartUploadMixin,
+)
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+
+_MAX_CACHE_SECONDS = 300
+
+
+def _get_utcnow_timestamp():
+    return datetime.utcnow().timestamp()
+
+
+@lru_cache(maxsize=64)
+def _cached_get_s3_client(
+    signature_version,
+    addressing_style,
+    s3_endpoint_url,
+    verify,
+    timestamp,
+    access_key_id=None,
+    secret_access_key=None,
+    session_token=None,
+    region_name=None,
+):
+    """Returns a boto3 client, caching to avoid extra boto3 verify calls.
+
+    This method is outside of the S3ArtifactRepository as it is
+    agnostic and could be used by other instances.
+
+    `maxsize` set to avoid excessive memory consumption in the case
+    a user has dynamic endpoints (intentionally or as a bug).
+
+    Some of the boto3 endpoint urls, in very edge cases, might expire
+    after twelve hours as that is the current expiration time. To ensure
+    we throw an error on verification instead of using an expired endpoint
+    we utilise the `timestamp` parameter to invalidate cache.
+    """
+    import boto3
+    from botocore.client import Config
+
+    # Making it possible to access public S3 buckets
+    # Workaround for https://github.com/boto/botocore/issues/2442
+    if signature_version.lower() == "unsigned":
+        from botocore import UNSIGNED
+
+        signature_version = UNSIGNED
+
+    return boto3.client(
+        "s3",
+        config=Config(
+            signature_version=signature_version, s3={"addressing_style": addressing_style}
+        ),
+        endpoint_url=s3_endpoint_url,
+        verify=verify,
+        aws_access_key_id=access_key_id,
+        aws_secret_access_key=secret_access_key,
+        aws_session_token=session_token,
+        region_name=region_name,
+    )
+
+
+def _get_s3_client(
+    addressing_style=None,
+    access_key_id=None,
+    secret_access_key=None,
+    session_token=None,
+    region_name=None,
+    s3_endpoint_url=None,
+):
+    if not s3_endpoint_url:
+        s3_endpoint_url = MLFLOW_S3_ENDPOINT_URL.get()
+    do_verify = not MLFLOW_S3_IGNORE_TLS.get()
+
+    # The valid verify argument value is None/False/path to cert bundle file, See
+    # https://github.com/boto/boto3/blob/73865126cad3938ca80a2f567a1c79cb248169a7/
+    # boto3/session.py#L212
+    verify = None if do_verify else False
+
+    # NOTE: If you need to specify this env variable, please file an issue at
+    # https://github.com/mlflow/mlflow/issues so we know your use-case!
+    signature_version = os.environ.get("MLFLOW_EXPERIMENTAL_S3_SIGNATURE_VERSION", "s3v4")
+
+    # Invalidate cache every `_MAX_CACHE_SECONDS`
+    timestamp = int(_get_utcnow_timestamp() / _MAX_CACHE_SECONDS)
+
+    if not addressing_style:
+        addressing_style = MLFLOW_BOTO_CLIENT_ADDRESSING_STYLE.get()
+
+    return _cached_get_s3_client(
+        signature_version,
+        addressing_style,
+        s3_endpoint_url,
+        verify,
+        timestamp,
+        access_key_id=access_key_id,
+        secret_access_key=secret_access_key,
+        session_token=session_token,
+        region_name=region_name,
+    )
+
+
+class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """Stores artifacts on Amazon S3."""
+
+    def __init__(
+        self,
+        artifact_uri: str,
+        tracking_uri: Optional[str] = None,
+        access_key_id=None,
+        secret_access_key=None,
+        session_token=None,
+    ) -> None:
+        super().__init__(artifact_uri, tracking_uri)
+        self._access_key_id = access_key_id
+        self._secret_access_key = secret_access_key
+        self._session_token = session_token
+
+    def _get_s3_client(self):
+        return _get_s3_client(
+            access_key_id=self._access_key_id,
+            secret_access_key=self._secret_access_key,
+            session_token=self._session_token,
+        )
+
+    def parse_s3_compliant_uri(self, uri):
+        """Parse an S3 URI, returning (bucket, path)"""
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "s3":
+            raise Exception(f"Not an S3 URI: {uri}")
+        path = parsed.path
+        if path.startswith("/"):
+            path = path[1:]
+        return parsed.netloc, path
+
+    @staticmethod
+    def get_s3_file_upload_extra_args():
+        s3_file_upload_extra_args = MLFLOW_S3_UPLOAD_EXTRA_ARGS.get()
+        if s3_file_upload_extra_args:
+            return json.loads(s3_file_upload_extra_args)
+        else:
+            return None
+
+    def _upload_file(self, s3_client, local_file, bucket, key):
+        extra_args = {}
+        guessed_type, guessed_encoding = guess_type(local_file)
+        if guessed_type is not None:
+            extra_args["ContentType"] = guessed_type
+        if guessed_encoding is not None:
+            extra_args["ContentEncoding"] = guessed_encoding
+        environ_extra_args = self.get_s3_file_upload_extra_args()
+        if environ_extra_args is not None:
+            extra_args.update(environ_extra_args)
+        s3_client.upload_file(Filename=local_file, Bucket=bucket, Key=key, ExtraArgs=extra_args)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        self._upload_file(
+            s3_client=self._get_s3_client(), local_file=local_file, bucket=bucket, key=dest_path
+        )
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        s3_client = self._get_s3_client()
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                upload_path = posixpath.join(dest_path, rel_path)
+
+            for f in filenames:
+                self._upload_file(
+                    s3_client=s3_client,
+                    local_file=os.path.join(root, f),
+                    bucket=bucket,
+                    key=posixpath.join(upload_path, f),
+                )
+
+    def list_artifacts(self, path=None):
+        (bucket, artifact_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        dest_path = artifact_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        infos = []
+        prefix = dest_path + "/" if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/")
+        for result in results:
+            # Subdirectories will be listed as "common prefixes" due to the way we made the request
+            for obj in result.get("CommonPrefixes", []):
+                subdir_path = obj.get("Prefix")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=subdir_path, artifact_path=artifact_path
+                )
+                subdir_rel_path = posixpath.relpath(path=subdir_path, start=artifact_path)
+                if subdir_rel_path.endswith("/"):
+                    subdir_rel_path = subdir_rel_path[:-1]
+                infos.append(FileInfo(subdir_rel_path, True, None))
+            # Objects listed directly will be files
+            for obj in result.get("Contents", []):
+                file_path = obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=artifact_path
+                )
+                file_rel_path = posixpath.relpath(path=file_path, start=artifact_path)
+                file_size = int(obj.get("Size"))
+                infos.append(FileInfo(file_rel_path, False, file_size))
+        return sorted(infos, key=lambda f: f.path)
+
+    @staticmethod
+    def _verify_listed_object_contains_artifact_path_prefix(listed_object_path, artifact_path):
+        if not listed_object_path.startswith(artifact_path):
+            raise MlflowException(
+                "The path of the listed S3 object does not begin with the specified"
+                f" artifact path. Artifact path: {artifact_path}. Object path:"
+                f" {listed_object_path}."
+            )
+
+    def _download_file(self, remote_file_path, local_path):
+        (bucket, s3_root_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        s3_full_path = posixpath.join(s3_root_path, remote_file_path)
+        s3_client = self._get_s3_client()
+        s3_client.download_file(bucket, s3_full_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=bucket, Prefix=dest_path)
+        for result in results:
+            keys = []
+            for to_delete_obj in result.get("Contents", []):
+                file_path = to_delete_obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=dest_path
+                )
+                keys.append({"Key": file_path})
+            if keys:
+                s3_client.delete_objects(Bucket=bucket, Delete={"Objects": keys})
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        s3_client = self._get_s3_client()
+        create_response = s3_client.create_multipart_upload(
+            Bucket=bucket,
+            Key=dest_path,
+        )
+        upload_id = create_response["UploadId"]
+        credentials = []
+        for i in range(1, num_parts + 1):  # part number must be in [1, 10000]
+            url = s3_client.generate_presigned_url(
+                "upload_part",
+                Params={
+                    "Bucket": bucket,
+                    "Key": dest_path,
+                    "PartNumber": i,
+                    "UploadId": upload_id,
+                },
+            )
+            credentials.append(
+                MultipartUploadCredential(
+                    url=url,
+                    part_number=i,
+                    headers={},
+                )
+            )
+        return CreateMultipartUploadResponse(
+            credentials=credentials,
+            upload_id=upload_id,
+        )
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        parts = [{"PartNumber": part.part_number, "ETag": part.etag} for part in parts]
+        s3_client = self._get_s3_client()
+        s3_client.complete_multipart_upload(
+            Bucket=bucket, Key=dest_path, UploadId=upload_id, MultipartUpload={"Parts": parts}
+        )
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        s3_client = self._get_s3_client()
+        s3_client.abort_multipart_upload(
+            Bucket=bucket,
+            Key=dest_path,
+            UploadId=upload_id,
+        )

mlflow/store/artifact/sftp_artifact_repo.py

@@ -0,0 +1,141 @@
+import os
+import posixpath
+import sys
+import urllib.parse
+from contextlib import contextmanager
+from queue import Queue
+from typing import Optional
+
+from mlflow.entities import FileInfo
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+
+
+# Based on: https://stackoverflow.com/a/58466685
+def _put_r_for_windows(sftp, local_dir, remote_dir, preserve_mtime=False):
+    for entry in os.listdir(local_dir):
+        local_path = os.path.join(local_dir, entry)
+        remote_path = posixpath.join(remote_dir, entry)
+        if os.path.isdir(local_path):
+            sftp.mkdir(remote_path)
+            _put_r_for_windows(sftp, local_path, remote_path, preserve_mtime)
+        else:
+            sftp.put(local_path, remote_path, preserve_mtime=preserve_mtime)
+
+
+class _SftpPool:
+    def __init__(self, connections):
+        self._idle_connection_queue = Queue()
+        for c in connections:
+            self._idle_connection_queue.put(c)
+
+    @contextmanager
+    def get_sftp_connection(self):
+        c = self._idle_connection_queue.get(block=True)
+        try:
+            yield c
+        finally:
+            self._idle_connection_queue.put(c)
+
+
+class SFTPArtifactRepository(ArtifactRepository):
+    """Stores artifacts as files in a remote directory, via sftp."""
+
+    def __init__(self, artifact_uri: str, tracking_uri: Optional[str] = None) -> None:
+        super().__init__(artifact_uri, tracking_uri)
+        parsed = urllib.parse.urlparse(artifact_uri)
+        self.config = {
+            "host": parsed.hostname,
+            "port": parsed.port,
+            "username": parsed.username,
+            "password": parsed.password,
+        }
+        self.path = parsed.path or "/"
+
+        import paramiko
+        import pysftp
+
+        if self.config["host"] is None:
+            self.config["host"] = "localhost"
+
+        ssh_config = paramiko.SSHConfig()
+        user_config_file = os.path.expanduser("~/.ssh/config")
+        if os.path.exists(user_config_file):
+            with open(user_config_file) as f:
+                ssh_config.parse(f)
+
+        user_config = ssh_config.lookup(self.config["host"])
+
+        if "hostname" in user_config:
+            self.config["host"] = user_config["hostname"]
+
+        if self.config.get("username", None) is None and "user" in user_config:
+            self.config["username"] = user_config["user"]
+
+        if self.config.get("port", None) is None:
+            if "port" in user_config:
+                self.config["port"] = int(user_config["port"])
+            else:
+                self.config["port"] = 22
+
+        if "identityfile" in user_config:
+            self.config["private_key"] = user_config["identityfile"][0]
+
+        connections = [pysftp.Connection(**self.config) for _ in range(self.max_workers)]
+        self.pool = _SftpPool(connections)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            sftp.put(local_file, posixpath.join(artifact_dir, os.path.basename(local_file)))
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            if sys.platform == "win32":
+                _put_r_for_windows(sftp, local_dir, artifact_dir)
+            else:
+                sftp.put_r(local_dir, artifact_dir)
+
+    def _is_directory(self, artifact_path):
+        artifact_dir = self.path
+        path = posixpath.join(artifact_dir, artifact_path) if artifact_path else artifact_dir
+        with self.pool.get_sftp_connection() as sftp:
+            return sftp.isdir(path)
+
+    def list_artifacts(self, path=None):
+        artifact_dir = self.path
+        list_dir = posixpath.join(artifact_dir, path) if path else artifact_dir
+        with self.pool.get_sftp_connection() as sftp:
+            if not sftp.isdir(list_dir):
+                return []
+            artifact_files = sftp.listdir(list_dir)
+            infos = []
+            for file_name in artifact_files:
+                file_path = file_name if path is None else posixpath.join(path, file_name)
+                full_file_path = posixpath.join(list_dir, file_name)
+                if sftp.isdir(full_file_path):
+                    infos.append(FileInfo(file_path, True, None))
+                else:
+                    infos.append(FileInfo(file_path, False, sftp.stat(full_file_path).st_size))
+            return infos
+
+    def _download_file(self, remote_file_path, local_path):
+        remote_full_path = posixpath.join(self.path, remote_file_path)
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.get(remote_full_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            self._delete_inner(artifact_dir, sftp)
+
+    def _delete_inner(self, artifact_path, sftp):
+        if sftp.isdir(artifact_path):
+            with sftp.cd(artifact_path):
+                for element in sftp.listdir():
+                    self._delete_inner(element, sftp)
+            sftp.rmdir(artifact_path)
+        elif sftp.isfile(artifact_path):
+            sftp.remove(artifact_path)

mlflow/store/artifact/uc_volume_artifact_repo.py

@@ -0,0 +1,76 @@
+from typing import Optional
+
+import mlflow.utils.databricks_utils
+from mlflow.environment_variables import MLFLOW_ENABLE_UC_VOLUME_FUSE_ARTIFACT_REPO
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.databricks_sdk_artifact_repo import DatabricksSdkArtifactRepository
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.utils.uri import (
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_model_registry_artifacts_uri,
+    is_valid_uc_volumes_uri,
+    remove_databricks_profile_info_from_artifact_uri,
+    strip_scheme,
+)
+
+
+class UCVolumesArtifactRepository(DatabricksSdkArtifactRepository):
+    """
+    Stores artifacts on UC Volumes using the Files REST API.
+    """
+
+    def __init__(self, artifact_uri: str, tracking_uri: Optional[str] = None) -> None:
+        if not is_valid_uc_volumes_uri(artifact_uri):
+            raise MlflowException(
+                message=(
+                    f"UC volumes URI must be of the form "
+                    f"dbfs:/Volumes/<catalog>/<schema>/<volume>/<path>: {artifact_uri}"
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        uri = remove_databricks_profile_info_from_artifact_uri(artifact_uri)
+        super().__init__("/" + strip_scheme(uri).strip("/"), tracking_uri)
+
+
+def uc_volume_artifact_repo_factory(artifact_uri: str, tracking_uri: Optional[str] = None):
+    """
+    Returns an ArtifactRepository subclass for storing artifacts on Volumes.
+
+    This factory method is used with URIs of the form ``dbfs:/Volumes/<path>``. Volume-backed
+    artifact storage can only be used together with the RestStore.
+
+    Args:
+        artifact_uri: Volume root artifact URI.
+        tracking_uri: The tracking URI.
+
+    Returns:
+        Subclass of ArtifactRepository capable of storing artifacts on DBFS.
+    """
+    if not is_valid_uc_volumes_uri(artifact_uri):
+        raise MlflowException(
+            message=(
+                f"UC volumes URI must be of the form "
+                f"dbfs:/Volumes/<catalog>/<schema>/<volume>/<path>: {artifact_uri}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    artifact_uri = artifact_uri.rstrip("/")
+    db_profile_uri = get_databricks_profile_uri_from_artifact_uri(artifact_uri)
+    if (
+        mlflow.utils.databricks_utils.is_uc_volume_fuse_available()
+        and MLFLOW_ENABLE_UC_VOLUME_FUSE_ARTIFACT_REPO.get()
+        and not is_databricks_model_registry_artifacts_uri(artifact_uri)
+        and (db_profile_uri is None or db_profile_uri == "databricks")
+    ):
+        # If the UC Volume FUSE mount is available, write artifacts directly to
+        # /Volumes/... using local filesystem APIs.
+        # Note: it is possible for a named Databricks profile to point to the current workspace,
+        # but we're going to avoid doing a complex check and assume users will use `databricks`
+        # to mean the current workspace. Using `UCVolumesArtifactRepository` to access
+        # the current workspace's Volumes should still work; it just may be slower.
+        uri_without_profile = remove_databricks_profile_info_from_artifact_uri(artifact_uri)
+        path = strip_scheme(uri_without_profile).lstrip("/")
+        return LocalArtifactRepository(f"file:///{path}", tracking_uri)
+    return UCVolumesArtifactRepository(artifact_uri, tracking_uri)