Merge pull request #48 from rekcurd/feature/gcs-support

Merged

Author: keigohtr

rekcurd/data_servers/__init__.py

@@ -12,6 +12,7 @@
 from .local_handler import LocalHandler
 from .ceph_handler import CephHandler
 from .aws_s3_handler import AwsS3Handler
+from .gcs_handler import GcsHandler
 
 
 class DataServer(object):
@@ -27,6 +28,8 @@ def __init__(self, config: RekcurdConfig):
             self._api_handler = CephHandler(config)
         elif config.MODEL_MODE_ENUM == ModelModeEnum.AWS_S3:
             self._api_handler = AwsS3Handler(config)
+        elif config.MODEL_MODE_ENUM == ModelModeEnum.GCS:
+            self._api_handler = GcsHandler(config)
         else:
             raise ValueError("Invalid ModelModeEnum value.")
 

rekcurd/data_servers/gcs_handler.py

@@ -0,0 +1,28 @@
+# coding: utf-8
+
+
+import boto3
+
+from .data_handler import DataHandler
+from rekcurd.utils import RekcurdConfig
+
+
+class GcsHandler(DataHandler):
+    """GcsHandler
+    """
+    def __init__(self, config: RekcurdConfig):
+        super(GcsHandler, self).__init__(config)
+        self._resource = boto3.resource(
+            's3',
+            region_name="auto",
+            endpoint_url="https://storage.googleapis.com",
+            aws_access_key_id=config.GCS_ACCESS_KEY,
+            aws_secret_access_key=config.GCS_SECRET_KEY,
+        )
+        self._bucket = config.GCS_BUCKET_NAME
+
+    def download(self, remote_filepath: str, local_filepath: str) -> None:
+        self._resource.Bucket(self._bucket).download_file(remote_filepath, local_filepath)
+
+    def upload(self, remote_filepath: str, local_filepath: str) -> None:
+        self._resource.Bucket(self._bucket).upload_file(local_filepath, remote_filepath)

rekcurd/template/settings.yml-tpl

@@ -28,5 +28,8 @@ model:
     access_key: xxxxx               #  AWS access key.
     secret_key: xxxxx               #  AWS secret key.
     bucket: xxxxx                   #  AWS bucket name.
-  gcs:                              # :TODO: GCS parameters.
+  gcs:                              # GCS access parameters. Required for "gcs" mode. To generate keys, see "https://cloud.google.com/storage/docs/migrating#keys"
     filepath: model/default.model   #  ML model file. Default "model/default.model"
+    access_key: xxxxx               #  GCS access key.
+    secret_key: xxxxx               #  GCS secret key.
+    bucket: xxxxx                   #  GCS bucket name.

rekcurd/utils/rekcurd_config.py

@@ -56,7 +56,9 @@ class RekcurdConfig:
     AWS_ACCESS_KEY: str = None
     AWS_SECRET_KEY: str = None
     AWS_BUCKET_NAME: str = None
-    # TODO: GCS
+    GCS_ACCESS_KEY: str = None
+    GCS_SECRET_KEY: str = None
+    GCS_BUCKET_NAME: str = None
 
     def __init__(self, config_file: str = None):
         self.KUBERNETES_MODE = os.getenv("REKCURD_KUBERNETES_MODE")
@@ -74,6 +76,7 @@ def set_configurations(
             ceph_port: int = None, ceph_is_secure: bool = None,
             ceph_bucket_name: str = None, aws_access_key: str = None,
             aws_secret_key: str = None, aws_bucket_name: str = None,
+            gcs_access_key: str = None, gcs_secret_key: str = None, gcs_bucket_name: str = None,
             **options):
         self.DEBUG_MODE = debug_mode if debug_mode is not None else self.DEBUG_MODE
         self.APPLICATION_NAME = application_name or self.APPLICATION_NAME
@@ -92,7 +95,9 @@ def set_configurations(
         self.AWS_ACCESS_KEY = aws_access_key or self.AWS_ACCESS_KEY
         self.AWS_SECRET_KEY = aws_secret_key or self.AWS_SECRET_KEY
         self.AWS_BUCKET_NAME = aws_bucket_name or self.AWS_BUCKET_NAME
-        # TODO: GCS
+        self.GCS_ACCESS_KEY = gcs_access_key or self.GCS_ACCESS_KEY
+        self.GCS_SECRET_KEY = gcs_secret_key or self.GCS_SECRET_KEY
+        self.GCS_BUCKET_NAME = gcs_bucket_name or self.GCS_BUCKET_NAME
 
     def __load_from_file(self, config_file: str):
         if config_file is not None:
@@ -128,8 +133,13 @@ def __load_from_file(self, config_file: str):
             self.AWS_ACCESS_KEY = config_model_mode.get("access_key")
             self.AWS_SECRET_KEY = config_model_mode.get("secret_key")
             self.AWS_BUCKET_NAME = config_model_mode.get("bucket")
+        elif self.MODEL_MODE_ENUM == ModelModeEnum.GCS:
+            config_model_mode = config_model.get(model_mode, dict())
+            self.MODEL_FILE_PATH = config_model_mode.get("filepath", "model/default.model")
+            self.GCS_ACCESS_KEY = config_model_mode.get("access_key")
+            self.GCS_SECRET_KEY = config_model_mode.get("secret_key")
+            self.GCS_BUCKET_NAME = config_model_mode.get("bucket")
         else:
-            # TODO: GCS
             raise ValueError("'{}' is not supported as ModelModeEnum".format(model_mode))
 
     def __load_from_env(self):
@@ -151,4 +161,6 @@ def __load_from_env(self):
         self.AWS_ACCESS_KEY = os.getenv("REKCURD_AWS_ACCESS_KEY")
         self.AWS_SECRET_KEY = os.getenv("REKCURD_AWS_SECRET_KEY")
         self.AWS_BUCKET_NAME = os.getenv("REKCURD_AWS_BUCKET_NAME")
-        # TODO: GCS
+        self.GCS_ACCESS_KEY = os.getenv("REKCURD_GCS_ACCESS_KEY")
+        self.GCS_SECRET_KEY = os.getenv("REKCURD_GCS_SECRET_KEY")
+        self.GCS_BUCKET_NAME = os.getenv("REKCURD_GCS_BUCKET_NAME")

requirements.txt

@@ -4,4 +4,4 @@ grpcio>=1.13.0 # Apache-2.0
 grpcio-tools>=1.13.0 # Apache-2.0
 PyYAML>=3.12 # MIT
 boto>=2.49.0 # MIT
-boto3>=1.9.0 # Apache-2.0
+boto3>=1.9.38 # Apache-2.0

test/data_servers/__init__.py

@@ -16,6 +16,10 @@ def inner_method(*args, **kwargs):
                           new=Mock(return_value=None)) as _, \
                     patch('rekcurd.data_servers.AwsS3Handler.upload',
                           new=Mock(return_value=None)) as _, \
+                    patch('rekcurd.data_servers.GcsHandler.download',
+                          new=Mock(return_value=None)) as _, \
+                    patch('rekcurd.data_servers.GcsHandler.upload',
+                          new=Mock(return_value=None)) as _, \
                     patch('builtins.open', new_callable=mock_open) as _:
                 return func(*args, **kwargs)
         return inner_method

test/data_servers/test_gcs_handler.py

@@ -0,0 +1,26 @@
+import unittest
+
+from rekcurd.utils import RekcurdConfig, ModelModeEnum
+from rekcurd.data_servers import GcsHandler
+
+from . import patch_predictor
+
+
+class GcsHandlerTest(unittest.TestCase):
+    """Tests for GcsHandlerTest.
+    """
+
+    def setUp(self):
+        config = RekcurdConfig("./test/test-settings.yml")
+        config.set_configurations(
+            model_mode=ModelModeEnum.GCS.value, gcs_access_key="xxx",
+            gcs_secret_key="xxx", gcs_bucket_name="xxx")
+        self.handler = GcsHandler(config)
+
+    @patch_predictor()
+    def test_download(self):
+        self.assertIsNone(self.handler.download("remote","local"))
+
+    @patch_predictor()
+    def test_upload(self):
+        self.assertIsNone(self.handler.upload("remote","local"))