update matrix script

Author: ohmayr

.github/scripts/matrix_generator.py

@@ -1,94 +1,60 @@
 import os
 import json
 import argparse
-from typing import List, Dict
 
-# ---------------------------------------------------------
-# THE VIP LIST: These will get their own dedicated isolated VMs
-# ---------------------------------------------------------
-MAMMOTH_OVERRIDES = {
-    "google-cloud-spanner": 9999,
-    "google-cloud-compute": 9999,
-    "google-cloud-compute-v1beta": 9999,
-    "google-cloud-discoveryengine": 9999
+VIPS = {
+    "google-cloud-spanner", 
+    "google-cloud-compute",
+    "google-cloud-compute-v1beta", 
+    "google-cloud-discoveryengine"
 }
 
-def calculate_package_weight(pkg_path: str) -> int:
-    pkg_name = os.path.basename(os.path.normpath(pkg_path))
-    if pkg_name in MAMMOTH_OVERRIDES:
-        return MAMMOTH_OVERRIDES[pkg_name]
+def get_valid_packages(directories):
+    return [p for p in directories if os.path.isfile(os.path.join(p, "noxfile.py"))]
 
-    base_weight = 1
-    meta_path = os.path.join(pkg_path, ".repo-metadata.json")
-    if os.path.isfile(meta_path):
-        try:
-            with open(meta_path, 'r') as f:
-                if json.load(f).get("library_type") == "GAPIC_AUTO":
-                    return base_weight
-        except Exception:
-            pass 
-            
-    test_dir = os.path.join(pkg_path, "tests")
-    if os.path.isdir(test_dir):
-        for root, _, files in os.walk(test_dir):
-            base_weight += sum(1 for f in files if f.endswith(".py"))
-            
-    return base_weight
-
-def create_balanced_buckets(packages: List[str], max_buckets: int) -> List[Dict]:
-    valid_pkgs = [p for p in packages if os.path.isfile(os.path.join(p, "noxfile.py"))]
-    if not valid_pkgs:
-        return []
-        
-    pkg_weights = [(pkg, calculate_package_weight(pkg)) for pkg in valid_pkgs]
-    pkg_weights.sort(key=lambda x: x[1], reverse=True)
-    
-    num_buckets = min(len(valid_pkgs), max_buckets)
-    buckets = [{"id": i + 1, "packages": [], "total_weight": 0} for i in range(num_buckets)]
+def distribute_packages(packages, max_buckets):
+    # VIPs jump to the front of the line
+    packages.sort(key=lambda p: os.path.basename(p) not in VIPS)
 
-    for pkg, weight in pkg_weights:
-        lightest = min(buckets, key=lambda b: b["total_weight"])
-        lightest["packages"].append(pkg)
-        lightest["total_weight"] += weight
-        
-    # Build the final output payload
-    final_output = []
-    for b in buckets:
-        if not b["packages"]:
-            continue
-            
-        # Clean UI labels for isolated VIPs
-        first_pkg_name = os.path.basename(os.path.normpath(b["packages"][0]))
-        if first_pkg_name in MAMMOTH_OVERRIDES:
-            ui_label = first_pkg_name.replace("google-cloud-", "")
-        else:
-            ui_label = f"Bucket {b['id']}"
-            
-        final_output.append({
-            "id": ui_label, 
-            "packages": " ".join(b["packages"])
-        })
-        
-    return final_output
+    buckets = [{"weight": 0, "pkgs": []} for _ in range(min(len(packages), max_buckets))]
+
+    for pkg in packages:
+        # Find the bucket with the lowest weight, add the package, and update its weight
+        lightest = min(buckets, key=lambda b: b["weight"])
+        lightest["pkgs"].append(pkg)
+        lightest["weight"] += 9999 if os.path.basename(pkg) in VIPS else 1
+
+    return [b["pkgs"] for b in buckets]
+
+def build_github_actions_jobs(buckets):
+    jobs = []
+    for bucket in buckets:
+        # Create a clean UI label for the GitHub presubmit check
+        base_name = os.path.basename(bucket[0]).replace("google-cloud-", "")
+        job_label = f"{base_name} + {len(bucket) - 1}" if len(bucket) > 1 else base_name
+        jobs.append({"id": job_label, "packages": " ".join(bucket)})
+    return jobs
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--matrix-multiplier", type=int, required=True)
-    parser.add_argument("--max-vms", type=int, default=40)
+    parser.add_argument("--max-vms", type=int, default=20)
     args = parser.parse_args()
 
-    safe_github_limit = 250 // args.matrix_multiplier
-    max_allowed_buckets = min(safe_github_limit, args.max_vms)
+    packages = get_valid_packages(os.environ.get("CHANGED_DIRS", "").split())
+    if not packages:
+        return
 
-    changed_dirs = os.environ.get("CHANGED_DIRS", "").split()
-    buckets = create_balanced_buckets(changed_dirs, max_allowed_buckets)
+    max_buckets = min(250 // args.matrix_multiplier, args.max_vms)
+    buckets = distribute_packages(packages, max_buckets)
 
-    github_output = os.environ.get("GITHUB_OUTPUT")
-    if github_output:
-        with open(github_output, "a") as f:
-            f.write(f"buckets={json.dumps(buckets)}\n")
+    jobs_json = json.dumps(build_github_actions_jobs(buckets))
+
+    if "GITHUB_OUTPUT" in os.environ:
+        with open(os.environ["GITHUB_OUTPUT"], "a") as f:
+            f.write(f"buckets={jobs_json}\n")
     else:
-        print(json.dumps(buckets, indent=2))
+        print(jobs_json)
 
 if __name__ == "__main__":
     main()