Tile sizing

patchworks.auto_tile_shape(shape: tuple[int, ...], dtype: Any, target_bytes: int = 64 * 1024 ** 2, use_gpu: bool = False, gpu_memory: int | None = None, available_memory: int | None = None, n_workers: int | None = None, n_channels: int = 1, verbose: bool = False) -> tuple[int, ...]

Balanced tile shape for general-purpose 3-D processing.

Sizes the last three axes (spatial) to stay within the memory budget while keeping the shape as cubic as possible. Leading axes (t, c) are always 1.

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Full array shape, e.g. (z, y, x) or (t, c, z, y, x).	required
dtype	Any	Array dtype.	required
target_bytes	int	Memory ceiling per tile. Default 64 MiB.	64 * 1024 ** 2
use_gpu	bool	Size tiles against GPU VRAM rather than host RAM. Still capped by host RAM too -- see available_memory.	False
gpu_memory	int | None	Available GPU VRAM in bytes; auto-queried when None.	None
available_memory	int | None	Available host RAM in bytes; auto-queried when None. Consulted even when use_gpu is set: the tile has to be loaded into host memory before (and often after) it reaches the GPU, so it must fit both budgets, not just VRAM.	None
n_workers	int | None	Number of parallel workers (divides the RAM budget).	None
n_channels	int	Channels each tile carries (default 1). Above 1 the per-voxel cost scales with it, so the tile shrinks accordingly -- e.g. the workflow's nuclei_channel hands Cellpose a cyto+nuclei pair.	1
verbose	bool	Log the chosen shape and estimated tile size.	False

Returns:

Type	Description
tuple[int, ...]	Tile shape with the same number of dimensions as shape.

Examples:

>>> tile = auto_tile_shape((128, 2048, 2048), "uint16")
>>> tile
(128, 512, 512)

Source code in src/patchworks/_chunks.py

def auto_tile_shape(
    shape: tuple[int, ...],
    dtype: Any,
    target_bytes: int = 64 * 1024**2,
    use_gpu: bool = False,
    gpu_memory: int | None = None,
    available_memory: int | None = None,
    n_workers: int | None = None,
    n_channels: int = 1,
    verbose: bool = False,
) -> tuple[int, ...]:
    """Balanced tile shape for general-purpose 3-D processing.

    Sizes the last three axes (spatial) to stay within the memory budget while
    keeping the shape as cubic as possible. Leading axes (t, c) are always 1.

    Parameters
    ----------
    shape:
        Full array shape, e.g. ``(z, y, x)`` or ``(t, c, z, y, x)``.
    dtype:
        Array dtype.
    target_bytes:
        Memory ceiling per tile. Default 64 MiB.
    use_gpu:
        Size tiles against GPU VRAM rather than host RAM. Still capped by
        host RAM too -- see ``available_memory``.
    gpu_memory:
        Available GPU VRAM in bytes; auto-queried when None.
    available_memory:
        Available host RAM in bytes; auto-queried when None. Consulted even
        when ``use_gpu`` is set: the tile has to be loaded into host memory
        before (and often after) it reaches the GPU, so it must fit both
        budgets, not just VRAM.
    n_workers:
        Number of parallel workers (divides the RAM budget).
    n_channels:
        Channels each tile carries (default 1). Above 1 the per-voxel cost
        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
    verbose:
        Log the chosen shape and estimated tile size.

    Returns
    -------
    tuple[int, ...]
        Tile shape with the same number of dimensions as *shape*.

    Examples
    --------
    >>> tile = auto_tile_shape((128, 2048, 2048), "uint16")
    >>> tile
    (128, 512, 512)
    """
    n_workers = n_workers or cpu_allocation()
    # A tile holds n_channels planes per voxel (e.g. Cellpose's
    # cyto+nuclei pair), so the per-voxel cost -- and every budget
    # derived from it below -- scales with them.
    if n_channels < 1:
        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
    itemsize = np.dtype(dtype).itemsize * n_channels
    n_spatial = min(3, len(shape))

    if use_gpu:
        # A tile that fits VRAM still has to be decoded into (and often
        # copied back out of) host RAM first, so a GPU with more memory
        # than the job's own host allocation must not produce a tile the
        # job itself can't hold -- take whichever budget is tighter.
        gpu_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
        host_mem = available_memory or _get_available_memory()
        mem = min(gpu_mem, host_mem)
        budget = min(target_bytes * 2, mem // 2)
    else:
        mem = available_memory or _get_available_memory()
        budget = min(target_bytes, mem // (n_workers * 4))

    budget = max(32 * 1024**2, budget)

    leading = [1] * (len(shape) - n_spatial)
    spatial = list(shape[-n_spatial:])
    target_voxels = budget / itemsize
    target_side = int(target_voxels ** (1.0 / n_spatial))
    chunk_spatial = [min(s, target_side) for s in spatial]

    capped = [
        i for i, (c, s) in enumerate(zip(chunk_spatial, spatial)) if c == s
    ]
    uncapped = [i for i in range(n_spatial) if i not in capped]
    if uncapped:
        used_by_capped = (
            np.prod([chunk_spatial[i] for i in capped]) if capped else 1
        )
        remaining = target_voxels / max(1, used_by_capped)
        new_side = int(remaining ** (1.0 / len(uncapped)))
        for i in uncapped:
            chunk_spatial[i] = min(spatial[i], new_side)

    result = tuple(leading + chunk_spatial)

    if verbose:
        mib = np.prod(result) * itemsize / 1024**2
        logger.info(
            "auto_tile_shape: shape=%s dtype=%s → tiles=%s (~%.0f MiB/tile)",
            shape,
            np.dtype(dtype).name,
            result,
            mib,
        )

    return result

patchworks.auto_tile_shape_cellpose(shape: tuple[int, ...], dtype: Any, diameter: float | None = None, do_3D: bool = False, use_gpu: bool = False, gpu_memory: int | None = None, available_memory: int | None = None, n_workers: int | None = None, model_memory_bytes: int = 2 * 1024 ** 3, cellpose_memory_factor: int = 20, n_channels: int = 1, anisotropy: float | None = None, verbose: bool = False) -> tuple[int, ...]

Cellpose-optimised tile shape.

Cellpose is fundamentally 2-D: even in 3-D mode it runs 2-D segmentation on orthogonal planes and takes a consensus.

do_3D=False (default) z is set to 1. Each tile is one 2-D (y, x) slice.

do_3D=True z is kept at its full extent per tile. y and x are tiled based on the available memory, accounting for the 3× overhead of three plane orientations.

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Spatial shape, e.g. (z, y, x).	required
dtype	Any	Array dtype.	required
diameter	float | None	Expected cell diameter in pixels. Tile will be at least 4 × diameter.	None
do_3D	bool	Whether Cellpose will run in 3-D mode.	False
use_gpu	bool	Size tiles for GPU VRAM. Still capped by host RAM too -- see available_memory.	False
gpu_memory	int | None	Memory parameters (auto-queried when None). available_memory is consulted even when use_gpu is set: Cellpose's do_3D preprocessing (orthogonal-view flows, mask assembly) copies the tile through host memory, so VRAM headroom alone doesn't guarantee the job's own host allocation can hold it.	None
available_memory	int | None	Memory parameters (auto-queried when None). available_memory is consulted even when use_gpu is set: Cellpose's do_3D preprocessing (orthogonal-view flows, mask assembly) copies the tile through host memory, so VRAM headroom alone doesn't guarantee the job's own host allocation can hold it.	None
n_workers	int | None	Memory parameters (auto-queried when None). available_memory is consulted even when use_gpu is set: Cellpose's do_3D preprocessing (orthogonal-view flows, mask assembly) copies the tile through host memory, so VRAM headroom alone doesn't guarantee the job's own host allocation can hold it.	None
model_memory_bytes	int	Memory consumed by the Cellpose model weights (default 2 GiB).	2 * 1024 ** 3
cellpose_memory_factor	int	Cellpose allocates roughly this multiple of raw input bytes (default 20×). Applied against whichever of VRAM/host RAM is tighter -- ponytail: this reuses one factor for both budgets as an approximation; if it proves off in practice, calibrate a separate host-side factor from a SLURM job's peak RSS (seff/ /usr/bin/time -v) and pass it alongside this one.	20
n_channels	int	Channels each tile carries (default 1). Above 1 the per-voxel cost scales with it, so the tile shrinks accordingly -- e.g. the workflow's nuclei_channel hands Cellpose a cyto+nuclei pair.	1
anisotropy	float | None	The anisotropy Cellpose will be given (do_3D only). It is not merely a hint to the model: Cellpose resizes the tile to z * anisotropy planes before the net runs, so a 2.2x anisotropy costs 2.2x the z it was handed. Budgeting against the unscaled z under-counts by exactly that factor. None (the default) assumes isotropic, i.e. no resize.	None
verbose	bool	Log the chosen shape and memory estimates.	False

Returns:

Type	Description
tuple[int, ...]	Tile shape with the same number of dimensions as shape.

Examples:

>>> tile = auto_tile_shape_cellpose((128, 2048, 2048), "uint16", diameter=30)
>>> tile
(1, 2048, 2048)

Source code in src/patchworks/_chunks.py

def auto_tile_shape_cellpose(
    shape: tuple[int, ...],
    dtype: Any,
    diameter: float | None = None,
    do_3D: bool = False,
    use_gpu: bool = False,
    gpu_memory: int | None = None,
    available_memory: int | None = None,
    n_workers: int | None = None,
    model_memory_bytes: int = 2 * 1024**3,
    cellpose_memory_factor: int = 20,
    n_channels: int = 1,
    anisotropy: float | None = None,
    verbose: bool = False,
) -> tuple[int, ...]:
    """Cellpose-optimised tile shape.

    Cellpose is fundamentally 2-D: even in 3-D mode it runs 2-D segmentation
    on orthogonal planes and takes a consensus.

    **do_3D=False (default)**
        z is set to 1. Each tile is one 2-D ``(y, x)`` slice.

    **do_3D=True**
        z is kept at its full extent per tile. y and x are tiled based on the
        available memory, accounting for the 3× overhead of three plane orientations.

    Parameters
    ----------
    shape:
        Spatial shape, e.g. ``(z, y, x)``.
    dtype:
        Array dtype.
    diameter:
        Expected cell diameter in pixels. Tile will be at least ``4 × diameter``.
    do_3D:
        Whether Cellpose will run in 3-D mode.
    use_gpu:
        Size tiles for GPU VRAM. Still capped by host RAM too -- see
        ``available_memory``.
    gpu_memory, available_memory, n_workers:
        Memory parameters (auto-queried when None). ``available_memory`` is
        consulted even when ``use_gpu`` is set: Cellpose's ``do_3D``
        preprocessing (orthogonal-view flows, mask assembly) copies the tile
        through host memory, so VRAM headroom alone doesn't guarantee the
        job's own host allocation can hold it.
    model_memory_bytes:
        Memory consumed by the Cellpose model weights (default 2 GiB).
    cellpose_memory_factor:
        Cellpose allocates roughly this multiple of raw input bytes (default
        20×). Applied against whichever of VRAM/host RAM is tighter --
        ``ponytail:`` this reuses one factor for both budgets as an
        approximation; if it proves off in practice, calibrate a separate
        host-side factor from a SLURM job's peak RSS (``seff``/
        ``/usr/bin/time -v``) and pass it alongside this one.
    n_channels:
        Channels each tile carries (default 1). Above 1 the per-voxel cost
        scales with it, so the tile shrinks accordingly -- e.g. the workflow's
        ``nuclei_channel`` hands Cellpose a cyto+nuclei pair.
    anisotropy:
        The ``anisotropy`` Cellpose will be given (``do_3D`` only). It is not
        merely a hint to the model: Cellpose resizes the tile to
        ``z * anisotropy`` planes before the net runs, so a 2.2x anisotropy
        costs 2.2x the z it was handed. Budgeting against the unscaled ``z``
        under-counts by exactly that factor. ``None`` (the default) assumes
        isotropic, i.e. no resize.
    verbose:
        Log the chosen shape and memory estimates.

    Returns
    -------
    tuple[int, ...]
        Tile shape with the same number of dimensions as *shape*.

    Examples
    --------
    >>> tile = auto_tile_shape_cellpose((128, 2048, 2048), "uint16", diameter=30)
    >>> tile
    (1, 2048, 2048)
    """
    n_workers = n_workers or cpu_allocation()
    # Cellpose resizes the tile before the net runs -- by `rescale`
    # (= 30 / diameter) on every axis, and by `anisotropy` on z as well -- so
    # the array it actually holds is bigger than the one it was handed, and a
    # budget computed from the unresized tile under-counts by that factor. A
    # diameter half the model's 30 px means a 2x upsample per axis: 8x the
    # voxels, enough to turn a comfortable tile into an OOM.
    #
    # Both only ever *shrink* the tile. A predicted downsample (diameter > 30)
    # would license a bigger one, but these are a safety margin against a
    # rough memory model, not a measurement to spend headroom on.
    rescale = max(1.0, 30.0 / diameter) if diameter else 1.0
    z_resize = max(1.0, anisotropy or 1.0)
    # A tile holds n_channels planes per voxel (e.g. Cellpose's
    # cyto+nuclei pair), so the per-voxel cost -- and every budget
    # derived from it below -- scales with them.
    if n_channels < 1:
        raise ValueError(f"n_channels must be >= 1; got {n_channels!r}")
    itemsize = np.dtype(dtype).itemsize * n_channels

    if use_gpu:
        # A tile that fits VRAM still has to be decoded into (and often
        # copied back out of) host RAM first, so a GPU with more memory
        # than the job's own host allocation must not produce a tile the
        # job itself can't hold -- take whichever budget is tighter.
        gpu_mem = gpu_memory if gpu_memory is not None else _get_gpu_memory()
        host_mem = available_memory or _get_available_memory()
        total_mem = min(gpu_mem, host_mem)
    else:
        total_mem = (available_memory or _get_available_memory()) // n_workers

    usable = max(32 * 1024**2, total_mem - model_memory_bytes)
    max_raw_bytes = usable // cellpose_memory_factor

    n_spatial = min(3, len(shape))
    leading = [1] * (len(shape) - n_spatial)
    min_tile = int(4 * diameter) if diameter is not None else 1

    if n_spatial == 2 or not do_3D:
        # Two axes resized, so the cost per configured pixel is rescale**2.
        max_pixels_2d = max(1, int(max_raw_bytes // (itemsize * rescale**2)))
        tile_side = max(min_tile, int(max_pixels_2d**0.5))
        if n_spatial == 2:
            y, x = shape[-2], shape[-1]
            chunk_spatial = [min(y, tile_side), min(x, tile_side)]
        else:
            z, y, x = shape[-3], shape[-2], shape[-1]
            chunk_spatial = [1, min(y, tile_side), min(x, tile_side)]
    else:
        z, y, x = shape[-3], shape[-2], shape[-1]
        # All three axes are resized: z by anisotropy * rescale, y/x by
        # rescale each -- so one configured voxel costs
        # anisotropy * rescale**3 of them.
        effective_z = z * z_resize * rescale**3
        max_pixels_per_slice = max(
            1, int((max_raw_bytes // 3) // (effective_z * itemsize))
        )
        tile_side = max(min_tile, int(max_pixels_per_slice**0.5))
        chunk_spatial = [z, min(y, tile_side), min(x, tile_side)]

    result = tuple(leading + chunk_spatial)

    if verbose:
        raw_mib = np.prod(result) * itemsize / 1024**2
        logger.info(
            "auto_tile_shape_cellpose: shape=%s dtype=%s do_3D=%s "
            "→ tiles=%s (~%.0f MiB raw, ~%.0f MiB Cellpose estimate)",
            shape,
            np.dtype(dtype).name,
            do_3D,
            result,
            raw_mib,
            raw_mib * cellpose_memory_factor,
        )

    return result

patchworks.auto_overlap(diameter: float, safety: float = 1.0, voxel_size: Union[Sequence[float], None] = None) -> Union[int, tuple[int, ...]]

Recommended overlap (halo) for a given cell diameter.

Rule: overlap >= diameter so the segmentation function always sees at least one full cell's worth of context on every tile edge. Cells near tile boundaries are then segmented correctly and only genuinely split cells produce touching labels at the boundary → correct merge.

With voxel_size the halo is returned per axis instead of as one number. That matters on anisotropic stacks: a halo big enough laterally is far more than one cell deep in z, and the extra planes are read and segmented only to be trimmed away again.

Parameters:

Name	Type	Description	Default
diameter	float	Expected cell diameter in lateral pixels (same unit as your image's x/y).	required
safety	float	Multiplier on top of diameter. Default 1.0 (= one cell width). Use 1.5–2.0 for elongated or irregularly-shaped cells.	1.0
voxel_size	Union[Sequence[float], None]	Physical size per axis, in any single unit (e.g. (2.0, 0.1, 0.1) for a 2 µm z-step and 100 nm pixels). None → one isotropic number, as before.	None

Returns:

Type	Description
int or tuple of int	Overlap depth to pass to tile_process(..., overlap=...). A tuple (one entry per axis) when voxel_size is given.

Examples:

>>> from patchworks import auto_overlap, tile_process
>>> from patchworks.plugins.cellpose import cellpose_fn
>>>
>>> fn = cellpose_fn("cyto3", gpu=True, diameter=30)
>>> result = tile_process("image.zarr", fn,
...                       tile_shape=(1, 2048, 2048),
...                       overlap=auto_overlap(30))
>>> auto_overlap(15, voxel_size=(2.0, 0.1, 0.1))
(1, 15, 15)

Source code in src/patchworks/_chunks.py

def auto_overlap(
    diameter: float,
    safety: float = 1.0,
    voxel_size: Union[Sequence[float], None] = None,
) -> Union[int, tuple[int, ...]]:
    """Recommended overlap (halo) for a given cell diameter.

    Rule: overlap >= diameter so the segmentation function always sees at
    least one full cell's worth of context on every tile edge. Cells near
    tile boundaries are then segmented correctly and only genuinely split
    cells produce touching labels at the boundary → correct merge.

    With *voxel_size* the halo is returned per axis instead of as one number.
    That matters on anisotropic stacks: a halo big enough laterally is far
    more than one cell deep in z, and the extra planes are read and
    segmented only to be trimmed away again.

    Parameters
    ----------
    diameter:
        Expected cell diameter in **lateral** pixels (same unit as your
        image's x/y).
    safety:
        Multiplier on top of diameter. Default 1.0 (= one cell width).
        Use 1.5–2.0 for elongated or irregularly-shaped cells.
    voxel_size:
        Physical size per axis, in any single unit (e.g. ``(2.0, 0.1, 0.1)``
        for a 2 µm z-step and 100 nm pixels). ``None`` → one isotropic
        number, as before.

    Returns
    -------
    int or tuple of int
        Overlap depth to pass to ``tile_process(..., overlap=...)``. A tuple
        (one entry per axis) when *voxel_size* is given.

    Examples
    --------
    >>> from patchworks import auto_overlap, tile_process  # doctest: +SKIP
    >>> from patchworks.plugins.cellpose import cellpose_fn  # doctest: +SKIP
    >>>
    >>> fn = cellpose_fn("cyto3", gpu=True, diameter=30)  # doctest: +SKIP
    >>> result = tile_process("image.zarr", fn,  # doctest: +SKIP
    ...                       tile_shape=(1, 2048, 2048),
    ...                       overlap=auto_overlap(30))
    >>> auto_overlap(15, voxel_size=(2.0, 0.1, 0.1))
    (1, 15, 15)
    """
    lateral = max(1, int(np.ceil(diameter * safety)))
    if voxel_size is None:
        return lateral
    # Convert the lateral halo to a physical distance, then back into pixels
    # along each axis using that axis' own voxel size.
    physical = diameter * safety * float(voxel_size[-1])
    return tuple(max(1, int(np.ceil(physical / float(v)))) for v in voxel_size)

patchworks.normalize_overlap(overlap: Overlap, ndim: int, tile_shape: 'Sequence[int] | None' = None) -> tuple[int, ...]

Expand an overlap spec to one halo width per axis.

A scalar applies the same halo to every axis (the historical behaviour). A sequence gives the halo per axis, which matters for anisotropic tiles: a (16, 1024, 1024) tile with a scalar overlap of 30 reads 76 x 1084 x 1084 to keep 16 x 1024 x 1024 -- 5.3x more voxels than it uses, nearly all of it in z.

With tile_shape, an axis only one voxel thick gets no halo. There is no context to gather along an axis the tile does not span, and a 2-D method handed the extra planes would read them as channels. This is the tile_shape: "auto" + do_3D: false case, where tiles come out one plane thick: a z-overlap of 4 would otherwise read 9 planes per tile to keep 1.

Parameters:

Name	Type	Description	Default
overlap	int or sequence of int	Halo width, shared or per-axis.	required
ndim	int	Number of axes the halo is applied to.	required
tile_shape	sequence of int	Tile extent per axis. Used to drop halos an axis has no room for.	None

Returns:

Type	Description
tuple of int	One non-negative halo width per axis.

Source code in src/patchworks/_distributed.py

def normalize_overlap(
    overlap: Overlap,
    ndim: int,
    tile_shape: "Sequence[int] | None" = None,
) -> tuple[int, ...]:
    """Expand an overlap spec to one halo width per axis.

    A scalar applies the same halo to every axis (the historical behaviour).
    A sequence gives the halo per axis, which matters for anisotropic tiles:
    a ``(16, 1024, 1024)`` tile with a scalar overlap of 30 reads
    ``76 x 1084 x 1084`` to keep ``16 x 1024 x 1024`` -- 5.3x more voxels than
    it uses, nearly all of it in z.

    With *tile_shape*, an axis only one voxel thick gets **no** halo. There is
    no context to gather along an axis the tile does not span, and a 2-D
    method handed the extra planes would read them as channels. This is the
    ``tile_shape: "auto"`` + ``do_3D: false`` case, where tiles come out one
    plane thick: a z-overlap of 4 would otherwise read 9 planes per tile to
    keep 1.

    Parameters
    ----------
    overlap : int or sequence of int
        Halo width, shared or per-axis.
    ndim : int
        Number of axes the halo is applied to.
    tile_shape : sequence of int, optional
        Tile extent per axis. Used to drop halos an axis has no room for.

    Returns
    -------
    tuple of int
        One non-negative halo width per axis.
    """
    if isinstance(overlap, (int, np.integer)):
        values = (int(overlap),) * ndim
    else:
        values = tuple(int(o) for o in overlap)
        if len(values) != ndim:
            raise ValueError(
                f"overlap has {len(values)} entries but the tile is {ndim}-D"
            )
    if any(o < 0 for o in values):
        raise ValueError(f"overlap must be non-negative, got {values}")

    if tile_shape is not None:
        clipped = tuple(
            0 if int(t) <= 1 else o for o, t in zip(values, tile_shape)
        )
        if clipped != values:
            dropped = [
                i for i, (a, b) in enumerate(zip(values, clipped)) if a != b
            ]
            logger.info(
                "dropping the halo on axis %s: the tile is 1 voxel thick "
                "there, so there is no context to read.",
                dropped,
            )
        values = clipped
    return values

Cluster resource detection

On a shared node the machine's core count and free RAM say nothing about what this job was granted. These read the allocation instead, and everything that sizes a worker pool goes through them.

patchworks.cpu_allocation() -> int

Return the number of CPUs this process may actually use.

os.cpu_count() reports the machine's cores, which on a shared cluster node is wildly more than a job was granted -- a 4-core allocation on a 128-core node would size itself for 128. Prefer what the scheduler says, then the process' CPU affinity mask, and only then the machine.

Returns:

Type	Description
int	Usable CPU count (always >= 1).

Source code in src/patchworks/_chunks.py

def cpu_allocation() -> int:
    """Return the number of CPUs this process may actually use.

    ``os.cpu_count()`` reports the machine's cores, which on a shared cluster
    node is wildly more than a job was granted -- a 4-core allocation on a
    128-core node would size itself for 128. Prefer what the scheduler says,
    then the process' CPU affinity mask, and only then the machine.

    Returns
    -------
    int
        Usable CPU count (always >= 1).
    """
    for var in ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"):
        try:
            value = int(os.environ[var])
        except (KeyError, ValueError):
            continue
        if value > 0:
            return value
    try:
        n = max(1, len(os.sched_getaffinity(0)))
    except AttributeError:  # not POSIX
        n = max(1, os.cpu_count() or 1)
    quota = _cgroup_cpu_limit()
    return min(n, quota) if quota is not None else n

patchworks.safe_worker_count(tile_nbytes: int, *, use_gpu: bool = False, fn_overhead: int = 4, ram_fraction: float = 0.8) -> int

Concurrent tiles that fit the machine without OOM or a CPU freeze.

Bounds the threaded scheduler by two limits and takes the smaller:

• CPU — leaves at least one core free so the box stays responsive (never pins every core).
• RAM — at most ram_fraction of available memory, assuming each in-flight tile needs fn_overhead copies (halo + output + temporaries).

On GPU the answer is always 1: one evaluation at a time so concurrent tiles can never exhaust VRAM. Without psutil it returns a conservative default rather than guessing high.

Parameters:

Name	Type	Description	Default
tile_nbytes	int	Size of one tile in bytes (prod(tile_shape) * dtype.itemsize).	required
use_gpu	bool	Whether tiles are processed on the GPU.	False
fn_overhead	int	Assumed peak number of tile-sized buffers alive per worker.	4
ram_fraction	float	Fraction of available RAM the staging step may use.	0.8

Returns:

Type	Description
int	Worker-thread count (always >= 1).

Source code in src/patchworks/_chunks.py

def safe_worker_count(
    tile_nbytes: int,
    *,
    use_gpu: bool = False,
    fn_overhead: int = 4,
    ram_fraction: float = 0.8,
) -> int:
    """Concurrent tiles that fit the machine without OOM or a CPU freeze.

    Bounds the threaded scheduler by two limits and takes the smaller:

    * **CPU** — leaves at least one core free so the box stays responsive
      (never pins every core).
    * **RAM** — at most ``ram_fraction`` of available memory, assuming each
      in-flight tile needs ``fn_overhead`` copies (halo + output + temporaries).

    On GPU the answer is always 1: one evaluation at a time so concurrent
    tiles can never exhaust VRAM. Without ``psutil`` it returns a conservative
    default rather than guessing high.

    Parameters
    ----------
    tile_nbytes : int
        Size of one tile in bytes (``prod(tile_shape) * dtype.itemsize``).
    use_gpu : bool, optional
        Whether tiles are processed on the GPU.
    fn_overhead : int, optional
        Assumed peak number of tile-sized buffers alive per worker.
    ram_fraction : float, optional
        Fraction of available RAM the staging step may use.

    Returns
    -------
    int
        Worker-thread count (always >= 1).
    """
    cpu_cap = max(1, cpu_allocation() - 1)
    if use_gpu:
        return 1
    avail = _get_available_memory()
    per_tile = max(1, int(tile_nbytes) * max(1, fn_overhead))
    mem_cap = max(1, int(avail * ram_fraction) // per_tile)
    return max(1, min(cpu_cap, mem_cap))