add sagemaker predict PoC to eland

opensearch_py_ml/ml_commons_integration/ml_common_utils.py

@@ -3,4 +3,4 @@
 SPDX-License-Identifier: Apache-2.0
  """
 
-ML_BASE_URI = '_plugins/_ml'
+ML_BASE_URI = '/_plugins/_ml'

opensearch_py_ml/ml_commons_integration/upload/ml_common_upload_client.py

@@ -5,14 +5,11 @@
 
 from opensearchpy import OpenSearch
 import os
-import math
+from math import ceil
 from typing import Iterable
-import base64
 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:
@@ -28,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/{model_name}/{version_number}/{i}",
+                url=f"{ML_BASE_URI}/custom_model/upload_chunk/{model_name}/{version_number}/{i}/{total_num_chunks}",
                 body=chunk
-            )
+            )

opensearch_py_ml/sagemaker_tools.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 from opensearch_py_ml import DataFrame
-from typing import List, Optional
+from typing import List, Optional, Dict, Tuple, Any
 from math import ceil
 
 from sagemaker import RealTimePredictor, Session
@@ -21,7 +21,7 @@ def make_sagemaker_prediction(endpoint_name: str,
                               column_order: Optional[List[str]] = None,
                               chunksize: int = None,
                               sort_index: Optional[str] = '_doc'
-                              )-> np.array:
+                              ) -> Tuple[List[Any], Dict[Any, Any]]:
     """
     Make a prediction on an opensearch_py_ml dataframe using a deployed SageMaker model endpoint.
 
@@ -44,7 +44,7 @@ def make_sagemaker_prediction(endpoint_name: str,
 
     Returns
     ----------
-    np.array representing the output of the model on input data
+    list representing the indices, dictionary representing the output of the model on input data
     """
     predictor = RealTimePredictor(endpoint=endpoint_name, sagemaker_session=sagemaker_session, content_type='text/csv')
     data = data.drop(columns=target_column)