Remove unused device parameter from all rolling batch classes

engines/python/setup/djl_python/deepspeed.py

@@ -131,9 +131,8 @@ def initialize(self, properties: dict):
             }
             if "output_formatter" in properties:
                 kwargs["output_formatter"] = properties.get("output_formatter")
-            self.rolling_batch = DeepSpeedRollingBatch(self.model,
-                                                       self.properties.device,
-                                                       properties, **kwargs)
+            self.rolling_batch = DeepSpeedRollingBatch(self.model, properties,
+                                                       **kwargs)
         else:
             self.create_model_pipeline()
         self.logger.info(

engines/python/setup/djl_python/huggingface.py

@@ -147,8 +147,8 @@ def initialize(self, properties: dict):
                 self.hf_configs.rolling_batch.value, self.hf_configs.is_mpi,
                 self.model_config)
             self.rolling_batch = _rolling_batch_cls(
-                self.hf_configs.model_id_or_path, self.hf_configs.device,
-                properties, **self.hf_configs.kwargs)
+                self.hf_configs.model_id_or_path, properties,
+                **self.hf_configs.kwargs)
             self.initialized = True
             return
         elif is_streaming_enabled(self.hf_configs.enable_streaming):

engines/python/setup/djl_python/rolling_batch/deepspeed_rolling_batch.py

@@ -23,12 +23,11 @@
 
 class DeepSpeedRollingBatch(RollingBatch):
 
-    def __init__(self, model: InferenceEngine, device, properties, **kwargs):
+    def __init__(self, model: InferenceEngine, properties, **kwargs):
         """
-        Initializes the LmiDistRollingBatch.
+        Initializes the DeepSpeedRollingBatch.
 
         :param model_id_or_path: model id or path
-        :param device: model loaded device
         :param properties: other properties of the model, such as decoder strategy
         :param kwargs passed while loading the model
         """

engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py

@@ -30,12 +30,11 @@
 
 class LmiDistRollingBatch(RollingBatch):
 
-    def __init__(self, model_id_or_path, device, properties, **kwargs):
+    def __init__(self, model_id_or_path, properties, **kwargs):
         """
         Initializes the LmiDistRollingBatch.
 
         :param model_id_or_path: model id or path
-        :param device: model loaded device
         :param properties: other properties of the model, such as decoder strategy
         :param kwargs passed while loading the model
         """

engines/python/setup/djl_python/rolling_batch/scheduler_rolling_batch.py

@@ -41,12 +41,11 @@ def enable_flash():
 
 class SchedulerRollingBatch(RollingBatch):
 
-    def __init__(self, model_id_or_path, device, properties, **kwargs):
+    def __init__(self, model_id_or_path, properties, **kwargs):
         """
         Initializes the rolling batch scheduler.
 
         :param model_id_or_path: model id or path
-        :param device: model loaded device
         :param properties: other properties of the model, such as decoder strategy
         :param kwargs passed while loading the model
         """

engines/python/setup/djl_python/rolling_batch/trtllm_rolling_batch.py

@@ -17,7 +17,7 @@
 
 class TRTLLMRollingBatch(RollingBatch):
 
-    def __init__(self, model_id_or_path, device, properties, **kwargs):
+    def __init__(self, model_id_or_path, properties, **kwargs):
         """
         Initializes the TRTLLMRollingBatch.
         :param model_id_or_path: model id or path

engines/python/setup/djl_python/rolling_batch/vllm_rolling_batch.py

@@ -36,7 +36,7 @@
 class VLLMRollingBatch(RollingBatch):
 
     # TODO: Make properties is the only parameter, after refactoring all rolling batch handlers
-    def __init__(self, model_id_or_path, device, properties, **kwargs):
+    def __init__(self, model_id_or_path, properties, **kwargs):
         """
         Initializes the VLLMRollingBatch.
         :param properties: other properties of the model, such as decoder strategy

engines/python/setup/djl_python/tensorrt_llm.py

@@ -31,7 +31,7 @@ def initialize(self, properties: dict):
         self.trt_configs = TensorRtLlmProperties(**properties)
 
         self.rolling_batch = TRTLLMRollingBatch(
-            self.trt_configs.model_id_or_path, None, properties, **properties)
+            self.trt_configs.model_id_or_path, properties, **properties)
         self.initialized = True
         return
 

engines/python/setup/djl_python/tests/rolling_batch_test_scripts/run_rolling_batch_alone.py

@@ -53,7 +53,7 @@ def init_rolling_batch(rolling_batch_type: str, model_id: str,
     else:
         rolling_batcher_cls = get_rolling_batch_class_from_str(
             rolling_batch_type)
-        return rolling_batcher_cls(model_id, device, properties, **properties)
+        return rolling_batcher_cls(model_id, properties, **properties)
 
 
 def print_rank0(content):

engines/python/setup/djl_python/tests/rolling_batch_test_scripts/test_rolling_batch.py

@@ -41,8 +41,8 @@
 
 # ===================== lmi ============================
 print("=========== before =========")
-rolling_batch = SchedulerRollingBatch(model_id, device, properties)
-# rolling_batch = LmiDistRollingBatch(model_id, device, properties)
+rolling_batch = SchedulerRollingBatch(model_id, properties)
+# rolling_batch = LmiDistRollingBatch(model_id, properties)
 rolling_batch.output_formatter = None
 print("reach here")
 

engines/python/setup/djl_python/tests/test_lmi_dist.py

@@ -98,7 +98,7 @@ def test_models(self):
             device = int(os.environ.get("RANK", 0))
             properties["device"] = int(os.environ.get("RANK", 0))
 
-            rolling_batch = LmiDistRollingBatch(model_id, device, properties)
+            rolling_batch = LmiDistRollingBatch(model_id, properties)
             rolling_batch.output_formatter = None
 
             gen = Generator(rolling_batch=rolling_batch)