Checkpoints and recovery

import json
from pathlib import Path


job_id = training_client.get_info().model_data.model_id
state = training_client.save_state(f"step-{step:06d}").result()
record = {
    "job_id": job_id,
    "state_path": state.path,
    "model": model_name,
    "dataset_revision": dataset_revision,
    "optimizer": optimizer_config,
    "step": step,
    "dataset_cursor": dataset_cursor,
    "reward_version": reward_version,
}

with Path(log_path).open("a", encoding="utf-8") as checkpoint_log:
    checkpoint_log.write(json.dumps(record, sort_keys=True) + "\n")

def get_last_checkpoint(log_path: str) -> dict[str, object] | None:
    path = Path(log_path)
    if not path.exists():
        return None

    for line in reversed(path.read_text(encoding="utf-8").splitlines()):
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(record, dict) and isinstance(record.get("state_path"), str):
            return record
    return None

training_client = (
    service_client.create_training_client_from_state_with_optimizer(
        str(resume["state_path"])
    )
)

state = training_client.save_state("step-000100").result()
sampler = training_client.save_weights_for_sampler(
    "step-000100"
).result()