Pretrained loading with typed OOM

Call model_class.from_pretrained(...) and convert CUDA OOM into the substrate’s typed CapabilityResourceError for CR-7 reactive retry.

load_pretrained_with_oom is the thin convenience wrapper over the from_pretrained convention (transformers, qwen_asr, …): it forwards every kwarg and converts only torch.cuda.OutOfMemoryError into a typed CapabilityResourceError (via cjm-substrate-torch-utils’ primitive). Every other exception propagates unchanged for the substrate’s CR-5 classifier.

Consumers whose load is a plain constructor (e.g. LavaSR’s LavaEnhance2(path)) should call cuda_oom_to_capability_resource_error directly in a try/except rather than this wrapper.

---

load_pretrained_with_oom

def load_pretrained_with_oom(
    model_class:Type, # A class exposing `.from_pretrained(repo_id, **kwargs)`
    repo_id:str, # Model repo id or local path passed to from_pretrained
    label:Optional=None, # OOM-message context (default: f"loading {repo_id!r}")
    kwargs:VAR_KEYWORD
)->Any: # The loaded model

Call model_class.from_pretrained(repo_id, **kwargs), converting CUDA OOM to a typed error.

On torch.cuda.OutOfMemoryError, re-raises as CapabilityResourceError (via cuda_oom_to_capability_resource_error) so the substrate’s CR-7 reactive-retry path can evict + reload + retry. All other exceptions propagate unchanged.

Wrap the call in self.heartbeat(...) at the capability site to cover the silent construction phase; pre-download with snapshot_download_with_progress for real download progress.

Tests. (Model class is mocked — no real model download.)

from cjm_substrate.core.errors import CapabilityResourceError


class _FakeModel:
    @classmethod
    def from_pretrained(cls, repo_id, **kwargs):
        return {"repo_id": repo_id, "kwargs": kwargs}


# Success path returns the model and forwards kwargs verbatim.
_m = load_pretrained_with_oom(_FakeModel, "org/m", device_map="cuda", dtype="bf16")
assert _m["repo_id"] == "org/m"
assert _m["kwargs"]["device_map"] == "cuda" and _m["kwargs"]["dtype"] == "bf16"


class _OOMModel:
    @classmethod
    def from_pretrained(cls, repo_id, **kwargs):
        raise torch.cuda.OutOfMemoryError("CUDA out of memory")


# CUDA OOM becomes a typed CapabilityResourceError with the cause preserved.
try:
    load_pretrained_with_oom(_OOMModel, "org/m")
    assert False, "expected CapabilityResourceError"
except CapabilityResourceError as e:
    assert e.resource_shortfall is not None
    assert e.resource_shortfall.resource == "gpu_vram_mb"
    assert isinstance(e.__cause__, torch.cuda.OutOfMemoryError)


class _ValueErrModel:
    @classmethod
    def from_pretrained(cls, repo_id, **kwargs):
        raise ValueError("bad config")


# Non-OOM errors propagate unchanged (not converted to CapabilityResourceError).
try:
    load_pretrained_with_oom(_ValueErrModel, "org/m")
    assert False, "expected ValueError"
except CapabilityResourceError:
    assert False, "ValueError must not convert to CapabilityResourceError"
except ValueError:
    pass