Fix spawner job id parser for SLURM multi cluster config

This commit is contained in:
Benoît Seignovert 2024-02-21 15:11:14 +01:00
parent a038d46231
commit 8574a75e37
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 38 additions and 8 deletions

View file

@ -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)

View file

@ -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})"