fix: apply dynamo config on GPU worker thread (#167)

[gitcommit90]

Summary

Apply torch._dynamo.config on the GPU worker thread so torch 2.13+ ContextVar settings are not silently ignored.

Changes

• Extract RECOMPILE_LIMIT + apply_dynamo_config() and call at import and as GPU executor initializer
• Pass recompile_limit into torch.compile(...) call sites
• Focused regression coverage for thread-local dynamo config

Closes #167

[NSagan271]

RECOMPILE_LIMIT should be able to be overridden as an environment variable, with reasonable minimum and maximum limits set in the code (and then documented in docs/environment_variables.rst)

[NSagan271]

Heads up: a bare RECOMPILE_LIMIT constant fixes today's config set but stays fragile. Dynamo keys the recompile counter on the underlying code object, not the torch.compile wrapper, so in the cuda graph setting, every slot, every batch size, every config sharing a capture_forward_method, plus the eager path that compiles the same forward_batched, all accumulate onto one counter.

With dynamic=False (which we use to get maximum performance out of the cuda graph path), each static shape is one specialization, so the entries on that counter = sum(batch_sizes × token_buckets) over those configs + eager runtime shapes. A fixed constant has no relationship to that sum, so adding a batch size later can silently push it over, resulting eager fallback mid-capture. Right now, the RECOMPILE_LIMIT is high enough that we don't see that in practice

Suggest deriving it instead: recompile_limit = min(max(FLOOR, capture_buckets_for_this_code_object + EAGER_HEADROOM), CEIL) (with FLOOR, CEIL, and EAGER_HEADROOM being the knobs rather than RECOMPILE_LIMIT. I think the easiest way to get "capture_buckets_for_this_code_object" is to have a function on the base submodule that computes it based on the cuda graph configs.

[NSagan271]

Actually, I realize that my suggested fix in #167 of setting recompile_limit on torch.compile only works for torch 2.13; before then, the only lever for setting the recompile limit is the global one. I think in that case, since we only pin torch>=2.9.1 in pyproject.toml, we should rely on the global lever torch._dynamo.config.recompile_limit only for now (instead of, e.g., including the kwarg contingent on the right value of torch; we can change the default behavior down the road after pinning a newer torch).

With respect to my comment https://github.com/mstar-project/mstar/pull/178/changes#r3612126831 about computing the recompile limit based on the cuda graph configurations, I think the best way to proceed is to put apply_dynamo_config onto EngineManager, with the following additional changes needed to make it work:

1. Have EngineManager now hold the node_submodules: dict[str, NodeSubmodule] mapping computed when building the engine manager. node_submodules is already computed as a local in build() today, so this is just promoting an existing local to a field, not new plumbing.
2. Add a function on the engine manager to compute the recompile limit (as the maximum recompile limit across submodules). EngineManager.apply_dynamo_config will call this function.
3. Call self.engine_manager.apply_dynamo_config in worker.pyafter EngineManager.build in worker.py, and also use that function as the initializer for the thread pool executor. Note that both call sites are needed: the former because the warmup (cuda graph capture) occurs on the main thread, the latter for request execution in the async worker.

[NSagan271]

Although torch.set_float32_matmul_precision("high") doesn't need to be set per-thread, I still think it might be cleaner (and more future-proof) to group it in with the apply_dynamo_config function (possibly giving the function a different name if that makes sense to do).