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