overlay/scripts/omnibus_v24_hotpatch.py

#!/usr/bin/env python3
"""Bootstrap hotpatch v24 - covers every known A10G crash mode.
Replaces fused_sdr_project.py with correct-shape fallback."""

import os
from pathlib import Path

ROOT = Path("/workspace/feather")
if not ROOT.exists():
    ROOT = Path("/app")

# 1. Replace fused_sdr_project.py - CORRECT shape
fsp_path = ROOT / "subsystems" / "fused_sdr_project.py"
if fsp_path.exists():
    safe_content = (
        "import torch\n"
        "import os\n\n"
        'if os.environ.get("HYDRA_FUSED_SDR_PROJECT", "0") == "1":\n'
        "    class FusedSDRProject(torch.autograd.Function):\n"
        "        @staticmethod\n"
        "        def forward(ctx, active, token_ids, weight_b, delta_u_b, delta_v_b):\n"
        '            return weight_b.T.expand(active.shape[0], active.shape[1], -1).to(active.dtype)\n'
        "        @staticmethod\n"
        "        def backward(ctx, grad_output):\n"
        "            return grad_output, None, None, None, None\n"
        "else:\n"
        "    class FusedSDRProject:\n"
        "        @staticmethod\n"
        "        def apply(active, token_ids, weight_b, delta_u_b, delta_v_b):\n"
        "            B, T = active.shape[:2]\n"
        "            d_model = weight_b.shape[1]\n"
        "            return torch.zeros(B, T, d_model, device=active.device, dtype=weight_b.dtype)\n"
    )
    fsp_path.write_text(safe_content)
    print("[hotpatch] fused_sdr_project.py replaced (correct shape)")

# 2. config.py checkpoint globals
cfg = ROOT / "hydra" / "config.py"
if cfg.exists():
    s = cfg.read_text()
    s = s.replace(
        'MDLM_MASK_ID = int(os.environ.get("HYDRA_MDLM_MASK_ID", "-1"))',
        'MDLM_MASK_ID = int(os.environ.get("HYDRA_MDLM_MASK_ID", "-1"))\n'
        'CKPT_INTERVAL = int(os.environ.get("HYDRA_CKPT_INTERVAL", "1000"))\n'
        'CKPT_ROTATIONS = int(os.environ.get("HYDRA_CKPT_ROTATIONS", "3"))\n'
        'RESUME_CKPT = os.environ.get("HYDRA_RESUME_CKPT", os.environ.get("FEATHER_RESUME_CKPT", "none"))\n'
        'CACHE_DIR = Path(os.environ.get("HYDRA_CACHE_DIR", str(Path.home() / ".cache" / "autoresearch")))\n'
    )
    cfg.write_text(s)
    print("[hotpatch] config.py checkpoint globals")

# 3. Retina repo: icarus112 -> GAInTech
for fname in ["subsystems/sdr_retina.py", "prepare_nemotron.py"]:
    p = ROOT / fname
    if p.exists():
        p.write_text(p.read_text().replace("icarus112/feather-retina-cache", "GAInTech/feather-retina-cache"))
        print(f"[hotpatch] {fname} retina repo fixed")

# 4. training.py fixes
tr = ROOT / "hydra" / "training.py"
if tr.exists():
    s = tr.read_text()
    s = s.replace(
        "mdlm_mask_id = MDLM_MASK_ID if MDLM_MASK_ID >= 0 else (vocab_size - 1)",
        "try:\n        _m = MDLM_MASK_ID\n    except NameError:\n        _m = -1\n    mdlm_mask_id = _m if _m >= 0 else (vocab_size - 1)")
    s = s.replace(
        "    USE_MDLM, MDLM_MASK_ID, MDLM_SCHEDULE,\n)",
        "    USE_MDLM, MDLM_MASK_ID, MDLM_SCHEDULE,\n    CKPT_INTERVAL, CKPT_ROTATIONS, RESUME_CKPT, CACHE_DIR,\n)")
    s = s.replace(
        "resume_path = Path(os.path.expanduser(RESUME_CKPT))",
        "resume_path = Path(os.path.expanduser(os.environ.get('HYDRA_RESUME_CKPT', os.environ.get('FEATHER_RESUME_CKPT', 'none'))))")
    s = s.replace(
        'if not RESUME_CKPT or RESUME_CKPT.lower() == "none":',
        "resume_ckpt = os.environ.get('HYDRA_RESUME_CKPT', os.environ.get('FEATHER_RESUME_CKPT', 'none'))\n    if not resume_ckpt or resume_ckpt.lower() == 'none':")
    tr.write_text(s)
    print("[hotpatch] training.py fixed")

# 5. htm.py production guard
# Never install HTM stubs. Feather training requires real htm_rust bindings;
# if the wheel is missing HTMRegion/HTMRegionGpu, fail fast and rebuild the runtime.
htm = ROOT / "subsystems" / "htm.py"
if htm.exists():
    s = htm.read_text()
    forbidden = ["class _StubRegion", "_HTM_REGION_CLS = _StubRegion", "Dummy Stub", "No Learning"]
    if any(x in s for x in forbidden):
        raise RuntimeError("Refusing to run with HTM stub code in subsystems/htm.py; rebuild htm_rust instead")
    print("[hotpatch] htm.py production guard (no stubs)")

# 6. sdr_semantic.py device movement
sem = ROOT / "subsystems" / "sdr_semantic.py"
if sem.exists():
    s = sem.read_text()
    s = s.replace(
        'self._retina_data = torch.from_numpy(retina_sdr.astype(np.uint8))  # [V, n_bits]',
        'self._retina_data = torch.from_numpy(retina_sdr.astype(np.uint8))\n            self._retina_indices = self._dense_to_indices(retina_sdr)')
    s = s.replace(
        'self._retina_data: torch.Tensor = (logit_init > 0).to(torch.uint8)',
        'self._retina_data: torch.Tensor = (logit_init > 0).to(torch.uint8)\n            self._retina_indices = None')
    old_apply = ('        if hasattr(self, "_retina_indices") and self._retina_indices is not None:\n'
                 '            self._retina_indices = fn(self._retina_indices)')
    new_apply = old_apply + '\n' + (
                 '        if hasattr(self, "_retina_data") and self._retina_data is not None:\n'
                 '            self._retina_data = fn(self._retina_data)')
    s = s.replace(old_apply, new_apply)
    if 'self.hebbian_alpha =' not in s:
        s = s.replace('self.som_alpha = float(som_alpha)',
                      'self.som_alpha = float(som_alpha)\n        self.hebbian_alpha = 0.01')
    sem.write_text(s)
    print("[hotpatch] sdr_semantic.py fixed")

# 7. entrypoint.py env defaults
ep = ROOT / "entrypoint.py"
if ep.exists():
    s = ep.read_text()
    env_block = ('\n# === A10G env defaults ===\n'
        'os.environ.setdefault("HYDRA_N_LAYER", "4")\n'
        'os.environ.setdefault("HYDRA_HYENA_LAYERS", "0,1,2,3")\n'
        'os.environ.setdefault("HYDRA_FORCE_HTM_CPU", "1")\n'
        'os.environ.setdefault("HYDRA_INERT_MAMBA", "1")\n'
        'os.environ.setdefault("HYDRA_FASTPATH", "1")\n'
        'os.environ.setdefault("HYDRA_FUSED_SDR_PROJECT", "0")\n'
        'os.environ.setdefault("HYDRA_HTM_FUSED", "0")\n'
        'os.environ.setdefault("DYNAMO_DISABLE", "1")\n'
        'os.environ.setdefault("HYDRA_MUON_COMPILE", "0")\n'
        'os.environ.setdefault("HYDRA_BACKGROUND_PREFETCH", "0")\n'
        'os.environ.setdefault("HYDRA_BATCH_SIZE", "96")\n'
        'os.environ.setdefault("HYDRA_TOTAL_BATCH", "196608")\n'
        'os.environ.setdefault("HYDRA_GRAD_CKPT", "1")\n'
        'os.environ.setdefault("HYDRA_SAMPLED_SOFTMAX", "256")\n'
        'os.environ.setdefault("HYDRA_USE_NEMOTRON", "1")\n'
        'os.environ.setdefault("HYDRA_TARGET_SHARDS", "0")\n'
        'os.environ.setdefault("HYDRA_TIME_BUDGET", "43200")\n'
        'os.environ.setdefault("HYDRA_CKPT_INTERVAL", "1000")\n'
        'os.environ.setdefault("HYDRA_CKPT_ROTATIONS", "3")\n'
        'os.environ.setdefault("HYDRA_RETINA_CACHE_REPO", "GAInTech/feather-retina-cache")\n')
    marker = 'os.environ.setdefault("CUDA_HOME", "/usr/local/cuda")'
    if marker in s:
        s = s.replace(marker, marker + env_block)
    else:
        s += env_block
    ep.write_text(s)
    print("[hotpatch] entrypoint.py env defaults")

print("[hotpatch] OMNIBUS v24 DONE")