Refactor get_*_model_name to avoid caching fallback model name

Author: charlesbluca

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'.",