class DataRobotFileSystemCheckpointSaver(BaseCheckpointSaver[str]):
"""Persist LangGraph checkpoints on an fsspec filesystem.
For example :class:`datarobot.fs.DataRobotFileSystem`.
**Root and effective prefix**
Provide at most one of ``root`` or ``catalog_id`` (defaults to bare ``dr://`` when neither
is set):
- ``catalog_id="<catalog_item_id>"``: builds the root
``dr://<catalog_id>/langgraph_checkpoints`` so objects are stored under
``<catalog_id>/langgraph_checkpoints/checkpoints/...``. The catalog item must already
exist (create one with :meth:`datarobot.fs.DataRobotFileSystem.create_catalog_item_dir`).
- ``root="dr://<catalog_id>/langgraph_checkpoints"`` (or any path the filesystem accepts):
the ``checkpoints/`` subtree is appended under that path. With
:class:`datarobot.fs.DataRobotFileSystem` the first path segment is the catalog item id.
- neither: defaults to ``root="dr://"``.
**Layout** (paths below are relative to the effective prefix ending in
``.../checkpoints/``):
- ``<b64(thread_id)>/ns/<b64(checkpoint_ns)>/cpts/<checkpoint_id>.bin``
(empty ``thread_id`` or ``checkpoint_ns`` uses ``dr_genai_langgraph_empty_segment``
instead of an empty path segment)
- ``.../writes/<checkpoint_id>/w_<sha256(task_id)>.bin`` (per-task writes shard;
avoids lost updates when parallel branches call ``put_writes`` concurrently)
- legacy ``.../writes/<checkpoint_id>.bin`` (monolithic map) is still read and merged
when present
- ``.../blobs/<sha256(channel,version)>.bin`` (single ``serde.dumps_typed`` frame)
On-disk format is concatenated length-prefixed ``struct`` segments (no magic header, no
pickle). Layout is implied by path (``blobs/`` vs ``cpts/`` vs ``writes/``).
"""
def __init__(
self,
*,
fs: AbstractFileSystem,
root: str | None = None,
catalog_id: str | None = None,
serde: SerializerProtocol | None = None,
) -> None:
"""Initialize the saver.
At most one of ``root`` or ``catalog_id`` may be provided. If neither is given, the
root defaults to the bare ``dr://`` scheme.
:param fs: fsspec filesystem to persist on, for example
:class:`datarobot.fs.DataRobotFileSystem`.
:param root: explicit root prefix the filesystem accepts, for example
``dr://<catalog_id>/langgraph_checkpoints``. The ``checkpoints/`` subtree is
appended under it. Defaults to ``dr://`` when neither ``root`` nor ``catalog_id``
is set.
:param catalog_id: DataRobot catalog item id to store checkpoints in. When given,
``root`` is built as ``dr://<catalog_id>/langgraph_checkpoints`` (the catalog item
must already exist). Cannot be combined with ``root``.
:param serde: optional serializer; defaults to LangGraph's.
:raises ValueError: if both ``root`` and ``catalog_id`` are provided.
"""
if root is not None and catalog_id is not None:
msg = "Provide at most one of 'root' or 'catalog_id'."
raise ValueError(msg)
super().__init__(serde=serde)
self.fs = fs
if catalog_id is not None:
resolved_root = _root_from_catalog_id(catalog_id)
elif root is not None:
resolved_root = root
else:
resolved_root = "dr://"
self.root = _normalize_dr_fs_root(resolved_root)
def _ns_root(self, thread_id: str, checkpoint_ns: str) -> str:
return _path_join(
self.root,
_CHECKPOINT_TREE_DIR,
_encoder_segment(thread_id),
"ns",
_encoder_segment(checkpoint_ns),
)
def _blob_path(
self,
thread_id: str,
checkpoint_ns: str,
channel: str,
version: str | int | float,
) -> str:
return _path_join(
self._ns_root(thread_id, checkpoint_ns),
"blobs",
f"{_blob_filename(channel, version)}{_CHECKPOINT_ARTIFACT_EXT}",
)
def _cpt_path(self, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> str:
return _path_join(
self._ns_root(thread_id, checkpoint_ns),
"cpts",
f"{checkpoint_id}{_CHECKPOINT_ARTIFACT_EXT}",
)
def _writes_legacy_path(self, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> str:
"""Pre-shard monolithic writes map (still merged on read)."""
return _path_join(
self._ns_root(thread_id, checkpoint_ns),
"writes",
f"{checkpoint_id}{_CHECKPOINT_ARTIFACT_EXT}",
)
def _writes_shard_dir(self, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> str:
return _path_join(
self._ns_root(thread_id, checkpoint_ns),
"writes",
checkpoint_id,
)
def _task_writes_shard_path(
self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, task_id: str
) -> str:
digest = sha256(task_id.encode("utf-8")).hexdigest()
return _path_join(
self._writes_shard_dir(thread_id, checkpoint_ns, checkpoint_id),
f"w_{digest}{_CHECKPOINT_ARTIFACT_EXT}",
)
def _load_writes_map(
self, thread_id: str, checkpoint_ns: str, checkpoint_id: str
) -> dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]]:
merged: dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]] = {}
legacy = self._writes_legacy_path(thread_id, checkpoint_ns, checkpoint_id)
if self.fs.exists(legacy):
with self.fs.open(legacy, "rb") as f:
try:
merged.update(_unpack_writes_map_file(f.read()))
except (ValueError, struct.error, OSError):
pass
shard_dir = self._writes_shard_dir(thread_id, checkpoint_ns, checkpoint_id)
if self.fs.exists(shard_dir):
try:
entries = self.fs.ls(shard_dir, detail=False)
except OSError:
entries = []
for entry in entries:
name = entry.rsplit("/", 1)[-1]
if not name.endswith(_CHECKPOINT_ARTIFACT_EXT) or not name.startswith("w_"):
continue
with self.fs.open(entry, "rb") as f:
try:
shard = _unpack_writes_map_file(f.read())
except (ValueError, struct.error, OSError):
continue
merged.update(shard)
return merged
def _load_task_writes_shard(
self, shard_path: str
) -> dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]]:
if not self.fs.exists(shard_path):
return {}
with self.fs.open(shard_path, "rb") as f:
try:
return _unpack_writes_map_file(f.read())
except (ValueError, struct.error, OSError):
return {}
def _save_task_writes_shard(
self,
shard_path: str,
data: dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
) -> None:
with self.fs.open(shard_path, "wb") as f:
f.write(_pack_writes_map_file(data))
def _load_blobs(
self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
) -> dict[str, Any]:
channel_values: dict[str, Any] = {}
for k, v in versions.items():
bp = self._blob_path(thread_id, checkpoint_ns, k, v)
if not self.fs.exists(bp):
continue
with self.fs.open(bp, "rb") as f:
try:
typed = _unpack_blob_file(f.read())
except (ValueError, struct.error, OSError):
continue
if typed[0] != "empty":
channel_values[k] = self.serde.loads_typed(typed)
return channel_values
def _tuple_from_saved(
self,
*,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
saved: tuple[tuple[str, bytes], tuple[str, bytes], str | None],
response_config: RunnableConfig,
metadata: CheckpointMetadata | None = None,
) -> CheckpointTuple:
checkpoint_bytes, metadata_bytes, parent_checkpoint_id = saved
writes_map = self._load_writes_map(thread_id, checkpoint_ns, checkpoint_id)
writes = writes_map.values()
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint_bytes)
metadata_ = metadata if metadata is not None else self.serde.loads_typed(metadata_bytes)
return CheckpointTuple(
config=response_config,
checkpoint={
**checkpoint_,
"channel_values": self._load_blobs(
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
),
},
metadata=metadata_,
pending_writes=[(wid, c, self.serde.loads_typed(v)) for wid, c, v, _ in writes],
parent_config=(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None
),
)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
thread_id: str = config["configurable"]["thread_id"]
checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
ns_root = self._ns_root(thread_id, checkpoint_ns)
cpts_dir = _path_join(ns_root, "cpts")
if checkpoint_id := get_checkpoint_id(config):
path = self._cpt_path(thread_id, checkpoint_ns, checkpoint_id)
if not self.fs.exists(path):
return None
with self.fs.open(path, "rb") as f:
saved = _unpack_cpt_file(f.read())
return self._tuple_from_saved(
thread_id=thread_id,
checkpoint_ns=checkpoint_ns,
checkpoint_id=checkpoint_id,
saved=saved,
response_config=config,
)
if not self.fs.exists(cpts_dir):
return None
entries = self.fs.ls(cpts_dir, detail=False)
if not entries:
return None
ids: list[str] = []
for e in entries:
name = e.rsplit("/", 1)[-1]
if name.endswith(_CHECKPOINT_ARTIFACT_EXT):
ids.append(name[: -len(_CHECKPOINT_ARTIFACT_EXT)])
if not ids:
return None
checkpoint_id = max(ids)
path = self._cpt_path(thread_id, checkpoint_ns, checkpoint_id)
with self.fs.open(path, "rb") as f:
saved = _unpack_cpt_file(f.read())
return self._tuple_from_saved(
thread_id=thread_id,
checkpoint_ns=checkpoint_ns,
checkpoint_id=checkpoint_id,
saved=saved,
response_config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
)
def list(
self,
config: RunnableConfig | None,
*,
filter: dict[str, Any] | None = None,
before: RunnableConfig | None = None,
limit: int | None = None,
) -> Iterator[CheckpointTuple]:
checkpoints_base = _path_join(self.root, _CHECKPOINT_TREE_DIR)
if config:
thread_ids = [config["configurable"]["thread_id"]]
else:
if not self.fs.exists(checkpoints_base):
return
thread_ids = []
for e in self.fs.ls(checkpoints_base, detail=False):
seg = e.rstrip("/").rsplit("/", 1)[-1]
thread_ids.append(_decoder_segment(seg))
config_checkpoint_ns = config["configurable"].get("checkpoint_ns") if config else None
config_checkpoint_id = get_checkpoint_id(config) if config else None
for thread_id in thread_ids:
troot = _path_join(checkpoints_base, _encoder_segment(thread_id), "ns")
if not self.fs.exists(troot):
continue
for ns_entry in self.fs.ls(troot, detail=False):
enc_ns = ns_entry.rstrip("/").rsplit("/", 1)[-1]
checkpoint_ns = _decoder_segment(enc_ns)
if config_checkpoint_ns is not None and checkpoint_ns != config_checkpoint_ns:
continue
cpts_dir = _path_join(ns_entry.rstrip("/"), "cpts")
if not self.fs.exists(cpts_dir):
continue
checkpoint_entries: list[tuple[str, str]] = []
for cp_entry in self.fs.ls(cpts_dir, detail=False):
fname = cp_entry.rsplit("/", 1)[-1]
if not fname.endswith(_CHECKPOINT_ARTIFACT_EXT):
continue
cid = fname[: -len(_CHECKPOINT_ARTIFACT_EXT)]
checkpoint_entries.append((cid, cp_entry))
for checkpoint_id, path in sorted(
checkpoint_entries,
key=lambda x: x[0],
reverse=True,
):
if config_checkpoint_id and checkpoint_id != config_checkpoint_id:
continue
if (
before
and (before_checkpoint_id := get_checkpoint_id(before))
and checkpoint_id >= before_checkpoint_id
):
continue
with self.fs.open(path, "rb") as f:
saved = _unpack_cpt_file(f.read())
parsed_metadata: CheckpointMetadata | None = None
if filter:
parsed_metadata = self.serde.loads_typed(saved[1])
if not all(
query_value == parsed_metadata.get(query_key)
for query_key, query_value in filter.items()
):
continue
if limit is not None and limit <= 0:
return
elif limit is not None:
limit -= 1
yield self._tuple_from_saved(
thread_id=thread_id,
checkpoint_ns=checkpoint_ns,
checkpoint_id=checkpoint_id,
saved=saved,
response_config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
metadata=parsed_metadata,
)
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
c = checkpoint.copy()
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc]
for k, v in new_versions.items():
blob = self.serde.dumps_typed(values[k]) if k in values else ("empty", b"")
bp = self._blob_path(thread_id, checkpoint_ns, k, v)
with self.fs.open(bp, "wb") as f:
f.write(_pack_blob_file(blob))
saved = (
self.serde.dumps_typed(c),
self.serde.dumps_typed(get_checkpoint_metadata(config, metadata)),
config["configurable"].get("checkpoint_id"),
)
cpp = self._cpt_path(thread_id, checkpoint_ns, checkpoint["id"])
with self.fs.open(cpp, "wb") as f:
f.write(_pack_cpt_file(saved))
return {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint["id"],
}
}
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"]["checkpoint_id"]
shard_path = self._task_writes_shard_path(thread_id, checkpoint_ns, checkpoint_id, task_id)
with _locked_task_writes_shard(shard_path):
writes_map = self._load_task_writes_shard(shard_path)
for idx, (c, v) in enumerate(writes):
inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
if inner_key[1] >= 0 and inner_key in writes_map:
continue
writes_map[inner_key] = (
task_id,
c,
self.serde.dumps_typed(v),
task_path,
)
self._save_task_writes_shard(shard_path, writes_map)
def delete_thread(self, thread_id: str) -> None:
tdir = _path_join(self.root, _CHECKPOINT_TREE_DIR, _encoder_segment(thread_id))
if self.fs.exists(tdir):
self.fs.rm(tdir, recursive=True)
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
return await asyncio.to_thread(self.get_tuple, config)
async def alist(
self,
config: RunnableConfig | None,
*,
filter: dict[str, Any] | None = None,
before: RunnableConfig | None = None,
limit: int | None = None,
) -> AsyncIterator[CheckpointTuple]:
tuples = await asyncio.to_thread(
list,
self.list(config, filter=filter, before=before, limit=limit),
)
for item in tuples:
yield item
async def aput(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
return await asyncio.to_thread(self.put, config, checkpoint, metadata, new_versions)
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
await asyncio.to_thread(self.put_writes, config, writes, task_id, task_path)
async def adelete_thread(self, thread_id: str) -> None:
await asyncio.to_thread(self.delete_thread, thread_id)
def get_next_version(
self,
current: str | int | float | None,
channel: ChannelProtocol[Any, Any, Any] | None,
) -> str:
if current is None:
current_v = 0
elif isinstance(current, (int, float)):
current_v = int(current)
else:
current_v = int(str(current).split(".")[0])
next_v = current_v + 1
# Fixed-width decimal suffix (not ``random.random()`` + float format, which inserts an
# extra dot and breaks zero-padding for some values).
suffix = random.randrange(10**16)
return f"{next_v:032}.{suffix:016d}"