Merge pull request #8 from pypr/fix-cluster-manager-scripts

Improve cluster manager extensibility.
pypr · Sep 2, 2018 · efa0e17 · efa0e17
2 parents 19333d5 + f0b0e7d
commit efa0e17
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 8 deletions.
diff --git a/automan/cluster_manager.py b/automan/cluster_manager.py
@@ -68,6 +68,11 @@ class ClusterManager(object):
     update scripts are put here and may be edited by the user for any new
     hosts.
 
+    One may override the `_get_python, _get_helper_scripts`, and
+    `_get_bootstrap_code, _get_update_code` methods to change this to use other
+    package managers like edm or conda. See the conda_cluster_manager for an
+    example.
+
     """
 
     #######################################################
@@ -202,6 +207,9 @@ def _bootstrap(self, host, home):
         else:
             print("Bootstrapping {host} succeeded!".format(host=host))
 
+    def _get_bootstrap_code(self):
+        return self.BOOTSTRAP.format(project_name=self.project_name)
+
     def _get_python(self, host, home):
         return os.path.join(
             home, self.root,
@@ -210,6 +218,9 @@ def _get_python(self, host, home):
             )
         )
 
+    def _get_update_code(self):
+        return self.UPDATE.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.
@@ -291,8 +302,8 @@ def _update_sources(self, host, home):
             self._sync_dir(host, local_dir, remote_dir)
 
         scripts_dir = self.scripts_dir
-        bootstrap_code = self.BOOTSTRAP.format(project_name=self.project_name)
-        update_code = self.UPDATE.format(project_name=self.project_name)
+        bootstrap_code = self._get_bootstrap_code()
+        update_code = self._get_update_code()
         scripts = {'bootstrap.sh': bootstrap_code,
                    'update.sh': update_code}
         for script, code in scripts.items():

diff --git a/automan/conda_cluster_manager.py b/automan/conda_cluster_manager.py
@@ -4,7 +4,7 @@
 from .cluster_manager import ClusterManager
 
 
-class CondaClusterManager(ClusterManager):  # pragma: no cover
+class CondaClusterManager(ClusterManager):
 
     # The path to conda root on the remote, this is a relative path
     # and is relative to the home directory.
@@ -14,7 +14,7 @@ class CondaClusterManager(ClusterManager):  # pragma: no cover
     #!/bin/bash
 
     set -e
-    CONDA_ROOT=%s
+    CONDA_ROOT={conda_root}
     ENV_FILE="{project_name}/environments.yml"
     if [ -f $ENV_FILE ] ; then
         ~/$CONDA_ROOT/bin/conda env create -q -f $ENV_FILE -n {project_name}
@@ -29,13 +29,13 @@ class CondaClusterManager(ClusterManager):  # pragma: no cover
     if [ -f "requirements.txt" ] ; then
         pip install -r requirements.txt
     fi
-    """ % CONDA_ROOT)
+    """)
 
     UPDATE = dedent("""\
     #!/bin/bash
 
     set -e
-    CONDA_ROOT=%s
+    CONDA_ROOT={conda_root}
     ENV_FILE="{project_name}/environments.yml"
     if [ -f $ENV_FILE ] ; then
         ~/$CONDA_ROOT/bin/conda env update -q -f $ENV_FILE -n {project_name}
@@ -47,8 +47,12 @@ class CondaClusterManager(ClusterManager):  # pragma: no cover
     if [ -f "requirements.txt" ] ; then
         pip install -r requirements.txt
     fi
+    """)
 
-    """ % CONDA_ROOT)
+    def _get_bootstrap_code(self):
+        return self.BOOTSTRAP.format(
+            project_name=self.project_name, conda_root=self.CONDA_ROOT
+        )
 
     def _get_python(self, host, home):
         return os.path.join(
@@ -58,6 +62,11 @@ def _get_python(self, host, home):
             )
         )
 
+    def _get_update_code(self):
+        return self.UPDATE.format(
+            project_name=self.project_name, conda_root=self.CONDA_ROOT
+        )
+
     def _get_helper_scripts(self):
         """Return a space separated string of script files that you need copied over to
         the remote host.

diff --git a/automan/tests/test_cluster_manager.py b/automan/tests/test_cluster_manager.py
@@ -16,6 +16,7 @@
 
 from automan.jobs import Job
 from automan.cluster_manager import ClusterManager
+from automan.conda_cluster_manager import CondaClusterManager
 from .test_jobs import wait_until
 
 
@@ -36,7 +37,7 @@ class MyClusterManager(ClusterManager):
          echo "update"
          """)
 
-    def _get_virtualenv(self):
+    def _get_helper_scripts(self):
         return None
 
 
@@ -184,3 +185,36 @@ def test_remote_bootstrap_and_sync(self):
         # Then
         dest = os.path.join(self.root, 'automan', project_name, 'script.py')
         self.assertTrue(os.path.exists(dest))
+
+
+class TestCondaClusterManager(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.conda_cluster_manager.CondaClusterManager.CONDA_ROOT",
+                'TEST_ROOT')
+    def test_overloaded_methods(self):
+        # Given
+        cm = CondaClusterManager()
+
+        # When/Then
+        python = cm._get_python('foo', 'blah')
+        name = os.path.basename(self.root)
+        self.assertEqual(python, os.path.join(
+            'blah', '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(), '')