Merge remote-tracking branch 'origin/main' into chtruong/re-add-libsox

Signed-off-by: Charlie Truong <[email protected]>

Author: chtruong814

.github/scripts/components_to_run.py

@@ -76,9 +76,6 @@ def main(source_sha: str, target_sha: str):
 
     test_modules = list(set(test_modules))
 
-    if len(test_modules) == 0:
-        test_modules = ["nemo2", "automodel", "export-deploy", "speech"]
-
     with open("test_modules.json", "w", encoding="utf-8") as f:
         json.dump(test_modules, f)
 

.github/scripts/nemo_dependencies.py

@@ -205,7 +205,7 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
         if parts[1] in top_level_packages and parts[1] != 'collections' and parts[0] != 'tests':
             dependencies[module_path] = list(set(analyze_imports(nemo_root, file_path)))
         elif parts[0] == 'tests':
-            dependencies[module_path] = [relative_path]
+            dependencies[module_path] = [relative_path.replace("/", ".").replace(".py", "")]
         elif parts[1] == 'collections':
             dependencies[module_path] = list(set(analyze_imports(nemo_root, file_path)))
 
@@ -252,7 +252,9 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
     for package, deps in dependencies.items():
         package_parts = package.split('.')
 
-        if os.path.isfile((file_path := f"{os.path.join(*package_parts[:-1])}.py")):
+        if package_parts[0] == "tests":
+            simplified_package_path = f"{os.path.join(*package_parts)}.py"
+        elif os.path.isfile((file_path := f"{os.path.join(*package_parts[:-1])}.py")):
             simplified_package_path = file_path
         elif os.path.isdir((file_path := f"{os.path.join(*package_parts[:-1])}")):
             simplified_package_path = file_path
@@ -267,13 +269,14 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
 
             if (
                 len(dep_parts) >= 2
-                and dep_parts[1] in find_top_level_packages(nemo_root)
+                and (dep_parts[1] in find_top_level_packages(nemo_root))
                 and dep_parts[1] != 'collections'
             ):
                 simplified_dependencies[simplified_package_path].append(f"{dep_parts[0]}.{dep_parts[1]}")
-
+            elif dep_parts[0] == "tests":
+                simplified_dependencies[simplified_package_path].append(".".join(dep_parts))
             elif len(dep_parts) >= 3 and (
-                simplified_name := f"{dep_parts[0]}.{dep_parts[1]}.{dep_parts[2]}"
+                simplified_name := f"nemo.{dep_parts[1]}.{dep_parts[2]}"
             ) in find_collection_modules(nemo_root):
                 simplified_dependencies[simplified_package_path].append(simplified_name)
 
@@ -353,7 +356,7 @@ def main():
 
     # Output as JSON
     data = json.dumps(dependencies, indent=4)
-    # print(data)
+
     with open('nemo_dependencies.json', 'w', encoding='utf-8') as f:
         f.write(data)
 

.github/workflows/cicd-main-automodel.yml

@@ -57,6 +57,7 @@ jobs:
           tests_to_run: ${{ inputs.test_to_run }}
           image: ${{ inputs.image-name }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}
 
   e2e-tests:
     strategy:

.github/workflows/cicd-main-nemo2.yml

@@ -47,6 +47,8 @@ jobs:
             runner: self-hosted-azure
           - script: L2_NeMo_2_llama3_straggler_detection
             runner: self-hosted-azure
+          - script: L2_NeMo_2_llama3_local_ckpt
+            runner: self-hosted-azure
           - script: L2_NeMo_2_GPT_DDP_Param_Parity_check
             runner: self-hosted-azure
           - script: L2_NeMo_2_Hyena_Conversion_from_HF
@@ -240,9 +242,20 @@ jobs:
             script: L2_NeMo_2_LLAMA4_MOCK_FINETUNE_CP2
           - runner: self-hosted-azure
             script: L2_NeMo_2_LLAMA4_ENERGON_FINETUNE_EP2
+          - runner: self-hosted-azure
+            script: L2_NeMo_2_Diffusion_Recipe_Test
+          - runner: self-hosted-azure
+            script: L2_NeMo_2_Diffusion_Taskencoder_Test
+          - runner: self-hosted-azure
+            script: L2_NeMo_2_Flux_Import_Test
+          - runner: self-hosted-azure
+            script: L2_NeMo_2_Flux_Inference_Test
+          - runner: self-hosted-azure
+            script: L2_NeMo_2_Flux_Training_Test
+          
     needs: [build]
     runs-on: ${{ matrix.runner }}
-    name: ${{ matrix.script }}
+    name: ${{ matrix.is_optional && 'PLEASEFIXME_' || '' }}${{ matrix.script }}
     steps:
       - name: Checkout
         uses: actions/checkout@v4

.github/workflows/cicd-main-speech.yml

@@ -70,6 +70,7 @@ jobs:
           image: ${{ inputs.image-name }}
           timeout: ${{ matrix.timeout || 10 }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}
 
   e2e-tests:
     strategy:

.github/workflows/cicd-main-unit-tests.yml

@@ -71,6 +71,7 @@ jobs:
           is_unit_test: true
           tests_to_run: ${{ inputs.test_to_run }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}
 
   collections-multimodal-tests:
     strategy:
@@ -97,7 +98,7 @@ jobs:
           is_unit_test: true
           tests_to_run: ${{ inputs.test_to_run }}
           cpu-only: ${{ matrix.cpu-only || false }}
-
+          is_optional: ${{ matrix.is-optional || false }}
   collections-vlm-tests:
     strategy:
       fail-fast: false
@@ -123,6 +124,7 @@ jobs:
           is_unit_test: true
           tests_to_run: ${{ inputs.test_to_run }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}
 
   core-tests:
     strategy:
@@ -180,6 +182,7 @@ jobs:
           is_unit_test: true
           tests_to_run: ${{ inputs.test_to_run }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}
 
   other-tests:
     strategy:
@@ -206,3 +209,4 @@ jobs:
           is_unit_test: true
           tests_to_run: ${{ inputs.test_to_run }}
           cpu-only: ${{ matrix.cpu-only || false }}
+          is_optional: ${{ matrix.is-optional || false }}

.github/workflows/cicd-main.yml

@@ -21,9 +21,6 @@ on:
       - r**
       - weekly-bump*
     types: [labeled]
-  push:
-    branches:
-      - main
   workflow_dispatch:
     inputs:
       test_to_run:
@@ -134,11 +131,12 @@ jobs:
     uses: ./.github/workflows/code-linting.yml
 
   cicd-wait-in-queue:
-    needs: [pre-flight]
+    needs: [pre-flight, code-linting]
     runs-on: ubuntu-latest
     environment: test
     if: |
       needs.pre-flight.outputs.test_to_run != '[]' 
+      && needs.pre-flight.outputs.components_to_run != '[]'
       && needs.pre-flight.outputs.is_ci_workload == 'false'
     steps:
       - name: Running CI tests
@@ -150,6 +148,7 @@ jobs:
     needs: [pre-flight, code-linting, cicd-wait-in-queue]
     if: |
       needs.pre-flight.outputs.test_to_run != '[]' 
+      && needs.pre-flight.outputs.components_to_run != '[]'
       && (
         success() 
         || (
@@ -165,6 +164,7 @@ jobs:
   cicd-import-tests:
     if: |
       needs.pre-flight.outputs.test_to_run != '[]'
+      && needs.pre-flight.outputs.components_to_run != '[]'
       && (
         success()
         || (
@@ -214,6 +214,7 @@ jobs:
     runs-on: self-hosted-azure
     if: |
       needs.pre-flight.outputs.test_to_run != '[]'
+      && needs.pre-flight.outputs.components_to_run != '[]'
       && (
         success()
         || (
@@ -240,6 +241,7 @@ jobs:
     uses: ./.github/workflows/cicd-main-unit-tests.yml
     if: |
       needs.pre-flight.outputs.test_to_run != '[]'
+      && needs.pre-flight.outputs.components_to_run != '[]'
       && (
         success()
         || (
@@ -360,12 +362,14 @@ jobs:
           GH_TOKEN: ${{ github.token }}
           RUN_ID: ${{ github.run_id }}
           HAS_LABEL: ${{ github.event.label.name == 'Run CICD' }}
+          IS_SCHEDULED: ${{ github.event_name == 'schedule' }}
         run: |
           # Get workflow run details and check job conclusions
+          LATEST_ATTEMPT=$(gh run view $RUN_ID --json jobs -q '[.jobs[] | select(.conclusion != null) | .conclusion] | last')
           NUM_FAILED=$(gh run view $RUN_ID --json jobs -q '[.jobs[] | select(.conclusion == "failure") | .name] | length')
           NUM_CANCELLED=$(gh run view $RUN_ID --json jobs -q '[.jobs[] | select(.conclusion == "cancelled") | .name] | length')
 
-          if [[ $NUM_FAILED -eq 0 && $NUM_CANCELLED -eq 0 && "$HAS_LABEL" == "true" ]]; then
+          if [[ $NUM_FAILED -eq 0 && $NUM_CANCELLED -eq 0 && ("$HAS_LABEL" == "true" || "$IS_SCHEDULED" == "true") ]]; then
             RESULT="success"
           elif [[ $NUM_CANCELLED -gt 0 ]]; then
             RESULT="cancelled"
@@ -445,9 +449,11 @@ jobs:
 
   Coverage:
     runs-on: ubuntu-latest
-    needs: [Nemo_CICD_Test]
+    needs: [pre-flight, Nemo_CICD_Test]
     if: |
-      (
+      needs.pre-flight.outputs.test_to_run != '[]'
+      && needs.pre-flight.outputs.components_to_run != '[]'
+      && (
         success() 
         || needs.Nemo_CICD_Test.result == 'success'
       )

Dockerfile.ci

@@ -16,8 +16,8 @@
 
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.07-py3
 
-FROM ${BASE_IMAGE}
-
+FROM ${BASE_IMAGE} as base-image
+ENV PIP_CONSTRAINT="" 
 ENV TRANSFORMERS_OFFLINE=0
 ENV HYDRA_FULL_ERROR=1
 ENV PYTHONUNBUFFERED=1
@@ -29,6 +29,17 @@ apt-get install -y bc libsox-fmt-all
 apt-get clean
 EOF
 
+FROM base-image as te-wheel
+ARG TE_REPO
+ARG TE_TAG
+RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh \
+  --mount=type=bind,source=external/patches,target=/tmp/NeMo/external/patches <<"EOF" bash -ex
+  
+  bash /tmp/NeMo/reinstall.sh --library te --mode build
+  ls -al /tmp/Megatron-LM || true
+EOF
+
+FROM base-image as mcore-wheel
 ARG MLM_REPO
 ARG MLM_TAG
 RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
@@ -37,8 +48,11 @@ RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF"
   ls -al /tmp/Megatron-LM || true
 EOF
 
+FROM base-image 
 WORKDIR /workspace
-RUN \
+RUN \  
+  --mount=type=bind,from=te-wheel,source=/opt/wheels/te,target=/opt/wheels/te \
+  --mount=type=bind,from=mcore-wheel,source=/opt/wheels/mcore,target=/opt/wheels/mcore \
   --mount=type=bind,source=requirements,target=/tmp/NeMo/requirements \
   --mount=type=bind,source=tools/ctc_segmentation/requirements.txt,target=/tmp/NeMo/tools/ctc_segmentation/requirements.txt \
   --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh \
@@ -48,6 +62,7 @@ RUN \
   --mount=type=bind,source=nemo/package_info.py,target=/tmp/NeMo/nemo/package_info.py \
   --mount=type=bind,source=nemo/__init__.py,target=/tmp/NeMo/nemo/__init__.py <<"EOF" bash -ex
 
+  bash /tmp/NeMo/reinstall.sh --library te --mode install
   bash /tmp/NeMo/reinstall.sh --library mcore --mode install
   bash /tmp/NeMo/reinstall.sh --library nemo --mode install
   rm -rf $NEMO_DIR || true

Dockerfile.ci.export_deploy

@@ -16,7 +16,8 @@
 
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.07-py3
 
-FROM ${BASE_IMAGE} AS trt-base
+FROM ${BASE_IMAGE} AS base-image
+ENV PIP_CONSTRAINT="" 
 ARG IMAGE_LABEL
 LABEL "nemo.library"=${IMAGE_LABEL}
 
@@ -39,18 +40,25 @@ RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF"
 EOF
 
 
-FROM trt-base AS trt-llm-wheel
-
+FROM base-image AS trt-llm-wheel
 ARG TRTLLM_REPO
 ARG TRTLLM_TAG
 RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
 
   bash /tmp/NeMo/reinstall.sh --library trtllm --mode build
 EOF
 
+FROM base-image as te-wheel
+ARG TE_REPO
+ARG TE_TAG
+RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh \
+  --mount=type=bind,source=external/patches,target=/tmp/NeMo/external/patches <<"EOF" bash -ex
+  
+  bash /tmp/NeMo/reinstall.sh --library te --mode build
+  ls -al /tmp/Megatron-LM || true
+EOF
 
-FROM trt-base as final
-
+FROM base-image as mcore-wheel
 ARG MLM_REPO
 ARG MLM_TAG
 RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF" bash -ex
@@ -59,9 +67,13 @@ RUN --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh <<"EOF"
   ls -al /tmp/Megatron-LM || true
 EOF
 
+
+FROM base-image
 WORKDIR /workspace
 RUN \
   --mount=type=bind,from=trt-llm-wheel,source=/opt/wheels/trtllm,target=/opt/wheels/trtllm \
+  --mount=type=bind,from=te-wheel,source=/opt/wheels/te,target=/opt/wheels/te \
+  --mount=type=bind,from=mcore-wheel,source=/opt/wheels/mcore,target=/opt/wheels/mcore \
   --mount=type=bind,source=requirements,target=/tmp/NeMo/requirements \
   --mount=type=bind,source=tools/ctc_segmentation/requirements.txt,target=/tmp/NeMo/tools/ctc_segmentation/requirements.txt \
   --mount=type=bind,source=reinstall.sh,target=/tmp/NeMo/reinstall.sh \

docs/links_needing_review.json

@@ -102,14 +102,6 @@
   "uri": "https://docs.nvidia.com/nemo-framework/user-guide/latest/multimodalmodels/index.html",
   "info": "404 Client Error: Not Found for url: https://docs.nvidia.com/nemo-framework/user-guide/latest/multimodalmodels/index.html"
 }
-{
-  "filename": "core/api.rst",
-  "lineno": 2,
-  "status": "broken",
-  "code": 0,
-  "uri": "https://github.com/Lightning-AI/lightning/blob/master/src/pytorch_lightning/core/mixins/device_dtype_mixin.py#L113",
-  "info": "404 Client Error: Not Found for url: https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/pytorch_lightning/core/mixins/device_dtype_mixin.py"
-}
 {
   "filename": "multimodal/vlm/clip.rst",
   "lineno": 140,
@@ -302,14 +294,6 @@
   "uri": "https://github.com/webdataset/webdataset#multinode-training",
   "info": "Anchor 'multinode-training' not found"
 }
-{
-  "filename": "core/api.rst",
-  "lineno": 6,
-  "status": "broken",
-  "code": 0,
-  "uri": "https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py#L728",
-  "info": "Anchor 'L728' not found"
-}
 {
   "filename": "nlp/question_answering.rst",
   "lineno": 196,