genesis-flow 1.0.0__py3-none-any.whl

mlflow/store/artifact/http_artifact_repo.py

@@ -0,0 +1,218 @@
+import logging
+import os
+import posixpath
+
+import requests
+from requests import HTTPError
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+    MultipartUploadPart,
+)
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_PROXY_MULTIPART_UPLOAD,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+    MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE,
+)
+from mlflow.exceptions import MlflowException, _UnsupportedMultipartUploadException
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    MultipartUploadMixin,
+    verify_artifact_path,
+)
+from mlflow.store.artifact.cloud_artifact_repo import _complete_futures, _compute_num_chunks
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.file_utils import read_chunk, relative_path_to_artifact_path
+from mlflow.utils.mime_type_utils import _guess_mime_type
+from mlflow.utils.rest_utils import augmented_raise_for_status, http_request
+from mlflow.utils.uri import validate_path_is_safe
+
+_logger = logging.getLogger(__name__)
+
+
+class HttpArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """Stores artifacts in a remote artifact storage using HTTP requests"""
+
+    @property
+    def _host_creds(self):
+        return get_default_host_creds(self.artifact_uri)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        verify_artifact_path(artifact_path)
+
+        # Try to perform multipart upload if the file is large.
+        # If the server does not support, or if the upload failed, revert to normal upload.
+        if (
+            MLFLOW_ENABLE_PROXY_MULTIPART_UPLOAD.get()
+            and os.path.getsize(local_file) >= MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE.get()
+        ):
+            try:
+                self._try_multipart_upload(local_file, artifact_path)
+                return
+            except _UnsupportedMultipartUploadException:
+                pass
+
+        file_name = os.path.basename(local_file)
+        mime_type = _guess_mime_type(file_name)
+        paths = (artifact_path, file_name) if artifact_path else (file_name,)
+        endpoint = posixpath.join("/", *paths)
+        extra_headers = {"Content-Type": mime_type}
+        with open(local_file, "rb") as f:
+            resp = http_request(
+                self._host_creds, endpoint, "PUT", data=f, extra_headers=extra_headers
+            )
+            augmented_raise_for_status(resp)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            if root == local_dir:
+                artifact_dir = artifact_path
+            else:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                artifact_dir = (
+                    posixpath.join(artifact_path, rel_path) if artifact_path else rel_path
+                )
+            for f in filenames:
+                self.log_artifact(os.path.join(root, f), artifact_dir)
+
+    def list_artifacts(self, path=None):
+        endpoint = "/mlflow-artifacts/artifacts"
+        url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
+        root = tail.lstrip("/")
+        params = {"path": posixpath.join(root, path) if path else root}
+        host_creds = get_default_host_creds(url)
+        resp = http_request(host_creds, endpoint, "GET", params=params)
+        augmented_raise_for_status(resp)
+        file_infos = []
+        for f in resp.json().get("files", []):
+            validated_path = validate_path_is_safe(f["path"])
+            file_info = FileInfo(
+                posixpath.join(path, validated_path) if path else validated_path,
+                f["is_dir"],
+                int(f["file_size"]) if ("file_size" in f) else None,
+            )
+            file_infos.append(file_info)
+
+        return sorted(file_infos, key=lambda f: f.path)
+
+    def _download_file(self, remote_file_path, local_path):
+        endpoint = posixpath.join("/", remote_file_path)
+        resp = http_request(self._host_creds, endpoint, "GET", stream=True)
+        augmented_raise_for_status(resp)
+        with open(local_path, "wb") as f:
+            chunk_size = 1024 * 1024  # 1 MB
+            for chunk in resp.iter_content(chunk_size=chunk_size):
+                f.write(chunk)
+
+    def delete_artifacts(self, artifact_path=None):
+        endpoint = posixpath.join("/", artifact_path) if artifact_path else "/"
+        resp = http_request(self._host_creds, endpoint, "DELETE", stream=True)
+        augmented_raise_for_status(resp)
+
+    def _construct_mpu_uri_and_path(self, base_endpoint, artifact_path):
+        uri, path = self.artifact_uri.split("/mlflow-artifacts/artifacts", maxsplit=1)
+        path = path.strip("/")
+        endpoint = (
+            posixpath.join(base_endpoint, path, artifact_path)
+            if artifact_path
+            else posixpath.join(base_endpoint, path)
+        )
+        return uri, endpoint
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/create", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "num_parts": num_parts,
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+        return CreateMultipartUploadResponse.from_dict(resp.json())
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/complete", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "upload_id": upload_id,
+            "parts": [part.to_dict() for part in parts],
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/abort", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "upload_id": upload_id,
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+
+    @staticmethod
+    def _upload_part(credential: MultipartUploadCredential, local_file, size, start_byte):
+        data = read_chunk(local_file, size, start_byte)
+        response = requests.put(credential.url, data=data, headers=credential.headers)
+        augmented_raise_for_status(response)
+        return MultipartUploadPart(
+            part_number=credential.part_number,
+            etag=response.headers.get("ETag", ""),
+            url=credential.url,
+        )
+
+    def _try_multipart_upload(self, local_file, artifact_path=None):
+        """
+        Attempts to perform multipart upload to log an artifact.
+        Returns if the multipart upload is successful.
+        Raises UnsupportedMultipartUploadException if multipart upload is unsupported.
+        """
+        chunk_size = MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+        num_parts = _compute_num_chunks(local_file, chunk_size)
+
+        try:
+            create = self.create_multipart_upload(local_file, num_parts, artifact_path)
+        except HTTPError as e:
+            # return False if server does not support multipart upload
+            error_message = e.response.json().get("message", "")
+            if isinstance(error_message, str) and error_message.startswith(
+                _UnsupportedMultipartUploadException.MESSAGE
+            ):
+                raise _UnsupportedMultipartUploadException()
+            raise
+
+        try:
+            futures = {}
+            for i, credential in enumerate(create.credentials):
+                future = self.thread_pool.submit(
+                    self._upload_part,
+                    credential=credential,
+                    local_file=local_file,
+                    size=chunk_size,
+                    start_byte=chunk_size * i,
+                )
+                futures[future] = credential.part_number
+
+            parts, errors = _complete_futures(futures, local_file)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {local_file}. Errors: {errors}"
+                )
+
+            parts = sorted(parts.values(), key=lambda part: part.part_number)
+            self.complete_multipart_upload(local_file, create.upload_id, parts, artifact_path)
+        except Exception as e:
+            self.abort_multipart_upload(local_file, create.upload_id, artifact_path)
+            _logger.warning(f"Failed to upload file {local_file} using multipart upload: {e}")
+            raise

mlflow/store/artifact/local_artifact_repo.py

@@ -0,0 +1,142 @@
+import os
+import shutil
+from typing import Any, Optional
+
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    try_read_trace_data,
+    verify_artifact_path,
+)
+from mlflow.tracing.utils.artifact_utils import TRACE_DATA_FILE_NAME
+from mlflow.utils.file_utils import (
+    get_file_info,
+    list_all,
+    local_file_uri_to_path,
+    mkdir,
+    relative_path_to_artifact_path,
+)
+from mlflow.utils.uri import validate_path_is_safe
+
+
+class LocalArtifactRepository(ArtifactRepository):
+    """Stores artifacts as files in a local directory."""
+
+    def __init__(self, artifact_uri: str, tracking_uri: Optional[str] = None) -> None:
+        super().__init__(artifact_uri, tracking_uri)
+        self._artifact_dir = local_file_uri_to_path(self.artifact_uri)
+
+    @property
+    def artifact_dir(self):
+        return self._artifact_dir
+
+    def log_artifact(self, local_file, artifact_path=None):
+        verify_artifact_path(artifact_path)
+        # NOTE: The artifact_path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if artifact_path:
+            artifact_path = os.path.normpath(artifact_path)
+
+        artifact_dir = (
+            os.path.join(self.artifact_dir, artifact_path) if artifact_path else self.artifact_dir
+        )
+        if not os.path.exists(artifact_dir):
+            mkdir(artifact_dir)
+        try:
+            shutil.copy2(local_file, os.path.join(artifact_dir, os.path.basename(local_file)))
+        except shutil.SameFileError:
+            pass
+
+    def _is_directory(self, artifact_path):
+        # NOTE: The path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        path = os.path.normpath(artifact_path) if artifact_path else ""
+        list_dir = os.path.join(self.artifact_dir, path) if path else self.artifact_dir
+        return os.path.isdir(list_dir)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        verify_artifact_path(artifact_path)
+        # NOTE: The artifact_path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if artifact_path:
+            artifact_path = os.path.normpath(artifact_path)
+        artifact_dir = (
+            os.path.join(self.artifact_dir, artifact_path) if artifact_path else self.artifact_dir
+        )
+        if not os.path.exists(artifact_dir):
+            mkdir(artifact_dir)
+        shutil.copytree(src=local_dir, dst=artifact_dir, dirs_exist_ok=True)
+
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Artifacts tracked by ``LocalArtifactRepository`` already exist on the local filesystem.
+        If ``dst_path`` is ``None``, the absolute filesystem path of the specified artifact is
+        returned. If ``dst_path`` is not ``None``, the local artifact is copied to ``dst_path``.
+
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            dst_path: Absolute path of the local filesystem destination directory to which to
+                download the specified artifacts. This directory must already exist. If
+                unspecified, the absolute path of the local artifact will be returned.
+
+        Returns:
+            Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        if dst_path:
+            return super().download_artifacts(artifact_path, dst_path)
+        # NOTE: The artifact_path is expected to be a relative path in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        artifact_path = validate_path_is_safe(artifact_path)
+        local_artifact_path = os.path.join(self.artifact_dir, os.path.normpath(artifact_path))
+        if not os.path.exists(local_artifact_path):
+            raise OSError(f"No such file or directory: '{local_artifact_path}'")
+        return os.path.abspath(local_artifact_path)
+
+    def list_artifacts(self, path=None):
+        # NOTE: The path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if path:
+            path = os.path.normpath(path)
+        list_dir = os.path.join(self.artifact_dir, path) if path else self.artifact_dir
+        if os.path.isdir(list_dir):
+            artifact_files = list_all(list_dir, full_path=True)
+            infos = [
+                get_file_info(
+                    f, relative_path_to_artifact_path(os.path.relpath(f, self.artifact_dir))
+                )
+                for f in artifact_files
+            ]
+            return sorted(infos, key=lambda f: f.path)
+        else:
+            return []
+
+    def _download_file(self, remote_file_path, local_path):
+        # NOTE: The remote_file_path is expected to be a relative path in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        remote_file_path = validate_path_is_safe(remote_file_path)
+        remote_file_path = os.path.join(self.artifact_dir, os.path.normpath(remote_file_path))
+        shutil.copy2(remote_file_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        artifact_path = local_file_uri_to_path(
+            os.path.join(self._artifact_dir, artifact_path) if artifact_path else self._artifact_dir
+        )
+
+        if os.path.exists(artifact_path):
+            if os.path.isfile(artifact_path):
+                os.remove(artifact_path)
+            else:
+                shutil.rmtree(artifact_path)
+
+    def download_trace_data(self) -> dict[str, Any]:
+        """
+        Download the trace data.
+
+        Returns:
+            The trace data as a dictionary.
+
+        Raises:
+            - `MlflowTraceDataNotFound`: The trace data is not found.
+            - `MlflowTraceDataCorrupted`: The trace data is corrupted.
+        """
+        trace_data_path = os.path.join(self.artifact_dir, TRACE_DATA_FILE_NAME)
+        return try_read_trace_data(trace_data_path)

mlflow/store/artifact/mlflow_artifacts_repo.py

@@ -0,0 +1,94 @@
+import re
+from typing import Optional
+from urllib.parse import urlparse, urlunparse
+
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.http_artifact_repo import HttpArtifactRepository
+from mlflow.tracking._tracking_service.utils import get_tracking_uri
+
+
+def _check_if_host_is_numeric(hostname):
+    if hostname:
+        try:
+            float(hostname)
+            return True
+        except ValueError:
+            return False
+    else:
+        return False
+
+
+def _validate_port_mapped_to_hostname(uri_parse):
+    # This check is to catch an mlflow-artifacts uri that has a port designated but no
+    # hostname specified. `urllib.parse.urlparse` will treat such a uri as a filesystem
+    # definition, mapping the provided port as a hostname value if this condition is not
+    # validated.
+    if uri_parse.hostname and _check_if_host_is_numeric(uri_parse.hostname) and not uri_parse.port:
+        raise MlflowException(
+            "The mlflow-artifacts uri was supplied with a port number: "
+            f"{uri_parse.hostname}, but no host was defined."
+        )
+
+
+def _validate_uri_scheme(parsed_uri):
+    allowable_schemes = {"http", "https"}
+    if parsed_uri.scheme not in allowable_schemes:
+        raise MlflowException(
+            "When an mlflow-artifacts URI was supplied, the tracking URI must be a valid "
+            f"http or https URI, but it was currently set to {parsed_uri.geturl()}. "
+            "Perhaps you forgot to set the tracking URI to the running MLflow server. "
+            "To set the tracking URI, use either of the following methods:\n"
+            "1. Set the MLFLOW_TRACKING_URI environment variable to the desired tracking URI. "
+            "`export MLFLOW_TRACKING_URI=http://localhost:5000`\n"
+            "2. Set the tracking URI programmatically by calling `mlflow.set_tracking_uri`. "
+            "`mlflow.set_tracking_uri('http://localhost:5000')`"
+        )
+
+
+class MlflowArtifactsRepository(HttpArtifactRepository):
+    """Scheme wrapper around HttpArtifactRepository for mlflow-artifacts server functionality"""
+
+    def __init__(self, artifact_uri: str, tracking_uri: Optional[str] = None) -> None:
+        effective_tracking_uri = tracking_uri or get_tracking_uri()
+        super().__init__(self.resolve_uri(artifact_uri, effective_tracking_uri), tracking_uri)
+
+    @classmethod
+    def resolve_uri(cls, artifact_uri, tracking_uri):
+        base_url = "/api/2.0/mlflow-artifacts/artifacts"
+
+        track_parse = urlparse(tracking_uri)
+
+        uri_parse = urlparse(artifact_uri)
+
+        # Check to ensure that a port is present with no hostname
+        _validate_port_mapped_to_hostname(uri_parse)
+
+        # Check that tracking uri is http or https
+        _validate_uri_scheme(track_parse)
+
+        if uri_parse.path == "/":  # root directory; build simple path
+            resolved = f"{base_url}{uri_parse.path}"
+        elif uri_parse.path == base_url:  # for operations like list artifacts
+            resolved = base_url
+        else:
+            resolved = f"{track_parse.path}/{base_url}/{uri_parse.path}"
+        resolved = re.sub("//+", "/", resolved)
+
+        resolved_artifacts_uri = urlunparse(
+            (
+                # scheme
+                track_parse.scheme,
+                # netloc
+                uri_parse.netloc if uri_parse.netloc else track_parse.netloc,
+                # path
+                resolved,
+                # params
+                "",
+                # query
+                "",
+                # fragment
+                "",
+            )
+        )
+
+        return resolved_artifacts_uri.replace("///", "/").rstrip("/")