RetainZ: Reclaim-Time-Aware Placement for AI Checkpoints on Zoned SSDs

Abstract

Solid State Drives (SSDs) are becoming the dominant medium for performance-critical storage. Zoned Namespace (ZNS) SSDs are getting more and more attractive because sequential writes and explicit zone resets reduce address-mapping, over-provisioning, and internal garbage-collection costs. However, reclaiming space requires resetting an entire zone, so deleting one file does not free its space while other files there must be kept. This erase-before-write constraint becomes a bottleneck for write-intensive workloads, particularly AI checkpointing during model training, where frequent saves of model and optimizer state compete for space with long-retained checkpoints. We observe that retention policies inherently encode when data can be retired, avoiding the need for a learned lifetime predictor. We present RetainZ, a storage backend that translates these policies into explicit reclaim epochs. Its allocator isolates long-retained data and groups objects with nearby reclaim epochs, balancing delayed space reclamation against unused space at zone ends. Evaluated on FEMU with multi-tenant streams and checkpoint manifests up to Pythia-1B, RetainZ reduces space occupied by deleted data and unused zone tails by 41.7% and mean per-input checkpoint p99 latency by 47.5% at 84% target load compared with the allocator of ZenFS, an established lifetime-aware ZNS backend, and lowers host write amplification from 1.165 to 1.003.