fix: add envs to script op template

fix: add workflow actions: terminate, delete, resubmit, resume, retry, stop, suspend

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

Author: zjgemi

src/dflow/op_template.py

@@ -9,7 +9,7 @@
                                        V1alpha1ResourceTemplate,
                                        V1alpha1ScriptTemplate,
                                        V1alpha1Template,
-                                       V1ConfigMapKeySelector,
+                                       V1ConfigMapKeySelector, V1EnvVar,
                                        V1ResourceRequirements, V1Volume,
                                        V1VolumeMount)
     from argo.workflows.client.configuration import Configuration
@@ -108,6 +108,7 @@ class ScriptOPTemplate(OPTemplate):
         annotations: annotations for the OP template
         requests: a dict of resource requests
         limits: a dict of resource limits
+        envs: environment variables
     """
 
     def __init__(
@@ -130,6 +131,7 @@ def __init__(
             image_pull_policy: str = None,
             requests: Dict[str, str] = None,
             limits: Dict[str, str] = None,
+            envs: Dict[str, str] = None,
             **kwargs,
     ) -> None:
         super().__init__(name=name, inputs=inputs, outputs=outputs,
@@ -153,6 +155,7 @@ def __init__(
         self.image_pull_policy = image_pull_policy
         self.requests = requests
         self.limits = limits
+        self.envs = envs
 
     def convert_to_argo(self, memoize_prefix=None,
                         memoize_configmap="dflow"):
@@ -170,6 +173,10 @@ def convert_to_argo(self, memoize_prefix=None,
                                     volumes=self.volumes,
                                     resource=self.resource)
         else:
+            if self.envs is not None:
+                env = [V1EnvVar(name=k, value=v) for k, v in self.envs.items()]
+            else:
+                env = None
             return \
                 V1alpha1Template(name=self.name,
                                  metadata=V1alpha1Metadata(
@@ -187,7 +194,8 @@ def convert_to_argo(self, memoize_prefix=None,
                                      volume_mounts=self.mounts,
                                      resources=V1ResourceRequirements(
                                          limits=self.limits,
-                                         requests=self.requests)))
+                                         requests=self.requests),
+                                     env=env))
 
 
 class ShellOPTemplate(ScriptOPTemplate):
@@ -212,6 +220,7 @@ class ShellOPTemplate(ScriptOPTemplate):
         annotations: annotations for the OP template
         requests: a dict of resource requests
         limits: a dict of resource limits
+        envs: environment variables
     """
 
     def __init__(
@@ -233,6 +242,7 @@ def __init__(
         image_pull_policy: str = None,
         requests: Dict[str, str] = None,
         limits: Dict[str, str] = None,
+        envs: Dict[str, str] = None,
         **kwargs,
     ) -> None:
         if command is None:
@@ -243,7 +253,7 @@ def __init__(
             init_progress=init_progress, timeout=timeout,
             retry_strategy=retry_strategy, memoize_key=memoize_key, pvcs=pvcs,
             image_pull_policy=image_pull_policy, annotations=annotations,
-            requests=requests, limits=limits)
+            requests=requests, limits=limits, envs=envs)
 
 
 class PythonScriptOPTemplate(ScriptOPTemplate):
@@ -268,6 +278,7 @@ class PythonScriptOPTemplate(ScriptOPTemplate):
         annotations: annotations for the OP template
         requests: a dict of resource requests
         limits: a dict of resource limits
+        envs: environment variables
     """
 
     def __init__(
@@ -289,6 +300,7 @@ def __init__(
         image_pull_policy: str = None,
         requests: Dict[str, str] = None,
         limits: Dict[str, str] = None,
+        envs: Dict[str, str] = None,
         **kwargs,
     ) -> None:
         if command is None:
@@ -299,4 +311,4 @@ def __init__(
             init_progress=init_progress, timeout=timeout,
             retry_strategy=retry_strategy, memoize_key=memoize_key, pvcs=pvcs,
             image_pull_policy=image_pull_policy, annotations=annotations,
-            requests=requests, limits=limits)
+            requests=requests, limits=limits, envs=envs)

src/dflow/python/python_op_template.py

@@ -101,6 +101,7 @@ class PythonOPTemplate(PythonScriptOPTemplate):
         image_pull_policy: Always, IfNotPresent, Never
         requests: a dict of resource requests
         limits: a dict of resource limits
+        envs: environment variables
     """
 
     def __init__(self,
@@ -130,6 +131,7 @@ def __init__(self,
                  requests: Dict[str, str] = None,
                  limits: Dict[str, str] = None,
                  upload_dflow: bool = True,
+                 envs: Dict[str, str] = None,
                  **kwargs,
                  ) -> None:
         op = None
@@ -167,7 +169,7 @@ def __init__(self,
         super().__init__(name="%s-%s" % (class_name, "".join(random.sample(
             string.digits + string.ascii_lowercase, 5))), inputs=Inputs(),
             outputs=Outputs(), volumes=volumes, mounts=mounts,
-            requests=requests, limits=limits)
+            requests=requests, limits=limits, envs=envs)
         self.slices = slices
         if timeout is not None:
             self.timeout = "%ss" % timeout

src/dflow/workflow.py

@@ -355,3 +355,72 @@ def query_keys_of_steps(
             a list of keys
         """
         return [step.key for step in self.query_step() if step.key is not None]
+
+    def terminate(self) -> None:
+        """
+        Terminate the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/terminate' % (self.namespace, self.id),
+            'PUT')
+
+    def delete(self) -> None:
+        """
+        Delete the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s' % (self.namespace, self.id), 'DELETE')
+
+    def resubmit(self) -> None:
+        """
+        Resubmit the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/resubmit' % (self.namespace, self.id),
+            'PUT')
+
+    def resume(self) -> None:
+        """
+        Resume the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/resume' % (self.namespace, self.id),
+            'PUT')
+
+    def retry(self) -> None:
+        """
+        Retry the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/retry' % (self.namespace, self.id),
+            'PUT')
+
+    def stop(self) -> None:
+        """
+        Stop the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/stop' % (self.namespace, self.id),
+            'PUT')
+
+    def suspend(self) -> None:
+        """
+        Suspend the workflow
+        """
+        if self.id is None:
+            raise RuntimeError("Workflow ID is None")
+        self.api_instance.api_client.call_api(
+            '/api/v1/workflows/%s/%s/suspend' % (self.namespace, self.id),
+            'PUT')