Add JSON output to benchmarks and upload as artifacts

- Update benchmark-hardware.py to save results to JSON
- Update benchmark-jupyter.ipynb to save results to JSON
- Update benchmark-jupyterbook.md to save results to JSON
- Add CI step to collect and display benchmark results
- Add CI step to upload benchmark results as artifact

Author: mmcky

.github/workflows/ci.yml

@@ -31,12 +31,17 @@ jobs:
         run: |
           echo "=== Bare Metal Python Script Execution ==="
           python scripts/benchmark-hardware.py
+          mkdir -p benchmark_results
+          mv benchmark_results_baremetal.json benchmark_results/
       - name: Run Jupyter Notebook Benchmark (via nbconvert)
         shell: bash -l {0}
         run: |
           echo "=== Jupyter Kernel Execution ==="
-          jupyter nbconvert --to notebook --execute scripts/benchmark-jupyter.ipynb --output benchmark-jupyter-executed.ipynb
+          cd scripts
+          jupyter nbconvert --to notebook --execute benchmark-jupyter.ipynb --output benchmark-jupyter-executed.ipynb
           echo "Notebook executed successfully"
+          cd ..
+          mv scripts/benchmark_results_jupyter.json benchmark_results/
       - name: Run Jupyter-Book Benchmark
         shell: bash -l {0}
         run: |
@@ -54,6 +59,26 @@ jobs:
           # Build
           jb build benchmark_test --path-output benchmark_build/
           echo "Jupyter-Book build completed successfully"
+          # Move JSON results if generated
+          if [ -f benchmark_test/benchmark_results_jupyterbook.json ]; then
+            mv benchmark_test/benchmark_results_jupyterbook.json benchmark_results/
+          elif [ -f benchmark_results_jupyterbook.json ]; then
+            mv benchmark_results_jupyterbook.json benchmark_results/
+          fi
+      - name: Collect and Upload Benchmark Results
+        uses: actions/upload-artifact@v5
+        with:
+          name: benchmark-results
+          path: benchmark_results/
+      - name: Display Benchmark Results
+        shell: bash -l {0}
+        run: |
+          echo "=== Benchmark Results Summary ==="
+          for f in benchmark_results/*.json; do
+            echo "--- $f ---"
+            cat "$f"
+            echo ""
+          done
       - name: Install latex dependencies
         run: |
           sudo apt-get -qq update

scripts/benchmark-hardware.py

@@ -1,11 +1,21 @@
 """
-Hardware benchmark script for CI runners.
+"""Hardware benchmark script for CI runners.
 Compares CPU and GPU performance to diagnose slowdowns.
 Works on both CPU-only (GitHub Actions) and GPU (RunsOn) runners.
 """
 import time
 import platform
 import os
+import json
+from datetime import datetime
+
+# Global results dictionary for JSON output
+RESULTS = {
+    "timestamp": datetime.now().isoformat(),
+    "execution_method": "bare_metal",
+    "system": {},
+    "benchmarks": {}
+}
 
 def get_cpu_info():
     """Get CPU information."""
@@ -16,6 +26,10 @@ def get_cpu_info():
     print(f"Processor: {platform.processor()}")
     print(f"Python: {platform.python_version()}")
     
+    RESULTS["system"]["platform"] = platform.platform()
+    RESULTS["system"]["processor"] = platform.processor()
+    RESULTS["system"]["python_version"] = platform.python_version()
+    
     # Try to get CPU frequency
     try:
         with open('/proc/cpuinfo', 'r') as f:
@@ -38,18 +52,23 @@ def get_cpu_info():
     
     # CPU count
     print(f"CPU Count: {os.cpu_count()}")
+    RESULTS["system"]["cpu_count"] = os.cpu_count()
     
     # Check for GPU
     try:
         import subprocess
         result = subprocess.run(['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader'], 
                               capture_output=True, text=True, timeout=5)
         if result.returncode == 0:
-            print(f"GPU: {result.stdout.strip()}")
+            gpu_info = result.stdout.strip()
+            print(f"GPU: {gpu_info}")
+            RESULTS["system"]["gpu"] = gpu_info
         else:
             print("GPU: None detected")
+            RESULTS["system"]["gpu"] = None
     except:
         print("GPU: None detected (nvidia-smi not available)")
+        RESULTS["system"]["gpu"] = None
     
     print()
 
@@ -59,11 +78,14 @@ def benchmark_cpu_pure_python():
     print("CPU BENCHMARK: Pure Python")
     print("=" * 60)
     
+    results = {}
+    
     # Integer computation
     start = time.perf_counter()
     total = sum(i * i for i in range(10_000_000))
     elapsed = time.perf_counter() - start
     print(f"Integer sum (10M iterations): {elapsed:.3f} seconds")
+    results["integer_sum_10M"] = elapsed
     
     # Float computation
     start = time.perf_counter()
@@ -72,6 +94,9 @@ def benchmark_cpu_pure_python():
         total += (i * 0.1) ** 0.5
     elapsed = time.perf_counter() - start
     print(f"Float sqrt (1M iterations): {elapsed:.3f} seconds")
+    results["float_sqrt_1M"] = elapsed
+    
+    RESULTS["benchmarks"]["pure_python"] = results
     print()
 
 def benchmark_cpu_numpy():
@@ -82,6 +107,8 @@ def benchmark_cpu_numpy():
     print("CPU BENCHMARK: NumPy")
     print("=" * 60)
     
+    results = {}
+    
     # Matrix multiplication
     n = 3000
     A = np.random.randn(n, n)
@@ -91,6 +118,7 @@ def benchmark_cpu_numpy():
     C = A @ B
     elapsed = time.perf_counter() - start
     print(f"Matrix multiply ({n}x{n}): {elapsed:.3f} seconds")
+    results["matmul_3000x3000"] = elapsed
     
     # Element-wise operations
     x = np.random.randn(50_000_000)
@@ -99,6 +127,9 @@ def benchmark_cpu_numpy():
     y = np.cos(x**2) + np.sin(x)
     elapsed = time.perf_counter() - start
     print(f"Element-wise ops (50M elements): {elapsed:.3f} seconds")
+    results["elementwise_50M"] = elapsed
+    
+    RESULTS["benchmarks"]["numpy"] = results
     print()
 
 def benchmark_gpu_jax():
@@ -107,12 +138,18 @@ def benchmark_gpu_jax():
         import jax
         import jax.numpy as jnp
         
+        results = {}
+        
         devices = jax.devices()
         default_backend = jax.default_backend()
         
         # Check if GPU is available
         has_gpu = any('cuda' in str(d).lower() or 'gpu' in str(d).lower() for d in devices)
         
+        results["has_gpu"] = has_gpu
+        results["default_backend"] = default_backend
+        results["devices"] = [str(d) for d in devices]
+        
         print("=" * 60)
         if has_gpu:
             print("JAX BENCHMARK: GPU")
@@ -141,12 +178,14 @@ def matmul(a, b):
         C = matmul(A, B).block_until_ready()
         warmup_time = time.perf_counter() - start
         print(f"Warm-up (includes JIT compile, {n}x{n}): {warmup_time:.3f} seconds")
+        results["matmul_1000x1000_warmup"] = warmup_time
         
         # Actual benchmark (compiled)
         start = time.perf_counter()
         C = matmul(A, B).block_until_ready()
         elapsed = time.perf_counter() - start
         print(f"Matrix multiply compiled ({n}x{n}): {elapsed:.3f} seconds")
+        results["matmul_1000x1000_compiled"] = elapsed
         
         # Larger matrix
         n = 3000
@@ -158,12 +197,14 @@ def matmul(a, b):
         C = matmul(A, B).block_until_ready()
         warmup_time = time.perf_counter() - start
         print(f"Warm-up (recompile for {n}x{n}): {warmup_time:.3f} seconds")
+        results["matmul_3000x3000_warmup"] = warmup_time
         
         # Benchmark compiled
         start = time.perf_counter()
         C = matmul(A, B).block_until_ready()
         elapsed = time.perf_counter() - start
         print(f"Matrix multiply compiled ({n}x{n}): {elapsed:.3f} seconds")
+        results["matmul_3000x3000_compiled"] = elapsed)
         
         # Element-wise GPU benchmark
         x = jax.random.normal(key, (50_000_000,))
@@ -177,26 +218,33 @@ def elementwise_ops(x):
         y = elementwise_ops(x).block_until_ready()
         warmup_time = time.perf_counter() - start
         print(f"Element-wise warm-up (50M): {warmup_time:.3f} seconds")
+        results["elementwise_50M_warmup"] = warmup_time
         
         # Compiled
         start = time.perf_counter()
         y = elementwise_ops(x).block_until_ready()
         elapsed = time.perf_counter() - start
         print(f"Element-wise compiled (50M): {elapsed:.3f} seconds")
+        results["elementwise_50M_compiled"] = elapsed
         
+        RESULTS["benchmarks"]["jax"] = results
         print()
         
     except ImportError as e:
         print(f"JAX not available: {e}")
+        RESULTS["benchmarks"]["jax"] = {"error": str(e)}
     except Exception as e:
         print(f"JAX benchmark failed: {e}")
+        RESULTS["benchmarks"]["jax"] = {"error": str(e)})
 
 def benchmark_numba():
     """Numba CPU benchmark."""
     try:
         import numba
         import numpy as np
 
+        results = {}
+        
         print("=" * 60)
         print("CPU BENCHMARK: Numba")
         print("=" * 60)
@@ -213,12 +261,14 @@ def numba_sum(n):
         result = numba_sum(10_000_000)
         warmup_time = time.perf_counter() - start
         print(f"Integer sum warm-up (includes compile): {warmup_time:.3f} seconds")
+        results["integer_sum_10M_warmup"] = warmup_time
 
         # Compiled run
         start = time.perf_counter()
         result = numba_sum(10_000_000)
         elapsed = time.perf_counter() - start
         print(f"Integer sum compiled (10M): {elapsed:.3f} seconds")
+        results["integer_sum_10M_compiled"] = elapsed)
 
         @numba.jit(nopython=True, parallel=True)
         def numba_parallel_sum(arr):
@@ -234,19 +284,24 @@ def numba_parallel_sum(arr):
         result = numba_parallel_sum(arr)
         warmup_time = time.perf_counter() - start
         print(f"Parallel sum warm-up (50M): {warmup_time:.3f} seconds")
+        results["parallel_sum_50M_warmup"] = warmup_time
 
         # Compiled
         start = time.perf_counter()
         result = numba_parallel_sum(arr)
         elapsed = time.perf_counter() - start
         print(f"Parallel sum compiled (50M): {elapsed:.3f} seconds")
+        results["parallel_sum_50M_compiled"] = elapsed
 
+        RESULTS["benchmarks"]["numba"] = results
         print()
 
     except ImportError as e:
         print(f"Numba not available: {e}")
+        RESULTS["benchmarks"]["numba"] = {"error": str(e)}
     except Exception as e:
         print(f"Numba benchmark failed: {e}")
+        RESULTS["benchmarks"]["numba"] = {"error": str(e)})
 
 if __name__ == "__main__":
     print("\n" + "=" * 60)
@@ -262,3 +317,9 @@ def numba_parallel_sum(arr):
     print("=" * 60)
     print("BENCHMARK COMPLETE")
     print("=" * 60)
+    
+    # Save results to JSON
+    output_file = "benchmark_results_baremetal.json"
+    with open(output_file, 'w') as f:
+        json.dump(RESULTS, f, indent=2)
+    print(f"\nResults saved to {output_file}")