Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / OURS /tests /test_sft_runtime.py

EmergentTwilight

24 days ago

download

raw

17.6 kB

import hashlib
import copy
import json
import random
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

import numpy as np
import torch
from sft.config import SFTConfig
from sft.runtime import (
    acquire_run_writer,
    cache_fingerprint,
    claim_run_directory,
    file_sha256,
    load_training_state,
    resolve_cuda_device_index,
    resolve_run_directory,
    save_checkpoint,
)


class _SavedObject:
    def save_pretrained(self, path, **_kwargs):
        Path(path).mkdir(parents=True, exist_ok=True)
        (Path(path) / "fixture.txt").write_text("fixture\n", encoding="utf-8")


class SFTRuntimeTest(unittest.TestCase):
    def test_run_writer_lock_rejects_a_concurrent_writer(self):
        with tempfile.TemporaryDirectory() as directory:
            first = acquire_run_writer(directory)
            try:
                with self.assertRaisesRegex(RuntimeError, "active writer"):
                    acquire_run_writer(directory)
            finally:
                first.close()
            second = acquire_run_writer(directory)
            second.close()

    @staticmethod
    def _prepare_submission(run: Path, job_id: str = "123") -> None:
        submission = run / "_submission"
        submission.mkdir(parents=True)
        ready = submission / "READY"
        ready.write_text(
            json.dumps({"schema": 1, "metadata_sha256": {}}) + "\n",
            encoding="utf-8",
        )
        (submission / "manifest.json").write_text(
            json.dumps({"copied_hash": "source-hash"}) + "\n", encoding="utf-8"
        )
        (submission / f"runtime-verification-{job_id}.json").write_text(
            json.dumps(
                {
                    "schema": "spatial_mmcot_runtime_verification_v1",
                    "slurm_job_id": job_id,
                    "source_hash": "source-hash",
                    "ready_sha256": hashlib.sha256(ready.read_bytes()).hexdigest(),
                }
            )
            + "\n",
            encoding="utf-8",
        )

    def test_claim_run_directory_creates_a_fresh_run(self):
        with tempfile.TemporaryDirectory() as directory:
            run = Path(directory) / "run"
            self.assertEqual(claim_run_directory(run), run)
            self.assertTrue(run.is_dir())
            self.assertFalse((run / ".run-claim.json").exists())

    def test_claim_run_directory_accepts_only_a_ready_submission(self):
        with tempfile.TemporaryDirectory() as directory, patch.dict(
            "os.environ", {"SLURM_JOB_ID": "123"}, clear=False
        ):
            run = Path(directory) / "run"
            self._prepare_submission(run)
            (run / "slurm-123.out").touch()
            claim_run_directory(run)
            payload = (run / ".run-claim.json").read_text(encoding="utf-8")
            self.assertIn('"slurm_job_id": "123"', payload)
            with self.assertRaisesRegex(RuntimeError, "already claimed"):
                claim_run_directory(run)
            self.assertEqual(claim_run_directory(run, resume=True), run)

    def test_claim_run_directory_rejects_unready_or_dirty_directories(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            unready = root / "unready"
            unready.mkdir()
            with self.assertRaisesRegex(FileExistsError, "lacks"):
                claim_run_directory(unready)
            dirty = root / "dirty"
            with patch.dict("os.environ", {"SLURM_JOB_ID": "123"}, clear=False):
                self._prepare_submission(dirty)
            (dirty / "metrics.json").write_text("{}\n", encoding="utf-8")
            with (
                patch.dict("os.environ", {"SLURM_JOB_ID": "123"}, clear=False),
                self.assertRaisesRegex(FileExistsError, "experiment artifacts"),
            ):
                claim_run_directory(dirty)

    def test_claim_run_directory_rejects_unattested_ready_marker(self):
        with tempfile.TemporaryDirectory() as directory, patch.dict(
            "os.environ", {"SLURM_JOB_ID": "123"}, clear=False
        ):
            run = Path(directory) / "run"
            (run / "_submission").mkdir(parents=True)
            (run / "_submission/READY").write_text("ready\n", encoding="utf-8")
            with self.assertRaisesRegex(FileExistsError, "invalid submission"):
                claim_run_directory(run)

    def test_model_first_run_directory_is_repository_relative(self):
        config = SFTConfig()
        config.run.dataset = "zebra-cot-small"
        config.run.category = "sweeps"
        config.run.name = "lambda_grid"
        with tempfile.TemporaryDirectory() as directory, patch(
            "os.getcwd", return_value=directory
        ):
            run = resolve_run_directory(config)
        repository = Path(__file__).resolve().parents[2]
        self.assertEqual(run.parents[7], repository)
        self.assertEqual(
            run.relative_to(repository / "experiments/output").parts[:5],
            ("uniar-sft", "zebra-cot-small", "unversioned", "sweeps", "sft"),
        )

    def test_probe_override_routes_to_smoke(self):
        config = SFTConfig()
        config.run.dataset = "zebra-cot-small"
        run = resolve_run_directory(config, "sft_probe", category="smoke")
        self.assertEqual(run.parent.name, "sft_probe")
        self.assertEqual(run.parent.parent.name, "smoke")

    def test_versioned_run_directory_places_version_after_dataset(self):
        config = SFTConfig()
        config.run.dataset = "zebra-cot-rpm"
        config.run.version = "v1"
        config.run.category = "sweeps"
        config.run.name = "lambda_grid"
        run = resolve_run_directory(config)
        repository = Path(__file__).resolve().parents[2]
        self.assertEqual(
            run.relative_to(repository / "experiments/output").parts[:5],
            ("uniar-sft", "zebra-cot-rpm", "v1", "sweeps", "sft"),
        )

    def test_queue_physical_gpu_maps_to_visible_logical_ordinal(self):
        with patch.dict(
            "os.environ",
            {"GPUQ_ASSIGNED_ID": "5", "CUDA_VISIBLE_DEVICES": "4,5"},
            clear=False,
        ):
            self.assertEqual(resolve_cuda_device_index(local_rank=0, world_size=1), 1)
        with (
            patch.dict(
                "os.environ",
                {"GPUQ_ASSIGNED_ID": "5", "CUDA_VISIBLE_DEVICES": "GPU-uuid"},
                clear=False,
            ),
            patch("torch.cuda.device_count", return_value=1),
        ):
            self.assertEqual(resolve_cuda_device_index(local_rank=0, world_size=1), 0)
        self.assertEqual(resolve_cuda_device_index(local_rank=3, world_size=4), 3)

    def test_cache_fingerprint_does_not_hash_tensor_shard_contents(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "READY").write_text("ready\n", encoding="utf-8")
            (root / "manifest.jsonl").write_text("{}\n", encoding="utf-8")
            (root / "shard.pt").write_bytes(b"tensor-payload")
            with patch("shared.artifacts.file_sha256", wraps=file_sha256) as hashed:
                cache_fingerprint(root, [{"cache_shard": "shard.pt"}])
            hashed_paths = {call.args[0].name for call in hashed.call_args_list}
            self.assertEqual(hashed_paths, {"READY", "manifest.jsonl"})

    def test_atomic_checkpoint_restores_world_one_rng_and_state(self):
        with (
            tempfile.TemporaryDirectory() as directory,
            patch("torch.cuda.is_available", return_value=False),
        ):
            root = Path(directory)
            checkpoints = root / "checkpoints"
            checkpoints.mkdir()
            parameter = torch.nn.Parameter(torch.tensor(1.0))
            optimizer = torch.optim.AdamW([parameter], lr=1e-3)
            scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
            config = SFTConfig()
            trainer_state = {
                "attempted_updates": 3,
                "successful_updates": 2,
                "nonfinite_updates": 1,
                "microbatch_cursor": 12,
            }
            random.seed(9)
            np.random.seed(9)
            torch.manual_seed(9)
            checkpoint = save_checkpoint(
                checkpoints_dir=checkpoints,
                model=_SavedObject(),
                processor=_SavedObject(),
                optimizer=optimizer,
                scheduler=scheduler,
                trainer_state=trainer_state,
                config=config,
                provenance_fingerprints={
                    "cache": "cache-v1",
                    "model": "model-v1",
                    "source": "source-v1",
                },
                rank=0,
                world_size=1,
                ema_state={"schema": "fixture", "update_count": 2},
            )
            expected = (random.random(), float(np.random.rand()), float(torch.rand(())))
            for _ in range(5):
                random.random()
                np.random.rand()
                torch.rand(())
            restored_ema = []
            restored = load_training_state(
                checkpoint,
                optimizer=optimizer,
                scheduler=scheduler,
                config=config,
                provenance_fingerprints={
                    "cache": "cache-v1",
                    "model": "model-v1",
                    "source": "source-v1",
                },
                rank=0,
                world_size=1,
                ema_state_loader=restored_ema.append,
            )
            actual = (random.random(), float(np.random.rand()), float(torch.rand(())))
            self.assertEqual(restored, trainer_state)
            self.assertEqual(
                restored_ema, [{"schema": "fixture", "update_count": 2}]
            )
            self.assertEqual(actual, expected)
            self.assertTrue((checkpoint / "COMPLETE").is_file())
            self.assertFalse(any(checkpoints.glob(".tmp-*")))

    def test_parent_linked_ddp_execution_change_can_restore_exact_state(self):
        with (
            tempfile.TemporaryDirectory() as directory,
            patch("torch.cuda.is_available", return_value=False),
        ):
            root = Path(directory)
            parent = root / "parent"
            checkpoints = parent / "checkpoints"
            checkpoints.mkdir(parents=True)
            parameter = torch.nn.Parameter(torch.tensor(1.0))
            optimizer = torch.optim.AdamW([parameter], lr=1e-3)
            scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
            original = SFTConfig()
            checkpoint = save_checkpoint(
                checkpoints_dir=checkpoints,
                model=_SavedObject(),
                processor=_SavedObject(),
                optimizer=optimizer,
                scheduler=scheduler,
                trainer_state={"successful_updates": 3},
                config=original,
                provenance_fingerprints={"cache": "c", "model": "m", "source": "old"},
                rank=0,
                world_size=1,
            )
            (parent / "config.json").write_text(
                json.dumps(
                    {
                        "provenance_fingerprints": {
                            "cache": "c",
                            "model": "m",
                            "source": "old",
                        }
                    }
                ),
                encoding="utf-8",
            )
            child = root / "child"
            (child / "_submission").mkdir(parents=True)
            (parent / "_submission").mkdir(parents=True)
            (parent / "_submission/manifest.json").write_text(
                json.dumps({"copied_hash": "parent-source"}), encoding="utf-8"
            )
            (child / "_submission/submission.json").write_text(
                json.dumps(
                    {
                        "parent_run": {
                            "run_dir": str(parent),
                            "source_hash": "parent-source",
                        }
                    }
                ),
                encoding="utf-8",
            )
            migrated = copy.deepcopy(original)
            migrated.trainer.ddp_find_unused_parameters = False
            migrated.checkpoint.resume_mode = "ddp_execution_change"
            restored = load_training_state(
                checkpoint,
                optimizer=optimizer,
                scheduler=scheduler,
                config=migrated,
                provenance_fingerprints={"cache": "c", "model": "m", "source": "new"},
                rank=0,
                world_size=1,
                run_dir=child,
            )
            self.assertEqual(restored, {"successful_updates": 3})

            control_migration = copy.deepcopy(original)
            control_migration.checkpoint.resume_mode = (
                "distributed_execution_change"
            )
            restored = load_training_state(
                checkpoint,
                optimizer=optimizer,
                scheduler=scheduler,
                config=control_migration,
                provenance_fingerprints={
                    "cache": "c",
                    "model": "m",
                    "source": "new",
                },
                rank=0,
                world_size=1,
                run_dir=child,
            )
            self.assertEqual(restored, {"successful_updates": 3})

            migrated.objective.text.weight = 2.0
            with self.assertRaisesRegex(ValueError, "non-DDP configuration"):
                load_training_state(
                    checkpoint,
                    optimizer=optimizer,
                    scheduler=scheduler,
                    config=migrated,
                    provenance_fingerprints={"cache": "c", "model": "m", "source": "new"},
                    rank=0,
                    world_size=1,
                    run_dir=child,
                )

    def test_parent_linked_epoch_extension_restores_final_state(self):
        with (
            tempfile.TemporaryDirectory() as directory,
            patch("torch.cuda.is_available", return_value=False),
        ):
            root = Path(directory)
            parent = root / "parent"
            checkpoints = parent / "checkpoints"
            checkpoints.mkdir(parents=True)
            parameter = torch.nn.Parameter(torch.tensor(1.0))
            optimizer = torch.optim.AdamW([parameter], lr=1e-3)
            scheduler = torch.optim.lr_scheduler.LambdaLR(
                optimizer, lambda _: 1.0
            )
            original = SFTConfig()
            original.run.name = "two_epoch"
            original.scheduler.kind = "constant"
            original.trainer.epochs = 2
            state = {
                "attempted_updates": 20,
                "successful_updates": 20,
                "nonfinite_updates": 0,
                "microbatch_cursor": 160,
            }
            checkpoint = save_checkpoint(
                checkpoints_dir=checkpoints,
                model=_SavedObject(),
                processor=_SavedObject(),
                optimizer=optimizer,
                scheduler=scheduler,
                trainer_state=state,
                config=original,
                provenance_fingerprints={
                    "cache": "c",
                    "model": "m",
                    "source": "old",
                },
                rank=0,
                world_size=1,
            )
            (parent / "config.json").write_text(
                json.dumps(
                    {
                        "provenance_fingerprints": {
                            "cache": "c",
                            "model": "m",
                            "source": "old",
                        }
                    }
                ),
                encoding="utf-8",
            )
            (parent / "summary.json").write_text(
                json.dumps({"successful_updates": 20, "target_updates": 20}),
                encoding="utf-8",
            )
            (parent / "_submission").mkdir()
            (parent / "_submission/manifest.json").write_text(
                json.dumps({"copied_hash": "parent-source"}), encoding="utf-8"
            )
            child = root / "child"
            (child / "_submission").mkdir(parents=True)
            (child / "_submission/submission.json").write_text(
                json.dumps(
                    {
                        "parent_run": {
                            "run_dir": str(parent),
                            "source_hash": "parent-source",
                        }
                    }
                ),
                encoding="utf-8",
            )
            extended = copy.deepcopy(original)
            extended.run.name = "three_epoch"
            extended.trainer.epochs = 3
            extended.checkpoint.resume_mode = "epoch_extension"

            restored = load_training_state(
                checkpoint,
                optimizer=optimizer,
                scheduler=scheduler,
                config=extended,
                provenance_fingerprints={
                    "cache": "c",
                    "model": "m",
                    "source": "new",
                },
                rank=0,
                world_size=1,
                run_dir=child,
            )

            self.assertEqual(restored, state)

if __name__ == "__main__":
    unittest.main()

Xet Storage Details

Size:: 17.6 kB
Xet hash:: a75e80d23222de782d40b49a3b3b0016709632a1a6eec58671854d116ed74898

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import hashlib
	import copy
	import json
	import random
	import tempfile
	import unittest
	from pathlib import Path
	from unittest.mock import patch

	import numpy as np
	import torch
	from sft.config import SFTConfig
	from sft.runtime import (
	acquire_run_writer,
	cache_fingerprint,
	claim_run_directory,
	file_sha256,
	load_training_state,
	resolve_cuda_device_index,
	resolve_run_directory,
	save_checkpoint,
	)


	class _SavedObject:
	def save_pretrained(self, path, **_kwargs):
	Path(path).mkdir(parents=True, exist_ok=True)
	(Path(path) / "fixture.txt").write_text("fixture\n", encoding="utf-8")


	class SFTRuntimeTest(unittest.TestCase):
	def test_run_writer_lock_rejects_a_concurrent_writer(self):
	with tempfile.TemporaryDirectory() as directory:
	first = acquire_run_writer(directory)
	try:
	with self.assertRaisesRegex(RuntimeError, "active writer"):
	acquire_run_writer(directory)
	finally:
	first.close()
	second = acquire_run_writer(directory)
	second.close()

	@staticmethod
	def _prepare_submission(run: Path, job_id: str = "123") -> None:
	submission = run / "_submission"
	submission.mkdir(parents=True)
	ready = submission / "READY"
	ready.write_text(
	json.dumps({"schema": 1, "metadata_sha256": {}}) + "\n",
	encoding="utf-8",
	)
	(submission / "manifest.json").write_text(
	json.dumps({"copied_hash": "source-hash"}) + "\n", encoding="utf-8"
	)
	(submission / f"runtime-verification-{job_id}.json").write_text(
	json.dumps(
	{
	"schema": "spatial_mmcot_runtime_verification_v1",
	"slurm_job_id": job_id,
	"source_hash": "source-hash",
	"ready_sha256": hashlib.sha256(ready.read_bytes()).hexdigest(),
	}
	)
	+ "\n",
	encoding="utf-8",
	)

	def test_claim_run_directory_creates_a_fresh_run(self):
	with tempfile.TemporaryDirectory() as directory:
	run = Path(directory) / "run"
	self.assertEqual(claim_run_directory(run), run)
	self.assertTrue(run.is_dir())
	self.assertFalse((run / ".run-claim.json").exists())

	def test_claim_run_directory_accepts_only_a_ready_submission(self):
	with tempfile.TemporaryDirectory() as directory, patch.dict(
	"os.environ", {"SLURM_JOB_ID": "123"}, clear=False
	):
	run = Path(directory) / "run"
	self._prepare_submission(run)
	(run / "slurm-123.out").touch()
	claim_run_directory(run)
	payload = (run / ".run-claim.json").read_text(encoding="utf-8")
	self.assertIn('"slurm_job_id": "123"', payload)
	with self.assertRaisesRegex(RuntimeError, "already claimed"):
	claim_run_directory(run)
	self.assertEqual(claim_run_directory(run, resume=True), run)

	def test_claim_run_directory_rejects_unready_or_dirty_directories(self):
	with tempfile.TemporaryDirectory() as directory:
	root = Path(directory)
	unready = root / "unready"
	unready.mkdir()
	with self.assertRaisesRegex(FileExistsError, "lacks"):
	claim_run_directory(unready)
	dirty = root / "dirty"
	with patch.dict("os.environ", {"SLURM_JOB_ID": "123"}, clear=False):
	self._prepare_submission(dirty)
	(dirty / "metrics.json").write_text("{}\n", encoding="utf-8")
	with (
	patch.dict("os.environ", {"SLURM_JOB_ID": "123"}, clear=False),
	self.assertRaisesRegex(FileExistsError, "experiment artifacts"),
	):
	claim_run_directory(dirty)

	def test_claim_run_directory_rejects_unattested_ready_marker(self):
	with tempfile.TemporaryDirectory() as directory, patch.dict(
	"os.environ", {"SLURM_JOB_ID": "123"}, clear=False
	):
	run = Path(directory) / "run"
	(run / "_submission").mkdir(parents=True)
	(run / "_submission/READY").write_text("ready\n", encoding="utf-8")
	with self.assertRaisesRegex(FileExistsError, "invalid submission"):
	claim_run_directory(run)

	def test_model_first_run_directory_is_repository_relative(self):
	config = SFTConfig()
	config.run.dataset = "zebra-cot-small"
	config.run.category = "sweeps"
	config.run.name = "lambda_grid"
	with tempfile.TemporaryDirectory() as directory, patch(
	"os.getcwd", return_value=directory
	):
	run = resolve_run_directory(config)
	repository = Path(__file__).resolve().parents[2]
	self.assertEqual(run.parents[7], repository)
	self.assertEqual(
	run.relative_to(repository / "experiments/output").parts[:5],
	("uniar-sft", "zebra-cot-small", "unversioned", "sweeps", "sft"),
	)

	def test_probe_override_routes_to_smoke(self):
	config = SFTConfig()
	config.run.dataset = "zebra-cot-small"
	run = resolve_run_directory(config, "sft_probe", category="smoke")
	self.assertEqual(run.parent.name, "sft_probe")
	self.assertEqual(run.parent.parent.name, "smoke")

	def test_versioned_run_directory_places_version_after_dataset(self):
	config = SFTConfig()
	config.run.dataset = "zebra-cot-rpm"
	config.run.version = "v1"
	config.run.category = "sweeps"
	config.run.name = "lambda_grid"
	run = resolve_run_directory(config)
	repository = Path(__file__).resolve().parents[2]
	self.assertEqual(
	run.relative_to(repository / "experiments/output").parts[:5],
	("uniar-sft", "zebra-cot-rpm", "v1", "sweeps", "sft"),
	)

	def test_queue_physical_gpu_maps_to_visible_logical_ordinal(self):
	with patch.dict(
	"os.environ",
	{"GPUQ_ASSIGNED_ID": "5", "CUDA_VISIBLE_DEVICES": "4,5"},
	clear=False,
	):
	self.assertEqual(resolve_cuda_device_index(local_rank=0, world_size=1), 1)
	with (
	patch.dict(
	"os.environ",
	{"GPUQ_ASSIGNED_ID": "5", "CUDA_VISIBLE_DEVICES": "GPU-uuid"},
	clear=False,
	),
	patch("torch.cuda.device_count", return_value=1),
	):
	self.assertEqual(resolve_cuda_device_index(local_rank=0, world_size=1), 0)
	self.assertEqual(resolve_cuda_device_index(local_rank=3, world_size=4), 3)

	def test_cache_fingerprint_does_not_hash_tensor_shard_contents(self):
	with tempfile.TemporaryDirectory() as directory:
	root = Path(directory)
	(root / "READY").write_text("ready\n", encoding="utf-8")
	(root / "manifest.jsonl").write_text("{}\n", encoding="utf-8")
	(root / "shard.pt").write_bytes(b"tensor-payload")
	with patch("shared.artifacts.file_sha256", wraps=file_sha256) as hashed:
	cache_fingerprint(root, [{"cache_shard": "shard.pt"}])
	hashed_paths = {call.args[0].name for call in hashed.call_args_list}
	self.assertEqual(hashed_paths, {"READY", "manifest.jsonl"})

	def test_atomic_checkpoint_restores_world_one_rng_and_state(self):
	with (
	tempfile.TemporaryDirectory() as directory,
	patch("torch.cuda.is_available", return_value=False),
	):
	root = Path(directory)
	checkpoints = root / "checkpoints"
	checkpoints.mkdir()
	parameter = torch.nn.Parameter(torch.tensor(1.0))
	optimizer = torch.optim.AdamW([parameter], lr=1e-3)
	scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
	config = SFTConfig()
	trainer_state = {
	"attempted_updates": 3,
	"successful_updates": 2,
	"nonfinite_updates": 1,
	"microbatch_cursor": 12,
	}
	random.seed(9)
	np.random.seed(9)
	torch.manual_seed(9)
	checkpoint = save_checkpoint(
	checkpoints_dir=checkpoints,
	model=_SavedObject(),
	processor=_SavedObject(),
	optimizer=optimizer,
	scheduler=scheduler,
	trainer_state=trainer_state,
	config=config,
	provenance_fingerprints={
	"cache": "cache-v1",
	"model": "model-v1",
	"source": "source-v1",
	},
	rank=0,
	world_size=1,
	ema_state={"schema": "fixture", "update_count": 2},
	)
	expected = (random.random(), float(np.random.rand()), float(torch.rand(())))
	for _ in range(5):
	random.random()
	np.random.rand()
	torch.rand(())
	restored_ema = []
	restored = load_training_state(
	checkpoint,
	optimizer=optimizer,
	scheduler=scheduler,
	config=config,
	provenance_fingerprints={
	"cache": "cache-v1",
	"model": "model-v1",
	"source": "source-v1",
	},
	rank=0,
	world_size=1,
	ema_state_loader=restored_ema.append,
	)
	actual = (random.random(), float(np.random.rand()), float(torch.rand(())))
	self.assertEqual(restored, trainer_state)
	self.assertEqual(
	restored_ema, [{"schema": "fixture", "update_count": 2}]
	)
	self.assertEqual(actual, expected)
	self.assertTrue((checkpoint / "COMPLETE").is_file())
	self.assertFalse(any(checkpoints.glob(".tmp-*")))

	def test_parent_linked_ddp_execution_change_can_restore_exact_state(self):
	with (
	tempfile.TemporaryDirectory() as directory,
	patch("torch.cuda.is_available", return_value=False),
	):
	root = Path(directory)
	parent = root / "parent"
	checkpoints = parent / "checkpoints"
	checkpoints.mkdir(parents=True)
	parameter = torch.nn.Parameter(torch.tensor(1.0))
	optimizer = torch.optim.AdamW([parameter], lr=1e-3)
	scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0)
	original = SFTConfig()
	checkpoint = save_checkpoint(
	checkpoints_dir=checkpoints,
	model=_SavedObject(),
	processor=_SavedObject(),
	optimizer=optimizer,
	scheduler=scheduler,
	trainer_state={"successful_updates": 3},
	config=original,
	provenance_fingerprints={"cache": "c", "model": "m", "source": "old"},
	rank=0,
	world_size=1,
	)
	(parent / "config.json").write_text(
	json.dumps(
	{
	"provenance_fingerprints": {
	"cache": "c",
	"model": "m",
	"source": "old",
	}
	}
	),
	encoding="utf-8",
	)
	child = root / "child"
	(child / "_submission").mkdir(parents=True)
	(parent / "_submission").mkdir(parents=True)
	(parent / "_submission/manifest.json").write_text(
	json.dumps({"copied_hash": "parent-source"}), encoding="utf-8"
	)
	(child / "_submission/submission.json").write_text(
	json.dumps(
	{
	"parent_run": {
	"run_dir": str(parent),
	"source_hash": "parent-source",
	}
	}
	),
	encoding="utf-8",
	)
	migrated = copy.deepcopy(original)
	migrated.trainer.ddp_find_unused_parameters = False
	migrated.checkpoint.resume_mode = "ddp_execution_change"
	restored = load_training_state(
	checkpoint,
	optimizer=optimizer,
	scheduler=scheduler,
	config=migrated,
	provenance_fingerprints={"cache": "c", "model": "m", "source": "new"},
	rank=0,
	world_size=1,
	run_dir=child,
	)
	self.assertEqual(restored, {"successful_updates": 3})

	control_migration = copy.deepcopy(original)
	control_migration.checkpoint.resume_mode = (
	"distributed_execution_change"
	)
	restored = load_training_state(
	checkpoint,
	optimizer=optimizer,
	scheduler=scheduler,
	config=control_migration,
	provenance_fingerprints={
	"cache": "c",
	"model": "m",
	"source": "new",
	},
	rank=0,
	world_size=1,
	run_dir=child,
	)
	self.assertEqual(restored, {"successful_updates": 3})

	migrated.objective.text.weight = 2.0
	with self.assertRaisesRegex(ValueError, "non-DDP configuration"):
	load_training_state(
	checkpoint,
	optimizer=optimizer,
	scheduler=scheduler,
	config=migrated,
	provenance_fingerprints={"cache": "c", "model": "m", "source": "new"},
	rank=0,
	world_size=1,
	run_dir=child,
	)

	def test_parent_linked_epoch_extension_restores_final_state(self):
	with (
	tempfile.TemporaryDirectory() as directory,
	patch("torch.cuda.is_available", return_value=False),
	):
	root = Path(directory)
	parent = root / "parent"
	checkpoints = parent / "checkpoints"
	checkpoints.mkdir(parents=True)
	parameter = torch.nn.Parameter(torch.tensor(1.0))
	optimizer = torch.optim.AdamW([parameter], lr=1e-3)
	scheduler = torch.optim.lr_scheduler.LambdaLR(
	optimizer, lambda _: 1.0
	)
	original = SFTConfig()
	original.run.name = "two_epoch"
	original.scheduler.kind = "constant"
	original.trainer.epochs = 2
	state = {
	"attempted_updates": 20,
	"successful_updates": 20,
	"nonfinite_updates": 0,
	"microbatch_cursor": 160,
	}
	checkpoint = save_checkpoint(
	checkpoints_dir=checkpoints,
	model=_SavedObject(),
	processor=_SavedObject(),
	optimizer=optimizer,
	scheduler=scheduler,
	trainer_state=state,
	config=original,
	provenance_fingerprints={
	"cache": "c",
	"model": "m",
	"source": "old",
	},
	rank=0,
	world_size=1,
	)
	(parent / "config.json").write_text(
	json.dumps(
	{
	"provenance_fingerprints": {
	"cache": "c",
	"model": "m",
	"source": "old",
	}
	}
	),
	encoding="utf-8",
	)
	(parent / "summary.json").write_text(
	json.dumps({"successful_updates": 20, "target_updates": 20}),
	encoding="utf-8",
	)
	(parent / "_submission").mkdir()
	(parent / "_submission/manifest.json").write_text(
	json.dumps({"copied_hash": "parent-source"}), encoding="utf-8"
	)
	child = root / "child"
	(child / "_submission").mkdir(parents=True)
	(child / "_submission/submission.json").write_text(
	json.dumps(
	{
	"parent_run": {
	"run_dir": str(parent),
	"source_hash": "parent-source",
	}
	}
	),
	encoding="utf-8",
	)
	extended = copy.deepcopy(original)
	extended.run.name = "three_epoch"
	extended.trainer.epochs = 3
	extended.checkpoint.resume_mode = "epoch_extension"

	restored = load_training_state(
	checkpoint,
	optimizer=optimizer,
	scheduler=scheduler,
	config=extended,
	provenance_fingerprints={
	"cache": "c",
	"model": "m",
	"source": "new",
	},
	rank=0,
	world_size=1,
	run_dir=child,
	)

	self.assertEqual(restored, state)

	if __name__ == "__main__":
	unittest.main()