diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py index 99cffa78..d6883440 100644 --- a/optuna_dashboard/artifact/_backend.py +++ b/optuna_dashboard/artifact/_backend.py @@ -9,6 +9,7 @@ import uuid from bottle import BaseRequest from bottle import Bottle +from bottle import HTTPResponse from bottle import request from bottle import response import optuna @@ -58,7 +59,7 @@ def register_artifact_route( app: Bottle, storage: BaseStorage, artifact_backend: Optional[ArtifactBackend] ) -> None: @app.get("/artifacts///") - def proxy_artifact(study_id: int, trial_id: int, artifact_id: str) -> bytes: + def proxy_artifact(study_id: int, trial_id: int, artifact_id: str) -> HTTPResponse | bytes: if artifact_backend is None: response.status = 400 # Bad Request return b"Cannot access to the artifacts." @@ -66,13 +67,13 @@ def register_artifact_route( if artifact_dict is None: response.status = 404 return b"Not Found" - response.set_header("Content-Type", artifact_dict["mimetype"]) - if artifact_dict.get("encoding"): - response.set_header("Content-Encodings", artifact_dict.get("encoding")) + headers = {"Content-Type": artifact_dict["mimetype"]} + encoding = artifact_dict.get("encoding") + if encoding: + headers["Content-Encodings"] = encoding - with artifact_backend.open(artifact_id) as f: - body = f.read() - return body + fp = artifact_backend.open(artifact_id) + return HTTPResponse(fp, headers=headers) @app.post("/api/artifacts//") @json_api_view diff --git a/optuna_dashboard/artifact/boto3.py b/optuna_dashboard/artifact/boto3.py index f3fc1550..b248d5dc 100644 --- a/optuna_dashboard/artifact/boto3.py +++ b/optuna_dashboard/artifact/boto3.py @@ -9,7 +9,9 @@ import boto3 if TYPE_CHECKING: from typing import BinaryIO + from typing import IO from typing import Optional + from typing import TypeGuard from _typeshed import SupportsRead from mypy_boto3_s3 import S3Client @@ -45,6 +47,11 @@ class Boto3Backend: return body # type: ignore def write(self, artifact_id: str, content_body: SupportsRead[bytes]) -> None: + if _is_file_like_obj(content_body): + self.client.upload_fileobj(content_body, self.bucket, artifact_id) + return + + # Convert SupportsRead[bytes] to file-like object buf = io.BytesIO() shutil.copyfileobj(content_body, buf) buf.seek(0) @@ -54,6 +61,15 @@ class Boto3Backend: self.client.delete_object(Bucket=self.bucket, Key=artifact_id) +def _is_file_like_obj(obj: SupportsRead[bytes]) -> TypeGuard[IO[bytes]]: + return ( + isinstance(obj, io.TextIOBase) + or isinstance(obj, io.BufferedIOBase) + or isinstance(obj, io.RawIOBase) + or isinstance(obj, io.IOBase) + ) + + if TYPE_CHECKING: # A mypy-runtime assertion to ensure that Boto3Backend # implements all abstract methods in ArtifactBackendProtocol.