feat(harness): respect "Use Default Heuristics" preset option in harness run (#1804)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: charlesbluca

nemo_retriever/src/nemo_retriever/harness/config.py

@@ -82,6 +82,7 @@ class HarnessConfig:
     extract_page_as_image: bool = True
     extract_infographics: bool = False
     write_detection_file: bool = False
+    use_heuristics: bool = False
 
     pdf_extract_workers: int = 8
     pdf_extract_num_cpus: float = 2.0
@@ -158,12 +159,15 @@ def validate(self) -> list[str]:
         if self.embed_granularity not in VALID_EMBED_GRANULARITIES:
             errors.append(f"embed_granularity must be one of {sorted(VALID_EMBED_GRANULARITIES)}")
 
+        _ZERO_ALLOWED_WORKERS = {f for f in TUNING_FIELDS if f.endswith("_workers")} if self.use_heuristics else set()
         for name in TUNING_FIELDS:
             val = getattr(self, name)
             if name.startswith("gpu_") and float(val) < 0.0:
                 errors.append(f"{name} must be >= 0.0")
-            elif name.endswith("_workers") and int(val) < 1:
-                errors.append(f"{name} must be >= 1")
+            elif name.endswith("_workers"):
+                min_val = 0 if name in _ZERO_ALLOWED_WORKERS else 1
+                if int(val) < min_val:
+                    errors.append(f"{name} must be >= {min_val}")
 
         return errors
 
@@ -281,6 +285,7 @@ def _apply_env_overrides(config_dict: dict[str, Any]) -> None:
         "HARNESS_EXTRACT_PAGE_AS_IMAGE": ("extract_page_as_image", _parse_bool),
         "HARNESS_EXTRACT_INFOGRAPHICS": ("extract_infographics", _parse_bool),
         "HARNESS_WRITE_DETECTION_FILE": ("write_detection_file", _parse_bool),
+        "HARNESS_USE_HEURISTICS": ("use_heuristics", _parse_bool),
     }
 
     for key in TUNING_FIELDS:

nemo_retriever/src/nemo_retriever/harness/run.py

@@ -206,38 +206,45 @@ def _build_command(cfg: HarnessConfig, artifact_dir: Path, run_id: str) -> tuple
         cfg.input_type,
         "--evaluation-mode",
         cfg.evaluation_mode,
-        "--pdf-extract-tasks",
-        str(cfg.pdf_extract_workers),
-        "--pdf-extract-cpus-per-task",
-        str(cfg.pdf_extract_num_cpus),
-        "--pdf-extract-batch-size",
-        str(cfg.pdf_extract_batch_size),
-        "--pdf-split-batch-size",
-        str(cfg.pdf_split_batch_size),
-        "--page-elements-batch-size",
-        str(cfg.page_elements_batch_size),
-        "--page-elements-actors",
-        str(cfg.page_elements_workers),
-        "--ocr-actors",
-        str(cfg.ocr_workers),
-        "--ocr-batch-size",
-        str(cfg.ocr_batch_size),
-        "--embed-actors",
-        str(cfg.embed_workers),
-        "--embed-batch-size",
-        str(cfg.embed_batch_size),
-        "--page-elements-cpus-per-actor",
-        str(cfg.page_elements_cpus_per_actor),
-        "--ocr-cpus-per-actor",
-        str(cfg.ocr_cpus_per_actor),
-        "--embed-cpus-per-actor",
-        str(cfg.embed_cpus_per_actor),
-        "--page-elements-gpus-per-actor",
-        str(cfg.gpu_page_elements),
-        "--ocr-gpus-per-actor",
-        str(cfg.gpu_ocr),
-        "--embed-gpus-per-actor",
-        str(cfg.gpu_embed),
+    ]
+
+    if not cfg.use_heuristics:
+        cmd += [
+            "--pdf-extract-tasks",
+            str(cfg.pdf_extract_workers),
+            "--pdf-extract-cpus-per-task",
+            str(cfg.pdf_extract_num_cpus),
+            "--pdf-extract-batch-size",
+            str(cfg.pdf_extract_batch_size),
+            "--pdf-split-batch-size",
+            str(cfg.pdf_split_batch_size),
+            "--page-elements-batch-size",
+            str(cfg.page_elements_batch_size),
+            "--page-elements-actors",
+            str(cfg.page_elements_workers),
+            "--ocr-actors",
+            str(cfg.ocr_workers),
+            "--ocr-batch-size",
+            str(cfg.ocr_batch_size),
+            "--embed-actors",
+            str(cfg.embed_workers),
+            "--embed-batch-size",
+            str(cfg.embed_batch_size),
+            "--page-elements-cpus-per-actor",
+            str(cfg.page_elements_cpus_per_actor),
+            "--ocr-cpus-per-actor",
+            str(cfg.ocr_cpus_per_actor),
+            "--embed-cpus-per-actor",
+            str(cfg.embed_cpus_per_actor),
+            "--page-elements-gpus-per-actor",
+            str(cfg.gpu_page_elements),
+            "--ocr-gpus-per-actor",
+            str(cfg.gpu_ocr),
+            "--embed-gpus-per-actor",
+            str(cfg.gpu_embed),
+        ]
+
+    cmd += [
         "--embed-model-name",
         cfg.embed_model_name,
         "--embed-modality",
@@ -470,6 +477,7 @@ def _run_single(cfg: HarnessConfig, artifact_dir: Path, run_id: str, tags: list[
             "extract_page_as_image": cfg.extract_page_as_image,
             "extract_infographics": cfg.extract_infographics,
             "write_detection_file": cfg.write_detection_file,
+            "use_heuristics": cfg.use_heuristics,
             "lancedb_uri": _resolve_lancedb_uri(cfg, artifact_dir),
             "tuning": {field: getattr(cfg, field) for field in sorted(TUNING_FIELDS)},
         },

nemo_retriever/tests/test_harness_run.py

@@ -78,14 +78,6 @@ def test_build_command_uses_hidden_detection_file_by_default(tmp_path: Path) ->
     assert "element" in cmd
     assert "--extract-page-as-image" in cmd
     assert "--no-extract-page-as-image" not in cmd
-    assert "--pdf-extract-workers" not in cmd
-    assert "--pdf-extract-num-cpus" not in cmd
-    assert "--page-elements-workers" not in cmd
-    assert "--ocr-workers" not in cmd
-    assert "--embed-workers" not in cmd
-    assert "--gpu-page-elements" not in cmd
-    assert "--gpu-ocr" not in cmd
-    assert "--gpu-embed" not in cmd
     assert detection_file.parent == runtime_dir
     assert detection_file.name == ".detection_summary.json"
     assert effective_query_csv == query_csv
@@ -232,6 +224,42 @@ def test_build_command_passes_audio_recall_options(tmp_path: Path) -> None:
     assert cmd[cmd.index("--audio-split-interval") + 1] == "45"
 
 
+def test_build_command_omits_tuning_flags_when_use_heuristics(tmp_path: Path) -> None:
+    dataset_dir = tmp_path / "dataset"
+    dataset_dir.mkdir()
+    query_csv = tmp_path / "query.csv"
+    query_csv.write_text("q,s,p\nx,y,1\n", encoding="utf-8")
+
+    cfg = HarnessConfig(
+        dataset_dir=str(dataset_dir),
+        dataset_label="jp20",
+        preset="single_gpu",
+        query_csv=str(query_csv),
+        use_heuristics=True,
+    )
+    cmd, _runtime_dir, _detection_file, _effective_query_csv = _build_command(cfg, tmp_path, run_id="r1")
+
+    assert "--pdf-extract-tasks" not in cmd
+    assert "--pdf-extract-cpus-per-task" not in cmd
+    assert "--pdf-extract-batch-size" not in cmd
+    assert "--pdf-split-batch-size" not in cmd
+    assert "--page-elements-batch-size" not in cmd
+    assert "--page-elements-actors" not in cmd
+    assert "--ocr-actors" not in cmd
+    assert "--ocr-batch-size" not in cmd
+    assert "--embed-actors" not in cmd
+    assert "--embed-batch-size" not in cmd
+    assert "--page-elements-cpus-per-actor" not in cmd
+    assert "--ocr-cpus-per-actor" not in cmd
+    assert "--embed-cpus-per-actor" not in cmd
+    assert "--page-elements-gpus-per-actor" not in cmd
+    assert "--ocr-gpus-per-actor" not in cmd
+    assert "--embed-gpus-per-actor" not in cmd
+    # non-tuning flags still present
+    assert "--embed-model-name" in cmd
+    assert "--evaluation-mode" in cmd
+
+
 def test_normalize_recall_metric_key_removes_duplicate_prefix() -> None:
     assert _normalize_recall_metric_key("recall@1") == "recall_1"
     assert _normalize_recall_metric_key("recall@10") == "recall_10"
@@ -565,6 +593,7 @@ def _fake_run_subprocess(_cmd: list[str], metrics) -> int:
             "extract_page_as_image": cfg.extract_page_as_image,
             "extract_infographics": cfg.extract_infographics,
             "write_detection_file": True,
+            "use_heuristics": cfg.use_heuristics,
             "lancedb_uri": str((artifact_dir / "lancedb").resolve()),
             "tuning": {field: getattr(cfg, field) for field in sorted(harness_run.TUNING_FIELDS)},
         },