Replace shell matrix script with robust Python solution

Author: PointKernel

.github/actions/compute-matrix/compute-matrix.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import yaml
+import sys
+import os
+from typing import Dict, List, Any
+
+def write_output(key: str, value: str):
+    """Write GitHub Actions output."""
+    github_output = os.environ.get('GITHUB_OUTPUT', '/dev/null')
+    with open(github_output, 'a') as f:
+        f.write(f"{key}={value}\n")
+    print(f"{key}={value}")
+
+def write_json_file(filename: str, data: Any):
+    """Write data to a JSON file."""
+    with open(filename, 'w') as f:
+        json.dump(data, f, indent=2)
+
+def write_text_file(filename: str, content: str):
+    """Write content to a text file."""
+    with open(filename, 'w') as f:
+        f.write(content)
+
+def explode_std_versions(matrix_entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Explode std arrays into individual entries."""
+    result = []
+    for entry in matrix_entries:
+        if 'std' in entry and isinstance(entry['std'], list):
+            for std in entry['std']:
+                new_entry = entry.copy()
+                new_entry['std'] = std
+                result.append(new_entry)
+        else:
+            result.append(entry)
+    return result
+
+def extract_matrix(file_path: str, matrix_type: str):
+    """Extract and process the matrix configuration."""
+    try:
+        with open(file_path, 'r') as f:
+            data = yaml.safe_load(f)
+        
+        if matrix_type not in data:
+            print(f"Error: Matrix type '{matrix_type}' not found in {file_path}", file=sys.stderr)
+            sys.exit(1)
+        
+        matrix = data[matrix_type]
+        
+        # Write devcontainer version
+        devcontainer_version = data.get('devcontainer_version', '25.08')
+        write_output("DEVCONTAINER_VERSION", devcontainer_version)
+        
+        # Process nvcc matrix
+        if 'nvcc' not in matrix:
+            print(f"Error: 'nvcc' section not found in {matrix_type} matrix", file=sys.stderr)
+            sys.exit(1)
+        
+        nvcc_matrix = matrix['nvcc']
+        nvcc_full_matrix = explode_std_versions(nvcc_matrix)
+        
+        write_output("NVCC_FULL_MATRIX", json.dumps(nvcc_full_matrix))
+        
+        # Extract unique CUDA versions
+        cuda_versions = list(set(entry['cuda'] for entry in nvcc_full_matrix))
+        cuda_versions.sort()
+        write_output("CUDA_VERSIONS", json.dumps(cuda_versions))
+        
+        # Extract unique host compilers
+        host_compilers = list(set(entry['compiler']['name'] for entry in nvcc_full_matrix))
+        host_compilers.sort()
+        write_output("HOST_COMPILERS", json.dumps(host_compilers))
+        
+        # Create per-cuda-compiler matrix
+        per_cuda_compiler = {}
+        for entry in nvcc_full_matrix:
+            key = f"{entry['cuda']}-{entry['compiler']['name']}"
+            if key not in per_cuda_compiler:
+                per_cuda_compiler[key] = []
+            per_cuda_compiler[key].append(entry)
+        
+        write_output("PER_CUDA_COMPILER_MATRIX", json.dumps(per_cuda_compiler))
+        
+        # Create output directory and write detailed files (CCCL approach)
+        os.makedirs("workflow", exist_ok=True)
+        
+        # Write individual output files for debugging and artifacts
+        write_json_file("workflow/devcontainer_version.json", {"version": devcontainer_version})
+        write_json_file("workflow/nvcc_full_matrix.json", nvcc_full_matrix)
+        write_json_file("workflow/cuda_versions.json", cuda_versions)
+        write_json_file("workflow/host_compilers.json", host_compilers)
+        write_json_file("workflow/per_cuda_compiler_matrix.json", per_cuda_compiler)
+        
+        # Write summary
+        summary = {
+            "total_matrix_entries": len(nvcc_full_matrix),
+            "cuda_compiler_combinations": len(per_cuda_compiler),
+            "cuda_versions": cuda_versions,
+            "host_compilers": host_compilers
+        }
+        write_json_file("workflow/matrix_summary.json", summary)
+        
+        # Write human-readable summary
+        summary_text = f"Matrix Summary:\n"
+        summary_text += f"  Total matrix entries: {len(nvcc_full_matrix)}\n"
+        summary_text += f"  CUDA versions: {', '.join(cuda_versions)}\n"
+        summary_text += f"  Host compilers: {', '.join(host_compilers)}\n"
+        summary_text += f"  CUDA-compiler combinations: {len(per_cuda_compiler)}\n\n"
+        summary_text += "Combinations:\n"
+        for key, entries in per_cuda_compiler.items():
+            summary_text += f"  {key}: {len(entries)} entries\n"
+        
+        write_text_file("workflow/matrix_summary.txt", summary_text)
+        
+        print(f"Successfully processed {len(nvcc_full_matrix)} matrix entries", file=sys.stderr)
+        print(f"Generated {len(per_cuda_compiler)} cuda-compiler combinations", file=sys.stderr)
+        print("Matrix data written to workflow/ directory", file=sys.stderr)
+        
+    except FileNotFoundError:
+        print(f"Error: Matrix file '{file_path}' not found", file=sys.stderr)
+        sys.exit(1)
+    except yaml.YAMLError as e:
+        print(f"Error parsing YAML file '{file_path}': {e}", file=sys.stderr)
+        sys.exit(1)
+    except KeyError as e:
+        print(f"Error: Missing required key in matrix file: {e}", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"Unexpected error processing matrix: {e}", file=sys.stderr)
+        sys.exit(1)
+
+def main():
+    if len(sys.argv) != 3:
+        print("Usage: compute-matrix.py MATRIX_FILE MATRIX_TYPE", file=sys.stderr)
+        print("  MATRIX_FILE   : The path to the matrix file.", file=sys.stderr)
+        print("  MATRIX_TYPE   : The desired matrix. Supported values: 'pull_request'", file=sys.stderr)
+        sys.exit(1)
+    
+    matrix_file = sys.argv[1]
+    matrix_type = sys.argv[2]
+    
+    if matrix_type != "pull_request":
+        print(f"Error: Unsupported matrix type '{matrix_type}'. Only 'pull_request' is supported.", file=sys.stderr)
+        sys.exit(1)
+    
+    print(f"Input matrix file: {matrix_file}", file=sys.stderr)
+    print(f"Matrix Type: {matrix_type}", file=sys.stderr)
+    
+    # Show matrix file content for debugging
+    try:
+        with open(matrix_file, 'r') as f:
+            content = f.read()
+            print("Matrix file content:", file=sys.stderr)
+            print(content, file=sys.stderr)
+            print("=" * 50, file=sys.stderr)
+    except Exception as e:
+        print(f"Warning: Could not read matrix file for debugging: {e}", file=sys.stderr)
+    
+    extract_matrix(matrix_file, matrix_type)
+
+if __name__ == "__main__":
+    main() 

.github/actions/compute-matrix/compute-matrix.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

.github/workflows/pr.yml

@@ -62,23 +62,28 @@ jobs:
     steps:
       - name: Checkout repo
         uses: actions/checkout@v3
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.x'
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install PyYAML
       - name: Compute matrix outputs
         id: set-outputs
         run: |
-          echo "=== Debug Info ==="
-          echo "yq version: $(yq --version)"
-          echo "jq version: $(jq --version)"
-          echo ""
-          echo "=== Matrix File Content ==="
-          cat ci/matrix.yml
-          echo ""
-          echo "=== Running compute-matrix script ==="
-          .github/actions/compute-matrix/compute-matrix.sh ci/matrix.yml pull_request
+          echo "=== Running Python matrix computation ==="
+          python .github/actions/compute-matrix/compute-matrix.py ci/matrix.yml pull_request
           echo ""
-          echo "=== Output Variables ==="
-          echo "CUDA_VERSIONS: ${{ steps.set-outputs.outputs.CUDA_VERSIONS }}"
-          echo "HOST_COMPILERS: ${{ steps.set-outputs.outputs.HOST_COMPILERS }}"
-          echo "PER_CUDA_COMPILER_MATRIX: ${{ steps.set-outputs.outputs.PER_CUDA_COMPILER_MATRIX }}"
+          echo "=== Matrix Summary ==="
+          cat workflow/matrix_summary.txt
+      - name: Upload matrix artifacts
+        uses: actions/upload-artifact@v4
+        with:
+                     name: matrix-data
+           path: workflow/
+           compression-level: 0
 
   ci:
     name: CUDA${{ matrix.cuda_version }} ${{ matrix.compiler }}