Repository navigation
perf(datasources): parse each YAML text once and reuse the package probe (ENG-3362) - #558
Merged
Merged
Conversation
…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
7 of 8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
DatasourceRegistry()re-parsed every block of the built-indatasources.md, and the user's, each time it was built. A turn builds it at least four times: once fromrestore_namespaced_envat request entry, twice while the harness builds the session and opens its turn scope, and again each timecollect_datasource_catalogbuilds the system prompt.probe_packages.ScratchpadManager.__init__read every installed distribution'sMETADATAon every turn.list_summariesparses on every turn.What changes
anton/core/utils/yaml_cache.pyprovidessafe_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_fileandparse_skill_dirboth use it. Both still read their files on every call, so an edited or newdatasources.mdorSKILL.mdis picked up on its next parse.AliasGrowthErrororAliasLoopError(bothyaml.YAMLError).parse_skill_dirreads at most 64 KiB of frontmatter. It finds the closing---without splitting the whole file into lines. It skips aSKILL.mdwhose YAML can't be loaded or read as text, as_parse_fileskips such adatasources.mdblock, instead of raising. A frontmatternamethat isn't text is read withstr(). Its warnings keepstaging's format: the error class and position, never the YAML.probe_packageskeeps its last answer, together withsys.pathand each entry's mtime, and probes again when either changes. An installer adds or removes a whole*.dist-infodirectory, which moves its entry's mtime, andimportlib.metadatakeys its own listing on that same mtime. The probe only ever covered the host'ssys.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]Measured
CPU time in
anton's benchmark, the median of 40 runs, with cowork-server's 110 installed distributions:DatasourceRegistry()restore_namespaced_env, as cowork-server calls itcollect_datasource_catalogprobe_packages()list_summaries()The first call in each process still pays the full parse, about 22 ms.
Acceptance criteria
datasources.mdis seen on the next construction and onreload(), even when the edit keeps the file's size and mtime.name_fromor auth methods doesn't reach another registry.sys.pathis seen on the next probe.staging, apart from the intended refusals andstr()names.How to test
pytest tests/on this branch.staging.~/.anton/datasources.md, then ask again. The edit shows up on the next turn.Notes for the reviewer
name_fromcan be a YAML list, and_parse_filesetsrequired = Falseon 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.yaml.CSafeLoader(2.3 ms instead of 21 ms) but left it out. libyaml can accept or reject hand-written YAML differently fromSafeLoader, and with the cache it would save about 19 ms once per process.staging's redaction: an undecodableSKILL.md, and a value that can't be read as text.Verified locally
pytest tests/ --ignore=tests/e2eon this commitpytest tests/e2e(stub)ruff checkon the changed filesstaging: 3,132SKILL.mdinputs, each in its own containerstr()namesShips with
Refs: ENG-3362