| 1 |
"""Shared tiny-model-aware planning helpers for integration tests.""" |
| 2 |
|
| 3 |
from __future__ import annotations |
| 4 |
|
| 5 |
from typing import TYPE_CHECKING |
| 6 |
|
| 7 |
import pytest |
| 8 |
|
| 9 |
if TYPE_CHECKING: |
| 10 |
from dlm.base_models.schema import BaseModelSpec |
| 11 |
from dlm.doc.parser import ParsedDlm |
| 12 |
from dlm.hardware.capabilities import Capabilities |
| 13 |
from dlm.hardware.plan import TrainingPlan |
| 14 |
|
| 15 |
|
| 16 |
def resolve_spec_and_plan( |
| 17 |
parsed: ParsedDlm, |
| 18 |
*, |
| 19 |
accept_license: bool = False, |
| 20 |
skip_reason: str = "doctor() returned no viable training plan on this host", |
| 21 |
) -> tuple[BaseModelSpec, TrainingPlan, Capabilities]: |
| 22 |
"""Resolve the base spec and a host-appropriate plan for `parsed`. |
| 23 |
|
| 24 |
Slow integration tests should not call bare `doctor()` for tiny-model |
| 25 |
fixtures; that falls back to a generic host heuristic and can reject |
| 26 |
CPU-capable tiny bases that should run fine on CI. Mirror the real |
| 27 |
train path instead: feed the parsed training config, the base-model |
| 28 |
parameter count, and the bounded effective context length. |
| 29 |
""" |
| 30 |
|
| 31 |
from dlm.base_models import resolve as resolve_base_model |
| 32 |
from dlm.hardware import doctor |
| 33 |
|
| 34 |
spec = resolve_base_model(parsed.frontmatter.base_model, accept_license=accept_license) |
| 35 |
doctor_result = doctor( |
| 36 |
training_config=parsed.frontmatter.training, |
| 37 |
base_params=spec.params, |
| 38 |
seq_len=min(parsed.frontmatter.training.sequence_len, spec.effective_context_length), |
| 39 |
) |
| 40 |
plan = doctor_result.plan |
| 41 |
if plan is None: |
| 42 |
pytest.skip(skip_reason) |
| 43 |
return spec, plan, doctor_result.capabilities |