61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
"""GLiCID spawner module."""
|
|
|
|
from batchspawner import SlurmSpawner
|
|
from jinja2 import Environment, PackageLoader, select_autoescape
|
|
from traitlets import Unicode
|
|
|
|
from .micromamba import get_envs
|
|
from .resources import CPU, GPU, RAM
|
|
|
|
|
|
class GlicidSpawner(SlurmSpawner):
|
|
"""Glicid SLURM Spawner."""
|
|
|
|
batchspawner_singleuser_cmd = Unicode(
|
|
'glicid-spawner-singleuser',
|
|
help='Glicid spawner singleuser command.',
|
|
).tag(config=True)
|
|
|
|
def _options_form_default(self) -> str:
|
|
"""JupyterHub rendered form template."""
|
|
environment = Environment(
|
|
loader=PackageLoader('glicid_spawner'),
|
|
autoescape=select_autoescape(),
|
|
)
|
|
template = environment.get_template('interactive.html')
|
|
|
|
return template.render(
|
|
username=self.user.name,
|
|
python_envs=get_envs(self.user.name),
|
|
cpu_available=CPU,
|
|
ram_available=RAM,
|
|
gpu_available=GPU,
|
|
)
|
|
|
|
def options_from_form(self, formdata) -> dict:
|
|
"""Export options from form."""
|
|
# Index of user resources choices
|
|
i_cpu = int(formdata['cpu'][0])
|
|
i_ram = int(formdata['ram'][0])
|
|
i_gpu = int(formdata['gpu'][0])
|
|
|
|
duration = min(
|
|
CPU[i_cpu].max_duration,
|
|
RAM[i_ram].max_duration,
|
|
GPU[i_gpu].max_duration,
|
|
)
|
|
|
|
# Export options
|
|
options = {}
|
|
|
|
options['pyenv'] = formdata['python-env'][0]
|
|
|
|
options['cpus-per-task'] = CPU[i_cpu].description
|
|
options['mem'] = RAM[i_ram].description.replace(' ', '')
|
|
|
|
if i_gpu:
|
|
options['gres'] = 'gpu:' + GPU[i_gpu].description
|
|
|
|
options['time'] = f'{duration:02d}:00:00'
|
|
|
|
return options
|