hegm 0.3.3

HeGM: Heterogeneous GPU Migration framework for deep learning (PyTorch today, TensorFlow-ready)

HeGM - Heterogeneous GPU Migration Framework

HeGM lets you live-migrate long-running deep learning jobs between GPUs using CRIU (Checkpoint/Restore in Userspace), without requiring any code changes in your training script. Your existing for step, batch in enumerate(dataloader) loop works as-is -- HeGM transparently resumes from the correct batch and step number after migration.

It targets PyTorch today and is structured to support additional backends (e.g. TensorFlow) in the future.

Installation

Install the Python package

The HeGM Python package (launcher + runtime hooks) is published on PyPI as hegm.

On any machine/container where you run training:

python -m pip install --upgrade pip
pip install hegm

This installs:

• the hegm package (transparent enumerate() patching, hooks, backends)
• a sitecustomize.py module in site-packages that automatically loads HeGM
• the hegm-launcher CLI

You can then start training under HeGM with:

hegm-launcher python -u train.py

In Kubernetes Pods, your command block typically becomes:

hegm-launcher python -u /workspace/train.py

Runtime prerequisites (cluster)

HeGM relies on a CRIU and CRI-O build that understand GPU checkpoint/restore and the /checkpoint/*/lock convention. Install these first on your nodes:

• CRIU (GPU migration fork)
  leehun-criu branch 2026-01-26/gpu-migration-support
  See the upstream README for build/install instructions:

  • Repo: lehuannhatrang/leehun-criu

• CRI-O (restore-from-file fork)
  leehun-cri-o branch 2026-02-03/support-restore-from-file
  Install and configure it as your Kubernetes container runtime:

  • Repo: lehuannhatrang/leehun-cri-o

In addition you need:

• A Kubernetes cluster with GPU nodes and the NVIDIA drivers/runtime configured.
• kubectl access to the cluster.

Deploying HeGM example

From this repository:

1. Create the demo namespace and storage (if not already present):
   • examples/dra/ns.yaml
   • examples/dra/storage.yaml
2. Create the ConfigMaps:
   • examples/dra/hegm-scripts.yaml (Launcher + HeGM payload)
   • examples/dra/training-script.yaml (your train.py)
3. Create the resource claim(s):
   • examples/dra/resource-claim.yaml (or resource-claim-restore.yaml)
4. Launch a training pod:
   • Single worker: examples/dra/training-pod.yaml
   • Two workers in one Pod (for PID-isolation testing): examples/dra/multiple-training-pod.yaml
5. Trigger checkpoint from outside the pod:
   • Prefer using the hegm-ctrl CLI, which also exposes a REST API inside the Pod:
     • All workers using the GPU:
       kubectl exec ... -- hegm-ctrl checkpoint --all
     • Specific PID(s):
       kubectl exec ... -- hegm-ctrl checkpoint --pid <pid1> <pid2> ...
   • Under the hood this sends SIGUSR1 to the worker PID(s) and waits until the corresponding /checkpoint/<PPID>/lock file(s) appear, indicating that the checkpoint is ready for CRIU/CRI-O to snapshot.

Control API (CLI + REST)

When at least one hegm-launcher is active, HeGM starts a small HTTP server inside the container (default 0.0.0.0:8298, configurable via HEGM_CTRL_PORT). It is implemented by the hegm-ctrl module.

• CLI (inside the Pod):

  • Trigger checkpoint for all GPU-using or worker processes:

    hegm-ctrl checkpoint --all
    

  • Trigger checkpoint for specific worker PID(s):

    hegm-ctrl checkpoint --pid 1234 5678
    

  • Resume from checkpoints:

    # All pending checkpoints (all /checkpoint/*/lock)
    hegm-ctrl resume --all
    
    # Specific parent PID(s)
    hegm-ctrl resume --ppid 1 2 3
    

  checkpoint waits until the corresponding lock file(s) /checkpoint/<PPID>/lock appear (or timeout), so when the command returns the checkpoint is ready for CRIU/CRI-O.

• REST API (from outside the Pod):

  Assuming HEGM_CTRL_PORT=8298 and using the Pod IP:

  • POST /checkpoint

    // Signal specific worker PIDs
    {
      "pids": [1234, 5678]
    }
    
    // Or auto-detect GPU/worker PIDs
    {
      "all": true
    }
    

    Response:

    {
      "results": {
        "1234": "checkpoint-ready",
        "5678": "signalled (timeout-waiting-lock)"
      }
    }
    

  • POST /resume

    // Resume specific parent PIDs
    {
      "ppids": [1, 2, 3]
    }
    
    // Or resume all pending checkpoints
    {
      "all": true
    }
    

    Response:

    {
      "results": {
        "1": "lock-removed",
        "2": "no-lock"
      }
    }
    

6. Use your CRI-O / CRIU integration to checkpoint and restore the container, pointing CRI-O at the CRIU checkpoint tarball and re-using the same hegm-scripts and train-script ConfigMaps (see examples/dra/restore-pod.yaml).

High-level architecture

HeGM is split into two main pieces:

• Launcher (launcher.py / hegm-launcher): a small supervisor process that:

  • spawns your training script as a child process (the Worker)
  • injects sitecustomize.py via PYTHONPATH
  • watches the Worker's exit code:
    • 0 → training finished, exit normally
    • 99 → Worker saved a checkpoint and exited for migration
  • buffers the checkpoint file into RAM so CRIU can carry it across nodes
  • creates a per-process lock file so an external controller knows when it is safe to snapshot
  • starts a lightweight in-container REST API (via the hegm-ctrl module) so external controllers can trigger checkpoint/resume over HTTP

• Payload (sitecustomize.py + hegm/ package): automatically loaded into the Worker by Python. It:

  • installs a PEP‑451 import hook for torch
  • monkey-patches torch.nn.Module and torch.optim.Optimizer to track models and optimizers
  • patches builtins.enumerate so enumerate(DataLoader) transparently resumes from the correct batch after a checkpoint restore
  • tracks global training steps and RNG state
  • handles SIGUSR1 by saving a checkpoint and exiting with code 99

All of this is delivered to your Pods via a ConfigMap (examples/dra/hegm-scripts.yaml).

Key files

• launcher.py

  • Entry point you run instead of python train.py.
  • For each Launcher process, checkpoints and lock files are isolated by PID:
    • /checkpoint/<PID>/latest.pt
    • /checkpoint/<PID>/lock
  • The external controller can discover all ready instances via /checkpoint/*/lock.

• sitecustomize.py

  • Thin bootstrap that simply does import hegm.
  • Needs to stay at the top level of your PYTHONPATH so Python’s sitecustomize mechanism can find it.

• hegm/

  • __init__.py: patches enumerate() for transparent DataLoader resume, public API (hegm.global_step()), SIGUSR1 handler, import hooks.
  • _config.py: env‑driven configuration and lightweight logging.
  • _hook.py: generic PEP‑451 import hook logic.
  • backends/__init__.py: AbstractBackend interface and registry.
  • backends/pytorch.py: PyTorch backend (tracking, checkpointing, RNG, GPU teardown).
  • ctrl.py: external control helpers (hegm-ctrl CLI and REST API).

Why launcher.py and sitecustomize.py live at the top level

The package code lives under hegm/, but we intentionally keep:

• sitecustomize.py at the top level so that Python’s automatic sitecustomize import works as soon as /opt/hegm is on PYTHONPATH.
• launcher.py at the top level so Kubernetes manifests can invoke it as /opt/hegm/launcher.py without needing to worry about Python module import paths.

Internally, both of these files are very thin:

• sitecustomize.py just imports the hegm package.
• launcher.py is a small script that wires the environment and hands off most logic to the shared configuration/checkpointing scheme.

In other words, the reusable logic already lives in hegm/; the two top‑level scripts are just ergonomic entry points for Python and Kubernetes. We can still add more entry points later (e.g. hegm.launcher console script) without changing this layout.

Examples

Working DRA examples live under examples/dra/:

• training-pod.yaml – single training job.
• multiple-training-pod.yaml – two launcher.py processes in one Pod, useful for testing PID‑isolated checkpoints and locks.
• restore-pod.yaml – example restore workflow using CRIU.

See docs/HeGM_Architecture.md for a more detailed architectural walk‑through.