Kheiss/qa review3

.github/workflows/release-helm.yml

@@ -39,9 +39,12 @@ jobs:
     runs-on: ubuntu-latest
     env:
       NGC_CLI_API_KEY: ${{ secrets.NVIDIA_API_KEY }}
+      NGC_CLI_ORG: ${{ inputs.ngc-org }}
+      NGC_CLI_TEAM: ${{ inputs.ngc-team }}
+      NGC_CLI_FORMAT_TYPE: json
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: ${{ inputs.source-ref }}
 
@@ -54,23 +57,13 @@ jobs:
           curl -sSL "https://github.com/norwoodj/helm-docs/releases/download/v${HELM_DOCS_VERSION}/helm-docs_${HELM_DOCS_VERSION}_Linux_x86_64.tar.gz" \
             | tar xz -C /usr/local/bin helm-docs
 
-      - name: Install NGC CLI
-        run: |
-          curl -sSL "https://api.ngc.nvidia.com/v2/resources/nvidia/ngc-apps/ngc_cli/versions/3.55.0/files/ngccli_linux.zip" -o /tmp/ngccli.zip
-          unzip -q /tmp/ngccli.zip -d /tmp
-          sudo mv /tmp/ngc-cli/ngc /usr/local/bin/ngc
-          sudo chmod +x /usr/local/bin/ngc
+      - name: Setup Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: '3.12'
 
-      - name: Configure and verify NGC CLI
-        run: |
-          ngc config set <<EOF
-          $NGC_CLI_API_KEY
-          json
-          ${{ inputs.ngc-org }}
-          ${{ inputs.ngc-team }}
-          EOF
-          echo "NGC CLI configured. Verifying authentication..."
-          ngc config current
+      - name: Install Python dependencies
+        run: pip install ngcsdk pyyaml
 
       - name: Update Helm README
         run: helm/update_helm_readme.sh
@@ -88,14 +81,6 @@ jobs:
           helm dependency update helm/
           helm dependency build helm/
 
-      - name: Setup Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.12'
-
-      - name: Install Python dependencies
-        run: pip install pyyaml
-
       - name: Release Helm chart
         run: |
           DRY_RUN_FLAG=""

.github/workflows/retriever-unit-tests.yml

@@ -19,11 +19,15 @@ jobs:
         with:
           python-version: "3.12"
 
+      - name: Install uv
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.local/bin" >> "$GITHUB_PATH"
+
       - name: Install unit test dependencies
         run: |
-          python -m pip install --upgrade pip
-          python -m pip install pytest pandas pydantic pyyaml typer scikit-learn
-          python -m pip install api/
+          uv pip install --system -e src/ -e api/ -e client/
+          uv pip install --system -e nemo_retriever
 
       - name: Run retriever unit tests
         env:

.github/workflows/reusable-pypi-build.yml

@@ -18,6 +18,11 @@ on:
         required: false
         type: string
         default: 'main'
+      workflow-ref:
+        description: 'Git ref of the workflow branch (used to overlay pyproject.toml files)'
+        required: false
+        type: string
+        default: ''
       runner:
         description: 'GitHub runner to use'
         required: false
@@ -36,10 +41,16 @@ jobs:
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: ${{ inputs.source-ref }}
 
+      - name: Overlay build config from workflow branch
+        if: ${{ inputs.workflow-ref != '' && inputs.workflow-ref != inputs.source-ref }}
+        run: |
+          git fetch --depth=1 origin "${{ inputs.workflow-ref }}"
+          git checkout FETCH_HEAD -- api/pyproject.toml client/pyproject.toml src/pyproject.toml nemo_retriever/pyproject.toml
+
       - name: Determine version
         id: set-version
         run: |
@@ -52,7 +63,7 @@ jobs:
           echo "Building version: $VERSION"
 
       - name: Setup Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: '3.12'
 
@@ -103,12 +114,12 @@ jobs:
           PY
           RETRIEVER_RELEASE_TYPE=${{ inputs.release-type }} \
           RETRIEVER_VERSION=${{ steps.set-version.outputs.version }} \
-          RETRIEVER_BUILD_NUMBER=${{ github.run_number }} \
+          RETRIEVER_BUILD_NUMBER=${{ inputs.release-type == 'release' && '0' || github.run_number }} \
           RETRIEVER_GIT_SHA=${{ github.sha }} \
           python -m build
 
       - name: Upload wheel artifacts
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: python-wheels
           path: |

.github/workflows/reusable-pypi-publish.yml

@@ -16,13 +16,18 @@ jobs:
 
     steps:
       - name: Download wheel artifacts
-        uses: actions/download-artifact@v4
+        uses: actions/download-artifact@v5
         with:
           name: python-wheels
           path: ./dist
 
+      - name: Setup Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: '3.12'
+
       - name: Install twine
-        run: pip install twine
+        run: pip install 'twine>=6.1'
 
       - name: Publish wheels to Artifactory
         env:
@@ -31,7 +36,7 @@ jobs:
           ARTIFACTORY_PASSWORD: ${{ secrets.ARTIFACTORY_PASSWORD }}
         run: |
           # Publish all wheels
-          twine upload \
+          twine upload --verbose \
             --repository-url $ARTIFACTORY_URL \
             -u $ARTIFACTORY_USERNAME \
             -p $ARTIFACTORY_PASSWORD \

Dockerfile

@@ -39,13 +39,13 @@ RUN chmod +x scripts/install_ffmpeg.sh \
 # For GPL-licensed components, we provide their source code in the container
 # via `apt-get source` below to satisfy GPL requirements.
 ARG GPL_LIBS="\
+    libfreetype6 \
     libltdl7 \
     libhunspell-1.7-0 \
     libhyphen0 \
     libdbus-1-3 \
 "
 ARG FORCE_REMOVE_PKGS="\
-    libfreetype6 \
     ucf \
     liblangtag-common \
     libjbig0 \

api/pyproject.toml

@@ -13,10 +13,10 @@ readme = "README.md"
 authors = [
     {name = "Jeremy Dyer", email = "jdyer@nvidia.com"}
 ]
-license = {file = "LICENSE"}
+license = "Apache-2.0"
+license-files = ["LICENSE"]
 classifiers = [
     "Programming Language :: Python :: 3",
-    "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
 ]
 dependencies = [

api/src/nv_ingest_api/internal/primitives/nim/model_interface/ocr.py

@@ -16,7 +16,8 @@
 import tritonclient.grpc as grpcclient
 
 from nv_ingest_api.internal.primitives.nim import ModelInterface
-from nv_ingest_api.internal.primitives.nim.model_interface.decorators import multiprocessing_cache
+from nv_ingest_api.internal.primitives.nim.model_interface.decorators import global_cache
+from nv_ingest_api.internal.primitives.nim.model_interface.decorators import lock
 from nv_ingest_api.internal.primitives.nim.model_interface.helpers import preprocess_image_for_paddle
 from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
 
@@ -752,12 +753,11 @@ def _format_single_batch(
             raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
 
 
-@multiprocessing_cache(max_calls=100)  # Cache results first to avoid redundant retries from backoff
 @backoff.on_predicate(backoff.expo, max_time=30)
 def get_ocr_model_name(ocr_grpc_endpoint=None, default_model_name=DEFAULT_OCR_MODEL_NAME):
     """
     Determines the OCR model name by checking the environment, querying the gRPC endpoint,
-    or falling back to a default.
+    or falling back to a default. Only caches when the repository is successfully queried.
     """
     # 1. Check for an explicit override from the environment variable first.
     ocr_model_name = os.getenv("OCR_MODEL_NAME", None)
@@ -769,14 +769,25 @@ def get_ocr_model_name(ocr_grpc_endpoint=None, default_model_name=DEFAULT_OCR_MO
         logger.debug(f"No OCR gRPC endpoint provided. Falling back to default model name '{default_model_name}'.")
         return default_model_name
 
-    # 3. Attempt to query the gRPC endpoint to discover the model name.
+    # 3. Check cache (only populated on successful repository query).
+    key = (
+        "get_ocr_model_name",
+        (ocr_grpc_endpoint,),
+        frozenset({"default_model_name": default_model_name}.items()),
+    )
+    with lock:
+        if key in global_cache:
+            return global_cache[key]
+
+    # 4. Attempt to query the gRPC endpoint to discover the model name.
     try:
         client = grpcclient.InferenceServerClient(ocr_grpc_endpoint)
         model_index = client.get_model_repository_index(as_json=True)
         model_names = [x["name"] for x in model_index.get("models", [])]
         ocr_model_name = model_names[0]
+        with lock:
+            global_cache[key] = ocr_model_name
+        return ocr_model_name
     except Exception:
         logger.warning(f"Failed to get ocr model name after 30 seconds. Falling back to '{default_model_name}'.")
-        ocr_model_name = default_model_name
-
-    return ocr_model_name
+        return default_model_name

api/src/nv_ingest_api/internal/primitives/nim/model_interface/yolox.py

@@ -20,6 +20,8 @@
 
 from nv_ingest_api.internal.primitives.nim import ModelInterface
 import tritonclient.grpc as grpcclient
+from nv_ingest_api.internal.primitives.nim.model_interface.decorators import global_cache
+from nv_ingest_api.internal.primitives.nim.model_interface.decorators import lock
 from nv_ingest_api.internal.primitives.nim.model_interface.decorators import multiprocessing_cache
 from nv_ingest_api.internal.primitives.nim.model_interface.helpers import get_model_name
 from nv_ingest_api.util.image_processing import scale_image_to_encoding_size
@@ -135,10 +137,36 @@ def __init__(
         self.class_labels = class_labels
 
         if endpoints:
-            self.model_name = get_yolox_model_name(endpoints[0], default_model_name="yolox_ensemble")
-            self._grpc_uses_bls = self.model_name == "pipeline"
+            self._yolox_grpc_endpoint = endpoints[0]
+            self._model_name = None
+            self._grpc_uses_bls_value = None  # Resolved on first use
         else:
-            self._grpc_uses_bls = False
+            self._yolox_grpc_endpoint = None
+            self._model_name = None
+            self._grpc_uses_bls_value = False
+
+    def _resolve_yolox_model_name_if_needed(self) -> None:
+        """Resolve model name and BLS flag from the gRPC endpoint on first use. Cached on the instance."""
+        if self._yolox_grpc_endpoint is None:
+            return
+        if self._model_name is not None:
+            return
+        self._model_name = get_yolox_model_name(self._yolox_grpc_endpoint, default_model_name="yolox_ensemble")
+        self._grpc_uses_bls_value = self._model_name == "pipeline"
+
+    @property
+    def model_name(self) -> Optional[str]:
+        self._resolve_yolox_model_name_if_needed()
+        return self._model_name
+
+    @model_name.setter
+    def model_name(self, value: Optional[str]) -> None:
+        self._model_name = value
+
+    @property
+    def _grpc_uses_bls(self) -> bool:
+        self._resolve_yolox_model_name_if_needed()
+        return bool(self._grpc_uses_bls_value)
 
     def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
@@ -2117,7 +2145,6 @@ def postprocess_included_texts(boxes, confs, labels, classes):
     return boxes, labels, confs
 
 
-@multiprocessing_cache(max_calls=100)  # Cache results first to avoid redundant retries from backoff
 @backoff.on_predicate(backoff.expo, max_time=30)
 def get_yolox_model_name(yolox_grpc_endpoint, default_model_name="yolox"):
     # If a gRPC endpoint isn't provided (common when using HTTP-only NIM endpoints),
@@ -2131,6 +2158,15 @@ def get_yolox_model_name(yolox_grpc_endpoint, default_model_name="yolox"):
     ):
         return default_model_name
 
+    key = (
+        "get_yolox_model_name",
+        (yolox_grpc_endpoint,),
+        frozenset({"default_model_name": default_model_name}.items()),
+    )
+    with lock:
+        if key in global_cache:
+            return global_cache[key]
+
     try:
         client = grpcclient.InferenceServerClient(yolox_grpc_endpoint)
         model_index = client.get_model_repository_index(as_json=True)
@@ -2148,14 +2184,23 @@ def get_yolox_model_name(yolox_grpc_endpoint, default_model_name="yolox"):
             "nemoretriever-page-elements-v2",
         ):
             if preferred in model_names:
-                return preferred
+                result = preferred
+                with lock:
+                    global_cache[key] = result
+                return result
 
         # Otherwise pick a best-effort match for newer model names.
         candidates = [m for m in model_names if isinstance(m, str) and ("yolox" in m or "page-elements" in m)]
         if candidates:
-            return sorted(candidates)[0]
-
-        return default_model_name
+            result = sorted(candidates)[0]
+            with lock:
+                global_cache[key] = result
+            return result
+
+        result = default_model_name
+        with lock:
+            global_cache[key] = result
+        return result
     except Exception as e:
         logger.warning(
             "Failed to inspect YOLOX model repository at '%s' (%s). Falling back to '%s'.",

api/src/nv_ingest_api/internal/transform/split_text.py

@@ -56,14 +56,8 @@ def _get_tokenizer(
         if cache_key in _tokenizer_cache:
             return _tokenizer_cache[cache_key]
 
-        from nemo_retriever.utils.hf_model_registry import get_hf_revision
-
         logger.info("Loading and caching tokenizer: %s", tokenizer_identifier)
-        tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_identifier,
-            revision=get_hf_revision(tokenizer_identifier),
-            token=token,
-        )
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_identifier, token=token)
         _tokenizer_cache[cache_key] = tokenizer
         return tokenizer