Shard Utilities

mmirage.shard_process

Main entry point for processing a single shard.

Main script for processing dataset shards with MMIRAGE.

Supports both text-only and multimodal (vision-language) processing.

mmirage.shard_process.rewrite_batch(batch, mapper, renderer, image_base_path=None)[source]

Rewrite a batch of samples by applying transformations. :param batch: Dictionary mapping column names to lists of values. :param mapper: MMIRAGEMapper for processing transformations. :param renderer: TemplateRenderer for generating output. :param image_base_path: Optional base directory for resolving relative image paths.

Returns:

Dictionary mapping output keys to lists of rendered values.

Raises:

ValueError – If variables are not computable given the configuration.

Parameters:
Return type:

Dict[str, List[Any]]

mmirage.shard_process.main()[source]

Process a single shard of the dataset. Loads configuration, datasets, processes the shard using MMIRAGE transformations (including multimodal), and saves the result to disk.

mmirage.shard_utils

Low-level helpers for shard state management, atomic saves, and status markers.

Utility functions for shard and merge processing.

This module contains helper functions for dataset sharding, state management, and file operations used in the MMIRAGE shard processing pipeline.

mmirage.shard_utils.format_duration(seconds)[source]

Format a duration given in seconds as a human-readable string.

Parameters:

seconds (float | None)

Return type:

str | None

class mmirage.shard_utils.ShardStats(runtime_seconds=None, rows_processed=None, throughput_rows_per_sec=None, gpu_util_mean=None, gpu_util_min=None, gpu_util_max=None, gpu_util_samples=None, input_tokens=None, output_tokens=None, num_gpus=None, model_load_seconds=None)[source]

Bases: object

Per-shard benchmark statistics recorded at completion.

Parameters:

  • runtime_seconds (float | None)

  • rows_processed (int | None)

  • throughput_rows_per_sec (float | None)

  • gpu_util_mean (float | None)

  • gpu_util_min (float | None)

  • gpu_util_max (float | None)

  • gpu_util_samples (int | None)

  • input_tokens (int | None)

  • output_tokens (int | None)

  • num_gpus (int | None)

  • model_load_seconds (float | None)

runtime_seconds: float | None = None
rows_processed: int | None = None
throughput_rows_per_sec: float | None = None
gpu_util_mean: float | None = None
gpu_util_min: float | None = None
gpu_util_max: float | None = None
gpu_util_samples: int | None = None
input_tokens: int | None = None
output_tokens: int | None = None
num_gpus: int | None = None
model_load_seconds: float | None = None
classmethod from_dict(data)[source]

Build a ShardStats from a JSON payload, or return None if data is missing.

Parameters:

data (Dict[str, Any] | None)

Return type:

ShardStats | None

to_dict()[source]
Return type:

Dict[str, Any]

class mmirage.shard_utils.GpuUtilizationPoller(interval_seconds=5.0, gpu_indices=None)[source]

Bases: object

Polls nvidia-smi in a background daemon thread.

Usage:

poller = GpuUtilizationPoller()
poller.start()
# ... do work ...
gpu_info = poller.stop()  # {"mean": 85.2, "min": 70.0, "max": 98.0, "samples": 24}

If nvidia-smi is unavailable all values are None and samples is 0.

Parameters:
start()[source]

Start background polling.

Return type:

None

stop()[source]

Stop polling and return a summary dict.

Return type:

Dict[str, Any]

class mmirage.shard_utils.ShardStatus(status='unknown', retry_count=0, shard_id=None, started_at=None, finished_at=None, error=None, hostname=None, pid=None, slurm_job_id=None, slurm_array_task_id=None, datasets=None, stats=None)[source]

Bases: object

Typed representation of the shard status.json payload.

Parameters:

  • status (str)

  • retry_count (int)

  • shard_id (int | None)

  • started_at (str | None)

  • finished_at (str | None)

  • error (str | None)

  • hostname (str | None)

  • pid (int | None)

  • slurm_job_id (str | None)

  • slurm_array_task_id (str | None)

  • datasets (List[Dict[str, Any]] | None)

  • stats (ShardStats | None)

status: str = 'unknown'
retry_count: int = 0
shard_id: int | None = None
started_at: str | None = None
finished_at: str | None = None
error: str | None = None
hostname: str | None = None
pid: int | None = None
slurm_job_id: str | None = None
slurm_array_task_id: str | None = None
datasets: List[Dict[str, Any]] | None = None
stats: ShardStats | None = None
classmethod from_dict(payload)[source]

Build a status object from a JSON payload.

Parameters:

payload (Dict[str, Any])

Return type:

ShardStatus

to_dict()[source]

Serialize status to the JSON payload written on disk.

Return type:

Dict[str, Any]

class mmirage.shard_utils.MergeReport(dataset_name, input_dir, output_dir, used_shards, merged_rows, skipped_invalid_dirs, skipped_zero_rows)[source]

Bases: object

Summary of a merge operation for one dataset directory.

Parameters:

  • dataset_name (str)

  • input_dir (str)

  • output_dir (str)

  • used_shards (int)

  • merged_rows (int)

  • skipped_invalid_dirs (int)

  • skipped_zero_rows (int)

dataset_name: str
input_dir: str
output_dir: str
used_shards: int
merged_rows: int
skipped_invalid_dirs: int
skipped_zero_rows: int
mmirage.shard_utils.shard_state_dir(shard_idx, state_root)[source]

Get central state directory for a logical shard.

Parameters:

  • shard_idx (int)

  • state_root (str)

Return type:

str

mmirage.shard_utils.read_status(state_dir)[source]

Read status.json if present.

Parameters:

state_dir (str)

Return type:

ShardStatus

mmirage.merge_shards

Dataset shard merging logic.

Merge processed dataset shards.

mmirage.merge_shards.merge_dataset_dir(dataset_dir, output_dir)[source]

Merge one dataset directory containing shard_* folders.

Parameters:

  • dataset_dir (str) – Input directory containing shard_* folders.

  • output_dir (str) – Destination directory for merged dataset.

Returns:

MergeReport with summary details.

Return type:

MergeReport

mmirage.merge_shards.merge_input_dir(input_dir, output_dir)[source]

Merge all shard datasets found under an input directory.

The input can be either: - one dataset dir containing shard_* folders directly - a parent dir containing multiple dataset subdirectories, each with shard_*

Parameters:

  • input_dir (str)

  • output_dir (str)

Return type:

List[MergeReport]

mmirage.merge_shards.merge_from_config(cfg, output_root=None)[source]

Merge shard outputs described in config.loading_params.datasets.

Parameters:

  • cfg (MMirageConfig) – Loaded MMIRAGE config.

  • output_root (str | None) – Optional destination root. If omitted, each dataset writes into <dataset.output_dir>/merged.

Returns:

Merge reports for each dataset entry.

Return type:

List[MergeReport]

mmirage.merge_shards.main()[source]

CLI entrypoint for directory-based shard merging. Scans –input-dir for dataset subdirectories containing shard_* folders. For each dataset directory, merges shard datasets and writes directly to the provided –output-dir.

mmirage.cli_utils.status

Shard status checking and retry helpers.

Shard status and retry helpers for the MMIRAGE CLI.

class mmirage.cli_utils.status.ShardSummary(total, successful, running, failed, max_retries_exceeded)[source]

Bases: object

Compact status summary for shard execution.

Parameters:

  • total (int)

  • successful (int)

  • running (int)

  • failed (int)

  • max_retries_exceeded (int)

total: int
successful: int
running: int
failed: int
max_retries_exceeded: int
mmirage.cli_utils.status.max_allowed_attempts(max_retries)[source]

Return max allowed total attempts for a shard.

Total attempts = initial attempt + max_retries.

Parameters:

max_retries (int)

Return type:

int

mmirage.cli_utils.status.is_retry_budget_exceeded(attempt_count, max_retries)[source]

Return whether a shard has exceeded the retry budget.

Parameters:

  • attempt_count (int)

  • max_retries (int)

Return type:

bool

mmirage.cli_utils.status.get_shard_status(state_dir)[source]

Read the current status and attempt counter for a shard.

Parameters:

state_dir (str)

Return type:

Tuple[str, int]

mmirage.cli_utils.status.check_failed_shards(cfg)[source]

Return retryable failed shards and a compact summary.

Parameters:

cfg (MMirageConfig)

Return type:

Tuple[List[int], ShardSummary]

mmirage.cli_utils.status.confirm_retry(count, confirm_mode)[source]

Return whether retry submission is confirmed.

Modes: - prompt: ask the user interactively - yes: submit without prompting

Parameters:

  • count (int)

  • confirm_mode (Literal['prompt', 'yes'])

Return type:

bool

mmirage.cli_utils.status.status_exit_code(failed_shards, summary)[source]

Map shard status to an exit code.

Parameters:
Return type:

int

mmirage.cli_utils.status.submit_failed_shards(cfg, config_path, failed_shards, confirm_mode, collect_stats=False)[source]

Submit retry jobs for failed shards when requested.

Parameters:
Return type:

int

mmirage.cli_utils.status.collect_bench_stats(cfg)[source]

Collect per-shard benchmark statistics and compute aggregate totals.

Returns a dict with two keys:

  • per_shard: list of dicts, one per shard, each containing the full ShardStatus payload plus a flattened stats sub-dict.

  • aggregate: rolled-up totals across all completed shards.

Shards without stats (e.g. still running or from older runs) are included in per_shard but excluded from aggregate calculations.

Parameters:

cfg (MMirageConfig)

Return type:

Dict[str, Any]

mmirage.cli_utils.slurm

SLURM submission and job monitoring helpers.

SLURM helpers for the MMIRAGE CLI.

mmirage.cli_utils.slurm.build_sbatch_script(cfg, config_path, collect_stats=False, shard_ids=None)[source]

Build the sbatch payload for shard processing.

Parameters:
Return type:

str

mmirage.cli_utils.slurm.submit_slurm_job(cfg, config_path, shard_ids=None, collect_stats=False)[source]

Submit a SLURM shard-processing job and return its job ID.

Parameters:
Return type:

int | None

mmirage.cli_utils.slurm.wait_for_slurm_job(job_id, cfg)[source]

Wait for a SLURM job array to leave the queue.

Parameters:
Return type:

None

mmirage.cli_utils.slurm.require_slurm(cfg, command_name)[source]

Ensure command can only run in SLURM mode.

Parameters:
Return type:

int

mmirage.cli_utils.runtime

Runtime environment setup helpers.

Runtime/path helpers for the MMIRAGE CLI.

mmirage.cli_utils.runtime.non_empty_path(value)[source]

Reject a path argument given as an empty or blank string.

An empty value reads as “flag not passed” further down, which would turn a dry run into a real submission.

Parameters:

value (str)

Return type:

str

mmirage.cli_utils.runtime.expand_path(path, project_root=None)[source]

Expand environment variables, user home and relative paths.

Parameters:

  • path (str)

  • project_root (str | None)

Return type:

str

mmirage.cli_utils.runtime.get_project_root(cfg)[source]

Return the configured project root, or the current working directory.

Parameters:

cfg (MMirageConfig)

Return type:

str

mmirage.cli_utils.runtime.create_directories(paths)[source]

Create directories if they do not already exist.

Parameters:

paths (Sequence[str])

Return type:

None

mmirage.cli_utils.runtime.validate_edf_env_path(cfg)[source]

Validate the optional EDF environment file path.

Parameters:

cfg (MMirageConfig)

Return type:

None

mmirage.cli_utils.runtime.add_file_logging(log_file, level)[source]

Add a file handler so logs are also written to disk.

Parameters:
Return type:

None

mmirage.cli_utils.runtime.setup_runtime(cfg, log_level)[source]

Initialize runtime-level logging.

Parameters:
Return type:

None