diff --git a/pyproject.toml b/pyproject.toml index 7dbccd9..9ae9b79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "gen-worker" -version = "0.40.6" +version = "0.40.7" description = "A library used to build custom functions in Cozy Creator's serverless function platform." readme = "README.md" license = "MIT" diff --git a/src/gen_worker/compile_cache.py b/src/gen_worker/compile_cache.py index caa5913..acc1dac 100644 --- a/src/gen_worker/compile_cache.py +++ b/src/gen_worker/compile_cache.py @@ -505,7 +505,20 @@ def _disable_aot_autograd_cache() -> None: requires the FX cache to be the lookup surface: disable the AOT layer symmetrically for producer capture and consumer seeding. Costs a cheap AOT re-analysis per fresh process; the expensive inductor compile still - serves from the (portable) FX entries.""" + serves from the (portable) FX entries. + + LIVE DISPROOF of the 0.40.4/0.40.5 shape (2026-07-21, B200 pods, + gen-worker 0.40.5): the mint capture still packed 8 ASLR-keyed + ``aotautograd/`` entries and the store-served sibling still failed 8/8 + — because in torch 2.13 ``ConfigModule`` user overrides are a + ContextVar, i.e. THREAD-LOCAL: the assignment below ran on the arming + thread while the warmup compile ran on another thread that still saw + the default True. The env var is no rescue post-import + (``env_name_force`` is read once at config install). Process-global + disable therefore needs BOTH: the pre-torch-import env in the + entrypoint (fresh processes, incl. compile-worker subprocesses) and + the installed config entry's ``env_value_force`` mutated here (torch + already imported — tools, tests, embedders).""" os.environ["TORCHINDUCTOR_AUTOGRAD_CACHE"] = "0" import sys @@ -514,7 +527,11 @@ def _disable_aot_autograd_cache() -> None: try: import torch._functorch.config as fconf - fconf.enable_autograd_cache = False + fconf.enable_autograd_cache = False # this thread (fast path, public API) + # Process-global: user overrides are thread-local ContextVars in + # torch>=2.13; the entry-level env force is consulted by every + # thread with top precedence. + fconf._config["enable_autograd_cache"].env_value_force = False # type: ignore[attr-defined] except Exception: logger.debug("compile-cache: AOT autograd cache disable unavailable", exc_info=True) diff --git a/src/gen_worker/entrypoint.py b/src/gen_worker/entrypoint.py index 46555ce..0b4ac86 100644 --- a/src/gen_worker/entrypoint.py +++ b/src/gen_worker/entrypoint.py @@ -23,6 +23,15 @@ # https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") +# gw#608: compile-cell portability requires the (portable) FxGraphCache to be +# the lookup surface — the AOTAutogradCache key embeds a process memory +# address (ASLR) and can never hit across pods. TORCHINDUCTOR_AUTOGRAD_CACHE +# is an `env_name_force` config read ONCE at torch import, and runtime config +# assignments are thread-local (ContextVar) in torch>=2.13, so this must be +# set here, before any torch import, to bind every thread and every +# compile-worker subprocess. See compile_cache._disable_aot_autograd_cache. +os.environ.setdefault("TORCHINDUCTOR_AUTOGRAD_CACHE", "0") + import json import logging import sys diff --git a/tests/test_compile_cache.py b/tests/test_compile_cache.py index e4dd9a0..65b9aa7 100644 --- a/tests/test_compile_cache.py +++ b/tests/test_compile_cache.py @@ -1267,3 +1267,43 @@ class _Cfg: # must already have happened (it precedes every compile decision). cc.apply(_P(), _Cfg(), cache_ready=False) assert fconf.enable_autograd_cache is False + + +def test_aot_autograd_cache_disabled_across_threads(monkeypatch, tmp_path): + """gw#608 residual (live-disproven 0.40.5, B200, 2026-07-21): torch>=2.13 + config user overrides are ContextVars — THREAD-LOCAL. The arming thread's + ``enable_autograd_cache = False`` was invisible to the warmup thread that + actually compiled, so the mint still packed ASLR-keyed AOT entries and the + store-served sibling still missed 8/8. The disable must bind EVERY thread + (entry-level env force), not just the caller's.""" + import threading + + import torch._functorch.config as fconf + + entry = fconf._config["enable_autograd_cache"] + had_force = "env_value_force" in entry.__dict__ + old_force = entry.__dict__.get("env_value_force") + monkeypatch.delenv("TORCHINDUCTOR_AUTOGRAD_CACHE", raising=False) + monkeypatch.setattr(fconf, "enable_autograd_cache", True) + try: + entry.__dict__.pop("env_value_force", None) # simulate no pre-import env + cc.capture_env(tmp_path / "cap") + + seen: dict = {} + + def probe() -> None: + seen["value"] = fconf.enable_autograd_cache + + t = threading.Thread(target=probe) + t.start() + t.join() + assert seen["value"] is False, ( + "AOT autograd cache still enabled on a sibling thread — the " + "warmup/compile thread would repack ASLR-keyed AOT entries and " + "consumers would miss 8/8 (gw#608)" + ) + finally: + if had_force: + entry.env_value_force = old_force + else: + entry.__dict__.pop("env_value_force", None) diff --git a/uv.lock b/uv.lock index 0abb88b..7395301 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "gen-worker" -version = "0.40.6" +version = "0.40.7" source = { editable = "." } dependencies = [ { name = "blake3" },