Skip to content

Commit 806f9d6

Browse files
Python: serialize concurrent checkpoint deletes
1 parent bf7d2ee commit 806f9d6

2 files changed

Lines changed: 56 additions & 16 deletions

File tree

‎python/packages/core/agent_framework/_workflows/_checkpoint.py‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,13 @@
77
import json
88
import logging
99
import os
10+
import threading
1011
import uuid
1112
from collections.abc import Mapping
1213
from dataclasses import dataclass, field, fields
1314
from datetime import datetime, timezone
1415
from pathlib import Path
15-
from typing import TYPE_CHECKING, Any, Protocol, TypeAlias
16+
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias
1617

1718
from ..exceptions import WorkflowCheckpointException
1819

@@ -272,6 +273,11 @@ class FileCheckpointStorage:
272273
)
273274
"""
274275

276+
_DELETE_LOCK_STRIPE_COUNT: ClassVar[int] = 64
277+
_DELETE_LOCKS: ClassVar[tuple[threading.Lock, ...]] = tuple(
278+
threading.Lock() for _ in range(_DELETE_LOCK_STRIPE_COUNT)
279+
)
280+
275281
def __init__(
276282
self,
277283
storage_path: str | Path,
@@ -312,6 +318,11 @@ def _validate_file_path(self, checkpoint_id: CheckpointID) -> Path:
312318
raise WorkflowCheckpointException(f"Invalid checkpoint ID: {checkpoint_id}")
313319
return file_path
314320

321+
@classmethod
322+
def _delete_lock(cls, file_path: Path) -> threading.Lock:
323+
"""Return the process-local deletion lock for a checkpoint file."""
324+
return cls._DELETE_LOCKS[hash(file_path) % cls._DELETE_LOCK_STRIPE_COUNT]
325+
315326
async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
316327
"""Save a checkpoint and return its ID.
317328
@@ -412,12 +423,14 @@ async def delete(self, checkpoint_id: CheckpointID) -> bool:
412423
True if the checkpoint was successfully deleted, False if no checkpoint with the given ID exists.
413424
"""
414425
file_path = self._validate_file_path(checkpoint_id)
426+
file_lock = self._delete_lock(file_path)
415427

416428
def _delete() -> bool:
417-
try:
418-
file_path.unlink()
419-
except FileNotFoundError:
420-
return False
429+
with file_lock:
430+
try:
431+
file_path.unlink()
432+
except FileNotFoundError:
433+
return False
421434
logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}")
422435
return True
423436

‎python/packages/core/tests/workflow/test_checkpoint.py‎

Lines changed: 38 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1262,8 +1262,10 @@ async def test_file_checkpoint_storage_delete():
12621262

12631263

12641264
async def test_file_checkpoint_storage_concurrent_delete():
1265+
"""Serialize deletes when overlapping unlink calls could both report success."""
12651266
with tempfile.TemporaryDirectory() as temp_dir:
12661267
storage = FileCheckpointStorage(temp_dir)
1268+
other_storage = FileCheckpointStorage(temp_dir)
12671269
checkpoint = WorkflowCheckpoint(
12681270
workflow_name="test-workflow",
12691271
graph_signature_hash="test-hash",
@@ -1272,19 +1274,44 @@ async def test_file_checkpoint_storage_concurrent_delete():
12721274
await storage.save(checkpoint)
12731275

12741276
file_path = (Path(temp_dir) / "same.json").resolve()
1275-
original_exists = Path.exists
1276-
barrier = threading.Barrier(2)
1277-
1278-
def synchronized_exists(path: Path) -> bool:
1279-
result = original_exists(path)
1280-
if path.resolve() == file_path:
1281-
barrier.wait(timeout=5)
1282-
return result
1283-
1284-
with patch.object(Path, "exists", synchronized_exists):
1277+
original_to_thread = asyncio.to_thread
1278+
original_unlink = Path.unlink
1279+
worker_barrier = threading.Barrier(2)
1280+
unlink_barrier = threading.Barrier(2)
1281+
unlink_guard = threading.Lock()
1282+
overlapping_delete_completed = False
1283+
1284+
async def synchronized_to_thread(function: Any, /, *args: Any, **kwargs: Any) -> Any:
1285+
def synchronized_call() -> Any:
1286+
worker_barrier.wait(timeout=5)
1287+
return function(*args, **kwargs)
1288+
1289+
return await original_to_thread(synchronized_call)
1290+
1291+
def macos_style_unlink(path: Path, missing_ok: bool = False) -> None:
1292+
nonlocal overlapping_delete_completed
1293+
if path.resolve() != file_path:
1294+
original_unlink(path, missing_ok=missing_ok)
1295+
return
1296+
1297+
try:
1298+
unlink_barrier.wait(timeout=0.5)
1299+
except threading.BrokenBarrierError:
1300+
original_unlink(path, missing_ok=missing_ok)
1301+
return
1302+
1303+
with unlink_guard:
1304+
if not overlapping_delete_completed:
1305+
original_unlink(path, missing_ok=missing_ok)
1306+
overlapping_delete_completed = True
1307+
1308+
with (
1309+
patch.object(asyncio, "to_thread", synchronized_to_thread),
1310+
patch.object(Path, "unlink", macos_style_unlink),
1311+
):
12851312
results = await asyncio.gather(
12861313
storage.delete(checkpoint.checkpoint_id),
1287-
storage.delete(checkpoint.checkpoint_id),
1314+
other_storage.delete(checkpoint.checkpoint_id),
12881315
)
12891316

12901317
assert sorted(results) == [False, True]

0 commit comments

Comments
 (0)