[bugfix] fail fast when graphlearn sampler server dies

[tiankongdeguiji]

Problem

The graphlearn negative-sampler server runs as a daemon process of LOCAL_RANK 0. graphlearn's launch_server start()s it but never join()s it and returns no handle, so when the server dies — most often an OOM kill when many dataloader workers each buffer large sampled batches, but also any crash — every sampler RPC client blocks forever (graphlearn retries the dead connection rather than erroring). The failure then surfaces only ~10 min later as an opaque ALLTOALL collective (NCCL) timeout that gives no hint of the real cause, and the dead server lingers as a <defunct> process.

Fix

Inline graphlearn's server launch (it uses only public Graph.init/server_get_stats/close, and tzrec never reads graphlearn's STATS_DICT/SERVER_LAUNCHED/stats) so we hold the exact server Process handle. A background daemon thread on rank 0 then watches that process and, on an abnormal exit, logs the cause and os._exit(1)s so torchrun tears the job down in seconds instead of hanging until the collective timeout.

A thread is needed because rank 0's main thread is blocked inside the dataloader fetch when the server dies and never returns to self-check (SIGCHLD was avoided since PyTorch's DataLoader already installs its own SIGCHLD handler). A threading.Event set from an atexit hook disarms the watchdog before normal interpreter shutdown SIGTERMs the daemon server, so a clean teardown is never mistaken for a crash (no false positives during healthy training).

Verification

• New regression test test_server_liveness_watchdog_aborts_on_server_death launches a real GL server, kills it, and asserts the watchdog aborts the rank; it passes alongside the existing test_negative_sampler (which confirms the inlined launch still serves correctly).
• On a 4×V100 reproduction of the original hang, kill -9 of the server triggered the watchdog in ~10 s and the job was torn down in ~15 s (vs a ~600 s NCCL-timeout hang previously), with rank 0 reported as the root cause.

Note

This makes the failure loud and fast; it does not prevent the OOM itself. The operational mitigation when keeping a large negative-sample count is to lower data_config.num_workers, which bounds the number of worker processes buffering large sampled batches.

Also bumps the package version to 1.2.23.

🤖 Generated with Claude Code

https://claude.ai/code/session_012uPyFcxcLAdNpgFFCYSZ2f

[github-actions]

The new test covers the firing path (abnormal death → os._exit(1)) well, but the disarm / false-positive-avoidance path is untested — i.e. the stop.is_set() or proc.exitcode == 0 branch in _watch, plus the atexit/__del__ disarm. That logic is what prevents a healthy job from being killed on a normal teardown, which is arguably the more dangerous regression: an inverted check (e.g. exitcode != 0) would still pass the whole suite.

Consider a fast, GL-free unit test that calls _watch's logic against a stub proc (is_alive()→False, exitcode==0) and asserts it returns without exiting, covering all three branches (alive / clean-exit / abnormal-exit) deterministically.

[github-actions]

Inlining drops graphlearn's SERVER_LAUNCHED duplicate-launch guard (which raised on a second launch). It's latent today (one launch per process: train launches, eval reuses the cluster), but the guard removal turns a previously-idempotent call into a footgun — a second launch_server() would now start a second server Process + watchdog + atexit registration. Cheap insurance:

if self._server_proc is not None:
    return

at the top of the LOCAL_RANK == 0 block.

[github-actions]

Minor: "(inlined _server_manager)" implies a faithful copy, but this intentionally drops the counts_dict.update(graph.server_get_stats()) bookkeeping (tzrec never reads the stats). Worth noting the simplification so a future reader doesn't "restore" gl.get_counts() — e.g. "simplified from graphlearn's _server_manager; the server-stats return value is intentionally discarded."

[github-actions]

Review summary

Solid, well-targeted fix. The watchdog design is sound and I verified the trickiest parts hold up:

• Disarm correctness (no false positives): I reproduced the shutdown path — at normal interpreter exit atexit runs stop.set() before multiprocessing SIGTERMs the daemon server (atexit is LIFO and mp's _exit_function registers at import, so stop.set always runs first). The watchdog observes stop.is_set() and does not fire. Confirmed empirically: clean exit 0, no spurious abort.
• _run_gl_server faithfully reproduces graphlearn's _server_manager (minus the unused counts_dict), and LOCAL_RANK==0-only arming + os._exit(1) correctly leans on torchrun for job-wide teardown (matches the PR's measured ~15s).
• Log diagnostic survives os._exit: tzrec uses stdlib logging, whose StreamHandler.emit flushes per record, so the "exited unexpectedly" line is not lost.

Noteworthy feedback left inline (none blocking):

1. Test gap — only the firing path is covered; the disarm / exitcode==0 branch (the part that protects healthy jobs) has no test.
2. Lost duplicate-launch guard — latent today, but worth a one-line idempotency check.
3. Docstring nit — "inlined _server_manager" understates the intentional stats-channel drop.

I dismissed a few flagged items as non-issues after checking them: the atexit ordering is robust (not fragile), the del utils.STATS_DICT teardown in dataset_test.py does not break (those are graphlearn's still-defined globals, not tzrec's), and the log message is not lost to os._exit.