|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import json |
| 4 | +import os |
4 | 5 | import tempfile |
5 | 6 | import unittest |
6 | 7 | from pathlib import Path |
7 | 8 |
|
8 | 9 | from tools.cosyvoice_resume_contract import ( |
9 | 10 | LOCK_NAME, |
| 11 | + OPTIMIZER_STATE_NAME, |
10 | 12 | ResumeContractError, |
| 13 | + SCHEDULER_STATE_NAME, |
11 | 14 | acquire_output_lock, |
12 | 15 | build_contract, |
13 | 16 | checkpoint_children, |
14 | 17 | epoch_checkpoint_name, |
| 18 | + evaluator_lora_artifact_paths, |
15 | 19 | prune_owned_checkpoints, |
16 | 20 | publish_checkpoint, |
17 | 21 | require_fresh_output, |
@@ -72,6 +76,8 @@ def _adapter_saver(directory: Path) -> None: |
72 | 76 | @staticmethod |
73 | 77 | def _runtime_saver(path: Path) -> None: |
74 | 78 | path.write_bytes(b"trusted-pickle-fixture") |
| 79 | + (path.parent / OPTIMIZER_STATE_NAME).write_bytes(b"optimizer-fixture") |
| 80 | + (path.parent / SCHEDULER_STATE_NAME).write_bytes(b"scheduler-fixture") |
75 | 81 |
|
76 | 82 | def _checkpoint( |
77 | 83 | self, |
@@ -110,6 +116,97 @@ def test_exact_trusted_checkpoint_validates(self) -> None: |
110 | 116 | self.assertEqual(state["completed_epoch"], 1) |
111 | 117 | self.assertEqual(state["monitor_state"]["cv_no_improve_epochs"], 1) |
112 | 118 |
|
| 119 | + def test_evaluator_lora_artifact_roles_are_independent(self) -> None: |
| 120 | + checkpoint = self._checkpoint(1) |
| 121 | + artifacts = evaluator_lora_artifact_paths(checkpoint) |
| 122 | + self.assertEqual( |
| 123 | + {role: path.name for role, path in artifacts.items()}, |
| 124 | + { |
| 125 | + "model_state": "adapter_model.safetensors", |
| 126 | + "optimizer_state": OPTIMIZER_STATE_NAME, |
| 127 | + "scheduler_state": SCHEDULER_STATE_NAME, |
| 128 | + "trainer_state": "training-state.json", |
| 129 | + "rng_state": "runtime-state.pt", |
| 130 | + }, |
| 131 | + ) |
| 132 | + |
| 133 | + def test_legacy_combined_runtime_state_remains_resumable(self) -> None: |
| 134 | + checkpoint = publish_checkpoint( |
| 135 | + output_dir=self.output, |
| 136 | + completed_epoch=1, |
| 137 | + completed_step=10, |
| 138 | + contract=self.contract, |
| 139 | + adapter_saver=self._adapter_saver, |
| 140 | + runtime_state_saver=lambda path: path.write_bytes(b"legacy-combined-state"), |
| 141 | + monitor_state={"best_cv_loss": 3.0}, |
| 142 | + ) |
| 143 | + selected, _ = validate_checkpoint( |
| 144 | + checkpoint, |
| 145 | + output_dir=self.output, |
| 146 | + expected_contract=self.contract, |
| 147 | + trust_resume_state=True, |
| 148 | + world_size=1, |
| 149 | + train_engine="torch_ddp", |
| 150 | + ) |
| 151 | + self.assertEqual(selected, checkpoint) |
| 152 | + with self.assertRaisesRegex(ResumeContractError, "omits decomposed state"): |
| 153 | + evaluator_lora_artifact_paths(checkpoint) |
| 154 | + |
| 155 | + def test_evaluator_mapping_rejects_ambiguous_model_and_hardlinks(self) -> None: |
| 156 | + def ambiguous_adapter(directory: Path) -> None: |
| 157 | + self._adapter_saver(directory) |
| 158 | + (directory / "adapter_model.bin").write_bytes(b"second-adapter") |
| 159 | + |
| 160 | + ambiguous = publish_checkpoint( |
| 161 | + output_dir=self.output, |
| 162 | + completed_epoch=1, |
| 163 | + completed_step=10, |
| 164 | + contract=self.contract, |
| 165 | + adapter_saver=ambiguous_adapter, |
| 166 | + runtime_state_saver=self._runtime_saver, |
| 167 | + monitor_state={}, |
| 168 | + ) |
| 169 | + with self.assertRaisesRegex(ResumeContractError, "exactly one adapter"): |
| 170 | + evaluator_lora_artifact_paths(ambiguous) |
| 171 | + |
| 172 | + other_output = self.root / "hardlink-output" |
| 173 | + other_output.mkdir() |
| 174 | + hardlink_contract = build_contract( |
| 175 | + output_dir=other_output, |
| 176 | + base_checkpoint=self.base, |
| 177 | + config_file=self.config, |
| 178 | + qwen_pretrain=self.qwen, |
| 179 | + data_files={"train": [self.train_list], "cross_validation": [self.cv_list]}, |
| 180 | + source_files=[self.source], |
| 181 | + training_config={"max_epoch": 5, "train_engine": "torch_ddp"}, |
| 182 | + runtime={"python": "fixture", "world_size": 1}, |
| 183 | + ) |
| 184 | + |
| 185 | + def hardlinked_runtime(path: Path) -> None: |
| 186 | + path.write_bytes(b"runtime") |
| 187 | + optimizer = path.parent / OPTIMIZER_STATE_NAME |
| 188 | + optimizer.write_bytes(b"shared-state") |
| 189 | + os.link(optimizer, path.parent / SCHEDULER_STATE_NAME) |
| 190 | + |
| 191 | + hardlinked = publish_checkpoint( |
| 192 | + output_dir=other_output, |
| 193 | + completed_epoch=1, |
| 194 | + completed_step=10, |
| 195 | + contract=hardlink_contract, |
| 196 | + adapter_saver=self._adapter_saver, |
| 197 | + runtime_state_saver=hardlinked_runtime, |
| 198 | + monitor_state={}, |
| 199 | + ) |
| 200 | + with self.assertRaisesRegex(ResumeContractError, "must not share hardlinks"): |
| 201 | + evaluator_lora_artifact_paths(hardlinked) |
| 202 | + |
| 203 | + def test_trainer_writes_decomposed_state_before_publication(self) -> None: |
| 204 | + source = ( |
| 205 | + Path(__file__).parents[1] / "tools" / "train_cosyvoice3_lora.py" |
| 206 | + ).read_text(encoding="utf-8") |
| 207 | + self.assertIn("path.parent / OPTIMIZER_STATE_NAME", source) |
| 208 | + self.assertIn("path.parent / SCHEDULER_STATE_NAME", source) |
| 209 | + |
113 | 210 | def test_resume_requires_explicit_trust(self) -> None: |
114 | 211 | checkpoint = self._checkpoint(1) |
115 | 212 | with self.assertRaisesRegex(ResumeContractError, "pickle-capable"): |
|
0 commit comments