Repository: opensearch-project/opensearch-py-ml
Pull request #475: Fix integ test flaky test
Author: login: zhichao-aws
Created: 2025-04-18T02:25:16Z | Merged: 2025-04-23T14:44:47Z

### Description
We see integ test failing https://github.com/opensearch-project/opensearch-py-ml/actions/runs/14520924174/job/40742876990, and the error message is 
```
E           huggingface_hub.errors.HfHubHTTPError: 429 Client Error: Too Many Requests for url: https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill/resolve/main/config.json
```

The error is caused by huggingface throttling. This PR makes all sparse model tests use one model id, therefore we can reuse the model file in cache, to reduce the probability of flaky test.
 
### Issues Resolved
[List any issues this PR will resolve]
 
### Check List
- [ ] New functionality includes testing.
  - [ ] All tests pass
- [ ] New functionality has been documented.
  - [ ] New functionality has javadoc added
- [ ] Commits are signed per the DCO using --signoff 

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check [here](https://github.com/opensearch-project/OpenSearch/blob/main/CONTRIBUTING.md#developer-certificate-of-origin).



Diff:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0efed5fb..1213800c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -70,6 +70,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - fix CVE vulnerability by @rawwar in ([#383](https://github.com/opensearch-project/opensearch-py-ml/pull/383))
 - refactor: replace 'payload' with 'body' in `create_standalone_connector` by @yerzhaisang ([#424](https://github.com/opensearch-project/opensearch-py-ml/pull/424))
 - Fix CVE vulnerability by @nathaliellenaa in ([#447](https://github.com/opensearch-project/opensearch-py-ml/pull/447))
+- Use one model id for sparse model integ test to reduce flaky error caused by throttling ([#475](https://github.com/opensearch-project/opensearch-py-ml/pull/475))
 
 ## [1.1.0]
 
diff --git a/tests/ml_models/test_sparse_tokenize_model_pytest.py b/tests/ml_models/test_sparse_tokenize_model_pytest.py
index da91f7aa..86edce1d 100644
--- a/tests/ml_models/test_sparse_tokenize_model_pytest.py
+++ b/tests/ml_models/test_sparse_tokenize_model_pytest.py
@@ -56,11 +56,11 @@ def test_check_attribute():
     clean_test_folder(TEST_FOLDER)
     test_model1 = SparseTokenizeModel(
         folder_path=TEST_FOLDER,
-        model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill",
+        model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
     )
     assert (
         test_model1.model_id
-        == "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+        == "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     )
 
 
@@ -92,7 +92,7 @@ def test_save_as_pt():
 def test_make_model_config_json_for_torch_script():
     model_format = "TORCH_SCRIPT"
     expected_model_description = "This is a sparse encoding model for opensearch-neural-sparse-encoding-doc-v3-distill."
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     clean_test_folder(TEST_FOLDER)
     test_model3 = SparseTokenizeModel(model_id=model_id, folder_path=TEST_FOLDER)
     test_model3.save_as_pt(model_id=model_id, sentences=["today is sunny"])
@@ -112,7 +112,7 @@ def test_make_model_config_json_for_torch_script():
 
 
 def test_overwrite_description():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     expected_model_description = "Expected Description"
 
@@ -143,7 +143,7 @@ def test_overwrite_description():
 
 
 def test_long_description():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     expected_model_description = (
         "This is a sparce encoding model: It generate lots of tokens with different weight "
@@ -178,10 +178,10 @@ def test_long_description():
 
 
 def test_save_as_pt_with_license():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     torch_script_zip_file_path = os.path.join(
-        TEST_FOLDER, "opensearch-neural-sparse-encoding-doc-v3-distill.zip"
+        TEST_FOLDER, "opensearch-neural-sparse-encoding-doc-v2-distill.zip"
     )
     torch_script_expected_filenames = {
         "idf.json",
@@ -207,7 +207,7 @@ def test_save_as_pt_with_license():
 
 
 def test_default_description():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     expected_model_description = "This is a neural sparse tokenizer model: It tokenize input sentence into tokens and assign pre-defined weight from IDF to each. It serves only in query."
 
@@ -238,7 +238,7 @@ def test_default_description():
 
 
 def test_process_sparse_encoding():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
 
     test_model8 = SparseTokenizeModel(
         folder_path=TEST_FOLDER,
@@ -252,7 +252,7 @@ def test_process_sparse_encoding():
 
 
 def test_save():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
 
     test_model9 = SparseTokenizeModel(
         folder_path=TEST_FOLDER,
diff --git a/tests/ml_models/test_sparseencondingmodel_pytest.py b/tests/ml_models/test_sparseencondingmodel_pytest.py
index b273dd01..f336ce04 100644
--- a/tests/ml_models/test_sparseencondingmodel_pytest.py
+++ b/tests/ml_models/test_sparseencondingmodel_pytest.py
@@ -8,6 +8,8 @@
 import json
 import os
 import shutil
+import time
+from functools import wraps
 from zipfile import ZipFile
 
 import pytest
@@ -25,6 +27,35 @@
 )
 
 
+def retry(n_retries=3, delay_sec=10, exceptions=(Exception,)):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            for attempt in range(1, n_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except exceptions as e:
+                    if attempt == n_retries:
+                        raise
+                    print(
+                        f"Attempt {attempt} failed: {e}. Retrying in {delay_sec}sec..."
+                    )
+                    time.sleep(delay_sec)
+            return None
+
+        return wrapper
+
+    return decorator
+
+
+@retry(n_retries=5, delay_sec=60, exceptions=(ConnectionError,))
+def init_model():
+    return SparseEncodingModel(
+        folder_path=TEST_FOLDER,
+        model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+    )
+
+
 def clean_test_folder(TEST_FOLDER):
     if os.path.exists(TEST_FOLDER):
         for files in os.listdir(TEST_FOLDER):
@@ -104,8 +135,8 @@ def compare_model_zip_file(zip_file_path, expected_filenames, model_format):
 
 
 clean_test_folder(TEST_FOLDER)
-# test model with a default model id opensearch-project/opensearch-neural-sparse-encoding-v1
-test_model = SparseEncodingModel(folder_path=TEST_FOLDER)
+# test model with model id "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
+test_model = init_model()
 
 
 def test_check_attribute():
@@ -117,13 +148,16 @@ def test_check_attribute():
 
     assert test_model.folder_path == TEST_FOLDER
     assert (
-        test_model.model_id == "opensearch-project/opensearch-neural-sparse-encoding-v1"
+        test_model.model_id
+        == "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     )
 
     default_folder = os.path.join(os.getcwd(), "opensearch_neural_sparse_model_files")
 
     clean_test_folder(default_folder)
-    test_model0 = SparseEncodingModel()
+    test_model0 = SparseEncodingModel(
+        model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
+    )
     assert test_model0.folder_path == default_folder
     clean_test_folder(default_folder)
 
@@ -139,7 +173,11 @@ def test_folder_path():
         test_non_empty_path = os.path.join(
             os.path.dirname(os.path.abspath("__file__")), "tests"
         )
-        SparseEncodingModel(folder_path=test_non_empty_path, overwrite=False)
+        SparseEncodingModel(
+            model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+            folder_path=test_non_empty_path,
+            overwrite=False,
+        )
     assert exc_info.type is Exception
     assert "The default folder path already exists" in exc_info.value.args[0]
 
@@ -161,12 +199,13 @@ def test_save_as_pt():
 
 def test_make_model_config_json_for_torch_script():
     model_format = "TORCH_SCRIPT"
-    expected_model_description = (
-        "This is a sparse encoding model for opensearch-neural-sparse-encoding-v1."
-    )
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-v1"
+    expected_model_description = "This is a sparse encoding model for opensearch-neural-sparse-encoding-doc-v2-distill."
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     clean_test_folder(TEST_FOLDER)
-    test_model3 = SparseEncodingModel(folder_path=TEST_FOLDER)
+    test_model3 = SparseEncodingModel(
+        model_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill",
+        folder_path=TEST_FOLDER,
+    )
     test_model3.save_as_pt(model_id=model_id, sentences=["today is sunny"])
     model_config_path_torch = test_model3.make_model_config_json(
         model_format="TORCH_SCRIPT", description=expected_model_description
@@ -214,7 +253,7 @@ def test_overwrite_description():
 
 
 def test_long_description():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-v1"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     expected_model_description = (
         "This is a sparce encoding model: It generate lots of tokens with different weight "
@@ -249,13 +288,13 @@ def test_long_description():
 
 
 def test_save_as_pt_with_license():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-v1"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     torch_script_zip_file_path = os.path.join(
-        TEST_FOLDER, "opensearch-neural-sparse-encoding-v1.zip"
+        TEST_FOLDER, "opensearch-neural-sparse-encoding-doc-v2-distill.zip"
     )
     torch_script_expected_filenames = {
-        "opensearch-neural-sparse-encoding-v1.pt",
+        "opensearch-neural-sparse-encoding-doc-v2-distill.pt",
         "tokenizer.json",
         "LICENSE",
     }
@@ -278,7 +317,7 @@ def test_save_as_pt_with_license():
 
 
 def test_default_description():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
     model_format = "TORCH_SCRIPT"
     expected_model_description = "This is a neural sparse encoding model: It transfers text into sparse vector, and then extract nonzero index and value to entry and weights. It serves only in ingestion and customer should use tokenizer model in query."
 
@@ -309,7 +348,7 @@ def test_default_description():
 
 
 def test_process_sparse_encoding():
-    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-distill"
+    model_id = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill"
 
     test_model8 = SparseEncodingModel(
         folder_path=TEST_FOLDER,
@@ -317,10 +356,10 @@ def test_process_sparse_encoding():
     )
 
     encoding_result = test_model8.process_sparse_encoding(["hello world", "hello"])
-    assert len(encoding_result[0]) == 73
-    check_value(1.3667216300964355, encoding_result[0]["hello"], 0.001)
-    assert len(encoding_result[1]) == 46
-    check_value(1.4557286500930786, encoding_result[1]["hello"], 0.001)
+    assert len(encoding_result[0]) == 298
+    check_value(1.6551653146743774, encoding_result[0]["hello"], 0.001)
+    assert len(encoding_result[1]) == 260
+    check_value(1.9172965288162231, encoding_result[1]["hello"], 0.001)
 
     test_model8 = SparseEncodingModel(
         folder_path=TEST_FOLDER,
@@ -329,10 +368,10 @@ def test_process_sparse_encoding():
         activation="l0",
     )
     encoding_result = test_model8.process_sparse_encoding(["hello world", "hello"])
-    assert len(encoding_result[0]) == 33
-    check_value(0.8615057468414307, encoding_result[0]["hello"], 0.001)
-    assert len(encoding_result[1]) == 30
-    check_value(0.8984234929084778, encoding_result[1]["hello"], 0.001)
+    assert len(encoding_result[0]) == 77
+    check_value(0.9765068888664246, encoding_result[0]["hello"], 0.001)
+    assert len(encoding_result[1]) == 64
+    check_value(1.0706572532653809, encoding_result[1]["hello"], 0.001)
 
 
 clean_test_folder(TEST_FOLDER)


Reviews:
user: login: dhrubo-os
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-18T05:51:56Z
user: login: zhichao-aws
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-18T06:24:00Z
user: login: dhrubo-os
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-23T02:55:43Z
user: login: b4sjoo
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-23T06:14:07Z

Review comments:
diff_hunk: @@ -25,6 +27,35 @@
 )
 
 
+def retry(n_retries=3, delay_sec=10, exceptions=(Exception,)):
path: tests/ml_models/test_sparseencondingmodel_pytest.py
user: login: dhrubo-os
body: Will that increase the total integration test run time?
created_at: 2025-04-18T05:51:55Z
author_association: COLLABORATOR
diff_hunk: @@ -25,6 +27,35 @@
 )
 
 
+def retry(n_retries=3, delay_sec=10, exceptions=(Exception,)):
path: tests/ml_models/test_sparseencondingmodel_pytest.py
user: login: zhichao-aws
body: If we success to download the model at first time, it won't have impact on test time.

But if the first time get throttled, then yes, we'll wait extra times for retry.
created_at: 2025-04-18T06:24:00Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-04-18T03:11:57Z
body: ## [Codecov](https://app.codecov.io/gh/opensearch-project/opensearch-py-ml/pull/475?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 91.52%. Comparing base [(`529ee34`)](https://app.codecov.io/gh/opensearch-project/opensearch-py-ml/commit/529ee3441ffc0bee59ca72adf214bf61599e0901?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) to head [(`540a501`)](https://app.codecov.io/gh/opensearch-project/opensearch-py-ml/commit/540a501eb1392e47cdc278ee81a6b73aeb94ef51?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).
> Report is 24 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #475      +/-   ##
==========================================
- Coverage   91.53%   91.52%   -0.01%     
==========================================
  Files          42       44       +2     
  Lines        4395     4757     +362     
==========================================
+ Hits         4023     4354     +331     
- Misses        372      403      +31     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/opensearch-project/opensearch-py-ml/pull/475?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: dhrubo-os
created_at: 2025-04-18T05:52:58Z
body: Why do we need to download the model every time?
author_association: COLLABORATOR
user: login: zhichao-aws
created_at: 2025-04-18T06:25:11Z
body: > Why do we need to download the model every time?

Because every time we trigger integ test CI, it starts a new instance and run all test cases. On this instance, it doesn't have local file cache of these models, so it need to download the model file when first called to init the model.

But for this instance, it only download the model once. The following init code just read from local cache.
author_association: MEMBER