xai-org/x-algorithm · critical · RuntimeError

Orbax version is too old

Error message

Orbax version is too old

What it means

get_checkpointer builds an Orbax AsyncCheckpointer and requires the private _post_finalization_callback attribute to exist; older Orbax releases lack it, so the code refuses to run rather than silently skipping finalization callbacks. It's a hard version gate on the orbax-checkpoint dependency.

Source

Thrown at phoenix/xrex/utils/checkpointing.py:97

    if _CHECKPOINTER is None:
        handler_kwargs: dict[str, Any] = {"use_zarr3": True, "restore_concurrent_gb": 1}
        if save_concurrent_gb is not None:
            handler_kwargs["save_concurrent_gb"] = save_concurrent_gb
            handler_kwargs["restore_concurrent_gb"] = save_concurrent_gb
        _CHECKPOINTER_SAVE_CONCURRENT_GB = save_concurrent_gb
        rank_logger.info(
            "Creating AsyncCheckpointer: PyTreeCheckpointHandler(use_zarr3=True, "
            "save_concurrent_gb=%s, restore_concurrent_gb=%s) "
            "[None => Orbax write-limiter default 96GB; D2H still all-at-once "
            "unless save_checkpoint registers ThrottledD2HArrayHandler]",
            save_concurrent_gb,
            save_concurrent_gb,
        )
        _CHECKPOINTER = ocp.AsyncCheckpointer(
            ocp.PyTreeCheckpointHandler(**handler_kwargs), timeout_secs
        )
        if not hasattr(_CHECKPOINTER, "_post_finalization_callback"):
            raise RuntimeError("Orbax version is too old")
    elif save_concurrent_gb is not None and _CHECKPOINTER_SAVE_CONCURRENT_GB != save_concurrent_gb:
        rank_logger.warning(
            "get_checkpointer(save_concurrent_gb=%s) ignored; checkpointer already "
            "created with save_concurrent_gb=%s. D2H throttle (if enabled) still "
            "uses the value passed to save_checkpoint.",
            save_concurrent_gb,
            _CHECKPOINTER_SAVE_CONCURRENT_GB,
        )
    return _CHECKPOINTER


class NoCompressionArrayHandler(ocp.type_handlers.ArrayHandler):
    def _get_json_tspec_write(self, *args, **kwargs):
        spec = super()._get_json_tspec_write(*args, **kwargs)
        for codec in spec["metadata"]["codecs"]:
            cfg = codec["configuration"]
            cfg["codecs"] = [c for c in cfg["codecs"] if c["name"] != "zstd"]
        return spec

View on GitHub (pinned to 24c60942c5)

Solutions

1. Upgrade orbax-checkpoint to the version pinned in the repo's requirements/pyproject (one that has _post_finalization_callback)
2. Verify with: python -c "import orbax.checkpoint as ocp; print(ocp.__version__)" and compare against the pin
3. Recreate the environment from the project's lockfile

Example fix

# before
pip install orbax-checkpoint  # resolves to old version
# after
pip install 'orbax-checkpoint>=0.5.0'  # or the repo's pinned version

import orbax.checkpoint as ocp
c = ocp.AsyncCheckpointer(ocp.PyTreeCheckpointHandler())
assert hasattr(c, '_post_finalization_callback'), 'orbax too old; upgrade'

When it happens

Trigger: Installing a version of orbax-checkpoint older than the one this repo was developed against, then calling save_checkpoint (which calls get_checkpointer).

Common situations: Fresh environment resolving to an old orbax from a transitive dependency; pip install orbax-checkpoint without a lower bound; environment drift after a base-image update.

Related errors

• concurrent_bytes must be > 0, got {concurrent_bytes}
• Unknown remat policy: {policy}

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/50be67d0ee643d6c. Report an issue: GitHub.