[Bugfix][KV Offload] Clean up resources after initialization failure (#51227)
Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: OpenAI Codex <noreply@openai.com>
AlexHuang committed
Aug 6, 2026 at 19:42 UTC
46e6a83ce12b5968d956279a1bd4611de16d69eb
2 files changed
+85
-53
vllm/v1/kv_offload/cpu/spec.py
+11
-6
@@ -165,12 +165,17 @@ class CPUOffloadingSpec(OffloadingSpec):
165
kv_bytes_per_block=self.kv_bytes_per_chunk,
166
cpu_page_size=self.cpu_page_size_per_worker,
167
)
168
- return CPUOffloadingWorker(
169
- kv_caches=kv_caches,
170
- blocks_per_chunk=self.blocks_per_chunk,
171
- num_cpu_blocks=self.num_blocks,
172
- mmap_region=mmap_region,
173
- )
168
+ try:
169
+ return CPUOffloadingWorker(
170
+ kv_caches=kv_caches,
171
+ blocks_per_chunk=self.blocks_per_chunk,
172
+ num_cpu_blocks=self.num_blocks,
173
+ mmap_region=mmap_region,
174
+ )
175
+ except Exception:
176
+ if mmap_region is not None:
177
+ mmap_region.cleanup()
178
+ raise
179
180
@override
181
def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
vllm/v1/kv_offload/tiering/spec.py
+74
-47
@@ -185,31 +185,38 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
185
TieringOffloadingManager instance
186
"""
187
if not self._manager:
188
- # Create scheduler-side SharedOffloadRegion (rank=None) so the
189
- # primary tier can eagerly create a memoryview over _base.
190
- scheduler_mmap = SharedOffloadRegion(
191
- engine_id=self._engine_id,
192
- num_blocks=self.num_blocks,
193
- rank=None,
194
- kv_bytes_per_block=self.kv_bytes_per_chunk,
195
- cpu_page_size=self.cpu_page_size_per_worker,
196
- )
197
- self._scheduler_mmap = scheduler_mmap
198
-
199
- # Create primary tier (CPU-based)
200
- primary_tier = CPUPrimaryTierOffloadingManager(
201
- num_blocks=self.num_blocks,
202
- cache_policy=self.eviction_policy,
203
- cache_policy_module_path=self.cache_policy_module_path,
204
- enable_events=self.kv_events_config.enable_kv_cache_events,
205
- mmap_region=scheduler_mmap,
206
- )
188
+ if int(self.extra_config.get("store_threshold", 0)) >= 2:
189
+ raise ValueError(
190
+ "store_threshold is not supported for TieringOffloadingSpec"
191
+ )
192
208
- # Create secondary tiers
209
- primary_kv_view = primary_tier.get_kv_memoryview()
193
+ scheduler_mmap: SharedOffloadRegion | None = None
194
+ primary_tier: CPUPrimaryTierOffloadingManager | None = None
195
secondary_tiers = []
211
- for i, tier_config in enumerate(self.secondary_tier_configs):
212
- try:
196
+ try:
197
+ # Create scheduler-side SharedOffloadRegion (rank=None) so the
198
+ # primary tier can eagerly create a memoryview over _base.
199
+ scheduler_mmap = SharedOffloadRegion(
200
+ engine_id=self._engine_id,
201
+ num_blocks=self.num_blocks,
202
+ rank=None,
203
+ kv_bytes_per_block=self.kv_bytes_per_chunk,
204
+ cpu_page_size=self.cpu_page_size_per_worker,
205
+ )
206
+ self._scheduler_mmap = scheduler_mmap
207
+
208
+ # Create primary tier (CPU-based)
209
+ primary_tier = CPUPrimaryTierOffloadingManager(
210
+ num_blocks=self.num_blocks,
211
+ cache_policy=self.eviction_policy,
212
+ cache_policy_module_path=self.cache_policy_module_path,
213
+ enable_events=self.kv_events_config.enable_kv_cache_events,
214
+ mmap_region=scheduler_mmap,
215
+ )
216
+
217
+ # Create secondary tiers
218
+ primary_kv_view = primary_tier.get_kv_memoryview()
219
+ for i, tier_config in enumerate(self.secondary_tier_configs):
220
tier = SecondaryTierFactory.create_secondary_tier(
221
tier_config, primary_kv_view, self
222
)
@@ -219,26 +226,42 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
226
i,
227
tier.tier_type,
228
)
222
- except Exception as e:
223
- logger.error(
224
- "Failed to create secondary tier from config index %i: %s",
225
- i,
226
- e,
227
- )
228
- raise
229
230
- # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited
231
- # get_worker(). Secondary tier transfers are handled by the
232
- # secondary tier managers and need no additional workers here.
233
- tiering_manager = TieringOffloadingManager(
234
- primary_tier=primary_tier,
235
- secondary_tiers=secondary_tiers,
236
- )
237
- if int(self.extra_config.get("store_threshold", 0)) >= 2:
238
- raise ValueError(
239
- "store_threshold is not supported for TieringOffloadingSpec"
230
+ # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited
231
+ # get_worker(). Secondary tier transfers are handled by the
232
+ # secondary tier managers and need no additional workers here.
233
+ tiering_manager = TieringOffloadingManager(
234
+ primary_tier=primary_tier,
235
+ secondary_tiers=secondary_tiers,
236
)
241
- self._manager = tiering_manager
237
+ self._manager = tiering_manager
238
+ except Exception:
239
+ for tier in reversed(secondary_tiers):
240
+ try:
241
+ tier.shutdown()
242
+ except Exception:
243
+ logger.exception(
244
+ "Failed to shut down secondary tier during "
245
+ "initialization cleanup"
246
+ )
247
+ if primary_tier is not None:
248
+ try:
249
+ primary_tier.shutdown()
250
+ except Exception:
251
+ logger.exception(
252
+ "Failed to shut down primary tier during "
253
+ "initialization cleanup"
254
+ )
255
+ elif scheduler_mmap is not None:
256
+ try:
257
+ scheduler_mmap.cleanup()
258
+ except Exception:
259
+ logger.exception(
260
+ "Failed to clean up scheduler mmap during "
261
+ "initialization cleanup"
262
+ )
263
+ self._scheduler_mmap = None
264
+ raise
265
266
logger.info(
267
"Created TieringOffloadingManager with primary tier "
@@ -273,9 +296,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
296
kv_bytes_per_block=self.kv_bytes_per_chunk,
297
cpu_page_size=self.cpu_page_size_per_worker,
298
)
276
- return CPUOffloadingWorker(
277
- kv_caches=kv_caches,
278
- blocks_per_chunk=self.blocks_per_chunk,
279
- num_cpu_blocks=self.num_blocks,
280
- mmap_region=worker_mmap,
281
- )
299
+ try:
300
+ return CPUOffloadingWorker(
301
+ kv_caches=kv_caches,
302
+ blocks_per_chunk=self.blocks_per_chunk,
303
+ num_cpu_blocks=self.num_blocks,
304
+ mmap_region=worker_mmap,
305
+ )
306
+ except Exception:
307
+ worker_mmap.cleanup()
308
+ raise