Merge branch 'main' of https://github.com/NVIDIA/NeMo into lgrigoryan/fix-eval_beamsearch_ngram_ctc

Author: lilithgrigoryan

.coveragerc

@@ -4,9 +4,29 @@ omit =
     /tmp/*
     /home/TestData/*
     /workspace/Megatron-LM/*
-    nemo/collections/nlp/*
     nemo/collections/multimodal/*
+    nemo/collections/multimodal_autoregressive/*
     nemo/collections/vision/*
+    nemo/collections/diffusion/*
+    nemo/collections/nlp/*
+
+    nemo/collections/asr/*
+    nemo/collections/speechlm/*
+    nemo/collections/tts/*
+
+    # omit from audio
+    nemo/collections/audio/data/data_simulation.py
+    nemo/collections/audio/metrics/squim.py
+    nemo/collections/audio/losses/maxine/*
+    nemo/collections/audio/models/maxine/*
+    nemo/collections/audio/parts/utils/maxine.py
+
+    nemo/core/*
+    nemo/collections/common/*
+
+    /workspace/config-3.12.py
+    /workspace/config-3.py
+    /workspace/config.py
 
 [paths]
 source =

.github/CODEOWNERS

@@ -1,4 +1,4 @@
 .github/ @pablo-garay @ko3n1g @thomasdhc @chtruong814
-Dockerfile.ci @pablo-garay @ko3n1g @thomasdhc @chtruong814
+docker/Dockerfile.ci @pablo-garay @ko3n1g @thomasdhc @chtruong814
 .pylintrc.* @pablo-garay @ko3n1g @thomasdhc @chtruong814
-.flake8.* @pablo-garay @ko3n1g @thomasdhc @chtruong814
+.flake8.* @pablo-garay @ko3n1g @thomasdhc @chtruong814

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

@@ -72,11 +72,6 @@ runs:
       shell: bash
       run: |
 
-    - name: Install jq
-      shell: bash
-      run: |
-        curl -sS https://webi.sh/jq | sh
-
     - name: Create UUID
       id: uuid
       shell: bash

.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

@@ -44,6 +44,53 @@ def find_python_files(directory: str) -> List[str]:
 def analyze_imports(nemo_root: str, file_path: str) -> Set[str]:
     """Analyze a Python file and return its NeMo package dependencies using AST parsing."""
     imports = set()
+    visited = set()  # Track visited modules to prevent circular imports
+
+    def get_init_imports(module_path: str, depth: int = 0) -> Dict[str, str]:
+        """Recursively analyze imports from __init__.py files and map them to their final destinations."""
+        # Prevent infinite recursion
+        if depth > 10 or module_path in visited:  # Limit depth to 10 levels
+            return {}
+
+        visited.add(module_path)
+        init_path = os.path.join(module_path, '__init__.py')
+        if not os.path.exists(init_path):
+            return {}
+
+        try:
+            with open(init_path, 'r', encoding='utf-8') as f:
+                init_tree = ast.parse(f.read(), filename=init_path)
+
+            import_map = {}
+            for node in ast.walk(init_tree):
+                if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith('nemo.'):
+                    if node.names:
+                        for name in node.names:
+                            if name.name == '*':
+                                continue
+
+                            # Get the full module path for the import
+                            module_parts = node.module.split('.')
+                            module_dir = os.path.join(nemo_root, *module_parts)
+
+                            # If the imported module has an __init__.py, recursively analyze it
+                            if os.path.exists(os.path.join(module_dir, '__init__.py')):
+                                sub_imports = get_init_imports(module_dir, depth + 1)
+                                if name.name in sub_imports:
+                                    import_map[name.name] = sub_imports[name.name]
+                                else:
+                                    # If not found in sub-imports, it might be from the module itself
+                                    module_file = os.path.join(module_dir, f"{module_parts[-1]}.py")
+                                    if os.path.exists(module_file):
+                                        import_map[name.name] = f"{node.module}.{name.name}"
+                            else:
+                                # Direct module import
+                                import_map[name.name] = f"{node.module}.{name.name}"
+
+            return import_map
+        except Exception as e:
+            print(f"Error analyzing {init_path}: {e}")
+            return {}
 
     try:
         with open(file_path, 'r', encoding='utf-8') as f:
@@ -68,14 +115,31 @@ def analyze_imports(nemo_root: str, file_path: str) -> Set[str]:
                                 if name.name == '*':
                                     continue
 
-                                imports.add(f"{node.module}.{name.name}")
+                                # Check if this is an __init__ import
+                                module_path = os.path.join(nemo_root, *parts)
+                                init_imports = get_init_imports(module_path)
+
+                                if name.name in init_imports:
+                                    # Use the mapped import path
+                                    imports.add(init_imports[name.name])
+                                else:
+                                    imports.add(f"{node.module}.{name.name}")
 
                     elif module_type in find_top_level_packages(nemo_root):
                         if node.names:
                             for name in node.names:
                                 if name.name == '*':
                                     continue
-                                imports.add(f"{node.module}.{name.name}")
+
+                                # Check if this is an __init__ import
+                                module_path = os.path.join(nemo_root, *parts)
+                                init_imports = get_init_imports(module_path)
+
+                                if name.name in init_imports:
+                                    # Use the mapped import path
+                                    imports.add(init_imports[name.name])
+                                else:
+                                    imports.add(f"{node.module}.{name.name}")
 
     except Exception as e:
         print(f"Error analyzing {file_path}: {e}")
@@ -87,12 +151,16 @@ def find_top_level_packages(nemo_root: str) -> List[str]:
     """Find all top-level packages under nemo directory."""
     packages: List[str] = []
     nemo_dir = os.path.join(nemo_root, 'nemo')
+    tests_dir = os.path.join(nemo_root, 'tests')
 
     if not os.path.exists(nemo_dir):
         print(f"Warning: nemo directory not found at {nemo_dir}")
         return packages
+    if not os.path.exists(tests_dir):
+        print(f"Warning: nemo directory not found at {nemo_dir}")
+        return packages
 
-    for item in os.listdir(nemo_dir):
+    for item in os.listdir(nemo_dir) + os.listdir(tests_dir):
         item_path = os.path.join(nemo_dir, item)
         if os.path.isdir(item_path) and not item.startswith('__'):
             packages.append(item)
@@ -125,17 +193,19 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
 
     dependencies: Dict[str, List[str]] = {}
 
-    # Second pass: analyze imports and build reverse dependencies
     for file_path in find_python_files(nemo_root):
         relative_path = os.path.relpath(file_path, nemo_root)
+
         parts = relative_path.split(os.sep)
 
-        if len(parts) == 1 or parts[-1] == "__init__.py" or parts[0] != "nemo":
+        if len(parts) == 1 or (parts[0] != "nemo" and parts[0] != "tests"):
             continue
 
         module_path = relative_path.replace(".py", "").replace("/", ".")
-        if parts[1] in top_level_packages and parts[1] != 'collections':
+        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.replace("/", ".").replace(".py", "")]
         elif parts[1] == 'collections':
             dependencies[module_path] = list(set(analyze_imports(nemo_root, file_path)))
 
@@ -181,8 +251,10 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
     simplified_dependencies: Dict[str, List[str]] = {}
     for package, deps in dependencies.items():
         package_parts = package.split('.')
-        print(f"{os.path.join(*package_parts[:-1])}.py")
-        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
@@ -197,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)
 
@@ -218,22 +291,90 @@ def build_dependency_graph(nemo_root: str) -> Dict[str, List[str]]:
     for package, deps in dependencies.items():
         new_deps = []
         for dep in deps:
-            if "asr" in dep or "tts" in dep or "speechlm" in dep or "audio" in dep:
+            if (
+                "nemo.collections.asr" in dep
+                or "nemo.collections.tts" in dep
+                or "nemo.collections.speechlm" in dep
+                or "nemo.collections.audio" in dep
+                or "tests.collections.asr" in dep
+                or "tests.collections.tts" in dep
+                or "tests.collections.speechlm" in dep
+                or "tests.collections.audio" in dep
+            ):
                 new_deps.append("speech")
+                new_deps.append("unit-tests")
 
-            if "export" in dep or "deploy" in dep:
+            if "nemo.export" in dep or "nemo.deploy" in dep or "tests.export" in dep or "tests.deploy" in dep:
                 new_deps.append("export-deploy")
+                new_deps.append("unit-tests")
 
-            if "llm" in dep or "vlm" in dep or "automodel" in dep:
+            if (
+                "nemo.collections.llm" in dep
+                or "nemo.collections.vlm" in dep
+                or "nemo.automodel" in dep
+                or "tests.collections.llm" in dep
+                or "tests.collections.vlm" in dep
+                or "tests.automodel" in dep
+            ):
                 new_deps.append("automodel")
+                new_deps.append("unit-tests")
+
+            if "tests" in dep and "tests.functional_tests" not in dep:
+                new_deps.append("unit-tests")
 
-            if "collections" in dep and not ("asr" in dep or "tts" in dep or "speechlm" in dep or "audio" in dep):
+            if (
+                "nemo.collections" in dep
+                and "nemo.collections.asr" not in dep
+                and "nemo.collections.tts" not in dep
+                and "nemo.collections.speechlm" not in dep
+                and "nemo.collections.audio" not in dep
+                and "tests.collections.asr" not in dep
+                and "tests.collections.tts" not in dep
+                and "tests.collections.speechlm" not in dep
+                and "tests.collections.audio" not in dep
+            ):
                 new_deps.append("nemo2")
+                new_deps.append("unit-tests")
 
         bucket_deps[package] = sorted(list(set(new_deps)))
 
     dependencies = bucket_deps
 
+    # Additional dependencies
+    # Add all files in requirements/ directory
+    requirements_dir = os.path.join(nemo_root, "requirements")
+    if os.path.exists(requirements_dir):
+        for filename in os.listdir(requirements_dir):
+            filepath = os.path.join("requirements", filename)
+            relative_path = os.path.relpath(filepath, nemo_root)
+
+            dependencies[relative_path] = [
+                "nemo2",
+                "unit-tests",
+                "speech",
+                "automodel",
+                "export-deploy",
+            ]
+
+    # Add all Dockerfile files
+    for root, _, files in os.walk(nemo_root):
+        for file_path in files:
+            full_path = os.path.join(root, file_path)
+            relative_path = os.path.relpath(full_path, nemo_root)
+
+            if "cicd-main-export-deploy" in file_path:
+                dependencies[relative_path] = ["export-deploy"]
+            if "cicd-main-nemo2" in file_path:
+                dependencies[relative_path] = ["nemo2"]
+            if "cicd-main-speech" in file_path:
+                dependencies[relative_path] = ["speech"]
+            if "cicd-main-automodel" in file_path:
+                dependencies[relative_path] = ["automodel"]
+            if "cicd-main-unit-tests" in file_path:
+                dependencies[relative_path] = ["unit-tests"]
+            if "Dockerfile" in file_path:
+                dependencies[relative_path] = ["nemo2", "unit-tests", "speech", "automodel", "export-deploy"]
+
     # Sort dependencies by length of values (number of dependencies)
     dependencies = dict(sorted(dependencies.items(), key=lambda x: len(x[1]), reverse=True))
 
@@ -250,7 +391,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/_build_container.yml

@@ -68,7 +68,7 @@ jobs:
           echo "EOF" >> $GITHUB_OUTPUT
 
   build:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/[email protected]
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     needs: [pre-flight]
     with:
       image-name: ${{ inputs.image-name }}

.github/workflows/_bump_mcore_tag.yml

@@ -42,7 +42,7 @@ jobs:
           git push -u origin $SOURCE_BRANCH
 
   mcore:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/[email protected]
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     needs: [update-branch]
     with:
       source-repository: NVIDIA/Megatron-LM

.github/workflows/build-test-publish-wheel.yml

@@ -26,7 +26,7 @@ defaults:
 
 jobs:
   build-test-publish-wheel:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/[email protected]
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     with:
       dry-run: true
       python-package: nemo

.github/workflows/cherry-pick-release-commit.yml

@@ -7,7 +7,7 @@ on:
 
 jobs:
   cherry-pick:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/[email protected]
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     secrets:
       PAT: ${{ secrets.PAT }}
       SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }}