Customize shebang 739

.travis-data/computer-setup-input.txt

@@ -4,6 +4,7 @@ localhost
 True
 ssh
 torque
+#!/bin/bash
 /scratch/{username}/aiida_run
 mpirun -np {tot_num_mpiprocs}
 1

aiida/backends/djsite/db/models.py

@@ -1426,21 +1426,28 @@ def get_aiida_class(self):
         from aiida.orm.computer import Computer
         return Computer(dbcomputer=self)
 
-    def get_workdir(self):
+
+    def _get_val_from_metadata(self, key):
         import json
 
         try:
             metadata = json.loads(self.metadata)
         except ValueError:
             raise DbContentError(
-                "Error while reading metadata for DbComputer {} ({})".format(
-                    self.name, self.hostname))
-
+                "Error while reading metadata for DbComputer {} ({})".format(self.name, self.hostname))
         try:
-            return metadata['workdir']
+            return metadata[key]
         except KeyError:
-            raise ConfigurationError('No workdir found for DbComputer {} '.format(
-                self.name))
+            raise ConfigurationError('No {} found for DbComputer {} '.format(key,self.name))
+
+    def get_workdir(self):
+        return self._get_val_from_metadata('workdir')
+
+    def get_shebang(self):
+        """
+        Return the shebang line
+        """
+        return self._get_val_from_metadata('shebang')
 
     def __str__(self):
         if self.enabled:

aiida/backends/sqlalchemy/models/computer.py

@@ -98,6 +98,13 @@ def get_workdir(self):
             raise ConfigurationError('No workdir found for DbComputer {} '.format(
                 self.name))
 
+    def get_shebang(self):
+        try:
+            return self._metadata['shebang']
+        except KeyError:
+            raise ConfigurationError('No shebang found for DbComputer {} '.format(
+                self.name))
+
     @property
     def pk(self):
         return self.id

aiida/backends/tests/verdi_commands.py

@@ -24,6 +24,7 @@
     "True",
     "ssh",
     "torque",
+    "#!/bin/bash",
     "/scratch/{username}/aiida_run",
     "mpirun -np {tot_num_mpiprocs}",
     "1",

aiida/orm/implementation/django/computer.py

@@ -104,6 +104,7 @@ def full_text_info(self):
         ret_lines.append(
             " * Scheduler type: {}".format(self.get_scheduler_type()))
         ret_lines.append(" * Work directory: {}".format(self.get_workdir()))
+        ret_lines.append(" * Shebang:        {}".format(self.get_shebang()))
         ret_lines.append(" * mpirun command: {}".format(" ".join(
             self.get_mpirun_command())))
         def_cpus_machine = self.get_default_mpiprocs_per_machine()
@@ -220,6 +221,14 @@ def get_workdir(self):
             # This happens the first time: I provide a reasonable default value
             return "/scratch/{username}/aiida_run/"
 
+    def get_shebang(self):
+        try:
+            return self.dbcomputer.get_shebang()
+        except ConfigurationError:
+            # This happens the first time: I provide a reasonable default value
+            return "#!/bin/bash"
+
+
     def set_workdir(self, val):
         # if self.to_be_stored:
         if not isinstance(val, basestring):
@@ -228,8 +237,7 @@ def set_workdir(self, val):
         metadata = self._get_metadata()
         metadata['workdir'] = val
         self._set_metadata(metadata)
-        # else:
-        #    raise ModificationNotAllowed("Cannot set a property after having stored the entry")
+
 
     def get_name(self):
         return self.dbcomputer.name

aiida/orm/implementation/general/calculation/job/__init__.py

@@ -1453,6 +1453,7 @@ def _presubmit(self, folder, use_unstored_links=False):
 
         # I create the job template to pass to the scheduler
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = computer.get_shebang()
         ## TODO: in the future, allow to customize the following variables
         job_tmpl.submit_as_hold = False
         job_tmpl.rerunnable = False

aiida/orm/implementation/general/computer.py

@@ -102,6 +102,11 @@ def _conf_attributes(self):
                  ",".join(Scheduler.get_valid_schedulers())),
              False,
              ),
+            ("shebang",
+             "shebang line at the beginning of the submission script",
+             "this line specifies the first line of the submission script for this computer",
+             False,
+             ),
             ("workdir",
              "AiiDA work directory",
              "The absolute path of the directory on the computer where AiiDA will\n"
@@ -449,6 +454,18 @@ def _set_workdir_string(self, string):
         self._workdir_validator(string)
         self.set_workdir(string)
 
+
+    def _get_shebang_string(self):
+        return self.get_shebang()
+
+    def _set_shebang_string(self, string):
+        """
+        Set the shebang line.
+        """
+        # self._shebang_validator(string)
+        # Should we validate?
+        self.set_shebang(string)
+
     @classmethod
     def _workdir_validator(cls, workdir):
         """
@@ -682,10 +699,26 @@ def set_transport_params(self, val):
     def get_workdir(self):
         pass
 
+    @abstractmethod
+    def get_shebang(self):
+        pass
+
     @abstractmethod
     def set_workdir(self, val):
         pass
 
+    def set_shebang(self, val):
+        """
+        :param str val: A valid shebang line
+        """
+        if not isinstance(val, basestring):
+            raise ValueError("{} is invalid. Input has to be a string".format(val))
+        if not val.startswith('#!'):
+            raise ValueError("{} is invalid. A shebang line has to start with #!".format(val))
+        metadata = self._get_metadata()
+        metadata['shebang'] = val
+        self._set_metadata(metadata)
+
     @abstractmethod
     def get_name(self):
         pass
@@ -809,4 +842,4 @@ class Util(object):
 
     @abstractmethod
     def delete_computer(self, pk):
-        pass
+        pass

aiida/orm/implementation/sqlalchemy/computer.py

@@ -111,6 +111,7 @@ def full_text_info(self):
         ret_lines.append(" * Transport type: {}".format(self.get_transport_type()))
         ret_lines.append(" * Scheduler type: {}".format(self.get_scheduler_type()))
         ret_lines.append(" * Work directory: {}".format(self.get_workdir()))
+        ret_lines.append(" * Shebang:        {}".format(self.get_shebang()))
         ret_lines.append(" * mpirun command: {}".format(" ".join(
             self.get_mpirun_command())))
         def_cpus_machine = self.get_default_mpiprocs_per_machine()
@@ -232,8 +233,13 @@ def set_workdir(self, val):
         metadata = self._get_metadata()
         metadata['workdir'] = val
         self._set_metadata(metadata)
-        #else:
-        #    raise ModificationNotAllowed("Cannot set a property after having stored the entry")
+
+    def get_shebang(self):
+        try:
+            return self.dbcomputer.get_shebang()
+        except ConfigurationError:
+            # This happens the first time: I provide a reasonable default value
+            return "#!/bin/bash"
 
     def get_name(self):
         return self.dbcomputer.name

aiida/scheduler/__init__.py

@@ -122,7 +122,7 @@ def get_submit_script(self, job_tmpl):
 
         The plugin returns something like
 
-        #!/bin/bash <- this shebang line could be configurable in the future
+        #!/bin/bash <- this shebang line is configurable to some extent
         scheduler_dependent stuff to choose numnodes, numcores, walltime, ...
         prepend_computer [also from calcinfo, joined with the following?]
         prepend_code [from calcinfo]
@@ -145,13 +145,20 @@ def get_submit_script(self, job_tmpl):
 
         empty_line = ""
 
-        shebang = "#!/bin/bash"
 
         # I fill the list with the lines, and finally join them and return
         script_lines = []
-        script_lines.append(shebang)
-        script_lines.append(empty_line)
 
+        if job_tmpl.shebang:
+            script_lines.append(job_tmpl.shebang)
+        elif job_tmpl.shebang == '':
+            # Here I check whether the shebang was set explicitly as an empty line.
+            # In such a case, the first line is empty, if that's what the user wants:
+            script_lines.append(job_tmpl.shebang)
+        elif job_tmpl.shebang is None:
+            script_lines.append('#!/bin/bash')
+        else:
+            raise ValueError("Invalid shebang set: {}".format(job_tmpl.shebang))
         script_lines.append(self._get_submit_script_header(job_tmpl))
         script_lines.append(empty_line)
 

aiida/scheduler/datastructures.py

@@ -298,6 +298,7 @@ class JobTemplate(DefaultFieldsAttributeDict):
 
     Fields:
 
+      * ``shebang line``: The first line of the submission script
       * ``submit_as_hold``: if set, the job will be in a 'hold' status right
         after the submission
       * ``rerunnable``: if the job is rerunnable (boolean)
@@ -377,6 +378,7 @@ class JobTemplate(DefaultFieldsAttributeDict):
     #        place then.
 
     _default_fields = (
+        'shebang',
         'submit_as_hold',
         'rerunnable',
         'job_environment',

aiida/scheduler/plugins/test_lsf.py

@@ -142,6 +142,7 @@ def test_submit_script(self):
         s = LsfScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.uuid = str(uuid.uuid4())
         job_tmpl.job_resource = s.create_job_resource(tot_num_mpiprocs=2,
                                                       parallel_env='b681e480bd.cern.ch')
@@ -154,6 +155,7 @@ def test_submit_script(self):
 
         submit_script_text = s.get_submit_script(job_tmpl)
 
+
         self.assertTrue( submit_script_text.startswith('#!/bin/bash') )
 
         self.assertTrue( '#BSUB -rn' in submit_script_text )

aiida/scheduler/plugins/test_pbspro.py

@@ -908,6 +908,7 @@ def test_submit_script(self):
         s = PbsproScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash -l'
         job_tmpl.job_resource = s.create_job_resource(num_machines=1, num_mpiprocs_per_machine=1)
         job_tmpl.uuid = str(uuid.uuid4())
         job_tmpl.max_wallclock_seconds = 24 * 3600
@@ -920,11 +921,39 @@ def test_submit_script(self):
         submit_script_text = s.get_submit_script(job_tmpl)
 
         self.assertTrue('#PBS -r n' in submit_script_text)
-        self.assertTrue(submit_script_text.startswith('#!/bin/bash'))
+        self.assertTrue(submit_script_text.startswith('#!/bin/bash -l'))
         self.assertTrue('#PBS -l walltime=24:00:00' in submit_script_text)
         self.assertTrue('#PBS -l select=1' in submit_script_text)
         self.assertTrue("'mpirun' '-np' '23' 'pw.x' '-npool' '1'" + \
                         " < 'aiida.in'" in submit_script_text)
+    def test_submit_script_bad_shebang(self):
+        """
+        Test to verify if scripts works fine with default options
+        """
+        from aiida.scheduler.datastructures import JobTemplate
+        from aiida.common.datastructures import CodeInfo, code_run_modes
+
+        s = PbsproScheduler()
+        code_info = CodeInfo()
+        code_info.cmdline_params = ["mpirun", "-np", "23", "pw.x", "-npool", "1"]
+        code_info.stdin_name = 'aiida.in'
+
+        for (shebang, expected_first_line) in ((None, '#!/bin/bash'), ("",""), ("NOSET", '#!/bin/bash')):
+            job_tmpl = JobTemplate()
+            if shebang == "NOSET":
+                pass
+            else:
+                job_tmpl.shebang = shebang
+            job_tmpl.job_resource = s.create_job_resource(num_machines=1, num_mpiprocs_per_machine=1)
+            job_tmpl.codes_info = [code_info]
+            job_tmpl.codes_run_mode = code_run_modes.SERIAL
+
+            submit_script_text = s.get_submit_script(job_tmpl)
+
+            # This tests if the implementation correctly chooses the default:
+            self.assertEquals(submit_script_text.split('\n')[0], expected_first_line)
+
+
 
     def test_submit_script_with_num_cores_per_machine(self):
         """
@@ -937,6 +966,7 @@ def test_submit_script_with_num_cores_per_machine(self):
         s = PbsproScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=2,
@@ -973,6 +1003,7 @@ def test_submit_script_with_num_cores_per_mpiproc(self):
         s = PbsproScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=1,
@@ -1013,6 +1044,7 @@ def test_submit_script_with_num_cores_per_machine_and_mpiproc1(self):
         s = PbsproScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=1,

aiida/scheduler/plugins/test_slurm.py

@@ -177,6 +177,7 @@ def test_submit_script(self):
         s = SlurmScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.uuid = str(uuid.uuid4())
         job_tmpl.job_resource = s.create_job_resource(num_machines=1, num_mpiprocs_per_machine=1)
         job_tmpl.max_wallclock_seconds = 24 * 3600 
@@ -197,6 +198,31 @@ def test_submit_script(self):
         self.assertTrue( "'mpirun' '-np' '23' 'pw.x' '-npool' '1'" + \
                          " < 'aiida.in'" in submit_script_text )
 
+    def test_submit_script_bad_shebang(self):
+        from aiida.scheduler.datastructures import JobTemplate
+        from aiida.common.datastructures import CodeInfo, code_run_modes
+
+        s = SlurmScheduler()
+        code_info = CodeInfo()
+        code_info.cmdline_params = ["mpirun", "-np", "23", "pw.x", "-npool", "1"]
+        code_info.stdin_name = 'aiida.in'
+
+        for (shebang, expected_first_line) in ((None, '#!/bin/bash'), ("",""), ("NOSET", '#!/bin/bash')):
+            job_tmpl = JobTemplate()
+            if shebang == "NOSET":
+                pass
+            else:
+                job_tmpl.shebang = shebang
+            job_tmpl.job_resource = s.create_job_resource(num_machines=1, num_mpiprocs_per_machine=1)
+            job_tmpl.codes_info = [code_info]
+            job_tmpl.codes_run_mode = code_run_modes.SERIAL
+
+            submit_script_text = s.get_submit_script(job_tmpl)
+
+            # This tests if the implementation correctly chooses the default:
+            self.assertEquals(submit_script_text.split('\n')[0], expected_first_line)
+
+
     def test_submit_script_with_num_cores_per_machine(self):
         """
         Test to verify if script works fine if we specify only
@@ -208,6 +234,7 @@ def test_submit_script_with_num_cores_per_machine(self):
         s = SlurmScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=2,
@@ -244,6 +271,7 @@ def test_submit_script_with_num_cores_per_mpiproc(self):
         s = SlurmScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=1,
@@ -284,6 +312,7 @@ def test_submit_script_with_num_cores_per_machine_and_mpiproc1(self):
         s = SlurmScheduler()
 
         job_tmpl = JobTemplate()
+        job_tmpl.shebang = '#!/bin/bash'
         job_tmpl.job_resource = s.create_job_resource(
             num_machines=1,
             num_mpiprocs_per_machine=1,