spawner/tests/test_form.py

"""Test template form module."""

import re
from pathlib import Path

from glicid_spawner import form
from glicid_spawner.form import options_attrs, options_form, options_from_form
from glicid_spawner.micromamba import MicromambaEnv
from glicid_spawner.slurm import sinfo_from_file
from pytest import fixture

DATA = Path(__file__).parent / 'data'
SINFO = sinfo_from_file(DATA / 'sinfo.txt')

SLURM_SINGLE_CLUSTER = {'N/A': SINFO.pop('N/A')}
SLURM_MULTI_CLUSTER = SINFO


@fixture
def mock_python_envs(monkeypatch):
    """Mock python environments list."""
    monkeypatch.setattr(
        form,
        'get_envs',
        lambda username: [
            MicromambaEnv('USER', 'foo', f'/{username}/envs/foo'),
            MicromambaEnv('USER', 'bar', f'/{username}/envs/bar'),
            MicromambaEnv('GLOBAL', 'baz', '/global/envs/baz'),
        ],
    )


@fixture
def mock_cluster(monkeypatch, mock_python_envs):
    """Mock multi cluster configuration (default)."""
    monkeypatch.setattr(form, 'sinfo', lambda _: SLURM_MULTI_CLUSTER)


@fixture
def mock_single_cluster(monkeypatch, mock_python_envs):
    """Mock multi cluster configuration."""
    monkeypatch.setattr(form, 'sinfo', lambda _: SLURM_SINGLE_CLUSTER)


def test_options_attrs(mock_cluster):
    """Test form options attributes."""
    options = options_attrs('john-doe')

    assert options['username'] == 'john-doe'

    assert options['folders'] == [
        '/home/john-doe',
        '/scratch/nautilus/users/john-doe',
        '/scratch/waves/users/john-doe',
        '/scratch/nautilus/projects',
        '/scratch/waves/projects',
        '/LAB-DATA/',
    ]

    assert [env.path for env in options['envs']] == [
        '/john-doe/envs/foo',
        '/john-doe/envs/bar',
        '/global/envs/baz',
    ]

    cpu = options['cpus']

    assert cpu[1] == 24
    assert cpu[24] == 1

    mem = options['mems']

    assert mem[4] == 24
    assert mem[96] == 1

    gpu = options['gpus']

    assert gpu['None'] == 24
    assert gpu['A100'] == 1

    # Multi cluster configuration (default)
    sinfo = options['sinfo']

    assert 'N/A' not in sinfo
    assert 'nautilus' in sinfo
    assert 'waves' in sinfo

    node = sinfo['nautilus']['gpu']['gnode1']

    assert node == 'gnode1'
    assert node.cpu.idle == 92
    assert node.gpu.name == 'A100'


def test_options_attrs_single_cluster(mock_single_cluster):
    """Test form options attributes in single cluster configuration."""
    options = options_attrs('john-doe')

    # Single cluster configuration
    sinfo = options['sinfo']

    assert 'N/A' in sinfo
    assert 'nautilus' not in sinfo
    assert 'waves' not in sinfo

    node = sinfo['N/A']['Devel']['nazare001']

    assert node == 'nazare001'
    assert node.cpu.idle == 20
    assert node.gpu.name == 'None'


def test_options_form_resources(mock_cluster):
    """Test options form render."""
    html = re.sub(r'\n\s+', ' ', options_form('john-doe'))  # trim line breaks

    # Username
    assert '<div class="form-control-static">john-doe</div>' in html

    # Python environments
    assert '<option value="/john-doe/envs/foo">foo (USER)</option>' in html
    assert '<option value="/john-doe/envs/bar">bar (USER)</option>' in html
    assert '<option value="/global/envs/baz">baz (GLOBAL)</option>' in html

    # CPU
    assert (
        '<input type="radio" name="cpu" id="cpu_1" value="1" data-max-duration="24" checked>'
        in html
    )
    assert '<label for="cpu_1" class="btn btn-default btn-block"> 1 </label>' in html
    assert '<input type="radio" name="cpu" id="cpu_24" value="24" data-max-duration="1">' in html
    assert '<label for="cpu_24" class="btn btn-default btn-block"> 24 </label>' in html

    # Memory
    assert (
        '<input type="radio" name="mem" id="mem_4" value="4" data-max-duration="24" checked>'
        in html
    )
    assert '<label for="mem_4" class="btn btn-default btn-block"> 4 GB </label>' in html
    assert '<input type="radio" name="mem" id="mem_96" value="96" data-max-duration="1">' in html
    assert '<label for="mem_96" class="btn btn-default btn-block"> 96 GB </label>' in html

    # GPU
    assert (
        '<input type="radio" name="gpu" id="gpu_None" value="None" data-max-duration="24" checked>'
        in html
    )
    assert '<label for="gpu_None" class="btn btn-default btn-block"> None </label>' in html
    assert (
        '<input type="radio" name="gpu" id="gpu_A100" value="A100" data-max-duration="1">' in html
    )
    assert '<label for="gpu_A100" class="btn btn-default btn-block"> A100 </label>' in html


def test_options_form_slurm(mock_cluster):
    """Test options form render."""
    html = re.sub(r'\n\s+', ' ', options_form('john-doe'))  # trim line breaks

    # Cluster (multi-cluster by default)
    assert '<label for="cluster" class="col-sm-3 control-label">Cluster:</label>' in html

    assert '<div class="flex-item-2 slurm-cluster" data-cluster="nautilus">' in html
    assert '<div class="flex-item-2 slurm-cluster" data-cluster="waves">' in html

    # The 1st cluster is always selected when present…
    assert (
        '<input type="radio" name="cluster" id="cluster_nautilus" value="nautilus" checked>' in html
    )
    assert (
        '<label for="cluster_nautilus" class="btn btn-default btn-block"> Nautilus </label>' in html
    )

    # … not the second one
    assert '<input type="radio" name="cluster" id="cluster_waves" value="waves" >' in html

    # Partitions
    assert '<label for="partition" class="col-sm-3 control-label">Partition:</label>' in html

    assert (
        '<div class="flex-item-4 slurm-partition" '
        'data-cluster="nautilus" data-partition="gpu" '
        'data-cpu="96" data-mem="768" data-gpu="A100">' in html
    )

    assert '<input type="radio" name="partition" id="partition_nautilus_gpu" value="gpu">' in html
    assert (
        '<label for="partition_nautilus_gpu" class="btn btn-default btn-block"> Gpu </label>'
        in html
    )

    # Nodes (hidden by default)
    assert '<div class="form-group nodes hidden">' in html
    assert '<label for="node" class="col-sm-3 control-label">Node:</label>' in html

    assert (
        '<div class="flex-item-4 slurm-node" '
        'data-cluster="nautilus" data-partition="gpu" data-node="gnode1" '
        'data-cpu="92" data-mem="768" data-gpu="A100">' in html
    )

    assert '<input type="radio" name="node" id="node_nautilus_gpu_gnode1" value="gnode1">' in html
    assert (
        '<label for="node_nautilus_gpu_gnode1" class="btn btn-default btn-block"> Gnode1 </label>'
        in html
    )


def test_options_form_slurm_single_cluster(mock_single_cluster):
    """Test options form render."""
    html = re.sub(r'\n\s+', ' ', options_form('john-doe'))  # trim line breaks

    # No cluster
    assert '<label for="cluster" class="col-sm-3 control-label">Cluster:</label>' not in html

    # Partitions (not hidden by default for single-cluster)
    assert '<div class="form-group partitions">' in html
    assert '<label for="partition" class="col-sm-3 control-label">Partition:</label>' in html

    assert (
        '<div class="flex-item-4 slurm-partition" '
        'data-cluster="N/A" data-partition="GPU-short" '
        'data-cpu="20" data-mem="184" data-gpu="T4">' in html
    )

    assert (
        '<input type="radio" name="partition" id="partition_N/A_GPU-short" value="GPU-short">'
        in html
    )
    assert (
        '<label for="partition_N/A_GPU-short" class="btn btn-default btn-block"> Gpu-short </label>'
        in html
    )

    # Nodes (hidden by default)
    assert (
        '<div class="flex-item-4 slurm-node" '
        'data-cluster="N/A" data-partition="GPU-short" data-node="budbud001" '
        'data-cpu="20" data-mem="184" data-gpu="T4">' in html
    )

    assert (
        '<input type="radio" name="node" id="node_N/A_GPU-short_budbud001" value="budbud001">'
        in html
    )
    assert (
        '<label for="node_N/A_GPU-short_budbud001" class="btn btn-default btn-block"> Budbud001 </label>'
        in html
    )


def test_options_from_form():
    """Test options from form parser."""
    # No GPU
    formdata = {
        'workdir': ['/home/john-doe/'],
        'python-env': ['/john-doe/envs/foo'],
        'cpu': ['1'],
        'mem': ['4'],
        'gpu': ['None'],
    }

    assert options_from_form(formdata) == {
        'workdir': '/home/john-doe/',
        'pyenv': '/john-doe/envs/foo',
        'nprocs': 1,
        'memory': '4GB',
        'runtime': '24:00:00',
    }

    # With GPU (in defaults list)
    formdata = {
        'workdir': ['/scratch/nautilus/users/john-doe/'],
        'python-env': ['/john-doe/envs/bar'],
        'cpu': ['2'],
        'mem': ['16'],
        'gpu': ['A40'],  # -> 2 h
        'cluster': ['nautilus'],
    }

    assert options_from_form(formdata) == {
        'workdir': '/scratch/nautilus/users/john-doe/',
        'pyenv': '/john-doe/envs/bar',
        'nprocs': 2,
        'memory': '16GB',
        'runtime': '02:00:00',
        'gres': 'gpu:A40',
        'cluster': 'nautilus',
    }

    # With unknown GPU (default 1h allocation)
    formdata = {
        'workdir': ['/scratch/waves/projects/'],
        'python-env': ['/global/envs/baz'],
        'cpu': ['8'],
        'mem': ['16'],
        'gpu': ['T4'],
        'partition': ['GPU-short'],
    }

    assert options_from_form(formdata) == {
        'workdir': '/scratch/waves/projects/',
        'pyenv': '/global/envs/baz',
        'nprocs': 8,
        'memory': '16GB',
        'runtime': '01:00:00',
        'gres': 'gpu:T4',
        'partition': 'GPU-short',
    }

    # Invalid CPU request (0h allocated)
    formdata = {
        'workdir': ['/LAD-DATA/'],
        'python-env': ['/global/envs/qux'],
        'cpu': ['128'],
        'mem': ['4096'],
        'gpu': ['A100'],
        'cluster': ['waves'],  # Gres should be lowercased for this cluster
        'node': ['cribbar001'],
    }

    assert options_from_form(formdata) == {
        'workdir': '/LAD-DATA/',
        'pyenv': '/global/envs/qux',
        'nprocs': 128,
        'memory': '4096GB',
        'runtime': '00:00:00',
        'gres': 'gpu:a100',
        'cluster': 'waves',
        'node': 'cribbar001',
    }