Lilith cluster helper bash script, calls gen-witness and prove and polls

gen-witness-prove.sh

#!/usr/bin/env bash
set -euo pipefail

# Requires: curl, jq, uuidgen
# Usage example:
# ./gwp.sh \
#   --user-id YOUR_USER_ID \
#   --api-key YOUR_API_KEY \
#   --artifact YOUR_ARTIFACT \
#   --deployment YOUR_DEPLOYMENT \
#   --input-file input.json \

ARCHON_URL="https://archon-v0.ezkl.xyz"
ARCHON_USER_ID=""
ARCHON_API_KEY=""
ARCHON_ARTIFACT=""
ARCHON_DEPLOYMENT=""
INPUT_FILE=""
COMPILED_CIRCUIT="model.compiled"
PK_PATH="pk.key"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --archon-url)        ARCHON_URL="$2"; shift 2 ;;
    --user-id)           ARCHON_USER_ID="$2"; shift 2 ;;
    --api-key)           ARCHON_API_KEY="$2"; shift 2 ;;
    --artifact)          ARCHON_ARTIFACT="$2"; shift 2 ;;
    --deployment)        ARCHON_DEPLOYMENT="$2"; shift 2 ;;
    --input-file)        INPUT_FILE="$2"; shift 2 ;;
    --compiled-circuit)  COMPILED_CIRCUIT="$2"; shift 2 ;;
    --pk-path)          PK_PATH="$2"; shift 2 ;;
    -h|--help)
      echo "See script header for usage."
      exit 0
      ;;
    *)
      echo "Unknown arg: $1" >&2; exit 1 ;;
  esac
done

# Simple validation
for v in ARCHON_URL ARCHON_USER_ID ARCHON_API_KEY ARCHON_ARTIFACT ARCHON_DEPLOYMENT INPUT_FILE COMPILED_CIRCUIT PK_PATH; do
  if [[ -z "${!v}" ]]; then
    echo "Missing required flag for $v" >&2
    exit 1
  fi
done

if ! command -v jq >/dev/null; then echo "jq is required" >&2; exit 1; fi
if ! command -v uuidgen >/dev/null; then echo "uuidgen is required" >&2; exit 1; fi

LATEST_UUID="$(uuidgen)"

INPUT_BASENAME="input_${LATEST_UUID}.json"
WITNESS_BASENAME="witness_${LATEST_UUID}.json"
PROOF_BASENAME="proof_${LATEST_UUID}.json"

# Build JSON payload using jq to avoid quoting issues
POST_BODY="$(jq -n \
  --arg artifact "$ARCHON_ARTIFACT" \
  --arg deployment "$ARCHON_DEPLOYMENT" \
  --arg compiled "$COMPILED_CIRCUIT" \
  --arg pk "$PK_PATH" \
  --arg input_path "$INPUT_BASENAME" \
  --arg witness_path "$WITNESS_BASENAME" \
  --arg proof_path "$PROOF_BASENAME" \
  --argjson input_data "$(cat "$INPUT_FILE")" '
{
  commands: [
    {
      artifact: $artifact,
      binary: "ezkl",
      deployment: $deployment,
      command: [
        "gen-witness",
        ("--data " + $input_path),
        ("--compiled-circuit " + $compiled),
        ("--output " + $witness_path)
      ]
    },
    {
      artifact: $artifact,
      binary: "ezkl",
      deployment: $deployment,
      command: [
        "prove",
        ("--witness " + $witness_path),
        ("--compiled-circuit " + $compiled),
        ("--pk-path " + $pk),
        ("--proof-path " + $proof_path)
      ],
      output_path: [ $proof_path ]
    }
  ],
  data: [
    {
      target_path: $input_path,
      data: $input_data
    }
  ]
}
')"

# Submit job
CREATE_URL="${ARCHON_URL%/}/recipe?user_id=${ARCHON_USER_ID}"
CREATE_RESP="$(curl -sS -X POST "$CREATE_URL" \
  -H "X-API-KEY: $ARCHON_API_KEY" \
  -H "Content-Type: application/json" \
  -d "$POST_BODY")"

# Check HTTP errors by retrying with -w if needed; here we assume non-error if jq can parse id.
CLUSTER_ID="$(printf '%s' "$CREATE_RESP" | jq -r '.id // empty')"
if [[ -z "$CLUSTER_ID" ]]; then
  echo "Create request failed or missing id. Full response:" >&2
  echo "$CREATE_RESP" | jq .
  exit 1
fi
echo "Submitted. cluster_id=$CLUSTER_ID"

# Poll for status up to 120 times with 10s sleep
STATUS_URL="${ARCHON_URL%/}/recipe/${CLUSTER_ID}?user_id=${ARCHON_USER_ID}"
ATTEMPTS=0
MAX_ATTEMPTS=120
SLEEP_SECS=10

while (( ATTEMPTS < MAX_ATTEMPTS )); do
  RESP="$(curl -sS -X GET "$STATUS_URL" -H "X-API-KEY: $ARCHON_API_KEY")"
  echo "$RESP"

  # Expect an array with at least two elements; second element is the prove step
  STATUS="$(
    printf '%s' "$RESP" | jq -r '
      def norm: (if type=="string" then fromjson? else . end) // .;
      # Normalize once
      norm
      # If array, pick step 1, else keep object
      | (if type=="array" then .[1]? else . end)
      # Extract status or return empty
      | (.status // empty)
    '
  )"
  if [[ -z "$STATUS" ]]; then
    echo "Unexpected status payload:" >&2
    echo "$RESP" | jq .
    exit 1
  fi

  echo "Attempt $((ATTEMPTS+1)) status=$STATUS"

  if [[ "$STATUS" == "Complete" ]]; then
    # Decode nested JSON from utf8_string then extract hex_proof and instances
    NESTED="$(
    printf '%s' "$RESP" | jq -r '
      def norm: (if type=="string" then fromjson? else . end) // .;
      norm
      | (if type=="array" then .[1]? else . end)
      | .output[0].utf8_string
      | fromjson
    '
  )"

    HEX_PROOF="$(printf '%s' "$NESTED" | jq -r '.hex_proof')"

    # Flatten inputs and outputs then concatenate
    INSTANCES_JSON="$(printf '%s' "$NESTED" | jq -c '
      ((.pretty_public_inputs.inputs // []) | flatten)
      + ((.pretty_public_inputs.outputs // []) | flatten)
    ')"

    # Print everything
    echo "$RESP"

    # Pretty output
    echo "hex_proof:"
    echo "$HEX_PROOF"
    echo
    echo "instances (JSON array):"
    echo "$INSTANCES_JSON"
    exit 0
  fi

  if [[ "$STATUS" == "Errored" ]]; then
    ERR_FILE="lilith_error_${LATEST_UUID}.json"
    printf '%s\n' "$RESP" | jq . > "$ERR_FILE"
    echo "Job errored. Saved details to $ERR_FILE" >&2
    exit 2
  fi

  ATTEMPTS=$((ATTEMPTS+1))
  sleep "$SLEEP_SECS"
done

echo "Timeout waiting for completion." >&2
exit 3