refine multi-backend setup.py (#2880)

* refine multi-backend setup.py

* fix ci requirments

* add maca requirements

* add maca runtime requirements

* Update Dockerfile_aarch64_ascend

* Update runtime_maca.txt

* change env name

* fix cuda requirements

.github/workflows/pr_ete_test.yml

@@ -10,7 +10,7 @@ on:
       - "3rdparty/**"
       - "lmdeploy/**"
       - "requirements/**"
-      - "requirements.txt"
+      - "requirements_cuda.txt"
       - "CMakeLists.txt"
       - "setup.py"
   workflow_dispatch:
@@ -68,7 +68,7 @@ jobs:
           export PATH=$PATH:/usr/local/openmpi/bin
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/openmpi/lib
           python3 -m pip install cmake packaging wheel transformers_stream_generator transformers datasets openai einops timm decord
-          python3 -m pip install -r requirements.txt -r requirements/test.txt -r requirements/build.txt
+          python3 -m pip install -r requirements_cuda.txt -r requirements/test.txt -r requirements/build.txt
           mkdir -p build && cd build &&\
           sh ../generate.sh &&\
           ninja -j$(nproc) && ninja install &&\

.github/workflows/unit-test.yml

@@ -10,7 +10,7 @@ on:
       - "3rdparty/**"
       - "lmdeploy/**"
       - "requirements/**"
-      - "requirements.txt"
+      - "requirements_cuda.txt"
       - "CMakeLists.txt"
       - "setup.py"
   push:
@@ -24,7 +24,7 @@ on:
       - "3rdparty/**"
       - "lmdeploy/**"
       - "requirements/**"
-      - "requirements.txt"
+      - "requirements_cuda.txt"
       - "CMakeLists.txt"
       - "setup.py"
     tags:
@@ -78,7 +78,7 @@ jobs:
           python3 -m pip install pynvml packaging protobuf transformers_stream_generator
           # manually install flash attn
           python3 -m pip install /root/packages/flash_attn-2.6.3+cu118torch2.3cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
-          python3 -m pip install -r requirements.txt -r requirements/test.txt
+          python3 -m pip install -r requirements_cuda.txt -r requirements/test.txt
           python3 -m pip install .
       - name: Check env
         run: |

docker/Dockerfile_aarch64_ascend

@@ -122,4 +122,4 @@ WORKDIR /opt/lmdeploy
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     sed -i '/triton/d' requirements/runtime.txt && \
-    pip3 install -v --no-build-isolation -e .
+    LMDEPLOY_TARGET_DEVICE=ascend pip3 install -v --no-build-isolation -e .

requirements/runtime_ascend.txt

@@ -16,7 +16,8 @@ safetensors
 sentencepiece
 shortuuid
 tiktoken
-torch<=2.4.0,>=2.0.0
+torch<=2.4.0,>=2.3.1
+torch-npu==2.3.1
 torchvision<=0.19.0,>=0.15.0
 transformers
 uvicorn

requirements/runtime_maca.txt

@@ -0,0 +1,22 @@
+accelerate==0.32.1
+einops
+fastapi
+fire
+mmengine-lite
+numpy<2.0.0
+openai
+outlines<0.1.0
+peft<=0.11.1
+pillow
+protobuf
+pydantic>2.0.0
+pynvml
+safetensors
+sentencepiece
+shortuuid
+tiktoken
+torch<=2.4.0,>=2.0.0
+torchvision<=0.19.0,>=0.15.0
+transformers
+triton>=2.1.0; sys_platform == "linux"
+uvicorn

requirements.txt → requirements_cuda.txt

@@ -1,4 +1,4 @@
 -r requirements/build.txt
--r requirements/runtime.txt
+-r requirements/runtime_cuda.txt
 -r requirements/lite.txt
 -r requirements/serve.txt

requirements_maca.txt

@@ -0,0 +1,4 @@
+-r requirements/build.txt
+-r requirements/runtime_maca.txt
+-r requirements/lite.txt
+-r requirements/serve.txt

setup.py

@@ -4,18 +4,14 @@
 
 from setuptools import find_packages, setup
 
-npu_available = False
-try:
-    import torch_npu
-
-    npu_available = torch_npu.npu.is_available()
-except ImportError:
-    pass
-
 pwd = os.path.dirname(__file__)
 version_file = 'lmdeploy/version.py'
 
 
+def get_target_device():
+    return os.getenv('LMDEPLOY_TARGET_DEVICE', 'cuda')
+
+
 def readme():
     with open(os.path.join(pwd, 'README.md'), encoding='utf-8') as f:
         content = f.read()
@@ -154,16 +150,12 @@ def gen_packages_items():
         setup_requires=parse_requirements('requirements/build.txt'),
         tests_require=parse_requirements('requirements/test.txt'),
         install_requires=parse_requirements(
-            'requirements/runtime_ascend.txt'
-            if npu_available else 'requirements/runtime.txt'),
+            f'requirements/runtime_{get_target_device()}.txt'),
         extras_require={
             'all':
-            parse_requirements('requirements_ascend.txt'
-                               if npu_available else 'requirements.txt'),
-            'lite':
-            parse_requirements('requirements/lite.txt'),
-            'serve':
-            parse_requirements('requirements/serve.txt')
+            parse_requirements(f'requirements_{get_target_device()}.txt'),
+            'lite': parse_requirements('requirements/lite.txt'),
+            'serve': parse_requirements('requirements/serve.txt')
         },
         has_ext_modules=check_ext_modules,
         classifiers=[