From 19ad87e54bfa1d8730a8bd9f01def424dde5f3a2 Mon Sep 17 00:00:00 2001 From: huangkuanlin Date: Thu, 9 Jul 2026 19:23:09 -0400 Subject: [PATCH] perf(workspace): bound + cache directory scans so large workspaces can't wedge boot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Large AD-Workbench workspaces (many files on a slow Azure Files / blobfuse mount) froze agent bootstrap and crash-looped the pod ("Agent could not start" / "could not reach server"). Root cause: on_chat_start walked the user-data mount ~4x via unbounded os.walk (the max_items caps only limited the preview, not the traversal), and _build_user_data_sidebar_elements ran synchronously on the asyncio event loop — starving the /healthz liveness probe until Kubernetes restarted the pod. Fine with few files; fatal at scale. Introduce biomni/fs_scan.py: one bounded, cached, thread-safe directory scanner, and route every workspace walk through it. - Bounded by BOTH a file-count cap (BIOMNI_WORKSPACE_MAX_FILES, default 10000) and a wall-clock deadline (BIOMNI_WORKSPACE_SCAN_TIMEOUT_S, default 15s) — a scan can never hang regardless of size or mount latency. - TTL-cached (BIOMNI_WORKSPACE_SCAN_TTL_S, default 60s) behind a per-key lock so the several inventory/sidebar builders of one chat start traverse the tree once, not ~4x. - Sidebar build moved off the event loop into run_in_executor and guarded, so a slow/failed scan degrades to no panel instead of wedging the server. - chainlit inventory/sidebar/overview/welcome and agent _get_data_root_items / _get_user_data_resources / _get_data_lake_items all go through the scanner; no raw unbounded os.walk on a user/network mount remains on the boot path. - Honest UX: bounded scans render counts as a floor ("N+", "listing bounded"), a timed-out scan reads "listing unavailable" (not "no files found"), and the agent is told to enumerate on demand. - runs/ excluded from scans (matches biomni.artifact); non-positive env values fall back to defaults (the guard cannot be disabled). Tests: tests/test_fs_scan.py (cap/deadline bounds, cache identity, depth + exclusion pruning, env fallbacks). Verified end-to-end in docker: a 50k-file workspace boots with /healthz=200, 0 restarts, scan capped at 10k in ~28ms. Co-Authored-By: Claude Opus 4.8 (1M context) --- biomni/agent/a1.py | 116 +++++----------- biomni/fs_scan.py | 304 ++++++++++++++++++++++++++++++++++++++++++ chainlit_app.py | 303 +++++++++++++++++++++-------------------- tests/test_fs_scan.py | 198 +++++++++++++++++++++++++++ 4 files changed, 687 insertions(+), 234 deletions(-) create mode 100644 biomni/fs_scan.py create mode 100644 tests/test_fs_scan.py diff --git a/biomni/agent/a1.py b/biomni/agent/a1.py index f4562af68..6856b423e 100644 --- a/biomni/agent/a1.py +++ b/biomni/agent/a1.py @@ -30,6 +30,7 @@ summarize_topic_for_run_id as _shared_summarize_topic_for_run_id, ) from biomni.config import default_config +from biomni.fs_scan import scan_directory from biomni.know_how import KnowHowLoader from biomni.llm import SourceType, get_llm, resolve_source from biomni.llm_resilience import LLMUsageTracker, prepare_messages_for_cache @@ -597,62 +598,35 @@ def _prioritize_local_tools(self, tools: list[Any]) -> list[Any]: return local_tools + network_tools def _get_data_lake_items(self) -> list[str]: - """Return all files currently present in the data lake (recursive).""" - items: list[str] = [] + """Return all files currently present in the data lake (recursive). + Backed by the bounded/cached scanner: on AKS ``data_lake_dir`` may be a + network-backed mount, and this runs on the boot/system-prompt path + (``_get_data_lake_resources`` and AD1's local-data-priority build), so an + unbounded walk here would re-introduce the boot stall via a different + entry point. + """ if not os.path.isdir(self.data_lake_dir): - return items - - for root, _dirs, files in os.walk(self.data_lake_dir): - for file_name in files: - full_path = os.path.join(root, file_name) - relative_path = os.path.relpath(full_path, self.data_lake_dir).replace(os.sep, "/") - if relative_path == "_custom_data_index.json": - continue - items.append(relative_path) + return [] - return sorted(set(items)) + result = scan_directory(self.data_lake_dir) + return sorted(f for f in result.files if f != "_custom_data_index.json") def _get_data_root_items(self, max_depth: int = 3) -> list[str]: """Return files and directories under the configured BIOMNI_DATA_PATH root. This captures datasets placed directly under the data root (e.g. /mnt/dataset1/files) that are outside the standard data_lake sub-tree. + + Backed by the bounded/cached scanner so it cannot hang while building the + system prompt on a large or network-backed workspace. """ root_dir = getattr(self, "data_root_dir", None) if not root_dir or not os.path.isdir(root_dir): return [] - excluded = { - ".git", - "__pycache__", - ".venv", - "venv", - "env", - ".chainlit", - "node_modules", - "site-packages", - } - items: list[str] = [] - - for root, dirs, files in os.walk(root_dir): - # Respect max depth - depth = root.replace(root_dir, "").count(os.sep) - if depth >= max_depth: - dirs[:] = [] - continue - # Prune hidden and excluded dirs - dirs[:] = [d for d in dirs if not d.startswith(".") and d not in excluded] - for file_name in files: - if file_name.startswith("."): - continue - full_path = os.path.join(root, file_name) - rel = os.path.relpath(full_path, root_dir).replace(os.sep, "/") - if rel == "_custom_data_index.json": - continue - items.append(rel) - - return sorted(set(items)) + result = scan_directory(root_dir, max_depth=max_depth) + return sorted(f for f in result.files if f != "_custom_data_index.json") def _get_user_data_resources(self, max_depth: int = 6, max_items: int = 500) -> list[dict[str, str]]: """Build user-data resources from data_root_dir for retrieval indexing. @@ -661,51 +635,31 @@ def _get_user_data_resources(self, max_depth: int = 6, max_items: int = 500) -> ``data_lake/`` directory, this method indexes files under the user-mounted data root (e.g. /app/user-data or /mnt) so the tool retriever can surface VM-mounted datasets when relevant. + + Bounded by the shared scanner (file cap + wall-clock deadline); the + built-in data-lake subtree is pruned so its files are not double-listed. """ root_dir = getattr(self, "data_root_dir", None) if not root_dir or not os.path.isdir(root_dir): return [] - # Avoid double-counting files already under data_lake_dir - data_lake_abs = os.path.abspath(self.data_lake_dir) if hasattr(self, "data_lake_dir") else "" + # Avoid double-counting files already under data_lake_dir (which may be + # nested beneath the user root). + prune = [self.data_lake_dir] if getattr(self, "data_lake_dir", "") else None - excluded = { - ".git", - "__pycache__", - ".venv", - "venv", - "env", - ".chainlit", - "node_modules", - "site-packages", - } - resources: list[dict[str, str]] = [] - - for root, dirs, files in os.walk(root_dir): - # Skip anything already covered by the built-in data lake - if data_lake_abs and os.path.abspath(root).startswith(data_lake_abs): - dirs[:] = [] - continue - depth = root.replace(root_dir, "").count(os.sep) - if depth >= max_depth: - dirs[:] = [] - continue - dirs[:] = [d for d in dirs if not d.startswith(".") and d not in excluded] - for file_name in files: - if file_name.startswith("."): - continue - full_path = os.path.join(root, file_name) - rel = os.path.relpath(full_path, root_dir).replace(os.sep, "/") - resources.append( - { - "name": f"user-data:{rel}", - "description": f"User dataset file at {root_dir}/{rel}", - } - ) - if len(resources) >= max_items: - return resources - - return resources + result = scan_directory( + root_dir, + max_depth=max_depth, + prune_subtrees=prune, + max_files_override=max_items, + ) + return [ + { + "name": f"user-data:{rel}", + "description": f"User dataset file at {root_dir}/{rel}", + } + for rel in result.files + ] def _resolve_data_path(self, data_path: str) -> str: """Resolve a data path to an absolute path with data-lake-first semantics. diff --git a/biomni/fs_scan.py b/biomni/fs_scan.py new file mode 100644 index 000000000..26c87b7ea --- /dev/null +++ b/biomni/fs_scan.py @@ -0,0 +1,304 @@ +"""Bounded, cached filesystem scanning for workspace inventories. + +Agent bootstrap and the Chainlit sidebar need a picture of the user's workspace +(file tree, counts, extensions). The naive approach — ``os.walk`` the whole tree +on every chat start — is fine for a handful of files but catastrophic for a +large workspace on a network-backed mount (Azure Files / blobfuse), where every +``readdir``/``stat`` is a round-trip. A multi-minute walk on the asyncio event +loop wedges the whole server: the ``/healthz`` liveness probe cannot be served, +so Kubernetes restarts the pod and the user sees "Agent could not start". + +Every scan here is bounded by BOTH, whichever trips first: + +* a hard **file-count cap** (``BIOMNI_WORKSPACE_MAX_FILES``, default 10000), and +* a wall-clock **deadline** (``BIOMNI_WORKSPACE_SCAN_TIMEOUT_S``, default 15s), + +so a scan can never hang regardless of workspace size or mount latency. Results +are cached per (root, depth, exclusions, cap) for a short TTL +(``BIOMNI_WORKSPACE_SCAN_TTL_S``, default 60s) behind a per-key lock, so the +several inventory/sidebar builders that fire during one chat start traverse the +tree at most once instead of ~4 times. + +When either bound trips, the scan sets ``truncated`` / ``timed_out`` and reports +counts as a *floor* (``file_count`` means "at least this many"). Callers should +render such totals with a trailing ``+`` and a note rather than as exact figures +— see :meth:`ScanResult.count_label`. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path + +logger = logging.getLogger(__name__) + +_DEFAULT_MAX_FILES = 10_000 +_DEFAULT_TIMEOUT_S = 15.0 +_DEFAULT_TTL_S = 60.0 + +# Directories never worth walking for a user-facing data inventory. Kept here so +# every scanner call (Chainlit sidebar + agent system prompt) prunes identically. +EXCLUDED_DIRS = frozenset( + { + ".git", + "__pycache__", + ".venv", + "venv", + "env", + ".chainlit", + "node_modules", + "site-packages", + ".ipynb_checkpoints", + ".cache", + ".mypy_cache", + ".ruff_cache", + ".pytest_cache", + # Biomni's per-run output history (also pruned by + # biomni.artifact.DEFAULT_EXCLUDED_DIRS). It can grow without bound and + # is never user data, so keep the file cap / deadline for real datasets. + "runs", + } +) + + +def _int_env(name: str, default: int) -> int: + raw = os.getenv(name, "").strip() + if not raw: + return default + try: + value = int(raw) + except ValueError: + logger.warning("invalid %s=%r; using default %d", name, raw, default) + return default + # <=0 is rejected on purpose: the bounds exist to keep scans finite, so + # "0"/"unlimited" must not silently disable the guard — fall back to default. + if value <= 0: + logger.warning("%s=%r is <=0; ignoring (guard cannot be disabled), using default %d", name, raw, default) + return default + return value + + +def _float_env(name: str, default: float) -> float: + raw = os.getenv(name, "").strip() + if not raw: + return default + try: + value = float(raw) + except ValueError: + logger.warning("invalid %s=%r; using default %s", name, raw, default) + return default + if value <= 0: + logger.warning("%s=%r is <=0; ignoring (guard cannot be disabled), using default %s", name, raw, default) + return default + return value + + +def max_files() -> int: + """Hard cap on files visited per scan (env ``BIOMNI_WORKSPACE_MAX_FILES``).""" + return _int_env("BIOMNI_WORKSPACE_MAX_FILES", _DEFAULT_MAX_FILES) + + +def scan_timeout_s() -> float: + """Wall-clock budget per scan (env ``BIOMNI_WORKSPACE_SCAN_TIMEOUT_S``).""" + return _float_env("BIOMNI_WORKSPACE_SCAN_TIMEOUT_S", _DEFAULT_TIMEOUT_S) + + +def cache_ttl_s() -> float: + """How long a scan result stays fresh (env ``BIOMNI_WORKSPACE_SCAN_TTL_S``).""" + return _float_env("BIOMNI_WORKSPACE_SCAN_TTL_S", _DEFAULT_TTL_S) + + +@dataclass +class ScanResult: + """Outcome of a bounded directory walk. + + ``file_count`` / ``dir_count`` are exact when neither bound tripped; when + ``truncated`` or ``timed_out`` is set they are a *floor* (the real tree is at + least this big). ``files`` holds relative POSIX paths, itself capped at the + file cap so it is safe to embed in a prompt or render in a panel. + + READ-ONLY: instances are cached and handed to multiple threads by reference + (see :func:`scan_directory`). Never mutate ``files`` / ``top_level_counts`` / + ``extension_counts`` in place — copy first (``list(...)`` / ``dict(...)``). + """ + + root: str + files: list[str] = field(default_factory=list) + file_count: int = 0 + dir_count: int = 0 + truncated: bool = False # hit the file-count cap + timed_out: bool = False # hit the wall-clock deadline + top_level_counts: dict[str, int] = field(default_factory=dict) + extension_counts: dict[str, int] = field(default_factory=dict) + + @property + def bounded(self) -> bool: + """True if the scan stopped early (counts are a lower bound).""" + return self.truncated or self.timed_out + + def count_label(self) -> str: + """Human file total, e.g. ``"328"`` or ``"10000+"`` when cut short.""" + return f"{self.file_count}+" if self.bounded else str(self.file_count) + + +# Deadline is re-checked inside a directory every this-many files so a single +# pathologically large directory cannot blow the wall-clock budget. +_DEADLINE_CHECK_EVERY = 2000 + + +def _walk( + path: str, + max_depth: int, + exclude_top: frozenset[str], + cap: int, + deadline: float, + prune_subtrees: tuple[str, ...], +) -> ScanResult: + """Single bounded ``os.walk`` pass. Stops at ``cap`` files or ``deadline``.""" + result = ScanResult(root=path) + + for root, dirs, files in os.walk(path): + # os.walk is lazy: pruning ``dirs`` / breaking here actually stops + # further readdir/stat traffic, which is what makes the deadline bite. + if time.monotonic() >= deadline: + result.timed_out = True + break + + # Skip entire nested subtrees (e.g. a built-in data lake mounted under + # the user root) — matched on absolute path prefix, so this covers dirs + # deeper than the top level that ``exclude_top`` cannot reach. + if prune_subtrees and any(root == p or root.startswith(p + os.sep) for p in prune_subtrees): + dirs[:] = [] + continue + + rel_root = os.path.relpath(root, path) + depth = 0 if rel_root == "." else rel_root.count(os.sep) + 1 + if depth > max_depth: + dirs[:] = [] + continue + + visible = [d for d in dirs if not d.startswith(".") and d not in EXCLUDED_DIRS] + if depth == 0 and exclude_top: + visible = [d for d in visible if d not in exclude_top] + dirs[:] = sorted(visible) + result.dir_count += len(dirs) + + top = "[root]" if rel_root == "." else rel_root.split(os.sep, 1)[0] + for file_name in sorted(files): + if file_name.startswith("."): + continue + result.file_count += 1 + ext = Path(file_name).suffix.lower() or "[no_ext]" + result.extension_counts[ext] = result.extension_counts.get(ext, 0) + 1 + result.top_level_counts[top] = result.top_level_counts.get(top, 0) + 1 + if len(result.files) < cap: + rel = file_name if rel_root == "." else f"{rel_root}/{file_name}" + result.files.append(rel.replace(os.sep, "/")) + if result.file_count >= cap: + result.truncated = True + return result + if result.file_count % _DEADLINE_CHECK_EVERY == 0 and time.monotonic() >= deadline: + result.timed_out = True + return result + + return result + + +# Cache maps (abspath, depth, excludes, cap) -> (monotonic_stamp, ScanResult). +# Cardinality is tiny (a few roots x a couple of depths), so no eviction is +# needed beyond TTL freshness; entries are simply overwritten on refresh. +_cache: dict[tuple, tuple[float, ScanResult]] = {} +_cache_lock = threading.Lock() # guards _cache and _key_locks +_key_locks: dict[tuple, threading.Lock] = {} + + +def scan_directory( + path: str | None, + *, + max_depth: int = 10, + exclude_top_subdirs: set[str] | frozenset[str] | None = None, + prune_subtrees: set[str] | frozenset[str] | list[str] | None = None, + max_files_override: int | None = None, +) -> ScanResult: + """Return a bounded, cached inventory of ``path``. + + Thread-safe: concurrent callers for the same key block on a per-key lock so + exactly one thread walks a cold tree while the rest reuse its result. This + collapses the several sidebar/inventory walks of one chat start into one. + + Args: + path: Directory to scan. Non-existent / non-dir returns an empty result. + max_depth: Maximum directory depth (top-level files are depth 0). + exclude_top_subdirs: Top-level subdirectory names to skip entirely. + prune_subtrees: Absolute paths whose subtrees are skipped entirely + (e.g. a built-in data lake nested under the user root). + max_files_override: Override the global file cap for this call (e.g. a + retrieval index that only wants the first few hundred files). + """ + if not path or not os.path.isdir(path): + return ScanResult(root=path or "") + + exclude_top = frozenset(exclude_top_subdirs) if exclude_top_subdirs else frozenset() + prune = tuple(sorted(os.path.abspath(p) for p in prune_subtrees)) if prune_subtrees else () + cap = max_files_override if (max_files_override and max_files_override > 0) else max_files() + abspath = os.path.abspath(path) + key = (abspath, int(max_depth), exclude_top, prune, int(cap)) + ttl = cache_ttl_s() + + # Fast path: fresh cache hit under the shared lock. + with _cache_lock: + hit = _cache.get(key) + if hit and (time.monotonic() - hit[0]) < ttl: + return hit[1] + klock = _key_locks.get(key) + if klock is None: + klock = threading.Lock() + _key_locks[key] = klock + + # Only one thread scans a given key at a time; others block then read the + # value it just cached rather than re-walking the same (slow) tree. + with klock: + with _cache_lock: + hit = _cache.get(key) + if hit and (time.monotonic() - hit[0]) < ttl: + return hit[1] + + started = time.monotonic() + deadline = started + scan_timeout_s() + result = _walk(abspath, int(max_depth), exclude_top, cap, deadline, prune) + elapsed = time.monotonic() - started + + if result.bounded: + logger.warning( + "workspace scan of %s bounded after %.1fs (files>=%d, dirs>=%d, truncated=%s, timed_out=%s); " + "listing is partial — set BIOMNI_WORKSPACE_MAX_FILES / BIOMNI_WORKSPACE_SCAN_TIMEOUT_S to tune", + path, + elapsed, + result.file_count, + result.dir_count, + result.truncated, + result.timed_out, + ) + else: + logger.info( + "workspace scan of %s: %d files, %d dirs in %.2fs", + path, + result.file_count, + result.dir_count, + elapsed, + ) + + with _cache_lock: + _cache[key] = (time.monotonic(), result) + return result + + +def clear_cache() -> None: + """Drop all cached scans (call after a known bulk change, or in tests).""" + with _cache_lock: + _cache.clear() + _key_locks.clear() diff --git a/chainlit_app.py b/chainlit_app.py index 1f8e02efb..50550478e 100644 --- a/chainlit_app.py +++ b/chainlit_app.py @@ -71,6 +71,7 @@ import chainlit as cl from biomni.artifact import build_run_id, get_all_files from biomni.config import default_config, resolve_default_llm +from biomni.fs_scan import scan_directory from biomni.health import register_health_routes from biomni.observability import ( RunHeartbeat, @@ -236,38 +237,14 @@ def _list_path_entries_recursive( Returns a (preview_items, total_file_count) tuple. Preview items are relative POSIX-style paths suitable for UI display. - ``exclude_top_subdirs`` names top-level subdirectories to skip entirely. + Thin adapter over :func:`biomni.fs_scan.scan_directory` — the scan is + bounded (file cap + wall-clock deadline) and cached, so this never hangs on + a large or network-backed workspace. ``total_file_count`` is a lower bound + when the underlying scan was truncated; callers that display it should treat + it as approximate (see ``_build_user_data_tree_content``). """ - if not path or not os.path.isdir(path): - return [], 0 - - excluded_dirs = {".git", "__pycache__", ".venv", "venv", "env", "node_modules", "site-packages"} - - preview: list[str] = [] - total_count = 0 - - for root, dirs, files in os.walk(path): - rel_root = os.path.relpath(root, path) - depth = 0 if rel_root == "." else rel_root.count(os.sep) + 1 - if depth > max_depth: - dirs[:] = [] - continue - - visible = [d for d in dirs if not d.startswith(".") and d not in excluded_dirs] - if depth == 0 and exclude_top_subdirs: - visible = [d for d in visible if d not in exclude_top_subdirs] - dirs[:] = sorted(visible) - - for file_name in sorted(files): - if file_name.startswith("."): - continue - total_count += 1 - rel = file_name if rel_root == "." else f"{rel_root}/{file_name}" - rel = rel.replace(os.sep, "/") - if len(preview) < max_items: - preview.append(rel) - - return preview, total_count + result = scan_directory(path, max_depth=max_depth, exclude_top_subdirs=exclude_top_subdirs) + return result.files[:max_items], result.file_count def _collect_path_stats(path: str, max_depth: int = 10, exclude_top_subdirs: set[str] | None = None) -> dict: @@ -275,49 +252,18 @@ def _collect_path_stats(path: str, max_depth: int = 10, exclude_top_subdirs: set ``exclude_top_subdirs`` names top-level subdirectories to skip entirely (useful for counting the datalake root without the biomniAD subfolder). + + Backed by the bounded/cached scanner; ``truncated`` is True when the counts + are a floor (workspace larger than the scan budget). """ - stats = { - "total_files": 0, - "total_dirs": 0, - "top_level_counts": {}, - "extension_counts": {}, + result = scan_directory(path, max_depth=max_depth, exclude_top_subdirs=exclude_top_subdirs) + return { + "total_files": result.file_count, + "total_dirs": result.dir_count, + "top_level_counts": dict(result.top_level_counts), + "extension_counts": dict(result.extension_counts), + "truncated": result.bounded, } - if not path or not os.path.isdir(path): - return stats - - excluded_dirs = {".git", "__pycache__", ".venv", "venv", "env", "node_modules", "site-packages"} - - for root, dirs, files in os.walk(path): - rel_root = os.path.relpath(root, path) - depth = 0 if rel_root == "." else rel_root.count(os.sep) + 1 - if depth > max_depth: - dirs[:] = [] - continue - - visible_dirs = [d for d in dirs if not d.startswith(".") and d not in excluded_dirs] - if depth == 0 and exclude_top_subdirs: - visible_dirs = [d for d in visible_dirs if d not in exclude_top_subdirs] - dirs[:] = sorted(visible_dirs) - stats["total_dirs"] += len(visible_dirs) - - for file_name in files: - if file_name.startswith("."): - continue - - stats["total_files"] += 1 - - ext = Path(file_name).suffix.lower() or "[no_ext]" - ext_counts = stats["extension_counts"] - ext_counts[ext] = ext_counts.get(ext, 0) + 1 - - if rel_root == ".": - top = "[root]" - else: - top = rel_root.split(os.sep, 1)[0] - top_counts = stats["top_level_counts"] - top_counts[top] = top_counts.get(top, 0) + 1 - - return stats def _format_compact_counts(counts: dict[str, int], max_items: int = 8) -> str: @@ -420,20 +366,43 @@ def _render(node: dict[str, object], prefix: str, depth: int) -> bool: def _build_user_data_tree_content(root_path: str, preview_files: int = 300) -> tuple[str, int]: """Build concise per-root content for sidebar readability.""" - preview, total_files = _list_path_entries_recursive(root_path, max_items=preview_files) - if total_files == 0: + result = scan_directory(root_path, max_depth=10) + if result.file_count == 0: + # Distinguish a genuinely empty root from a scan that hit the + # time/size budget before reading any file (slow network mount) — the + # latter must not masquerade as "no files". + if result.bounded: + return ( + "(listing unavailable — the workspace scan hit its time/size budget before any file " + "was read; the folder is likely very large or on a slow mount. Open it directly to browse.)", + 0, + ) return "(no files found)", 0 + preview = result.files[:preview_files] tree_lines = _build_tree_preview_lines(preview, max_lines=70, max_depth=3) lines: list[str] = [ f"Directory structure preview (showing first {len(preview)} files)", "", *tree_lines, ] - if total_files > len(preview): - lines.append(f"... and {total_files - len(preview)} more files") + if result.bounded: + # Workspace exceeded the scan budget: the tree is a partial sample. + lines.append( + f"... workspace is large — listing capped at {result.count_label()} files " + "(not fully indexed; scan bounded for responsiveness)" + ) + elif result.file_count > len(preview): + lines.append(f"... and {result.file_count - len(preview)} more files") - return "\n".join(lines), total_files + return "\n".join(lines), result.file_count + + +def _root_count_label(root_path: str, exclude_top_subdirs: set[str] | None = None) -> str: + """Cached file count for a root, with a trailing ``+`` when the scan was + bounded (a floor, not an exact total). Reuses the cached scan so calling it + for a tab label right after building the tree is free.""" + return scan_directory(root_path, max_depth=10, exclude_top_subdirs=exclude_top_subdirs).count_label() def _build_sidebar_overview_content() -> str: @@ -451,6 +420,21 @@ def _build_sidebar_overview_content() -> str: lines: list[str] = ["At-a-glance overview", ""] grand_files = 0 grand_dirs = 0 + grand_truncated = False + + def _add(label: str, stats: dict) -> None: + # Render one section line and fold its counts into the grand totals. A + # trailing "+" signals the scan was bounded (workspace larger than the + # scan budget), so the number is a floor, not an exact count. + nonlocal grand_files, grand_dirs, grand_truncated + if stats["total_files"] <= 0: + return + grand_files += stats["total_files"] + grand_dirs += stats["total_dirs"] + grand_truncated = grand_truncated or stats.get("truncated", False) + plus = "+" if stats.get("truncated") else "" + lines.append(f"{label}: {stats['total_files']}{plus} files, {stats['total_dirs']}{plus} folders") + lines.append("") # --- 1. AD Workbench Datasets ------------------------------------------- # When HOST_PATH is defined this is an AD-Workbench deployment. The host @@ -465,53 +449,35 @@ def _build_sidebar_overview_content() -> str: elif os.path.isdir(user_data_host_path): ad_path = user_data_host_path if ad_path: - stats = _collect_path_stats(ad_path) - if stats["total_files"] > 0: - grand_files += stats["total_files"] - grand_dirs += stats["total_dirs"] - lines.append(f"AD Workbench Datasets: {stats['total_files']} files, {stats['total_dirs']} folders") - lines.append("") + _add("AD Workbench Datasets", _collect_path_stats(ad_path)) else: # Local / non-AD-Workbench deployment: show regular user data roots. for env_name, root in _resolve_user_data_roots(): - stats = _collect_path_stats(root) - if stats["total_files"] > 0: - grand_files += stats["total_files"] - grand_dirs += stats["total_dirs"] - display = _display_data_root_label(env_name) - lines.append(f"{display}: {stats['total_files']} files, {stats['total_dirs']} folders") - lines.append("") + _add(_display_data_root_label(env_name), _collect_path_stats(root)) # --- 2. Biomni-AD Datalake ---------------------------------------------- builtin_root = _resolve_builtin_data_lake_root() biomni_ad_root = os.path.join(builtin_root, "biomniAD") if os.path.isdir(biomni_ad_root): - stats = _collect_path_stats(biomni_ad_root) - if stats["total_files"] > 0: - grand_files += stats["total_files"] - grand_dirs += stats["total_dirs"] - lines.append(f"Biomni-AD Datalake: {stats['total_files']} files, {stats['total_dirs']} folders") - lines.append("") + _add("Biomni-AD Datalake", _collect_path_stats(biomni_ad_root)) # --- 3. Biomni Datalake (root of data_lake, excluding biomniAD) ---------- if os.path.isdir(builtin_root): - stats = _collect_path_stats(builtin_root, exclude_top_subdirs={"biomniAD"}) - if stats["total_files"] > 0: - grand_files += stats["total_files"] - grand_dirs += stats["total_dirs"] - lines.append(f"Biomni Datalake: {stats['total_files']} files, {stats['total_dirs']} folders") - lines.append("") + _add("Biomni Datalake", _collect_path_stats(builtin_root, exclude_top_subdirs={"biomniAD"})) if grand_files == 0 and grand_dirs == 0: return "No local data roots found." + plus = "+" if grand_truncated else "" lines.extend( [ "Combined totals", - f"Files: {grand_files}", - f"Folders: {grand_dirs}", + f"Files: {grand_files}{plus}", + f"Folders: {grand_dirs}{plus}", ] ) + if grand_truncated: + lines.append("(partial — workspace exceeded the scan budget; counts are a lower bound)") return "\n".join(lines) @@ -532,18 +498,10 @@ def _list_local_data_lake_files(base_path: str, max_items: int = 30) -> list[str if data_lake_dir is None: return [] - items: list[str] = [] - for root, _dirs, files in os.walk(data_lake_dir): - for file_name in files: - if file_name.startswith("."): - continue - full_path = Path(root) / file_name - rel = full_path.relative_to(data_lake_dir).as_posix() - if rel == "_custom_data_index.json": - continue - items.append(rel) - - items = sorted(set(items)) + # Bounded/cached scan so this can't stall process startup even if the + # built-in data lake is pointed at a large volume. + result = scan_directory(str(data_lake_dir)) + items = sorted(f for f in result.files if f != "_custom_data_index.json") return items[:max_items] @@ -557,15 +515,19 @@ def _build_welcome_local_dataset_section() -> str: # User data can come from BIOMNI_USER_DATA_PATH and legacy BIOMNI_DATA_PATH/BIOMNI_PATH. user_roots = _resolve_user_data_roots() user_total_files = 0 + user_count_label = "0" user_preview: list[str] = [] if user_roots: first_root = user_roots[0][1] - user_preview, user_total_files = _list_path_entries_recursive(first_root, max_items=20) + _user_scan = scan_directory(first_root, max_depth=10) + user_preview = _user_scan.files[:20] + user_total_files = _user_scan.file_count + user_count_label = _user_scan.count_label() # One-line summary for the collapsed header summary_parts = [f"{len(data_lake_files)} data lake files"] if user_total_files: - summary_parts.append(f"{user_total_files} user data files") + summary_parts.append(f"{user_count_label} user data files") lines: list[str] = [] lines.append(f"
📊 {' · '.join(summary_parts)} available — click to expand") @@ -592,7 +554,7 @@ def _build_welcome_local_dataset_section() -> str: lines.append(f"- `{_display_data_root_label(env_name)}`: `{root}`") lines.append("") if user_preview: - lines.append(f"Detected files: {user_total_files}") + lines.append(f"Detected files: {user_count_label}") for name in user_preview: lines.append(f"- `{name}`") if user_total_files > len(user_preview): @@ -805,6 +767,32 @@ def _build_full_user_data_inventory() -> str: user_data_host_path = os.getenv("BIOMNI_USER_DATA_HOST_PATH", "").strip() user_data_path = os.getenv("BIOMNI_USER_DATA_PATH", "").strip() + def _section(label: str, root: str) -> None: + result = scan_directory(root, max_depth=10) + if result.file_count == 0: + # A timed-out scan (0 files but bounded) must still tell the agent the + # path exists but wasn't indexed, so it enumerates on demand rather + # than concluding the workspace is empty. + if result.bounded: + sections.append( + f"{label} ({root}) — listing unavailable (scan hit its time/size budget before " + "any file was read; enumerate with os.listdir()/glob on the path)" + ) + return + preview = result.files[:500] + tree_lines = _build_tree_preview_lines(preview, max_lines=300, max_depth=6) + section = f"{label} ({root}) — {result.count_label()} files:\n" + "\n".join(tree_lines) + if result.bounded: + # Tell the agent the listing is partial so it enumerates on demand + # (os.listdir/glob) rather than trusting this as the full inventory. + section += ( + f"\n ... workspace is large — only the first {len(preview)} files are listed " + "(scan bounded for responsiveness; use os.listdir()/glob on the path for the rest)" + ) + elif result.file_count > len(preview): + section += f"\n ... and {result.file_count - len(preview)} more files" + sections.append(section) + # --- User / AD Workbench data --- if user_data_host_path: ad_path = "" @@ -813,21 +801,10 @@ def _build_full_user_data_inventory() -> str: elif os.path.isdir(user_data_host_path): ad_path = user_data_host_path if ad_path: - preview, total = _list_path_entries_recursive(ad_path, max_items=500, max_depth=10) - if total > 0: - tree_lines = _build_tree_preview_lines(preview, max_lines=300, max_depth=6) - sections.append(f"AD Workbench / User Data ({ad_path}) — {total} files:\n" + "\n".join(tree_lines)) - if total > len(preview): - sections[-1] += f"\n ... and {total - len(preview)} more files" + _section("AD Workbench / User Data", ad_path) else: for env_name, root in _resolve_user_data_roots(): - preview, total = _list_path_entries_recursive(root, max_items=500, max_depth=10) - if total > 0: - label = _display_data_root_label(env_name) - tree_lines = _build_tree_preview_lines(preview, max_lines=300, max_depth=6) - sections.append(f"{label} ({root}) — {total} files:\n" + "\n".join(tree_lines)) - if total > len(preview): - sections[-1] += f"\n ... and {total - len(preview)} more files" + _section(_display_data_root_label(env_name), root) return "\n\n".join(sections) if sections else "" @@ -860,13 +837,17 @@ def _build_user_data_sidebar_elements() -> list[cl.Text]: if total_files > 0: content = f"Path: {ad_path}\n\n{tree_content}" elements.append( - cl.Text(name=f"Tree [AD Workbench Datasets] ({total_files})", content=content, display="page") + cl.Text( + name=f"Tree [AD Workbench Datasets] ({_root_count_label(ad_path)})", + content=content, + display="page", + ) ) else: for env_name, root in _resolve_user_data_roots(): tree_content, total_files = _build_user_data_tree_content(root) if total_files > 0: - label = f"Tree [{_display_data_root_label(env_name)}] ({total_files})" + label = f"Tree [{_display_data_root_label(env_name)}] ({_root_count_label(root)})" content = f"Path: {root}\n\n{tree_content}" elements.append(cl.Text(name=label, content=content, display="page")) @@ -879,25 +860,33 @@ def _build_user_data_sidebar_elements() -> list[cl.Text]: tree_content, total_files = _build_user_data_tree_content(biomni_ad_root) if total_files > 0: content = f"Path: {biomni_ad_root}\n\n{tree_content}" - elements.append(cl.Text(name=f"Tree [Biomni-AD Datalake] ({total_files})", content=content, display="page")) + elements.append( + cl.Text( + name=f"Tree [Biomni-AD Datalake] ({_root_count_label(biomni_ad_root)})", + content=content, + display="page", + ) + ) # Biomni Datalake (root, excluding biomniAD) if os.path.isdir(builtin_root): - preview, total_files = _list_path_entries_recursive( - builtin_root, max_items=300, exclude_top_subdirs={"biomniAD"} - ) - if total_files > 0: + lake = scan_directory(builtin_root, max_depth=10, exclude_top_subdirs={"biomniAD"}) + if lake.file_count > 0: + preview = lake.files[:300] tree_lines = _build_tree_preview_lines(preview, max_lines=70, max_depth=3) lake_lines: list[str] = [ f"Directory structure preview (showing first {len(preview)} files)", "", *tree_lines, ] - if total_files > len(preview): - lake_lines.append(f"... and {total_files - len(preview)} more files") - tree_content = "\n".join(lake_lines) - content = f"Path: {builtin_root}\n\n{tree_content}" - elements.append(cl.Text(name=f"Tree [Biomni Datalake] ({total_files})", content=content, display="page")) + if lake.bounded: + lake_lines.append(f"... listing capped at {lake.count_label()} files (scan bounded for responsiveness)") + elif lake.file_count > len(preview): + lake_lines.append(f"... and {lake.file_count - len(preview)} more files") + content = f"Path: {builtin_root}\n\n" + "\n".join(lake_lines) + elements.append( + cl.Text(name=f"Tree [Biomni Datalake] ({lake.count_label()})", content=content, display="page") + ) return elements @@ -1087,15 +1076,23 @@ async def on_chat_start(): await cl.Message(content=f"Failed to initialize {label}: {exc}").send() return - # Render local-data panel in the native sidebar at startup, - # keeping the center welcome/search screen unchanged. - sidebar_content = cl.user_session.get("dataset_listing") or _build_dataset_listing(agent) - sidebar_elements = _build_user_data_sidebar_elements() - if not sidebar_elements: - sidebar_elements = [cl.Text(name="Local Datasets", content=sidebar_content)] - await cl.ElementSidebar.set_title("Local Datasets") - await cl.ElementSidebar.set_elements(sidebar_elements) - cl.user_session.set("dataset_panel_shown", True) + # Render local-data panel in the native sidebar at startup, keeping the + # center welcome/search screen unchanged. The tree build walks the user-data + # mount, so it MUST run in a worker thread — running it inline would block + # the asyncio event loop for the duration of the walk, starving the /healthz + # liveness probe and getting the pod restarted on a large workspace. Guarded + # so a scan/sidebar hiccup degrades to no panel rather than failing the + # already-usable session. + try: + sidebar_content = cl.user_session.get("dataset_listing") or _build_dataset_listing(agent) + sidebar_elements = await run_in_executor(_build_user_data_sidebar_elements) + if not sidebar_elements: + sidebar_elements = [cl.Text(name="Local Datasets", content=sidebar_content)] + await cl.ElementSidebar.set_title("Local Datasets") + await cl.ElementSidebar.set_elements(sidebar_elements) + cl.user_session.set("dataset_panel_shown", True) + except Exception: + logger.exception("Failed to render local-data sidebar (session remains usable)") # --------------------------------------------------------------------------- diff --git a/tests/test_fs_scan.py b/tests/test_fs_scan.py new file mode 100644 index 000000000..96c34ed24 --- /dev/null +++ b/tests/test_fs_scan.py @@ -0,0 +1,198 @@ +"""Tests for biomni.fs_scan — the bounded/cached workspace scanner. + +These lock in the properties that keep agent bootstrap from hanging on a large +or network-backed workspace: a hard file-count cap, a wall-clock deadline, and a +short-TTL cache that collapses the several sidebar/inventory walks of one chat +start into a single traversal. +""" + +from __future__ import annotations + +import pytest +from biomni import fs_scan + + +@pytest.fixture(autouse=True) +def _reset(monkeypatch: pytest.MonkeyPatch) -> None: + """Start every test from a cold cache and default (unset) env knobs.""" + for k in ( + "BIOMNI_WORKSPACE_MAX_FILES", + "BIOMNI_WORKSPACE_SCAN_TIMEOUT_S", + "BIOMNI_WORKSPACE_SCAN_TTL_S", + ): + monkeypatch.delenv(k, raising=False) + fs_scan.clear_cache() + yield + fs_scan.clear_cache() + + +def _make_tree(root, n_files: int, *, subdir: str = "") -> None: + base = root / subdir if subdir else root + base.mkdir(parents=True, exist_ok=True) + for i in range(n_files): + (base / f"f{i}.txt").write_text("x") + + +# --------------------------------------------------------------------------- # +# Basic correctness +# --------------------------------------------------------------------------- # + + +def test_missing_path_returns_empty() -> None: + result = fs_scan.scan_directory("/no/such/path/really") + assert result.file_count == 0 + assert result.files == [] + assert not result.bounded + + +def test_none_path_returns_empty() -> None: + assert fs_scan.scan_directory(None).file_count == 0 + + +def test_small_tree_is_exact(tmp_path) -> None: + _make_tree(tmp_path, 5) + _make_tree(tmp_path, 3, subdir="nested") + result = fs_scan.scan_directory(str(tmp_path)) + assert result.file_count == 8 + assert not result.bounded + assert result.count_label() == "8" + assert len(result.files) == 8 + assert "nested/f0.txt" in result.files + + +def test_hidden_files_and_excluded_dirs_skipped(tmp_path) -> None: + (tmp_path / "keep.txt").write_text("x") + (tmp_path / ".secret").write_text("x") # hidden file + (tmp_path / ".git").mkdir() + (tmp_path / ".git" / "config").write_text("x") # excluded dir + (tmp_path / "__pycache__").mkdir() + (tmp_path / "__pycache__" / "m.pyc").write_text("x") + result = fs_scan.scan_directory(str(tmp_path)) + assert result.file_count == 1 + assert result.files == ["keep.txt"] + + +def test_max_depth_prunes_deep_files(tmp_path) -> None: + (tmp_path / "top.txt").write_text("x") + deep = tmp_path / "a" / "b" / "c" + deep.mkdir(parents=True) + (deep / "deep.txt").write_text("x") + shallow = fs_scan.scan_directory(str(tmp_path), max_depth=1) + assert "top.txt" in shallow.files + assert "a/b/c/deep.txt" not in shallow.files + deep_scan = fs_scan.scan_directory(str(tmp_path), max_depth=10) + assert "a/b/c/deep.txt" in deep_scan.files + + +def test_exclude_top_subdirs(tmp_path) -> None: + _make_tree(tmp_path, 2, subdir="biomniAD") + (tmp_path / "root.txt").write_text("x") + result = fs_scan.scan_directory(str(tmp_path), exclude_top_subdirs={"biomniAD"}) + assert result.file_count == 1 + assert result.files == ["root.txt"] + + +def test_prune_subtrees_skips_nested_tree(tmp_path) -> None: + (tmp_path / "user.txt").write_text("x") + lake = tmp_path / "data" / "data_lake" + _make_tree(lake, 4) + result = fs_scan.scan_directory(str(tmp_path), prune_subtrees=[str(lake)]) + assert result.file_count == 1 + assert result.files == ["user.txt"] + + +# --------------------------------------------------------------------------- # +# Bounding: the core anti-hang guarantees +# --------------------------------------------------------------------------- # + + +def test_file_cap_truncates(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("BIOMNI_WORKSPACE_MAX_FILES", "10") + _make_tree(tmp_path, 50) + result = fs_scan.scan_directory(str(tmp_path)) + assert result.truncated is True + assert result.bounded is True + assert result.file_count == 10 # a floor: stops exactly at the cap + assert len(result.files) == 10 + assert result.count_label() == "10+" + + +def test_max_files_override(tmp_path) -> None: + _make_tree(tmp_path, 50) + result = fs_scan.scan_directory(str(tmp_path), max_files_override=7) + assert result.truncated is True + assert result.file_count == 7 + + +def test_wall_clock_deadline_trips(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + _make_tree(tmp_path, 20) + + # Deterministic clock: first reading is the scan start, every later reading + # is well past the deadline, so the walk aborts on its first check. + calls = {"n": 0} + + def fake_monotonic() -> float: + calls["n"] += 1 + return 1000.0 if calls["n"] == 1 else 1_000_000.0 + + monkeypatch.setattr(fs_scan.time, "monotonic", fake_monotonic) + result = fs_scan.scan_directory(str(tmp_path)) + assert result.timed_out is True + assert result.bounded is True + + +# --------------------------------------------------------------------------- # +# Caching: collapse repeated walks within one chat start +# --------------------------------------------------------------------------- # + + +def test_repeated_scan_is_cached(tmp_path) -> None: + _make_tree(tmp_path, 3) + first = fs_scan.scan_directory(str(tmp_path)) + second = fs_scan.scan_directory(str(tmp_path)) + assert first is second # served from cache, tree walked once + + +def test_clear_cache_forces_rescan(tmp_path) -> None: + _make_tree(tmp_path, 3) + first = fs_scan.scan_directory(str(tmp_path)) + fs_scan.clear_cache() + second = fs_scan.scan_directory(str(tmp_path)) + assert first is not second + + +def test_distinct_params_do_not_collide(tmp_path) -> None: + _make_tree(tmp_path, 3) + a = fs_scan.scan_directory(str(tmp_path), max_depth=1) + b = fs_scan.scan_directory(str(tmp_path), max_depth=5) + assert a is not b # different key => independent cache entries + + +# --------------------------------------------------------------------------- # +# Review-fix regressions +# --------------------------------------------------------------------------- # + + +def test_runs_dir_is_excluded(tmp_path) -> None: + """`runs/` (Biomni per-run output history) must not consume the scan budget.""" + (tmp_path / "dataset.csv").write_text("x") + _make_tree(tmp_path, 20, subdir="runs") + result = fs_scan.scan_directory(str(tmp_path)) + assert result.file_count == 1 + assert result.files == ["dataset.csv"] + + +@pytest.mark.parametrize( + ("var", "getter", "default"), + [ + ("BIOMNI_WORKSPACE_MAX_FILES", fs_scan.max_files, 10_000), + ("BIOMNI_WORKSPACE_SCAN_TIMEOUT_S", fs_scan.scan_timeout_s, 15.0), + ], +) +def test_nonpositive_env_falls_back_to_default(monkeypatch, var, getter, default) -> None: + """0 / negative must NOT disable the guard — it falls back to the default.""" + for bad in ("0", "-5"): + monkeypatch.setenv(var, bad) + assert getter() == default + monkeypatch.setenv(var, "garbage") + assert getter() == default