polish code

FlagOpen · Dec 27, 2024 · 8a1bd3e · 8a1bd3e
1 parent 008416d
commit 8a1bd3e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/flagscale/serve/core/dag.py b/flagscale/serve/core/dag.py
@@ -95,7 +95,7 @@ def build_task(self):
         self.check_dag()
         return
 
-    def run_task(self, input_data):
+    def run_task(self, input_data=None):
         assert len(self.tasks) > 0
         ray.init(
             num_gpus=6,
@@ -134,7 +134,10 @@ def run_task(self, input_data):
                                 model_nodes[dependencies[0]]
                             )
                     else:
-                        model_nodes[model_alias] = self.tasks[model_alias].bind(input_data)
+                        if input_data is None:
+                            model_nodes[model_alias] = self.tasks[model_alias].bind()
+                        else:
+                            model_nodes[model_alias] = self.tasks[model_alias].bind(input_data)
                     models_to_process.remove(model_alias)
                     progress = True
             if not progress:

diff --git a/flagscale/serve/run_serve.py b/flagscale/serve/run_serve.py
@@ -13,7 +13,7 @@ def main():
     if config["deploy"].get("router"):
         builder.run_router_task()
     else:
-        result = builder.run_task(input_data="Introduce Bruce Lee")
+        result = builder.run_task()
 
 if __name__ == "__main__":
     main()