Function check_state_structure

axlearn/common/checkpointer.py:127–178

(
    ckpt_structure: list[tuple[str, Any]],
    target_structure: list[tuple[str, Any]],
    *,
    validation: CheckpointValidationType = CheckpointValidationType.EXACT,
)

def check_state_structure(
    ckpt_structure: list[tuple[str, Any]],
    target_structure: list[tuple[str, Any]],
    *,
    validation: CheckpointValidationType = CheckpointValidationType.EXACT,
):
    # Maybe filter structure before comparison.
    def filter_for_validation(structure):
        filtered_structure = []
        for key, value in structure:
            if validation in [
                CheckpointValidationType.EXACT_UP_TO_DTYPE,
                CheckpointValidationType.CONTAINS_STATE_UP_TO_DTYPE,
            ] and isinstance(value, dict):
                # Drop dtype if it's in the value.
                value = {k: v for k, v in value.items() if k != "dtype"}
            filtered_structure.append((key, value))
        return filtered_structure

    filtered_ckpt_structure = sorted(
        [f"{key}={value}" for key, value in filter_for_validation(ckpt_structure)]
    )
    filtered_target_structure = sorted(
        [f"{key}={value}" for key, value in filter_for_validation(target_structure)]
    )

    msg = ""
    if validation in [CheckpointValidationType.EXACT, CheckpointValidationType.EXACT_UP_TO_DTYPE]:
        is_compatible = filtered_ckpt_structure == filtered_target_structure
    elif validation in [
        CheckpointValidationType.CONTAINS_STATE,
        CheckpointValidationType.CONTAINS_STATE_UP_TO_DTYPE,
    ]:
        # Allow checkpoint to contain additional information.
        is_compatible = set(filtered_ckpt_structure) >= set(filtered_target_structure)
        if not is_compatible:
            msg = (
                "Missing:\n"
                + "\n".join(sorted(set(filtered_target_structure) - set(filtered_ckpt_structure)))
                + "\n"
            )
    else:
        raise ValueError(f"Unknown validation type: {validation}")
    if not is_compatible:
        msg += "Diff:\n" + "\n".join(
            difflib.ndiff(sorted(filtered_ckpt_structure), sorted(filtered_target_structure))
        )
        raise ValueError(
            f"Unable to restore checkpoint ({validation}). A mismatch between the saved "
            "checkpoint tree dtypes or shapes and the current one has been detected:\n"
            f"{msg}"
        )


def _upload_dir(src_dir_handle: tempfile.TemporaryDirectory, *, dst_dir: str):