Avoid to load an entire artifact file on memory

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/<study_id:int>/<trial_id:int>/<artifact_id:re:[0-9a-fA-F-]+>")
-    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/<study_id:int>/<trial_id:int>")
     @json_api_view

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.