"""Marker-driven join: launch fusion when both extraction halves complete. Two event shapes invoke this handler: 1. ``s3:ObjectCreated`` for a completion marker (``*.complete.json``) in the output bucket. Each extraction job's last write is its marker, so a marker event means that job's artifacts are fully in place. On every marker event the handler checks whether the *other* half's marker is also present: - one marker only -> exit; the second marker's event will act. - both present, both good -> claim the ``fuse-started`` lock and launch the Fargate fuse task. - either says ``failed`` -> hold the document and alert. Never fuse a partial document. The lock is an S3 conditional create (``If-None-Match: *``): exactly one concurrent invocation wins, which is what prevents a double fuse launch when both markers land near-simultaneously or S3 delivers an event twice. 2. A one-shot EventBridge schedule (``{"deadline_check": ...}``), created by the inbound Lambda when it submitted the parse job. If the parse marker is still absent at the deadline the job died without a trace — the handler asks RunPod what happened and raises the same hold-and-alert. Holding a document writes two objects and one alert: ``hold.json`` with the reason, an empty pointer under ``hold/`` so held documents are listable in one place, and an SNS publish. Release = fix the cause, delete the pointer and hold.json, re-drop the file; the inbound Lambda's per-engine dedup re-runs only the missing half. Required environment (set by Terraform): ECS_CLUSTER, TASK_DEFINITION, CONTAINER_NAME, SUBNETS, SECURITY_GROUPS, OUTPUT_BUCKET, INPUT_BUCKET, ALERT_TOPIC_ARN, RUNPOD_ENDPOINT_ID, RUNPOD_API_KEY_PARAM; optional INPUT_PREFIX (default "inbound/"). """ from __future__ import annotations import json import logging import os import urllib.parse import urllib.request from typing import Any import boto3 log = logging.getLogger() log.setLevel(logging.INFO) ecs = boto3.client("ecs") s3 = boto3.client("s3") sns = boto3.client("sns") ssm = boto3.client("ssm") CLUSTER = os.environ["ECS_CLUSTER"] TASK_DEFINITION = os.environ["TASK_DEFINITION"] CONTAINER_NAME = os.environ["CONTAINER_NAME"] SUBNETS = [s for s in os.environ["SUBNETS"].split(",") if s] SECURITY_GROUPS = [s for s in os.environ["SECURITY_GROUPS"].split(",") if s] OUTPUT_BUCKET = os.environ["OUTPUT_BUCKET"] INPUT_BUCKET = os.environ["INPUT_BUCKET"] INPUT_PREFIX = os.environ.get("INPUT_PREFIX", "inbound/") ALERT_TOPIC_ARN = os.environ["ALERT_TOPIC_ARN"] RUNPOD_ENDPOINT_ID = os.environ["RUNPOD_ENDPOINT_ID"] RUNPOD_API_KEY_PARAM = os.environ["RUNPOD_API_KEY_PARAM"] runpod_api_key: str | None = None def get_runpod_api_key() -> str: global runpod_api_key if runpod_api_key is None: response = ssm.get_parameter(Name=RUNPOD_API_KEY_PARAM, WithDecryption=True) runpod_api_key = str(response["Parameter"]["Value"]) return runpod_api_key def read_json(key: str) -> dict[str, object] | None: try: body = s3.get_object(Bucket=OUTPUT_BUCKET, Key=key)["Body"].read() except s3.exceptions.NoSuchKey: return None except s3.exceptions.ClientError as error: if error.response["Error"]["Code"] in ("404", "NoSuchKey"): return None raise return json.loads(body) def claim(key: str) -> bool: """Create-only put: True for exactly one concurrent caller.""" try: s3.put_object(Bucket=OUTPUT_BUCKET, Key=key, Body=b"", IfNoneMatch="*") return True except s3.exceptions.ClientError as error: if error.response["Error"]["Code"] in ("PreconditionFailed", "412"): return False raise def source_pdf(prefix: str) -> str: """Invert the path mirroring: /// -> the inbound key.""" return f"s3://{INPUT_BUCKET}/{INPUT_PREFIX}{prefix.rstrip('/')}.pdf" def launch_fuse(prefix: str) -> list[str]: dst = f"s3://{OUTPUT_BUCKET}/{prefix}" resp = ecs.run_task( cluster=CLUSTER, taskDefinition=TASK_DEFINITION, launchType="FARGATE", count=1, networkConfiguration={ "awsvpcConfiguration": { "subnets": SUBNETS, "securityGroups": SECURITY_GROUPS, "assignPublicIp": "ENABLED", } }, overrides={ "containerOverrides": [ { "name": CONTAINER_NAME, "command": [ "fuse", source_pdf(prefix), "--artifacts-dir", dst, "-o", dst, "--llm-backend", "api", ], } ] }, ) failures = resp.get("failures", []) if failures: log.error("fuse RunTask failures for %s: %s", prefix, failures) return [t["taskArn"] for t in resp.get("tasks", [])] def hold_document(prefix: str, base: str, reason: dict[str, object]) -> None: """Park the document for a human: hold.json + listable pointer + alert. The hold.json is claimed conditionally so concurrent invocations (both markers failing close together, duplicate event delivery) produce one alert, not several. """ hold_key = f"{prefix}hold.json" if not claim(hold_key): log.info("already held: %s", prefix) return s3.put_object( Bucket=OUTPUT_BUCKET, Key=hold_key, Body=json.dumps(reason, indent=2).encode("utf-8"), ) s3.put_object(Bucket=OUTPUT_BUCKET, Key=f"hold/{prefix.rstrip('/')}", Body=b"") sns.publish( TopicArn=ALERT_TOPIC_ARN, Subject=f"quber document held: {base}"[:100], Message=json.dumps( {"prefix": prefix, "hold": f"s3://{OUTPUT_BUCKET}/{hold_key}", **reason}, indent=2 ), ) log.error("held %s: %s", prefix, reason) def join_on_marker(prefix: str, base: str) -> dict[str, object]: parse_marker = read_json(f"{prefix}{base}.parse.complete.json") table_marker = read_json(f"{prefix}{base}.tables.complete.json") if parse_marker is None or table_marker is None: log.info( "waiting for other half: %s (parse=%s, table=%s)", prefix, bool(parse_marker), bool(table_marker) ) return {"action": "waiting"} failed = { name: marker for name, marker in (("parse", parse_marker), ("table", table_marker)) if marker.get("status") != "complete" } if failed: hold_document(prefix, base, {"cause": "extraction failed", "failed": failed}) return {"action": "held", "failed": sorted(failed)} if not claim(f"{prefix}fuse-started"): log.info("fuse already claimed: %s", prefix) return {"action": "already_claimed"} arns = launch_fuse(prefix) log.info("fuse launched for %s: %s", prefix, arns) return {"action": "fused", "tasks": arns} def deadline_check(prefix: str, base: str, job_id: str) -> dict[str, object]: if read_json(f"{prefix}{base}.parse.complete.json") is not None: return {"action": "marker_present"} request = urllib.request.Request( f"https://api.runpod.ai/v2/{RUNPOD_ENDPOINT_ID}/status/{job_id}", headers={"Authorization": f"Bearer {get_runpod_api_key()}"}, ) try: with urllib.request.urlopen(request, timeout=30) as response: status = json.load(response) except Exception as error: # noqa: BLE001 — the alert must carry whatever we know status = {"status": "unknown", "error": str(error)} hold_document( prefix, base, {"cause": "parse deadline expired with no marker", "job_id": job_id, "runpod_status": status}, ) return {"action": "held", "runpod_status": status.get("status")} def handler(event: dict[str, Any], context: object) -> dict[str, Any]: if "deadline_check" in event: check = event["deadline_check"] return deadline_check(check["prefix"], check["base"], check["job_id"]) results = [] for record in event.get("Records", []): key = urllib.parse.unquote_plus(record["s3"]["object"]["key"]) if key.startswith("hold/") or key.endswith(".fuse.complete.json"): continue if not (key.endswith(".parse.complete.json") or key.endswith(".tables.complete.json")): log.info("skip non-join marker: %s", key) continue prefix, marker_name = key.rsplit("/", 1) prefix += "/" base = marker_name.removesuffix(".parse.complete.json").removesuffix(".tables.complete.json") results.append(join_on_marker(prefix, base)) return {"results": results}