Skip to content

costs

tit.jobs.costs

Per-kind default resource costs (TODO.md §2.4), data + one lookup function.

Placeholders (documented as such in TODO.md) until the resource-usage spike lands; a project can override any kind's default in settings later. A job's config may also carry its own cpus/memory_gb (QSIPrep-shaped configs do) — that always wins over the table below.

default_cost

default_cost(kind: str, config: dict[str, Any] | None = None) -> Cost

The default :class:~tit.jobs.spec.Cost for one job.

A config carrying its own cpus/memory_gb (or mem_gb) — QSIPrep/QSIRecon-shaped configs, and FlexConfig.cpus for the cpus half — overrides the table. sim scales memory by the number of montages in the config (sim 1 cpu / 4 GB per montage, run sequentially by the runner so cpus stays flat).

Source code in tit/jobs/costs.py
def default_cost(kind: str, config: dict[str, Any] | None = None) -> Cost:
    """The default :class:`~tit.jobs.spec.Cost` for one job.

    A config carrying its own ``cpus``/``memory_gb`` (or ``mem_gb``) — QSIPrep/QSIRecon-shaped
    configs, and ``FlexConfig.cpus`` for the ``cpus`` half — overrides the table.  ``sim`` scales
    memory by the number of montages in the config (``sim 1 cpu / 4 GB per montage``, run
    sequentially by the runner so cpus stays flat).
    """
    config = config or {}
    base = DEFAULT_COSTS.get(kind, _FALLBACK)
    if kind == "pre" and config.get("run_fastsurfer"):
        # Upstream minimum system memory for segmentation; keep other pre stages unchanged.
        base = Cost(
            cpus=effective_threads("fastsurfer", config.get("fastsurfer_threads")),
            mem_gb=8,
        )
    if kind == "pre" and config.get("run_freesurfer"):
        base = Cost(
            cpus=max(
                base.cpus,
                effective_threads("freesurfer", config.get("freesurfer_threads")),
            ),
            mem_gb=16,
        )
    if kind == "pre" and config.get("create_m2m"):
        base = Cost(
            cpus=max(
                base.cpus, effective_threads("charm", config.get("charm_threads"))
            ),
            mem_gb=base.mem_gb,
        )
    if kind == "pre":
        for tool, field in (
            ("qsiprep", "qsiprep_config"),
            ("qsirecon", "qsi_recon_config"),
        ):
            if config.get(f"run_{tool}"):
                resources = config.get(field) or {}
                base = Cost(
                    cpus=max(base.cpus, effective_threads(tool, resources.get("cpus"))),
                    mem_gb=max(
                        base.mem_gb, _num(resources.get("memory_gb")) or base.mem_gb
                    ),
                )
    if kind == "blender" and config.get("_type") in (
        "VectorConfig",
        "RegionConfig",
        "SubcorticalConfig",
    ):
        # These geometry exports do not launch Blender or evaluate montage subdivision.
        base = Cost(cpus=1, mem_gb=2)
    cpus = _num(config.get("cpus"))
    mem = _num(config.get("memory_gb"))
    if mem is None:
        mem = _num(config.get("mem_gb"))

    if kind in ("ex", "mex", "stats"):
        # These fork one worker per CPU (exhaustive searches; cluster-permutation stats):
        # `n_jobs < 1` means "the whole budget this job is admitted with", so the plan claims the
        # user's global CPU limit, and an explicit `n_jobs` is clamped to it
        # (tit.opt.ex.parallel.resolve_n_jobs does the same at run time).
        limit = cpu_limit()
        n_jobs = _num(config.get("n_jobs"))
        if n_jobs is not None and n_jobs >= 1:
            return Cost(cpus=min(int(n_jobs), limit), mem_gb=base.mem_gb)
        return Cost(cpus=limit, mem_gb=base.mem_gb)

    if kind == "sim":
        n_montages = _n_montages(config)
        return Cost(cpus=base.cpus, mem_gb=base.mem_gb * max(n_montages, 1))

    if cpus is not None or mem is not None:
        return Cost(
            cpus=cpus if cpus is not None else base.cpus,
            mem_gb=mem if mem is not None else base.mem_gb,
        )

    return base