Repository: pypr/automan

Issue:
Add ability to specify inter-simulation dependencies
With tasks one can have complex dependencies but there isn't an easy way to spell this out when creating simulations.

Comments:

Patch:
diff --git a/automan/automation.py b/automan/automation.py
index bf7b07f..bd2f330 100644
--- a/automan/automation.py
+++ b/automan/automation.py
@@ -39,7 +39,7 @@ class Task(object):
         """Return iterable of tasks this task requires.
 
         It is important that one either return tasks that are idempotent or
-        return the same instance as this method is called repeateadly.
+        return the same instance as this method is called repeatedly.
 
         """
         return []
@@ -67,6 +67,8 @@ class TaskRunner(object):
         self.scheduler = scheduler
         self.todo = []
         self.task_status = dict()
+        self.task_outputs = set()
+        self.repeat_tasks = set()
         for task in tasks:
             self.add_task(task)
 
@@ -93,7 +95,22 @@ class TaskRunner(object):
         return complete
 
     def _get_tasks_with_status(self, status):
-        return [t for t, s in self.task_status.items() if s == status]
+        return [
+            t for t, s in self.task_status.items()
+            if s == status and t not in self.repeat_tasks
+        ]
+
+    def _is_output_registered(self, task):
+        # Note, this has a side-effect of registering the task's output
+        # when called.
+        output = task.output()
+        output_str = str(output)
+        if output and output_str in self.task_outputs:
+            self.repeat_tasks.add(task)
+            return True
+        else:
+            self.task_outputs.add(output_str)
+            return False
 
     def _run(self, task):
         try:
@@ -130,6 +147,11 @@ class TaskRunner(object):
     # #### Public protocol  ##############################################
 
     def add_task(self, task):
+        if task in self.task_status or self._is_output_registered(task):
+            # This task is already added or another task produces exactly
+            # the same output, so do nothing.
+            return
+
         if not task.complete():
             self.todo.append(task)
             self.task_status[task] = 'not started'
@@ -165,14 +187,14 @@ class TaskRunner(object):
 
 
 class CommandTask(Task):
-    """Convenience class to run a command via the framework. The class provides a
-    method to run the simulation and also check if the simulation is completed.
-    The command should ideally produce all of its outputs inside an output
-    directory that is specified.
+    """Convenience class to run a command via the framework. The class provides
+    a method to run the simulation and also check if the simulation is
+    completed. The command should ideally produce all of its outputs inside an
+    output directory that is specified.
 
     """
 
-    def __init__(self, command, output_dir, job_info=None):
+    def __init__(self, command, output_dir, job_info=None, depends=None):
         """Constructor
 
         **Parameters**
@@ -180,6 +202,7 @@ class CommandTask(Task):
         command: str or list: command to run $output_dir is substituted.
         output_dir: str : path of output directory.
         job_info: dict: dictionary of job information.
+        depends: list: list of tasks this depends on.
 
         """
         if isinstance(command, str):
@@ -190,6 +213,7 @@ class CommandTask(Task):
                         for x in self.command]
         self.output_dir = output_dir
         self.job_info = job_info if job_info is not None else {}
+        self.depends = depends if depends is not None else []
         self.job_proxy = None
         self._copy_proc = None
         # This is a sentinel set to true when the job is finished
@@ -231,6 +255,14 @@ class CommandTask(Task):
         if os.path.exists(self.output_dir):
             shutil.rmtree(self.output_dir)
 
+    def output(self):
+        """Return list of output paths.
+        """
+        return [self.output_dir]
+
+    def requires(self):
+        return self.depends
+
     # #### Private protocol ###########################################
 
     @property
@@ -299,7 +331,7 @@ class PySPHTask(CommandTask):
 
     """
 
-    def __init__(self, command, output_dir, job_info=None):
+    def __init__(self, command, output_dir, job_info=None, depends=None):
         """Constructor
 
         **Parameters**
@@ -307,9 +339,10 @@ class PySPHTask(CommandTask):
         command: str or list: command to run $output_dir is substituted.
         output_dir: str : path of output directory.
         job_info: dict: dictionary of job information.
+        depends: list: list of tasks this depends on.
 
         """
-        super(PySPHTask, self).__init__(command, output_dir, job_info)
+        super(PySPHTask, self).__init__(command, output_dir, job_info, depends)
         self.command += ['-d', output_dir]
 
     # #### Private protocol ###########################################
@@ -351,8 +384,8 @@ class Problem(object):
        results and simulations are collected inside a directory with
        this name.
      - `get_commands(self)`: returns a sequence of (directory_name,
-       command_string, job_info) tuples.  These are to be exeuted before the
-       `run` method is called.
+       command_string, job_info, depends) tuples. These are to be executed
+       before the `run` method is called.
      - `get_requires(self)`: returns a sequence of (name, task) tuples. These
        are to be exeuted before the `run` method is called.
      - `run(self)`: Processes the completed simulations to make plots etc.
@@ -379,6 +412,32 @@ class Problem(object):
         self.cases = None
         self.setup()
 
+    def _make_depends(self, depends):
+        if not depends:
+            return []
+        deps = []
+        for x in depends:
+            if isinstance(x, Task):
+                deps.append(x)
+            elif isinstance(x, Simulation):
+                if x.depends:
+                    my_depends = self._make_depends(x.depends)
+                else:
+                    my_depends = None
+                task = self.task_cls(
+                    x.command, self.input_path(x.name), x.job_info,
+                    depends=my_depends
+                )
+                deps.append(task)
+            else:
+                raise RuntimeError(
+                    'Invalid dependency: {0} for problem {1}'.format(
+                        x, self
+                    )
+                )
+
+        return deps
+
     # #### Public protocol ###########################################
 
     def input_path(self, *args):
@@ -413,20 +472,27 @@ class Problem(object):
         return self.__class__.__name__
 
     def get_commands(self):
-        """Return a sequence of (name, command_string, job_info_dict).
+        """Return a sequence of (name, command_string, job_info_dict)
+        or (name, command_string, job_info_dict, depends).
 
-         The name represents the command being run and is used as
-         a subdirectory for generated output.
+        The name represents the command being run and is used as a subdirectory
+        for generated output.
 
         The command_string is the command that needs to be run.
 
         The job_info_dict is a dictionary with any additional info to be used
         by the job, these are additional arguments to the
-        `automan.jobss.Job` class. It may be None if nothing special need
+        `automan.jobs.Job` class. It may be None if nothing special need
         be passed.
+
+        The depends is any dependencies this simulation has in terms of other
+        simulations/tasks.
+
         """
         if self.cases is not None:
-            return [(x.name, x.command, x.job_info) for x in self.cases]
+            return [
+                (x.name, x.command, x.job_info, x.depends) for x in self.cases
+            ]
         else:
             return []
 
@@ -440,9 +506,14 @@ class Problem(object):
         """
         base = self.get_name()
         result = []
-        for name, cmd, job_info in self.get_commands():
+        for cmd_info in self.get_commands():
+            name, cmd, job_info = cmd_info[:3]
+            deps = cmd_info[3] if len(cmd_info) == 4 else []
             sim_output_dir = self.input_path(name)
-            task = self.task_cls(cmd, sim_output_dir, job_info)
+            depends = self._make_depends(deps)
+            task = self.task_cls(
+                cmd, sim_output_dir, job_info, depends=depends
+            )
             task_name = '%s.%s' % (base, name)
             result.append((task_name, task))
         return result
@@ -457,7 +528,7 @@ class Problem(object):
         """Run any analysis code for the simulations completed.  This
         is usually run after the simulation commands are completed.
         """
-        raise NotImplementedError()
+        pass
 
     def clean(self):
         """Cleanup any generated output from the analysis code.  This does not
@@ -549,7 +620,7 @@ class Simulation(object):
     this is an extremely powerful way to automate and compare results.
 
     """
-    def __init__(self, root, base_command, job_info=None, **kw):
+    def __init__(self, root, base_command, job_info=None, depends=None, **kw):
         """Constructor
 
         **Parameters**
@@ -560,6 +631,8 @@ class Simulation(object):
             Base command to run.
         job_info: dict
             Extra arguments to the `automan.jobs.Job` class.
+        depends: list
+            List of other simulations/tasks this simulation depends on.
         **kw: dict
             Additional parameters to pass to command.
         """
@@ -567,6 +640,7 @@ class Simulation(object):
         self.name = os.path.basename(root)
         self.base_command = base_command
         self.job_info = job_info
+        self.depends = depends if depends is not None else []
         self.params = dict(kw)
         self._results = None
 
@@ -699,11 +773,26 @@ class SolveProblem(Task):
         self.problem = problem
         self.match = match
         self._requires = [
-            task
+            self._make_task(task)
             for name, task in self.problem.get_requires()
             if len(match) == 0 or fnmatch(name, match)
         ]
 
+    def _make_task(self, obj):
+        if isinstance(obj, Task):
+            return obj
+        elif isinstance(obj, Problem):
+            return SolveProblem(problem=obj, match=self.match)
+        elif isinstance(obj, type) and issubclass(obj, Problem):
+            problem = obj(self.problem.sim_dir, self.problem.out_dir)
+            return SolveProblem(problem=problem, match=self.match)
+        else:
+            raise RuntimeError(
+                'Unknown requirement: {0}, for problem: {1}.'.format(
+                    obj, self.problem
+                )
+            )
+
     def __str__(self):
         return 'Problem named %s' % self.problem.get_name()
 
@@ -846,8 +935,8 @@ class Automator(object):
                     self.parser.exit(1)
 
     def _get_exclude_paths(self):
-        """Returns a list of exclude paths suitable for passing on to rsync to exclude
-        syncing some directories on remote machines.
+        """Returns a list of exclude paths suitable for passing on to rsync to
+        exclude syncing some directories on remote machines.
         """
         paths = []
         for path in [self.simulation_dir, self.output_dir]:
diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst
index 55255f3..a65d9fc 100644
--- a/docs/source/tutorial.rst
+++ b/docs/source/tutorial.rst
@@ -84,16 +84,16 @@ Let us execute this to see what it does::
   Writing config.json
   4 tasks pending and 0 tasks running
 
-  Running task <automan.automation.CommandTask object at 0x10628d978>...
+  Running task CommandTask with output directory: 2 ...
   Starting worker on localhost.
   Job run by localhost
   Running python square.py 2
 
-  Running task <automan.automation.CommandTask object at 0x10628d9b0>...
+  Running task CommandTask with output directory: 1 ...
   Job run by localhost
   Running python square.py 1
   2 tasks pending and 2 tasks running
-  Running task <automan.automation.SolveProblem object at 0x10628d940>...
+  Running task Problem named squares...
 
   Running task <automan.automation.RunAll object at 0x10628d908>...
   Finished!
@@ -499,12 +499,12 @@ availability.  For example::
    $ python automate4.py
    14 tasks pending and 0 tasks running
 
-   Running task <automan.automation.CommandTask object at 0x1141da748>...
+   Running task CommandTask with output directory: 4 ...
    Starting worker on localhost.
    Job run by localhost
    Running python powers.py --output-dir outputs/powers/4 --power=4.0
 
-   Running task <automan.automation.CommandTask object at 0x1141da6d8>...
+   Running task CommandTask with output directory: 3 ...
    Starting worker on 10.1.10.242.
    Job run by 10.1.10.242
    Running python powers.py --output-dir outputs/powers/3 --power=3.0
@@ -653,6 +653,76 @@ https://github.com/pypr/automan/tree/master/examples/edm_conda_cluster
 The README in the directory tells you how to run the examples.
 
 
+Specifying simulation dependencies
+-----------------------------------
+
+There are times when one simulation uses the output from another and you wish
+to execute them in the right order. This can be quite easily achieved. Here is
+a simple example from the test suite that illustrates this::
+
+  class MyProblem(Problem):
+      def setup(self):
+          cmd = 'python -c "import time; print(time.time())"'
+          s1 = Simulation(self.input_path('1'), cmd)
+          s2 = Simulation(self.input_path('2'), cmd, depends=[s1])
+          s3 = Simulation(self.input_path('3'), cmd, depends=[s1, s2])
+          self.cases = [s1, s2, s3]
+
+Notice the extra keyword argument, ``depends=`` which specifies a list of
+other simulations. In the above case, we could have also used ``self.cases =
+[s3]`` and that would have automatically picked up the other simulations.
+
+When this problem is run, ``s1`` will run first followed by ``s2`` and then by
+``s3``. Note that this will only execute ``s1`` once even though it is
+declared as a dependency for two other simulations. This makes it possible to
+easily define inter-dependent tasks/simulations. In general, the dependencies
+could be any :py:class:`automan.automation.Simulation` or
+:py:class:`automan.automation.Task` instance.
+
+Internally, these simulations create suitable task instances that support
+dependencies see :py:class:`automan.automation.CommandTask`
+
+
+Specifying inter-problem dependencies
+--------------------------------------
+
+Sometimes you may have a situation where one problem depends on the output of
+another. These may be done by overriding the ``Problem.get_requires`` method.
+Here is an example from the test suite::
+
+  class A(Problem):
+      def get_requires(self):
+          cmd = 'python -c "print(1)"'
+          ct = CommandTask(cmd, output_dir=self.sim_dir)
+          return [('task1', ct)]
+
+  class B(Problem):
+      def get_requires(self):
+          # or return Problem instances ...
+          return [('a', A(self.sim_dir, self.out_dir))]
+
+  class C(Problem):
+      def get_requires(self):
+          # ... or Problem subclasses
+          return [('a', A), ('b', B)]
+
+Normally, the ``get_requires`` method automatically creates tasks from the
+simulations specified but in the above example we show a case (problem ``A``)
+where we explicitly create command tasks. In the above example, the problem
+``B`` depends on the problem ``A`` and simply returns an instance of ``A``. On
+the other hand ``C`` only returns the problem class and not an instance. This
+shows how one can specify inter problem dependencies.
+
+Note that if the problem performs some simulations (by setting
+``self.cases``), you should call the parent method (via ``super``) and add
+your additional dependencies to this.
+
+Also note that the dependencies are resolved based on the "outputs" of a task.
+So two tasks with the same outputs are treated as the same. This is consistent
+with the design of automan where each simulation's output goes in its own
+directory.
+
+
 Using docker
 ------------
 


Test patch:
diff --git a/automan/tests/test_automation.py b/automan/tests/test_automation.py
index 4e07c93..f544148 100644
--- a/automan/tests/test_automation.py
+++ b/automan/tests/test_automation.py
@@ -1,7 +1,6 @@
 from __future__ import print_function
 
 import os
-import shutil
 import sys
 import tempfile
 import unittest
@@ -12,8 +11,8 @@ except ImportError:
     import mock
 
 from automan.automation import (
-    Automator, CommandTask, PySPHProblem, Simulation, SolveProblem,
-    TaskRunner, compare_runs, filter_cases
+    Automator, CommandTask, Problem, PySPHProblem, RunAll, Simulation,
+    SolveProblem, TaskRunner, compare_runs, filter_cases
 )
 try:
     from automan.jobs import Scheduler, RemoteWorker
@@ -35,8 +34,8 @@ class MySimulation(Simulation):
 
 
 class EllipticalDrop(PySPHProblem):
-    """We define a simple example problem which we will run using the automation
-    framework.
+    """We define a simple example problem which we will run using the
+    automation framework.
 
     In this case we run two variants of the elliptical drop problem.
 
@@ -102,6 +101,168 @@ class TestAutomationBase(unittest.TestCase):
             safe_rmtree(self.root)
 
 
+class TestTaskRunner(TestAutomationBase):
+    def _make_scheduler(self):
+        worker = dict(host='localhost')
+        s = Scheduler(root='.', worker_config=[worker])
+        return s
+
+    def _get_time(self, path):
+        with open(os.path.join(path, 'stdout.txt')) as f:
+            t = float(f.read())
+        return t
+
+    def test_task_runner_does_not_add_repeated_tasks(self):
+        # Given
+        s = self._make_scheduler()
+        cmd = 'python -c "print(1)"'
+        ct1 = CommandTask(cmd, output_dir=self.sim_dir)
+        ct2 = CommandTask(cmd, output_dir=self.sim_dir)
+
+        # When
+        t = TaskRunner(tasks=[ct1, ct2, ct1], scheduler=s)
+
+        # Then
+        self.assertEqual(len(t.todo), 1)
+
+    def test_problem_depending_on_other_problems(self):
+        # Given
+        class A(Problem):
+            def get_requires(self):
+                cmd = 'python -c "print(1)"'
+                # Can return tasks ...
+                ct = CommandTask(cmd, output_dir=self.sim_dir)
+                return [('task1', ct)]
+
+            def run(self):
+                self.make_output_dir()
+
+        class B(Problem):
+            def get_requires(self):
+                # or return Problem instances ...
+                return [('a', A(self.sim_dir, self.out_dir))]
+
+            def run(self):
+                self.make_output_dir()
+
+        class C(Problem):
+            def get_requires(self):
+                # ... or Problem subclasses
+                return [('a', A), ('b', B)]
+
+            def run(self):
+                self.make_output_dir()
+
+        s = self._make_scheduler()
+
+        # When
+        task = RunAll(
+            simulation_dir=self.sim_dir, output_dir=self.output_dir,
+            problem_classes=[A, B, C]
+        )
+        t = TaskRunner(tasks=[task], scheduler=s)
+
+        # Then
+        self.assertEqual(len(t.todo), 5)
+        # Basically only one instance of CommandTask should be created.
+        names = [x.__class__.__name__ for x in t.todo]
+        problems = [x.problem for x in t.todo if isinstance(x, SolveProblem)]
+        self.assertEqual(names.count('RunAll'), 1)
+        self.assertEqual(names.count('CommandTask'), 1)
+        self.assertEqual(names.count('SolveProblem'), 3)
+        self.assertEqual(len(problems), 3)
+        self.assertEqual(
+            sorted(x.__class__.__name__ for x in problems),
+            ['A', 'B', 'C']
+        )
+
+        # When
+        t.run(wait=0.1)
+
+        # Then.
+        self.assertEqual(t.todo, [])
+
+    def test_problem_with_bad_requires_raises_error(self):
+        # Given
+        class D(Problem):
+            def get_requires(self):
+                return [('a', 'A')]
+
+        # When
+        self.assertRaises(
+            RuntimeError,
+            SolveProblem, D(self.sim_dir, self.output_dir)
+        )
+
+    def test_tasks_with_dependencies(self):
+        # Given
+        s = self._make_scheduler()
+        cmd = 'python -c "import time; print(time.time())"'
+        ct1_dir = os.path.join(self.sim_dir, '1')
+        ct2_dir = os.path.join(self.sim_dir, '2')
+        ct3_dir = os.path.join(self.sim_dir, '3')
+        ct1 = CommandTask(cmd, output_dir=ct1_dir)
+        ct2 = CommandTask(cmd, output_dir=ct2_dir, depends=[ct1])
+        ct3 = CommandTask(cmd, output_dir=ct3_dir, depends=[ct1, ct2])
+
+        # When
+        t = TaskRunner(tasks=[ct1, ct2, ct3], scheduler=s)
+
+        # Then
+        self.assertEqual(len(t.todo), 3)
+
+        # When
+        t.run(wait=0.1)
+
+        wait_until(lambda: not ct3.complete())
+
+        # Then.
+        # Ensure that the tasks are run in the right order.
+        ct1_t, ct2_t, ct3_t = [
+            self._get_time(x) for x in (ct1_dir, ct2_dir, ct3_dir)
+        ]
+        self.assertTrue(ct2_t > ct1_t)
+        self.assertTrue(ct3_t > ct2_t)
+
+    def test_simulation_with_dependencies(self):
+        # Given
+        class A(Problem):
+            def setup(self):
+                cmd = 'python -c "import time; print(time.time())"'
+                s1 = Simulation(self.input_path('1'), cmd)
+                s2 = Simulation(self.input_path('2'), cmd, depends=[s1])
+                s3 = Simulation(self.input_path('3'), cmd, depends=[s1, s2])
+                self.cases = [s1, s2, s3]
+
+            def run(self):
+                self.make_output_dir()
+
+        s = self._make_scheduler()
+
+        # When
+        problem = A(self.sim_dir, self.output_dir)
+        task = SolveProblem(problem)
+        t = TaskRunner(tasks=[task], scheduler=s)
+
+        # Then
+        self.assertEqual(len(t.todo), 4)
+        # Basically only one instance of CommandTask should be created.
+        names = [x.__class__.__name__ for x in t.todo]
+        self.assertEqual(names.count('CommandTask'), 3)
+        self.assertEqual(names.count('SolveProblem'), 1)
+
+        # When
+        t.run(wait=0.1)
+        wait_until(lambda: not task.complete())
+
+        # Then
+        ct1_t, ct2_t, ct3_t = [
+            self._get_time(problem.input_path(x)) for x in ('1', '2', '3')
+        ]
+        self.assertTrue(ct2_t > ct1_t)
+        self.assertTrue(ct3_t > ct2_t)
+
+
 class TestLocalAutomation(TestAutomationBase):
     def _make_scheduler(self):
         worker = dict(host='localhost')