Merge pull request #60 from mgermain/helios_support

Added support for helios gpu cluster.

Author: MarcCote

smartdispatch/config/helios.json

@@ -0,0 +1,10 @@
+{
+    "maint": {
+        "ram": 128,
+        "modules": ["cuda/6.0.37"],
+        "cores": 20,
+        "max_walltime": "12:00:00",
+        "gpus": 8,
+        "nodes": 15
+    }
+}

smartdispatch/job_generator.py

@@ -11,6 +11,8 @@ def job_generator_factory(queue, commands, command_params={}, cluster_name=None)
         return GuilliminJobGenerator(queue, commands, command_params)
     elif cluster_name == "mammouth":
         return MammouthJobGenerator(queue, commands, command_params)
+    elif cluster_name == "helios":
+        return HeliosJobGenerator(queue, commands, command_params)
 
     return JobGenerator(queue, commands, command_params)
 
@@ -80,11 +82,23 @@ def write_pbs_files(self, pbs_dir="./"):
 
         return pbs_filenames
 
+    def generate_pbs_with_account_name_from_env(self, environment_variable_name):
+        pbs_list = JobGenerator.generate_pbs(self)
+
+        if environment_variable_name not in os.environ:
+            raise ValueError("Undefined environment variable: ${}. Please, provide your account name!".format(environment_variable_name))
+
+        account_name = os.path.basename(os.path.realpath(os.getenv(environment_variable_name)))
+        for pbs in pbs_list:
+            pbs.add_options(A=account_name)
+
+        return pbs_list
+
 
 class MammouthJobGenerator(JobGenerator):
 
-    def generate_pbs(self, *args, **kwargs):
-        pbs_list = JobGenerator.generate_pbs(self, *args, **kwargs)
+    def generate_pbs(self):
+        pbs_list = JobGenerator.generate_pbs(self)
 
         if self.queue.name.endswith("@mp2"):
             for pbs in pbs_list:
@@ -95,14 +109,23 @@ def generate_pbs(self, *args, **kwargs):
 
 class GuilliminJobGenerator(JobGenerator):
 
-    def generate_pbs(self, *args, **kwargs):
-        pbs_list = JobGenerator.generate_pbs(self, *args, **kwargs)
+    def generate_pbs(self):
+        return self.generate_pbs_with_account_name_from_env('HOME_GROUP')
+
 
-        if 'HOME_GROUP' not in os.environ:
-            raise ValueError("Undefined environment variable: $HOME_GROUP. Please, provide your account name if on Guillimin!")
+# https://wiki.calculquebec.ca/w/Ex%C3%A9cuter_une_t%C3%A2che#tab=tab6
+class HeliosJobGenerator(JobGenerator):
+
+    def generate_pbs(self):
+        pbs_list = self.generate_pbs_with_account_name_from_env('RAP')
 
-        account_name = os.path.split(os.getenv('HOME_GROUP', ''))[-1]
         for pbs in pbs_list:
-            pbs.add_options(A=account_name)
+            # Remove forbidden ppn option. Default is 5 cores per 2 gpu.
+            pbs.resources['nodes'] = re.sub(":ppn=[0-9]+", "", pbs.resources['nodes'])
+
+            # Nb of GPUs has to be a multiple of 2
+            nb_gpus = int(re.findall("gpus=([0-9]+)", pbs.resources['nodes'])[0])
+            if nb_gpus % 2 != 0:
+                pbs.resources['nodes'] = re.sub("gpus=[0-9]+", "gpus={0}".format(nb_gpus+1), pbs.resources['nodes'])
 
         return pbs_list

smartdispatch/tests/test_job_generator.py

@@ -1,8 +1,8 @@
-from nose.tools import assert_true, assert_equal, assert_raises
+from nose.tools import assert_true, assert_false, assert_equal, assert_raises
 
 import os
 from smartdispatch.queue import Queue
-from smartdispatch.job_generator import JobGenerator, GuilliminJobGenerator, MammouthJobGenerator
+from smartdispatch.job_generator import JobGenerator, GuilliminJobGenerator, MammouthJobGenerator, HeliosJobGenerator
 from smartdispatch.job_generator import job_generator_factory
 import unittest
 import tempfile
@@ -112,6 +112,7 @@ def test_generate_pbs(self):
 
 
 class TestMammouthQueue(unittest.TestCase):
+
     def setUp(self):
         self.commands = ["echo 1", "echo 2", "echo 3", "echo 4"]
         self.queue = Queue("qtest@mp2", "mammouth")
@@ -122,6 +123,48 @@ def test_generate_pbs(self):
         assert_true("ppn=1" in job_generator.generate_pbs()[0].__str__())
 
 
+class TestHeliosQueue(unittest.TestCase):
+
+    def setUp(self):
+        self.commands = ["echo 1", "echo 2", "echo 3", "echo 4"]
+        self.queue = Queue("maint", "helios")
+
+        self.env_val = 'RAP'
+
+        self.bak_env_home_group = os.environ.get(self.env_val)
+        if self.bak_env_home_group is not None:
+            del os.environ[self.env_val]
+        os.environ[self.env_val] = "/rap/group/"
+
+        self.job_generator = HeliosJobGenerator(self.queue, self.commands)
+
+    def tearDown(self):
+        if self.bak_env_home_group is not None:
+            os.environ[self.env_val] = self.bak_env_home_group
+
+    def test_generate_pbs_invalid_group(self):
+        del os.environ[self.env_val]
+
+        assert_raises(ValueError, self.job_generator.generate_pbs)
+
+    def test_generate_pbs_valid_group(self):
+        pbs = self.job_generator.generate_pbs()[0]
+
+        assert_equal(pbs.options['-A'], "group")
+
+    def test_generate_pbs_ppn_is_absent(self):
+        assert_false("ppn=" in self.job_generator.generate_pbs()[0].__str__())
+
+    def test_generate_pbs_even_nb_commands(self):
+        assert_true("gpus=4" in self.job_generator.generate_pbs()[0].__str__())
+
+    def test_generate_pbs_odd_nb_commands(self):
+        commands = ["echo 1", "echo 2", "echo 3", "echo 4", "echo 5"]
+        job_generator = HeliosJobGenerator(self.queue, commands)
+
+        assert_true("gpus=6" in job_generator.generate_pbs()[0].__str__())
+
+
 def test_job_generator_factory():
     queue = {"queue_name": "qtest"}
     commands = []
@@ -131,6 +174,11 @@ def test_job_generator_factory():
     job_generator = job_generator_factory(queue, commands, cluster_name="mammouth")
     assert_true(isinstance(job_generator, MammouthJobGenerator))
 
+    job_generator = job_generator_factory(queue, commands, cluster_name="helios")
+    assert_true(isinstance(job_generator, HeliosJobGenerator))
+
     job_generator = job_generator_factory(queue, commands, cluster_name=None)
     assert_true(isinstance(job_generator, JobGenerator))
     assert_true(not isinstance(job_generator, GuilliminJobGenerator))
+    assert_true(not isinstance(job_generator, MammouthJobGenerator))
+    assert_true(not isinstance(job_generator, HeliosJobGenerator))

smartdispatch/utils.py

@@ -74,4 +74,6 @@ def detect_cluster():
         cluster_name = "mammouth"
     elif server_name.split('.')[-1] == 'guil':
         cluster_name = "guillimin"
+    elif server_name.split('.')[-1] == 'helios':
+        cluster_name = "helios"
     return cluster_name