Add sys.modules registration - #105
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates MaxKernel/evaluation/harness_code.py to register a module in sys.modules before executing it, preventing failures during string annotation resolution. Feedback highlights a potential issue where unconditionally popping the module on failure could delete a pre-existing module, and suggests capturing and restoring the previous module state instead.
| sys.modules[module_name] = module | ||
| try: | ||
| spec.loader.exec_module(module) | ||
| except Exception: | ||
| sys.modules.pop(module_name, None) | ||
| raise |
There was a problem hiding this comment.
If module_name was already present in sys.modules before calling load_module_from_path, unconditionally popping it on failure will delete the pre-existing module instead of restoring it. It is safer to capture the previous value of sys.modules[module_name] and restore it if execution fails.
| sys.modules[module_name] = module | |
| try: | |
| spec.loader.exec_module(module) | |
| except Exception: | |
| sys.modules.pop(module_name, None) | |
| raise | |
| old_module = sys.modules.get(module_name)\n sys.modules[module_name] = module\n try:\n spec.loader.exec_module(module)\n except Exception:\n if old_module is None:\n sys.modules.pop(module_name, None)\n else:\n sys.modules[module_name] = old_module\n raise |
There was a problem hiding this comment.
If it gets to Exception, it will fail anyway.
| return avg_wall_time, xprof_time | ||
|
|
||
|
|
||
| def diff_metrics(b, o, chunk_elems=1 << 24): |
There was a problem hiding this comment.
Is this function added because the code can run into OOM during the diff calculation? It might be beneficial if we can make this comment easier to understand.
There was a problem hiding this comment.
Config 3 (csa_decode_bs512) has a uint8 cache of (32769, 256, 4, 128) = 4 GiB. b and o are host numpy arrays by this point, so (b - o) / b is a numpy true division → promotes to float64 on the host. Then jnp.abs(...) ships it to the TPU as f32: a 16 GiB argument plus a 16 GiB result = 32.00 GiB against a 31.25 GiB chip, overflowing by 773 M.
| @@ -266,8 +315,9 @@ def main(): | |||
| if b.shape != o.shape: | |||
| raise ValueError(f"Shape mismatch: {b.shape} vs {o.shape}") | |||
| is_correct = is_correct and bool(jnp.allclose(b, o, atol=curr_atol, rtol=curr_rtol)) | |||
There was a problem hiding this comment.
If the diff can make OOM, why this line will not make OOM?
There was a problem hiding this comment.
OOM happens in abs calculation. allclose survives because it's jitted and XLA fuses the promotion
2 participants
from future import annotations + a dataclass would fail without this change.