diff --git a/.ci/run-repository.sh b/.ci/run-repository.sh index ba0cea2b..2dee9683 100755 --- a/.ci/run-repository.sh +++ b/.ci/run-repository.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash -# Called by entry point `run-test` use this script to add your repository specific test commands +# Called by entry point `run-test` use this script to add your repository specific task commands # Once called opensearch is up and running and the following parameters are available to this script # OPENSEARCH_VERSION -- version e.g Major.Minor.Patch(-Prelease) @@ -16,7 +16,7 @@ set -e echo -e "\033[34;1mINFO:\033[0m URL ${opensearch_url}\033[0m" echo -e "\033[34;1mINFO:\033[0m EXTERNAL OS URL ${external_opensearch_url}\033[0m" echo -e "\033[34;1mINFO:\033[0m VERSION ${OPENSEARCH_VERSION}\033[0m" -echo -e "\033[34;1mINFO:\033[0m IS_DOC: ${IS_DOC}\033[0m" +echo -e "\033[34;1mINFO:\033[0m TASK_TYPE: ${TASK_TYPE}\033[0m" echo -e "\033[34;1mINFO:\033[0m TEST_SUITE ${TEST_SUITE}\033[0m" echo -e "\033[34;1mINFO:\033[0m PYTHON_VERSION ${PYTHON_VERSION}\033[0m" echo -e "\033[34;1mINFO:\033[0m PYTHON_CONNECTION_CLASS ${PYTHON_CONNECTION_CLASS}\033[0m" @@ -33,7 +33,8 @@ docker build \ echo -e "\033[1m>>>>> Run [opensearch-project/opensearch-py-ml container] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>\033[0m" -if [[ "$IS_DOC" == "false" ]]; then +if [[ "$TASK_TYPE" == "test" ]]; then + # Set up OpenSearch cluster & Run test (Invoked by integration.yml workflow) docker run \ --network=${network_name} \ --env "STACK_VERSION=${STACK_VERSION}" \ @@ -45,10 +46,11 @@ if [[ "$IS_DOC" == "false" ]]; then --name opensearch-py-ml-test-runner \ opensearch-project/opensearch-py-ml \ nox -s "test-${PYTHON_VERSION}(pandas_version='${PANDAS_VERSION}')" + docker cp opensearch-py-ml-test-runner:/code/opensearch-py-ml/junit/ ./junit/ - docker rm opensearch-py-ml-test-runner -else +elif [[ "$TASK_TYPE" == "doc" ]]; then + # Set up OpenSearch cluster & Run docs (Invoked by build_deploy_doc.yml workflow) docker run \ --network=${network_name} \ --env "STACK_VERSION=${STACK_VERSION}" \ @@ -60,7 +62,31 @@ else --name opensearch-py-ml-doc-runner \ opensearch-project/opensearch-py-ml \ nox -s docs + docker cp opensearch-py-ml-doc-runner:/code/opensearch-py-ml/docs/build/ ./docs/ - docker rm opensearch-py-ml-doc-runner -fi \ No newline at end of file +elif [[ "$TASK_TYPE" == "trace" ]]; 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 EMBEDDING_DIMENSION: ${EMBEDDING_DIMENSION:-N/A}\033[0m" + echo -e "\033[34;1mINFO:\033[0m POOLING_MODE: ${POOLING_MODE:-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 \ + /bin/bash -c "python -m pip install -r requirements-dev.txt --timeout 1500; + python -m pip install pandas~=${PANDAS_VERSION}; + python utils/model_uploader/model_autotracing.py ${MODEL_ID} ${MODEL_VERSION} ${TRACING_FORMAT} -ed ${EMBEDDING_DIMENSION} -pm ${POOLING_MODE}" + + docker cp opensearch-py-ml-trace-runner:/code/opensearch-py-ml/upload/ ./upload/ + docker rm opensearch-py-ml-trace-runner +fi diff --git a/.ci/run-tests b/.ci/run-tests index abfeac34..258da8a5 100755 --- a/.ci/run-tests +++ b/.ci/run-tests @@ -10,7 +10,7 @@ export PYTHON_CONNECTION_CLASS="${PYTHON_CONNECTION_CLASS:=Urllib3HttpConnection export CLUSTER="${1:-opensearch}" export SECURE_INTEGRATION="${2:-true}" export OPENSEARCH_VERSION="${3:-latest}" -export IS_DOC="${4:-false}" +export TASK_TYPE="${4:-test}" if [[ "$SECURE_INTEGRATION" == "true" ]]; then export OPENSEARCH_URL_EXTENSION="https" else diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e04a41ea..d043d2cb 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1 +1 @@ -* @dhrubo-os @greaa-aws @ylwu-amzn @b4sjoo @jngz-es @rbhavna \ No newline at end of file +* @thanawan-atc diff --git a/.github/workflows/build_deploy_doc.yml b/.github/workflows/build_deploy_doc.yml index 9321b32e..876a73ea 100644 --- a/.github/workflows/build_deploy_doc.yml +++ b/.github/workflows/build_deploy_doc.yml @@ -20,7 +20,7 @@ jobs: - name: Checkout Repository uses: actions/checkout@v2 - name: Integ ${{ matrix.cluster }} secured=${{ matrix.secured }} version=${{matrix.entry.opensearch_version}} - run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }} true" + run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }} doc" - name: Deploy uses: peaceiris/actions-gh-pages@v3 with: diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 607bd356..e36c7735 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -18,7 +18,7 @@ jobs: - name: Checkout uses: actions/checkout@v2 - name: Integ ${{ matrix.cluster }} secured=${{ matrix.secured }} version=${{matrix.entry.opensearch_version}} - run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }}" + run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }} test" - name: Upload coverage to Codecov uses: codecov/codecov-action@v2 with: diff --git a/.github/workflows/model_uploader.yml b/.github/workflows/model_uploader.yml new file mode 100644 index 00000000..e3c38c0f --- /dev/null +++ b/.github/workflows/model_uploader.yml @@ -0,0 +1,239 @@ +name: Model Auto-tracing & Uploading +on: + # Step 1: Initiate the workflow + workflow_dispatch: + inputs: + model_id: + description: "Model ID for auto-tracing and uploading (e.g. sentence-transformers/msmarco-distilbert-base-tas-b)" + required: true + type: string + model_version: + description: "Model version number (e.g. 1.0.1)" + required: true + type: string + tracing_format: + description: "Model format for auto-tracing (torch_script/onnx)" + required: true + type: choice + options: + - "BOTH" + - "TORCH_SCRIPT" + - "ONNX" + embedding_dimension: + description: "(Optional) Embedding Dimension (Specify here if it does not exist in original config.json file, or you want to overwrite it.)" + required: false + type: int + pooling_mode: + description: "(Optional) Pooling Mode (Specify here if it does not exist in original config.json file or you want to overwrite it.)" + required: false + type: choice + options: + - "" + - "CLS" + - "MEAN" + - "MAX" + - "MEAN_SQRT_LEN" + +jobs: + # Step 2: Check if the model already exists in the model hub + checking-out-model-hub: + runs-on: 'ubuntu-latest' + permissions: + id-token: write + contents: read + steps: + - name: Checkout Repository + uses: actions/checkout@v3 + - name: Set Up Python + uses: actions/setup-python@v2 + with: + python-version: '3.x' + - name: Configure AWS Credentials + uses: aws-actions/configure-aws-credentials@v2 + with: + aws-region: ${{ secrets.MODEL_UPLOADER_AWS_REGION }} + role-to-assume: ${{ secrets.MODEL_UPLOADER_ROLE }} + role-session-name: checking-out-model-hub + - name: Check if TORCH_SCRIPT Model Exists + if: github.event.inputs.tracing_format == 'TORCH_SCRIPT' || github.event.inputs.tracing_format == 'BOTH' + run: | + TORCH_FILE_PATH=$(python utils/model_uploader/save_model_file_path_to_env.py \ + ${{ github.event.inputs.model_id }} ${{ github.event.inputs.model_version }} TORCH_SCRIPT) + aws s3api head-object --bucket opensearch-exp --key $TORCH_FILE_PATH > /dev/null 2>&1 || TORCH_MODEL_NOT_EXIST=true + if [[ -z $TORCH_MODEL_NOT_EXIST ]]; + then + echo "TORCH_SCRIPT Model already exists on model hub." + exit 1 + fi + - name: Check if ONNX Model Exists + if: github.event.inputs.tracing_format == 'ONNX' || github.event.inputs.tracing_format == 'BOTH' + run: | + ONNX_FILE_PATH=$(python utils/model_uploader/save_model_file_path_to_env.py \ + ${{ github.event.inputs.model_id }} ${{ github.event.inputs.model_version }} ONNX) + aws s3api head-object --bucket opensearch-exp --key $ONNX_FILE_PATH > /dev/null 2>&1 || ONNX_MODEL_NOT_EXIST=true + if [[ -z $ONNX_MODEL_NOT_EXIST ]]; + then + echo "TORCH_SCRIPT Model already exists on model hub." + exit 1; + fi + + # Step 3: Trace the model, Verify the embeddings & Upload the model files as artifacts + model-auto-tracing: + needs: checking-out-model-hub + name: model-auto-tracing + runs-on: ubuntu-latest + permissions: + id-token: write + contents: read + strategy: + matrix: + cluster: ["opensearch"] + secured: ["true"] + entry: + - { opensearch_version: 2.7.0 } + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Export Arguments + run: | + echo "MODEL_ID=${{ github.event.inputs.model_id }}" >> $GITHUB_ENV + echo "MODEL_VERSION=${{ github.event.inputs.model_version }}" >> $GITHUB_ENV + echo "TRACING_FORMAT=${{ github.event.inputs.tracing_format }}" >> $GITHUB_ENV + echo "EMBEDDING_DIMENSION=${{ github.event.inputs.embedding_dimension }}" >> $GITHUB_ENV + echo "POOLING_MODE=${{ github.event.inputs.pooling_mode }}" >> $GITHUB_ENV + - name: Autotracing ${{ matrix.cluster }} secured=${{ matrix.secured }} version=${{matrix.entry.opensearch_version}} + run: "./.ci/run-tests ${{ matrix.cluster }} ${{ matrix.secured }} ${{ matrix.entry.opensearch_version }} trace" + - name: Upload Artifact + uses: actions/upload-artifact@v3 + with: + name: upload + path: ./upload/ + retention-days: 5 + if-no-files-found: error + - name: Configure AWS Credentials + uses: aws-actions/configure-aws-credentials@v2 + with: + aws-region: ${{ secrets.MODEL_UPLOADER_AWS_REGION }} + role-to-assume: ${{ secrets.MODEL_UPLOADER_ROLE }} + role-session-name: model-auto-tracing + - name: Dryrun model uploading + id: dryrun_model_uploading + run: | + aws s3 sync ./upload/ s3://opensearch-exp/ml-models/huggingface/sentence-transformers/ --dryrun + dryrun_output=$(aws s3 sync ./upload/ s3://opensearch-exp/ml-models/huggingface/sentence-transformers/ --dryrun) + echo "dryrun_output<> $GITHUB_OUTPUT + echo "${dryrun_output@E}" >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + echo "${dryrun_output@E}" + outputs: + dryrun_output: ${{ steps.dryrun_model_uploading.outputs.dryrun_output }} + + # Step 4: Ask for manual approval from the CODEOWNERS + manual-approval: + needs: model-auto-tracing + runs-on: 'ubuntu-latest' + permissions: + issues: write + steps: + - name: Checkout Repository + uses: actions/checkout@v3 + - name: Get Approvers + id: get_approvers + run: | + echo "approvers=$(cat .github/CODEOWNERS | grep @ | tr -d '* ' | sed 's/@/,/g' | sed 's/,//1')" >> $GITHUB_OUTPUT + - name: Create Issue Body + id: create_issue_body + run: | + embedding_dimension=${{ github.event.inputs.embedding_dimension }} + pooling_mode=${{ github.event.inputs.pooling_mode }} + issue_body="Please approve or deny opensearch-py-ml model uploading: + + ========= Workflow Details ========== + - Workflow Name: ${{ github.workflow }} + - Workflow Initiator: @${{ github.actor }} + + ========= Model Information ========= + - Model ID: ${{ github.event.inputs.model_id }} + - Model Version: ${{ github.event.inputs.model_version }} + - Tracing Format: ${{ github.event.inputs.tracing_format }} + - Embedding Dimension: ${embedding_dimension:-Default} + - Pooling Mode: ${pooling_mode:-Default} + + ===== Dry Run of Model Uploading ===== + ${{ needs.model-auto-tracing.outputs.dryrun_output }}" + + echo "issue_body<> $GITHUB_OUTPUT + echo "${issue_body@E}" >> $GITHUB_OUTPUT + echo "EOF" >> $GITHUB_OUTPUT + echo "${issue_body@E}" + - uses: trstringer/manual-approval@v1 + with: + secret: ${{ github.TOKEN }} + approvers: ${{ steps.get_approvers.outputs.approvers }} + minimum-approvals: 1 + issue-title: "Upload Model to OpenSearch Model Hub (${{ github.event.inputs.model_id }})" + issue-body: ${{ steps.create_issue_body.outputs.issue_body }} + exclude-workflow-initiator-as-approver: false + + # Step 5: Download the artifacts & Upload it to the S3 bucket + model-uploading: + needs: manual-approval + runs-on: 'ubuntu-latest' + permissions: + id-token: write + contents: read + steps: + - name: Download Artifact + uses: actions/download-artifact@v2 + with: + name: upload + path: ./upload/ + - name: Configure AWS Credentials + uses: aws-actions/configure-aws-credentials@v2 + with: + aws-region: ${{ secrets.MODEL_UPLOADER_AWS_REGION }} + role-to-assume: ${{ secrets.MODEL_UPLOADER_ROLE }} + role-session-name: model-uploading + - name: Copy Files to the Bucket + id: copying_to_bucket + run: | + aws s3 sync ./upload/ s3://opensearch-exp/ml-models/huggingface/sentence-transformers/ + echo "upload_time=$(TZ='America/Los_Angeles' date "+%Y-%m-%d %T")" >> $GITHUB_OUTPUT + outputs: + upload_time: ${{ steps.copying_to_bucket.outputs.upload_time }} + + # Step 6: Update MODEL_UPLOAD_HISTORY.md & supported_models.json + history-update: + needs: model-uploading + runs-on: 'ubuntu-latest' + permissions: + id-token: write + contents: write + concurrency: ${{ github.workflow }}-concurrency + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Set Up Python + uses: actions/setup-python@v2 + with: + python-version: '3.x' + - name: Install Packages + run: + python -m pip install mdutils + - name: Update MODEL_UPLOAD_HISTORY.md + run: | + python utils/model_uploader/update_models_upload_history_md.py \ + ${{ github.event.inputs.model_id }} \ + ${{ github.event.inputs.model_version }} \ + ${{ github.event.inputs.tracing_format }} \ + -ed ${{ github.event.inputs.embedding_dimension }} \ + -pm ${{ github.event.inputs.pooling_mode }} \ + -u ${{ github.actor }} -t "${{ needs.model-uploading.outputs.upload_time }}" + - name: Commit Updates + uses: stefanzweifel/git-auto-commit-action@v4 + id: commit + with: + commit_message: 'GitHub Actions Workflow - Update MODEL_UPLOAD_HISTORY.md (${{ github.event.inputs.model_id }})' + commit_options: '--signoff' + repository: ./utils/model_uploader/upload_history + file_pattern: MODEL_UPLOAD_HISTORY.md supported_models.json diff --git a/noxfile.py b/noxfile.py index 448c3990..c9a835d0 100644 --- a/noxfile.py +++ b/noxfile.py @@ -61,7 +61,7 @@ @nox.session(reuse_venv=True) def format(session): session.install("black", "isort", "flynt") - session.run("python", "utils/license-headers.py", "fix", *SOURCE_FILES) + session.run("python", "utils/lint/license-headers.py", "fix", *SOURCE_FILES) session.run("flynt", *SOURCE_FILES) session.run("black", "--target-version=py38", *SOURCE_FILES) session.run("isort", "--profile=black", *SOURCE_FILES) @@ -73,7 +73,7 @@ def lint(session): # Install numpy to use its mypy plugin # https://numpy.org/devdocs/reference/typing.html#mypy-plugin session.install("black", "flake8", "mypy", "isort", "numpy") - session.run("python", "utils/license-headers.py", "check", *SOURCE_FILES) + session.run("python", "utils/lint/license-headers.py", "check", *SOURCE_FILES) session.run("black", "--check", "--target-version=py38", *SOURCE_FILES) session.run("isort", "--check", "--profile=black", *SOURCE_FILES) session.run("flake8", "--ignore=E501,W503,E402,E712,E203", *SOURCE_FILES) diff --git a/utils/license-headers.py b/utils/lint/license-headers.py similarity index 100% rename from utils/license-headers.py rename to utils/lint/license-headers.py diff --git a/utils/model_uploader/model_autotracing.py b/utils/model_uploader/model_autotracing.py new file mode 100644 index 00000000..78b35704 --- /dev/null +++ b/utils/model_uploader/model_autotracing.py @@ -0,0 +1,454 @@ +# SPDX-License-Identifier: Apache-2.0 +# The OpenSearch Contributors require contributions made to +# this file be licensed under the Apache-2.0 license or a +# compatible open source license. +# Any modifications Copyright OpenSearch Contributors. See +# GitHub history for details. + +# This program is run by "Model Auto-tracing & Uploading" workflow +# (See model_uploader.yml) to perform model auto-tracing and prepare +# files for uploading to OpenSearch model hub. + +import argparse +import os +import shutil +import sys +import warnings +from typing import List, Optional, Tuple +from zipfile import ZipFile + +import numpy as np +from numpy.typing import DTypeLike +from sentence_transformers import SentenceTransformer + +ROOT_DIR = os.path.abspath(os.path.join("../opensearch-py-ml")) +TEST_DIR = os.path.join(ROOT_DIR, "tests") +LICENSE_PATH = os.path.join(ROOT_DIR, "LICENSE") +sys.path.append(ROOT_DIR) +sys.path.append(TEST_DIR) + +from opensearch_py_ml.ml_commons import MLCommonClient +from opensearch_py_ml.ml_models.sentencetransformermodel import SentenceTransformerModel +from tests import OPENSEARCH_TEST_CLIENT + +BOTH_FORMAT = "BOTH" +TORCH_SCRIPT_FORMAT = "TORCH_SCRIPT" +ONNX_FORMAT = "ONNX" + +ORIGINAL_FOLDER_PATH = "sentence-transformers-original/" +TORCHSCRIPT_FOLDER_PATH = "sentence-transformers-torchscript/" +ONNX_FOLDER_PATH = "sentence-transformers-onnx/" +UPLOAD_FOLDER_PATH = "upload/" +MODEL_CONFIG_FILE_NAME = "ml-commons_model_config.json" +TEST_SENTENCES = [ + "First test sentence", + "This is a very long sentence used for testing model embedding outputs.", +] +RTOL_TEST = 1e-03 +ATOL_TEST = 1e-05 +ML_BASE_URI = "/_plugins/_ml" + + +def trace_sentence_transformer_model( + model_id: str, + model_version: str, + model_format: str, + embedding_dimension: Optional[int] = None, + pooling_mode: Optional[str] = None, +) -> Tuple[str, str]: + """ + Trace the pretrained sentence transformer model, create a model config file, + and return a path to the model file and a path to the model config file required for model registration + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + :param model_format: Model format ("TORCH_SCRIPT" or "ONNX") + :type model_format: string + :param embedding_dimension: Embedding dimension input + :type embedding_dimension: int + :param pooling_mode: Pooling mode input ("CLS", "MEAN", "MAX", "MEAN_SQRT_LEN" or None) + :type pooling_mode: string + """ + folder_path = ( + TORCHSCRIPT_FOLDER_PATH + if model_format == TORCH_SCRIPT_FORMAT + else ONNX_FOLDER_PATH + ) + + # 1.) Initiate a sentence transformer model class object + pre_trained_model = None + try: + pre_trained_model = SentenceTransformerModel( + model_id=model_id, folder_path=folder_path, overwrite=True + ) + except Exception as e: + assert ( + False + ), f"Raised Exception in tracing {model_format} model\ + during initiating a sentence transformer model class object: {e}" + + # 2.) Save the model in the specified format + model_path = None + try: + if model_format == TORCH_SCRIPT_FORMAT: + model_path = pre_trained_model.save_as_pt( + model_id=model_id, sentences=TEST_SENTENCES + ) + else: + model_path = pre_trained_model.save_as_onnx(model_id=model_id) + except Exception as e: + assert False, f"Raised Exception during saving model as {model_format}: {e}" + + # 3.) Create a model config json file + try: + pre_trained_model.make_model_config_json( + version_number=model_version, + model_format=model_format, + embedding_dimension=embedding_dimension, + pooling_mode=pooling_mode, + ) + except Exception as e: + assert ( + False + ), f"Raised Exception during making model config file for {model_format} model: {e}" + + # 4.) Return model_path & model_config_path for model registration + model_config_path = folder_path + MODEL_CONFIG_FILE_NAME + return model_path, model_config_path + + +def register_and_deploy_sentence_transformer_model( + ml_client: "MLCommonClient", + model_path: str, + model_config_path: str, + model_format: str, +) -> List["DTypeLike"]: + """ + Register the pretrained sentence transformer model by using the model file and the model config file, + deploy the model to generate embeddings for the TEST_SENTENCES, + and return the embeddings for model verification + + :param ml_client: A client that communicates to the ml-common plugin for OpenSearch + :type ml_client: MLCommonClient + :param model_path: Path to model file + :type model_path: string + :param model_config_path: Path to model config file + :type model_config_path: string + :param model_format: Model format ("TORCH_SCRIPT" or "ONNX") + :type model_format: string + """ + embedding_data = None + + # 1.) Register & Deploy the model + model_id = "" + task_id = "" + try: + model_id = ml_client.register_model( + model_path=model_path, + model_config_path=model_config_path, + deploy_model=True, + isVerbose=True, + ) + print() + print(f"{model_format}_model_id:", model_id) + assert model_id != "" or model_id is not None + except Exception as e: + assert False, f"Raised Exception in {model_format} model registration/deployment: {e}" + + # 2.) Check model status + try: + ml_model_status = ml_client.get_model_info(model_id) + print() + print("Model Status:") + print(ml_model_status) + assert ml_model_status.get("model_format") == model_format + assert ml_model_status.get("algorithm") == "TEXT_EMBEDDING" + except Exception as e: + assert False, f"Raised Exception in getting {model_format} model info: {e}" + + # 3.) Generate embeddings + try: + embedding_output = ml_client.generate_embedding(model_id, TEST_SENTENCES) + assert len(embedding_output.get("inference_results")) == 2 + embedding_data = [ + embedding_output["inference_results"][i]["output"][0]["data"] + for i in range(len(TEST_SENTENCES)) + ] + except Exception as e: + assert ( + False + ), f"Raised Exception in generating sentence embedding with {model_format} model: {e}" + + # 4.) Undeploy the model + try: + ml_client.undeploy_model(model_id) + ml_model_status = ml_client.get_model_info(model_id) + assert ml_model_status.get("model_state") != "UNDEPLOY_FAILED" + except Exception as e: + assert False, f"Raised Exception in {model_format} model undeployment: {e}" + + # 5.) Delete the model + try: + delete_model_obj = ml_client.delete_model(model_id) + assert delete_model_obj.get("result") == "deleted" + except Exception as e: + assert False, f"Raised Exception in deleting {model_format} model: {e}" + + # 6.) Return embedding outputs for model verification + return embedding_data + + +def verify_embedding_data( + original_embedding_data: List["DTypeLike"], + tracing_embedding_data: List["DTypeLike"], + model_format: str, +) -> None: + """ + Verify the embeddings generated by the traced model with that of original model + + :param original_embedding_data: Embedding outputs of TEST_SENTENCES generated by the original model + :type original_embedding_data: List['DTypeLike'] + :param tracing_embedding_data: Embedding outputs of TEST_SENTENCES generated by the traced model + :type tracing_embedding_data: List['DTypeLike'] + :param model_format: Model format ("TORCH_SCRIPT" or "ONNX") + :type model_format: string + """ + failed_cases = [] + for i in range(len(TEST_SENTENCES)): + try: + np.testing.assert_allclose( + original_embedding_data[i], + tracing_embedding_data[i], + rtol=RTOL_TEST, + atol=ATOL_TEST, + ) + except Exception as e: + failed_cases.append((TEST_SENTENCES[i], e)) + + if len(failed_cases): + print() + print( + f"Original embeddings DOES NOT matches {model_format} embeddings in the following case(s):" + ) + for sentence, e in failed_cases: + print(sentence) + print(e) + assert False, f"Failed while verifying embeddings of {model_format} model" + else: + print() + print(f"Original embeddings matches {model_format} embeddings") + + +def prepare_files_for_uploading( + model_id: str, + model_version: str, + model_format: str, + src_model_path: str, + src_model_config_path: str, +) -> None: + """ + Prepare files for uploading by storing them in UPLOAD_FOLDER_PATH + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + :param model_format: Model format ("TORCH_SCRIPT" or "ONNX") + :type model_format: string + :param src_model_path: Path to model files for uploading + :type src_model_path: string + :param src_model_config_path: Path to model config files for uploading + :type src_model_config_path: string + """ + model_name = str(model_id.split("/")[-1]) + model_format = model_format.lower() + folder_to_delete = ( + TORCHSCRIPT_FOLDER_PATH if model_format == "torch_script" else ONNX_FOLDER_PATH + ) + + # Store to be uploaded files in UPLOAD_FOLDER_PATH + try: + dst_model_dir = ( + f"{UPLOAD_FOLDER_PATH}{model_name}/{model_version}/{model_format}" + ) + os.makedirs(dst_model_dir, exist_ok=True) + dst_model_filename = ( + f"sentence-transformers_{model_name}-{model_version}-{model_format}.zip" + ) + dst_model_path = dst_model_dir + "/" + dst_model_filename + with ZipFile(src_model_path, "a") as zipObj: + zipObj.write(filename=LICENSE_PATH, arcname="LICENSE") + shutil.copy(src_model_path, dst_model_path) + print() + print(f"Copied {src_model_path} to {dst_model_path}") + + dst_model_config_dir = ( + f"{UPLOAD_FOLDER_PATH}{model_name}/{model_version}/{model_format}" + ) + os.makedirs(dst_model_config_dir, exist_ok=True) + dst_model_config_filename = "config.json" + dst_model_config_path = dst_model_config_dir + "/" + dst_model_config_filename + shutil.copy(src_model_config_path, dst_model_config_path) + print(f"Copied {src_model_config_path} to {dst_model_config_path}") + except Exception as e: + assert ( + False + ), f"Raised Exception during preparing {model_format} files for uploading: {e}" + + # Delete model folder downloaded from HuggingFace during model tracing + try: + shutil.rmtree(folder_to_delete) + except Exception as e: + assert False, f"Raised Exception while deleting {folder_to_delete}: {e}" + + +def main( + model_id: str, + model_version: str, + tracing_format: str, + embedding_dimension: Optional[int] = None, + pooling_mode: Optional[str] = None, +) -> None: + """ + Perform model auto-tracing and prepare files for uploading to OpenSearch model hub + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + :param tracing_format: Tracing format ("TORCH_SCRIPT", "ONNX", or "BOTH") + :type tracing_format: string + :param embedding_dimension: Embedding dimension input + :type embedding_dimension: int + :param pooling_mode: Pooling mode input ("CLS", "MEAN", "MAX", "MEAN_SQRT_LEN" or None) + :type pooling_mode: string + """ + + print() + print("=== Begin running model_autotracing.py ===") + print("Model ID: ", model_id) + print("Model Version: ", model_version) + print("Tracing Format: ", tracing_format) + print("Embedding Dimension: ", embedding_dimension) + print("Pooling Mode: ", pooling_mode) + print("==========================================") + + ml_client = MLCommonClient(OPENSEARCH_TEST_CLIENT) + + pre_trained_model = SentenceTransformer(model_id) + original_embedding_data = list( + pre_trained_model.encode(TEST_SENTENCES, convert_to_numpy=True) + ) + + if tracing_format in [TORCH_SCRIPT_FORMAT, BOTH_FORMAT]: + print("--- Begin tracing a model in TORCH_SCRIPT ---") + ( + torchscript_model_path, + torchscript_model_config_path, + ) = trace_sentence_transformer_model( + model_id, + model_version, + TORCH_SCRIPT_FORMAT, + embedding_dimension, + pooling_mode, + ) + torch_embedding_data = register_and_deploy_sentence_transformer_model( + ml_client, + torchscript_model_path, + torchscript_model_config_path, + TORCH_SCRIPT_FORMAT, + ) + verify_embedding_data( + original_embedding_data, torch_embedding_data, TORCH_SCRIPT_FORMAT + ) + prepare_files_for_uploading( + model_id, + model_version, + TORCH_SCRIPT_FORMAT, + torchscript_model_path, + torchscript_model_config_path, + ) + print("--- Finished tracing a model in TORCH_SCRIPT ---") + + if tracing_format in [ONNX_FORMAT, BOTH_FORMAT]: + print("--- Begin tracing a model in ONNX ---") + onnx_model_path, onnx_model_config_path = trace_sentence_transformer_model( + model_id, + model_version, + ONNX_FORMAT, + embedding_dimension, + pooling_mode, + ) + onnx_embedding_data = register_and_deploy_sentence_transformer_model( + ml_client, onnx_model_path, onnx_model_config_path, ONNX_FORMAT + ) + + verify_embedding_data(original_embedding_data, onnx_embedding_data, ONNX_FORMAT) + prepare_files_for_uploading( + model_id, + model_version, + ONNX_FORMAT, + onnx_model_path, + onnx_model_config_path, + ) + print("--- Finished tracing a model in ONNX ---") + + print() + print("=== Finished running model_autotracing.py ===") + + +if __name__ == "__main__": + warnings.filterwarnings("ignore", category=DeprecationWarning) + warnings.filterwarnings("ignore", category=FutureWarning) + warnings.filterwarnings("ignore", message="Unverified HTTPS request") + warnings.filterwarnings("ignore", message="TracerWarning: torch.tensor") + warnings.filterwarnings( + "ignore", message="using SSL with verify_certs=False is insecure." + ) + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "model_id", + type=str, + help="Model ID for auto-tracing and uploading (e.g. sentence-transformers/msmarco-distilbert-base-tas-b)", + ) + parser.add_argument( + "model_version", type=str, help="Model version number (e.g. 1.0.1)" + ) + parser.add_argument( + "tracing_format", + choices=["BOTH", "TORCH_SCRIPT", "ONNX"], + help="Model format for auto-tracing", + ) + parser.add_argument( + "-ed", + "--embedding_dimension", + type=int, + nargs="?", + default=None, + const=None, + help="Embedding dimension of the model to use if it does not exist in original config.json", + ) + parser.add_argument( + "-pm", + "--pooling_mode", + type=str, + nargs="?", + default=None, + const=None, + choices=["CLS", "MEAN", "MAX", "MEAN_SQRT_LEN"], + help="Pooling mode if it does not exist in original config.json", + ) + args = parser.parse_args() + + main( + args.model_id, + args.model_version, + args.tracing_format, + args.embedding_dimension, + args.pooling_mode, + ) + + # TODO: Check if model exists in database diff --git a/utils/model_uploader/save_model_file_path_to_env.py b/utils/model_uploader/save_model_file_path_to_env.py new file mode 100644 index 00000000..1e6d5567 --- /dev/null +++ b/utils/model_uploader/save_model_file_path_to_env.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# The OpenSearch Contributors require contributions made to +# this file be licensed under the Apache-2.0 license or a +# compatible open source license. +# Any modifications Copyright OpenSearch Contributors. See +# GitHub history for details. + +# This program is run by "Model Auto-tracing & Uploading" workflow +# (See model_uploader.yml) to verify if the model already exists in +# model hub before continuing the workflow. + +import argparse +import re + +MODEL_ID_START_PATTERN = "sentence-transformers/" +VERSION_PATTERN = r"^([1-9]\d*|0)(\.(([1-9]\d*)|0)){0,3}$" +MODEL_HUB_PATH = "ml-models/huggingface/sentence-transformers/" + + +def verify_inputs(model_id: str, model_version: str) -> None: + """ + Verify the format of model_id and model_version + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + """ + if not model_id.startswith(MODEL_ID_START_PATTERN): + assert False, f"Invalid Model ID: {model_id}" + if re.fullmatch(VERSION_PATTERN, model_version) is None: + assert False, f"Invalid Model Version: {model_version}" + + +def get_model_file_path(model_id: str, model_version: str, model_format: str) -> str: + """ + Construct the expected model file path on model hub + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + :param model_format: Model format ("TORCH_SCRIPT" or "ONNX") + :type model_format: string + """ + model_name = str(model_id.split("/")[-1]) + model_format = model_format.lower() + model_dirname = f"{MODEL_HUB_PATH}{model_name}/{model_version}/{model_format}" + model_filename = ( + f"sentence-transformers_{model_name}-{model_version}-{model_format}.zip" + ) + model_file_path = model_dirname + "/" + model_filename + return model_file_path + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "model_id", + type=str, + help="Model ID for auto-tracing and uploading (e.g. sentence-transformers/msmarco-distilbert-base-tas-b)", + ) + parser.add_argument( + "model_version", type=str, help="Model version number (e.g. 1.0.1)" + ) + parser.add_argument( + "model_format", + choices=["TORCH_SCRIPT", "ONNX"], + help="Model format for auto-tracing", + ) + + args = parser.parse_args() + verify_inputs(args.model_id, args.model_version) + model_file_path = get_model_file_path( + args.model_id, args.model_version, args.model_format + ) + + # Print the model file path so that the workflow can store it in the variable (See model_uploader.yml) + print(model_file_path) diff --git a/utils/model_uploader/update_models_upload_history_md.py b/utils/model_uploader/update_models_upload_history_md.py new file mode 100644 index 00000000..ab08bf14 --- /dev/null +++ b/utils/model_uploader/update_models_upload_history_md.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# The OpenSearch Contributors require contributions made to +# this file be licensed under the Apache-2.0 license or a +# compatible open source license. +# Any modifications Copyright OpenSearch Contributors. See +# GitHub history for details. + +# This program is run by "Model Auto-tracing & Uploading" workflow +# (See model_uploader.yml) to update MODEL_UPLOAD_HISTORY.md & supported_models.json +# after uploading the model to our model hub. + +import argparse +import json +import os +from typing import Optional + +from mdutils.fileutils import MarkDownFile +from mdutils.tools.Table import Table + +MD_FILENAME = "MODEL_UPLOAD_HISTORY.md" +JSON_FILENAME = "supported_models.json" +DIRNAME = "utils/model_uploader/upload_history" +MODEL_JSON_FILEPATH = os.path.join(DIRNAME, JSON_FILENAME) +KEYS = [ + "Upload Time", + "Model Uploader", + "Model ID", + "Model Version", + "Tracing Format", + "Embedding Dimension", + "Pooling Mode", +] +MD_HEADER = "# Pretrained Model Upload History\n\nThe model-serving framework supports a variety of open-source pretrained models that can assist with a range of machine learning (ML) search and analytics use cases. \n\n\n## Uploaded Pretrained Models\n\n\n### Sentence transformers\n\nSentence transformer models map sentences and paragraphs across a dimensional dense vector space. The number of vectors depends on the model. Use these models for use cases such as clustering and semantic search. \n\nThe following table shows sentence transformer model upload history.\n\n[//]: # (This may be the most platform independent comment)\n" + + +def update_model_json_file( + model_id: str, + model_version: str, + tracing_format: str, + embedding_dimension: Optional[int] = None, + pooling_mode: Optional[str] = None, + model_uploader: Optional[str] = None, + uploader_time: Optional[str] = None, +) -> None: + """ + Update supported_models.json + + :param model_id: Model ID of the pretrained model + :type model_id: string + :param model_version: Version of the pretrained model for registration + :type model_version: string + :param tracing_format: Tracing format ("TORCH_SCRIPT", "ONNX", or "BOTH") + :type tracing_format: string + :param embedding_dimension: Embedding dimension input + :type embedding_dimension: int + :param pooling_mode: Pooling mode input ("CLS", "MEAN", "MAX", "MEAN_SQRT_LEN" or None) + :type pooling_mode: string + """ + models = [] + if os.path.exists(MODEL_JSON_FILEPATH): + with open(MODEL_JSON_FILEPATH, "r") as f: + models = json.load(f) + else: + os.makedirs(DIRNAME) + + new_model = { + "Model Uploader": "@" + model_uploader if model_uploader is not None else "-", + "Upload Time": uploader_time if uploader_time is not None else "-", + "Model ID": model_id, + "Model Version": model_version, + "Tracing Format": tracing_format, + "Embedding Dimension": str(embedding_dimension) + if embedding_dimension is not None + else "Default", + "Pooling Mode": pooling_mode if pooling_mode is not None else "Default", + } + + models.append(new_model) + models = [dict(t) for t in {tuple(m.items()) for m in models}] + models = sorted(models, key=lambda d: d["Upload Time"]) + with open(MODEL_JSON_FILEPATH, "w") as f: + json.dump(models, f, indent=4) + + +def update_md_file(): + """ + Update MODEL_UPLOAD_HISTORY.md + """ + models = [] + if os.path.exists(MODEL_JSON_FILEPATH): + with open(MODEL_JSON_FILEPATH, "r") as f: + models = json.load(f) + models = sorted(models, key=lambda d: d["Upload Time"]) + table_data = KEYS[:] + for m in models: + for k in KEYS: + if k == "Model ID": + table_data.append(f"`{m[k]}`") + else: + table_data.append(m[k]) + + table = Table().create_table( + columns=len(KEYS), rows=len(models) + 1, text=table_data, text_align="center" + ) + + mdFile = MarkDownFile(MD_FILENAME, dirname=DIRNAME) + mdFile.rewrite_all_file(data=MD_HEADER + table) + print(f"Finished updating {MD_FILENAME}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "model_id", + type=str, + help="Model ID for auto-tracing and uploading (e.g. sentence-transformers/msmarco-distilbert-base-tas-b)", + ) + parser.add_argument( + "model_version", type=str, help="Model version number (e.g. 1.0.1)" + ) + parser.add_argument( + "tracing_format", + choices=["BOTH", "TORCH_SCRIPT", "ONNX"], + help="Model format for auto-tracing", + ) + parser.add_argument( + "-ed", + "--embedding_dimension", + type=int, + nargs="?", + default=None, + const=None, + help="Embedding dimension of the model to use if it does not exist in original config.json", + ) + + parser.add_argument( + "-pm", + "--pooling_mode", + type=str, + nargs="?", + default=None, + const=None, + choices=["CLS", "MEAN", "MAX", "MEAN_SQRT_LEN"], + help="Pooling mode if it does not exist in original config.json", + ) + + parser.add_argument( + "-u", + "--model_uploader", + type=str, + nargs="?", + default=None, + const=None, + help="Model Uploader", + ) + + parser.add_argument( + "-t", + "--upload_time", + type=str, + nargs="?", + default=None, + const=None, + help="Upload Time", + ) + args = parser.parse_args() + + update_model_json_file( + args.model_id, + args.model_version, + args.tracing_format, + args.embedding_dimension, + args.pooling_mode, + args.model_uploader, + args.upload_time, + ) + + update_md_file() diff --git a/utils/model_uploader/upload_history/MODEL_UPLOAD_HISTORY.md b/utils/model_uploader/upload_history/MODEL_UPLOAD_HISTORY.md new file mode 100644 index 00000000..81186de2 --- /dev/null +++ b/utils/model_uploader/upload_history/MODEL_UPLOAD_HISTORY.md @@ -0,0 +1,20 @@ +# Pretrained Model Upload History + +The model-serving framework supports a variety of open-source pretrained models that can assist with a range of machine learning (ML) search and analytics use cases. + + +## Uploaded Pretrained Models + + +### Sentence transformers + +Sentence transformer models map sentences and paragraphs across a dimensional dense vector space. The number of vectors depends on the model. Use these models for use cases such as clustering and semantic search. + +The following table shows sentence transformer model upload history. + +[//]: # (This may be the most platform independent comment) + +|Upload Time|Model Uploader|Model ID|Model Version|Tracing Format|Embedding Dimension|Pooling Mode| +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +|2023-07-17 15:35:25|@thanawan-atc|`sentence-transformers/msmarco-distilbert-base-tas-b`|1.0.1|ONNX|Default|Default| +|2023-07-18 22:32:24|@thanawan-atc|`sentence-transformers/msmarco-distilbert-base-tas-b`|1.0.1|TORCH_SCRIPT|Default|Default| diff --git a/utils/model_uploader/upload_history/supported_models.json b/utils/model_uploader/upload_history/supported_models.json new file mode 100644 index 00000000..9df2c435 --- /dev/null +++ b/utils/model_uploader/upload_history/supported_models.json @@ -0,0 +1,20 @@ +[ + { + "Model Uploader": "@thanawan-atc", + "Upload Time": "2023-07-17 15:35:25", + "Model ID": "sentence-transformers/msmarco-distilbert-base-tas-b", + "Model Version": "1.0.1", + "Tracing Format": "ONNX", + "Embedding Dimension": "Default", + "Pooling Mode": "Default" + }, + { + "Model Uploader": "@thanawan-atc", + "Upload Time": "2023-07-18 22:32:24", + "Model ID": "sentence-transformers/msmarco-distilbert-base-tas-b", + "Model Version": "1.0.1", + "Tracing Format": "TORCH_SCRIPT", + "Embedding Dimension": "Default", + "Pooling Mode": "Default" + } +] \ No newline at end of file