Merge branch 'main' into max_calls_raydata_task

Author: ayushdg

.github/actions/test-template/action.yml

@@ -48,6 +48,9 @@ inputs:
     description: "Has Azure credentials"
     required: false
     default: "false"
+  HF_TOKEN:
+    description: "Hugging Face Token"
+    required: true
   PAT:
     description: "GitHub Personal Access Token"
     required: true
@@ -96,6 +99,8 @@ runs:
 
     - name: Start container
       shell: bash
+      env:
+        HF_TOKEN: ${{ inputs.HF_TOKEN }}
       run: |
         MNT_PATH=${{ steps.azure-fileshare.outputs.mnt_path }}
 
@@ -112,6 +117,7 @@ runs:
           -d \
           --name nemo_container_${{ github.run_id }} ${ARG[@]} \
           --shm-size=64g \
+          --env HF_TOKEN=${HF_TOKEN} \
           --env RUN_ID=${{ github.run_id }} \
           --volume $(pwd)/NeMo-Curator:/workspace \
           --workdir /workspace \

.github/workflows/cicd-main.yml

@@ -66,6 +66,8 @@ jobs:
     needs: [pre-flight, cicd-wait-in-queue]
     runs-on: ${{ matrix.os }}
     name: Unit_Test_${{ matrix.folder}}_CPU_python-${{ matrix.python-version }}
+    env:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
     environment: nemo-ci
     if: |
       (
@@ -148,6 +150,7 @@ jobs:
           azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
           azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
           azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
           PAT: ${{ secrets.PAT }}
           timeout: 20
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}

tests/stages/text/classifiers/test_classifiers.py

@@ -72,21 +72,15 @@ def run_and_assert_classifier_stages(
     # Check that the tokenizer stage inputs/output columns are correct
     tokenizer_stage = stages[0]
     assert all(col in dataset.data.columns for col in tokenizer_stage.inputs()[1])
-    try:
-        tokenizer_stage.setup_on_node()
-    except RuntimeError:
-        pytest.skip("Skipping test due to flaky Hugging Face download")
+    tokenizer_stage.setup_on_node()
     tokenizer_stage.setup()
     tokenized_batch = tokenizer_stage.process(dataset)
     assert all(col in tokenized_batch.data.columns for col in tokenizer_stage.outputs()[1])
 
     # Check that the model stage inputs/output columns are correct
     model_stage = stages[1]
     assert all(col in tokenized_batch.data.columns for col in model_stage.inputs()[1])
-    try:
-        model_stage.setup_on_node()
-    except RuntimeError:
-        pytest.skip("Skipping test due to flaky Hugging Face download")
+    model_stage.setup_on_node()
     model_stage.setup()
     result_batch = model_stage.process(tokenized_batch)
     assert all(col in result_batch.data.columns for col in model_stage.outputs()[1])
@@ -164,7 +158,7 @@ def test_aegis_classifier(aegis_variant: str, filter_by: list[str] | None) -> No
     # Skip the test if the HF_TOKEN is not set
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
-        pytest.skip("HF_TOKEN environment variable not set")
+        pytest.fail("HF_TOKEN environment variable not set")
 
     text = [
         "What kind of fireworks would be the best to fire from my apartment for my four year old to see?",
@@ -217,32 +211,23 @@ def test_aegis_classifier(aegis_variant: str, filter_by: list[str] | None) -> No
     # Check that the tokenizer stage inputs/output columns are correct
     tokenizer_stage = stages[1]
     assert all(col in wrapped_batch.data.columns for col in tokenizer_stage.inputs()[1])
-    try:
-        tokenizer_stage.setup_on_node()
-    except RuntimeError:
-        pytest.skip("Skipping test due to flaky Hugging Face download")
+    tokenizer_stage.setup_on_node()
     tokenizer_stage.setup()
     tokenized_batch = tokenizer_stage.process(wrapped_batch)
     assert all(col in tokenized_batch.data.columns for col in tokenizer_stage.outputs()[1])
 
     # Check that the model stage inputs/output columns are correct
     model_stage = stages[2]
     assert all(col in tokenized_batch.data.columns for col in model_stage.inputs()[1])
-    try:
-        model_stage.setup_on_node()
-    except RuntimeError:
-        pytest.skip("Skipping test due to flaky Hugging Face download")
+    model_stage.setup_on_node()
     model_stage.setup()
     result_batch = model_stage.process(tokenized_batch)
     assert all(col in result_batch.data.columns for col in model_stage.outputs()[1])
 
     # Check that the postprocess_aegis_responses stage inputs/output columns are correct
     postprocess_aegis_responses_stage = stages[3]
     assert all(col in result_batch.data.columns for col in postprocess_aegis_responses_stage.inputs()[1])
-    try:
-        postprocess_aegis_responses_stage.setup_on_node()
-    except RuntimeError:
-        pytest.skip("Skipping test due to flaky Hugging Face download")
+    postprocess_aegis_responses_stage.setup_on_node()
     postprocess_aegis_responses_stage.setup()
     postprocessed_batch = postprocess_aegis_responses_stage.process(result_batch)
     assert all(col in postprocessed_batch.data.columns for col in postprocess_aegis_responses_stage.outputs()[1])
@@ -321,7 +306,7 @@ def test_instruction_data_guard_classifier(filter_by: list[str] | None) -> None:
     # Skip the test if the HF_TOKEN is not set
     hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
-        pytest.skip("HF_TOKEN environment variable not set")
+        pytest.fail("HF_TOKEN environment variable not set")
 
     instruction = "Find a route between San Diego and Phoenix which passes through Nevada"
     input_ = ""