[Refactor] Merge the sparse and dense model's ci branch.

Signed-off-by: conggguan <[email protected]>
opensearch-project · Jul 30, 2024 · 493f301 · 493f301
1 parent 6578981
commit 493f301
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 29 deletions.
diff --git a/.ci/run-repository.sh b/.ci/run-repository.sh
@@ -65,7 +65,7 @@ elif [[ "$TASK_TYPE" == "doc" ]]; then
 
   docker cp opensearch-py-ml-doc-runner:/code/opensearch-py-ml/docs/build/ ./docs/
   docker rm opensearch-py-ml-doc-runner
-elif [[ "$TASK_TYPE" == "SentenceTransformerTrace" ]]; then
+elif [[ "$TASK_TYPE" == "SentenceTransformerTrace" || "$TASK_TYPE" == "SparseTrace" ]]; then
   # Set up OpenSearch cluster & Run model autotracing (Invoked by model_uploader.yml workflow)
   echo -e "\033[34;1mINFO:\033[0m MODEL_ID: ${MODEL_ID}\033[0m"
   echo -e "\033[34;1mINFO:\033[0m MODEL_VERSION: ${MODEL_VERSION}\033[0m"
@@ -74,30 +74,14 @@ elif [[ "$TASK_TYPE" == "SentenceTransformerTrace" ]]; then
   echo -e "\033[34;1mINFO:\033[0m POOLING_MODE: ${POOLING_MODE:-N/A}\033[0m"
   echo -e "\033[34;1mINFO:\033[0m MODEL_DESCRIPTION: ${MODEL_DESCRIPTION:-N/A}\033[0m"
 
-  docker run \
-  --network=${network_name} \
-  --env "STACK_VERSION=${STACK_VERSION}" \
-  --env "OPENSEARCH_URL=${opensearch_url}" \
-  --env "OPENSEARCH_VERSION=${OPENSEARCH_VERSION}" \
-  --env "TEST_SUITE=${TEST_SUITE}" \
-  --env "PYTHON_CONNECTION_CLASS=${PYTHON_CONNECTION_CLASS}" \
-  --env "TEST_TYPE=server" \
-  --name opensearch-py-ml-trace-runner \
-  opensearch-project/opensearch-py-ml \
-  nox -s "trace-${PYTHON_VERSION}" -- ${MODEL_ID} ${MODEL_VERSION} ${TRACING_FORMAT} -ed ${EMBEDDING_DIMENSION} -pm ${POOLING_MODE} -md ${MODEL_DESCRIPTION:+"$MODEL_DESCRIPTION"}
-
-  # To upload a model, we need the model artifact, description, license files into local path
-  # trace_output should include description and license file.
-  docker cp opensearch-py-ml-trace-runner:/code/opensearch-py-ml/upload/ ./upload/
-  docker cp opensearch-py-ml-trace-runner:/code/opensearch-py-ml/trace_output/ ./trace_output/
-  docker rm opensearch-py-ml-trace-runner
-
-elif [[ "$TASK_TYPE" == "SparseTrace" ]]; then
-  # Set up OpenSearch cluster & Run model autotracing (Invoked by model_uploader.yml workflow)
-  echo -e "\033[34;1mINFO:\033[0m MODEL_ID: ${MODEL_ID}\033[0m"
-  echo -e "\033[34;1mINFO:\033[0m MODEL_VERSION: ${MODEL_VERSION}\033[0m"
-  echo -e "\033[34;1mINFO:\033[0m TRACING_FORMAT: ${TRACING_FORMAT}\033[0m"
-  echo -e "\033[34;1mINFO:\033[0m MODEL_DESCRIPTION: ${MODEL_DESCRIPTION:-N/A}\033[0m"
+  if [[ "$TASK_TYPE" == "SentenceTransformerTrace" ]]; then
+      NOX_TRACE_TYPE="trace"
+  elif [[ "$TASK_TYPE" == "SparseTrace" ]]; then
+      NOX_TRACE_TYPE="sparsetrace"
+  else
+      echo "Unknown TASK_TYPE: $TASK_TYPE"
+      exit 1
+  fi
 
   docker run \
   --network=${network_name} \
@@ -109,13 +93,11 @@ elif [[ "$TASK_TYPE" == "SparseTrace" ]]; then
   --env "TEST_TYPE=server" \
   --name opensearch-py-ml-trace-runner \
   opensearch-project/opensearch-py-ml \
-  nox -s "sparsetrace-${PYTHON_VERSION}" -- ${MODEL_ID} ${MODEL_VERSION} ${TRACING_FORMAT} -md ${MODEL_DESCRIPTION:+"$MODEL_DESCRIPTION"}
+  nox -s "${NOX_TRACE_TYPE}-${PYTHON_VERSION}" -- ${MODEL_ID} ${MODEL_VERSION} ${TRACING_FORMAT} -ed ${EMBEDDING_DIMENSION} -pm ${POOLING_MODE} -md ${MODEL_DESCRIPTION:+"$MODEL_DESCRIPTION"}
 
   # To upload a model, we need the model artifact, description, license files into local path
   # trace_output should include description and license file.
   docker cp opensearch-py-ml-trace-runner:/code/opensearch-py-ml/upload/ ./upload/
   docker cp opensearch-py-ml-trace-runner:/code/opensearch-py-ml/trace_output/ ./trace_output/
-
-  # Delete the docker image
   docker rm opensearch-py-ml-trace-runner
 fi
diff --git a/opensearch_py_ml/ml_commons/ml_commons_client.py b/opensearch_py_ml/ml_commons/ml_commons_client.py
@@ -509,7 +509,7 @@ def generate_sparse_encoding(self, model_id: str, sentences: List[str]) -> objec
         :return: Returns a JSON object `inference_results` containing the sparse encoding results for the given sentences.
         :rtype: object
         """
-        API_URL = f"{ML_BASE_URI}/_predict/sparse_encoding/{model_id}"
+        API_URL = f"{ML_BASE_URI}/models/{model_id}/_predict/"
         API_BODY = {"text_docs": sentences}
         return self._client.transport.perform_request(
             method="POST",