Repository: couler-proj/couler

Issue:
couler.set_dependencies do not set dependencies correctly
## Summary 
when `set_dependencies` None dep, have no deal with this issue
minimal Reproducible  code：
```
        couler.set_dependencies(
            lambda: job_a(message="A"), dependencies=None
        )
        couler.set_dependencies(
            lambda: job_b(message="B"), dependencies="A.Succeeded"
        )
        couler.set_dependencies(
            lambda: job_c(message="C"), dependencies=None
        )
```
What happened/what you expected to happen?
```
        wf_tree = couler.workflow_yaml()
        tasks = wf_tree["spec"]["templates"][0]["dag"]["tasks"]
        deps = tasks[2].get("dependencies")
```
Before my pr,  deps will be `B` ,but it should be None
## Diagnostics

What is the version of Couler you are using? 
master branch

I'll fix it
Impacted by this bug? Give it a 👍. We prioritize the issues with the most 👍.


Discussion:

Patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index efb73cb..aa86622 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -31,6 +31,7 @@ jobs:
         pre-commit run --all
     - name: Unit tests
       run: |
+        python setup.py proto
         python setup.py install
         python -m pytest
     - uses: opsgang/ga-setup-minikube@v0.1.1
diff --git a/.gitignore b/.gitignore
index 24f4e90..a1ab6f2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,6 +12,7 @@ __pycache__
 # Distribution / packaging
 .Python
 build/
+couler/proto/
 develop-eggs/
 dist/
 downloads/
diff --git a/couler/core/states.py b/couler/core/states.py
index e590a60..9ee8209 100644
--- a/couler/core/states.py
+++ b/couler/core/states.py
@@ -71,10 +71,16 @@ def _cleanup():
     """Cleanup the cached fields, just used for unit test.
     """
     global _secrets, _update_steps_lock, _dag_caller_line, _upstream_dag_task, _upstream_dag_depends_logic, workflow, _steps_outputs  # noqa: E501
+    global _exit_handler_enable, _when_prefix, _when_task, _while_steps, _concurrent_func_line  # noqa: E501
     _secrets = {}
     _update_steps_lock = True
     _dag_caller_line = None
     _upstream_dag_task = None
     _upstream_dag_depends_logic = None
+    _exit_handler_enable = False
+    _when_prefix = None
+    _when_task = None
+    _while_steps = OrderedDict()
+    _concurrent_func_line = -1
     _steps_outputs = OrderedDict()
     workflow.cleanup()
diff --git a/couler/core/syntax/dag.py b/couler/core/syntax/dag.py
index 99017e2..a4e28e4 100644
--- a/couler/core/syntax/dag.py
+++ b/couler/core/syntax/dag.py
@@ -65,7 +65,9 @@ def set_dependencies(step_function, dependencies=None):
             states._upstream_dag_task = None
         else:
             raise SyntaxError("dependencies must be a list or a string")
-
+    else:
+        states._upstream_dag_depends_logic = None
+        states._upstream_dag_task = None
     if not callable(step_function):
         raise SyntaxError("require step_function to a function")
 
diff --git a/proto/couler.proto b/proto/couler.proto
new file mode 100644
index 0000000..0413c9d
--- /dev/null
+++ b/proto/couler.proto
@@ -0,0 +1,47 @@
+syntax = "proto3";
+
+package couler.v1;
+
+message Parameter {
+  string name = 1;
+  string value = 2;
+}
+
+message Artifact {
+  string name = 1;
+  string type = 2;  // HTTP/GIT/...
+  string value = 3;
+}
+
+message StdOut {
+  string name = 1;
+}
+
+message StepIO {
+  string name = 1;
+  int32 source = 2;
+  oneof step_io {
+    Parameter parameter = 3;
+    Artifact  artifact = 4;
+    StdOut stdout = 5;
+  }
+}
+
+message Step {
+  int32 id = 1;
+  string name = 2;      // name for reference
+  string tmpl_name = 3; // name for generate template
+  string image = 4;
+  repeated string command = 5;
+  string script = 6;
+  repeated StepIO inputs = 7;
+  repeated StepIO outputs = 8;
+  repeated int32  dependencies = 9;
+}
+
+message Workflow {
+  repeated Step steps = 1;
+  int32 parallelism = 2;
+  string secret = 3;
+}
+
diff --git a/setup.py b/setup.py
index 9db4156..43def85 100644
--- a/setup.py
+++ b/setup.py
@@ -11,11 +11,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import distutils.cmd
 import io
 import os
 
 from setuptools import find_packages, setup
 
+
+class ProtocCommand(distutils.cmd.Command):
+    description = "run protoc to compile protobuf files"
+    user_options = []
+
+    def initialize_options(self):
+        """Set default values for options."""
+        pass
+
+    def finalize_options(self):
+        """Post-process options."""
+        pass
+
+    def run(self):
+        os.system("protoc --python_out=couler proto/couler.proto")
+        os.system("touch couler/proto/__init__.py")
+        self.announce(
+            "proto file generated under couler/proto.",
+            level=distutils.log.INFO,
+        )
+
+
 with open("requirements.txt") as f:
     required_deps = f.read().splitlines()
 
@@ -39,4 +62,5 @@ setup(
     python_requires=">=3.6",
     packages=find_packages(exclude=["*test*"]),
     package_data={"": ["requirements.txt"]},
+    cmdclass={"proto": ProtocCommand},
 )


Test patch:
diff --git a/couler/tests/argo_test.py b/couler/tests/argo_test.py
index a221e67..949274a 100644
--- a/couler/tests/argo_test.py
+++ b/couler/tests/argo_test.py
@@ -4,14 +4,44 @@ from collections import OrderedDict
 import yaml
 
 import couler.argo as couler
+from couler.core import states
 from couler.core.templates.volume import Volume, VolumeMount
 from couler.core.templates.volume_claim import VolumeClaimTemplate
 
 
-class ArgoTest(unittest.TestCase):
+class ArgoBaseTestCase(unittest.TestCase):
     def setUp(self):
+        self.state_keys = dir(states)
+        self.tc_dump = {}
+        for key in self.state_keys:
+            if key.startswith("__"):
+                continue
+            if key.startswith("_"):
+                value = getattr(states, key)
+                if callable(value):
+                    continue
+                self.tc_dump[key] = value
         couler._cleanup()
 
+    def tearDown(self):
+        couler._cleanup()
+        for key in self.state_keys:
+            if key.startswith("__"):
+                continue
+            if key.startswith("_"):
+                value = getattr(states, key)
+                if callable(value):
+                    continue
+                self.assertEqual(
+                    value, self.tc_dump[key], msg="state not cleanup:%s" % key
+                )
+
+
+class ArgoTest(ArgoBaseTestCase):
+    def setUp(self):
+        couler._cleanup()
+        super().setUp()
+
     def test_run_none_source(self):
         with self.assertRaises(ValueError):
             couler.run_script(image="image1", command="python")
@@ -44,7 +74,7 @@ class ArgoTest(unittest.TestCase):
             template["script"],
             image="image1",
             command=["python"],
-            source="\ncouler._cleanup()\n",
+            source="\ncouler._cleanup()\nsuper().setUp()\n",
             env=None,
         )
 
@@ -61,7 +91,7 @@ class ArgoTest(unittest.TestCase):
             template["script"],
             image="image1",
             command=["python"],
-            source="\ncouler._cleanup()\n",
+            source="\ncouler._cleanup()\nsuper().setUp()\n",
             env=None,
         )
 
@@ -181,7 +211,6 @@ class ArgoTest(unittest.TestCase):
         self.assertEqual(
             template["inputs"]["parameters"], [{"name": "para-B-0"}]
         )
-        couler._cleanup()
 
     def test_set_dependencies_with_exit_handler(self):
         def producer():
@@ -236,7 +265,6 @@ class ArgoTest(unittest.TestCase):
             wf["spec"]["templates"][4],
             OrderedDict([("name", "failure-exit"), expected_container_spec]),
         )
-        couler._cleanup()
 
     def test_create_job(self):
         success_condition = "status.succeeded > 0"
@@ -353,7 +381,6 @@ class ArgoTest(unittest.TestCase):
         self.assertTrue(
             "{{tasks.A.outputs.parameters.output-id-" in envs["value"]
         )
-        couler._cleanup()
 
     def test_run_job_with_dependency_implicit_params_passing_from_job(self):
         success_condition = "status.succeeded > 0"
@@ -440,7 +467,6 @@ class ArgoTest(unittest.TestCase):
             [{"name": "para-B-0"}, {"name": "para-B-1"}, {"name": "para-B-2"}],
             template["inputs"]["parameters"],
         )
-        couler._cleanup()
 
     def _verify_script_body(
         self, script_to_check, image, command, source, env
diff --git a/couler/tests/argo_yaml_test.py b/couler/tests/argo_yaml_test.py
index 2364da3..c0ba8b2 100644
--- a/couler/tests/argo_yaml_test.py
+++ b/couler/tests/argo_yaml_test.py
@@ -1,23 +1,17 @@
 import json
 import os
 import re
-import unittest
 
 import pyaml
 import yaml
 
 import couler.argo as couler
+from couler.tests.argo_test import ArgoBaseTestCase
 
 _test_data_dir = "test_data"
 
 
-class ArgoYamlTest(unittest.TestCase):
-    def setUp(self):
-        couler._cleanup()
-
-    def tearDown(self):
-        couler._cleanup()
-
+class ArgoYamlTest(ArgoBaseTestCase):
     @staticmethod
     def mock_dict(x, mock_str="pytest"):
         # The following fields are overwritten so we won't be
diff --git a/couler/tests/cluster_config_test.py b/couler/tests/cluster_config_test.py
index 6cd31e3..c557ee3 100644
--- a/couler/tests/cluster_config_test.py
+++ b/couler/tests/cluster_config_test.py
@@ -1,10 +1,10 @@
 import os
-import unittest
 
 import couler.argo as couler
+from couler.tests.argo_test import ArgoBaseTestCase
 
 
-class ArgoTest(unittest.TestCase):
+class ArgoClusterTest(ArgoBaseTestCase):
     def test_cluster_config(self):
 
         couler.config_workflow(
diff --git a/couler/tests/daemon_step_test.py b/couler/tests/daemon_step_test.py
index 9cc497e..5443b60 100644
--- a/couler/tests/daemon_step_test.py
+++ b/couler/tests/daemon_step_test.py
@@ -1,12 +1,14 @@
-import unittest
-
 import couler.argo as couler
+from couler.tests.argo_test import ArgoBaseTestCase
 
 
-class DaemonStepTest(unittest.TestCase):
+class DaemonStepTest(ArgoBaseTestCase):
     def setUp(self) -> None:
         couler._cleanup()
 
+    def tearDown(self):
+        couler._cleanup()
+
     def test_run_daemon_container(self):
         self.assertEqual(len(couler.workflow.templates), 0)
         couler.run_container(
diff --git a/couler/tests/dag_test.py b/couler/tests/dag_test.py
index 446f0ac..79085c4 100644
--- a/couler/tests/dag_test.py
+++ b/couler/tests/dag_test.py
@@ -1,3 +1,5 @@
+import pyaml
+
 import couler.argo as couler
 from couler.tests.argo_yaml_test import ArgoYamlTest
 
@@ -160,6 +162,65 @@ class DAGTest(ArgoYamlTest):
 
         self.check_argo_yaml("parameter_passing_golden.yaml")
 
+    def test_set_dependencies_none(self):
+        couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
+        couler.set_dependencies(lambda: job_b(message="B"), dependencies=["A"])
+        couler.set_dependencies(lambda: job_c(message="C"), dependencies=None)
+        pyaml.dump(couler.workflow_yaml())
+        wf_tree = couler.workflow_yaml()
+        tasks = wf_tree["spec"]["templates"][0]["dag"]["tasks"]
+        assert_deps = {"A": None, "B": ["A"], "C": None}
+        self.assertEqual(len(tasks), 3)
+        for task in tasks:
+            assert_dep = assert_deps[task["name"]]
+            self.assertEqual(task.get("dependencies"), assert_dep)
+
+    def test_set_depends_none(self):
+        couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
+        couler.set_dependencies(
+            lambda: job_b(message="B"), dependencies="A.Succeeded"
+        )
+        couler.set_dependencies(lambda: job_c(message="C"), dependencies=None)
+        content = pyaml.dump(couler.workflow_yaml())
+        self.assertNotIn("dependencies", content)
+        wf_tree = couler.workflow_yaml()
+        tasks = wf_tree["spec"]["templates"][0]["dag"]["tasks"]
+        assert_deps = {"A": None, "B": "A.Succeeded", "C": None}
+        self.assertEqual(len(tasks), 3)
+        for task in tasks:
+            assert_dep = assert_deps[task["name"]]
+            self.assertEqual(task.get("depends"), assert_dep)
+
+    def test_raise(self):
+        with self.assertRaises(ValueError):
+
+            def some_raise():
+                raise ValueError("test this")
+
+            couler.set_exit_handler(couler.WFStatus.Failed, some_raise)
+        couler._cleanup()
+        couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
+        content = pyaml.dump(couler.workflow_yaml())
+        self.assertNotIn("onExit", content)
+
+    # TODO: Fix exit_handler
+    # def test_set_dependencies_none2(self):
+    #     couler.set_dependencies(lambda: job_a(message="A"),
+    # dependencies=None)
+    #     couler.set_dependencies(lambda: job_b(message="B"),
+    # dependencies=["A"])
+
+    #     def job_exit():
+    #         return couler.run_container(
+    #             image="docker/whalesay:latest",
+    #             command=["cowsay"],
+    #             step_name="C",
+    #         )
+
+    #     couler.set_exit_handler(couler.WFStatus.Failed, job_exit)
+    #     content = pyaml.dump(couler.workflow_yaml())
+    #     self.assertIn("{{workflow.state}}", content)
+
     # TODO: Provide new test case without `tf.train`.
     # def test_set_dependencies_for_job(self):
     #
diff --git a/couler/tests/env_test.py b/couler/tests/env_test.py
index dd40502..01c2bf4 100644
--- a/couler/tests/env_test.py
+++ b/couler/tests/env_test.py
@@ -1,12 +1,12 @@
 import copy
-import unittest
 
 import couler.argo as couler
+from couler.tests.argo_test import ArgoBaseTestCase
 
 
-class EnvUnitTest(unittest.TestCase):
+class EnvUnitTest(ArgoBaseTestCase):
     def setUp(self):
-        couler._cleanup()
+        super().setUp()
         self.envs = {"str_env": "abc", "bool_env": False, "num_env": 1234}
 
     def run_a_container(self):
diff --git a/couler/tests/katib_step_test.py b/couler/tests/katib_step_test.py
index 2be8923..47a6a9e 100644
--- a/couler/tests/katib_step_test.py
+++ b/couler/tests/katib_step_test.py
@@ -1,7 +1,6 @@
-import unittest
-
 import couler.argo as couler
 import couler.steps.katib as katib
+from couler.tests.argo_test import ArgoBaseTestCase
 
 xgb_cmd = """ \
 python /opt/katib/train_xgboost.py \
@@ -36,7 +35,7 @@ xgboost_mainifest_template = """
 )
 
 
-class KatibTest(unittest.TestCase):
+class KatibTest(ArgoBaseTestCase):
     def test_katib_with_xgboost_training(self):
         katib.run(
             tuning_params=[
diff --git a/couler/tests/utils_test.py b/couler/tests/utils_test.py
index 2f522e9..1bd3fb8 100644
--- a/couler/tests/utils_test.py
+++ b/couler/tests/utils_test.py
@@ -1,10 +1,10 @@
 import base64
-import unittest
 
 from couler.core import utils
+from couler.tests.argo_test import ArgoBaseTestCase
 
 
-class PyfuncTest(unittest.TestCase):
+class PyfuncTest(ArgoBaseTestCase):
     def test_argo_safe_name(self):
         self.assertIsNone(utils.argo_safe_name(None))
         self.assertEqual(utils.argo_safe_name("a_b"), "a-b")