Merge pull request #7 from pypr/generalize-cluster-manager

Generalize cluster manager

MANIFEST.in

@@ -1,5 +1,5 @@
 include MANIFEST.in *.py *.rst *.txt *.yml
 recursive-include docs *.*
-recursive-include examples *.py
+recursive-include examples *.py *.yml *.txt *.md
 recursive-exclude examples/tutorial/.automan *.*
-recursive-exclude docs/build *.*
+recursive-exclude docs/build *.*

automan/cluster_manager.py

@@ -148,7 +148,7 @@ def __init__(self, root='automan', sources=None,
 
     # ### Private Protocol ########################################
     def _bootstrap(self, host, home):
-        venv_script = self._get_virtualenv()
+        helper_scripts = self._get_helper_scripts()
         base_cmd = ("cd {home}; mkdir -p {root}/envs; "
                     "mkdir -p {root}/{project_name}/.{root}").format(
                         home=home, root=self.root,
@@ -157,10 +157,10 @@ def _bootstrap(self, host, home):
         self._ssh_run_command(host, base_cmd)
 
         abs_root = os.path.join(home, self.root)
-        if venv_script:
+        if helper_scripts:
             real_host = '' if self.testing else '{host}:'.format(host=host)
-            cmd = "scp {venv_script} {host}{root}".format(
-                host=real_host, root=abs_root, venv_script=venv_script
+            cmd = "scp {helper_scripts} {host}{root}".format(
+                host=real_host, root=abs_root, helper_scripts=helper_scripts
             )
             self._run_command(cmd)
 
@@ -202,7 +202,21 @@ def _bootstrap(self, host, home):
         else:
             print("Bootstrapping {host} succeeded!".format(host=host))
 
-    def _get_virtualenv(self):
+    def _get_python(self, host, home):
+        return os.path.join(
+            home, self.root,
+            'envs/{project_name}/bin/python'.format(
+                project_name=self.project_name
+            )
+        )
+
+    def _get_helper_scripts(self):
+        """Return a space separated string of script files that you need copied over to
+        the remote host.
+
+        When overriding this, you can return None or '' if you do not need any.
+
+        """
         script = os.path.join(self.scripts_dir, 'virtualenv.py')
         if not os.path.exists(script):
             print("Downloading latest virtualenv.py")
@@ -318,19 +332,13 @@ def add_worker(self, host, home, nfs):
         if host == 'localhost':
             self.workers.append(dict(host=host, home=home, nfs=nfs))
         else:
-            root = self.root
             curdir = os.path.basename(os.getcwd())
             if nfs:
                 python = sys.executable
                 chdir = curdir
             else:
-                python = os.path.join(
-                    home, root,
-                    'envs/{project_name}/bin/python'.format(
-                        project_name=self.project_name
-                    )
-                )
-                chdir = os.path.join(home, root, curdir)
+                python = self._get_python(host, home)
+                chdir = os.path.join(home, self.root, curdir)
             self.workers.append(
                 dict(host=host, home=home, nfs=nfs, python=python, chdir=chdir)
             )

automan/conda_cluster_manager.py

@@ -1,7 +1,5 @@
 import os
-import sys
 from textwrap import dedent
-import subprocess
 
 from .cluster_manager import ClusterManager
 
@@ -52,42 +50,19 @@ class CondaClusterManager(ClusterManager):  # pragma: no cover
 
     """ % CONDA_ROOT)
 
-    def _get_virtualenv(self):
-        return None
-
-    def _get_conda_env_root(self, host):
-        cmd = [
-            'ssh', host,
-            '~/{conda_root}/bin/conda info --base'.format(
-                conda_root=self.CONDA_ROOT
-            )
-        ]
-        path = subprocess.check_output(cmd).strip()
-        if type(path) is bytes:
-            path = path.decode('utf-8')
-        return path
-
-    def add_worker(self, host, home, nfs):
-        if host == 'localhost':
-            self.workers.append(dict(host=host, home=home, nfs=nfs))
-        else:
-            root = self.root
-            curdir = os.path.basename(os.getcwd())
-            if nfs:
-                python = sys.executable
-                chdir = curdir
-            else:
-                conda_root = self._get_conda_env_root(host)
-                python = os.path.join(
-                    conda_root, 'envs/{project_name}/bin/python'.format(
-                        project_name=self.project_name
-                    )
-                )
-                chdir = os.path.join(home, root, curdir)
-            self.workers.append(
-                dict(host=host, home=home, nfs=nfs, python=python, chdir=chdir)
+    def _get_python(self, host, home):
+        return os.path.join(
+            home, self.CONDA_ROOT,
+            'envs/{project_name}/bin/python'.format(
+                project_name=self.project_name
             )
+        )
+
+    def _get_helper_scripts(self):
+        """Return a space separated string of script files that you need copied over to
+        the remote host.
+
+        When overriding this, you can return None or '' if you do not need any.
 
-        self._write_config()
-        if host != 'localhost' and not nfs:
-            self._bootstrap(host, home)
+        """
+        return ''

examples/conda_cluster/README.md

@@ -0,0 +1,21 @@
+# Conda cluster manager example
+
+This uses the example from the tutorial but shows how it can be used
+with the remote computers running conda.  The `environments.yml` is used
+to create the environment and there is also an optional
+`requirements.txt` which will install packages using `pip`.
+
+The example assumes that the remote computer has its conda root in `~/miniconda3`.
+
+To run the example, let us say you have a remote computer running Linux/Mac OS
+with password-less SSH setup, let us call this computer `remote_host`, then
+you can do the following:
+
+```
+   $ python automate.py -a remote_host
+   [...]
+   Bootstrapping remote_host succeeded!
+
+   $ python automate.py
+   [...]
+```

examples/conda_cluster/automate.py

@@ -0,0 +1,69 @@
+from automan.api import Problem, Automator, Simulation
+from automan.conda_cluster_manager import CondaClusterManager
+from matplotlib import pyplot as plt
+import numpy as np
+
+
+class Squares(Problem):
+    def get_name(self):
+        return 'squares'
+
+    def get_commands(self):
+        commands = [(str(i), 'python square.py %d' % i, None)
+                    for i in range(1, 8)]
+        return commands
+
+    def run(self):
+        self.make_output_dir()
+        data = []
+        for i in range(1, 8):
+            stdout = self.input_path(str(i), 'stdout.txt')
+            with open(stdout) as f:
+                values = [float(x) for x in f.read().split()]
+                data.append(values)
+
+        data = np.asarray(data)
+        plt.plot(data[:, 0], data[:, 1], 'o-')
+        plt.xlabel('x')
+        plt.ylabel('y')
+        plt.savefig(self.output_path('squares.pdf'))
+
+
+class Powers(Problem):
+    def get_name(self):
+        return 'powers'
+
+    def setup(self):
+        base_cmd = 'python powers.py --output-dir $output_dir'
+        self.cases = [
+            Simulation(
+                root=self.input_path(str(i)),
+                base_command=base_cmd,
+                power=float(i)
+            )
+            for i in range(1, 5)
+        ]
+
+    def run(self):
+        self.make_output_dir()
+        for case in self.cases:
+            data = np.load(case.input_path('results.npz'))
+            plt.plot(
+                data['x'], data['y'],
+                label=r'$x^{{%.2f}}$' % case.params['power']
+            )
+        plt.grid()
+        plt.xlabel('x')
+        plt.ylabel('y')
+        plt.legend()
+        plt.savefig(self.output_path('powers.pdf'))
+
+
+if __name__ == '__main__':
+    automator = Automator(
+        simulation_dir='outputs',
+        output_dir='manuscript/figures',
+        all_problems=[Squares, Powers],
+        cluster_manager_factory=CondaClusterManager
+    )
+    automator.run()

examples/conda_cluster/environments.yml

@@ -0,0 +1,3 @@
+name: conda_cluster
+dependencies:
+  - numpy

examples/conda_cluster/powers.py

@@ -0,0 +1,41 @@
+import argparse
+import os
+
+import numpy as np
+
+
+def compute_powers(r_max, power):
+    """Compute the powers of the integers upto r_max and return the result.
+    """
+    result = []
+    for i in range(0, r_max + 1):
+        result.append((i, i**power))
+    x = np.arange(0, r_max + 1)
+    y = np.power(x, power)
+    return x, y
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        '--power', type=float, default=2.0,
+        help='Power to calculate'
+    )
+    p.add_argument(
+        '--max', type=int, default=10,
+        help='Maximum integer that we must raise to the given power'
+    )
+    p.add_argument(
+        '--output-dir', type=str, default='.',
+        help='Output directory to generate file.'
+    )
+    opts = p.parse_args()
+
+    x, y = compute_powers(opts.max, opts.power)
+
+    fname = os.path.join(opts.output_dir, 'results.npz')
+    np.savez(fname, x=x, y=y)
+
+
+if __name__ == '__main__':
+    main()

examples/conda_cluster/requirements.txt

@@ -0,0 +1 @@
+numpy

examples/conda_cluster/square.py

@@ -0,0 +1,5 @@
+from __future__ import print_function
+import sys
+x = float(sys.argv[1])
+print(x, x*x)
+

examples/tutorial/README.md

@@ -0,0 +1,5 @@
+# Tutorial examples
+
+This directory contains the scripts and automation scripts that are
+discussed in the tutorials available here:
+https://automan.readthedocs.io/en/latest/tutorial.html