divide tests based on category

Author: ohmayr

.kokoro/system.sh

@@ -75,11 +75,6 @@ run_package_test() {
   # Export variables for the duration of this function's sub-processes
   export PROJECT_ID GOOGLE_APPLICATION_CREDENTIALS NOX_FILE NOX_SESSION
   export GOOGLE_CLOUD_PROJECT="${PROJECT_ID}"
-  
-  # Limit pytest-xdist to 4 workers per package. When 3 packages run in parallel,
-  # if each uses -n=auto they will spawn 16 workers each (48 total threads), 
-  # completely thrashing the CPU and causing all tests to hang/timeout.
-  export PYTEST_ADDOPTS="-n 4"
 
   # Isolate PIP cache to prevent concurrent pip file lock deadlocks
   export PIP_CACHE_DIR="/tmpfs/.pip_cache_$(basename ${package_name})"
@@ -164,18 +159,56 @@ for path in `find 'packages' \
 done
 
 if [ -n "$PACKAGES_TO_TEST" ]; then
-  mkdir -p .logs
   export -f run_package_test
   export system_test_script PROJECT_ROOT KOKORO_GFILE_DIR
-  
-  echo "Running system tests in parallel (3 workers)..."
-  # Use timeout to prevent infinite hangs, and < /dev/null to prevent stdin blocks
-  echo "$PACKAGES_TO_TEST" | tr ' ' '\n' | awk 'NF' | xargs -P 3 -I {} bash -c 'timeout 15m bash -c "run_package_test \"{}\" < /dev/null" > ".logs/{}.log" 2>&1 || touch ".logs/{}.failed"'
-  
-  for failed in .logs/*.failed; do
-    if [ -f "$failed" ]; then
-      echo "--- FAILED: ${failed%.failed} ---"
-      cat "${failed%.failed}.log"
+
+  # 1. DYNAMIC ROUTING: Automatically detect which packages are CPU hogs by checking if they install pytest-xdist or hardcode workers
+  LIGHT_TO_TEST=""
+  HEAVY_TO_TEST=""
+  for pkg in $PACKAGES_TO_TEST; do
+    if grep -qE "pytest-xdist|-n=auto|-n=[0-9]+" "packages/$pkg/noxfile.py" "packages/$pkg/setup.py" 2>/dev/null; then
+      HEAVY_TO_TEST="$HEAVY_TO_TEST $pkg"
+    else
+      LIGHT_TO_TEST="$LIGHT_TO_TEST $pkg"
+    fi
+  done
+
+  # 2. PARALLEL LANE (Live Streaming): Run light packages with a parallel job queue. 
+  # We prefix every line with the package name so output streams LIVE and remains readable.
+  if [ -n "$LIGHT_TO_TEST" ]; then
+    echo "============================================================"
+    echo "Running Lightweight Packages in Parallel (4 workers max)"
+    echo "============================================================"
+    for pkg in $LIGHT_TO_TEST; do
+      (
+        timeout 15m bash -c "run_package_test \"$pkg\" < /dev/null" 2>&1 | awk -v prefix="[$pkg]" '{print prefix, $0}'
+        if [ ${PIPESTATUS[0]} -ne 0 ]; then touch ".failed_$pkg"; fi
+      ) &
+      # Limit parallel background jobs to 4
+      while [ $(jobs -r | wc -l) -ge 4 ]; do sleep 1; done
+    done
+    wait # Wait for all parallel jobs to finish
+  fi
+
+  # 3. SEQUENTIAL VIP LANE: Run heavy packages one-by-one so they have 100% of the VM resources.
+  if [ -n "$HEAVY_TO_TEST" ]; then
+    echo "============================================================"
+    echo "Running CPU-Intensive Packages Sequentially"
+    echo "============================================================"
+    for pkg in $HEAVY_TO_TEST; do
+      if [ -n "$pkg" ]; then
+        echo "[$pkg] Starting sequential execution..."
+        timeout 25m bash -c "run_package_test \"$pkg\" < /dev/null" 2>&1 | awk -v prefix="[$pkg]" '{print prefix, $0}'
+        if [ ${PIPESTATUS[0]} -ne 0 ]; then touch ".failed_$pkg"; fi
+      fi
+    done
+  fi
+
+  # 4. FAIL STATE EVALUATION
+  for failed_marker in .failed_*; do
+    if [ -f "$failed_marker" ]; then
+      failed_pkg="${failed_marker#.failed_}"
+      echo "--- FAILED: $failed_pkg ---"
       RETVAL=1
     fi
   done