Skip to content

Evict older checkpoints in the shared Triton loader - #94

Open
haraschax wants to merge 1 commit into
masterfrom
codex/triton-checkpoint-retention
Open

haraschax wants to merge 1 commit into
masterfrom
codex/triton-checkpoint-retention

Conversation

@haraschax

Copy link
Copy Markdown
Contributor

Loading a newer checkpoint can leave many older epochs resident in Triton until the idle timeout, exhausting GPU memory.

Keep the newest TRITON_MAX_CHECKPOINTS_PER_EID epochs per EID, defaulting to 1, in the shared setup_triton_model loader. Artifacts of the same epoch share a slot. Serialize eviction and loading with an EID lock, including cached loads; reject requests outside the retained window. Names without an EID/epoch bypass the retention rule.

Wait for unloading to finish before deleting model files or allocating the replacement. If it does not finish within 60 seconds, raise instead of continuing. This avoids the old/new GPU allocation overlap observed with native Triton version replacement. Existing tasks using evicted checkpoints can fail on subsequent inference.

Validation: 24 temporary focused checks passed, including concurrency, cached backlog trimming, unload failures/timeouts, and numeric epoch ordering. An isolated Triton 2.69.0 GPU probe with Redis and ONNX/Python backends confirmed mixed-backend eviction, latest-N behavior, stale inference failures, old Python GPU allocations finalized before replacement allocation, and cleanup. All commit hooks, including Ruff and full Miniray ty, passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant