From bdc35ce3917ed47f9b2167da283055bdfa628ca1 Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Sat, 7 Jan 2023 23:56:33 +0900
Subject: [PATCH] Add boto3 backend

---
 .github/workflows/python-tests.yml  |  4 ++-
 optuna_dashboard/artifact/boto3.py  | 48 +++++++++++++++++++++++++++++
 python_tests/test_boto3_artifact.py | 41 ++++++++++++++++++++++++
 requirements.txt                    |  8 ++++-
 4 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 optuna_dashboard/artifact/boto3.py
 create mode 100644 python_tests/test_boto3_artifact.py

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index b927c705..3ba43912 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -22,7 +22,7 @@ jobs:
           python -m pip install --progress-bar off --upgrade pip setuptools
           pip install --progress-bar off "optuna>=3.0.0"
           pip install --progress-bar off .
-          pip install --progress-bar off flake8 black isort mypy
+          pip install --progress-bar off flake8 black isort mypy mypy-boto3-s3
       - run: flake8 . --show-source
       - run: black --check --diff .
       - run: isort --check --diff .
@@ -45,6 +45,7 @@ jobs:
         # python_tests requires optuna>=3.0.0 since it imports FloatDistribution
         run: |
           python -m pip install --progress-bar off --upgrade pip setuptools
+          pip install boto3 moto[s3]
           pip install --progress-bar off "optuna>=3.0.0"
           pip install --progress-bar off .
       - run: python -m unittest
@@ -60,6 +61,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --progress-bar off --upgrade pip setuptools
+          pip install boto3 moto[s3]
           pip install --progress-bar off .
           python -m pip install --progress-bar off --upgrade git+https://github.com/optuna/optuna.git
       - run: python -m unittest
diff --git a/optuna_dashboard/artifact/boto3.py b/optuna_dashboard/artifact/boto3.py
new file mode 100644
index 00000000..5af13189
--- /dev/null
+++ b/optuna_dashboard/artifact/boto3.py
@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import boto3
+
+
+if TYPE_CHECKING:
+    from typing import Optional
+    from typing import BinaryIO
+
+    from mypy_boto3_s3 import S3Client
+
+
+class Boto3Backend:
+    """An artifact backend for S3.
+
+    Example:
+       .. code-block:: python
+
+          import optuna
+          from optuna_dashboard.artifact import upload_artifact
+          from optuna_dashboard.artifact.boto3 import Boto3Backend
+
+          artifact_backend = Boto3Backend("my-bucket")
+
+          def objective(trial: optuna.Trial) -> float:
+              ... = trial.suggest_float("x", -10, 10)
+              file_path = generate_example_png(...)
+              upload_artifact(artifact_backend, trial, file_path)
+              return ...
+    """
+
+    def __init__(self, bucket_name: str, client: Optional[S3Client] = None) -> None:
+        self.bucket = bucket_name
+        self.client = client or boto3.client("s3")
+
+    def open(self, artifact_id: str) -> BinaryIO:
+        obj = self.client.get_object(Bucket=self.bucket, Key=artifact_id)
+        body = obj.get("Body")
+        assert body is not None
+        return body  # type: ignore
+
+    def write(self, artifact_id: str, content_body: BinaryIO) -> None:
+        self.client.upload_fileobj(content_body, self.bucket, artifact_id)
+
+    def remove(self, artifact_id: str) -> None:
+        self.client.delete_object(Bucket=self.bucket, Key=artifact_id)
diff --git a/python_tests/test_boto3_artifact.py b/python_tests/test_boto3_artifact.py
new file mode 100644
index 00000000..cafb9f91
--- /dev/null
+++ b/python_tests/test_boto3_artifact.py
@@ -0,0 +1,41 @@
+import io
+from unittest import TestCase
+
+import boto3
+from moto import mock_s3
+from optuna_dashboard.artifact.boto3 import Boto3Backend
+
+
+@mock_s3
+class Boto3BackendTestCase(TestCase):
+    def setUp(self) -> None:
+        self.s3_client = boto3.resource("s3")
+        self.bucket = self.s3_client.create_bucket(Bucket="moto-bucket")
+
+    def tearDown(self) -> None:
+        self.bucket.objects.all().delete()
+        self.bucket.delete()
+
+    def test_upload_download(self) -> None:
+        artifact_id = "dummy-uuid"
+        dummy_content = b"Hello World"
+
+        backend = Boto3Backend(self.bucket.name)
+        backend.write(artifact_id, io.BytesIO(dummy_content))
+
+        objects = [obj for obj in self.bucket.objects.all() if obj.key == artifact_id]
+        assert len(objects) == 1
+        assert objects[0].get()["Body"].read() == dummy_content
+
+        with backend.open(artifact_id) as f:
+            actual = f.read()
+        self.assertEqual(actual, dummy_content)
+
+    def test_remove(self) -> None:
+        artifact_id = "dummy-uuid"
+        backend = Boto3Backend(self.bucket.name)
+        backend.write(artifact_id, io.BytesIO(b"Hello"))
+        assert len([obj for obj in self.bucket.objects.all() if obj.key == artifact_id]) == 1
+
+        backend.remove(artifact_id)
+        assert len([obj for obj in self.bucket.objects.all() if obj.key == artifact_id]) == 0
diff --git a/requirements.txt b/requirements.txt
index 7545e1a9..57ecdd36 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,12 +2,18 @@
 optuna>=2.4
 bottle
 scikit-learn
+typing-extensions;python_version<"3.8"
 
 # lint
 black
 flake8
 isort
+mypy-boto3-s3
 mypy
 
-# test
+# unit test
+boto3
+moto[s3]
+
+# visual regression tests
 pyppeteer