From c8f438c10a64ad660b198ae857345ba65866fbaf Mon Sep 17 00:00:00 2001
From: u7-u7 <116776533+u7-u7@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:40:18 +0800
Subject: [PATCH] fix(environ): reject unsupported DEEPINFER_ACCELERATOR values

The variable picks the deepinfer/_so/<accelerator>-<arch> directory that holds
the extension module. The build scripts already reject anything but
nvidia/ascend (build.sh, scripts/build_vendor.sh), but the runtime accepted any
string: a typo such as "ascned" was silently coerced to the nvidia platform tag
in _so_platform_tag(), so the process looked for a directory that does not exist
and failed with a bare ModuleNotFoundError for deepinfer._C that never mentions
the variable.

Validate the value with EnvVar.choice so an unsupported value fails fast at
import time with the variable name. The documented behaviours are unchanged:
unset -> nvidia, "nvidia"/"ascend" -> themselves, and the empty value still
means "let deep_cluster detect the accelerator".
---
 deepinfer/environ.py              |  3 +-
 tests/test_environ_accelerator.py | 52 +++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_environ_accelerator.py

diff --git a/deepinfer/environ.py b/deepinfer/environ.py
index e3a55bb..78e919b 100644
--- a/deepinfer/environ.py
+++ b/deepinfer/environ.py
@@ -105,7 +105,8 @@ class envs:
     "环境变量的集中定义模块"
 
     # -- Platform --
-    DEEPINFER_ACCELERATOR = EnvVar.str(default="nvidia")
+    DEEPINFER_ACCELERATOR = EnvVar.choice(default="nvidia", choices=("", "nvidia", "ascend"))
+    "构建脚本（build.sh / build_vendor.sh）只接受 nvidia|ascend；空值表示运行时交给 deep_cluster 探测。"
 
     # -- Runtime --
     DEEPINFER_COMPUTE_SUM_SMALLER_PROBS = EnvVar.bool()
diff --git a/tests/test_environ_accelerator.py b/tests/test_environ_accelerator.py
new file mode 100644
index 0000000..9e6e148
--- /dev/null
+++ b/tests/test_environ_accelerator.py
@@ -0,0 +1,52 @@
+"""Strict validation of DEEPINFER_ACCELERATOR.
+
+The variable selects the `deepinfer/_so/<accelerator>-<arch>` directory that holds the
+extension module. `build.sh` and `scripts/build_vendor.sh` already reject any value other
+than `nvidia`/`ascend`; the runtime must reject it too instead of silently resolving to
+another platform's extension directory.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import pytest
+
+from deepinfer import environ, envs
+
+ACCELERATOR = "DEEPINFER_ACCELERATOR"
+
+
+@pytest.fixture
+def resolve(monkeypatch: pytest.MonkeyPatch) -> Callable[[str | None], str | None]:
+    """Resolve the descriptor against a given environment value.
+
+    The descriptor caches its value on the class, so it must be re-armed after the
+    environment changes and before every read.
+    """
+
+    def _resolve(value: str | None) -> str | None:
+        if value is None:
+            monkeypatch.delenv(ACCELERATOR, raising=False)
+        else:
+            monkeypatch.setenv(ACCELERATOR, value)
+        environ.rearm(ACCELERATOR)
+        return envs.DEEPINFER_ACCELERATOR
+
+    return _resolve
+
+
+def test_unset_defaults_to_nvidia(resolve: Callable[[str | None], str | None]) -> None:
+    assert resolve(None) == "nvidia"
+
+
+@pytest.mark.parametrize("value", ["nvidia", "ascend", ""])
+def test_supported_values_are_returned(resolve: Callable[[str | None], str | None], value: str) -> None:
+    # The empty value is the documented "let deep_cluster detect the accelerator" sentinel.
+    assert resolve(value) == value
+
+
+@pytest.mark.parametrize("value", ["ascned", "ASCEND", "Ascend", "gpu", "npu", "nvidia "])
+def test_unknown_values_are_rejected(resolve: Callable[[str | None], str | None], value: str) -> None:
+    with pytest.raises(ValueError, match=ACCELERATOR):
+        resolve(value)