fix some bug

Author: FlyPythons

DAGflow/DAG.py

@@ -125,7 +125,7 @@ def __init__(self, id, script, work_dir=".", type="sge", option=""):
         self.script = script
         self.type = type
         self._option = option
-        self.done = os.path.join(work_dir, "%s_done" % id)
+        self.done = os.path.join(self.work_dir, "%s_done" % id)
 
         self.depends = []
         self.status = None
@@ -157,9 +157,9 @@ def run_time(self):
         """
         if self.end_time and self.start_time:
             _time = self.end_time - self.start_time
-            return "%ss" % _time
+            return "%s" % int(_time)
         else:
-            return "0s"
+            return "0"
 
     """
     Functions to describe relationship between tasks
@@ -298,15 +298,21 @@ def kill(self):
         return 1
 
     def check_done(self):
+        """
+        check the status of done task
+        :return: success 1 or fail 0
+        """
         if os.path.isfile(self.done):
             self.status = "success"
             self.end_time = time.time()
             LOG.info("task %r finished by %s seconds" % (self.id, self.run_time))
+
+            return 1
         else:
             self.status = "failed"
             LOG.info("task %r run but failed" % self.id)
 
-        return 1
+            return 0
 
     """
     methods to read and write Task from json

DAGflow/do_DAG.py

@@ -23,7 +23,7 @@
 import time
 import json
 import signal
-from .DAG import Task, DAG
+from DAGflow import DAG
 
 
 LOG = logging.getLogger(__name__)
@@ -98,7 +98,7 @@ def ps():
     return r
 
 
-def update_task_status(tasks):
+def update_task_status(tasks, stop_on_failure):
     """
 
     :param tasks:
@@ -133,12 +133,19 @@ def update_task_status(tasks):
 
         # check recent done tasks on sge
         if task.type == "sge" and task.run_id not in sge_running_task:
-            task.check_done()
+            status = task.check_done()
+
+            if not status and stop_on_failure:
+                LOG.info("Task %r failed, stop all tasks" % task.id)
+                del_online_tasks()
             continue
         elif task.type == "local":
-            if task.run_id.poll():
-                task.check_done()
+            if not task.run_id.poll():
+                status = task.check_done()
 
+                if not status and stop_on_failure:
+                    LOG.info("Task %r failed, stop all tasks" % task.id)
+                    del_online_tasks()
             continue
         else:
             pass
@@ -192,7 +199,11 @@ def submit_tasks(tasks, concurrent_tasks):
     return tasks
 
 
-def qdel_online_tasks(signum, frame):
+def del_task_hander(signum, frame):
+    del_online_tasks()
+
+
+def del_online_tasks():
     LOG.info("delete all running jobs, please wait")
     time.sleep(3)
 
@@ -201,44 +212,37 @@ def qdel_online_tasks(signum, frame):
         if task.status == "running":
             task.kill()
 
-    write_tasks(TASKS, TASK_NAME + ".json")
+    write_tasks(TASKS)
 
     sys.exit("sorry, the program exit")
 
 
-def write_tasks(tasks, filename):
+def write_tasks(tasks):
     failed_tasks = []
 
-    tasks_json = OrderedDict()
-
     for id, task in tasks.items():
 
         if task.status != "success":
             failed_tasks.append(task.id)
 
-        tasks_json.update(task.to_json())
-
-    with open(filename, "w") as out:
-        json.dump(tasks_json, out, indent=2)
-
     if failed_tasks:
         LOG.info("""\
 The following tasks were failed:
 %s
-The tasks were save in %s, you can resub it.
-                    """ % ("\n".join([i for i in failed_tasks]), filename))
+""" % "\n".join([i for i in failed_tasks]))
         sys.exit("sorry, the program exit with some jobs failed")
     else:
         LOG.info("All jobs were done!")
 
 
-def do_dag(dag, concurrent_tasks, refresh_time, log_name=""):
+def do_dag(dag, concurrent_tasks=200, refresh_time=60, stop_on_failure=False):
 
+    dag.to_json()
     start = time.time()
 
     logging.basicConfig(level=logging.DEBUG,
                         format="[%(levelname)s] %(asctime)s  %(message)s",
-                        filename=log_name,
+                        filename="%s.log" % dag.id,
                         filemode='w',
                         )
 
@@ -253,15 +257,13 @@ def do_dag(dag, concurrent_tasks, refresh_time, log_name=""):
     global TASKS
     TASKS = dag.tasks
 
-    signal.signal(signal.SIGINT, qdel_online_tasks)
-    signal.signal(signal.SIGTERM, qdel_online_tasks)
+    signal.signal(signal.SIGINT, del_task_hander)
+    signal.signal(signal.SIGTERM, del_task_hander)
     # signal.signal(signal.SIGKILL, qdel_online_tasks)
 
     for id, task in TASKS.items():
         task.init()
 
-    failed_json = TASK_NAME + ".json"
-
     loop = 0
 
     while 1:
@@ -292,10 +294,10 @@ def do_dag(dag, concurrent_tasks, refresh_time, log_name=""):
         else:
             time.sleep(refresh_time)
             loop += 1
-            update_task_status(TASKS)
+            update_task_status(TASKS, stop_on_failure)
 
     # write failed
-    write_tasks(TASKS, failed_json)
+    write_tasks(TASKS)
     totalTime = time.time() - start
     LOG.info('Total time:' + time.strftime("%H:%M:%S", time.gmtime(totalTime)))
 
@@ -314,7 +316,8 @@ def get_args():
 
     parser.add_argument("json",  help="The json file contain DAG information")
     parser.add_argument("-m", "--max_task", type=int, default=200, help="concurrent_tasks")
-    parser.add_argument("-r", "--refresh", type=int, default=30, help="refresh time of task status (seconds)")
+    parser.add_argument("-r", "--refresh", type=int, default=60, help="refresh time of task status (seconds)")
+    parser.add_argument("-s", "--stopOnFailure", action="store_true", help="stop all tasks when any task failure")
     args = parser.parse_args()
 
     return args
@@ -327,7 +330,7 @@ def main():
     TASK_NAME = os.path.splitext(os.path.basename(args.json))[0]
     print(TASK_NAME)
     dag = DAG.from_json(args.json)
-    do_dag(dag, args.max_task, args.refresh, TASK_NAME+".log")
+    do_dag(dag, args.max_task, args.refresh, args.stopOnFailure)
 
 
 if __name__ == "__main__":

README.md

@@ -36,7 +36,8 @@ To complete this work, a workflow as following is needed
 At first, you should write your workflow script 
 ```python
 import os
-from DAGflow.DAG import DAG, Task
+from DAGflow import DAG, Task
+from DAGflow.do_DAG import do_dag
 
 
 inputs = ['1.fasta', "2.fasta", "3.fasta", "4.fasta"]
@@ -47,7 +48,7 @@ db = os.path.abspath(db)
 my_dag = DAG("blast")
 # create the first task 'make_db'
 make_db = Task(
-    task_id="make_db",  # your task id, should be unique
+    id="make_db",  # your task id, should be unique
     work_dir="",  # you task work directory
     type="local",  # the way your task run. if "sge", task will submit with qsub
     option={},  # the option of "sge" or "local"
@@ -64,7 +65,7 @@ n = 1
 for fn in inputs:
     task_id = "blast_%s" % n
     task = Task(
-        task_id= task_id,
+        id= task_id,
         work_dir=task_id,
         type="sge", 
         option={
@@ -82,7 +83,7 @@ for fn in inputs:
 
 # add blast_join task to join blast results
 blast_join = Task(
-    task_id="blast_join",
+    id="blast_join",
     work_dir="",
     type="local",  # option is default
     script="cat */*.m6 > blast.all.m6"
@@ -93,10 +94,8 @@ my_dag.add_task(blast_join)
 blast_join.set_upstream(*blast_tasks)
 
 # all of you tasks were added to you workflow, you can run it
-# write you workflow to a json file
-js = my_dag.to_json()
-# submit you workflow tasks with do_DAG
-os.system("python -m DAGflow.do_DAG %s" % js)
+do_dag(my_dag)
+
 ```
 Now, your workflow script is completed, you can name it as 'workflow.py'
 ### Run you workflow 
@@ -106,15 +105,11 @@ python workflow.py
 ```
 ### Re-run your workflow if it was break in the middle
 For some reason, you workflow was broken with some tasks undone.  
-You can use the following commands to re-run the undone jobs.
-```commandline
-python -m DAGflow.do_DAG blast.json 
-# note that the blast.json is in you work directory, 'blast' is your DAG id.
-```
+You can use the same command `python workflow.py`  to re-run the undone jobs.
 ### Add workflow to workflow
 Sometimes you may want to add a workflow to another workflow, this can be down as following:  
 ```python
-from DAGflow.DAG import *
+from DAGflow import *
 
 
 # two workflow wf1 and wf2

test/test_add_dag.py

@@ -4,39 +4,26 @@
 """
 
 
-from DAGflow.DAG import *
+from DAGflow import *
 
 workflow = DAG("test")
 
 task1 = Task(
-    task_id="task1",
+    id="task1",
     work_dir="",
     type="sge",
-    script="echo running workflow task 1"
+    script="echo running workflow task 1\nsleep 5"
 )
 
 workflow.add_task(task1)
 
 task2 = Task(
-    task_id="task2",
+    id="task2",
     work_dir="task2",
     type="local",
-    script="echo running workflow task 3"
+    script="echo running workflow task 2\n sleep 5"
 )
 workflow.add_task(task2)
 task2.set_upstream(task1)
 
-wf2 = DAG("test2")
-
-task3 = Task(
-    task_id="task3",
-    work_dir=".",
-    type="local",
-    script="echo running workflow task 3"
-)
-
-wf2.add_task(task3)
-
-workflow.add_dag(wf2)
-
-workflow.print_task()
+do_dag(workflow, 200, 10, "test")