Download tests/test_query_rocm_kb.py from lablab-ai-amd-developer-hackathon/gpu-goblin: direct link, hf CLI and curl.
- Browser
- Download file 13.4 kB
-
https://huggingface.co/spaces/lablab-ai-amd-developer-hackathon/gpu-goblin/resolve/main/tests/test_query_rocm_kb.py
- Command line
-
hf download hf://spaces/lablab-ai-amd-developer-hackathon/gpu-goblin/tests/test_query_rocm_kb.py
-
curl -L -o test_query_rocm_kb.py https://huggingface.co/spaces/lablab-ai-amd-developer-hackathon/gpu-goblin/resolve/main/tests/test_query_rocm_kb.py
13.4 kB
| """Tests for ``agent.tools.query_rocm_kb``. | |
| Coverage: | |
| * The shipped ``kb/rocm_rules.yaml`` is loadable, validates against | |
| :class:`Rule`, and contains every required-by-spec rule id. | |
| * Each rule's ``targets_bucket`` is a valid :class:`WasteBucket` and its | |
| ``category`` is a valid :class:`RuleCategory` (caught for free by | |
| pydantic, but assert here for a clearer failure when the YAML drifts). | |
| * Semantic search returns the bf16 rule first for an "fp16 on MI300X" | |
| query and the flash-attn / sdpa rules first for an attention query. | |
| * Bad inputs (empty symptom, ``top_k <= 0``, ``top_k`` larger than the | |
| rule count) are handled gracefully via ``ToolResult.ok=False`` or | |
| clamped to the rule count. | |
| * Embeddings cache: the cache file is created on first import, and a | |
| second ``_embed_rules`` call with the same YAML bytes hits the cache | |
| without re-encoding. | |
| * The frozen ``Tool`` definition retains its ``name``, ``description``, | |
| and ``input_schema`` shape — the agent registry depends on these. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import get_args | |
| import pytest | |
| import yaml | |
| from agent.schemas import Rule, RuleCategory, ToolResult, WasteBucket | |
| from agent.tools.query_rocm_kb import ( | |
| _KB_YAML, | |
| _RULES, | |
| QUERY_ROCM_KB, | |
| _cache_path, | |
| _embed_rules, | |
| _load_rules, | |
| _query_rocm_kb, | |
| ) | |
| REQUIRED_RULE_IDS = { | |
| "precision.bf16_over_fp16_on_mi300x", | |
| "attention.flash_rocm_over_eager", | |
| "attention.sdpa_over_eager", | |
| "memory.batch_too_small_for_192gb", | |
| "memory.gradient_checkpointing_for_long_seq", | |
| "data.dataloader_workers_zero", | |
| "data.pin_memory_false", | |
| "data.prefetch_factor_default", | |
| "data.persistent_workers_false", | |
| "compile.torch_compile_off", | |
| "env.nccl_min_nchannels", | |
| "env.numa_auto_balancing_disable", | |
| "env.hsa_force_fine_grain_pcie", | |
| "kernels.hipblaslt_hint_logging", | |
| "kernels.miopen_find_mode_2", | |
| "optimizer.bitsandbytes_not_supported_warning", | |
| "collectives.one_process_per_gpu", | |
| "topology.tp_within_xgmi_island", | |
| } | |
| # --------------------------------------------------------------------------- | |
| # YAML KB invariants | |
| # --------------------------------------------------------------------------- | |
| class TestYamlKB: | |
| def test_kb_yaml_exists(self) -> None: | |
| assert _KB_YAML.exists(), f"Missing {_KB_YAML}" | |
| def test_kb_loads_and_validates(self) -> None: | |
| rules, _raw = _load_rules(_KB_YAML) | |
| # Every rule pydantic-validated. | |
| for r in rules: | |
| assert isinstance(r, Rule) | |
| # Spec calls for 20-25 rules; allow a small wiggle room either way. | |
| assert 18 <= len(rules) <= 30, f"Expected 18-30 rules, got {len(rules)}" | |
| def test_required_rule_ids_present(self) -> None: | |
| ids = {r.id for r in _RULES} | |
| missing = REQUIRED_RULE_IDS - ids | |
| assert not missing, f"Required rule ids missing from KB: {sorted(missing)}" | |
| def test_rule_ids_unique(self) -> None: | |
| ids = [r.id for r in _RULES] | |
| dupes = {i for i in ids if ids.count(i) > 1} | |
| assert not dupes, f"Duplicate rule ids: {sorted(dupes)}" | |
| def test_categories_are_valid(self) -> None: | |
| valid = set(get_args(RuleCategory)) | |
| for r in _RULES: | |
| assert r.category in valid, f"{r.id}: invalid category {r.category!r}" | |
| def test_targets_bucket_valid(self) -> None: | |
| valid = set(get_args(WasteBucket)) | |
| for r in _RULES: | |
| assert r.targets_bucket in valid, ( | |
| f"{r.id}: invalid targets_bucket {r.targets_bucket!r}" | |
| ) | |
| def test_recovery_fraction_in_range(self) -> None: | |
| for r in _RULES: | |
| assert 0.0 <= r.expected_recovery_fraction <= 1.0, ( | |
| f"{r.id}: expected_recovery_fraction out of [0,1] " | |
| f"({r.expected_recovery_fraction})" | |
| ) | |
| def test_citations_non_empty(self) -> None: | |
| for r in _RULES: | |
| assert r.citation and r.citation.strip(), f"{r.id}: empty citation" | |
| def test_bitsandbytes_is_warning_only(self) -> None: | |
| rule = next(r for r in _RULES if r.id == "optimizer.bitsandbytes_not_supported_warning") | |
| # Warning rule has empty transform — propose_patch must not auto-fix. | |
| assert rule.transform == {} | |
| def test_categories_cover_spec(self) -> None: | |
| # Architecture §4 lists 10 categories. We expect rules in at least | |
| # the high-impact ones the spec calls out as required. | |
| cats = {r.category for r in _RULES} | |
| for required in ( | |
| "precision", | |
| "attention", | |
| "memory", | |
| "data", | |
| "compile", | |
| "env_vars", | |
| "kernels", | |
| "optimizer", | |
| "collectives", | |
| "topology", | |
| ): | |
| assert required in cats, f"No rule for category {required!r}" | |
| # --------------------------------------------------------------------------- | |
| # Skip-and-warn behaviour for invalid entries | |
| # --------------------------------------------------------------------------- | |
| class TestLoadRulesResilience: | |
| def test_invalid_entry_is_skipped_with_warning(self, tmp_path: Path) -> None: | |
| bad = tmp_path / "rules.yaml" | |
| bad.write_text( | |
| yaml.safe_dump( | |
| [ | |
| { | |
| "id": "good.rule", | |
| "category": "precision", | |
| "targets_bucket": "precision_path", | |
| "symptom": "fp16 used", | |
| "expected_impact": "switch to bf16", | |
| "citation": "ROCm guide", | |
| }, | |
| {"id": "bad.rule", "category": "not_a_real_category"}, | |
| ] | |
| ) | |
| ) | |
| with pytest.warns(UserWarning): | |
| rules, _raw = _load_rules(bad) | |
| assert [r.id for r in rules] == ["good.rule"] | |
| def test_top_level_must_be_list(self, tmp_path: Path) -> None: | |
| bad = tmp_path / "rules.yaml" | |
| bad.write_text("not_a_list: 1\n") | |
| with pytest.raises(ValueError, match="top-level"): | |
| _load_rules(bad) | |
| # --------------------------------------------------------------------------- | |
| # Semantic search behaviour | |
| # --------------------------------------------------------------------------- | |
| class TestQuery: | |
| def test_returns_ok_for_real_query(self) -> None: | |
| result = _query_rocm_kb("fp16 used on MI300X with eager attention", top_k=5) | |
| assert isinstance(result, ToolResult) | |
| assert result.ok, result.error | |
| rules = result.result["rules"] | |
| assert 1 <= len(rules) <= 5 | |
| def test_fp16_query_returns_bf16_rule_in_top_results(self) -> None: | |
| result = _query_rocm_kb("fp16 used on MI300X / CDNA3", top_k=3) | |
| ids = [r["id"] for r in result.result["rules"]] | |
| assert "precision.bf16_over_fp16_on_mi300x" in ids | |
| def test_eager_attention_query_returns_attention_rules(self) -> None: | |
| result = _query_rocm_kb( | |
| "eager attention with no flash kernel loaded on MI300X", top_k=3 | |
| ) | |
| ids = [r["id"] for r in result.result["rules"]] | |
| attention_ids = { | |
| "attention.flash_rocm_over_eager", | |
| "attention.sdpa_over_eager", | |
| } | |
| assert attention_ids & set(ids), ( | |
| f"Expected at least one of {attention_ids} in top 3, got {ids}" | |
| ) | |
| def test_dataloader_query_returns_data_rules(self) -> None: | |
| result = _query_rocm_kb( | |
| "DataLoader num_workers is zero, GPU starves waiting for batches", | |
| top_k=3, | |
| ) | |
| ids = [r["id"] for r in result.result["rules"]] | |
| data_ids = { | |
| "data.dataloader_workers_zero", | |
| "data.pin_memory_false", | |
| "data.prefetch_factor_default", | |
| "data.persistent_workers_false", | |
| } | |
| assert data_ids & set(ids), ( | |
| f"Expected at least one data.* rule in top 3, got {ids}" | |
| ) | |
| def test_top_k_bounds_returned(self) -> None: | |
| result = _query_rocm_kb("anything", top_k=2) | |
| assert len(result.result["rules"]) == 2 | |
| def test_top_k_clamped_to_rule_count(self) -> None: | |
| # Asking for more than we have should not crash; we should get every | |
| # rule back, ordered by score. | |
| result = _query_rocm_kb("anything", top_k=100) | |
| assert len(result.result["rules"]) == len(_RULES) | |
| def test_results_sorted_by_descending_score(self) -> None: | |
| # If two queries with different focus return different top rules, | |
| # ordering is real — top_1 differs by query. | |
| a = _query_rocm_kb("fp16 numerical instability", top_k=1).result["rules"][0] | |
| b = _query_rocm_kb("eager attention slow on long sequences", top_k=1).result[ | |
| "rules" | |
| ][0] | |
| assert a["id"] != b["id"], ( | |
| "Top rule should depend on query, but both queries returned " | |
| f"{a['id']} — semantic search is degenerate." | |
| ) | |
| def test_rule_payload_is_lite_shape(self) -> None: | |
| # The LLM-facing rule payload is intentionally trimmed from the full | |
| # Rule schema — only id / symptom / transform / expected_impact / | |
| # citation make it through. This shrinks the audit conversation enough | |
| # to fit Qwen2.5-7B's 8K window. Full Rule lookup happens server-side | |
| # in propose_patch via the loaded KB. | |
| result = _query_rocm_kb("any query", top_k=1) | |
| payload = result.result["rules"][0] | |
| assert set(payload.keys()) == { | |
| "id", | |
| "symptom", | |
| "transform", | |
| "expected_impact", | |
| "citation", | |
| } | |
| # And the id resolves against the loaded KB so propose_patch can | |
| # reconstruct the full Rule. | |
| from agent.tools.query_rocm_kb import _RULES | |
| kb_ids = {r.id for r in _RULES} | |
| assert payload["id"] in kb_ids | |
| # --------------------------------------------------------------------------- | |
| # Failure modes | |
| # --------------------------------------------------------------------------- | |
| class TestErrors: | |
| def test_empty_symptom(self) -> None: | |
| result = _query_rocm_kb("", top_k=3) | |
| assert result.ok is False | |
| assert "symptom" in (result.error or "") | |
| def test_whitespace_only_symptom(self) -> None: | |
| result = _query_rocm_kb(" \t\n ", top_k=3) | |
| assert result.ok is False | |
| def test_top_k_zero(self) -> None: | |
| result = _query_rocm_kb("fp16", top_k=0) | |
| assert result.ok is False | |
| assert "top_k" in (result.error or "") | |
| def test_top_k_negative(self) -> None: | |
| result = _query_rocm_kb("fp16", top_k=-3) | |
| assert result.ok is False | |
| # --------------------------------------------------------------------------- | |
| # Embeddings cache | |
| # --------------------------------------------------------------------------- | |
| class TestEmbeddingsCache: | |
| def test_cache_file_was_created_on_import(self) -> None: | |
| raw = _KB_YAML.read_bytes() | |
| cache = _cache_path(raw) | |
| assert cache.exists(), ( | |
| f"Expected embeddings cache at {cache}; cache write failed silently." | |
| ) | |
| def test_second_embed_call_uses_cache_without_recoding( | |
| self, monkeypatch: pytest.MonkeyPatch | |
| ) -> None: | |
| rules, raw = _load_rules(_KB_YAML) | |
| # Sentinel: replace the lazy model getter so any encode() call would | |
| # blow up. If the cache is hit, the model is never consulted. | |
| from agent.tools import query_rocm_kb as kb_module | |
| def explode() -> None: | |
| raise AssertionError( | |
| "_get_model() called on a cache-hit path; cache is not being used." | |
| ) | |
| monkeypatch.setattr(kb_module, "_get_model", explode) | |
| embeddings = _embed_rules(rules, raw) | |
| assert embeddings.shape[0] == len(rules) | |
| assert embeddings.dtype.kind == "f" | |
| # --------------------------------------------------------------------------- | |
| # Tool registry shape — the agent loop depends on these fields being stable. | |
| # --------------------------------------------------------------------------- | |
| class TestToolDefinition: | |
| def test_name_is_query_rocm_kb(self) -> None: | |
| assert QUERY_ROCM_KB.name == "query_rocm_kb" | |
| def test_description_unchanged_keywords(self) -> None: | |
| # Must still describe the search-by-symptom semantics; the system | |
| # prompt references this language. | |
| desc = QUERY_ROCM_KB.description | |
| assert "ROCm" in desc and "symptom" in desc | |
| def test_input_schema_shape(self) -> None: | |
| schema = QUERY_ROCM_KB.input_schema | |
| assert schema["type"] == "object" | |
| # Both single-query (`symptom`) and batched (`symptoms`) shapes are | |
| # advertised — either works at runtime, neither is strictly required | |
| # in the schema because the impl validates "at least one" itself. | |
| assert "symptom" in schema["properties"] | |
| assert "symptoms" in schema["properties"] | |
| assert "top_k" in schema["properties"] | |
| assert schema["properties"]["symptoms"]["type"] == "array" | |
| assert schema["properties"]["symptoms"]["items"] == {"type": "string"} | |
| def test_fn_is_module_query(self) -> None: | |
| assert QUERY_ROCM_KB.fn is _query_rocm_kb | |