Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
b92e02d
fix(batch): render untrusted console fields literally
yashrajp22 Oct 5, 2026
bd14216
fix(batch): enforce per-skill process timeouts
yashrajp22 Oct 5, 2026
5d8ee86
fix(batch): avoid signaling reaped worker process groups
yashrajp22 Oct 5, 2026
2c3e9ab
fix(batch): preserve literal terminal names and whitespace
yashrajp22 Oct 5, 2026
0144565
fix(batch): report interrupted scans and separate worker startup
yashrajp22 Oct 5, 2026
545dc4d
test(batch): leave room for worker shutdown in startup check
yashrajp22 Oct 5, 2026
704ffea
Merge reviewed terminal fixes from stacked base
yashrajp22 Oct 5, 2026
f13e98a
Merge reviewed Markdown fixes from stacked base
yashrajp22 Oct 5, 2026
c1a60e7
style: format batch worker regression tests
yashrajp22 Oct 5, 2026
901cb9e
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 5, 2026
2669280
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 5, 2026
2cefeac
fix(batch): keep warning names literal with Rich and plain output
yashrajp22 Oct 6, 2026
a80ef55
fix(batch): stop orphaned pool managers and correct failed scan summa…
yashrajp22 Oct 6, 2026
952ec04
Merge terminal warning follow-up into batch timeout fix
yashrajp22 Oct 6, 2026
bd6258d
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
6b47607
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
de8c857
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
c580fa6
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
a6d62fc
Merge pull request #743 from NVIDIA/yashraj/fix-batch-worker-timeout-…
rng1995 Oct 6, 2026
6ebbfb7
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
77f1cb6
Merge branch 'main' into yashraj/fix-batch-terminal-output-20261005
github-actions[bot] Oct 6, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion contrib/batch_scan/CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ python -m contrib.batch_scan.batch_scan ./tests/fixtures/ -f terminal --workers

```
contrib/batch_scan/
├── batch_scan.py # CLI entry + ThreadPoolExecutor (start here)
├── batch_scan.py # CLI entry + bounded process supervision (start here)
├── runner.py # graph.invoke() wrapper + 7 patches + pool wiring (core)
├── gap_fill.py # GapFillAnalyzer — LLM pass for 8 uncovered rules
├── api_pool.py # ApiKeyPool — multi-key scheduler + 429 backoff
Expand Down
33 changes: 30 additions & 3 deletions contrib/batch_scan/api_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,10 +159,12 @@ def __init__(self, keys: list[ApiKey]) -> None:
self._retry_successes: int = 0
self._total_requests_served: int = 0
self._peak_active_requests: int = 0
self._leases: dict[str, list[ApiKey]] = {}
self._cancelled_owners: set[str] = set()

# -- Public API -----------------------------------------------------------

def acquire(self, timeout: float | None = None) -> ApiKey:
def acquire(self, timeout: float | None = None, *, owner: str | None = None) -> ApiKey:
"""Acquire a slot on the least-loaded available key.

Scheduling priority:
Expand Down Expand Up @@ -193,6 +195,8 @@ def acquire(self, timeout: float | None = None) -> ApiKey:

with self._condition:
while True:
if owner in self._cancelled_owners:
raise RuntimeError("Scan worker has stopped")
now = time.monotonic()

# Step 1: recover rate-limited keys whose backoff has expired
Expand All @@ -208,6 +212,8 @@ def acquire(self, timeout: float | None = None) -> ApiKey:
_now_active = sum(k.active_requests for k in self._keys)
if _now_active > self._peak_active_requests:
self._peak_active_requests = _now_active
if owner is not None:
self._leases.setdefault(owner, []).append(key)
logger.debug(
"Pool: slot on key …%s (%d/%d active)",
key.key[-8:],
Expand Down Expand Up @@ -236,14 +242,16 @@ def acquire(self, timeout: float | None = None) -> ApiKey:
)
self._condition.wait(timeout=wait)

def try_acquire(self) -> ApiKey | None:
def try_acquire(self, *, owner: str | None = None) -> ApiKey | None:
"""Non-blocking acquire — returns a key immediately or ``None``.

Unlike :meth:`acquire`, this never blocks. If a slot is available
right now, return the least-loaded key; otherwise return ``None``.
Useful in async contexts where blocking would stall the event loop.
"""
with self._lock:
if owner in self._cancelled_owners:
raise RuntimeError("Scan worker has stopped")
self._recover_expired_keys(time.monotonic())
available = [k for k in self._keys if k.available]
if not available:
Expand All @@ -255,9 +263,11 @@ def try_acquire(self) -> ApiKey | None:
_now_active = sum(k.active_requests for k in self._keys)
if _now_active > self._peak_active_requests:
self._peak_active_requests = _now_active
if owner is not None:
self._leases.setdefault(owner, []).append(key)
return key

def release(self, key: ApiKey, *, success: bool = True) -> None:
def release(self, key: ApiKey, *, success: bool = True, owner: str | None = None) -> None:
"""Release a slot on *key* back to the pool.

Parameters
Expand All @@ -270,6 +280,15 @@ def release(self, key: ApiKey, *, success: bool = True) -> None:
rate-limited with exponential backoff.
"""
with self._condition:
if owner is not None:
# Manager RPC returns a copy: release the server's original lease.
leases = self._leases.get(owner, [])
for index, leased in enumerate(leases):
if (leased.key, leased.base_url, leased.model) == (key.key, key.base_url, key.model):
key = leases.pop(index)
break
else:
return
key.active_requests = max(0, key.active_requests - 1)

if success:
Expand Down Expand Up @@ -299,6 +318,14 @@ def release(self, key: ApiKey, *, success: bool = True) -> None:

self._condition.notify_all()

def release_owner(self, owner: str) -> None:
"""Reclaim slots and wake waiters after a scan process exits or is killed."""
with self._condition:
self._cancelled_owners.add(owner)
for key in self._leases.pop(owner, []):
key.active_requests = max(0, key.active_requests - 1)
self._condition.notify_all()

def record_retry_success(self) -> None:
"""Increment the retry-success counter for reporting.

Expand Down
Loading