Repository: MatterMiners/tardis

Issue:
[Slurm adapter] Need to specify custom (GRES) flags for sbatch
Some use cases in Slurm can only be covered with Generic Resource ([GRES](https://slurm.schedmd.com/gres.html)) scheduling. 
One such example is the reservation of scratch disk space:
https://bugs.schedmd.com/show_bug.cgi?id=2549#c9

Since the name of the GRES is completely arbitrary and up to the HPC-site, it is probably best to have a way to pass arbitrary parameters to `sbatch`. The best place (configuration-wise) is likely the `MachineTypeConfiguration`, so it is similar to the `jdl` funcitonality in HTCondor. 

I ma not yet sure if it is better to allow fully arbitrary parameters or have a dedicated `gres` configuration stanza which directly takes the `gres` to reserve. 

Pinging also @wiene 

Comments:

Patch:
diff --git a/docs/source/adapters/site.rst b/docs/source/adapters/site.rst
index f13b12b..ef9a436 100644
--- a/docs/source/adapters/site.rst
+++ b/docs/source/adapters/site.rst
@@ -16,9 +16,9 @@ Site Adapter
     :py:class:`~tardis.resources.dronestates.AvailableState` before draining it. If no value is given, infinite lifetime
     is assumed. Multiple sites are supported by using SequenceNodes.
 
-.. note::
-    Even a minimum lifetime is set, it is not guaranteed that the :py:class:`~tardis.resources.drone.Drone` is not
-    drained due to a dropping demand for it before its minimum lifetime is exceeded.
+    .. note::
+        Even if a minimum lifetime is set, it is not guaranteed that the :py:class:`~tardis.resources.drone.Drone` is not
+        drained due to its dropping demand before its minimum lifetime is exceeded.
 
 
 Generic Site Adapter Configuration
@@ -414,6 +414,23 @@ Available adapter configuration options
     |                | Default: ShellExecutor is used!                                                             |                 |
     +----------------+---------------------------------------------------------------------------------------------+-----------------+
 
+Available machine type configuration options
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. content-tabs:: left-col
+
+    +----------------+--------------------------------------------------------------------------------------------------+-----------------+
+    | Option         | Short Description                                                                                | Requirement     |
+    +================+==================================================================================================+=================+
+    | Walltime       | Expected walltime of drone                                                                       |  **Required**   |
+    +----------------+--------------------------------------------------------------------------------------------------+-----------------+
+    | Partition      | Name of the Slurm partition to run in                                                            |  **Required**   |
+    +----------------+--------------------------------------------------------------------------------------------------+-----------------+
+    | StartupCommand | The command to execute at job start                                                              |  **Required**   |
+    +----------------+--------------------------------------------------------------------------------------------------+-----------------+
+    | SubmitOptions  | Options to add to the `sbatch` command. `long` and `short` arguments are supported (see example) |  **Optional**   |
+    +----------------+--------------------------------------------------------------------------------------------------+-----------------+
+
 .. content-tabs:: right-col
 
     .. rubric:: Example configuration
@@ -440,6 +457,18 @@ Available adapter configuration options
               Walltime: '1440'
               Partition: normal
               StartupCommand: 'pilot_clean.sh'
+              SubmitOptions:
+                short:
+                  C: "intel"
+                long:
+                  gres: "gpu:2,mic:1"
+            six_hours:
+              Walltime: '360'
+              Partition: normal
+              StartupCommand: 'pilot_clean.sh'
+              SubmitOptions:
+                long:
+                  gres: "gpu:2,mic:1"
             twelve_hours:
               Walltime: '720'
               Partition: normal
@@ -453,6 +482,11 @@ Available adapter configuration options
               Cores: 20
               Memory: 62
               Disk: 480
+            six_hours:
+              Cores: 20
+              Memory: 62
+              Disk: 480
+
 
 .. content-tabs:: left-col
 
diff --git a/docs/source/api/tardis.plugins.prometheusmonitoring.rst b/docs/source/api/tardis.plugins.prometheusmonitoring.rst
new file mode 100644
index 0000000..0029134
--- /dev/null
+++ b/docs/source/api/tardis.plugins.prometheusmonitoring.rst
@@ -0,0 +1,7 @@
+tardis.plugins.prometheusmonitoring module
+==========================================
+
+.. automodule:: tardis.plugins.prometheusmonitoring
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/api/tardis.plugins.rst b/docs/source/api/tardis.plugins.rst
index 56d6aae..9c8ae6c 100644
--- a/docs/source/api/tardis.plugins.rst
+++ b/docs/source/api/tardis.plugins.rst
@@ -11,5 +11,6 @@ Submodules
 
 .. toctree::
 
+   tardis.plugins.prometheusmonitoring
    tardis.plugins.sqliteregistry
    tardis.plugins.telegrafmonitoring
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index d7667ed..b9a16fd 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,4 +1,4 @@
-.. Created by changelog.py at 2020-05-25, command
+.. Created by changelog.py at 2020-06-03, command
    '/Users/giffler/.cache/pre-commit/repont7o94ca/py_env-default/bin/changelog docs/source/changes compile --output=docs/source/changelog.rst'
    based on the format of 'https://keepachangelog.com/'
 
@@ -32,13 +32,14 @@ Fixed
 * Fix state transitions for jobs retried by HTCondor
 * Fix state transitions and refactoring of the SLURM site adapter
 
-[Unreleased] - 2020-05-25
+[Unreleased] - 2020-06-03
 =========================
 
 Added
 -----
 
 * Added an example HTCondor jdl for the HTCondor site adapter
+* Enable support for `sbatch` command line options in the Slurm site adapter
 * Add ssh connection sharing to `SSHExecutor` in order to re-use existing connection
 
 Changed
diff --git a/docs/source/changes/150.enable_sbatch_cmdline_options.yaml b/docs/source/changes/150.enable_sbatch_cmdline_options.yaml
new file mode 100644
index 0000000..0d5f281
--- /dev/null
+++ b/docs/source/changes/150.enable_sbatch_cmdline_options.yaml
@@ -0,0 +1,9 @@
+category: added
+summary: "Enable support for `sbatch` command line options in the Slurm site adapter"
+pull requests:
+  - 150
+issues:
+  - 147
+description: |
+  `sbatch` command line option can now be added to the `MachineTypeConfiguration` of the
+  Slurm site adapter. `short` and `long` option are supported via yaml MappingNodes.
diff --git a/tardis/adapters/sites/slurm.py b/tardis/adapters/sites/slurm.py
index a3a8fb5..9431ca6 100644
--- a/tardis/adapters/sites/slurm.py
+++ b/tardis/adapters/sites/slurm.py
@@ -11,6 +11,7 @@ from ...utilities.attributedict import convert_to_attribute_dict
 from ...utilities.executors.shellexecutor import ShellExecutor
 from ...utilities.asynccachemap import AsyncCacheMap
 from ...utilities.utils import htcondor_csv_parser
+from ...utilities.utils import slurm_cmd_option_formatter
 
 from asyncio import TimeoutError
 from contextlib import contextmanager
@@ -63,6 +64,10 @@ class SlurmAdapter(SiteAdapter):
             )
             self._startup_command = self._configuration.StartupCommand
 
+        self._sbatch_cmdline_option_string = slurm_cmd_option_formatter(
+            self.sbatch_cmdline_options
+        )
+
         self._executor = getattr(self._configuration, "executor", ShellExecutor())
 
         self._slurm_status = AsyncCacheMap(
@@ -107,15 +112,13 @@ class SlurmAdapter(SiteAdapter):
     async def deploy_resource(
         self, resource_attributes: AttributeDict
     ) -> AttributeDict:
+
         request_command = (
-            f"sbatch -p {self.machine_type_configuration.Partition} "
-            f"-N 1 -n {self.machine_meta_data.Cores} "
-            f"--mem={self.machine_meta_data.Memory}gb "
-            f"-t {self.machine_type_configuration.Walltime} "
-            f"--export=SLURM_Walltime="
-            f"{self.machine_type_configuration.Walltime} "
+            "sbatch "
+            f"{self._sbatch_cmdline_option_string} "
             f"{self._startup_command}"
         )
+
         result = await self._executor.run_command(request_command)
         logger.debug(f"{self.site_name} sbatch returned {result}")
         pattern = re.compile(r"^Submitted batch job (\d*)", flags=re.MULTILINE)
@@ -165,6 +168,27 @@ class SlurmAdapter(SiteAdapter):
             {"JobId": resource_attributes.remote_resource_uuid}, **resource_attributes
         )
 
+    @property
+    def sbatch_cmdline_options(self):
+        sbatch_options = self.machine_type_configuration.get(
+            "SubmitOptions", AttributeDict()
+        )
+
+        return AttributeDict(
+            short=AttributeDict(
+                **sbatch_options.get("short", AttributeDict()),
+                p=self.machine_type_configuration.Partition,
+                N=1,
+                n=self.machine_meta_data.Cores,
+                t=self.machine_type_configuration.Walltime,
+            ),
+            long=AttributeDict(
+                **sbatch_options.get("long", AttributeDict()),
+                mem=f"{self.machine_meta_data.Memory}gb",
+                export=f"SLURM_Walltime={self.machine_type_configuration.Walltime}",
+            ),
+        )
+
     async def stop_resource(self, resource_attributes: AttributeDict):
         logger.debug("Slurm jobs cannot be stopped gracefully. Terminating instead.")
         return await self.terminate_resource(resource_attributes)
diff --git a/tardis/utilities/utils.py b/tardis/utilities/utils.py
index 0525ff2..f17512e 100644
--- a/tardis/utilities/utils.py
+++ b/tardis/utilities/utils.py
@@ -1,12 +1,17 @@
+from .attributedict import AttributeDict
 from .executors.shellexecutor import ShellExecutor
 from ..exceptions.executorexceptions import CommandExecutionFailure
+from ..interfaces.executor import Executor
 
 from io import StringIO
+from typing import List, Tuple
 
 import csv
 
 
-async def async_run_command(cmd, shell_executor=ShellExecutor()):
+async def async_run_command(
+    cmd: str, shell_executor: Executor = ShellExecutor()
+) -> str:
     try:
         response = await shell_executor.run_command(cmd)
     except CommandExecutionFailure as ef:
@@ -22,16 +27,25 @@ async def async_run_command(cmd, shell_executor=ShellExecutor()):
         return response.stdout
 
 
-def htcondor_cmd_option_formatter(options):
+def cmd_option_formatter(options: AttributeDict, prefix: str, separator: str) -> str:
     options = (
-        f"-{name} {value}" if value is not None else f"-{name}"
+        f"{prefix}{name}{separator}{value}" if value is not None else f"{prefix}{name}"
         for name, value in options.items()
     )
 
     return " ".join(options)
 
 
-def htcondor_csv_parser(htcondor_input, fieldnames, delimiter="\t", replacements=None):
+def htcondor_cmd_option_formatter(options: AttributeDict) -> str:
+    return cmd_option_formatter(options, prefix="-", separator=" ")
+
+
+def htcondor_csv_parser(
+    htcondor_input: str,
+    fieldnames: [List, Tuple],
+    delimiter: str = "\t",
+    replacements: dict = None,
+):
     replacements = replacements or {}
     with StringIO(htcondor_input) as csv_input:
         cvs_reader = csv.DictReader(
@@ -42,3 +56,26 @@ def htcondor_csv_parser(htcondor_input, fieldnames, delimiter="\t", replacements
                 key: value if value not in replacements.keys() else replacements[value]
                 for key, value in row.items()
             }
+
+
+def slurm_cmd_option_formatter(options: AttributeDict) -> str:
+    option_prefix = dict(short="-", long="--")
+    option_separator = dict(short=" ", long="=")
+
+    option_string = ""
+
+    for option_type in ("short", "long"):
+        try:
+            tmp_option_string = cmd_option_formatter(
+                getattr(options, option_type),
+                prefix=option_prefix[option_type],
+                separator=option_separator[option_type],
+            )
+        except AttributeError:
+            pass
+        else:
+            if option_string:  # add additional space between short and long options
+                option_string += " "
+            option_string += tmp_option_string
+
+    return option_string


Test patch:
diff --git a/tests/adapters_t/sites_t/test_slurm.py b/tests/adapters_t/sites_t/test_slurm.py
index a8f0355..64c4dbf 100644
--- a/tests/adapters_t/sites_t/test_slurm.py
+++ b/tests/adapters_t/sites_t/test_slurm.py
@@ -176,7 +176,26 @@ class TestSlurmAdapter(TestCase):
         )
 
         self.mock_executor.return_value.run_command.assert_called_with(
-            "sbatch -p normal -N 1 -n 20 --mem=62gb -t 60 --export=SLURM_Walltime=60 pilot.sh"  # noqa: B950
+            "sbatch -p normal -N 1 -n 20 -t 60 --mem=62gb --export=SLURM_Walltime=60 pilot.sh"  # noqa: B950
+        )
+
+    @mock_executor_run_command(TEST_DEPLOY_RESOURCE_RESPONSE)
+    def test_deploy_resource_w_submit_options(self):
+        self.test_site_config.MachineTypeConfiguration.test2large.SubmitOptions = AttributeDict(  # noqa: B950
+            long=AttributeDict(gres="tmp:1G")
+        )
+
+        slurm_adapter = SlurmAdapter(machine_type="test2large", site_name="TestSite")
+
+        run_async(
+            slurm_adapter.deploy_resource,
+            resource_attributes=AttributeDict(
+                machine_type="test2large", site_name="TestSite"
+            ),
+        )
+
+        self.mock_executor.return_value.run_command.assert_called_with(
+            "sbatch -p normal -N 1 -n 20 -t 60 --gres=tmp:1G --mem=62gb --export=SLURM_Walltime=60 pilot.sh"  # noqa: B950
         )
 
     def test_machine_meta_data(self):
diff --git a/tests/utilities_t/test_utils.py b/tests/utilities_t/test_utils.py
index dcab165..76638d6 100644
--- a/tests/utilities_t/test_utils.py
+++ b/tests/utilities_t/test_utils.py
@@ -1,6 +1,9 @@
+from tardis.utilities.attributedict import AttributeDict
 from tardis.utilities.utils import async_run_command
 from tardis.utilities.utils import htcondor_cmd_option_formatter
 from tardis.utilities.utils import htcondor_csv_parser
+from tardis.utilities.utils import slurm_cmd_option_formatter
+
 from tardis.exceptions.executorexceptions import CommandExecutionFailure
 
 from ..utilities.utilities import run_async
@@ -21,10 +24,15 @@ class TestAsyncRunCommand(TestCase):
 
 class TestHTCondorCMDOptionFormatter(TestCase):
     def test_htcondor_cmd_option_formatter(self):
-        options = {"pool": "my-htcondor.local", "test": None}
-        options_string = htcondor_cmd_option_formatter(options)
+        options = AttributeDict(pool="my-htcondor.local", test=None)
+        option_string = htcondor_cmd_option_formatter(options)
+
+        self.assertEqual(option_string, "-pool my-htcondor.local -test")
 
-        self.assertEqual(options_string, "-pool my-htcondor.local -test")
+        options = AttributeDict()
+        option_string = htcondor_cmd_option_formatter(options)
+
+        self.assertEqual(option_string, "")
 
 
 class TestHTCondorCSVParser(TestCase):
@@ -63,3 +71,31 @@ class TestHTCondorCSVParser(TestCase):
                 Test2=None,
             ),
         )
+
+
+class TestSlurmCMDOptionFormatter(TestCase):
+    def test_slurm_cmd_option_formatter(self):
+        options = AttributeDict()
+        option_string = slurm_cmd_option_formatter(options)
+
+        self.assertEqual(option_string, "")
+
+        options = AttributeDict(short=AttributeDict(foo="bar", test=None))
+        option_string = slurm_cmd_option_formatter(options)
+
+        self.assertEqual(option_string, "-foo bar -test")
+
+        options = AttributeDict(long=AttributeDict(foo="bar", test=None))
+        option_string = slurm_cmd_option_formatter(options)
+
+        self.assertEqual(option_string, "--foo=bar --test")
+
+        options = AttributeDict(
+            short=AttributeDict(foo="bar", test=None),
+            long=AttributeDict(foo_long="bar_long", test_long=None),
+        )
+        option_string = slurm_cmd_option_formatter(options)
+
+        self.assertEqual(
+            option_string, "-foo bar -test --foo_long=bar_long --test_long"
+        )