1	#!/usr/bin/env python3
2	"""
3	Neucodec encoding worker.
4
5	Claims shards from Supabase, downloads audio.tar + metadata.parquet from R2,
6	encodes all audio with neucodec, uploads neucodec_tokens.parquet back to R2.
7
8	Usage:
9	    python worker.py                           # auto-generate worker ID
10	    python worker.py --worker-id my-worker-01  # explicit worker ID
11	    python worker.py --max-shards 5            # process at most 5 shards then exit
12	    python worker.py --dry-run                 # claim but don't process
13	"""
14
15	import argparse
16	import io
17	import logging
18	import os
19	import queue
20	import socket
21	import tarfile
22	import tempfile
23	import threading
24	import time
25	import uuid
26	from pathlib import Path
27
28	import boto3
29	import numpy as np
30	import pandas as pd
31	import psycopg2
32	import soundfile as sf
33	import torch
34
35	logging.basicConfig(
36	    level=logging.INFO,
37	    format="%(asctime)s [%(levelname)s] %(message)s",
38	    datefmt="%Y-%m-%d %H:%M:%S",
39	)
40	log = logging.getLogger("neucodec-worker")
41
42	# ── Config from environment ──────────────────────────────────────────────────
43	DATABASE_URL = os.environ["DATABASE_URL"]
44	R2_ENDPOINT_URL = os.environ["R2_ENDPOINT_URL"]
45	R2_ACCESS_KEY_ID = os.environ["R2_ACCESS_KEY_ID"]
46	R2_SECRET_ACCESS_KEY = os.environ["R2_SECRET_ACCESS_KEY"]
47	R2_BUCKET = os.environ.get("R2_BUCKET_SFT_DATA", "finalsftdata")
48
49	HEARTBEAT_INTERVAL = 60  # seconds
50	STALE_CLAIM_TIMEOUT = 3600  # 1 hour - reclaim shards stuck in 'claimed'/'processing'
51
52
53	def get_db():
54	    conn = psycopg2.connect(DATABASE_URL)
55	    conn.autocommit = True
56	    return conn
57
58
59	def get_s3():
60	    return boto3.client(
61	        "s3",
62	        endpoint_url=R2_ENDPOINT_URL,
63	        aws_access_key_id=R2_ACCESS_KEY_ID,
64	        aws_secret_access_key=R2_SECRET_ACCESS_KEY,
65	        region_name="auto",
66	    )
67
68
69	# ── Shard claim & status ────────────────────────────────────────────────────
70
71	def claim_shard(conn, worker_id: str):
72	    """Atomically claim the next pending shard. Returns (id, prefix, dataset, lang) or None."""
73	    cur = conn.cursor()
74	    cur.execute("SELECT * FROM claim_neucodec_shard(%s)", (worker_id,))
75	    row = cur.fetchone()
76	    cur.close()
77	    if row:
78	        return {"id": row[0], "prefix": row[1], "dataset": row[2], "language": row[3]}
79	    return None
80
81
82	def update_shard_status(conn, shard_id: int, status: str, **kwargs):
83	    sets = [f"status = %s"]
84	    vals = [status]
85	    for k, v in kwargs.items():
86	        sets.append(f"{k} = %s")
87	        vals.append(v)
88	    vals.append(shard_id)
89	    cur = conn.cursor()
90	    cur.execute(f"UPDATE neucodec_shards SET {', '.join(sets)} WHERE id = %s", vals)
91	    cur.close()
92
93
94	def heartbeat_loop(conn, shard_id: int, stop_event: threading.Event):
95	    """Periodically update claimed_at so other workers know we're alive."""
96	    while not stop_event.wait(HEARTBEAT_INTERVAL):
97	        try:
98	            cur = conn.cursor()
99	            cur.execute(
100	                "UPDATE neucodec_shards SET claimed_at = NOW() WHERE id = %s",
101	                (shard_id,),
102	            )
103	            cur.close()
104	        except Exception:
105	            pass
106
107
108	def reclaim_stale(conn):
109	    """Reset shards that have been claimed/processing too long without heartbeat."""
110	    cur = conn.cursor()
111	    cur.execute(
112	        """
113	        UPDATE neucodec_shards
114	        SET status = 'pending', worker_id = NULL, claimed_at = NULL
115	        WHERE status IN ('claimed', 'processing')
116	          AND claimed_at < NOW() - INTERVAL '%s seconds'
117	        RETURNING id, shard_prefix
118	        """,
119	        (STALE_CLAIM_TIMEOUT,),
120	    )
121	    reclaimed = cur.fetchall()
122	    cur.close()
123	    if reclaimed:
124	        log.warning("Reclaimed %d stale shards: %s", len(reclaimed), [r[1] for r in reclaimed])
125
126
127	# ── Download ─────────────────────────────────────────────────────────────────
128
129	def download_shard(s3, shard_prefix: str, work_dir: Path):
130	    """Download audio.tar and metadata.parquet for a shard."""
131	    tar_path = work_dir / "audio.tar"
132	    meta_path = work_dir / "metadata.parquet"
133
134	    tar_key = f"{shard_prefix}/audio.tar"
135	    meta_key = f"{shard_prefix}/metadata.parquet"
136
137	    log.info("Downloading %s ...", tar_key)
138	    t0 = time.time()
139
140	    # Multi-part download for large tars
141	    from boto3.s3.transfer import TransferConfig
142	    config = TransferConfig(
143	        multipart_threshold=64 * 1024 * 1024,  # 64MB
144	        max_concurrency=8,
145	        multipart_chunksize=64 * 1024 * 1024,
146	    )
147	    s3.download_file(R2_BUCKET, tar_key, str(tar_path), Config=config)
148	    s3.download_file(R2_BUCKET, meta_key, str(meta_path))
149
150	    tar_size = tar_path.stat().st_size
151	    dl_time = time.time() - t0
152	    speed = tar_size / dl_time / 1024 / 1024
153	    log.info("Downloaded %.1f GB in %.0fs (%.1f MB/s)", tar_size / 1e9, dl_time, speed)
154
155	    return tar_path, meta_path, tar_size
156
157
158	# ── Encode ───────────────────────────────────────────────────────────────────
159
160	MAX_AUDIO_SAMPLES = 30 * 16000  # 30s at 16kHz — chunk longer clips to avoid OOM on 24GB GPUs
161
162
163	def load_codec():
164	    """Load neucodec model to GPU with layer truncation optimization."""
165	    from neucodec import NeuCodec
166	    import torch.nn as nn
167
168	    log.info("Loading neucodec model...")
169	    t0 = time.time()
170	    codec = NeuCodec.from_pretrained("neuphonic/neucodec")
171
172	    # FIX: Upstream bug — line 34 of neucodec/model.py does `self.semantic_model.encoder.layers[:16]`
173	    # which is a no-op. The model only reads hidden_states[16], so layers 17-23 are wasted compute.
174	    # Truncating to 17 layers saves ~18% encode time and ~0.6GB VRAM, with bit-identical output.
175	    orig_layers = len(codec.semantic_model.encoder.layers)
176	    codec.semantic_model.encoder.layers = nn.ModuleList(
177	        list(codec.semantic_model.encoder.layers)[:17]
178	    )
179	    log.info("Truncated semantic model: %d → %d layers", orig_layers, len(codec.semantic_model.encoder.layers))
180
181	    _ = codec.eval().to("cuda")
182	    log.info("Neucodec loaded in %.1fs, VRAM: %.1f GB", time.time() - t0, torch.cuda.memory_allocated() / 1e9)
183	    return codec
184
185
186	def encode_shard(codec, tar_path: Path, meta_path: Path):
187	    """
188	    Encode all audio in a shard tar using pipelined CPU fbank + GPU encode.
189	    Returns a DataFrame with segment_id, neucodec_tokens (bytes), token_count.
190	    """
191	    metadata = pd.read_parquet(meta_path)
192
193	    tf = tarfile.open(str(tar_path), "r")
194	    members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
195	    log.info("Encoding %d audio files...", len(members))
196
197	    # Build segment_id lookup from tar member names
198	    # member.name is like 'iv_hi_00000000.flac' -> segment_id = 'iv_hi_00000000'
199	    member_to_seg = {}
200	    for m in members:
201	        seg_id = m.name.replace(".flac", "").split("/")[-1]
202	        member_to_seg[m.name] = seg_id
203
204	    # Pipelined: CPU fbank extraction in background thread
205	    fbank_queue = queue.Queue(maxsize=16)
206
207	    def fbank_worker():
208	        for member in members:
209	            try:
210	                f = tf.extractfile(member)
211	                audio_data, sr = sf.read(io.BytesIO(f.read()))
212	                # Chunk long audio to avoid OOM (30s max per chunk)
213	                if len(audio_data) > MAX_AUDIO_SAMPLES:
214	                    chunks = []
215	                    for start in range(0, len(audio_data), MAX_AUDIO_SAMPLES):
216	                        chunks.append(audio_data[start : start + MAX_AUDIO_SAMPLES])
217	                else:
218	                    chunks = [audio_data]
219
220	                seg_id = member_to_seg[member.name]
221	                chunk_feats = []
222	                for chunk in chunks:
223	                    y = torch.tensor(chunk, dtype=torch.float32)[None, None]
224	                    feats = codec.feature_extractor(
225	                        y[0, :].cpu(), sampling_rate=16_000, return_tensors="pt"
226	                    ).input_features
227	                    chunk_feats.append((y, feats))
228	                fbank_queue.put((seg_id, chunk_feats, len(audio_data) / sr))
229	            except Exception as e:
230	                seg_id = member_to_seg.get(member.name, member.name)
231	                log.warning("Failed to read/extract %s: %s", seg_id, e)
232	                fbank_queue.put((seg_id, None, 0))
233	        fbank_queue.put(None)  # sentinel
234
235	    cpu_thread = threading.Thread(target=fbank_worker, daemon=True)
236	    cpu_thread.start()
237
238	    results = []
239	    total_audio_s = 0.0
240	    total_tokens = 0
241	    processed = 0
242	    failed = 0
243	    t_start = time.time()
244
245	    while True:
246	        item = fbank_queue.get()
247	        if item is None:
248	            break
249
250	        seg_id, chunk_feats, dur = item
251	        if chunk_feats is None:
252	            failed += 1
253	            continue
254
255	        try:
256	            all_codes = []
257	            for y, feats in chunk_feats:
258	                feats_gpu = feats.to("cuda")
259	                y_gpu = y.to("cuda")
260
261	                with torch.no_grad():
262	                    acoustic_emb = codec.CodecEnc(y_gpu).transpose(1, 2)
263	                    semantic_output = (
264	                        codec.semantic_model(feats_gpu).hidden_states[16].transpose(1, 2)
265	                    )
266	                    semantic_encoded = codec.SemanticEncoder_module(semantic_output)
267
268	                    if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
269	                        ml = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
270	                        acoustic_emb = acoustic_emb[:, :, :ml]
271	                        semantic_encoded = semantic_encoded[:, :, :ml]
272
273	                    concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
274	                    concat_emb = codec.fc_prior(
275	                        concat_emb.transpose(1, 2)
276	                    ).transpose(1, 2)
277	                    _, codes, _ = codec.generator(concat_emb, vq=True)
278
279	                all_codes.append(codes[0, 0].cpu().numpy().astype(np.uint16))
280
281	                # Free GPU tensors between chunks
282	                del feats_gpu, y_gpu, acoustic_emb, semantic_output, semantic_encoded
283	                del concat_emb, codes
284
285	            token_array = np.concatenate(all_codes) if len(all_codes) > 1 else all_codes[0]
286	            results.append(
287	                {
288	                    "segment_id": seg_id,
289	                    "neucodec_tokens": token_array.tobytes(),
290	                    "token_count": len(token_array),
291	                }
292	            )
293	            total_audio_s += dur
294	            total_tokens += len(token_array)
295	        except Exception as e:
296	            log.warning("Encode failed for %s: %s", seg_id, e)
297	            failed += 1
298
299	        # Periodic VRAM cleanup to prevent fragmentation
300	        processed += 1
301	        if processed % 500 == 0:
302	            torch.cuda.empty_cache()
303	        if processed % 3000 == 0:
304	            elapsed = time.time() - t_start
305	            rtf = total_audio_s / elapsed if elapsed > 0 else 0
306	            vram_gb = torch.cuda.memory_allocated() / 1e9
307	            log.info(
308	                "[%d/%d] %.0f%% done, %.0fs elapsed, RTF=%.1fx, VRAM=%.1fGB",
309	                processed, len(members), processed / len(members) * 100, elapsed, rtf, vram_gb,
310	            )
311
312	    cpu_thread.join()
313	    tf.close()
314
315	    wall = time.time() - t_start
316	    rtf = total_audio_s / wall if wall > 0 else 0
317	    log.info(
318	        "Encoding done: %d segments, %d failed, %.1f hrs audio, %.0fs wall, RTF=%.1fx, %d tokens",
319	        processed, failed, total_audio_s / 3600, wall, rtf, total_tokens,
320	    )
321
322	    df = pd.DataFrame(results)
323	    return df, total_audio_s, total_tokens, wall
324
325
326	# ── Upload ───────────────────────────────────────────────────────────────────
327
328	def upload_results(s3, shard_prefix: str, tokens_df: pd.DataFrame, work_dir: Path):
329	    """Upload neucodec_tokens.parquet to the same shard directory in R2."""
330	    out_path = work_dir / "neucodec_tokens.parquet"
331	    tokens_df.to_parquet(str(out_path), index=False)
332
333	    upload_key = f"{shard_prefix}/neucodec_tokens.parquet"
334	    file_size = out_path.stat().st_size
335
336	    log.info("Uploading %s (%.1f MB)...", upload_key, file_size / 1e6)
337	    s3.upload_file(str(out_path), R2_BUCKET, upload_key)
338	    log.info("Upload complete: %s", upload_key)
339
340	    return upload_key
341
342
343	# ── Main loop ────────────────────────────────────────────────────────────────
344
345	def process_one_shard(codec, s3, conn, shard: dict, worker_id: str):
346	    """Full pipeline for one shard: download → encode → upload → update."""
347	    shard_id = shard["id"]
348	    prefix = shard["prefix"]
349
350	    log.info("═" * 60)
351	    log.info("Processing shard %d: %s", shard_id, prefix)
352	    log.info("═" * 60)
353
354	    # Start heartbeat
355	    stop_hb = threading.Event()
356	    hb_conn = get_db()
357	    hb_thread = threading.Thread(
358	        target=heartbeat_loop, args=(hb_conn, shard_id, stop_hb), daemon=True
359	    )
360	    hb_thread.start()
361
362	    try:
363	        update_shard_status(conn, shard_id, "processing", started_at="NOW()")
364
365	        with tempfile.TemporaryDirectory(prefix="neucodec_") as work_dir:
366	            work_dir = Path(work_dir)
367
368	            # Download
369	            tar_path, meta_path, tar_size = download_shard(s3, prefix, work_dir)
370
371	            # Encode
372	            tokens_df, total_audio_s, total_tokens, encode_wall = encode_shard(
373	                codec, tar_path, meta_path
374	            )
375
376	            if tokens_df.empty:
377	                raise RuntimeError("No segments encoded successfully")
378
379	            # Upload
380	            output_key = upload_results(s3, prefix, tokens_df, work_dir)
381
382	            # Mark complete
383	            update_shard_status(
384	                conn,
385	                shard_id,
386	                "completed",
387	                completed_at="NOW()",
388	                audio_tar_bytes=tar_size,
389	                segment_count=len(tokens_df),
390	                total_audio_seconds=total_audio_s,
391	                total_tokens=total_tokens,
392	                encode_wall_seconds=encode_wall,
393	                output_key=output_key,
394	            )
395	            log.info("Shard %d completed successfully", shard_id)
396
397	    except Exception as e:
398	        log.exception("Shard %d failed: %s", shard_id, e)
399	        try:
400	            update_shard_status(
401	                conn, shard_id, "failed", failed_at="NOW()", error_message=str(e)[:500]
402	            )
403	        except Exception:
404	            pass
405
406	    finally:
407	        stop_hb.set()
408	        hb_thread.join(timeout=5)
409	        hb_conn.close()
410
411
412	def main():
413	    parser = argparse.ArgumentParser(description="Neucodec encoding worker")
414	    parser.add_argument("--worker-id", default=None, help="Worker identifier")
415	    parser.add_argument("--max-shards", type=int, default=0, help="Max shards to process (0=unlimited)")
416	    parser.add_argument("--dry-run", action="store_true", help="Claim and print, don't process")
417	    args = parser.parse_args()
418
419	    worker_id = args.worker_id or f"{socket.gethostname()}-{uuid.uuid4().hex[:8]}"
420	    log.info("Worker starting: %s", worker_id)
421	    log.info("GPU: %s", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE")
422
423	    # Load codec once
424	    codec = load_codec()
425
426	    # Warmup
427	    with torch.no_grad():
428	        codec.encode_code(torch.randn(1, 1, 16000).cuda())
429	    log.info("Codec warmup complete")
430
431	    s3 = get_s3()
432	    conn = get_db()
433
434	    # Reclaim any stale shards on startup
435	    reclaim_stale(conn)
436
437	    shards_done = 0
438	    while True:
439	        shard = claim_shard(conn, worker_id)
440	        if shard is None:
441	            log.info("No more pending shards. Worker exiting.")
442	            break
443
444	        if args.dry_run:
445	            log.info("DRY RUN: would process shard %d: %s", shard["id"], shard["prefix"])
446	            update_shard_status(conn, shard["id"], "pending", worker_id=None, claimed_at=None)
447	            shards_done += 1
448	            if args.max_shards > 0 and shards_done >= args.max_shards:
449	                break
450	            continue
451
452	        process_one_shard(codec, s3, conn, shard, worker_id)
453	        shards_done += 1
454
455	        if args.max_shards > 0 and shards_done >= args.max_shards:
456	            log.info("Reached max_shards=%d, exiting.", args.max_shards)
457	            break
458
459	    conn.close()
460	    log.info("Worker %s finished after %d shards", worker_id, shards_done)
461
462
463	if __name__ == "__main__":
464	    main()
465

#!/usr/bin/env python3
"""
Neucodec encoding worker.

Claims shards from Supabase, downloads audio.tar + metadata.parquet from R2,
encodes all audio with neucodec, uploads neucodec_tokens.parquet back to R2.

Usage:
    python worker.py                           # auto-generate worker ID
    python worker.py --worker-id my-worker-01  # explicit worker ID
    python worker.py --max-shards 5            # process at most 5 shards then exit
    python worker.py --dry-run                 # claim but don't process
"""

import argparse
import io
import logging
import os
import queue
import socket
import tarfile
import tempfile
import threading
import time
import uuid
from pathlib import Path

import boto3
import numpy as np
import pandas as pd
import psycopg2
import soundfile as sf
import torch

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("neucodec-worker")

# ── Config from environment ──────────────────────────────────────────────────
DATABASE_URL = os.environ["DATABASE_URL"]
R2_ENDPOINT_URL = os.environ["R2_ENDPOINT_URL"]
R2_ACCESS_KEY_ID = os.environ["R2_ACCESS_KEY_ID"]
R2_SECRET_ACCESS_KEY = os.environ["R2_SECRET_ACCESS_KEY"]
R2_BUCKET = os.environ.get("R2_BUCKET_SFT_DATA", "finalsftdata")

HEARTBEAT_INTERVAL = 60  # seconds
STALE_CLAIM_TIMEOUT = 3600  # 1 hour - reclaim shards stuck in 'claimed'/'processing'

def get_db():
    conn = psycopg2.connect(DATABASE_URL)
    conn.autocommit = True
    return conn

def get_s3():
    return boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=R2_ACCESS_KEY_ID,
        aws_secret_access_key=R2_SECRET_ACCESS_KEY,
        region_name="auto",
    )

# ── Shard claim & status ────────────────────────────────────────────────────

def claim_shard(conn, worker_id: str):
    """Atomically claim the next pending shard. Returns (id, prefix, dataset, lang) or None."""
    cur = conn.cursor()
    cur.execute("SELECT * FROM claim_neucodec_shard(%s)", (worker_id,))
    row = cur.fetchone()
    cur.close()
    if row:
        return {"id": row[0], "prefix": row[1], "dataset": row[2], "language": row[3]}
    return None

def update_shard_status(conn, shard_id: int, status: str, **kwargs):
    sets = [f"status = %s"]
    vals = [status]
    for k, v in kwargs.items():
        sets.append(f"{k} = %s")
        vals.append(v)
    vals.append(shard_id)
    cur = conn.cursor()
    cur.execute(f"UPDATE neucodec_shards SET {', '.join(sets)} WHERE id = %s", vals)
    cur.close()

def heartbeat_loop(conn, shard_id: int, stop_event: threading.Event):
    """Periodically update claimed_at so other workers know we're alive."""
    while not stop_event.wait(HEARTBEAT_INTERVAL):
        try:
            cur = conn.cursor()
            cur.execute(
                "UPDATE neucodec_shards SET claimed_at = NOW() WHERE id = %s",
                (shard_id,),
            )
            cur.close()
        except Exception:
            pass

def reclaim_stale(conn):
    """Reset shards that have been claimed/processing too long without heartbeat."""
    cur = conn.cursor()
    cur.execute(
        """
        UPDATE neucodec_shards
        SET status = 'pending', worker_id = NULL, claimed_at = NULL
        WHERE status IN ('claimed', 'processing')
          AND claimed_at < NOW() - INTERVAL '%s seconds'
        RETURNING id, shard_prefix
        """,
        (STALE_CLAIM_TIMEOUT,),
    )
    reclaimed = cur.fetchall()
    cur.close()
    if reclaimed:
        log.warning("Reclaimed %d stale shards: %s", len(reclaimed), [r[1] for r in reclaimed])

# ── Download ─────────────────────────────────────────────────────────────────

def download_shard(s3, shard_prefix: str, work_dir: Path):
    """Download audio.tar and metadata.parquet for a shard."""
    tar_path = work_dir / "audio.tar"
    meta_path = work_dir / "metadata.parquet"

    tar_key = f"{shard_prefix}/audio.tar"
    meta_key = f"{shard_prefix}/metadata.parquet"

    log.info("Downloading %s ...", tar_key)
    t0 = time.time()

    # Multi-part download for large tars
    from boto3.s3.transfer import TransferConfig
    config = TransferConfig(
        multipart_threshold=64 * 1024 * 1024,  # 64MB
        max_concurrency=8,
        multipart_chunksize=64 * 1024 * 1024,
    )
    s3.download_file(R2_BUCKET, tar_key, str(tar_path), Config=config)
    s3.download_file(R2_BUCKET, meta_key, str(meta_path))

    tar_size = tar_path.stat().st_size
    dl_time = time.time() - t0
    speed = tar_size / dl_time / 1024 / 1024
    log.info("Downloaded %.1f GB in %.0fs (%.1f MB/s)", tar_size / 1e9, dl_time, speed)

    return tar_path, meta_path, tar_size

# ── Encode ───────────────────────────────────────────────────────────────────

MAX_AUDIO_SAMPLES = 30 * 16000  # 30s at 16kHz — chunk longer clips to avoid OOM on 24GB GPUs

def load_codec():
    """Load neucodec model to GPU with layer truncation optimization."""
    from neucodec import NeuCodec
    import torch.nn as nn

    log.info("Loading neucodec model...")
    t0 = time.time()
    codec = NeuCodec.from_pretrained("neuphonic/neucodec")

    # FIX: Upstream bug — line 34 of neucodec/model.py does `self.semantic_model.encoder.layers[:16]`
    # which is a no-op. The model only reads hidden_states[16], so layers 17-23 are wasted compute.
    # Truncating to 17 layers saves ~18% encode time and ~0.6GB VRAM, with bit-identical output.
    orig_layers = len(codec.semantic_model.encoder.layers)
    codec.semantic_model.encoder.layers = nn.ModuleList(
        list(codec.semantic_model.encoder.layers)[:17]
    )
    log.info("Truncated semantic model: %d → %d layers", orig_layers, len(codec.semantic_model.encoder.layers))

    _ = codec.eval().to("cuda")
    log.info("Neucodec loaded in %.1fs, VRAM: %.1f GB", time.time() - t0, torch.cuda.memory_allocated() / 1e9)
    return codec

def encode_shard(codec, tar_path: Path, meta_path: Path):
    """
    Encode all audio in a shard tar using pipelined CPU fbank + GPU encode.
    Returns a DataFrame with segment_id, neucodec_tokens (bytes), token_count.
    """
    metadata = pd.read_parquet(meta_path)

    tf = tarfile.open(str(tar_path), "r")
    members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
    log.info("Encoding %d audio files...", len(members))

    # Build segment_id lookup from tar member names
    # member.name is like 'iv_hi_00000000.flac' -> segment_id = 'iv_hi_00000000'
    member_to_seg = {}
    for m in members:
        seg_id = m.name.replace(".flac", "").split("/")[-1]
        member_to_seg[m.name] = seg_id

    # Pipelined: CPU fbank extraction in background thread
    fbank_queue = queue.Queue(maxsize=16)

    def fbank_worker():
        for member in members:
            try:
                f = tf.extractfile(member)
                audio_data, sr = sf.read(io.BytesIO(f.read()))
                # Chunk long audio to avoid OOM (30s max per chunk)
                if len(audio_data) > MAX_AUDIO_SAMPLES:
                    chunks = []
                    for start in range(0, len(audio_data), MAX_AUDIO_SAMPLES):
                        chunks.append(audio_data[start : start + MAX_AUDIO_SAMPLES])
                else:
                    chunks = [audio_data]

                seg_id = member_to_seg[member.name]
                chunk_feats = []
                for chunk in chunks:
                    y = torch.tensor(chunk, dtype=torch.float32)[None, None]
                    feats = codec.feature_extractor(
                        y[0, :].cpu(), sampling_rate=16_000, return_tensors="pt"
                    ).input_features
                    chunk_feats.append((y, feats))
                fbank_queue.put((seg_id, chunk_feats, len(audio_data) / sr))
            except Exception as e:
                seg_id = member_to_seg.get(member.name, member.name)
                log.warning("Failed to read/extract %s: %s", seg_id, e)
                fbank_queue.put((seg_id, None, 0))
        fbank_queue.put(None)  # sentinel

    cpu_thread = threading.Thread(target=fbank_worker, daemon=True)
    cpu_thread.start()

    results = []
    total_audio_s = 0.0
    total_tokens = 0
    processed = 0
    failed = 0
    t_start = time.time()

    while True:
        item = fbank_queue.get()
        if item is None:
            break

        seg_id, chunk_feats, dur = item
        if chunk_feats is None:
            failed += 1
            continue

        try:
            all_codes = []
            for y, feats in chunk_feats:
                feats_gpu = feats.to("cuda")
                y_gpu = y.to("cuda")

                with torch.no_grad():
                    acoustic_emb = codec.CodecEnc(y_gpu).transpose(1, 2)
                    semantic_output = (
                        codec.semantic_model(feats_gpu).hidden_states[16].transpose(1, 2)
                    )
                    semantic_encoded = codec.SemanticEncoder_module(semantic_output)

                    if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
                        ml = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
                        acoustic_emb = acoustic_emb[:, :, :ml]
                        semantic_encoded = semantic_encoded[:, :, :ml]

                    concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
                    concat_emb = codec.fc_prior(
                        concat_emb.transpose(1, 2)
                    ).transpose(1, 2)
                    _, codes, _ = codec.generator(concat_emb, vq=True)

                all_codes.append(codes[0, 0].cpu().numpy().astype(np.uint16))

                # Free GPU tensors between chunks
                del feats_gpu, y_gpu, acoustic_emb, semantic_output, semantic_encoded
                del concat_emb, codes

            token_array = np.concatenate(all_codes) if len(all_codes) > 1 else all_codes[0]
            results.append(
                {
                    "segment_id": seg_id,
                    "neucodec_tokens": token_array.tobytes(),
                    "token_count": len(token_array),
                }
            )
            total_audio_s += dur
            total_tokens += len(token_array)
        except Exception as e:
            log.warning("Encode failed for %s: %s", seg_id, e)
            failed += 1

        # Periodic VRAM cleanup to prevent fragmentation
        processed += 1
        if processed % 500 == 0:
            torch.cuda.empty_cache()
        if processed % 3000 == 0:
            elapsed = time.time() - t_start
            rtf = total_audio_s / elapsed if elapsed > 0 else 0
            vram_gb = torch.cuda.memory_allocated() / 1e9
            log.info(
                "[%d/%d] %.0f%% done, %.0fs elapsed, RTF=%.1fx, VRAM=%.1fGB",
                processed, len(members), processed / len(members) * 100, elapsed, rtf, vram_gb,
            )

    cpu_thread.join()
    tf.close()

    wall = time.time() - t_start
    rtf = total_audio_s / wall if wall > 0 else 0
    log.info(
        "Encoding done: %d segments, %d failed, %.1f hrs audio, %.0fs wall, RTF=%.1fx, %d tokens",
        processed, failed, total_audio_s / 3600, wall, rtf, total_tokens,
    )

    df = pd.DataFrame(results)
    return df, total_audio_s, total_tokens, wall

# ── Upload ───────────────────────────────────────────────────────────────────

def upload_results(s3, shard_prefix: str, tokens_df: pd.DataFrame, work_dir: Path):
    """Upload neucodec_tokens.parquet to the same shard directory in R2."""
    out_path = work_dir / "neucodec_tokens.parquet"
    tokens_df.to_parquet(str(out_path), index=False)

    upload_key = f"{shard_prefix}/neucodec_tokens.parquet"
    file_size = out_path.stat().st_size

    log.info("Uploading %s (%.1f MB)...", upload_key, file_size / 1e6)
    s3.upload_file(str(out_path), R2_BUCKET, upload_key)
    log.info("Upload complete: %s", upload_key)

    return upload_key

# ── Main loop ────────────────────────────────────────────────────────────────

def process_one_shard(codec, s3, conn, shard: dict, worker_id: str):
    """Full pipeline for one shard: download → encode → upload → update."""
    shard_id = shard["id"]
    prefix = shard["prefix"]

    log.info("═" * 60)
    log.info("Processing shard %d: %s", shard_id, prefix)
    log.info("═" * 60)

    # Start heartbeat
    stop_hb = threading.Event()
    hb_conn = get_db()
    hb_thread = threading.Thread(
        target=heartbeat_loop, args=(hb_conn, shard_id, stop_hb), daemon=True
    )
    hb_thread.start()

    try:
        update_shard_status(conn, shard_id, "processing", started_at="NOW()")

        with tempfile.TemporaryDirectory(prefix="neucodec_") as work_dir:
            work_dir = Path(work_dir)

            # Download
            tar_path, meta_path, tar_size = download_shard(s3, prefix, work_dir)

            # Encode
            tokens_df, total_audio_s, total_tokens, encode_wall = encode_shard(
                codec, tar_path, meta_path
            )

            if tokens_df.empty:
                raise RuntimeError("No segments encoded successfully")

            # Upload
            output_key = upload_results(s3, prefix, tokens_df, work_dir)

            # Mark complete
            update_shard_status(
                conn,
                shard_id,
                "completed",
                completed_at="NOW()",
                audio_tar_bytes=tar_size,
                segment_count=len(tokens_df),
                total_audio_seconds=total_audio_s,
                total_tokens=total_tokens,
                encode_wall_seconds=encode_wall,
                output_key=output_key,
            )
            log.info("Shard %d completed successfully", shard_id)

    except Exception as e:
        log.exception("Shard %d failed: %s", shard_id, e)
        try:
            update_shard_status(
                conn, shard_id, "failed", failed_at="NOW()", error_message=str(e)[:500]
            )
        except Exception:
            pass

    finally:
        stop_hb.set()
        hb_thread.join(timeout=5)
        hb_conn.close()

def main():
    parser = argparse.ArgumentParser(description="Neucodec encoding worker")
    parser.add_argument("--worker-id", default=None, help="Worker identifier")
    parser.add_argument("--max-shards", type=int, default=0, help="Max shards to process (0=unlimited)")
    parser.add_argument("--dry-run", action="store_true", help="Claim and print, don't process")
    args = parser.parse_args()

    worker_id = args.worker_id or f"{socket.gethostname()}-{uuid.uuid4().hex[:8]}"
    log.info("Worker starting: %s", worker_id)
    log.info("GPU: %s", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE")

    # Load codec once
    codec = load_codec()

    # Warmup
    with torch.no_grad():
        codec.encode_code(torch.randn(1, 1, 16000).cuda())
    log.info("Codec warmup complete")

    s3 = get_s3()
    conn = get_db()

    # Reclaim any stale shards on startup
    reclaim_stale(conn)

    shards_done = 0
    while True:
        shard = claim_shard(conn, worker_id)
        if shard is None:
            log.info("No more pending shards. Worker exiting.")
            break

        if args.dry_run:
            log.info("DRY RUN: would process shard %d: %s", shard["id"], shard["prefix"])
            update_shard_status(conn, shard["id"], "pending", worker_id=None, claimed_at=None)
            shards_done += 1
            if args.max_shards > 0 and shards_done >= args.max_shards:
                break
            continue

        process_one_shard(codec, s3, conn, shard, worker_id)
        shards_done += 1

        if args.max_shards > 0 and shards_done >= args.max_shards:
            log.info("Reached max_shards=%d, exiting.", args.max_shards)
            break

    conn.close()
    log.info("Worker %s finished after %d shards", worker_id, shards_done)

if __name__ == "__main__":
    main()


Now I'll rewrite worker.py with integrated worker-level heartbeat reporting, and update status.py accordingly.


Write

#!/usr/bin/env python3
"""
Neucodec encoding worker.

Claims shards from Supabase, downloads audio.tar + metadata.parquet from R2,
encodes all audio with neucodec, uploads neucodec_tokens.parquet back to R2.

Reports per-worker state to neucodec_workers table for monitoring.

Usage:
    python worker.py                           # auto-generate worker ID
    python worker.py --worker-id my-worker-01  # explicit worker ID
    python worker.py --max-shards 5            # process at most 5 shards then exit
    python worker.py --dry-run                 # claim but don't process
"""

import argparse
import io
import logging
import os
import queue
import socket
import tarfile
import tempfile
import threading
import time
import uuid
from pathlib import Path

import boto3
import numpy as np
import pandas as pd
import psycopg2
import soundfile as sf
import torch

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("neucodec-worker")

# ── Config from environment ──────────────────────────────────────────────────
DATABASE_URL = os.environ["DATABASE_URL"]
R2_ENDPOINT_URL = os.environ["R2_ENDPOINT_URL"]
R2_ACCESS_KEY_ID = os.environ["R2_ACCESS_KEY_ID"]
R2_SECRET_ACCESS_KEY = os.environ["R2_SECRET_ACCESS_KEY"]
R2_BUCKET = os.environ.get("R2_BUCKET_SFT_DATA", "finalsftdata")

HEARTBEAT_INTERVAL = 30  # seconds
STALE_CLAIM_TIMEOUT = 600  # 10 min - reclaim shards from dead workers
MAX_AUDIO_SAMPLES = 30 * 16000  # 30s at 16kHz — chunk longer clips to avoid OOM

def get_db():
    conn = psycopg2.connect(DATABASE_URL)
    conn.autocommit = True
    return conn

def get_s3():
    return boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=R2_ACCESS_KEY_ID,
        aws_secret_access_key=R2_SECRET_ACCESS_KEY,
        region_name="auto",
    )

# ── Worker state (shared mutable state for heartbeat thread) ─────────────────

class WorkerState:
    """Thread-safe worker state for heartbeat reporting."""

    def __init__(self, worker_id: str, gpu_name: str, gpu_vram_mb: int):
        self.worker_id = worker_id
        self.hostname = socket.gethostname()
        self.gpu_name = gpu_name
        self.gpu_vram_mb = gpu_vram_mb
        self.lock = threading.Lock()
        # Mutable fields updated by main thread, read by heartbeat thread
        self.state = "idle"
        self.current_shard_id = None
        self.current_shard_prefix = None
        self.shards_completed = 0
        self.shards_failed = 0
        self.total_audio_hours = 0.0
        self.total_tokens = 0
        self.progress_pct = 0.0
        self.progress_segments = 0
        self.progress_total_segments = 0
        self.progress_rtf = 0.0
        self.progress_vram_gb = 0.0
        self.progress_phase = "idle"
        self.last_error = None
        self.last_error_at = None

    def set_phase(self, phase: str, shard_id=None, shard_prefix=None):
        with self.lock:
            self.state = "working"
            self.progress_phase = phase
            if shard_id is not None:
                self.current_shard_id = shard_id
            if shard_prefix is not None:
                self.current_shard_prefix = shard_prefix

    def set_progress(self, processed: int, total: int, rtf: float):
        with self.lock:
            self.progress_segments = processed
            self.progress_total_segments = total
            self.progress_pct = processed / total * 100 if total > 0 else 0
            self.progress_rtf = rtf
            self.progress_vram_gb = torch.cuda.memory_allocated() / 1e9

    def mark_shard_done(self, audio_hours: float, tokens: int):
        with self.lock:
            self.shards_completed += 1
            self.total_audio_hours += audio_hours
            self.total_tokens += tokens
            self.state = "idle"
            self.current_shard_id = None
            self.current_shard_prefix = None
            self.progress_pct = 0
            self.progress_phase = "idle"

    def mark_shard_failed(self, error: str):
        with self.lock:
            self.shards_failed += 1
            self.last_error = error[:500]
            self.last_error_at = "NOW()"
            self.state = "idle"
            self.current_shard_id = None
            self.current_shard_prefix = None
            self.progress_pct = 0
            self.progress_phase = "idle"

    def snapshot(self):
        with self.lock:
            return {
                "worker_id": self.worker_id,
                "hostname": self.hostname,
                "gpu_name": self.gpu_name,
                "gpu_vram_mb": self.gpu_vram_mb,
                "state": self.state,
                "current_shard_id": self.current_shard_id,
                "current_shard_prefix": self.current_shard_prefix,
                "shards_completed": self.shards_completed,
                "shards_failed": self.shards_failed,
                "total_audio_hours": self.total_audio_hours,
                "total_tokens": self.total_tokens,
                "progress_pct": self.progress_pct,
                "progress_segments": self.progress_segments,
                "progress_total_segments": self.progress_total_segments,
                "progress_rtf": self.progress_rtf,
                "progress_vram_gb": self.progress_vram_gb,
                "progress_phase": self.progress_phase,
                "last_error": self.last_error,
            }

# ── Heartbeat ────────────────────────────────────────────────────────────────

def heartbeat_loop(wstate: WorkerState, stop_event: threading.Event):
    """Periodically write worker state to both neucodec_workers and neucodec_shards."""
    while not stop_event.wait(HEARTBEAT_INTERVAL):
        try:
            snap = wstate.snapshot()
            conn = get_db()
            cur = conn.cursor()

            # Upsert worker state
            cur.execute("""
                INSERT INTO neucodec_workers (
                    worker_id, hostname, gpu_name, gpu_vram_mb, state,
                    current_shard_id, current_shard_prefix,
                    shards_completed, shards_failed,
                    total_audio_hours, total_tokens,
                    progress_pct, progress_segments, progress_total_segments,
                    progress_rtf, progress_vram_gb, progress_phase,
                    last_heartbeat, last_error
                ) VALUES (
                    %(worker_id)s, %(hostname)s, %(gpu_name)s, %(gpu_vram_mb)s, %(state)s,
                    %(current_shard_id)s, %(current_shard_prefix)s,
                    %(shards_completed)s, %(shards_failed)s,
                    %(total_audio_hours)s, %(total_tokens)s,
                    %(progress_pct)s, %(progress_segments)s, %(progress_total_segments)s,
                    %(progress_rtf)s, %(progress_vram_gb)s, %(progress_phase)s,
                    NOW(), %(last_error)s
                ) ON CONFLICT (worker_id) DO UPDATE SET
                    state = EXCLUDED.state,
                    current_shard_id = EXCLUDED.current_shard_id,
                    current_shard_prefix = EXCLUDED.current_shard_prefix,
                    shards_completed = EXCLUDED.shards_completed,
                    shards_failed = EXCLUDED.shards_failed,
                    total_audio_hours = EXCLUDED.total_audio_hours,
                    total_tokens = EXCLUDED.total_tokens,
                    progress_pct = EXCLUDED.progress_pct,
                    progress_segments = EXCLUDED.progress_segments,
                    progress_total_segments = EXCLUDED.progress_total_segments,
                    progress_rtf = EXCLUDED.progress_rtf,
                    progress_vram_gb = EXCLUDED.progress_vram_gb,
                    progress_phase = EXCLUDED.progress_phase,
                    last_heartbeat = NOW(),
                    last_error = EXCLUDED.last_error
            """, snap)

            # Also update shard claimed_at for stale detection
            if snap["current_shard_id"]:
                cur.execute(
                    "UPDATE neucodec_shards SET claimed_at = NOW() WHERE id = %s",
                    (snap["current_shard_id"],),
                )

            cur.close()
            conn.close()
        except Exception as e:
            log.debug("Heartbeat failed: %s", e)

# ── Shard claim & status ────────────────────────────────────────────────────

def claim_shard(conn, worker_id: str):
    cur = conn.cursor()
    cur.execute("SELECT * FROM claim_neucodec_shard(%s)", (worker_id,))
    row = cur.fetchone()
    cur.close()
    if row:
        return {"id": row[0], "prefix": row[1], "dataset": row[2], "language": row[3]}
    return None

def update_shard_status(conn, shard_id: int, status: str, **kwargs):
    sets = ["status = %s"]
    vals = [status]
    for k, v in kwargs.items():
        sets.append(f"{k} = %s")
        vals.append(v)
    vals.append(shard_id)
    cur = conn.cursor()
    cur.execute(f"UPDATE neucodec_shards SET {', '.join(sets)} WHERE id = %s", vals)
    cur.close()

def reclaim_stale(conn):
    cur = conn.cursor()
    cur.execute("""
        UPDATE neucodec_shards
        SET status = 'pending', worker_id = NULL, claimed_at = NULL
        WHERE status IN ('claimed', 'processing')
          AND claimed_at < NOW() - INTERVAL '%s seconds'
        RETURNING id, shard_prefix
    """, (STALE_CLAIM_TIMEOUT,))
    reclaimed = cur.fetchall()
    cur.close()
    if reclaimed:
        log.warning("Reclaimed %d stale shards: %s", len(reclaimed), [r[1] for r in reclaimed])

# ── Download ─────────────────────────────────────────────────────────────────

def download_shard(s3, shard_prefix: str, work_dir: Path):
    tar_path = work_dir / "audio.tar"
    meta_path = work_dir / "metadata.parquet"

    log.info("Downloading %s ...", shard_prefix)
    t0 = time.time()

    from boto3.s3.transfer import TransferConfig
    config = TransferConfig(
        multipart_threshold=64 * 1024 * 1024,
        max_concurrency=8,
        multipart_chunksize=64 * 1024 * 1024,
    )
    s3.download_file(R2_BUCKET, f"{shard_prefix}/audio.tar", str(tar_path), Config=config)
    s3.download_file(R2_BUCKET, f"{shard_prefix}/metadata.parquet", str(meta_path))

    tar_size = tar_path.stat().st_size
    dl_time = time.time() - t0
    log.info("Downloaded %.1f GB in %.0fs (%.1f MB/s)", tar_size / 1e9, dl_time, tar_size / dl_time / 1e6)
    return tar_path, meta_path, tar_size

# ── Encode ───────────────────────────────────────────────────────────────────

def load_codec():
    from neucodec import NeuCodec
    import torch.nn as nn

    log.info("Loading neucodec model...")
    t0 = time.time()
    codec = NeuCodec.from_pretrained("neuphonic/neucodec")

    orig_layers = len(codec.semantic_model.encoder.layers)
    codec.semantic_model.encoder.layers = nn.ModuleList(
        list(codec.semantic_model.encoder.layers)[:17]
    )
    log.info("Truncated semantic model: %d → %d layers", orig_layers, 17)

    _ = codec.eval().to("cuda")
    log.info("Neucodec loaded in %.1fs, VRAM: %.1f GB", time.time() - t0, torch.cuda.memory_allocated() / 1e9)
    return codec

def encode_shard(codec, tar_path: Path, meta_path: Path, wstate: WorkerState):
    metadata = pd.read_parquet(meta_path)

    tf = tarfile.open(str(tar_path), "r")
    members = [m for m in tf.getmembers() if m.name.endswith(".flac")]
    total_members = len(members)
    log.info("Encoding %d audio files...", total_members)

    member_to_seg = {}
    for m in members:
        seg_id = m.name.replace(".flac", "").split("/")[-1]
        member_to_seg[m.name] = seg_id

    fbank_queue = queue.Queue(maxsize=16)

    def fbank_worker():
        for member in members:
            try:
                f = tf.extractfile(member)
                audio_data, sr = sf.read(io.BytesIO(f.read()))
                if len(audio_data) > MAX_AUDIO_SAMPLES:
                    chunks = [audio_data[s:s + MAX_AUDIO_SAMPLES]
                              for s in range(0, len(audio_data), MAX_AUDIO_SAMPLES)]
                else:
                    chunks = [audio_data]
                seg_id = member_to_seg[member.name]
                chunk_feats = []
                for chunk in chunks:
                    y = torch.tensor(chunk, dtype=torch.float32)[None, None]
                    feats = codec.feature_extractor(
                        y[0, :].cpu(), sampling_rate=16_000, return_tensors="pt"
                    ).input_features
                    chunk_feats.append((y, feats))
                fbank_queue.put((seg_id, chunk_feats, len(audio_data) / sr))
            except Exception as e:
                seg_id = member_to_seg.get(member.name, member.name)
                log.warning("Failed to read/extract %s: %s", seg_id, e)
                fbank_queue.put((seg_id, None, 0))
        fbank_queue.put(None)

    cpu_thread = threading.Thread(target=fbank_worker, daemon=True)
    cpu_thread.start()

    results = []
    total_audio_s = 0.0
    total_tokens = 0
    processed = 0
    failed = 0
    t_start = time.time()

    while True:
        item = fbank_queue.get()
        if item is None:
            break

        seg_id, chunk_feats, dur = item
        if chunk_feats is None:
            failed += 1
            processed += 1
            continue

        try:
            all_codes = []
            for y, feats in chunk_feats:
                feats_gpu = feats.to("cuda")
                y_gpu = y.to("cuda")

                with torch.no_grad():
                    acoustic_emb = codec.CodecEnc(y_gpu).transpose(1, 2)
                    semantic_output = (
                        codec.semantic_model(feats_gpu).hidden_states[16].transpose(1, 2)
                    )
                    semantic_encoded = codec.SemanticEncoder_module(semantic_output)

                    if acoustic_emb.shape[-1] != semantic_encoded.shape[-1]:
                        ml = min(acoustic_emb.shape[-1], semantic_encoded.shape[-1])
                        acoustic_emb = acoustic_emb[:, :, :ml]
                        semantic_encoded = semantic_encoded[:, :, :ml]

                    concat_emb = torch.cat([semantic_encoded, acoustic_emb], dim=1)
                    concat_emb = codec.fc_prior(concat_emb.transpose(1, 2)).transpose(1, 2)
                    _, codes, _ = codec.generator(concat_emb, vq=True)

                all_codes.append(codes[0, 0].cpu().numpy().astype(np.uint16))
                del feats_gpu, y_gpu, acoustic_emb, semantic_output, semantic_encoded
                del concat_emb, codes

            token_array = np.concatenate(all_codes) if len(all_codes) > 1 else all_codes[0]
            results.append({
                "segment_id": seg_id,
                "neucodec_tokens": token_array.tobytes(),
                "token_count": len(token_array),
            })
            total_audio_s += dur
            total_tokens += len(token_array)
        except Exception as e:
            log.warning("Encode failed for %s: %s", seg_id, e)
            failed += 1

        processed += 1
        if processed % 500 == 0:
            torch.cuda.empty_cache()
            elapsed = time.time() - t_start
            rtf = total_audio_s / elapsed if elapsed > 0 else 0
            wstate.set_progress(processed, total_members, rtf)
        if processed % 3000 == 0:
            elapsed = time.time() - t_start
            rtf = total_audio_s / elapsed if elapsed > 0 else 0
            vram_gb = torch.cuda.memory_allocated() / 1e9
            log.info(
                "[%d/%d] %.0f%% done, %.0fs elapsed, RTF=%.1fx, VRAM=%.1fGB",
                processed, total_members, processed / total_members * 100, elapsed, rtf, vram_gb,
            )

    cpu_thread.join()
    tf.close()

    wall = time.time() - t_start
    rtf = total_audio_s / wall if wall > 0 else 0
    log.info(
        "Encoding done: %d ok, %d failed, %.1f hrs audio, %.0fs wall, RTF=%.1fx, %d tokens",
        processed - failed, failed, total_audio_s / 3600, wall, rtf, total_tokens,
    )

    df = pd.DataFrame(results)
    return df, total_audio_s, total_tokens, wall

# ── Upload ───────────────────────────────────────────────────────────────────

def upload_results(s3, shard_prefix: str, tokens_df: pd.DataFrame, work_dir: Path):
    out_path = work_dir / "neucodec_tokens.parquet"
    tokens_df.to_parquet(str(out_path), index=False)
    upload_key = f"{shard_prefix}/neucodec_tokens.parquet"
    file_size = out_path.stat().st_size
    log.info("Uploading %s (%.1f MB)...", upload_key, file_size / 1e6)
    s3.upload_file(str(out_path), R2_BUCKET, upload_key)
    log.info("Upload complete: %s", upload_key)
    return upload_key

# ── Main loop ────────────────────────────────────────────────────────────────

def process_one_shard(codec, s3, conn, shard: dict, wstate: WorkerState):
    shard_id = shard["id"]
    prefix = shard["prefix"]

    log.info("═" * 60)
    log.info("Processing shard %d: %s", shard_id, prefix)
    log.info("═" * 60)

    wstate.set_phase("downloading", shard_id=shard_id, shard_prefix=prefix)

    try:
        update_shard_status(conn, shard_id, "processing", started_at="NOW()")

        with tempfile.TemporaryDirectory(prefix="neucodec_") as work_dir:
            work_dir = Path(work_dir)

            tar_path, meta_path, tar_size = download_shard(s3, prefix, work_dir)

            wstate.set_phase("encoding")
            tokens_df, total_audio_s, total_tokens, encode_wall = encode_shard(
                codec, tar_path, meta_path, wstate
            )

            if tokens_df.empty:
                raise RuntimeError("No segments encoded successfully")

            wstate.set_phase("uploading")
            output_key = upload_results(s3, prefix, tokens_df, work_dir)

            update_shard_status(
                conn, shard_id, "completed", completed_at="NOW()",
                audio_tar_bytes=tar_size,
                segment_count=len(tokens_df),
                total_audio_seconds=total_audio_s,
                total_tokens=total_tokens,
                encode_wall_seconds=encode_wall,
                output_key=output_key,
            )
            wstate.mark_shard_done(total_audio_s / 3600, total_tokens)
            log.info("Shard %d completed successfully", shard_id)

    except Exception as e:
        log.exception("Shard %d failed: %s", shard_id, e)
        wstate.mark_shard_failed(str(e))
        try:
            update_shard_status(conn, shard_id, "failed", failed_at="NOW()", error_message=str(e)[:500])
        except Exception:
            pass

def main():
    parser = argparse.ArgumentParser(description="Neucodec encoding worker")
    parser.add_argument("--worker-id", default=None)
    parser.add_argument("--max-shards", type=int, default=0)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    worker_id = args.worker_id or f"{socket.gethostname()}-{uuid.uuid4().hex[:8]}"

    gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "NONE"
    gpu_vram = torch.cuda.get_device_properties(0).total_memory // (1024 * 1024) if torch.cuda.is_available() else 0
    log.info("Worker starting: %s", worker_id)
    log.info("GPU: %s (%d MB)", gpu_name, gpu_vram)

    codec = load_codec()
    with torch.no_grad():
        codec.encode_code(torch.randn(1, 1, 16000).cuda())
    log.info("Codec warmup complete")

    # Initialize worker state and start heartbeat thread
    wstate = WorkerState(worker_id, gpu_name, gpu_vram)
    stop_hb = threading.Event()
    hb_thread = threading.Thread(target=heartbeat_loop, args=(wstate, stop_hb), daemon=True)
    hb_thread.start()

    s3 = get_s3()
    conn = get_db()
    reclaim_stale(conn)

    shards_done = 0
    try:
        while True:
            shard = claim_shard(conn, worker_id)
            if shard is None:
                log.info("No more pending shards. Worker exiting.")
                break

            if args.dry_run:
                log.info("DRY RUN: shard %d: %s", shard["id"], shard["prefix"])
                update_shard_status(conn, shard["id"], "pending", worker_id=None, claimed_at=None)
                shards_done += 1
                if args.max_shards > 0 and shards_done >= args.max_shards:
                    break
                continue

            process_one_shard(codec, s3, conn, shard, wstate)
            shards_done += 1

            if args.max_shards > 0 and shards_done >= args.max_shards:
                log.info("Reached max_shards=%d, exiting.", args.max_shards)
                break
    finally:
        # Mark worker as exited
        wstate.state = "exited"
        wstate.progress_phase = "exited"
        stop_hb.set()
        hb_thread.join(timeout=10)
        # Final heartbeat write
        try:
            snap = wstate.snapshot()
            fc = get_db()
            cur = fc.cursor()
            cur.execute("""
                UPDATE neucodec_workers SET state='exited', progress_phase='exited', last_heartbeat=NOW()
                WHERE worker_id = %s
            """, (worker_id,))
            cur.close()
            fc.close()
        except Exception:
            pass
        conn.close()

    log.info("Worker %s finished after %d shards", worker_id, shards_done)

if __name__ == "__main__":
    main()