Skip to content

Commit c20365d

Browse files
committed
MCM: fix auto-detection of vLLM binary cache
_has_artifact_compile_range_with_triton() only checked for a triton/ subdirectory, which cannot exist when the artifact is a packed binary file. Recognize binary artifact_compile_range_* files as valid vLLM cache indicators so detect_cache_mode() returns 'vllm' instead of falling through to 'triton'. Also: - sync requirements.txt with pyproject.toml (typer[all], structlog) - silence pylint R0903 on Pydantic data models - disable pylint import-error for declared but not-installed deps
1 parent 80bd9ac commit c20365d

File tree

5 files changed

+31
-7
lines changed

5 files changed

+31
-7
lines changed

mcm/model_cache_manager/data/kernel_validator.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,15 @@
1616
log = logging.getLogger(__name__)
1717

1818

19-
class KernelTarget(BaseModel):
19+
class KernelTarget(BaseModel): # pylint: disable=too-few-public-methods
2020
"""Validation model for the 'target' field of kernel metadata."""
2121

2222
backend: str = ""
2323
arch: int | str = 0
2424
warp_size: int = 0
2525

2626

27-
class KernelMetadata(BaseModel):
27+
class KernelMetadata(BaseModel): # pylint: disable=too-few-public-methods
2828
"""
2929
Pydantic model for validating kernel metadata.
3030

mcm/model_cache_manager/tests/utils/test_cache_mode_detection.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,23 @@ def test_detect_vllm_structure_no_rank_dirs(self):
244244
mode = detect_cache_mode(cache_dir)
245245
self.assertEqual(mode, MODE_TRITON)
246246

247+
def test_detect_vllm_binary_cache_structure(self):
248+
"""Test detection of new vLLM cache structure with binary artifacts."""
249+
# Binary artifacts are files (not directories) named artifact_compile_range_*
250+
vllm_dir = self.temp_dir / "vllm_binary"
251+
torch_compile = vllm_dir / "torch_compile_cache"
252+
hash_dir = torch_compile / "some_vllm_hash"
253+
rank_dir = hash_dir / "rank_0_0"
254+
backbone_dir = rank_dir / "backbone"
255+
backbone_dir.mkdir(parents=True)
256+
257+
# Create a binary artifact file (not a directory)
258+
binary_artifact = backbone_dir / "artifact_compile_range_0"
259+
binary_artifact.write_bytes(b"\x00binary data")
260+
261+
mode = detect_cache_mode(vllm_dir)
262+
self.assertEqual(mode, MODE_VLLM)
263+
247264
def test_detect_mixed_legacy_and_triton_vllm_takes_precedence(self):
248265
"""Test detection when both vllm-legacy and triton structures exist."""
249266
cache_dir = self.temp_dir / "mixed_cache"

mcm/model_cache_manager/utils/utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,8 +198,12 @@ def iter_artifact_compile_range_dirs(backbone_dir: Path):
198198

199199

200200
def _has_artifact_compile_range_with_triton(backbone_dir: Path) -> bool:
201-
"""Check if backbone directory has artifact_compile_range subdirectories with triton."""
201+
"""Check if backbone directory has artifact_compile_range artifacts (unpacked or binary)."""
202202
for artifact_dir in iter_artifact_compile_range_dirs(backbone_dir):
203+
# Binary artifact (file) — contains packed triton data
204+
if artifact_dir.is_file():
205+
return True
206+
# Unpacked artifact (directory) — check for triton subdirectory
203207
triton_dir = artifact_dir / "triton"
204208
if triton_dir.exists():
205209
return True

mcm/pyproject.toml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,21 @@ authors = [{name = "Alessandro Sangiorgi", email = "asangior@redhat.com"
1010
requires-python = ">=3.9"
1111

1212
dependencies = [
13-
"triton",
1413
"rich>=13.7",
1514
"typer[all]",
1615
"structlog",
1716
"pydantic>=2",
1817
"pydantic-settings>=2",
1918
"sqlalchemy>=2.0.40",
19+
"vllm==0.15.1",
2020
]
2121

2222
[project.scripts]
2323
mcm = "model_cache_manager.cli.main:run"
2424

25+
[tool.pylint."messages_control"]
26+
disable = ["import-error"]
27+
2528
[tool.setuptools]
2629
package-dir = {"" = "."}
2730

mcm/requirements.txt

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
typer>=0.9.0
1+
typer[all]
22
rich>=13.7
33
pydantic>=2.0.0
44
pydantic-settings>=2.0.0
5-
structlog>=23.0.0
5+
structlog
66
sqlalchemy>=2.0.40
77
pytest>=7.0.0
8-
vllm
8+
vllm==0.15.1

0 commit comments

Comments
 (0)