Pydantic models

Pydantic model for capturing SLURM node information.

Classes:

Name	Description
AllocTRES	Allocated Trackable Resources
CfgTRES	Configured Trackable Resources
GPU	GPU name and count.
Node	Slurm node model.
State	Node state

Functions:

Name	Description
gpu_mem_from_features	Get the GPU memory from features.

AllocTRES

Bases: BaseModel

Allocated Trackable Resources

CfgTRES

Bases: BaseModel

Configured Trackable Resources

GPU

Bases: BaseModel

GPU name and count.

Node

Bases: BaseModel

Slurm node model.

Methods:

Name	Description
gres_validator	Parameters

gres_validator

gres_validator(_value: str) -> list[GPU]

Parameters

_value : str The string to parse - cpu:192 - gpu:RTX6000:3 - gpu:a100:4(S:0-1),cpu:72 - gpu:a100_3g.20gb:8(S:0-1),cpu:72 - gpu:a100:4(S:0-1),nvme:3500 - a100_4g:2,a100_3g:2

Source code in src/slurm_viewer/data/node_model.py

@field_validator('gres', 'gres_used', mode='before')
@classmethod
def gres_validator(cls, _value: str) -> list[GPU]:
    """
    Parameters
    ----------
    _value : str
        The string to parse
         - cpu:192
         - gpu:RTX6000:3
         - gpu:a100:4(S:0-1),cpu:72
         - gpu:a100_3g.20gb:8(S:0-1),cpu:72
         - gpu:a100:4(S:0-1),nvme:3500
         - a100_4g:2,a100_3g:2
    """

    if len(_value) == 0:
        return []

    values = _value.split(',') if ',' in _value else [_value]

    gpus = []
    for value in values:
        data = value.split(':')
        if len(data) < 3:  # noqa: PLR2004
            continue

        name = data[1]
        try:
            num_gpus = int(re.split(r'\D+', data[2])[0])
        except ValueError:
            num_gpus = 0

        # Dirty fix for Alice
        if '4g.40gb' in name:
            name = 'a100_4g'
        if '3g.40gb' in name:
            name = 'a100_3g'
        if name == '1(S':
            name = 'tesla_t4'
            num_gpus = 1
        # End

        gpus.append(GPU(name=name, amount=num_gpus))
    return gpus

State

Bases: Enum

Node state

gpu_mem_from_features

gpu_mem_from_features(
    _features: list[str],
) -> MemoryUsed | None

Get the GPU memory from features.

Source code in src/slurm_viewer/data/node_model.py

def gpu_mem_from_features(_features: list[str]) -> MemoryUsed | None:
    """ Get the GPU memory from features. """
    for feature in _features:
        m = re.search(r'^.*.(?P<gpu_mem>\d{2,})[Gg]\w*$', feature)
        if m is None:
            continue

        try:
            return MemoryUsed.from_mb(int(m['gpu_mem']) * 1024)
        except ValueError:
            return None

    return None

Functions:

Name	Description
condense_string_list	Condenses a list of similar strings by grouping consecutive numbers.
format_group	Formats a group of numbers into a comma-separated string with ranges.

condense_string_list

condense_string_list(string_list: list[str]) -> list[str]

Condenses a list of similar strings by grouping consecutive numbers.

Args: string_list: A list of strings with a common prefix and a numerical suffix.

Returns: A list of strings where consecutive numerical suffixes are grouped using a hyphen.

Source code in src/slurm_viewer/data/partitions_model.py

def condense_string_list(string_list: list[str]) -> list[str]:
    """
    Condenses a list of similar strings by grouping consecutive numbers.

    Args:
        string_list: A list of strings with a common prefix and a numerical suffix.

    Returns:
        A list of strings where consecutive numerical suffixes are grouped using a hyphen.
    """
    if not string_list:
        return []

    string_list.sort()  # Sort the list for proper grouping

    grouped_strings = []
    current_prefix: str | None = 'None'
    current_numbers = []

    for s in string_list:
        match = re.match(r"([a-zA-Z-]+)(\d+)", s)
        if match:
            prefix = match.group(1)
            number = int(match.group(2))

            if prefix == current_prefix:
                current_numbers.append(number)
            else:
                if current_prefix is not None:
                    grouped_strings.append(format_group(current_prefix, current_numbers))
                current_prefix = prefix
                current_numbers = [number]
        else:
            # Handle strings that don't fit the expected pattern
            if current_prefix is not None:
                grouped_strings.append(format_group(current_prefix, current_numbers))
                current_prefix = None
                current_numbers = []
            grouped_strings.append(s)

    # Add the last group if any
    if current_prefix is not None:
        grouped_strings.append(format_group(current_prefix, current_numbers))

    return grouped_strings

format_group

format_group(prefix: str, numbers: list[int]) -> str

Formats a group of numbers into a comma-separated string with ranges.

Args: prefix: The common prefix of the strings. numbers: A sorted list of integers.

Returns: A formatted string representing the group.

Source code in src/slurm_viewer/data/partitions_model.py

def format_group(prefix: str, numbers: list[int]) -> str:
    """
    Formats a group of numbers into a comma-separated string with ranges.

    Args:
        prefix: The common prefix of the strings.
        numbers: A sorted list of integers.

    Returns:
        A formatted string representing the group.
    """
    if not numbers:
        return prefix  # Should not happen in the main logic

    formatted_parts = []
    i = 0
    while i < len(numbers):
        start = numbers[i]
        end = start
        j = i + 1
        while j < len(numbers) and numbers[j] == end + 1:
            end = numbers[j]
            j += 1

        if start == end:
            formatted_parts.append(f"{start:03d}")  # Ensure 3-digit formatting
        else:
            formatted_parts.append(f"{start:03d}-{end:03d}")  # Ensure 3-digit formatting
        i = j

    return f"{prefix}[{','.join(formatted_parts)}]"

Pydantic model for capturing SLURM queue information.

Classes:

Name	Description
JobStateCodes	Job state codes.
Queue	Slurm queue model.

JobStateCodes

Bases: Enum

Job state codes.

Queue

Bases: BaseModel

Slurm queue model.

Methods:

Name	Description
gres_validator	Parameters

gres_validator

gres_validator(_value: str) -> list[GPU]

Parameters

_value : str The string to parse - cpu:192 - gpu:RTX6000:3 - gpu:a100:4(S:0-1),cpu:72 - gpu:a100_3g.20gb:8(S:0-1),cpu:72 - gpu:a100:4(S:0-1),nvme:3500 - a100_4g:2,a100_3g:2

Source code in src/slurm_viewer/data/queue_model.py

@field_validator('gres_detail', mode='before')
@classmethod
def gres_validator(cls, _value: str) -> list[GPU]:
    """
    Parameters
    ----------
    _value : str
        The string to parse
         - cpu:192
         - gpu:RTX6000:3
         - gpu:a100:4(S:0-1),cpu:72
         - gpu:a100_3g.20gb:8(S:0-1),cpu:72
         - gpu:a100:4(S:0-1),nvme:3500
         - a100_4g:2,a100_3g:2
    """

    if len(_value) == 0:
        return []

    values = _value.split(',') if ',' in _value else [_value]

    gpus = []
    for value in values:
        if isinstance(value, list):
            if len(value) != 1:
                continue
            data = value[0].split(':')
        else:
            data = value.split(':')

        if len(data) < 3:  # noqa: PLR2004
            continue

        name = data[1]
        try:
            num_gpus = int(re.split(r'\D+', data[2])[0])
        except ValueError:
            num_gpus = 0

        # Dirty fix for Alice
        if '4g.40gb' in name:
            name = 'a100_4g'
        if '3g.40gb' in name:
            name = 'a100_3g'
        if name == '1(S':
            name = 'tesla_t4'
            num_gpus = 1
        # End

        gpus.append(GPU(name=name, amount=num_gpus))
    return gpus