fix: make exit_flag verification for ascend more general (#2588)

InternLM · Oct 14, 2024 · 88eccb2 · 88eccb2
1 parent 9b52f8d
commit 88eccb2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/lmdeploy/pytorch/engine/model_agent.py b/lmdeploy/pytorch/engine/model_agent.py
@@ -726,8 +726,7 @@ def _exit_by_sending_exit_flag(rank: int, agent: TPModelAgent):
         return
 
     import sys
-    if agent.backend_config.device_type == 'ascend' \
-            and 'uvicorn.server' in sys.modules:
+    if 'torch_npu' in sys.modules and 'uvicorn.server' in sys.modules:
         # Workaround for CLI serve mode with device_type ascend:
         # using uvicorn server causes ascend low-level backend of subprocesses
         # corrupted, and using _broadcast_inputs in this case leads to