[HWORKS-1171] Make chunk size configurable when uploading models to t… (logicalclocks#228)

* [HWORKS-1171] Make chunk size configurable when uploading models to the model registry from outside

* Revert black changes to predictor_state.py

Author: SirOibaf

python/hsml/client/istio/grpc/proto/grpc_predict_v2_pb2.pyi

@@ -90,6 +90,7 @@ class ModelInferRequest(_message.Message):
                 key: _Optional[str] = ...,
                 value: _Optional[_Union[InferParameter, _Mapping]] = ...,
             ) -> None: ...
+
         CONTENTS_FIELD_NUMBER: _ClassVar[int]
         DATATYPE_FIELD_NUMBER: _ClassVar[int]
         NAME_FIELD_NUMBER: _ClassVar[int]
@@ -123,6 +124,7 @@ class ModelInferRequest(_message.Message):
                 key: _Optional[str] = ...,
                 value: _Optional[_Union[InferParameter, _Mapping]] = ...,
             ) -> None: ...
+
         NAME_FIELD_NUMBER: _ClassVar[int]
         PARAMETERS_FIELD_NUMBER: _ClassVar[int]
         name: str
@@ -144,6 +146,7 @@ class ModelInferRequest(_message.Message):
             key: _Optional[str] = ...,
             value: _Optional[_Union[InferParameter, _Mapping]] = ...,
         ) -> None: ...
+
     ID_FIELD_NUMBER: _ClassVar[int]
     INPUTS_FIELD_NUMBER: _ClassVar[int]
     MODEL_NAME_FIELD_NUMBER: _ClassVar[int]
@@ -201,6 +204,7 @@ class ModelInferResponse(_message.Message):
                 key: _Optional[str] = ...,
                 value: _Optional[_Union[InferParameter, _Mapping]] = ...,
             ) -> None: ...
+
         CONTENTS_FIELD_NUMBER: _ClassVar[int]
         DATATYPE_FIELD_NUMBER: _ClassVar[int]
         NAME_FIELD_NUMBER: _ClassVar[int]
@@ -231,6 +235,7 @@ class ModelInferResponse(_message.Message):
             key: _Optional[str] = ...,
             value: _Optional[_Union[InferParameter, _Mapping]] = ...,
         ) -> None: ...
+
     ID_FIELD_NUMBER: _ClassVar[int]
     MODEL_NAME_FIELD_NUMBER: _ClassVar[int]
     MODEL_VERSION_FIELD_NUMBER: _ClassVar[int]
@@ -284,6 +289,7 @@ class ModelMetadataResponse(_message.Message):
             datatype: _Optional[str] = ...,
             shape: _Optional[_Iterable[int]] = ...,
         ) -> None: ...
+
     INPUTS_FIELD_NUMBER: _ClassVar[int]
     NAME_FIELD_NUMBER: _ClassVar[int]
     OUTPUTS_FIELD_NUMBER: _ClassVar[int]

python/hsml/core/dataset_api.py

@@ -38,17 +38,21 @@ class DatasetApi:
     def __init__(self):
         pass
 
-    DEFAULT_FLOW_CHUNK_SIZE = 1048576
+    DEFAULT_UPLOAD_FLOW_CHUNK_SIZE = 10
+    DEFAULT_UPLOAD_SIMULTANEOUS_UPLOADS = 3
+    DEFAULT_UPLOAD_MAX_CHUNK_RETRIES = 1
+
+    DEFAULT_DOWNLOAD_FLOW_CHUNK_SIZE = 1_048_576
     FLOW_PERMANENT_ERRORS = [404, 413, 415, 500, 501]
 
     def upload(
         self,
         local_path: str,
         upload_path: str,
         overwrite: bool = False,
-        chunk_size=1048576,
-        simultaneous_uploads=3,
-        max_chunk_retries=1,
+        chunk_size=DEFAULT_UPLOAD_FLOW_CHUNK_SIZE,
+        simultaneous_uploads=DEFAULT_UPLOAD_SIMULTANEOUS_UPLOADS,
+        max_chunk_retries=DEFAULT_UPLOAD_MAX_CHUNK_RETRIES,
         chunk_retry_interval=1,
     ):
         """Upload a file to the Hopsworks filesystem.
@@ -66,7 +70,7 @@ def upload(
             local_path: local path to file to upload
             upload_path: path to directory where to upload the file in Hopsworks Filesystem
             overwrite: overwrite file if exists
-            chunk_size: upload chunk size in bytes. Default 1048576 bytes
+            chunk_size: upload chunk size in megabytes. Default 10 MB
             simultaneous_uploads: number of simultaneous chunks to upload. Default 3
             max_chunk_retries: maximum retry for a chunk. Default is 1
             chunk_retry_interval: chunk retry interval in seconds. Default is 1sec
@@ -86,6 +90,7 @@ def upload(
         _, file_name = os.path.split(local_path)
 
         destination_path = upload_path + "/" + file_name
+        chunk_size_bytes = chunk_size * 1024 * 1024
 
         if self.path_exists(destination_path):
             if overwrite:
@@ -97,10 +102,10 @@ def upload(
                     )
                 )
 
-        num_chunks = math.ceil(file_size / chunk_size)
+        num_chunks = math.ceil(file_size / chunk_size_bytes)
 
         base_params = self._get_flow_base_params(
-            file_name, num_chunks, file_size, chunk_size
+            file_name, num_chunks, file_size, chunk_size_bytes
         )
 
         chunk_number = 1
@@ -119,7 +124,7 @@ def upload(
                 while True:
                     chunks = []
                     for _ in range(simultaneous_uploads):
-                        chunk = f.read(chunk_size)
+                        chunk = f.read(chunk_size_bytes)
                         if not chunk:
                             break
                         chunks.append(Chunk(chunk, chunk_number, "pending"))
@@ -241,7 +246,7 @@ def download(self, path, local_path):
                 downloaded = 0
                 # if not response.headers.get("Content-Length"), file is still downloading
                 for chunk in response.iter_content(
-                    chunk_size=self.DEFAULT_FLOW_CHUNK_SIZE
+                    chunk_size=self.DEFAULT_DOWNLOAD_FLOW_CHUNK_SIZE
                 ):
                     f.write(chunk)
                     downloaded += len(chunk)

python/hsml/engine/hopsworks_engine.py

@@ -33,7 +33,7 @@ def mkdir(self, remote_path: str):
     def delete(self, model_instance):
         self._model_api.delete(model_instance)
 
-    def upload(self, local_path: str, remote_path: str):
+    def upload(self, local_path: str, remote_path: str, upload_configuration=None):
         local_path = self._get_abs_path(local_path)
         remote_path = self._prepend_project_path(remote_path)
         self._native_hdfs_api.upload(local_path, remote_path)

python/hsml/engine/local_engine.py

@@ -32,10 +32,27 @@ def mkdir(self, remote_path: str):
     def delete(self, model_instance):
         self._model_api.delete(model_instance)
 
-    def upload(self, local_path: str, remote_path: str):
+    def upload(self, local_path: str, remote_path: str, upload_configuration=None):
         local_path = self._get_abs_path(local_path)
         remote_path = self._prepend_project_path(remote_path)
-        self._dataset_api.upload(local_path, remote_path)
+
+        # Initialize the upload configuration to empty dictionary if is None
+        upload_configuration = upload_configuration if upload_configuration else {}
+        self._dataset_api.upload(
+            local_path,
+            remote_path,
+            chunk_size=upload_configuration.get(
+                "chunk_size", self._dataset_api.DEFAULT_UPLOAD_FLOW_CHUNK_SIZE
+            ),
+            simultaneous_uploads=upload_configuration.get(
+                "simultaneous_uploads",
+                self._dataset_api.DEFAULT_UPLOAD_SIMULTANEOUS_UPLOADS,
+            ),
+            max_chunk_retries=upload_configuration.get(
+                "max_chunk_retries",
+                self._dataset_api.DEFAULT_UPLOAD_MAX_CHUNK_RETRIES,
+            ),
+        )
 
     def download(self, remote_path: str, local_path: str):
         local_path = self._get_abs_path(local_path)

python/hsml/engine/model_engine.py

@@ -203,6 +203,7 @@ def _upload_local_model(
         from_local_model_path,
         to_model_version_path,
         update_upload_progress,
+        upload_configuration=None,
     ):
         """Copy or upload model files from a local path to the model version folder in the Models dataset."""
         n_dirs, n_files = 0, 0
@@ -222,17 +223,30 @@ def _upload_local_model(
                     n_dirs += 1
                     update_upload_progress(n_dirs, n_files)
                 for f_name in files:
-                    self._engine.upload(root + "/" + f_name, remote_base_path)
+                    self._engine.upload(
+                        root + "/" + f_name,
+                        remote_base_path,
+                        upload_configuration=upload_configuration,
+                    )
                     n_files += 1
                     update_upload_progress(n_dirs, n_files)
         else:
             # if path is a file, upload file
-            self._engine.upload(from_local_model_path, to_model_version_path)
+            self._engine.upload(
+                from_local_model_path,
+                to_model_version_path,
+                upload_configuration=upload_configuration,
+            )
             n_files += 1
             update_upload_progress(n_dirs, n_files)
 
     def _save_model_from_local_or_hopsfs_mount(
-        self, model_instance, model_path, keep_original_files, update_upload_progress
+        self,
+        model_instance,
+        model_path,
+        keep_original_files,
+        update_upload_progress,
+        upload_configuration=None,
     ):
         """Save model files from a local path. The local path can be on hopsfs mount"""
         # check hopsfs mount
@@ -250,6 +264,7 @@ def _save_model_from_local_or_hopsfs_mount(
                 from_local_model_path=model_path,
                 to_model_version_path=model_instance.version_path,
                 update_upload_progress=update_upload_progress,
+                upload_configuration=upload_configuration,
             )
 
     def _set_model_version(
@@ -297,6 +312,7 @@ def save(
         model_path,
         await_registration=480,
         keep_original_files=False,
+        upload_configuration=None,
     ):
         _client = client.get_instance()
 
@@ -377,6 +393,7 @@ def update_upload_progress(n_dirs=0, n_files=0):
                             model_path=model_path,
                             keep_original_files=keep_original_files,
                             update_upload_progress=update_upload_progress,
+                            upload_configuration=upload_configuration,
                         )
                     # check local relative
                     elif os.path.exists(
@@ -387,6 +404,7 @@ def update_upload_progress(n_dirs=0, n_files=0):
                             model_path=os.path.join(os.getcwd(), model_path),
                             keep_original_files=keep_original_files,
                             update_upload_progress=update_upload_progress,
+                            upload_configuration=upload_configuration,
                         )
                     # check project relative
                     elif self._dataset_api.path_exists(

python/hsml/model.py

@@ -16,7 +16,7 @@
 
 import json
 import humps
-from typing import Union, Optional
+from typing import Union, Optional, Dict, Any
 
 
 from hsml import client, util
@@ -93,13 +93,25 @@ def __init__(
         self._feature_view = feature_view
         self._training_dataset_version = training_dataset_version
 
-    def save(self, model_path, await_registration=480, keep_original_files=False):
+    def save(
+        self,
+        model_path,
+        await_registration=480,
+        keep_original_files=False,
+        upload_configuration: Optional[Dict[str, Any]] = None,
+    ):
         """Persist this model including model files and metadata to the model registry.
 
         # Arguments
             model_path: Local or remote (Hopsworks file system) path to the folder where the model files are located, or path to a specific model file.
             await_registration: Awaiting time for the model to be registered in Hopsworks.
             keep_original_files: If the model files are located in hopsfs, whether to move or copy those files into the Models dataset. Default is False (i.e., model files will be moved)
+            upload_configuration: When saving a model from outside Hopsworks, the model is uploaded to the model registry using the REST APIs. Each model artifact is divided into
+                chunks and each chunk uploaded independently. This parameter can be used to control the upload chunk size, the parallelism and the number of retries.
+                `upload_configuration` can contain the following keys:
+                * key `chunk_size`: size of each chunk in megabytes. Default 10.
+                * key `simultaneous_uploads`: number of chunks to upload in parallel. Default 3.
+                * key `max_chunk_retries`: number of times to retry the upload of a chunk in case of failure. Default 1.
 
         # Returns
             `Model`: The model metadata object.
@@ -109,6 +121,7 @@ def save(self, model_path, await_registration=480, keep_original_files=False):
             model_path,
             await_registration=await_registration,
             keep_original_files=keep_original_files,
+            upload_configuration=upload_configuration,
         )
 
     def download(self):

python/hsml/resources.py

@@ -324,9 +324,9 @@ def to_dict(self):
         return {
             humps.camelize(self.NUM_INSTANCES_KEY): self._num_instances,
             humps.camelize(self.RESOURCES_CONFIG_KEY): {
-                "requests": self._requests.to_dict()
-                if self._requests is not None
-                else None,
+                "requests": (
+                    self._requests.to_dict() if self._requests is not None else None
+                ),
                 "limits": self._limits.to_dict() if self._limits is not None else None,
             },
         }
@@ -352,9 +352,9 @@ def to_dict(self):
         return {
             humps.camelize(self.NUM_INSTANCES_KEY): self._num_instances,
             humps.camelize(self.RESOURCES_CONFIG_KEY): {
-                "requests": self._requests.to_dict()
-                if self._requests is not None
-                else None,
+                "requests": (
+                    self._requests.to_dict() if self._requests is not None else None
+                ),
                 "limits": self._limits.to_dict() if self._limits is not None else None,
             },
         }