mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-04 07:30:21 +00:00
fix(auto-dream): checkpoint extraction and resume interrupted integration
This commit is contained in:
parent
c773fa7b77
commit
549b663ad0
7 changed files with 827 additions and 40 deletions
239
reme/steps/evolve/dream/checkpoint.py
Normal file
239
reme/steps/evolve/dream/checkpoint.py
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
"""Workspace-local, atomic checkpoints for interrupted Dream invocations."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from ....schema import DreamState, DreamUnit
|
||||
|
||||
|
||||
def _receipt(state: DreamState, index: int) -> dict:
|
||||
unit = state.units[index]
|
||||
return next(
|
||||
result
|
||||
for result in state.integrate_results
|
||||
if (result.get("unit"), result.get("bucket"), result.get("paths"))
|
||||
== (unit["name"], unit["bucket"], unit["paths"])
|
||||
)
|
||||
|
||||
|
||||
def integration_lock(step) -> asyncio.Lock:
|
||||
"""Serialize digest integration and catalog finalization within one Application."""
|
||||
if step.app_context is None:
|
||||
raise RuntimeError("dream integration requires app_context")
|
||||
key = "dream_integration_lock"
|
||||
lock = step.app_context.metadata.get(key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
step.app_context.metadata[key] = lock
|
||||
if not isinstance(lock, asyncio.Lock):
|
||||
raise TypeError(f"app_context.metadata[{key!r}] must be an asyncio.Lock")
|
||||
return lock
|
||||
|
||||
|
||||
class DreamCheckpoint(BaseModel):
|
||||
"""Versioned intermediate results; source watermarks remain in the dream catalog."""
|
||||
|
||||
version: Literal[1] = 1
|
||||
run_id: str = Field(default_factory=lambda: uuid4().hex)
|
||||
sources: dict[str, tuple[int, int, str]]
|
||||
state: DreamState
|
||||
completed: list[int] = Field(default_factory=list)
|
||||
targets: dict[str, str] = Field(default_factory=dict)
|
||||
active: int | None = None
|
||||
active_before: dict[str, tuple[int, int]] = Field(default_factory=dict)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_progress(self):
|
||||
"""Reject invalid progress rather than silently skipping unprocessed units."""
|
||||
for unit in self.state.units:
|
||||
parsed = DreamUnit.model_validate(unit)
|
||||
if not parsed.paths or any(
|
||||
p not in self.sources or p not in self.state.changed_paths for p in parsed.paths
|
||||
):
|
||||
raise ValueError("dream checkpoint unit references an unscanned source")
|
||||
indices = [*self.completed, *([] if self.active is None else [self.active])]
|
||||
if len(set(indices)) != len(indices) or any(i < 0 or i >= len(self.state.units) for i in indices):
|
||||
raise ValueError("invalid dream checkpoint unit indices")
|
||||
try:
|
||||
for index in self.completed:
|
||||
_receipt(self.state, index)
|
||||
except StopIteration as error:
|
||||
raise ValueError("completed dream unit has no integration receipt") from error
|
||||
return self
|
||||
|
||||
|
||||
def source_versions(workspace: Path, paths: dict[str, float]) -> dict[str, tuple[int, int, str]]:
|
||||
"""Fingerprint inputs before extraction, including edits that preserve mtimes."""
|
||||
versions = {}
|
||||
for rel in paths:
|
||||
path = (workspace / rel).resolve()
|
||||
if not path.is_relative_to(workspace):
|
||||
raise ValueError(f"dream source escapes workspace: {rel}")
|
||||
stat = path.stat()
|
||||
if stat.st_mtime != paths[rel]:
|
||||
raise RuntimeError(f"dream source changed after scanning; retry the job: {rel}")
|
||||
with path.open("rb") as stream:
|
||||
digest = hashlib.file_digest(stream, "sha256").hexdigest()
|
||||
after = path.stat()
|
||||
if (after.st_mtime_ns, after.st_size) != (stat.st_mtime_ns, stat.st_size):
|
||||
raise RuntimeError(f"dream source changed while fingerprinting; retry the job: {rel}")
|
||||
versions[rel] = (stat.st_mtime_ns, stat.st_size, digest)
|
||||
return versions
|
||||
|
||||
|
||||
class CheckpointStore:
|
||||
"""One checkpoint per scan window and extraction options, owned by this request."""
|
||||
|
||||
def __init__(self, step, state: DreamState, max_units: int):
|
||||
workspace = Path(state.workspace).resolve()
|
||||
self.workspace = workspace
|
||||
options = [state.dates, state.daily_dir, state.hint, max_units, step.config_value("digest_dir")]
|
||||
options.append(step.file_catalog.name)
|
||||
key = hashlib.sha256(json.dumps(options, ensure_ascii=False).encode("utf-8")).hexdigest()
|
||||
self.path = (workspace / step.config_value("metadata_dir") / "dream" / f"{key}.json").resolve()
|
||||
if not self.path.is_relative_to(workspace):
|
||||
raise ValueError("dream checkpoint escapes workspace")
|
||||
self.logger = step.logger
|
||||
self.record: DreamCheckpoint | None = None
|
||||
self.resumed = False
|
||||
|
||||
def read(self) -> DreamCheckpoint | None:
|
||||
"""Treat malformed/unsupported saved state as unusable, with a visible diagnostic."""
|
||||
try:
|
||||
return DreamCheckpoint.model_validate_json(self.path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except ValueError as error:
|
||||
self.logger.warning(f"[AutoDream] ignoring invalid checkpoint {self.path.name}: {error}")
|
||||
return None
|
||||
|
||||
def restore(self, sources: dict[str, tuple[int, int, str]]) -> DreamState | None:
|
||||
"""Reuse extraction only when the entire scanned input window is unchanged."""
|
||||
record = self.read()
|
||||
if record is None:
|
||||
return None
|
||||
if record.sources != sources:
|
||||
self.logger.info("[AutoDream] checkpoint invalidated: source files changed; extracting from scratch")
|
||||
return None
|
||||
if Path(record.state.workspace).resolve() != self.workspace:
|
||||
self.logger.warning("[AutoDream] ignoring checkpoint for another workspace")
|
||||
return None
|
||||
self.record = record
|
||||
self.resumed = True
|
||||
state = self._prepare()
|
||||
self.logger.info(
|
||||
f"[AutoDream] resuming: {len(record.completed)} of {len(state.units)} units completed from checkpoint, "
|
||||
f"{len(sources)} sources unchanged; {len(state.units) - len(record.completed)} units pending",
|
||||
)
|
||||
return state
|
||||
|
||||
def publish(self, state: DreamState, sources: dict[str, tuple[int, int, str]]) -> DreamState:
|
||||
"""First extraction wins when identical invocations overlap within an Application."""
|
||||
restored = self.restore(sources)
|
||||
if restored is not None:
|
||||
return restored
|
||||
self.record = DreamCheckpoint(sources=sources, state=state)
|
||||
self._write()
|
||||
return state
|
||||
|
||||
def refresh(self) -> DreamState:
|
||||
"""Reload progress after acquiring the existing application-wide integration lock."""
|
||||
self.record = self._current()
|
||||
return self._prepare()
|
||||
|
||||
def complete(self, index: int, state: DreamState) -> None:
|
||||
"""Acknowledge a unit together with the content version of its actual digest target."""
|
||||
assert self.record is not None
|
||||
target = _receipt(state, index)["target_path"]
|
||||
version = self._target_version(target)
|
||||
if version is None:
|
||||
raise RuntimeError(f"Dream digest target disappeared before completion: {target}")
|
||||
self.record.targets[target] = version
|
||||
self.record.completed.append(index)
|
||||
|
||||
def ensure_current(self, *, complete: bool) -> None:
|
||||
"""Reject stale finalizers and do not checkpoint missing or unacknowledged digest content."""
|
||||
current = self._current()
|
||||
if complete:
|
||||
if len(current.completed) != len(current.state.units):
|
||||
raise RuntimeError("Dream integration progress changed before Finish; retry the job")
|
||||
if self._changed_targets(current):
|
||||
raise RuntimeError("Dream digest changed before Finish; retry the job")
|
||||
|
||||
def _current(self) -> DreamCheckpoint:
|
||||
current = self.read()
|
||||
if current is None or self.record is None or current.run_id != self.record.run_id:
|
||||
raise RuntimeError("Dream checkpoint was replaced or completed by another invocation; retry the job")
|
||||
return current
|
||||
|
||||
def _changed_targets(self, record: DreamCheckpoint) -> set[str]:
|
||||
targets = {_receipt(record.state, index)["target_path"] for index in record.completed}
|
||||
return {
|
||||
path
|
||||
for path in targets
|
||||
if not record.targets.get(path) or record.targets[path] != self._target_version(path)
|
||||
}
|
||||
|
||||
def _target_version(self, relative: str) -> str | None:
|
||||
path = (self.workspace / relative).resolve()
|
||||
if not path.is_relative_to(self.workspace):
|
||||
raise ValueError("dream digest target escapes workspace")
|
||||
try:
|
||||
with path.open("rb") as stream:
|
||||
return hashlib.file_digest(stream, "sha256").hexdigest()
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
|
||||
def _prepare(self) -> DreamState:
|
||||
"""Recheck landed content before skipping completed units, including after waiting for the lock."""
|
||||
assert self.record is not None
|
||||
state = self.record.state
|
||||
if Path(state.workspace).resolve() != self.workspace:
|
||||
raise ValueError("dream checkpoint belongs to another workspace")
|
||||
invalid = self._changed_targets(self.record)
|
||||
if invalid:
|
||||
self.logger.info(f"[AutoDream] resume recheck: completed digest targets changed: {sorted(invalid)}")
|
||||
self.record.completed = [
|
||||
index for index in self.record.completed if _receipt(state, index)["target_path"] not in invalid
|
||||
]
|
||||
state.integrate_results = [
|
||||
result for result in state.integrate_results if result["target_path"] not in invalid
|
||||
]
|
||||
state.modified_paths = list(dict.fromkeys([*state.modified_paths, *sorted(invalid)]))
|
||||
self.resumed = True
|
||||
state.failed_units, state.failed_paths, state.errors, state.checkpoint_paths = [], [], [], []
|
||||
return state
|
||||
|
||||
def save(self, state: DreamState) -> None:
|
||||
"""Fail closed on persistence errors so Finish cannot acknowledge lost progress."""
|
||||
self._current()
|
||||
self.record.state = state
|
||||
self._write()
|
||||
|
||||
def remove(self) -> None:
|
||||
"""Only remove this generation, after the final catalog checkpoint is durable."""
|
||||
current = self.read()
|
||||
if current is not None and self.record is not None and current.run_id == self.record.run_id:
|
||||
self.path.unlink()
|
||||
|
||||
def _write(self) -> None:
|
||||
"""Publish complete JSON with atomic replacement; never truncate the last checkpoint."""
|
||||
assert self.record is not None
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
descriptor, temporary = tempfile.mkstemp(prefix=f".{self.path.stem}-", dir=self.path.parent)
|
||||
try:
|
||||
with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
|
||||
stream.write(self.record.model_dump_json(indent=2))
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary, self.path)
|
||||
finally:
|
||||
Path(temporary).unlink(missing_ok=True)
|
||||
|
|
@ -8,6 +8,7 @@ from .._evolve import agent_reply_result_text
|
|||
from ....components import R
|
||||
from ....enumeration import DreamBucketEnum
|
||||
from ....schema import DreamExtractOutput, DreamState
|
||||
from .checkpoint import CheckpointStore, source_versions
|
||||
from .utils import (
|
||||
clean_paths,
|
||||
daily_dir,
|
||||
|
|
@ -103,6 +104,20 @@ class DreamExtractStep(BaseStep):
|
|||
existing=existing,
|
||||
indexed=indexed,
|
||||
)
|
||||
checkpoint = CheckpointStore(self, state, max_units)
|
||||
saved = checkpoint.read()
|
||||
if changed or saved is not None:
|
||||
versions = source_versions(workspace, existing)
|
||||
restored = checkpoint.restore(versions)
|
||||
if restored is not None:
|
||||
self.context["_dream_checkpoint"] = checkpoint
|
||||
return self._finish(restored, True, f"Resumed {len(restored.units)} extracted unit(s) from checkpoint")
|
||||
if saved is not None:
|
||||
# A previous failed dump may have advanced only in-memory mtimes. Validate content, too.
|
||||
changed = list(dict.fromkeys([*changed, *(p for p in existing if versions[p] != saved.sources.get(p))]))
|
||||
state.changed_paths = changed
|
||||
state.unchanged_paths = [p for p in existing if p not in changed]
|
||||
state.files_changed, state.files_unchanged = len(changed), len(state.unchanged_paths)
|
||||
if not changed:
|
||||
self.logger.info(f"[{self.name}] skip no changed input dates={','.join(dates)}")
|
||||
return self._finish(state, True, f"No changed dream input for {', '.join(dates)}")
|
||||
|
|
@ -112,7 +127,7 @@ class DreamExtractStep(BaseStep):
|
|||
self.logger.warning(f"[{self.name}] skip no llm changed={len(changed)}")
|
||||
return self._finish(state, False, state.errors[-1])
|
||||
|
||||
self.logger.info(f"[{self.name}] agent start changed={len(changed)} dates={len(dates)}")
|
||||
self.logger.info(f"[{self.name}] extracting from scratch: changed={len(changed)} dates={len(dates)}")
|
||||
user_message = self.prompt_format(
|
||||
"extract_user_message",
|
||||
date=day,
|
||||
|
|
@ -162,6 +177,9 @@ class DreamExtractStep(BaseStep):
|
|||
return self._finish(state, False, error)
|
||||
|
||||
self.logger.info(f"[{self.name}] parse done keys={','.join(sorted(meta.keys())) if meta else '(none)'}")
|
||||
state = checkpoint.publish(state, versions)
|
||||
self.context["_dream_checkpoint"] = checkpoint
|
||||
self.logger.info(f"[{self.name}] extraction checkpoint saved: {len(state.units)} units")
|
||||
answer = f"Extracted {len(state.units)} unit(s) from {len(changed)} changed file(s) across {len(dates)} day(s)"
|
||||
return self._finish(state, True, answer)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from pathlib import Path
|
|||
from ...base_step import BaseStep
|
||||
from ....components import R
|
||||
from ....schema import DreamState, FileNode
|
||||
from .checkpoint import CheckpointStore, integration_lock
|
||||
from .utils import state_from_context, store_state, workspace_dir
|
||||
|
||||
|
||||
|
|
@ -19,6 +20,15 @@ class DreamFinishStep(BaseStep):
|
|||
async def execute(self):
|
||||
assert self.context is not None
|
||||
state = state_from_context(self)
|
||||
saved = self.context.get("_dream_checkpoint")
|
||||
if isinstance(saved, CheckpointStore):
|
||||
async with integration_lock(self):
|
||||
saved.ensure_current(complete=not state.errors and not state.failed_units)
|
||||
return await self._execute(state)
|
||||
return await self._execute(state)
|
||||
|
||||
async def _execute(self, state: DreamState):
|
||||
"""Persist catalog versions while holding the same lock as checkpointed integrations."""
|
||||
workspace = Path(state.workspace).resolve() if state.workspace else workspace_dir(self)
|
||||
if self.file_catalog is None:
|
||||
raise RuntimeError("dream_finish_step requires file_catalog")
|
||||
|
|
@ -46,6 +56,9 @@ class DreamFinishStep(BaseStep):
|
|||
self.logger.info(f"[{self.name}] catalog dump done")
|
||||
|
||||
state.checkpoint_paths = [n.path for n in upserts if n.path in checkpoint]
|
||||
saved = self.context.get("_dream_checkpoint")
|
||||
if isinstance(saved, CheckpointStore) and self.persist and not state.errors and not state.failed_units:
|
||||
saved.remove()
|
||||
state.summary = render_summary(state)
|
||||
store_state(self, state)
|
||||
self.context.response.success = not state.errors
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from .._evolve import agent_reply_result_text
|
|||
from ....components import R
|
||||
from ....enumeration import DreamBucketEnum
|
||||
from ....schema import IntegrateOutcome
|
||||
from .checkpoint import CheckpointStore, integration_lock
|
||||
from .utils import llm_available, pack_paths, parse_structured_reply, state_from_context, store_state, workspace_dir
|
||||
|
||||
_TOOLS = ("node_search", "read", "frontmatter_read", "write", "edit", "frontmatter_update")
|
||||
|
|
@ -62,13 +63,8 @@ class DreamIntegrateStep(BaseStep):
|
|||
if not state.units:
|
||||
self.logger.info(f"[{self.name}] skip no units")
|
||||
return self._finish(state, not state.errors, "No dream units to integrate")
|
||||
if not llm_available(self):
|
||||
err = "no llm configured; dream integrate requires an LLM"
|
||||
state.errors.append(err)
|
||||
state.failed_units = state.units
|
||||
state.failed_paths = sorted({p for u in state.units for p in u.get("paths", [])})
|
||||
self.logger.warning(f"[{self.name}] skip no llm units={len(state.units)}")
|
||||
return self._finish(state, False, err)
|
||||
if not isinstance(self.context.get("_dream_checkpoint"), CheckpointStore) and not llm_available(self):
|
||||
return self._missing_llm(state, state.units)
|
||||
|
||||
async with self._integration_lock():
|
||||
return await self._execute_locked(state)
|
||||
|
|
@ -76,6 +72,12 @@ class DreamIntegrateStep(BaseStep):
|
|||
async def _execute_locked(self, state):
|
||||
"""Integrate units while serializing digest writes for this Application."""
|
||||
|
||||
checkpoint = self.context.get("_dream_checkpoint")
|
||||
if isinstance(checkpoint, CheckpointStore):
|
||||
state = checkpoint.refresh()
|
||||
pending = [unit for i, unit in enumerate(state.units) if i not in checkpoint.record.completed]
|
||||
if pending and not llm_available(self):
|
||||
return self._missing_llm(state, pending)
|
||||
workspace = Path(state.workspace).resolve() if state.workspace else workspace_dir(self)
|
||||
digest_dir = self.config_value("digest_dir")
|
||||
self.logger.info(f"[{self.name}] start units={len(state.units)} workspace={workspace} digest_dir={digest_dir}")
|
||||
|
|
@ -85,12 +87,47 @@ class DreamIntegrateStep(BaseStep):
|
|||
# Compare the whole run so retries and later units updating a newly
|
||||
# created file still emit one "added" change for the final file.
|
||||
before = _snapshot_digest(workspace, digest_dir)
|
||||
interrupted_index = None
|
||||
if isinstance(checkpoint, CheckpointStore) and checkpoint.record.active is not None:
|
||||
interrupted_index = checkpoint.record.active
|
||||
# Preserve ambiguous writes even if an earlier unit now needs rechecking first.
|
||||
_record_modified_paths(state, _changed_digest_paths(checkpoint.record.active_before, before))
|
||||
for i, unit in enumerate(state.units, start=1):
|
||||
await self._integrate_one(state, unit, i, workspace, digest_dir)
|
||||
recovery_note = ""
|
||||
if isinstance(checkpoint, CheckpointStore):
|
||||
record = checkpoint.record
|
||||
if i - 1 in record.completed:
|
||||
self.logger.info(f"[{self.name}] resume skip completed unit {i}/{len(state.units)}")
|
||||
continue
|
||||
interrupted = interrupted_index == i - 1
|
||||
possible_writes = list(state.modified_paths)
|
||||
if interrupted or checkpoint.resumed:
|
||||
recovery_note = "\n\n" + self.prompt_format(
|
||||
"integrate_resume_message",
|
||||
possible_paths_json=json.dumps(sorted(set(possible_writes)), ensure_ascii=False),
|
||||
)
|
||||
self.logger.info(f"[{self.name}] resume reconcile pending unit {i}/{len(state.units)}")
|
||||
record.active = i - 1
|
||||
record.active_before = _snapshot_digest(workspace, digest_dir)
|
||||
checkpoint.save(state)
|
||||
failures = len(state.failed_units)
|
||||
await self._integrate_one(state, unit, i, workspace, digest_dir, recovery_note=recovery_note)
|
||||
if isinstance(checkpoint, CheckpointStore):
|
||||
if len(state.failed_units) == failures:
|
||||
checkpoint.complete(i - 1, state)
|
||||
checkpoint.record.active = None
|
||||
checkpoint.record.active_before = {}
|
||||
checkpoint.save(state)
|
||||
self.logger.info(
|
||||
f"[{self.name}] progress checkpoint saved: "
|
||||
f"{len(checkpoint.record.completed)}/{len(state.units)} units completed",
|
||||
)
|
||||
after = _snapshot_digest(workspace, digest_dir)
|
||||
changed_paths = set(_changed_digest_paths(before, after)) | set(state.modified_paths)
|
||||
self.context["changes"] = [
|
||||
{"change": "modified" if path in before else "added", "path": path}
|
||||
for path in _changed_digest_paths(before, after)
|
||||
{"change": "added" if path in state.nodes_created or path not in before else "modified", "path": path}
|
||||
for path in sorted(changed_paths)
|
||||
if (workspace / path).is_file()
|
||||
]
|
||||
state.failed_paths = sorted(set(state.failed_paths))
|
||||
answer = (
|
||||
|
|
@ -99,20 +136,29 @@ class DreamIntegrateStep(BaseStep):
|
|||
)
|
||||
return self._finish(state, not state.failed_units, answer)
|
||||
|
||||
def _missing_llm(self, state, pending):
|
||||
"""Keep pending sources eligible without invalidating already completed progress."""
|
||||
err = "no llm configured; dream integrate requires an LLM"
|
||||
state.errors.append(err)
|
||||
state.failed_units = pending
|
||||
state.failed_paths = sorted({p for unit in pending for p in unit.get("paths", [])})
|
||||
self.logger.warning(f"[{self.name}] skip no llm units={len(pending)}")
|
||||
return self._finish(state, False, err)
|
||||
|
||||
def _integration_lock(self) -> asyncio.Lock:
|
||||
"""Return the Application-wide lock used by AutoDream integration."""
|
||||
if self.app_context is None:
|
||||
raise RuntimeError("dream_integrate_step requires app_context")
|
||||
key = "dream_integration_lock"
|
||||
lock = self.app_context.metadata.get(key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self.app_context.metadata[key] = lock
|
||||
if not isinstance(lock, asyncio.Lock):
|
||||
raise TypeError(f"app_context.metadata[{key!r}] must be an asyncio.Lock")
|
||||
return lock
|
||||
return integration_lock(self)
|
||||
|
||||
async def _integrate_one(self, state, unit: dict, index: int, workspace: Path, digest_dir: str) -> None:
|
||||
async def _integrate_one(
|
||||
self,
|
||||
state,
|
||||
unit: dict,
|
||||
index: int,
|
||||
workspace: Path,
|
||||
digest_dir: str,
|
||||
*,
|
||||
recovery_note: str = "",
|
||||
) -> None:
|
||||
try:
|
||||
bucket = DreamBucketEnum(str(unit.get("bucket") or "")).value
|
||||
except ValueError:
|
||||
|
|
@ -137,7 +183,8 @@ class DreamIntegrateStep(BaseStep):
|
|||
unit_summary=unit.get("summary", ""),
|
||||
unit_paths_json=json.dumps(paths, ensure_ascii=False, indent=2),
|
||||
material_blob=pack_paths(workspace, paths),
|
||||
),
|
||||
)
|
||||
+ recovery_note,
|
||||
system_prompt=self.prompt_format(
|
||||
f"integrate_system_prompt_{bucket}",
|
||||
workspace_dir=str(workspace),
|
||||
|
|
@ -151,13 +198,17 @@ class DreamIntegrateStep(BaseStep):
|
|||
changed = _changed_digest_paths(before, after)
|
||||
_record_modified_paths(state, changed)
|
||||
created = len(changed) == 1 and changed[0] not in unit_before
|
||||
if len(changed) == 1 and self._valid_target(
|
||||
workspace,
|
||||
digest_dir,
|
||||
bucket,
|
||||
changed[0],
|
||||
action="CREATE" if created else "UPDATED",
|
||||
existed_before=changed[0] in unit_before,
|
||||
if (
|
||||
not recovery_note
|
||||
and len(changed) == 1
|
||||
and self._valid_target(
|
||||
workspace,
|
||||
digest_dir,
|
||||
bucket,
|
||||
changed[0],
|
||||
action="CREATE" if created else "UPDATED",
|
||||
existed_before=changed[0] in unit_before,
|
||||
)
|
||||
):
|
||||
self._record_recovered(state, unit, bucket, paths, changed[0], created, e)
|
||||
self.logger.warning(
|
||||
|
|
@ -193,13 +244,17 @@ class DreamIntegrateStep(BaseStep):
|
|||
changed = _changed_digest_paths(before, after)
|
||||
_record_modified_paths(state, changed)
|
||||
created = len(changed) == 1 and changed[0] not in unit_before
|
||||
if len(changed) == 1 and self._valid_target(
|
||||
workspace,
|
||||
digest_dir,
|
||||
bucket,
|
||||
changed[0],
|
||||
action="CREATE" if created else "UPDATED",
|
||||
existed_before=changed[0] in unit_before,
|
||||
if (
|
||||
not recovery_note
|
||||
and len(changed) == 1
|
||||
and self._valid_target(
|
||||
workspace,
|
||||
digest_dir,
|
||||
bucket,
|
||||
changed[0],
|
||||
action="CREATE" if created else "UPDATED",
|
||||
existed_before=changed[0] in unit_before,
|
||||
)
|
||||
):
|
||||
self._record_recovered(state, unit, bucket, paths, changed[0], created, e)
|
||||
self.logger.warning(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,15 @@
|
|||
integrate_resume_message: |
|
||||
An earlier Dream integration attempt was interrupted or failed. Digest writes may already have landed without a
|
||||
saved completion receipt. Inspect the existing digest and reconcile this unit before writing. Reuse existing
|
||||
content and do not append duplicate evidence or create a duplicate node. Verify that the unit is fully integrated
|
||||
before returning a successful receipt; do not assume that a changed file means the work is complete.
|
||||
Files possibly written by earlier attempts: {possible_paths_json}
|
||||
|
||||
integrate_resume_message_zh: |
|
||||
上一次 Dream 整合曾中断或失败,digest 文件可能已写入,但尚未保存完成回执。请在写入前检查已有 digest 并核对本 unit。
|
||||
复用已有内容,不要重复追加证据或创建重复节点。确认本 unit 已完整整合后再返回成功回执,不能仅凭文件有变化就认定完成。
|
||||
上一次尝试可能写入的文件:{possible_paths_json}
|
||||
|
||||
integrate_system_prompt_procedure: |
|
||||
You are the dream unit integration agent for the procedure bucket. The unit
|
||||
is a how-to-do-X abstraction: steps, methods, recipes, runbooks, executable
|
||||
|
|
|
|||
|
|
@ -153,7 +153,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older
|
|||
index = f"daily/{DAY}.md"
|
||||
(tmp_path / index).write_text("day index")
|
||||
|
||||
def agent_for(app, write_attempt):
|
||||
def agent_for(app, write_attempt, *, valid_receipt=False):
|
||||
calls = {"extract": 0, "integrate": 0}
|
||||
|
||||
async def reply(_message, **kwargs):
|
||||
|
|
@ -167,6 +167,8 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older
|
|||
if calls["integrate"] == write_attempt:
|
||||
written = await app.run_job("write", path=TARGET, content=OLD)
|
||||
assert written.success, written.answer
|
||||
if valid_receipt:
|
||||
return {"result": json.dumps({"action": "CREATE", "target_path": TARGET})}
|
||||
return {"result": "{}"}
|
||||
|
||||
agent = AsyncMock(spec=BaseAgentWrapper)
|
||||
|
|
@ -217,7 +219,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older
|
|||
catalog = app.context.components["file_catalog"]["dream"]
|
||||
if recover_on_attempt is None:
|
||||
assert {n.path: n.st_mtime for n in await catalog.get_nodes()} == previous
|
||||
agent, calls = agent_for(app, 1)
|
||||
agent, calls = agent_for(app, 1, valid_receipt=True)
|
||||
response = await app.run_job(
|
||||
"auto_dream",
|
||||
date=DAY,
|
||||
|
|
@ -226,7 +228,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older
|
|||
as_llm=Mock(spec=ChatModelBase),
|
||||
)
|
||||
assert response.success, response.answer
|
||||
assert calls == {"extract": int(recover_on_attempt is None), "integrate": int(recover_on_attempt is None)}
|
||||
assert calls == {"extract": 0, "integrate": int(recover_on_attempt is None)}
|
||||
assert {SOURCE, index} <= {n.path for n in await catalog.get_nodes()}
|
||||
assert OLD in (tmp_path / TARGET).read_text()
|
||||
assert source.read_text() == original
|
||||
|
|
|
|||
448
tests/unit/test_dream_resume.py
Normal file
448
tests/unit/test_dream_resume.py
Normal file
|
|
@ -0,0 +1,448 @@
|
|||
"""Exercise Dream recovery with fresh Applications and a controlled local provider."""
|
||||
|
||||
# pylint: disable=protected-access
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
from contextlib import asynccontextmanager, nullcontext
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from agentscope.model import ChatModelBase
|
||||
|
||||
from reme.application import Application
|
||||
from reme.components.agent_wrapper import BaseAgentWrapper
|
||||
from reme.config import resolve_app_config
|
||||
from reme.steps.evolve.dream.checkpoint import CheckpointStore, DreamCheckpoint, source_versions
|
||||
from reme.steps.evolve.dream.finish import DreamFinishStep
|
||||
from reme.steps.evolve.dream.integrate import DreamIntegrateStep
|
||||
|
||||
DAY = "2026-10-23"
|
||||
SOURCE = f"daily/{DAY}/source.md"
|
||||
TARGETS = ["digest/wiki/first.md", "digest/wiki/second.md"]
|
||||
UNITS = [
|
||||
{"name": name, "bucket": "wiki", "summary": f"{name} memory", "paths": [SOURCE]} for name in ("first", "second")
|
||||
]
|
||||
|
||||
|
||||
class DreamCase:
|
||||
"""Own isolated files, fresh application lifecycles, and provider call observations."""
|
||||
|
||||
def __init__(self, workspace):
|
||||
self.workspace, self.metadata = workspace, "metadata"
|
||||
self.source = workspace / SOURCE
|
||||
self.source.parent.mkdir(parents=True)
|
||||
self.source.write_text("Original source", encoding="utf-8")
|
||||
self.units, self.calls, self.prompts = copy.deepcopy(UNITS), [], []
|
||||
|
||||
@property
|
||||
def checkpoints(self):
|
||||
"""Return only committed JSON checkpoints, excluding atomic-write temporary files."""
|
||||
return list((self.workspace / self.metadata / "dream").glob("*.json"))
|
||||
|
||||
@property
|
||||
def saved(self):
|
||||
"""Read progress from disk, without borrowing any application state."""
|
||||
assert len(self.checkpoints) == 1
|
||||
return DreamCheckpoint.model_validate_json(self.checkpoints[0].read_text())
|
||||
|
||||
@asynccontextmanager
|
||||
async def open(self):
|
||||
"""Start real local catalogs and jobs without external services or model providers."""
|
||||
defaults = resolve_app_config(config="default", log_config=False)
|
||||
types = ("tokenizer", "keyword_index", "file_graph", "tag_index", "file_store", "file_catalog")
|
||||
app = Application(
|
||||
workspace_dir=str(self.workspace),
|
||||
metadata_dir=self.metadata,
|
||||
enable_logo=False,
|
||||
log_to_file=False,
|
||||
log_to_console=False,
|
||||
service={"backend": "http", "web_enabled": False, "mcp_enabled": False},
|
||||
components={name: copy.deepcopy(defaults["components"][name]) for name in types},
|
||||
jobs={name: copy.deepcopy(defaults["jobs"][name]) for name in ("auto_dream", "write")},
|
||||
)
|
||||
self.calls.clear()
|
||||
self.prompts.clear()
|
||||
await app.start()
|
||||
try:
|
||||
yield app
|
||||
finally:
|
||||
await app.close()
|
||||
|
||||
async def write(self, app, index, content):
|
||||
"""Exercise the normal workspace-contained write job."""
|
||||
response = await app.run_job("write", path=TARGETS[index], content=content)
|
||||
assert response.success, response.answer
|
||||
|
||||
@staticmethod
|
||||
def receipt(index, action):
|
||||
"""Return a valid integration receipt for a controlled target."""
|
||||
return {"result": json.dumps({"action": action, "target_path": TARGETS[index]})}
|
||||
|
||||
async def run(self, app, *, reply=None, interrupt="", fail=False, **options):
|
||||
"""Track extraction and integration independently; custom replies control only integration."""
|
||||
|
||||
async def provider(message, **kwargs):
|
||||
if "list_tags" in kwargs["job_tools"]:
|
||||
return {"result": "No tags added"}
|
||||
if kwargs["job_tools"] == ["read"]:
|
||||
self.calls.append("extract")
|
||||
return {"result": json.dumps({"units": self.units})}
|
||||
text = str(message)
|
||||
index = 0 if "first memory" in text else 1
|
||||
self.calls.append(UNITS[index]["name"])
|
||||
self.prompts.append(text)
|
||||
if reply:
|
||||
return await reply(index, text)
|
||||
if index == 1 and interrupt == "before_write":
|
||||
raise asyncio.CancelledError
|
||||
if index == 1 and fail:
|
||||
return {"result": "{}"}
|
||||
existed = (self.workspace / TARGETS[index]).exists()
|
||||
if not existed:
|
||||
await self.write(app, index, UNITS[index]["summary"])
|
||||
if index == 1 and interrupt == "after_write":
|
||||
raise asyncio.CancelledError
|
||||
return self.receipt(index, "CORROBORATE" if existed else "CREATE")
|
||||
|
||||
agent = AsyncMock(spec=BaseAgentWrapper)
|
||||
agent.reply.side_effect = provider
|
||||
return await app.run_job(
|
||||
"auto_dream",
|
||||
date=DAY,
|
||||
scan_days=1,
|
||||
agent_wrapper=agent,
|
||||
as_llm=Mock(spec=ChatModelBase),
|
||||
**options,
|
||||
)
|
||||
|
||||
async def interrupt(self, stage):
|
||||
"""Leave disk state at a named interruption boundary, then discard the Application."""
|
||||
boundary = {
|
||||
"after_extract": (DreamIntegrateStep, "execute"),
|
||||
"before_finish": (DreamFinishStep, "execute"),
|
||||
"after_commit": (CheckpointStore, "remove"),
|
||||
}.get(stage)
|
||||
async with self.open() as app:
|
||||
with (
|
||||
patch.object(*boundary, side_effect=asyncio.CancelledError) if boundary else nullcontext(),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await self.run(app, interrupt=stage)
|
||||
if stage != "after_commit":
|
||||
assert await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE]) == []
|
||||
|
||||
|
||||
@pytest.fixture(name="case")
|
||||
def dream_case(tmp_path):
|
||||
"""Each case starts with user-owned source material in a temporary workspace."""
|
||||
return DreamCase(tmp_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stage", ["after_extract", "before_write", "after_write", "before_finish", "no_llm"])
|
||||
async def test_restart_reuses_extraction_and_completed_units(case, stage):
|
||||
"""Resume only pending units; completed work and source content stay untouched."""
|
||||
original = case.source.read_bytes()
|
||||
await case.interrupt("before_finish" if stage == "no_llm" else stage)
|
||||
saved, logs = case.saved, []
|
||||
assert saved.state.units == UNITS
|
||||
assert saved.completed == ([] if stage == "after_extract" else [0] if "write" in stage else [0, 1])
|
||||
if "write" in stage:
|
||||
assert saved.active == 1
|
||||
restore = CheckpointStore.restore
|
||||
|
||||
def capture(store, sources):
|
||||
store.logger = Mock(wraps=store.logger)
|
||||
state = restore(store, sources)
|
||||
logs.extend(call.args[0] for call in store.logger.info.call_args_list)
|
||||
return state
|
||||
|
||||
async with case.open() as app:
|
||||
with (
|
||||
patch.object(CheckpointStore, "restore", capture),
|
||||
(
|
||||
patch("reme.steps.evolve.dream.integrate.llm_available", return_value=False)
|
||||
if stage == "no_llm"
|
||||
else nullcontext()
|
||||
),
|
||||
):
|
||||
response = await case.run(app)
|
||||
assert response.success, response.answer
|
||||
assert case.calls == (
|
||||
["first", "second"] if stage == "after_extract" else ["second"] if "write" in stage else []
|
||||
)
|
||||
assert len(response.metadata["dream"]["integrate_results"]) == 2
|
||||
assert SOURCE in response.metadata["dream"]["checkpoint_paths"]
|
||||
assert any(f"resuming: {len(saved.completed)} of 2 units completed" in line for line in logs)
|
||||
assert not case.checkpoints
|
||||
assert case.source.read_bytes() == original
|
||||
for index, path in enumerate(TARGETS):
|
||||
assert (case.workspace / path).read_text().strip() == UNITS[index]["summary"]
|
||||
if stage == "after_write":
|
||||
assert TARGETS[1] in case.prompts[0] and "do not append duplicate evidence" in case.prompts[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"change",
|
||||
["edit", "same_mtime", "add", "delete", "hint", "max_units", "corrupt", "version", "progress"],
|
||||
)
|
||||
async def test_incompatible_checkpoint_reextracts(case, change):
|
||||
"""Never silently reuse extraction for changed inputs or invalid saved progress."""
|
||||
await case.interrupt("after_extract")
|
||||
options, path = {}, case.checkpoints[0]
|
||||
if change in ("edit", "same_mtime"):
|
||||
stat = case.source.stat()
|
||||
case.source.write_text("Modified source")
|
||||
if change == "same_mtime":
|
||||
os.utime(case.source, ns=(stat.st_atime_ns, stat.st_mtime_ns))
|
||||
elif change == "add":
|
||||
(case.source.parent / "another.md").write_text("New source")
|
||||
elif change == "delete":
|
||||
case.source.unlink()
|
||||
elif change in ("hint", "max_units"):
|
||||
options[change] = "New guidance" if change == "hint" else 1
|
||||
elif change == "corrupt":
|
||||
path.write_text("{incomplete JSON")
|
||||
else:
|
||||
data = json.loads(path.read_text())
|
||||
data["version" if change == "version" else "completed"] = 999 if change == "version" else [0]
|
||||
path.write_text(json.dumps(data))
|
||||
case.units = []
|
||||
async with case.open() as app:
|
||||
response = await case.run(app, **options)
|
||||
assert response.success, response.answer
|
||||
assert case.calls == ["extract"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
async def test_restart_after_catalog_commit_only_cleans_up(case, empty):
|
||||
"""A durable final watermark needs no new model work, including empty extraction."""
|
||||
case.metadata = "custom-metadata"
|
||||
if empty:
|
||||
case.units = []
|
||||
await case.interrupt("after_commit")
|
||||
assert case.checkpoints and not (case.workspace / "metadata").exists()
|
||||
async with case.open() as app:
|
||||
response = await case.run(app)
|
||||
assert response.success, response.answer
|
||||
assert not case.calls and not case.checkpoints
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sources", ["shared", "independent", "missing_target"])
|
||||
async def test_partial_failure_retries_only_pending_work(case, sources):
|
||||
"""Independent source watermarks cannot hide saved units needing target repair."""
|
||||
if sources != "shared":
|
||||
other = f"daily/{DAY}/other.md"
|
||||
(case.workspace / other).write_text("Another source")
|
||||
case.units[1]["paths"] = [other]
|
||||
async with case.open() as app:
|
||||
response = await case.run(app, fail=True)
|
||||
assert not response.success
|
||||
dream = response.metadata["dream"]
|
||||
assert (SOURCE in dream["checkpoint_paths"]) is (sources != "shared")
|
||||
assert case.units[1]["paths"][0] not in dream["checkpoint_paths"]
|
||||
if sources == "missing_target":
|
||||
(case.workspace / TARGETS[0]).unlink()
|
||||
async with case.open() as app:
|
||||
response = await case.run(app)
|
||||
assert response.success, response.answer
|
||||
assert case.calls == (["first", "second"] if sources == "missing_target" else ["second"])
|
||||
assert len(response.metadata["dream"]["integrate_results"]) == 2
|
||||
assert not response.metadata["dream"]["failed_units"] and not response.metadata["dream"]["errors"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("bad_receipt", ["invalid", "exception"])
|
||||
@pytest.mark.parametrize("stage", ["after_extract", "before_write"])
|
||||
async def test_resume_requires_receipt_despite_partial_writes(case, bad_receipt, stage):
|
||||
"""Changed files without valid receipts must not become durable completion markers."""
|
||||
await case.interrupt(stage)
|
||||
async with case.open() as app:
|
||||
|
||||
async def reply(index, text):
|
||||
assert "do not append duplicate evidence" in text
|
||||
await case.write(app, index, f"Partial memory, attempt {len(case.calls)}")
|
||||
if bad_receipt == "exception":
|
||||
raise RuntimeError("incomplete integration")
|
||||
return {"result": "{}"}
|
||||
|
||||
response = await case.run(app, reply=reply)
|
||||
assert not response.success
|
||||
assert len(case.calls) == (4 if stage == "after_extract" else 2)
|
||||
assert SOURCE not in response.metadata["dream"]["checkpoint_paths"]
|
||||
assert case.saved.completed == ([] if stage == "after_extract" else [0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["delete", "rewrite", "later_interrupted_write"])
|
||||
async def test_recheck_only_affected_digest_units(case, change):
|
||||
"""Rechecking an earlier target also preserves evidence of a later interrupted write."""
|
||||
interrupted = change == "later_interrupted_write"
|
||||
await case.interrupt("after_write" if interrupted else "before_finish")
|
||||
first, second = (case.workspace / path for path in TARGETS)
|
||||
second_version = second.stat().st_mtime_ns
|
||||
if change == "delete":
|
||||
first.unlink()
|
||||
else:
|
||||
first.write_text("Overwritten memory")
|
||||
async with case.open() as app:
|
||||
|
||||
async def reply(index, text):
|
||||
assert "do not append duplicate evidence" in text
|
||||
if interrupted:
|
||||
assert TARGETS[1] in text
|
||||
if index == 0:
|
||||
existed = first.exists()
|
||||
await case.write(app, index, UNITS[index]["summary"])
|
||||
return case.receipt(index, "CORRECT" if existed else "CREATE")
|
||||
assert interrupted
|
||||
return case.receipt(index, "CORROBORATE")
|
||||
|
||||
response = await case.run(app, reply=reply)
|
||||
assert response.success, response.answer
|
||||
assert case.calls == (["first", "second"] if interrupted else ["first"])
|
||||
assert first.read_text().strip() == UNITS[0]["summary"]
|
||||
assert second.stat().st_mtime_ns == second_version
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ["replacement", "catalog"])
|
||||
async def test_failed_persistence_keeps_recoverable_checkpoint(case, failure):
|
||||
"""Neither atomic replacement failure nor repeated catalog failure can discard progress."""
|
||||
async with case.open() as app:
|
||||
catalog = app.context.components["file_catalog"]["dream"]
|
||||
original_write, writes = CheckpointStore._write, 0
|
||||
|
||||
def replace(store):
|
||||
nonlocal writes
|
||||
writes += 1
|
||||
with (
|
||||
(
|
||||
patch("reme.steps.evolve.dream.checkpoint.os.replace", side_effect=OSError("disk full"))
|
||||
if writes == 3
|
||||
else nullcontext()
|
||||
),
|
||||
):
|
||||
return original_write(store)
|
||||
|
||||
with (
|
||||
(
|
||||
patch.object(CheckpointStore, "_write", replace)
|
||||
if failure == "replacement"
|
||||
else patch.object(catalog, "dump", side_effect=OSError("disk full"))
|
||||
),
|
||||
):
|
||||
response = await case.run(app)
|
||||
assert not response.success and "disk full" in response.answer
|
||||
assert case.checkpoints
|
||||
if failure == "replacement":
|
||||
assert case.saved.active == 0 and case.saved.completed == []
|
||||
assert await catalog.get_nodes([SOURCE]) == []
|
||||
assert len(list(case.checkpoints[0].parent.iterdir())) == 1
|
||||
else:
|
||||
case.calls.clear()
|
||||
assert not (await case.run(app)).success
|
||||
assert not case.calls and case.checkpoints
|
||||
case.calls.clear()
|
||||
response = await case.run(app)
|
||||
assert response.success, response.answer
|
||||
assert "extract" not in case.calls and not case.checkpoints
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_cannot_escape_workspace(case):
|
||||
"""A metadata symlink must not allow checkpoint writes outside the workspace."""
|
||||
outside = case.workspace / "outside"
|
||||
outside.mkdir()
|
||||
nested = DreamCase(case.workspace / "nested")
|
||||
(nested.workspace / "metadata").symlink_to(outside, target_is_directory=True)
|
||||
async with nested.open() as app:
|
||||
response = await nested.run(app)
|
||||
assert not response.success and "checkpoint escapes workspace" in response.answer
|
||||
assert not nested.calls and not (outside / "dream").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_between_scan_and_extraction_keeps_source_eligible(case):
|
||||
"""Do not attach a newer extraction to the scan's older source watermark."""
|
||||
|
||||
def edit_then_hash(workspace, existing):
|
||||
previous = case.source.stat().st_mtime
|
||||
case.source.write_text("Modified after scanning")
|
||||
os.utime(case.source, (previous + 2, previous + 2))
|
||||
return source_versions(workspace, existing)
|
||||
|
||||
async with case.open() as app:
|
||||
with patch("reme.steps.evolve.dream.extract.source_versions", edit_then_hash):
|
||||
response = await case.run(app)
|
||||
assert not response.success and "changed after scanning" in response.answer
|
||||
assert not case.calls
|
||||
assert await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE]) == []
|
||||
assert (await case.run(app)).success
|
||||
assert case.calls == ["extract", "first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("overlap", ["integration", "finish"])
|
||||
async def test_overlapping_jobs_do_not_replay_or_regress_progress(case, overlap):
|
||||
"""Queued integrations refresh progress; stale finalizers cannot overwrite newer watermarks."""
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
boundary = DreamIntegrateStep if overlap == "integration" else DreamFinishStep
|
||||
method = "_execute_locked" if overlap == "integration" else "execute"
|
||||
original, invocations = getattr(boundary, method), 0
|
||||
|
||||
async def delay(step, *args):
|
||||
nonlocal invocations
|
||||
invocations += 1
|
||||
if invocations == 1:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return await original(step, *args)
|
||||
|
||||
async with case.open() as app:
|
||||
tasks = []
|
||||
try:
|
||||
with patch.object(boundary, method, delay):
|
||||
tasks.append(asyncio.create_task(case.run(app)))
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
if overlap == "finish":
|
||||
previous = case.source.stat().st_mtime
|
||||
case.source.write_text("New source version")
|
||||
os.utime(case.source, (previous + 2, previous + 2))
|
||||
assert (await case.run(app)).success
|
||||
else:
|
||||
restored = asyncio.Event()
|
||||
restore = CheckpointStore.restore
|
||||
|
||||
def signal(store, sources):
|
||||
result = restore(store, sources)
|
||||
restored.set()
|
||||
return result
|
||||
|
||||
with patch.object(CheckpointStore, "restore", signal):
|
||||
tasks.append(asyncio.create_task(case.run(app)))
|
||||
await asyncio.wait_for(restored.wait(), timeout=5)
|
||||
release.set()
|
||||
responses = await asyncio.gather(*tasks)
|
||||
if overlap == "finish":
|
||||
assert not responses[0].success and "another invocation" in responses[0].answer
|
||||
nodes = await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE])
|
||||
assert nodes[0].st_mtime == case.source.stat().st_mtime
|
||||
case.calls.clear()
|
||||
assert (await case.run(app)).success and not case.calls
|
||||
else:
|
||||
assert any(r.success for r in responses)
|
||||
assert all(r.success or "another invocation" in r.answer for r in responses)
|
||||
assert case.calls == ["extract", "first", "second"]
|
||||
finally:
|
||||
release.set()
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
Loading…
Add table
Reference in a new issue