fix: syntax sugar for python op template and task

Signed-off-by: zjgemi <liuxin_zijian@163.com>

Author: zjgemi

examples/test_dag.py

@@ -1,4 +1,4 @@
-from dflow import (DAG, InputArtifact, InputParameter, OutputArtifact,
+from dflow import (InputArtifact, InputParameter, OutputArtifact,
                    OutputParameter, ShellOPTemplate, Task, Workflow)
 
 if __name__ == "__main__":
@@ -21,14 +21,12 @@
     duplicate.inputs.artifacts = {"foo": InputArtifact(path="/tmp/foo.txt")}
     duplicate.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")}
 
-    dag = DAG()
+    wf = Workflow(name="dag")
     hello0 = Task(name="hello0", template=hello)
-    dag.add(hello0)
+    wf.add(hello0)
     hello1 = Task(name="hello1",
                   template=duplicate,
                   parameters={"msg": hello0.outputs.parameters["msg"]},
                   artifacts={"foo": hello0.outputs.artifacts["bar"]})
-    dag.add(hello1)
-
-    wf = Workflow(name="dag", dag=dag)
+    wf.add(hello1)
     wf.submit()

src/dflow/plugins/datasets.py

@@ -200,7 +200,7 @@ def render(self, template: PythonOPTemplate, name: str
         script += " && rclone mount %s@%s: /launching/%s" % (
             self.element, self.version, name)
         template.sidecars.append(V1alpha1UserContainer(
-            name="rclone-%s" % name,
+            name="rclone-%s" % name.replace("_", "-"),
             image=self.rclone_image,
             image_pull_policy=self.rclone_image_pull_policy,
             command=["sh", "-c"],

src/dflow/python/op.py

@@ -6,14 +6,16 @@
 import logging
 import os
 from abc import ABC
+from functools import partial
 from pathlib import Path
 from typing import Dict, List, Set, Union
 
 from typeguard import check_type
 
 from ..argo_objects import ArgoObjectDict
 from ..config import config
-from ..utils import get_key, s3_config
+from ..context_syntax import GLOBAL_CONTEXT
+from ..utils import get_key, randstr, s3_config
 from .opio import (OPIO, Artifact, BigParameter, OPIOSign, Parameter,
                    type_to_str)
 
@@ -181,7 +183,9 @@ def _check_signature(
                 check_type(ii, io, ss)
 
     @classmethod
-    def function(cls, func):
+    def function(cls, func=None, **kwargs):
+        if func is None:
+            return partial(cls.function, **kwargs)
         signature = func.__annotations__
         return_type = signature.get('return', None)
 
@@ -220,6 +224,22 @@ def execute(self, op_in):
                 return op_out
 
             def __call__(self, **op_in):
+                if GLOBAL_CONTEXT.in_context:
+                    from .python_op_template import PythonOPTemplate
+                    from ..task import Task
+                    input_sign = self.get_input_sign()
+                    parameters = {k: v for k, v in op_in.items()
+                                  if not isinstance(input_sign[k], Artifact)}
+                    artifacts = {k: v for k, v in op_in.items()
+                                 if isinstance(input_sign[k], Artifact)}
+                    name = func.__name__.lower().replace("_", "-") + "-" + \
+                        randstr()
+                    task = Task(name,
+                                template=PythonOPTemplate(self, **kwargs),
+                                parameters=parameters, artifacts=artifacts)
+                    op_out = {**task.outputs.parameters,
+                              **task.outputs.artifacts}
+                    return op_out
                 return self.execute(op_in)
 
         subclass.func = func

src/dflow/python/python_op_template.py

@@ -212,8 +212,8 @@ def __init__(self,
             for name, global_name in output_artifact_global_name.items():
                 output_sign[name].global_name = global_name
         super().__init__(
-            name="%s-%s" % (class_name, randstr()), inputs=Inputs(),
-            outputs=Outputs(), volumes=volumes, mounts=mounts,
+            name="%s-%s" % (class_name.replace("_", "-"), randstr()),
+            inputs=Inputs(), outputs=Outputs(), volumes=volumes, mounts=mounts,
             requests=requests, limits=limits, envs=envs,
             init_containers=init_containers, sidecars=sidecars)
         self.pre_script = pre_script

src/dflow/workflow.py

@@ -155,7 +155,7 @@ def __init__(
             assert isinstance(dag, DAG)
             self.entrypoint = dag
         else:
-            self.entrypoint = Steps(self.name + "-steps")
+            self.entrypoint = None
         self.templates = {}
         self.argo_templates = {}
         self.pvcs = {}
@@ -180,6 +180,7 @@ def __enter__(self) -> 'Workflow':
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         GLOBAL_CONTEXT.in_context = False
+        self.submit()
 
     def add(
             self,
@@ -192,6 +193,12 @@ def add(
             step: a step or a list of parallel steps to be added to the
             entrypoint of the workflow
         """
+        if self.entrypoint is None:
+            if isinstance(step, Task) or (isinstance(step, list) and all(
+                    [isinstance(s, Task) for s in step])):
+                self.entrypoint = DAG(self.name + "-dag")
+            else:
+                self.entrypoint = Steps(self.name + "-steps")
         self.entrypoint.add(step)
 
     def submit(

tutorials/dflow-function.ipynb

@@ -24,7 +24,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "##### In this tutorial, we will complete the task using `PythonOPTemplate` and function OP"
+    "##### In this tutorial, we will complete the task using function OP"
    ]
   },
   {
@@ -37,14 +37,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
+    "import sys\n",
     "from pathlib import Path\n",
     "\n",
-    "from dflow import Step, Workflow\n",
-    "from dflow.python import OP, Artifact, PythonOPTemplate"
+    "from dflow import Workflow\n",
+    "from dflow.python import OP, Artifact"
    ]
   },
   {
@@ -54,7 +55,7 @@
    "source": [
     "For step 1: \n",
     "\n",
-    "This class can achieve the function to write files. In the example, we try to write a file containing message of string format, and output a number of int format.\n",
+    "This OP is to write files. In the example, we try to write a file containing message of string format, and output a number of int format.\n",
     "- input:\n",
     "    - \"msg\": the input message\n",
     "- output:\n",
@@ -64,12 +65,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "@OP.function\n",
-    "def WriteFile(msg: str) -> {\"out_art\": Artifact(Path), \"length\": int}:\n",
+    "@OP.function(image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\")\n",
+    "def write_file(msg: str) -> {\"out_art\": Artifact(Path), \"length\": int}:\n",
     "    with open(\"msg.txt\",\"w\") as f:\n",
     "        f.write(msg)\n",
     "    \n",
@@ -91,12 +92,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "@OP.function\n",
-    "def Duplicate(in_art: Artifact(Path), in_num: int) -> {\"out_art\": Artifact(Path), \"out_num\": int}:\n",
+    "@OP.function(image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\")\n",
+    "def duplicate(in_art: Artifact(Path), in_num: int) -> {\"out_art\": Artifact(Path), \"out_num\": int}:\n",
     "    with open(in_art, \"r\") as f:\n",
     "        content = f.read()\n",
     "    with open(\"bar.txt\", \"w\") as f:\n",
@@ -113,52 +114,18 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "After defining the PythonOPTemplate, it is to define steps.\n",
-    "- step0 is to write a file and to output length of the content using WriteFile OP and python image\n",
-    "    - parameters:\n",
-    "        - \"msg\": \"HelloWorld!\"\n",
-    "- step1 is to duplicate the content in the file and to duplicate the number using Duplicate OP and python image\n",
-    "    - artifacts:\n",
-    "        - in_art file is from step0.outputs.artifacts[\"out_art\"]\n",
-    "    - parameters:\n",
-    "        - in_num is from step0.outputs.parameters[\"length\"]\n",
-    "\n",
-    "Finally, we need to set up a Workflow named \"python\" and then add step0 and step1.\n",
-    "\n",
-    "`wf.submit` is to submit this workflow to Argo."
+    "After defining OPs, call the OPs in series in the context of a workflow, which will connect them together as a workflow and submit it finally."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Workflow has been submitted (ID: python-wxj66, UID: 600f99da-21ee-40af-9a3b-c738e46556b8)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "import sys\n",
-    "step0 = Step(\n",
-    "    name=\"step0\",\n",
-    "    template=PythonOPTemplate(WriteFile, image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\"),\n",
-    "    parameters={\"msg\": \"HelloWorld!\"},\n",
-    ")\n",
-    "\n",
-    "step1 = Step(\n",
-    "    name=\"step1\",\n",
-    "    template=PythonOPTemplate(Duplicate, image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\"),\n",
-    "    parameters={\"in_num\": step0.outputs.parameters[\"length\"]},\n",
-    "    artifacts={\"in_art\": step0.outputs.artifacts[\"out_art\"]},\n",
-    ")\n",
-    "wf = Workflow(name=\"python\")\n",
-    "wf.add(step0)\n",
-    "wf.add(step1)\n",
-    "wf.submit();"
+    "with Workflow(name=\"python\") as wf:\n",
+    "    out = write_file(msg=\"HelloWorld!\")\n",
+    "    duplicate(in_num=out[\"length\"], in_art=out[\"out_art\"])"
    ]
   }
  ],