[ascend] add env to set rt visable by ray and disable warmup (#3894)

Author: tangzhiyi11

lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py

@@ -121,3 +121,8 @@ def prepare_inputs_for_generation(
             inputs_embeds=inputs_embeds,
             context=context,
         )
+
+    def get_capture_batch_sizes(self) -> List[int]:
+        """Capture batch sizes."""
+        # TODO: disable warmup now.
+        return []

lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from lmdeploy.pytorch import envs as _envs
 from lmdeploy.pytorch.config import BackendConfig, CacheConfig, ModelConfig
 from lmdeploy.utils import get_logger
 
@@ -339,5 +340,6 @@ def device_count():
     @staticmethod
     def support_ray():
         """Support ray."""
-        os.environ['RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES'] = '1'
+        if not _envs.ascend_set_rt_visable_devices_by_ray:
+            os.environ['RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES'] = '1'
         return True

lmdeploy/pytorch/engine/executor/ray_executor.py

@@ -560,17 +560,21 @@ def _init_distributed_environment_by_device(self, device_str: str):
     def _init_ascend_distributed_environment(self, driver_ip):
         """Init ascend distributed environment."""
         rank_table_file = _envs.ascend_rank_table_file
-        if not rank_table_file:
-            # if rank table file is not set, treat as single node
-            self.workers = self._sort_workers(driver_ip, self.workers)
-            # simply set device by index, this is for single node, multiple devices
-            ray.get([worker.set_device.remote(idx) for idx, worker in enumerate(self.workers)])
-        else:
+        set_rt_visable_devices_by_ray = _envs.ascend_set_rt_visable_devices_by_ray
+
+        if rank_table_file:
             # if rank table file is set, use it to get rank mapping, multiple nodes
             rank_mapping, worker_ips, envs = get_ascend_device_rank_mapping(driver_ip)
             self.workers = self._sort_workers_by_ip(worker_ips, self.workers)
             ray.get([worker.set_device.remote(rank_mapping[idx]) for idx, worker in enumerate(self.workers)])
             ray.get([worker.set_env.remote(envs) for worker in self.workers])
+        elif not set_rt_visable_devices_by_ray:
+            # if rank table file is not set, treat as single node
+            # simply set device by index, this is for single node, multiple devices
+            self.workers = self._sort_workers(driver_ip, self.workers)
+            ray.get([worker.set_device.remote(idx) for idx, worker in enumerate(self.workers)])
+        else:
+            self.workers = self._sort_workers(driver_ip, self.workers)
 
     """ PD Disaggregation API Begin """
 

lmdeploy/pytorch/envs.py

@@ -87,6 +87,7 @@ def _patched_get_env(
     ray_nsys_output_prefix = os.getenv('LMDEPLOY_RAY_NSYS_OUT_PREFIX', None)
 
     # ascend
+    ascend_set_rt_visable_devices_by_ray = env_to_bool('ASCEND_SET_RT_VISIBLE_DEVICES_BY_RAY', False)
     ascend_rank_table_file = os.getenv('ASCEND_RANK_TABLE_FILE_PATH')
 
     # dp