Merge pull request #6 from abhinavmuta/edm_cluster

A simple edm cluster manager

automan/api.py

@@ -1,11 +1,12 @@
-from .jobs import Job, Worker, LocalWorker, RemoteWorker, Scheduler # noqa
+from .jobs import Job, Worker, LocalWorker, RemoteWorker, Scheduler  # noqa
 
 from .automation import (  # noqa
     Automator, CommandTask, Problem, PySPHProblem, PySPHTask, RunAll,
     Simulation, SolveProblem, Task, TaskRunner, WrapperTask
 )
 
-from .automation import compare_runs, filter_by_name, filter_cases # noqa
+from .automation import compare_runs, filter_by_name, filter_cases  # noqa
 
-from .cluster_manager import ClusterManager # noqa
-from .conda_cluster_manager import CondaClusterManager # noqa
+from .cluster_manager import ClusterManager  # noqa
+from .conda_cluster_manager import CondaClusterManager  # noqa
+from .edm_cluster_manager import EDMClusterManager  # noqa

automan/edm_cluster_manager.py

@@ -0,0 +1,91 @@
+import os
+from textwrap import dedent
+
+from .cluster_manager import ClusterManager
+
+
+class EDMClusterManager(ClusterManager):
+
+    # The path to edm root on the remote, this is a relative path
+    # and is relative to the home directory.
+    EDM_ROOT = '.edm'
+
+    ENV_FILE = "bundled_env.json"
+
+    BOOTSTRAP = dedent("""\
+    #!/bin/bash
+
+    set -e
+    ENV_FILE="{project_name}/{env_file}"
+
+    if hash edm 2>/dev/null; then
+        EDM_EXE=edm
+    else
+        EDM_EXE=~/{edm_root}/bin/edm
+    fi
+
+    if [ -f $ENV_FILE ] ; then
+        $EDM_EXE -q envs import --force {project_name} -f $ENV_FILE
+    else
+        $EDM_EXE -q envs create --force {project_name} --version 3.6
+        $EDM_EXE -q install psutil execnet -y -e {project_name}
+    fi
+
+    $EDM_EXE run -e {project_name} -- pip install automan
+
+    cd {project_name}
+    if [ -f "requirements.txt" ] ; then
+        $EDM_EXE run -e {project_name} -- pip install -r requirements.txt
+    fi
+    """)
+
+    UPDATE = dedent("""\
+    #!/bin/bash
+
+    set -e
+    ENV_FILE="{project_name}/{env_file}"
+
+    if hash edm 2>/dev/null; then
+        EDM_EXE=edm
+    else
+        EDM_EXE=~/{edm_root}/bin/edm
+    fi
+
+    if [ -f $ENV_FILE ] ; then
+        $EDM_EXE -q envs import --force {project_name} -f $ENV_FILE
+    fi
+
+    cd {project_name}
+    if [ -f "requirements.txt" ] ; then
+        $EDM_EXE run -e {project_name} -- pip install -r requirements.txt
+    fi
+    """)
+
+    def _get_bootstrap_code(self):
+        return self.BOOTSTRAP.format(
+            project_name=self.project_name, edm_root=self.EDM_ROOT,
+            env_file=self.ENV_FILE
+        )
+
+    def _get_python(self, host, home):
+        return os.path.join(
+            home, self.EDM_ROOT,
+            'envs/{project_name}/bin/python'.format(
+                project_name=self.project_name
+            )
+        )
+
+    def _get_update_code(self):
+        return self.UPDATE.format(
+            project_name=self.project_name, edm_root=self.EDM_ROOT,
+            env_file=self.ENV_FILE
+        )
+
+    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.
+
+        """
+        return ''

automan/tests/test_cluster_manager.py

@@ -17,6 +17,7 @@
 from automan.jobs import Job
 from automan.cluster_manager import ClusterManager
 from automan.conda_cluster_manager import CondaClusterManager
+from automan.edm_cluster_manager import EDMClusterManager
 from .test_jobs import wait_until
 
 
@@ -218,3 +219,36 @@ def test_overloaded_methods(self):
         self.assertTrue('TEST_ROOT' in code)
 
         self.assertEqual(cm._get_helper_scripts(), '')
+
+
+class TestEDMClusterManager(unittest.TestCase):
+    def setUp(self):
+        self.cwd = os.getcwd()
+        self.root = tempfile.mkdtemp()
+        os.chdir(self.root)
+
+    def tearDown(self):
+        os.chdir(self.cwd)
+        if os.path.exists(self.root):
+            shutil.rmtree(self.root)
+
+    @mock.patch("automan.edm_cluster_manager.EDMClusterManager.EDM_ROOT",
+                'TEST_ROOT')
+    def test_overloaded_methods(self):
+        # Given
+        cm = EDMClusterManager()
+
+        # When/Then
+        python = cm._get_python('foo', 'bar')
+        name = os.path.basename(self.root)
+        self.assertEqual(python, os.path.join(
+            'bar', 'TEST_ROOT', 'envs/{name}/bin/python'.format(name=name)
+        ))
+
+        code = cm._get_bootstrap_code()
+        self.assertTrue('TEST_ROOT' in code)
+
+        code = cm._get_update_code()
+        self.assertTrue('TEST_ROOT' in code)
+
+        self.assertEqual(cm._get_helper_scripts(), '')

docs/source/reference/automan.rst

@@ -23,3 +23,7 @@ Cluster management module
 .. automodule:: automan.conda_cluster_manager
    :members:
    :undoc-members:
+
+.. automodule:: automan.edm_cluster_manager
+   :members:
+   :undoc-members:

docs/source/tutorial.rst

@@ -599,13 +599,19 @@ to the :py:class:`automan.automation.Automator` class as the
 you wish to use conda or some other tool to manage the Python environment on
 the remote computer.
 
-We also provide a simple
-:py:class:`automan.conda_cluster_manager.CondaClusterManager` which will setup
-a remote computer so long as it has conda_ on it. If your project directory
-has an ``environments.yml`` and/or a ``requirements.txt`` it will use those to
-setup the environment. This is really a prototype and you may feel free to
-customize this. To use the conda cluster manager you could do the following in
-the tutorial example::
+We provide two simple environment managers one is a based on anaconda's conda_
+and the other is on enthought's edm_, the following contains details on how to
+use them.
+
+.. _edm: https://docs.enthought.com/edm
+.. _conda: https://conda.io/
+
+A simple :py:class:`automan.conda_cluster_manager.CondaClusterManager` which
+will setup a remote computer so long as it has conda_ on it. If your project
+directory has an ``environments.yml`` and/or a ``requirements.txt`` it will use
+those to setup the environment. This is really a prototype and you may feel
+free to customize this. To use the conda cluster manager you could do the
+following in the tutorial example::
 
     from automan.api import CondaClusterManager
 
@@ -617,10 +623,33 @@ the tutorial example::
     )
     automator.run()
 
-You may also subclass this or customize the bootstrap code and use that.
-
 .. _conda: https://conda.io/
 
+A simple :py:class:`automan.edm_cluster_manager.EDMClusterManager` which will
+setup a remote computer so long as it has edm_ on it. If your project directory
+has an ``bundled_envs.json`` and/or a ``requirements.txt`` it will use those to
+setup the environment. You can change the file names by accessing ``ENV_FILE``
+class variable . By default this assumes the edm executable location to be in
+``~/.edm`` to change this point the ``EDM_ROOT`` variable to the correct
+location relative to ``~`` (the current user home folder) not including the
+symbol ``~``. To use the edm cluster manager you could do the following in the
+tutorial example::
+
+    from automan.api import EDMClusterManager
+
+    automator = Automator(
+        simulation_dir='outputs',
+        output_dir='manuscript/figures',
+        all_problems=[Squares, Powers],
+        cluster_manager_factory=EDMClusterManager
+    )
+    automator.run()
+
+You may also subclass these or customize the bootstrap code and use that.
+
+.. _edm: https://docs.enthought.com/edm
+
+
 
 Using docker
 ------------