fix: critical bugs — unsafe torch.load (RCE), UnboundLocalError crash, thread leak, broken tests

[Sunil56224972]

Summary

Deep manual code audit of the entire airllm codebase found 5 real bugs — a security vulnerability, crash bugs, a resource leak, and broken tests.

---

Bug 1 — 🔴 CRITICAL: Arbitrary Code Execution via Unsafe orch.load

File: �ir_llm/airllm/utils.py (lines 273, 434, 448)

orch.load() without weights_only=True uses Python's pickle module, which allows arbitrary code execution. A malicious .bin checkpoint can execute arbitrary Python when loaded — installing malware, exfiltrating data, or creating a reverse shell.

Fix: Added weights_only=True to all 3 orch.load() call sites.

Ref: https://pytorch.org/docs/stable/generated/torch.load.html

---

Bug 2 — 🔴 HIGH:

emove_real_and_linked_file() Crashes with UnboundLocalError

File: �ir_llm/airllm/utils.py (lines 198–204)

\\python
def remove_real_and_linked_file(to_delete):
if (os.path.realpath(to_delete) != to_delete):
targetpath = os.path.realpath(to_delete) # only assigned inside if
os.remove(to_delete)
if (targetpath): # ← UnboundLocalError when file is NOT a symlink
os.remove(targetpath)
\\

When the file is not a symlink, argetpath is never assigned, causing UnboundLocalError. This crashes the delete_original=True workflow (critical for 1.5TB+ models like Kimi K3).

Fix: Initialize argetpath = None before the conditional.

---

Bug 3 — 🟡 HIGH: ThreadPoolExecutor Resource Leak

File: �ir_llm/airllm/airllm_base.py (line 175)

The executor for prefetching is created but never shut down — no del, close(), or context manager. Each model instance leaks a daemon thread.

Fix: Added close() method and del safety net.

---

Bug 4 — 🟡 MEDIUM: Unit Tests Outdated After Refactor

File: �ir_llm/tests/test_automodel.py

Tests expect AirLLMLlama2, AirLLMMistral, AirLLMMixtral but these architectures are no longer in ARCH_OVERRIDES — they correctly fall through to AirLLMBaseModel. Tests fail on every run.

Fix: Updated expected class names to match current ARCH_OVERRIDES.

---

Bug 5 — 🟡 MEDIUM: clean_memory() Crashes on Non-CUDA Systems

File: �ir_llm/airllm/utils.py (line 83)

orch.cuda.empty_cache() called without checking 	orch.cuda.is_available(). Crashes on CPU-only systems used for checkpoint splitting.

Fix: Added orch.cuda.is_available() guard.

---

Files Changed

File	Changes
�ir_llm/airllm/utils.py	Bugs 1, 2, 5
�ir_llm/airllm/airllm_base.py	Bug 3
�ir_llm/tests/test_automodel.py	Bug 4

[Hotragn]

Heads up on Bug 2 — I ran your version against the real call site and it still crashes. Initialising targetpath = None removes the UnboundLocalError, but the double-delete is a separate failure and the realpath() != to_delete comparison is what causes it.

split_and_save_layers passes a Path, not a string:

to_delete = checkpoint_path / shard_num_to_file[shard]
remove_real_and_linked_file(to_delete)

os.path.realpath() returns a str, so str != Path is always true — even for an ordinary file. targetpath gets set to the same file under its resolved name, and the second os.remove() hits a file the first call already deleted.

I compared the four open variants of this helper on a temp file, using the exact Path call shape plus a relative-string call:

variant	Path object (real call site)	relative str
current main	FileNotFoundError	FileNotFoundError
this PR (#333)	FileNotFoundError	FileNotFoundError
#303 / #362 (os.path.islink)	PASS	PASS
#337 (islink + fspath + guards)	PASS	PASS

The fix is to ask the question you actually mean — is this a symlink whose target should also go? — rather than comparing paths:

targetpath = os.path.realpath(to_delete) if os.path.islink(to_delete) else None

Worth knowing there are now four PRs on this one helper: #303 (oldest, 2026-07-19), #337, #362, and this one. #303 and #362 are logically identical; #337 adds os.fspath normalisation and swallows FileNotFoundError on both removes.

Transparency: your Bug 1 (weights_only=True) overlaps my #305, which also reorders the loader to prefer safetensors over the pickle index. I'm not raising that to stake a claim — either fix closes the hole, and I'd rather the bug get fixed than mine get merged. Bug 2 is the one I'd genuinely change before this lands, since as written it would be reported again.

(I couldn't test the symlink case — this box won't let me create symlinks without elevation — so I can't speak to the HF-cache blob behaviour, which is where #337's extra guards would matter.)