utils.py

from typing import Literal
import os
import re
import subprocess
from pathlib import Path

# ENH: if this package becomes merged with noctua, need to replace this function since it's copied directly from there
def convert_si(value: str | float | int, 
               unit: Literal['base','K','M','G','T'] | None = None, 
               to_unit: Literal['base','K','M','G','T'] = 'base', 
               use_binary: bool = False
) -> float:
    """_summary_

    Parameters
    ----------
    value : str | float | int
        Input value to convert. If str, the unit can be embedded in the string.
    unit : Literal['base','K','M','G','T'] | None, optional
        The unit of the value if value is a numeric type, by default None
    to_unit : Literal['base','K','M','G','T'], optional
        The unit to convert the value to, by default 'base'
    use_binary : bool, optional
        Whether to use binary conversion (1024-based) or decimal conversion (1000-based), by default False

    Returns
    -------
    float
        The converted value
    """

    factor = 1024 if use_binary else 1000

    # Unit multipliers
    unit_multipliers = {
        'base': 1,
        'K': factor,
        'M': factor ** 2,
        'G': factor ** 3,
        'T': factor ** 4,
    }

    # If value is a string, extract the number and the unit
    if isinstance(value, str):
        # Extract numeric part and unit part
        value = value.strip()
        for suffix in ['K', 'M', 'G', 'T']:
            if value.upper().endswith(suffix):
                unit = suffix
                value = value[:-1]
                break
        else:
            unit = 'base'
        value = float(value)
    
    # If value is numeric, use the provided unit or default to 'base'
    if unit is None:
        unit = 'base'

    # Convert the input value to base
    base_value = value * unit_multipliers[unit]

    # Convert base value to the target unit
    converted_value = base_value / unit_multipliers[to_unit]
    
    return converted_value

def parse_scontrol():
    job_id = os.getenv('SLURM_JOB_ID')

    command = f"scontrol show job {job_id} | grep TRES="
    result = subprocess.run(command, shell=True, capture_output=True, text=True).stdout.strip()

    tres_pattern=r'.*cpu=(?P<cores>[\d]+),mem=(?P<mem>[\d]+[KMGT]?).*'
    cores,mem = re.search(tres_pattern,result).groupdict().values()
    
    cores = int(cores)
    mem = convert_si(mem,to_unit='G',use_binary=True)
    return [cores,mem]

def as_path(s: str | Path) -> Path:
    if not isinstance(s,Path):
        s = Path(s)
    return s