Skip to content

perf(datasources): parse each YAML text once and reuse the package probe (ENG-3362) - #558

Merged
lucas-koontz merged 1 commit into
stagingfrom
perf/eng-3362-cache-turn-setup
Oct 10, 2026
Merged

lucas-koontz merged 1 commit into
stagingfrom
perf/eng-3362-cache-turn-setup

Conversation

@lucas-koontz

Copy link
Copy Markdown
Contributor

User story

As a person asking Cowork questions while many colleagues ask at the same time,
I want the agent to skip setup work it already did on an earlier turn,
so that one cowork-server process serves more answers at once before my first text slows down.

Why

A py-spy profile of cowork-server's single process, taken while 25 testers asked at once, recorded only the thread that held the GIL. More than half of those samples were anton's per-turn setup, redoing work whose inputs hadn't changed:

  • About 48% was PyYAML's pure-Python parser. DatasourceRegistry() re-parsed every block of the built-in datasources.md, and the user's, each time it was built. A turn builds it at least four times: once from restore_namespaced_env at request entry, twice while the harness builds the session and opens its turn scope, and again each time collect_datasource_catalog builds the system prompt.
  • About 10% was probe_packages. ScratchpadManager.__init__ read every installed distribution's METADATA on every turn.
  • About 1.4% was skill frontmatter, which list_summaries parses on every turn.

What changes

  • One parse per YAML text. The new anton/core/utils/yaml_cache.py provides safe_load_cached. It parses each distinct YAML text once and gives every caller a deep copy, so no registry or skill can change what another one reads. DatasourceRegistry._parse_file and parse_skill_dir both use it. Both still read their files on every call, so an edited or new datasources.md or SKILL.md is picked up on its next parse.
  • Bounds on the cache. It keeps only texts up to 8 KiB whose parse stays under 64 KiB, and at most 512 of them.
  • Bounds on the loader. It refuses a document whose aliases add more than 65,536 nodes and characters, or refer to a node that holds them. It refuses before it builds any Python object, and raises AliasGrowthError or AliasLoopError (both yaml.YAMLError).
  • Skill frontmatter. parse_skill_dir reads at most 64 KiB of frontmatter. It finds the closing --- without splitting the whole file into lines. It skips a SKILL.md whose YAML can't be loaded or read as text, as _parse_file skips such a datasources.md block, instead of raising. A frontmatter name that isn't text is read with str(). Its warnings keep staging's format: the error class and position, never the YAML.
  • The package probe. probe_packages keeps its last answer, together with sys.path and each entry's mtime, and probes again when either changes. An installer adds or removes a whole *.dist-info directory, which moves its entry's mtime, and importlib.metadata keys its own listing on that same mtime. The probe only ever covered the host's sys.path; scratchpad venvs install elsewhere.
flowchart LR
    T[a turn] --> R1["DatasourceRegistry(), 4+ times"]
    T --> P[probe_packages]
    T --> S[list_summaries]
    R1 --> C{"safe_load_cached: text seen before?"}
    S --> C
    C -- yes --> D[deep copy of the kept parse]
    C -- no --> L[bounded loader parses once and keeps it]
    P --> M{"sys.path and mtimes unchanged?"}
    M -- yes --> K[last answer, copied]
    M -- no --> Q[probe again]
Loading

Measured

CPU time in anton's benchmark, the median of 40 runs, with cowork-server's 110 installed distributions:

Call Before After
DatasourceRegistry() 22.09 ms 0.66 ms
restore_namespaced_env, as cowork-server calls it 22.67 ms 0.98 ms
collect_datasource_catalog 22.01 ms 0.84 ms
probe_packages() 16.28 ms 0.013 ms
list_summaries() 2.26 ms 0.40 ms
A turn's setup, all of the above about 110 ms about 4.5 ms

The first call in each process still pays the full parse, about 22 ms.

Acceptance criteria

  • A registry built while the files are unchanged does no YAML parsing. An edited or new user datasources.md is seen on the next construction and on reload(), even when the edit keeps the file's size and mtime.
  • Changing one registry's engines, fields, name_from or auth methods doesn't reach another registry.
  • A malformed block or frontmatter warns on every parse, as before, and the warning names only the error class and position.
  • A package installed or removed on sys.path is seen on the next probe.
  • Every bundled skill, every test fixture and 2,958 generated edge cases parse identically to staging, apart from the intended refusals and str() names.
  • Under a 1 GB memory limit, no frontmatter under the cap takes the parse past about 31 MB or 0.5 s of CPU.
  • CI passes.
  • A staging release candidate carries this, and the cowork-server ramp on a 16-vCPU VM answers more turns a minute than 9 October's 166.

How to test

  1. pytest tests/ on this branch.
  2. Run cowork-server with this anton and ask several questions. The answers, the datasource catalog in the system prompt, and the Skills page should all match staging.
  3. Edit ~/.anton/datasources.md, then ask again. The edit shows up on the next turn.

Notes for the reviewer

  • The cache hands out deep copies, not shared objects. name_from can be a YAML list, and _parse_file sets required = False on custom engines' fields, so sharing could let one caller change another's data. A copy costs about 0.27 ms against the 22 ms parse it replaces.
  • It keys on the YAML text, not on the file's stat. A stat key can serve a stale parse after a same-size rewrite within one mtime tick, and a test pins that case.
  • I measured yaml.CSafeLoader (2.3 ms instead of 21 ms) but left it out. libyaml can accept or reject hand-written YAML differently from SafeLoader, and with the cache it would save about 19 ms once per process.
  • Two log lines on this path now name only the error class, matching staging's redaction: an undecodable SKILL.md, and a value that can't be read as text.

Verified locally

Check Result
pytest tests/ --ignore=tests/e2e on this commit 4,894 passed, 62 skipped
pytest tests/e2e (stub) 61 passed
ruff check on the changed files Clean
Parse parity against staging: 3,132 SKILL.md inputs, each in its own container Identical, apart from the intended refusals and str() names
Bounds under a 1 GB container limit, pure-Python PyYAML, about 27 shapes of YAML At most about 31 MB and 0.5 s of CPU each

Ships with

Refs: ENG-3362

…obe (ENG-3362)

A turn built DatasourceRegistry() at least four times, and each one ran
PyYAML's pure-Python parser over every block of the built-in
datasources.md and the user's. ScratchpadManager probed every installed
distribution's METADATA on each turn too. With 25 turns at once in one
cowork-server process, those took about 48% and 10% of the CPU samples
that held the GIL.

- safe_load_cached parses each distinct YAML text once and hands every
  caller a deep copy, so no registry or skill can change what another
  one reads. DatasourceRegistry and parse_skill_dir both use it; an
  edited or new datasources.md or SKILL.md is read again on its next
  parse.
- The cache keeps only texts up to 8 KiB whose parse stays under 64 KiB,
  and at most 512 of them. Its loader refuses a document whose aliases
  add more than 65,536 nodes and characters, or refer to a node that
  holds them, before it builds any Python object.
- parse_skill_dir reads at most 64 KiB of frontmatter, finds the closing
  "---" without splitting the file into lines, and skips a SKILL.md, as
  _parse_file skips a datasources.md block, when its YAML can't be
  loaded or read as text. A frontmatter name that isn't text is read
  with str(). Its warnings name only the error class and position.
- probe_packages keeps its last answer with sys.path and each entry's
  mtime, and probes again when either changes.

In the benchmark, DatasourceRegistry() drops from 22 to 0.7 ms of CPU and
probe_packages from 16 to 0.013 ms; a turn's setup from about 110 to 4.5 ms.

Lucas Koontz, ENG-3362: Cut the CPU cowork-server spends on each answer.

Refs: ENG-3362
@lucas-koontz
lucas-koontz merged commit 2e3a22f into staging Oct 10, 2026
9 checks passed
@lucas-koontz
lucas-koontz deleted the perf/eng-3362-cache-turn-setup branch October 10, 2026 22:55
@github-actions github-actions Bot locked and limited conversation to collaborators Oct 10, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant