added total chunks to model upload

opensearch-project · Aug 24, 2022 · 1b79f62 · 1b79f62
1 parent a3bf832
commit 1b79f62
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/opensearch_py_ml/ml_commons_integration/upload/ml_common_upload_client.py b/opensearch_py_ml/ml_commons_integration/upload/ml_common_upload_client.py
@@ -5,12 +5,11 @@
 
 from opensearchpy import OpenSearch
 import os
+from math import ceil
 from typing import Iterable
 from opensearch_py_ml.ml_commons_integration.ml_common_utils import ML_BASE_URI
 
-from tqdm.auto import tqdm # type: ignore
-
-DEFAULT_ML_COMMON_UPLOAD_CHUNK_SIZE = 10_000_000 # 10MB
+DEFAULT_ML_COMMON_UPLOAD_CHUNK_SIZE = 10_000_000  # 10MB
 
 
 class MLCommonUploadClient:
@@ -26,23 +25,23 @@ def put_model(self,
                   version_number: int,
                   chunk_size: int = DEFAULT_ML_COMMON_UPLOAD_CHUNK_SIZE,
                   verbose: bool = False) -> None:
-        total_model_size = os.stat(model_path).st_size
+        total_num_chunks = ceil(os.stat(model_path).st_size/chunk_size)
 
         def model_file_chunk_generator() -> Iterable[str]:
             with open(model_path, "rb") as f:
                 while True:
                     data = f.read(chunk_size)
                     if not data:
                         break
-                    yield data # check if we actually need to do base64 encoding
+                    yield data  # check if we actually need to do base64 encoding
 
         to_iterate_over = enumerate(model_file_chunk_generator())
-        if verbose:
-            to_iterate_over = tqdm(to_iterate_over)
 
         for i, chunk in to_iterate_over:
+            if verbose:
+                print(f"uploading chunk {i+1} of {total_num_chunks}")
             self._client.transport.perform_request(
                 method="POST",
-                url=f"{ML_BASE_URI}/custom_model/upload_chunk/{model_name}/{version_number}/{i}",
+                url=f"{ML_BASE_URI}/custom_model/upload_chunk/{model_name}/{version_number}/{i}/{total_num_chunks}",
                 body=chunk
-            )
+            )