fix: add parallelism to steps and dag

Signed-off-by: zjgemi <[email protected]>

Author: zjgemi

src/dflow/dag.py

@@ -22,6 +22,7 @@ class DAG(OPTemplate):
         tasks: a list of tasks
         memoize_key: memoized key of the dag
         annotations: annotations for the OP template
+        parallelism: maximum number of running pods for the OP template
         """
 
     def __init__(
@@ -32,9 +33,11 @@ def __init__(
             tasks: List[Task] = None,
             memoize_key: str = None,
             annotations: Dict[str, str] = None,
+            parallelism: int = None,
     ) -> None:
         super().__init__(name=name, inputs=inputs, outputs=outputs,
                          memoize_key=memoize_key, annotations=annotations)
+        self.parallelism = parallelism
         if tasks is not None:
             self.tasks = tasks
         else:
@@ -86,5 +89,6 @@ def convert_to_argo(self, memoize_prefix=None,
                              ),
                              inputs=self.inputs.convert_to_argo(),
                              outputs=self.outputs.convert_to_argo(),
-                             memoize=self.memoize)
+                             memoize=self.memoize,
+                             parallelism=self.parallelism)
         return argo_template, templates

src/dflow/steps.py

@@ -21,6 +21,7 @@ class Steps(OPTemplate):
         steps: a sequential list of steps
         memoize_key: memoized key of the steps
         annotations: annotations for the OP template
+        parallelism: maximum number of running pods for the OP template
         """
 
     def __init__(
@@ -31,9 +32,11 @@ def __init__(
             steps: List[Union[Step, List[Step]]] = None,
             memoize_key: str = None,
             annotations: Dict[str, str] = None,
+            parallelism: int = None,
     ) -> None:
         super().__init__(name=name, inputs=inputs, outputs=outputs,
                          memoize_key=memoize_key, annotations=annotations)
+        self.parallelism = parallelism
         if steps is not None:
             self.steps = steps
         else:
@@ -95,5 +98,6 @@ def convert_to_argo(self, memoize_prefix=None,
                              steps=argo_steps,
                              inputs=self.inputs.convert_to_argo(),
                              outputs=self.outputs.convert_to_argo(),
-                             memoize=self.memoize)
+                             memoize=self.memoize,
+                             parallelism=self.parallelism)
         return argo_template, templates