feat: configure Torch Dynamo recompile limit from environment (#7243)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

docs/configuration/environment_variables.md

@@ -130,6 +130,22 @@ settings.
 Omni-owned variable, even though vLLM-Omni supplies a persistent default when
 it is unset.
 
+### Torch compilation
+
+| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
+| --- | --- | --- | --- | --- |
+| `VLLM_OMNI_TORCH_DYNAMO_RECOMPILE_LIMIT` | Positive integer; unset preserves the current Torch setting | Process default for Torch Dynamo; read when `vllm_omni` is imported | Sets `torch._dynamo.config.recompile_limit` to the requested value. A non-positive or non-integer value raises `ValueError`. Later backend-specific overrides still take precedence. | Diagnostic |
+
+Set this variable before launching Omni, for example:
+
+```bash
+export VLLM_OMNI_TORCH_DYNAMO_RECOMPILE_LIMIT=64
+```
+
+The setting applies across model families and platforms. vLLM compilation
+contexts and backends can subsequently apply their own limits; this variable
+does not replace those backend-specific policies.
+
 ## Per-stage environment
 
 Deploy configurations can set arbitrary environment keys for one stage:

tests/config/test_env_override.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_IMPORT_SCRIPT = """
+import sys
+import types
+import torch._dynamo.config as dynamo_config
+
+# Isolate package startup from unrelated vLLM patches and model registration.
+version = types.ModuleType("vllm_omni.version")
+version.__version__ = "dev"
+version.__version_tuple__ = (0, 0, "dev")
+sys.modules[version.__name__] = version
+for name in ("vllm_omni.patch", "vllm_omni.transformers_utils.configs",
+             "vllm_omni.transformers_utils.parsers"):
+    sys.modules[name] = types.ModuleType(name)
+config = types.ModuleType("vllm_omni.config")
+config.OmniModelConfig = object
+sys.modules[config.__name__] = config
+
+dynamo_config.recompile_limit = 11
+import vllm_omni
+assert dynamo_config.recompile_limit == int(sys.argv[1])
+"""
+
+
+@pytest.mark.parametrize("value", [None, "2", "64", "0", "-1", "not-an-int"])
+def test_package_import_configures_recompile_limit(value: str | None) -> None:
+    env = os.environ.copy()
+    name = "VLLM_OMNI_TORCH_DYNAMO_RECOMPILE_LIMIT"
+    env.pop(name, None)
+    if value is not None:
+        env[name] = value
+    result = subprocess.run(
+        [sys.executable, "-c", _IMPORT_SCRIPT, value or "11"],
+        cwd=Path(__file__).resolve().parents[2],
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+
+    if value in {"0", "-1", "not-an-int"}:
+        assert result.returncode != 0
+        assert f"ValueError: {name} must be" in result.stderr
+    else:
+        assert result.returncode == 0, result.stderr

tests/config/test_environment_variables.py

@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     """Make an inventory expansion an explicit review decision."""
     category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
     assert category_counts == {
-        EnvironmentVariableCategory.PUBLIC_OMNI: 30,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 31,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 58,

vllm_omni/__init__.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 vLLM-Omni: Multi-modality models inference and serving with
 non-autoregressive structures.
@@ -38,6 +41,10 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 
 from .config import OmniModelConfig
 
+# Apply the process-wide default at vllm_omni import; later backend policies
+# take precedence.
+from . import env_override as _env_override  # isort:skip # noqa: F401, E402
+
 
 def __getattr__(name: str):
     # Lazy import for AsyncOmni and Omni to avoid pulling in heavy

vllm_omni/config/environment_variable_inventory.py

@@ -80,6 +80,7 @@ _PUBLIC_OMNI = (
     "VLLM_OMNI_SPEAKER_REGISTRATION_POLICY",
     "VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
     "VLLM_OMNI_STORAGE_PATH",
+    "VLLM_OMNI_TORCH_DYNAMO_RECOMPILE_LIMIT",
     "VLLM_OMNI_USE_QUACK_FP8",
     "VLLM_OMNI_VIDEO_SYNC_TIMEOUT",
     "VLLM_VIDEO_ASYNC_CHUNK",

vllm_omni/env_override.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Apply process-wide Torch settings configured through environment variables."""
+
+from __future__ import annotations
+
+import os
+
+_TORCH_DYNAMO_RECOMPILE_LIMIT = "VLLM_OMNI_TORCH_DYNAMO_RECOMPILE_LIMIT"
+
+
+def configure_torch_dynamo() -> None:
+    """Set Torch Dynamo's recompile limit when explicitly requested."""
+    value = os.environ.get(_TORCH_DYNAMO_RECOMPILE_LIMIT)
+    if value is None:
+        return
+
+    try:
+        recompile_limit = int(value)
+    except ValueError as exc:
+        raise ValueError(f"{_TORCH_DYNAMO_RECOMPILE_LIMIT} must be a positive integer, got {value!r}") from exc
+    if recompile_limit <= 0:
+        raise ValueError(f"{_TORCH_DYNAMO_RECOMPILE_LIMIT} must be positive, got {value!r}")
+
+    import torch._dynamo.config as dynamo_config
+
+    dynamo_config.recompile_limit = recompile_limit
+
+
+configure_torch_dynamo()