nemo_curator.backends.slurm_array

Module Contents

Classes

Name	Description
`SlurmArrayConfig`	Source-task sharding settings for one Slurm array task.
`SlurmArrayRetryPlan`	Outstanding shard IDs and the original logical shard configuration.
`SlurmArrayRetrySubmission`	Physical Slurm array indices and their logical shard offset.

Functions

Name	Description
`_ensure_slurm_array_run_config`	-
`_get_int_env_var`	Read an integer env var, with optional fallback/default.
`_read_slurm_array_run_config`	-
`_require_manifest_int`	-
`_slurm_array_completion_dir`	-
`build_slurm_array_completion_manifest`	Create durable completion tracking for one Slurm array shard.
`build_slurm_array_retry_submissions`	Map missing logical shards to one or more physical Slurm arrays.
`configure_slurm_array_source_filtering`	Set env vars consumed by source-stage filtering.
`filter_slurm_array_source_tasks`	Keep only source tasks assigned to the active Slurm array shard.
`find_slurm_array_retries`	Return expected shard IDs that have no completion manifest.
`format_slurm_array_indices`	Format shard indices as a compact Slurm `--array` expression.
`is_slurm_array_driver_process`	Return True for the process that owns retry metadata.
`resolve_slurm_array_config`	Resolve filtering config for source stages.
`slurm_array_shard_for_task`	Assign a task to a shard by hashing its deterministic task ID.

Data

SLURM_ARRAY_COMPLETION_DIRNAME

SLURM_ARRAY_COMPLETION_MANIFEST_NAMESPACE

SLURM_ARRAY_ENABLED_ENV_VAR

SLURM_ARRAY_MINIMUM_SHARD_INDEX_ENV_VAR

SLURM_ARRAY_RUN_CONFIG_FILENAME

SLURM_ARRAY_SHARD_INDEX_ENV_VAR

SLURM_ARRAY_TOTAL_SHARDS_ENV_VAR

_FALSE_ENV_VALUES

_TRUE_ENV_VALUES

API

class nemo_curator.backends.slurm_array.SlurmArrayConfig(
    shard_index: int,
    total_shards: int,
    minimum_shard_index: int = 0
)

Dataclass

Source-task sharding settings for one Slurm array task.

minimum_shard_index

int = 0

shard_index

int

total_shards

int

nemo_curator.backends.slurm_array.SlurmArrayConfig.from_env() -> nemo_curator.backends.slurm_array.SlurmArrayConfig | None

classmethod

Build config from Curator or Slurm env vars unless explicitly disabled.

class nemo_curator.backends.slurm_array.SlurmArrayRetryPlan(
    shard_indices: tuple[int, ...],
    total_shards: int,
    minimum_shard_index: int
)

Dataclass

Outstanding shard IDs and the original logical shard configuration.

minimum_shard_index

int

shard_indices

tuple[int, ...]

total_shards

int

class nemo_curator.backends.slurm_array.SlurmArrayRetrySubmission(
    array_indices: tuple[int, ...],
    shard_index_offset: int
)

Dataclass

Physical Slurm array indices and their logical shard offset.

array_indices

tuple[int, ...]

shard_index_offset

int

nemo_curator.backends.slurm_array._ensure_slurm_array_run_config(
    checkpoint_path: str | pathlib.Path,
    total_shards: int,
    minimum_shard_index: int
) -> pathlib.Path

nemo_curator.backends.slurm_array._get_int_env_var(
    env_var: str,
    fallback_name: str | None = None,
    default: int | None = None
) -> int

Read an integer env var, with optional fallback/default.

nemo_curator.backends.slurm_array._read_slurm_array_run_config(
    checkpoint_path: str | pathlib.Path
) -> tuple[pathlib.Path, dict[str, object]] | None

nemo_curator.backends.slurm_array._require_manifest_int(
    payload: collections.abc.Mapping[str, object],
    path: pathlib.Path,
    field: str
) -> int

nemo_curator.backends.slurm_array._slurm_array_completion_dir(
    checkpoint_path: str | pathlib.Path
) -> pathlib.Path

nemo_curator.backends.slurm_array.build_slurm_array_completion_manifest(
    checkpoint_path: str | pathlib.Path | None,
    shard_index: int,
    total_shards: int,
    minimum_shard_index: int
) -> nemo_curator.utils.retry_manifest.CompletionManifest | None

Create durable completion tracking for one Slurm array shard.

nemo_curator.backends.slurm_array.build_slurm_array_retry_submissions(
    retry_plan: nemo_curator.backends.slurm_array.SlurmArrayRetryPlan,
    max_array_size: int | None = None
) -> tuple[nemo_curator.backends.slurm_array.SlurmArrayRetrySubmission, ...]

Map missing logical shards to one or more physical Slurm arrays.

nemo_curator.backends.slurm_array.configure_slurm_array_source_filtering(
    shard_index: int,
    total_shards: int,
    minimum_shard_index: int
) -> None

Set env vars consumed by source-stage filtering.

nemo_curator.backends.slurm_array.filter_slurm_array_source_tasks(
    tasks: list[nemo_curator.tasks.Task],
    slurm_array: nemo_curator.backends.slurm_array.SlurmArrayConfig | None,
    stage_name: str
) -> list[nemo_curator.tasks.Task]

Keep only source tasks assigned to the active Slurm array shard.

nemo_curator.backends.slurm_array.find_slurm_array_retries(
    checkpoint_path: str | pathlib.Path
) -> nemo_curator.backends.slurm_array.SlurmArrayRetryPlan | None

Return expected shard IDs that have no completion manifest.

nemo_curator.backends.slurm_array.format_slurm_array_indices(
    indices: collections.abc.Iterable[int]
) -> str

Format shard indices as a compact Slurm --array expression.

nemo_curator.backends.slurm_array.is_slurm_array_driver_process() -> bool

Return True for the process that owns retry metadata.

The head node has SLURM_NODEID == 0; the variable is absent on local / single-node runs, which are also treated as head.

nemo_curator.backends.slurm_array.resolve_slurm_array_config(
    is_source_stage: bool
) -> nemo_curator.backends.slurm_array.SlurmArrayConfig | None

Resolve filtering config for source stages.

nemo_curator.backends.slurm_array.slurm_array_shard_for_task(
    task: nemo_curator.tasks.Task,
    slurm_array: nemo_curator.backends.slurm_array.SlurmArrayConfig
) -> int

Assign a task to a shard by hashing its deterministic task ID.

nemo_curator.backends.slurm_array.SLURM_ARRAY_COMPLETION_DIRNAME = '.slurm_array_completion'

nemo_curator.backends.slurm_array.SLURM_ARRAY_COMPLETION_MANIFEST_NAMESPACE = 'slurm_array'

nemo_curator.backends.slurm_array.SLURM_ARRAY_ENABLED_ENV_VAR = 'NEMO_CURATOR_SLURM_ARRAY_ENABLED'

nemo_curator.backends.slurm_array.SLURM_ARRAY_MINIMUM_SHARD_INDEX_ENV_VAR = 'NEMO_CURATOR_SLURM_ARRAY_MINIMUM_SHARD_INDEX'

nemo_curator.backends.slurm_array.SLURM_ARRAY_RUN_CONFIG_FILENAME = 'run.json'

nemo_curator.backends.slurm_array.SLURM_ARRAY_SHARD_INDEX_ENV_VAR = 'NEMO_CURATOR_SLURM_ARRAY_SHARD_INDEX'

nemo_curator.backends.slurm_array.SLURM_ARRAY_TOTAL_SHARDS_ENV_VAR = 'NEMO_CURATOR_SLURM_ARRAY_TOTAL_SHARDS'

nemo_curator.backends.slurm_array._FALSE_ENV_VALUES = {'0', 'false', 'no', 'off'}

nemo_curator.backends.slurm_array._TRUE_ENV_VALUES = {'1', 'true', 'yes', 'on'}