Source code for pyrevs.runner.config

from dataclasses import dataclass
from dataclasses import field
from pyrevs.core import MergePolicy


@dataclass(frozen=True)

[docs]
class DaskConfig:
    """Dask configuration."""

    __section__ = "dask"
    __merge_policy__ = MergePolicy.REPLACE


[docs]
    backend: str = field(
        default="local",
        metadata={
            "doc": "The backend to use. Currently only `local` and `slurm` are supported.",
        },
    )


[docs]
    slurm_config_file: str | None = field(
        default=None,
        metadata={
            "doc": "The path to the slurm config file.",
        },
    )


[docs]
    queue: str = field(
        default="regular",
        metadata={
            "doc": "The slurm queue to use.",
        },
    )


[docs]
    ntasks_per_job: int = field(
        default=1,
        metadata={
            "doc": "The number of tasks per job.",
        },
    )


[docs]
    ntasks_per_node: int = field(
        default=-1,
        metadata={
            "doc": "The number of tasks per node.",
        },
    )


[docs]
    ncores_per_worker: int = field(
        default=1,
        metadata={
            "doc": "The number of cores per worker.",
        },
    )


[docs]
    job_prologue: list[str] = field(
        default_factory=list,
        metadata={
            "doc": "The job prologue commands, included before srun in dask slurm script.",
        },
    )


[docs]
    worker_walltime: str = field(
        default="04:00:00",
        metadata={
            "doc": "The walltime for each worker, formatted as D:HH:MM:SS.",
        },
    )




@dataclass(frozen=True)

[docs]
class RunnerConfig:
    """Runner configuration."""

    __section__ = "runner"
    __merge_policy__ = MergePolicy.REPLACE


[docs]
    type: str = field(
        default="asyncio",
        metadata={
            "doc": "The type of runner to use. Currently only `asyncio` and `dask` are supported.",
        },
    )


[docs]
    nworkers: int = field(
        default=1,
        metadata={
            "doc": "The number of workers.",
        },
    )


[docs]
    dask_config: DaskConfig = field(
        default_factory=DaskConfig,
        metadata={
            "doc": "Dask configuration dictionary in runner.dask section.",
        },
    )