diff --git a/flow/util/uploadMetadata.py b/flow/util/uploadMetadata.py index 9f8dc63fc3..dbd4ab2e2c 100755 --- a/flow/util/uploadMetadata.py +++ b/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 ---