feat(sbom): add bounded concurrency for LLM enrichment (#197) - #246
feat(sbom): add bounded concurrency for LLM enrichment (#197)#246nikhilpatidar wants to merge 5 commits into
Conversation
SBOM enrichment calls the LLM sequentially inside _llm_enrich — most prominently in verify_uncertain_nodes which loops over per-node verification calls one at a time. On larger repos this underutilises available API throughput and lengthens scan wall-clock time. This change adds a configurable llm_concurrency (default 5) that bounds parallel LLM calls during enrichment: * nuguard/sbom/config.py — adds AiSbomConfig.llm_concurrency (1-64, default 5). Reads NUGUARD_LLM_CONCURRENCY / AISBOM_LLM_CONCURRENCY. * nuguard/sbom/core/verification.py — verify_uncertain_nodes now uses asyncio.Semaphore + asyncio.gather to run candidate verifications concurrently. Budget enforcement is preserved (over-budget calls are recorded as skipped under a shared lock). A concurrency < 1 value is coerced to 1 for backward compatibility. * nuguard/sbom/extractor/core.py — passes config.llm_concurrency to verify_uncertain_nodes. * nuguard/sbom/cli.py — adds --llm-concurrency flag. * nuguard/config.py — exposes sbom_generation.llm_concurrency via the nuguard.yaml flat key sbom_llm_concurrency. * nuguard/cli/commands/sbom.py — wires cfg.sbom_llm_concurrency into the AiSbomConfig built for each scan. * nuguard.yaml.example — documents sbom_generation.llm_concurrency. * tests/sbom/test_verification_concurrency.py — 7 tests covering: default sequential behaviour, semaphore cap, wall-clock speedup at concurrency > 1, correctness (same result set under parallelism), cost-budget enforcement, per-node exception isolation, and concurrency < 1 coercion. No behaviour change for the v1 path: callers that do not set llm_concurrency continue to run sequentially (default = 1 in the verify_uncertain_nodes signature). Tests: 256 passed across tests/sbom/ and tests/cli/.
KanishkThamman
left a comment
KanishkThamman
left a comment
There was a problem hiding this comment.
Concurrency/semaphore mechanism and config wiring are sound, but holding off approval for one correctness regression flagged inline — recommend fixing before merge rather than as a follow-up, since it weakens the LLM-failure fallback behavior this code path is supposed to guarantee.
| cost_used += actual_cost | ||
| async with shared_lock: | ||
| shared_cost[0] += actual_cost | ||
| result = parse_verification_response(response, node) |
There was a problem hiding this comment.
This now runs outside the try/except above. Previously the whole per-node call (including response parsing) was isolated; now an exception here (e.g. valid JSON that isn't a dict) propagates through asyncio.gather() with no return_exceptions=True and aborts the entire batch instead of just this node. Wrap this call too, or gather with return_exceptions=True.
There was a problem hiding this comment.
Fixed in d2a6c0f. Wrapped parse_verification_response in the same per-node try/except as the LLM call: any exception logs a warning, bumps stats.skipped_count under the shared lock, and returns None. The batch continues normally for sibling candidates.
The root cause is in parse_verification_response itself — it only catches (json.JSONDecodeError, KeyError, TypeError, ValueError), but valid JSON that isn't a dict (a list, number, or string) raises AttributeError on data.get(...). So 'valid JSON that isn't a dict' was the exact failure mode.
Added a regression test (test_verify_uncertain_nodes_skips_unparseable_response_under_concurrency) that returns '[1, 2, 3]' on the first call. Without the wrapper, asyncio.gather() aborts the whole batch with AttributeError; with the wrapper, 4 siblings verify normally and 1 is recorded as skipped.
The per-node try/except in verify_uncertain_nodes wrapped the llm_call_fn call but left parse_verification_response outside it. parse_verification_response only catches (json.JSONDecodeError, KeyError, TypeError, ValueError) — valid JSON that isn't a dict (e.g. a list, number, or string) raises AttributeError on `data.get(...)`. That AttributeError escaped the per-node scope, propagated through asyncio.gather() (which is called without `return_exceptions=True`), and aborted the *entire* verification batch instead of just the offending node. Fix: - Wrap the parse_verification_response call in the same per-node try/except so any unexpected exception logs a warning, bumps stats.skipped_count under the shared lock, and returns None. All other candidate calls continue to completion. - Add a regression test (5 candidates, first call returns valid JSON list '[1, 2, 3]') that fails without the wrapper — the AttributeError would escape asyncio.gather() and abort the test instead of leaving 4 successful verifications + 1 skipped.
KanishkThamman
left a comment
KanishkThamman
left a comment
There was a problem hiding this comment.
@nikhilpatidar Could you take a look at the lint checks and address the issues identified?
The original 'sbom_overrides: dict[str, object]' + '**sbom_overrides' pattern produced 6 mypy 'incompatible type' errors on the AiSbomConfig constructor — every field is a concrete type (int / bool / str / list[str] / set[str] / str | None) and dict[str, object] widens to object, which mypy refuses to silently unpack. Pass each field positionally instead, and use _default_llm_concurrency() when the config value is unset so the existing NUGUARD_LLM_CONCURRENCY / AISBOM_LLM_CONCURRENCY env-var precedence is preserved.
2 participants
PR Type
What
Adds a configurable
llm_concurrencysetting (default 5) that boundsparallel LLM calls during SBOM enrichment, with the biggest win coming
from parallelising the per-node LLM calls in
verify_uncertain_nodeswhich previously ran strictly sequentially.
Why
Issue #197 — SBOM extraction performs LLM enrichment calls sequentially
(gap-fill, uncertain-node verification, node-description enrichment).
On larger repos this underutilises available API throughput and
lengthens scan wall-clock time.
Root cause
AiSbomExtractor._llm_enrichruns 8 sequential LLM phases (gap-fill,verification, MCP annotation, use-case summary, IaC summary,
description enrichment, relationship graph, descriptive names). Within
those phases,
verify_uncertain_nodesinnuguard/sbom/core/verification.pylooped over per-node verification calls one at a time inside an
async for,so even though the function was
async, no calls actually overlapped.Changes
nuguard/sbom/config.py— addsAiSbomConfig.llm_concurrency(
int, default 5, range 1–64). ReadsNUGUARD_LLM_CONCURRENCY/AISBOM_LLM_CONCURRENCYfrom the environment.nuguard/sbom/core/verification.py—verify_uncertain_nodesnowuses
asyncio.Semaphore+asyncio.gatherto run candidateverifications concurrently. Per-call cost-budget enforcement is
preserved (over-budget calls are recorded as skipped under a shared
asyncio.Lock).concurrency < 1is coerced to1for backwardcompatibility.
nuguard/sbom/extractor/core.py— passesconfig.llm_concurrencyto
verify_uncertain_nodes.nuguard/sbom/cli.py— adds the--llm-concurrencyflag.nuguard/config.py— flattenssbom_generation.llm_concurrencyintothe
sbom_llm_concurrencypydantic field and adds the correspondingNuGuardConfig.sbom_llm_concurrencyentry.nuguard/cli/commands/sbom.py— passescfg.sbom_llm_concurrencyinto the
AiSbomConfigbuilt for each scan (only when explicitly setin
nuguard.yaml, so theAiSbomConfigdefault still appliesotherwise).
nuguard.yaml.example— documentssbom_generation.llm_concurrency.Tests
tests/sbom/test_verification_concurrency.py— 7 new tests covering:concurrency=1concurrency)concurrency > 1skipped and
stats.budget_exceededis set)the batch)
concurrency < 1to1All 256 existing tests in
tests/sbom/andtests/cli/still pass.ruff checkclean across all touched files.mypyclean on the changed modules.Backward compatibility
verify_uncertain_nodessignature gained a newconcurrencykwarg with default
1, so any external caller that does not pass itcontinues to run sequentially — identical observable behaviour to the
previous implementation.
AiSbomConfig.llm_concurrencyfield defaults to5; the existingv1 path is unaffected because no callers in this PR change the default
except by explicitly opting in via the new flag or nuguard.yaml key.
Closes #197