nvsnap: investigate reducing per-pid cuda-checkpoint restore-tid cost (prototype rolled back)

[balajinvda]

Summary

Investigate reducing the per-pid cuda-checkpoint --get-restore-tid cost during
restore. A prototype was built, measured, and then rolled back as too risky for
the value it delivered. This issue records the findings so the idea can be
revisited deliberately rather than rediscovered.

Status: rolled back, not shipped. Nothing is enabled by default.

The cost

The CUDA plugin resolves the driver's restore thread once per GPU pid by
spawning cuda-checkpoint --get-restore-tid. Each spawn pays a full CUDA driver
attach. Measured in-band from a restore log: one spawn took 2.724s, matching
the driver-attach floor seen elsewhere. On a workload with many GPU pids this
multiplies (inference NIM containers commonly run tens of GPU pids; one SGLang
configuration was observed with 103).

What was prototyped

Find the restore thread by scanning /proc/<pid>/task/*/comm instead of
spawning the tool, trusting the scan only when exactly one candidate exists and
otherwise falling back to asking the driver.

Why it was rolled back

1. It depends on an undocumented driver implementation detail: the restore
   worker thread's comm is cuda followed by hex digits (distinct from the
   driver's cuda-EvtHandlr and from application threads). This is a naming
   convention, not an API contract, and can change in any driver release.
2. The ambiguity guard is incomplete. Zero matches and multiple matches both
   fall back safely, but a single match that is the wrong thread cannot be
   detected and would silently resume the wrong tid.
3. It was validated against exactly one driver version on one cluster.
4. The payoff on typical workloads is small: about 5s of a ~26s restore for a
   2-3 pid workload.

Measured payoff/risk did not justify making the restore path depend on the
assumption.

What the measurements showed

Cross-check mode (resolve both ways on every lookup, compare, return the
driver's answer) across vLLM small/8B and an embedding model: 9 of 9 pid
resolutions agreed, zero mismatches, including the negative case where both the
scan and the driver report no restore thread. Separately, the scan and the
driver agreed across 14 checkpointed processes spanning four frameworks and both
single-GPU and TP=2 configurations, including individual tensor-parallel rank
workers.

So the approach is sound on the driver it was tested against. The risk is
version drift, not present-day correctness.

A hypothesis that did NOT hold

It was predicted that removing the per-pid spawns would rescue a 103-pid SGLang
checkpoint that fails on a ~10 minute dump timeout. It did not: that workload
fails at the same point with and without the optimization, so the failure is a
pre-existing dump-duration problem unrelated to this cost. Worth tracking
separately.

If revisited

• A more robust angle for many-pid workloads is likely raising or removing the
  fixed dump timeout and/or parallelizing the per-pid plugin calls, neither of
  which depends on undocumented naming.
• If the scan is revisited, keep it opt-in, and require a cross-check run
  against the target driver before enabling it.
• Prototype and measurements live on an unmerged CRIU fork branch
  (leverb-restore-tid-scan, three commits: the scan, the cross-check mode, and
  the change making it opt-in).