ppak10/Agentic-SLS-Telemetry

1#!/usr/bin/env python32"""Export new-DB builds into the export-format files with offset build IDs.3 4The Postgres DB was reset after build 40, so the new DB's build IDs restart at51. This script exports them with remapped IDs so they don't collide with the6existing dataset (builds 1-41 already processed).7 8Usage:9    uv run --with psycopg2-binary scripts/ticks/00_export_new_db.py 1:42 2:4310    uv run --with psycopg2-binary scripts/ticks/00_export_new_db.py 3:4411 12Each argument is old_db_id:new_dataset_id. Don't export an in-progress build13(no ended_at in the DB) — wait for it to finish and the spool to be imported.14 15Safe to re-run: telemetry/position_hf parquets are always overwritten; jsonl16files are checked first and skip if the new_id is already present.17"""18import json19import sys20from pathlib import Path21 22import psycopg223import psycopg2.extras24import pyarrow as pa25import pyarrow.parquet as pq26 27sys.path.insert(0, str(Path(__file__).parent.parent))28from _lib import EXPORTS_DIR29 30DATABASE_URL = "postgres://inova:inova@localhost:5432/inova"31BATCH_ROWS = 100_00032 33 34def get_conn():35    return psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.RealDictCursor)36 37 38def export_telemetry(conn, old_id: int, new_id: int):39    out = EXPORTS_DIR / "telemetry" / f"{new_id}.parquet"40    out.parent.mkdir(parents=True, exist_ok=True)41 42    schema = pa.schema([43        pa.field("ts",        pa.timestamp("us", tz="UTC")),44        pa.field("sensor_id", pa.string()),45        pa.field("kind",      pa.string()),46        pa.field("value",     pa.float64()),47    ])48 49    with conn.cursor() as cur:50        cur.execute(51            "SELECT COUNT(*) AS n FROM telemetry WHERE build_id = %s", (old_id,)52        )53        total = cur.fetchone()["n"]54        if total == 0:55            print(f"  [{old_id}→{new_id}] telemetry: 0 rows in DB, skipping")56            return57 58        cur.execute(59            "SELECT ts, sensor_id, kind, value FROM telemetry WHERE build_id = %s ORDER BY ts",60            (old_id,)61        )62        written = 063        with pq.ParquetWriter(out, schema, compression="zstd") as writer:64            while True:65                rows = cur.fetchmany(BATCH_ROWS)66                if not rows:67                    break68                table = pa.table({69                    "ts":        pa.array([r["ts"]        for r in rows], type=pa.timestamp("us", tz="UTC")),70                    "sensor_id": pa.array([r["sensor_id"] for r in rows], type=pa.string()),71                    "kind":      pa.array([r["kind"]      for r in rows], type=pa.string()),72                    "value":     pa.array([float(r["value"]) for r in rows], type=pa.float64()),73                }, schema=schema)74                writer.write_table(table)75                written += len(rows)76                print(f"  [{old_id}→{new_id}] telemetry: {written:,}/{total:,}...", end="\r")77        print(f"  [{old_id}→{new_id}] telemetry: {written:,} rows → {out.name}     ")78 79 80def export_position_hf(conn, old_id: int, new_id: int):81    out = EXPORTS_DIR / "position_hf" / f"{new_id}.parquet"82    out.parent.mkdir(parents=True, exist_ok=True)83 84    schema = pa.schema([85        pa.field("ts",        pa.timestamp("us", tz="UTC")),86        pa.field("x",         pa.float64()),87        pa.field("y",         pa.float64()),88        pa.field("z1",        pa.float64()),89        pa.field("z2",        pa.float64()),90        pa.field("r",         pa.float64()),91        pa.field("has_homed", pa.bool_()),92    ])93 94    with conn.cursor() as cur:95        cur.execute(96            "SELECT ts, x, y, z1, z2, r, has_homed FROM position_hf WHERE build_id = %s ORDER BY ts",97            (old_id,)98        )99        rows = cur.fetchall()100 101    if not rows:102        print(f"  [{old_id}→{new_id}] position_hf: 0 rows")103        return104 105    table = pa.table({106        "ts":        pa.array([r["ts"]        for r in rows], type=pa.timestamp("us", tz="UTC")),107        "x":         pa.array([float(r["x"])  for r in rows], type=pa.float64()),108        "y":         pa.array([float(r["y"])  for r in rows], type=pa.float64()),109        "z1":        pa.array([float(r["z1"]) for r in rows], type=pa.float64()),110        "z2":        pa.array([float(r["z2"]) for r in rows], type=pa.float64()),111        "r":         pa.array([float(r["r"])  for r in rows], type=pa.float64()),112        "has_homed": pa.array([bool(r["has_homed"]) for r in rows], type=pa.bool_()),113    }, schema=schema)114    pq.write_table(table, out, compression="zstd")115    print(f"  [{old_id}→{new_id}] position_hf: {len(rows):,} rows → {out.name}")116 117 118def append_builds_jsonl(conn, old_id: int, new_id: int):119    out = EXPORTS_DIR / "builds.jsonl"120 121    if out.exists():122        with out.open(encoding="utf-8") as f:123            for line in f:124                try:125                    r = json.loads(line)126                    if r.get("id") == new_id:127                        print(f"  [{old_id}→{new_id}] builds.jsonl: id={new_id} already present, skipping")128                        return129                except Exception:130                    pass131 132    with get_conn() as c, c.cursor() as cur:133        cur.execute("SELECT * FROM builds WHERE id = %s", (old_id,))134        row = cur.fetchone()135 136    if not row:137        print(f"  [{old_id}→{new_id}] builds.jsonl: build {old_id} not in DB, skipping")138        return139 140    entry = {141        "id":         new_id,142        "job_name":   row["job_name"],143        "started_at": row["started_at"].isoformat() if row["started_at"] else None,144        "ended_at":   row["ended_at"].isoformat()   if row["ended_at"]   else None,145        "phase":      row["phase"],146        "params":     row["params"],147        "notes":      (row["notes"] or "") + f" [originally DB id={old_id}, remapped to {new_id}]",148    }149    with out.open("a", encoding="utf-8") as f:150        f.write(json.dumps(entry) + "\n")151    print(f"  [{old_id}→{new_id}] builds.jsonl: appended id={new_id} ({entry['job_name']})")152 153 154def append_frames_jsonl(conn, old_id: int, new_id: int):155    out = EXPORTS_DIR / "frames.jsonl"156 157    if out.exists():158        with out.open(encoding="utf-8") as f:159            for line in f:160                try:161                    r = json.loads(line)162                    if r.get("build_id") == new_id:163                        print(f"  [{old_id}→{new_id}] frames.jsonl: build_id={new_id} already present, skipping")164                        return165                except Exception:166                    pass167 168    max_id = 0169    if out.exists():170        with out.open(encoding="utf-8") as f:171            for line in f:172                try:173                    r = json.loads(line)174                    max_id = max(max_id, int(r.get("id", 0)))175                except Exception:176                    pass177 178    with conn.cursor() as cur:179        cur.execute(180            "SELECT id, ts, kind, path FROM frames WHERE build_id = %s ORDER BY ts",181            (old_id,)182        )183        rows = cur.fetchall()184 185    if not rows:186        print(f"  [{old_id}→{new_id}] frames.jsonl: 0 frames in DB")187        return188 189    with out.open("a", encoding="utf-8") as f:190        for i, r in enumerate(rows):191            entry = {192                "id":       max_id + i + 1,193                "build_id": new_id,194                "ts":       r["ts"].isoformat(),195                "kind":     r["kind"],196                "path":     r["path"],197            }198            f.write(json.dumps(entry) + "\n")199    print(f"  [{old_id}→{new_id}] frames.jsonl: appended {len(rows):,} rows")200 201 202def append_events_jsonl(conn, old_id: int, new_id: int):203    """Append build_start event so load_build_to_profile_name() picks up the204    print profile name for this build."""205    out = EXPORTS_DIR / "events.jsonl"206 207    if out.exists():208        with out.open(encoding="utf-8") as f:209            for line in f:210                try:211                    r = json.loads(line)212                    if r.get("build_id") == new_id and r.get("kind") == "build_start":213                        print(f"  [{old_id}→{new_id}] events.jsonl: build_start already present, skipping")214                        return215                except Exception:216                    pass217 218    with conn.cursor() as cur:219        cur.execute(220            "SELECT * FROM events WHERE build_id = %s AND kind = 'build_start' LIMIT 1",221            (old_id,)222        )223        row = cur.fetchone()224 225    if not row:226        print(f"  [{old_id}→{new_id}] events.jsonl: no build_start event in DB")227        return228 229    entry = {230        "id":       -(new_id),231        "build_id": new_id,232        "ts":       row["ts"].isoformat() if row["ts"] else None,233        "kind":     "build_start",234        "message":  row["message"],235        "payload":  row["payload"],236    }237    with out.open("a", encoding="utf-8") as f:238        f.write(json.dumps(entry) + "\n")239    print(f"  [{old_id}→{new_id}] events.jsonl: appended build_start")240 241 242def main():243    if len(sys.argv) < 2:244        print(__doc__)245        sys.exit(1)246 247    mappings: list[tuple[int, int]] = []248    for arg in sys.argv[1:]:249        try:250            old_s, new_s = arg.split(":")251            mappings.append((int(old_s), int(new_s)))252        except ValueError:253            print(f"Bad argument {arg!r} — expected old_id:new_id")254            sys.exit(1)255 256    conn = get_conn()257    try:258        for old_id, new_id in mappings:259            print(f"\nExporting DB build {old_id} → dataset build {new_id}")260            export_telemetry(conn, old_id, new_id)261            export_position_hf(conn, old_id, new_id)262            append_builds_jsonl(conn, old_id, new_id)263            append_frames_jsonl(conn, old_id, new_id)264            append_events_jsonl(conn, old_id, new_id)265    finally:266        conn.close()267 268    new_ids = " ".join(str(n) for _, n in mappings)269    print(f"\nDone. Run:  uv run scripts/ticks/01_extract.py {new_ids}")270 271 272if __name__ == "__main__":273    main()274