""" Vertex AI Batch Prediction Eligibility Probe - gemini-3.1-pro-preview ===================================================================== §1.3f per PM-RATIFY-V5-RPD fallback Branch A (2026-04-24). Empirical question: does `publishers/google/models/gemini-3.1-pro-preview` accept Vertex Batch Prediction? If ELIGIBLE, unlocks an async path that sidesteps the 25 RPM online per-model cap (§1.3 probe root cause) and the current 250 RPD ceiling. Scope discipline: - No §11 frozen path touched (no runner, no judge-runner, no litellm- config.yaml, no failure-mode-judge.ts). - No benchmark adapter. Self-contained standalone script. - All artefacts emitted to `benchmarks/probes/vertex-batch-eligibility/`. Execution: 1. Build 5-instance Vertex Batch JSONL from the canonical LoCoMo fixture. 2. Upload to GCS. 3. Submit BatchPredictionJob via the newer Gemini `submit()` surface (vertexai.batch_prediction), with fallback to the legacy `aiplatform.BatchPredictionJob.create()` if submit is unavailable or rejects the preview model. 4. Poll up to `initial_poll_seconds` (default 450s = 7.5 min, within Bash 10-min timeout). On INCOMPLETE, emit resume-ready job name for a follow-up invocation with --resume . Usage: python probe-script.py # fresh submit + poll python probe-script.py --resume # poll-only resume of running job Judge prompt: verbatim reproduction of `packages/server/src/benchmarks/judge/failure-mode-judge.ts:93-140` (buildJudgePrompt). Identical to §1.3 + §1.3c probe payload so any eligibility difference is routing, not prompt content. """ from __future__ import annotations import argparse import json import sys import time import traceback from datetime import datetime, timezone from pathlib import Path # ── Bound parameters (all from PM direktiva; no placeholders) ───────────── PROJECT_ID = "gen-lang-client-0674908699" LOCATION = "us-central1" GCS_BUCKET = "egzakta-vertex-batch-probe-2026-04" GCS_INPUT = f"gs://{GCS_BUCKET}/probe-input.jsonl" GCS_OUTPUT_PREFIX = f"gs://{GCS_BUCKET}/probe-output/" MODEL_ID = "publishers/google/models/gemini-3.1-pro-preview" JOB_DISPLAY_NAME = "locomo-mini-vertex-batch-eligibility-probe-2026-04-24" FIXTURE_PATH = Path("D:/Projects/waggle-os/benchmarks/data/locomo/locomo-1540.jsonl") PROBE_DIR = Path("D:/Projects/waggle-os/benchmarks/probes/vertex-batch-eligibility") LOG_PATH = PROBE_DIR / "job-state-trace.log" INPUT_PATH = PROBE_DIR / "probe-input.jsonl" OUTPUT_PATH = PROBE_DIR / "job-output.jsonl" TERMINAL_STATES = { "JobState.JOB_STATE_SUCCEEDED", "JobState.JOB_STATE_FAILED", "JobState.JOB_STATE_CANCELLED", "JobState.JOB_STATE_EXPIRED", } # Verbatim from failure-mode-judge.ts:93-140 (buildJudgePrompt). JUDGE_PROMPT_TEMPLATE = "\n".join([ "You are evaluating whether an LLM's answer is correct against ground truth.", "", "## Question", "{question}", "", "## Ground-truth answer", "{ground_truth}", "", "## Ground-truth supporting context (excerpt shown to the model)", "{context}", "", "## Model's answer", "{model_answer}", "", "## Your task", "", "Step 1: Determine if the model's answer is correct.", "- \"correct\" means the model's answer contains all required facts from ground truth, with no additional incorrect claims.", "- Minor phrasing differences, synonyms, or alternative but equivalent formulations are acceptable.", "- Extra detail is acceptable ONLY if it is factually correct.", "", "Step 2: If incorrect, assign exactly one failure mode using this decision tree:", "", "1. Does the model explicitly refuse or say it does not know? ->F1 (ABSTAIN)", "2. Does the model answer a DIFFERENT question than was asked (coherent but off-topic)? ->F5 (OFF-TOPIC)", "3. Does the model rely on entities, names, dates, or claims that do NOT appear in the ground-truth context (fabrication)? ->F4 (HALLUCINATED)", "4. Does the model correctly state SOME required facts but miss others, without stating any incorrect facts? ->F2 (PARTIAL)", "5. Otherwise (model states facts derived from the context but gets them wrong): ->F3 (INCORRECT)", "", "Step 3: Return JSON only, no prose, in this exact schema:", "", "{{", " \"verdict\": \"correct\" | \"incorrect\",", " \"failure_mode\": null | \"F1\" | \"F2\" | \"F3\" | \"F4\" | \"F5\",", " \"rationale\": \"one sentence explaining the verdict\"", "}}", "", "If verdict is \"correct\", failure_mode MUST be null.", "If verdict is \"incorrect\", failure_mode MUST be one of F1-F5.", ]) def ts() -> str: return datetime.now(timezone.utc).isoformat() # Windows cp1252 stdout chokes on Unicode arrows; reconfigure once at import. try: sys.stdout.reconfigure(encoding="utf-8", errors="replace") sys.stderr.reconfigure(encoding="utf-8", errors="replace") except Exception: pass def logline(text: str) -> None: line = f"{ts()} {text}" try: print(line, flush=True) except UnicodeEncodeError: print(line.encode("ascii", errors="replace").decode("ascii"), flush=True) with LOG_PATH.open("a", encoding="utf-8") as f: f.write(line + "\n") def build_input_jsonl() -> int: PROBE_DIR.mkdir(parents=True, exist_ok=True) count = 0 with FIXTURE_PATH.open("r", encoding="utf-8") as f, INPUT_PATH.open("w", encoding="utf-8") as g: for line in f: if count >= 5: break line = line.strip() if not line: continue rec = json.loads(line) prompt = JUDGE_PROMPT_TEMPLATE.format( question=rec["question"], ground_truth=rec.get("gold_answer", ""), context=rec.get("context", ""), model_answer="unknown", ) # Vertex Batch JSONL request shape for Gemini models. req = { "request": { "contents": [ {"role": "user", "parts": [{"text": prompt}]}, ], "generationConfig": {"temperature": 0.0, "maxOutputTokens": 512}, } } g.write(json.dumps(req, ensure_ascii=False) + "\n") count += 1 return count def upload_to_gcs(local_path: Path, gcs_uri: str) -> None: from google.cloud import storage assert gcs_uri.startswith("gs://") rest = gcs_uri[len("gs://"):] bucket_name, _, blob_name = rest.partition("/") client = storage.Client(project=PROJECT_ID) bucket = client.bucket(bucket_name) blob = bucket.blob(blob_name) blob.upload_from_filename(str(local_path)) logline(f"[upload] {local_path} ->{gcs_uri} ({local_path.stat().st_size} bytes)") def download_first_output_jsonl(gcs_prefix: str, local_path: Path) -> int: from google.cloud import storage assert gcs_prefix.startswith("gs://") rest = gcs_prefix[len("gs://"):] bucket_name, _, prefix = rest.partition("/") client = storage.Client(project=PROJECT_ID) bucket = client.bucket(bucket_name) rows = 0 with local_path.open("w", encoding="utf-8") as out: for blob in client.list_blobs(bucket, prefix=prefix): if blob.name.endswith(".jsonl") or "predictions" in blob.name: data = blob.download_as_text() out.write(data) rows += sum(1 for line in data.splitlines() if line.strip()) logline(f"[download] {blob.name} ->{local_path} ({len(data)} bytes)") return rows def submit_via_preview_api(model_variant: str) -> object | None: """Try the Gemini-native BatchPredictionJob.submit() API with a given model identifier form. The SDK may want bare name (`gemini-3.1-pro- preview`) vs fully-qualified (`publishers/google/models/gemini-3.1-pro- preview`); caller loops through variants.""" try: import vertexai from vertexai.batch_prediction import BatchPredictionJob except ImportError as e: logline(f"[submit:preview:{model_variant}] ImportError {e}") return None try: vertexai.init(project=PROJECT_ID, location=LOCATION) job = BatchPredictionJob.submit( source_model=model_variant, input_dataset=GCS_INPUT, output_uri_prefix=GCS_OUTPUT_PREFIX, job_display_name=JOB_DISPLAY_NAME, ) logline(f"[submit:preview:{model_variant}] OK job_resource_name={job.resource_name}") return job except Exception as e: logline(f"[submit:preview:{model_variant}] ERROR {type(e).__name__}: {str(e)[:500]}") return None def submit_via_legacy_api() -> object | None: """Fall back to the legacy aiplatform.BatchPredictionJob.create() API.""" try: from google.cloud import aiplatform except ImportError as e: logline(f"[submit:legacy] ImportError {e}") return None try: aiplatform.init(project=PROJECT_ID, location=LOCATION) job = aiplatform.BatchPredictionJob.create( job_display_name=JOB_DISPLAY_NAME, model_name=MODEL_ID, instances_format="jsonl", gcs_source=GCS_INPUT, gcs_destination_prefix=GCS_OUTPUT_PREFIX, predictions_format="jsonl", sync=False, ) logline(f"[submit:legacy] OK - job_resource_name={job.resource_name}") return job except Exception as e: logline(f"[submit:legacy] ERROR {type(e).__name__}: {str(e)[:500]}") return None def poll_until_terminal(job, max_seconds: int, cadence_s: int = 60) -> str: t0 = time.time() last_state = None while time.time() - t0 < max_seconds: try: # Refresh from Vertex. job._sync_gca_resource() except Exception as e: logline(f"[poll] sync error {type(e).__name__}: {str(e)[:200]}") state = str(job.state) if hasattr(job, "state") else "UNKNOWN" if state != last_state: logline(f"[poll] t+{int(time.time()-t0):>4}s state={state}") last_state = state if state in TERMINAL_STATES: return state time.sleep(cadence_s) logline(f"[poll] INCOMPLETE after {max_seconds}s - state={last_state}") return "INCOMPLETE" def resume_job(resource_name: str): """Re-attach to an existing BatchPredictionJob by its resource name.""" logline(f"[resume] reattaching to {resource_name}") try: from vertexai.batch_prediction import BatchPredictionJob as PreviewBPJ import vertexai vertexai.init(project=PROJECT_ID, location=LOCATION) return PreviewBPJ(resource_name) except Exception as e: logline(f"[resume:preview] fail {type(e).__name__}: {e}") try: from google.cloud import aiplatform aiplatform.init(project=PROJECT_ID, location=LOCATION) return aiplatform.BatchPredictionJob(resource_name) except Exception as e: logline(f"[resume:legacy] fail {type(e).__name__}: {e}") return None def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--resume", default=None, help="resource_name of an already-submitted BPJ") ap.add_argument("--max-seconds", type=int, default=450, help="max wall-clock in poll loop") args = ap.parse_args() logline(f"[probe] START - PROJECT={PROJECT_ID} LOCATION={LOCATION} MODEL={MODEL_ID}") # Package version stamp for audit. try: import google.cloud.aiplatform as aip import google.cloud.storage as gcs logline(f"[probe] sdk versions: aiplatform={aip.__version__} storage={gcs.__version__}") except Exception: logline("[probe] sdk version stamp failed") if args.resume: job = resume_job(args.resume) if job is None: logline("[probe] resume failed - exit 2") return 2 else: # Phase 1: build input. n = build_input_jsonl() logline(f"[build] wrote {n} requests to {INPUT_PATH}") # Phase 2: upload. try: upload_to_gcs(INPUT_PATH, GCS_INPUT) except Exception as e: logline(f"[upload] FATAL {type(e).__name__}: {str(e)[:500]}") logline("[probe] VERDICT=INCONCLUSIVE reason=gcs_upload_fail") return 3 # Phase 3: submit. Try model-ID variants in order (SDK preferences # vary — bare name vs fully-qualified). The Vertex publisher catalog # GET confirmed `gemini-3.1-pro-preview` exists at HTTP 200; if all # variants via the SDK still 404, that's evidence of # "model exists but batch-prediction not enabled for it" -> INFEASIBLE. model_variants = [ "gemini-3.1-pro-preview", # bare name "publishers/google/models/gemini-3.1-pro-preview", # fully-qualified ] job = None for variant in model_variants: job = submit_via_preview_api(variant) if job is not None: break if job is None: logline("[submit:preview] all variants exhausted; trying legacy API") job = submit_via_legacy_api() if job is None: logline("[probe] SUBMISSION FAILED on all preview model variants + legacy API") logline("[probe] VERDICT=INFEASIBLE reason=submission_rejected_all_variants") return 4 # Phase 4: poll. final_state = poll_until_terminal(job, max_seconds=args.max_seconds) logline(f"[probe] TERMINAL_STATE={final_state} job_resource_name={job.resource_name}") if final_state == "INCOMPLETE": logline(f"[probe] RESUMABLE - re-run with: python probe-script.py --resume {job.resource_name}") logline("[probe] VERDICT=PARTIAL_PENDING wall_clock_exhausted") return 5 # Phase 5: classify. if final_state == "JobState.JOB_STATE_SUCCEEDED": logline("[probe] downloading output JSONL(s)...") rows = download_first_output_jsonl(GCS_OUTPUT_PREFIX, OUTPUT_PATH) logline(f"[probe] downloaded {rows} output rows ->{OUTPUT_PATH}") logline("[probe] VERDICT=ELIGIBLE") return 0 elif final_state == "JobState.JOB_STATE_FAILED": err = getattr(job, "error", None) logline(f"[probe] JOB FAILED - error={err}") logline("[probe] VERDICT=INFEASIBLE reason=job_state_failed") return 6 else: logline(f"[probe] VERDICT=INCONCLUSIVE reason=unexpected_state:{final_state}") return 7 if __name__ == "__main__": try: sys.exit(main()) except Exception as e: logline(f"[FATAL] {type(e).__name__}: {str(e)[:500]}") logline(f"[FATAL] traceback: {traceback.format_exc()[:2000]}") logline("[probe] VERDICT=INCONCLUSIVE reason=fatal_exception") sys.exit(99)