Add an option to launch cacheflow without ray (#51)

cacheflow/worker/controller.py

@@ -1,6 +1,9 @@
 from typing import Dict, List, Union, Tuple
 
-import ray
+try:
+    import ray
+except ImportError:
+    ray = None
 
 from cacheflow.master.scheduler import Scheduler
 from cacheflow.sequence import SequenceGroupInputs
@@ -29,6 +32,7 @@ class Controller:
         model_path: str,
         use_dummy_weights: bool,
         max_num_batched_tokens: int,
+        use_ray: bool,
     ) -> None:
         self.stage_id = stage_id
         self.stage_devices = stage_devices
@@ -36,6 +40,7 @@ class Controller:
         self.block_size = block_size
         self.num_gpu_blocks = num_gpu_blocks
         self.num_cpu_blocks = num_cpu_blocks
+        self.use_ray = use_ray
 
         # Which pipeline stage is this node assigned to?
         self.is_first_stage = stage_id == 0
@@ -43,10 +48,13 @@ class Controller:
 
         self.workers: List[Worker] = []
         for rank, node_resource, device_id in stage_devices:
-            worker_cls = ray.remote(num_cpus=0,
-                                    num_gpus=1,
-                                    resources={node_resource: 1e-5})(Worker)
-            worker = worker_cls.remote(
+            if self.use_ray:
+                worker_cls = ray.remote(num_cpus=0,
+                                        num_gpus=1,
+                                        resources={node_resource: 1e-5})(Worker).remote
+            else:
+                worker_cls = Worker
+            worker = worker_cls(
                 model_name=model_name,
                 block_size=block_size,
                 num_gpu_blocks=num_gpu_blocks,
@@ -78,17 +86,21 @@ class Controller:
         blocks_to_swap_out: Dict[int, int],
         blocks_to_copy: Dict[int, List[int]],
     ) -> None:
-        futures = []
+        all_outputs = []
         for worker in self.workers:
-            future = worker.execute_stage.remote(
+            executor = (worker.execute_stage.remote
+                        if self.use_ray else worker.execute_stage)
+            output = executor(
                 input_seq_groups,
                 blocks_to_swap_in,
                 blocks_to_swap_out,
                 blocks_to_copy,
             )
-            futures.append(future)
+            all_outputs.append(output)
+
+        if self.use_ray:
+            all_outputs = ray.get(all_outputs)
 
-        all_outputs = ray.get(futures)
         # Make sure all workers have the same results.
         output = all_outputs[0]
         for other_output in all_outputs[1:]: