fix coverage

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/test_biencoder_checkpoint_restoration.py

@@ -80,8 +80,11 @@ def _run_training(recipe_yaml: str, checkpoint_dir: str) -> Path:
     """Launch biencoder training as a subprocess and return the checkpoint dir."""
     cmd = [
         sys.executable,
-        "-m",
-        "nemo_automodel.recipes.biencoder.train_biencoder",
+        "-m", "coverage", "run",
+        "--data-file=/workspace/.coverage",
+        "--source=/workspace/",
+        "--parallel-mode",
+        "-m", "nemo_automodel.recipes.biencoder.train_biencoder",
         "--config",
         recipe_yaml,
     ]

tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/test_biencoder_finetuning_inline.sh

@@ -15,29 +15,33 @@
 
 set -exo pipefail
 
+COVERAGE_ARGS="--data-file=/workspace/.coverage --source=/workspace/ --parallel-mode"
+
 # Run the biencoder recipe (uses nemo_automodel/recipes/biencoder/train_biencoder.py via module entrypoint).
-python3 -m nemo_automodel.recipes.biencoder.train_biencoder \
+python3 -m coverage run ${COVERAGE_ARGS} \
+    -m nemo_automodel.recipes.biencoder.train_biencoder \
     --config \
     tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/recipe.yaml \
     --model.pretrained_model_name_or_path $TEST_DATA_DIR/llama-nemotron-embed-1b-v2/ \
     --tokenizer.pretrained_model_name_or_path $TEST_DATA_DIR/llama-nemotron-embed-1b-v2/ \
     --dataloader.dataset.data_dir_list $TEST_DATA_DIR/embedding_testdata/training.jsonl \
 
 # Compare baseline vs finetuned biencoder checkpoint (pos-neg separation should not degrade).
-python3 \
+python3 -m coverage run --append ${COVERAGE_ARGS} \
     tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/compare_biencoder_models.py \
     $TEST_DATA_DIR/llama-nemotron-embed-1b-v2 \
     /workspace/output/biencoder_inline/checkpoints/epoch_0_step_31/ \
     $TEST_DATA_DIR/embedding_testdata/testing.jsonl \
     true
 
-# ---- Checkpoint restoration tests ----
+# Checkpoint restoration tests
 # Test 1: Full-model checkpoint restoration (NeMo -> save -> transformers load)
 # Test 2: PEFT (LoRA) checkpoint restoration (NeMo -> save -> transformers + safetensors load)
 BASE_MODEL_PATH=$TEST_DATA_DIR/llama-nemotron-embed-1b-v2 \
 CHECKPOINT_DIR=/workspace/output/biencoder_ckpt_restore/checkpoints \
 PEFT_CHECKPOINT_DIR=/workspace/output/biencoder_ckpt_restore_peft/checkpoints \
 RECIPE_YAML=tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/recipe_ckpt_restore.yaml \
 PEFT_RECIPE_YAML=tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/recipe_peft.yaml \
-python3 -m pytest -xvs \
+python3 -m coverage run --append ${COVERAGE_ARGS} \
+    -m pytest -xvs \
     tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/test_biencoder_checkpoint_restoration.py

tests/functional_tests/llm_pretrain_and_kd/customizer_retrieval/test_customizer_retrieval.py

@@ -73,7 +73,11 @@ def _run_training() -> Path:
     """Launch the biencoder training recipe as a subprocess and return the
     checkpoint directory produced by the run."""
     cmd = [
-        sys.executable, "-m", "nemo_automodel.recipes.biencoder.train_biencoder",
+        sys.executable, "-m", "coverage", "run",
+        "--data-file=/workspace/.coverage",
+        "--source=/workspace/",
+        "--parallel-mode",
+        "-m", "nemo_automodel.recipes.biencoder.train_biencoder",
         "--config", RECIPE_YAML,
     ]
     result = subprocess.run(cmd, cwd=str(_REPO_ROOT), check=True)