CI: consolidate Pub/Sub publishing into a single pipeline-level message

flow/util/uploadMetadata.py

@@ -2,7 +2,6 @@
 
 import json
 import argparse
-import re
 import os
 from datetime import datetime, timezone
 
@@ -38,6 +37,13 @@
 parser.add_argument("--variant", type=str, default="base")
 
 # --- PUBSUB args ---
+parser.add_argument(
+    "--jenkinsEnv",
+    type=str,
+    default="unknown",
+    choices=["public", "secure", "unknown"],
+    help="Jenkins environment (public or secure)",
+)
 parser.add_argument("--pubsubProjectID", type=str, help="GCP project ID for Pub/Sub")
 parser.add_argument(
     "--pubsubTopicID",
@@ -81,7 +87,7 @@ def upload_data(db, dataFile, platform, design, variant, args, rules):
     excludes = ["run", "commit", "total_time", "constraints"]
     gen_date = datetime.now()
     for k, v in data.items():
-        new_key = re.sub(":", "__", k)  # replace ':' with '__'
+        new_key = k.replace(":", "__")  # replace ':' with '__'
         new_data[new_key] = v
         stage_name = k.split("__")[0]
         if stage_name not in excludes:
@@ -194,34 +200,103 @@ def upload_data(db, dataFile, platform, design, variant, args, rules):
 
 
 # --- PUBSUB ---
-def publish_to_pubsub(
-    publisher, topic_path, dataFile, platform, design, variant, args, rules
-):
-    """Publish a single design's metrics to Pub/Sub as a JSON message."""
+# Pub/Sub hard cap is 10 MB. Stay under with safety margin to leave room for
+# attribute overhead and future payload growth.
+MAX_PUBSUB_BYTES = 8 * 1024 * 1024
+
+
+def build_design_record(dataFile, platform, design, variant, rules):
+    """Return a dict for one design to be included in the pipeline-level payload."""
     with open(dataFile) as f:
         data = json.load(f)
+    metrics = {k.replace(":", "__"): v for k, v in data.items()}
+    return {
+        "platform": platform,
+        "design": design,
+        "variant": variant,
+        "rules": rules,
+        "metrics": metrics,
+    }
 
-    # Build the payload: CLI args + metrics with ':' replaced by '__'
-    payload = {
+
+def build_pipeline_payload(design_records, args):
+    """Return the v2 pipeline-level payload dict."""
+    return {
+        "payload_schema_version": 2,
+        "jenkins_env": args.jenkinsEnv,
         "build_id": args.buildID,
         "branch_name": args.branchName,
         "pipeline_id": args.pipelineID,
         "change_branch": args.changeBranch,
         "commit_sha": args.commitSHA,
         "jenkins_url": args.jenkinsURL,
-        "rules": rules,
+        "designs": design_records,
     }
 
-    for k, v in data.items():
-        new_key = re.sub(":", "__", k)
-        payload[new_key] = v
 
-    message_data = json.dumps(payload).encode("utf-8")
-    future = publisher.publish(topic_path, data=message_data)
-    message_id = future.result()
+def publish_pipeline_report(publisher, topic_path, message_data, design_count, args):
+    """Publish a pre-encoded v2 pipeline message."""
+    size_kb = len(message_data) / 1024
     print(
-        f"[INFO] Published to Pub/Sub (message ID: {message_id}) for {platform} {design} {variant}."
+        f"[INFO] Publishing pipeline report ({design_count} designs, {size_kb:.1f} KB) to Pub/Sub."
+    )
+    future = publisher.publish(
+        topic_path,
+        data=message_data,
+        payload_schema_version="2",
+        jenkins_env=args.jenkinsEnv,
     )
+    message_id = future.result()
+    print(f"[INFO] Published pipeline report to Pub/Sub (message ID: {message_id}).")
+
+
+def publish_v1_per_design(publisher, topic_path, design_records, args):
+    """Fallback path used when the v2 pipeline payload exceeds MAX_PUBSUB_BYTES.
+
+    Emits one v1-format message per design (no payload_schema_version, metrics
+    flattened at the root), matching the legacy schema the ingestion service
+    still supports.
+    """
+    futures = []
+    for d in design_records:
+        payload = {
+            "build_id": args.buildID,
+            "branch_name": args.branchName,
+            "pipeline_id": args.pipelineID,
+            "change_branch": args.changeBranch,
+            "commit_sha": args.commitSHA,
+            "jenkins_url": args.jenkinsURL,
+            "jenkins_env": args.jenkinsEnv,
+            "rules": d["rules"],
+        }
+        payload.update(d["metrics"])
+
+        message_data = json.dumps(payload, default=str).encode("utf-8")
+        try:
+            future = publisher.publish(
+                topic_path,
+                data=message_data,
+                jenkins_env=args.jenkinsEnv,
+            )
+            futures.append((d, future))
+        except Exception as e:
+            print(
+                f"[WARN] Pub/Sub v1 fallback publish failed for "
+                f"{d['platform']} {d['design']} {d['variant']}: {e}"
+            )
+
+    for d, future in futures:
+        try:
+            message_id = future.result()
+            print(
+                f"[INFO] Published v1 fallback message (ID: {message_id}) for "
+                f"{d['platform']} {d['design']} {d['variant']}."
+            )
+        except Exception as e:
+            print(
+                f"[WARN] Pub/Sub v1 fallback publish failed for "
+                f"{d['platform']} {d['design']} {d['variant']}: {e}"
+            )
 
 
 # --- END PUBSUB ---
@@ -264,6 +339,10 @@ def get_rules(dataFile):
 
 RUN_FILENAME = "metadata.json"
 
+# --- PUBSUB ---
+design_records = []
+# --- END PUBSUB ---
+
 for reportDir, dirs, files in sorted(os.walk("reports", topdown=False)):
     dirList = reportDir.split(os.sep)
     if len(dirList) != 4:
@@ -293,12 +372,30 @@ def get_rules(dataFile):
 
     # --- PUBSUB ---
     if publisher:
+        design_records.append(
+            build_design_record(dataFile, platform, design, variant, rules)
+        )
+    # --- END PUBSUB ---
+
+# --- PUBSUB ---
+if publisher and design_records:
+    payload = build_pipeline_payload(design_records, args)
+    message_data = json.dumps(payload, default=str).encode("utf-8")
+
+    if len(message_data) > MAX_PUBSUB_BYTES:
+        print(
+            f"[WARN] v2 payload size {len(message_data) / 1024:.1f} KB exceeds "
+            f"{MAX_PUBSUB_BYTES // 1024} KB cap. Falling back to v1 per-design publish "
+            f"({len(design_records)} messages)."
+        )
+        publish_v1_per_design(publisher, topic_path, design_records, args)
+    else:
         try:
-            publish_to_pubsub(
-                publisher, topic_path, dataFile, platform, design, variant, args, rules
+            publish_pipeline_report(
+                publisher, topic_path, message_data, len(design_records), args
             )
         except Exception as e:
-            print(
-                f"[WARN] Pub/Sub publish failed for {platform} {design} {variant}: {e}"
-            )
-    # --- END PUBSUB ---
+            print(f"[WARN] Pub/Sub publish failed for pipeline report: {e}")
+elif publisher and not design_records:
+    print("[WARN] Pub/Sub publisher initialized but no design records were collected.")
+# --- END PUBSUB ---