diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index fd8d648..7073489 100644 --- a/src/glicid_spawner/spawner.py +++ b/src/glicid_spawner/spawner.py @@ -1,5 +1,7 @@ """GLiCID spawner module.""" +import re + from batchspawner import SlurmSpawner from traitlets import Unicode, default @@ -9,16 +11,33 @@ from .form import options_form, options_from_form class GlicidSpawner(SlurmSpawner): """Glicid SLURM Spawner.""" - glicid_singleuser_cmd = Unicode( + batchspawner_singleuser_cmd = Unicode( 'glicid-spawner-singleuser', - help='Glicid spawner singleuser command.', + help='Spawner singleuser command.', ).tag(config=True) - @default('batchspawner_singleuser_cmd') - def _batchspawner_singleuser_cmd_default(self): - return self.glicid_singleuser_cmd + req_qos = Unicode( + 'short', + help='QoS name to submit job to resource manager', + ).tag(config=True) - def options_form(self) -> str: + slurm_job_id_re = Unicode(r'(\d+)(?:;(\w+))?').tag(config=True) + + def parse_job_id(self, output): + """Parse job id with cluster name support. + + If cluster name is present, `job_id` will be a string + and suffix with `-M job_cluster` name. + + """ + for job_id, job_cluster in re.findall(self.slurm_job_id_re, output): + return f'{job_id} -M {job_cluster}' if job_cluster else int(job_id) + + self.log.error(f'GlicidSpawner unable to parse job ID from text: {output}') + return None + + @default('options_form') + def _options_form_default(self) -> str: """JupyterHub rendered form template.""" return options_form(self.user.name) diff --git a/tests/test_spawner.py b/tests/test_spawner.py index c71574e..59aa2b7 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -16,9 +16,20 @@ def test_spawner_config(): assert isinstance(spawner, GlicidSpawner) assert isinstance(spawner, SlurmSpawner) - assert spawner.glicid_singleuser_cmd == 'glicid-spawner-singleuser' assert spawner.batchspawner_singleuser_cmd == 'glicid-spawner-singleuser' + assert spawner.req_qos == 'short' + + +def test_spawner_parse_job_id(): + """Test spawner job id parser.""" + spawner = GlicidSpawner() + + assert spawner.parse_job_id('123') == 123 + assert spawner.parse_job_id('456;nautilus') == '456 -M nautilus' + + assert spawner.parse_job_id('') is None + def test_spawner_options_form(monkeypatch): """Test spawner options form.""" @@ -34,5 +45,5 @@ def test_spawner_options_form(monkeypatch): spawner = GlicidSpawner(user=User('john-doe')) assert spawner.user.name == 'john-doe' - assert spawner.options_form() == "options_form('john-doe')" + assert spawner.options_form == "options_form('john-doe')" assert spawner.options_from_form({'foo': 123}) == "options_from_form({'foo': 123})"