File size: 7,180 Bytes
94cbe85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | import importlib.util
import json
from pathlib import Path
import pytest
def _row(row_id: str, task_type: str, category: str, repair_scope: str | None = None) -> dict:
metadata = {"task_type": task_type}
if repair_scope:
metadata["repair_scope"] = repair_scope
return {
"case_id": row_id,
"uuid": row_id,
"category": category,
"messages": [
{"role": "user", "content": f"prompt {row_id}"},
{"role": "assistant", "content": '{"protocol_urgency":"routine"}'},
],
"metadata": metadata,
}
def _write_jsonl(path: Path, rows: list[dict]) -> None:
path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8")
def _load_modal_module():
module_path = Path(__file__).resolve().parents[1] / "modal" / "finetune_figment_nemotron.py"
spec = importlib.util.spec_from_file_location("figment_modal_finetune", module_path)
assert spec is not None
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def test_prepare_modal_dataset_stratifies_and_preserves_rows(tmp_path):
from scripts.prepare_modal_finetune_dataset import prepare_dataset
rows = []
rows.extend(_row(f"full-{index}", "navigator_full", "missing_observation_cues") for index in range(8))
rows.extend(_row(f"repair-schema-{index}", "focused_repair", "focused_repair:schema", "schema") for index in range(6))
rows.extend(
_row(
f"repair-sbar-{index}",
"focused_repair",
"focused_repair:handoff_note_sbar",
"handoff_note_sbar",
)
for index in range(6)
)
dataset = tmp_path / "input.jsonl"
output_dir = tmp_path / "prepared"
_write_jsonl(dataset, rows)
manifest = prepare_dataset(
dataset_path=dataset,
output_dir=output_dir,
dataset_version="figment_sft_test",
validation_fraction=0.2,
seed="unit-test",
min_validation_group_size=5,
)
train_rows = [json.loads(line) for line in (output_dir / "train.jsonl").read_text().splitlines()]
validation_rows = [json.loads(line) for line in (output_dir / "validation.jsonl").read_text().splitlines()]
train_ids = {row["uuid"] for row in train_rows}
validation_ids = {row["uuid"] for row in validation_rows}
assert train_ids.isdisjoint(validation_ids)
assert train_ids | validation_ids == {row["uuid"] for row in rows}
assert manifest["row_count"] == 20
assert manifest["train_count"] + manifest["validation_count"] == 20
assert manifest["validation_group_counts"]["navigator_full:missing_observation_cues"] >= 1
assert manifest["validation_group_counts"]["focused_repair:schema"] >= 1
assert manifest["validation_group_counts"]["focused_repair:handoff_note_sbar"] >= 1
def test_prepare_modal_dataset_rejects_rows_outside_chat_shape(tmp_path):
from scripts.prepare_modal_finetune_dataset import DatasetPrepError
from scripts.prepare_modal_finetune_dataset import prepare_dataset
dataset = tmp_path / "bad.jsonl"
bad_row = _row("bad-1", "navigator_full", "missing_observation_cues")
bad_row["messages"] = [{"role": "user", "content": "prompt only"}]
_write_jsonl(dataset, [bad_row])
with pytest.raises(DatasetPrepError, match="expected user/assistant messages"):
prepare_dataset(
dataset_path=dataset,
output_dir=tmp_path / "prepared",
dataset_version="figment_sft_test",
)
def test_modal_smoke_config_is_small_and_namespaced():
module = _load_modal_module()
config = module.build_train_config(
dataset_version="figment_sft_v1",
output_name="unit",
smoke=True,
max_steps=100,
max_seq_length=12288,
)
paths = module.dataset_volume_paths("figment_sft_v1")
assert config["max_steps"] == 5
assert config["max_seq_length"] == 2048
assert config["output_dir"].endswith("/figment_sft_v1/unit-smoke")
assert paths["train"].endswith("/figment_sft_v1/train.jsonl")
assert paths["validation"].endswith("/figment_sft_v1/validation.jsonl")
def test_modal_default_config_matches_first_run_plan():
module = _load_modal_module()
config = module.build_train_config(dataset_version="figment_sft_v1", output_name="pilot")
assert config["learning_rate"] == 1e-4
assert config["max_steps"] == 40
assert config["max_seq_length"] == 16384
assert config["gradient_accumulation_steps"] == 8
assert config["validation_steps"] == 25
assert config["save_steps"] == 40
def test_modal_v4_config_can_use_lower_lr_and_lora_controls():
module = _load_modal_module()
config = module.build_train_config(
dataset_version="figment_sft_v4",
output_name="figment-sft-v4-lora",
max_steps=900,
learning_rate=2e-5,
lora_r=16,
lora_alpha=32,
lora_dropout=0.05,
gradient_accumulation_steps=8,
validation_steps=50,
save_steps=100,
)
assert config["dataset_version"] == "figment_sft_v4"
assert config["learning_rate"] == 2e-5
assert config["lora_r"] == 16
assert config["lora_alpha"] == 32
assert config["lora_dropout"] == 0.05
assert config["gradient_accumulation_steps"] == 8
assert config["validation_steps"] == 50
assert config["save_steps"] == 100
def test_modal_v5_config_can_resume_from_v4_adapter():
module = _load_modal_module()
config = module.build_train_config(
dataset_version="figment_sft_v5",
output_name="figment-sft-v5-lora",
resume_adapter_name="figment-sft-v4-lora",
resume_adapter_dataset_version="figment_sft_v4",
)
assert config["resume_adapter_name"] == "figment-sft-v4-lora"
assert config["resume_adapter_dataset_version"] == "figment_sft_v4"
assert config["resume_adapter_dir"] == "/checkpoints/figment_sft_v4/figment-sft-v4-lora"
assert config["output_dir"] == "/checkpoints/figment_sft_v5/figment-sft-v5-lora"
def test_modal_entrypoint_exposes_v4_training_knobs():
source = (Path(__file__).resolve().parents[1] / "modal" / "finetune_figment_nemotron.py").read_text(
encoding="utf-8"
)
for parameter in (
"learning_rate",
"lora_r",
"lora_alpha",
"lora_dropout",
"gradient_accumulation_steps",
"validation_steps",
"save_steps",
"resume_adapter_name",
"resume_adapter_dataset_version",
):
assert f"{parameter}:" in source
assert f"{parameter}={parameter}" in source or f'"{parameter}": {parameter}' in source
def test_modal_merge_config_points_at_checkpoint_volume():
module = _load_modal_module()
config = module.build_merge_config(
dataset_version="figment_sft_v1",
adapter_name="pilot-20260608",
)
assert config["adapter_dir"] == "/checkpoints/figment_sft_v1/pilot-20260608"
assert config["output_dir"] == "/checkpoints/figment_sft_v1/pilot-20260608-merged-bf16"
assert config["output_name"] == "pilot-20260608-merged-bf16"
|