API Reference

Data

Image

dl_utils.data.image.byte_imread(data: bytes) → ndarray

dl_utils.data.image.byte_imwrite(image: ndarray, format='PNG', **kwargs) → bytes

Video

dl_utils.data.video.save_video(frames: ndarray | Tensor | list | tuple, save_path: str | PathLike, fps: int | float = 30, codec: str = 'avc1')

Save video frames to a video file utilizing opencv.

Parameters:

• frames – Video frames in shape (F, H, W, C). The pixel values should be in range [0, 255].

• save_path – Path to save video.

• fps – FPS of video, default 30.

• codec – Codec of video, default avc1.

dl_utils.data.video.load_video(video_path: str | PathLike, resize: Tuple[int, int] | int = None, center_crop: Tuple[int, int] | int = None, max_frames: int = None) → ndarray

Load a video file.

Parameters:

• video_path – Path to the video file.

• resize – Resize frames to the specified size. If None, no resizing. Accepts (width, height) or int.

• center_crop – Center crop frames to the specified size. If None, no cropping. Accepts (width, height) or int.

• max_frames – Maximum number of frames to load. If None, load all frames.

Returns:

Frames as a NumPy array with shape (F, H, W, C). Pixel values are in [0, 255], color order is RGB.

Note

• If the video is grayscale, the color channel will be replicated to 3.

dl_utils.data.video.get_video_fps(video_path: str | PathLike) → float

Retrieve the FPS of a video.

Parameters:

video_path – Path to the video file.

Returns:

The FPS of the video.

dl_utils.data.video.get_video_frame_count(video_path: str | PathLike) → int

Retrieve the total number of frames in a video.

Parameters:

video_path – Path to the video file.

Returns:

The number of frames in the video.

dl_utils.data.video.get_video_duration(video_path: str | PathLike) → Tuple[float, int, float]

Retrieve the FPS, frame count, and duration (in seconds) of a video.

Parameters:

video_path – Path to the video file.

Returns:

A tuple containing FPS, frame count, and duration in seconds.

dl_utils.data.video.get_video_duration_batch(video_paths: List[str | PathLike]) → List[float]

Get duration of videos in batch.

Parameters:

video_paths – List of paths to videos.

Returns:

A list of tuples, each containing FPS, frame count, and duration in seconds.

dl_utils.data.video.convert_to_h265(input_file: AnyStr, output_file: AnyStr, ffmpeg_exec: AnyStr = '/usr/bin/ffmpeg', keyint: int = None, overwrite: bool = False, verbose: bool = False) → None

convert video to h265 format using ffmpeg @param input_file: input path @param output_file: output path @param ffmpeg_exec: @param keyint: @param overwrite: overwrite the existing file @param verbose: show ffmpeg output

dl_utils.data.video.convert_to_h264(input_file: AnyStr, output_file: AnyStr, ffmpeg_exec: AnyStr = '/usr/bin/ffmpeg', keyint: int = None, overwrite: bool = False, verbose: bool = False) → None

Array

dl_utils.data.array.to_numpy(array: ndarray | Tensor | list | tuple | int | float) → ndarray

Convert array-like object or scalar to NumPy array.

Parameters:

array – Array-like object or scalar to be converted.

dl_utils.data.array.to_tensor(array: ndarray | Tensor | list | tuple | int | float, *args, **kwargs) → Tensor

Convert a scalar or array-like object to a PyTorch tensor.

Parameters:

• array – Numeric scalar or array-like object to convert.

• *args – Additional arguments passed to torch.Tensor.to()

• **kwargs – Additional keyword arguments passed to torch.Tensor.to()

dl_utils.data.array.to_original(array: ndarray | Tensor | list | tuple, ori_dtype) → ndarray | Tensor | list | tuple

Convert array-like object to original type.

dl_utils.data.array.is_tensor(x)

Is torch.Tensor

dl_utils.data.array.is_ndarray(x)

Is np.ndarray

Normalization

dl_utils.data.normalize.normalize(data: ndarray | Tensor | list | tuple, mean: ndarray | Tensor | list | tuple | int | float, std: ndarray | Tensor | list | tuple | int | float, dim: int = -1) → ndarray | Tensor | list | tuple

Normalize the input array (usually image or video).

Parameters:

• data – Input array, can be a NumPy array or a PyTorch tensor.

• mean – Numeric scalar or vector of means for each channel.

• std – Numeric scalar or vector of standard deviations for each channel.

• dim – The channel dimension to normalize along. Default is -1 (last dimension).

Returns:

Normalized image or video in the same type as input (NumPy array or PyTorch tensor).

Examples

>>> import numpy as np
>>> from dl_utils import normalize
>>> img = np.array([[[0, 128, 255]]], dtype=np.float32)
>>> normalize(img, mean=128, std=64)
array([[[-2.      ,  0.      ,  1.984375]]])

>>> import torch
>>> from dl_utils import normalize
>>> img_t = torch.tensor([[[0, 128, 255]]], dtype=torch.float32)
>>> normalize(img_t, mean=torch.tensor([0, 128, 255]), std=torch.tensor([1, 64, 255]))
tensor([[[0., 0., 0.]]], dtype=torch.float64)

dl_utils.data.normalize.inv_normalize(data: ndarray | Tensor | list | tuple, mean: ndarray | Tensor | list | tuple | int | float, std: ndarray | Tensor | list | tuple | int | float, dim=-1) → ndarray | Tensor | list | tuple

Inverse normalize the input array (usually image or video).

Parameters:

• data – Input array, can be a NumPy array or a PyTorch tensor, which has been previously normalized.

• mean – Numeric scalar or vector of means used in the original normalization.

• std – Numeric scalar or vector of standard deviations used in the original normalization.

• dim – The channel dimension along which normalization was applied. Default is -1 (last dimension).

Returns:

Denormalized image or video in the same type as input (NumPy array or PyTorch tensor).

dl_utils.data.normalize.convert_image_video_range(data: ndarray | Tensor, pattern: str)

Convert torch/numpy images/videos between flexible dtype and range.

Note

• When converting from uint8, inputs may be integers or floating point values.

• When converting from -1_1 or 0_1, inputs must be floating point.

• When converting from uint8 to -1_1 or 0_1, floating-point inputs preserve their dtype; integer inputs are promoted to float32.

• When converting between -1_1 and 0_1, the output keeps the input’s dtype.

Examples

>>> img = np.array([[[0, 128, 255]]], dtype=np.uint8)
>>> convert_image_video_range(img, "uint8->0_1")
array([[[0.       , 0.5019608, 1.       ]]], dtype=float32)
>>> img = np.array([[[0.0, 0.5, 1.0]]], dtype=np.float32)
>>> convert_image_video_range(img, "0_1->uint8")
array([[[0, 128, 255]]], dtype=uint8)
>>> img = np.array([[[0.0, 0.5, 1.0]]], dtype=np.float32)
>>> convert_image_video_range(img, "0_1->-1_1")
array([[[-1.,  0.,  1.]]], dtype=float32)

Save & Load

dl_utils.data.save_and_load.save_text(text: str, file)

Save text content to a UTF-8 text file.

Parent directories are created automatically before writing.

Parameters:

• text – Text content to write.

• file – Destination file path.

dl_utils.data.save_and_load.load_text(file) → str

Load the full content of a text file.

Parameters:

file – Source file path.

Returns:

The text content read from file.

dl_utils.data.save_and_load.save_bytes(data: bytes, file)

Save raw bytes to a binary file.

Parent directories are created automatically before writing.

Parameters:

• data – Bytes content to write.

• file – Destination file path.

dl_utils.data.save_and_load.load_bytes(file) → bytes

Load the full content of a binary file.

Parameters:

file – Source file path.

Returns:

The bytes read from file.

dl_utils.data.save_and_load.save_pickle(obj, file)

Serialize an object to a pickle file.

Warning

Pickle is not safe for untrusted data. Only load pickle files from trusted sources.

Parameters:

• obj – Python object to serialize.

• file – Destination pickle file path.

dl_utils.data.save_and_load.load_pickle(file)

Load an object from a pickle file.

Warning

Pickle can execute arbitrary code while loading. Only load pickle files from trusted sources.

Parameters:

file – Source pickle file path.

Returns:

The Python object deserialized from file.

dl_utils.data.save_and_load.save_json(data, file, save_pretty=False, **kwargs)

Save an object as JSON.

bytes values are converted to UTF-8 strings by the default custom JSON encoder. The JSON string is fully serialized before writing, so serialization errors do not leave a partially written file.

Parameters:

• data – JSON-serializable object to save.

• file – Destination file path.

• save_pretty – If True, write human-readable JSON with indentation and ensure_ascii=False.

• **kwargs – Extra keyword arguments forwarded to json.dumps().

dl_utils.data.save_and_load.load_json(file)

Load a JSON file.

Parameters:

file – Source JSON file path.

Returns:

The Python object decoded from the JSON file.

dl_utils.data.save_and_load.save_jsonl(data, file, **kwargs)

Save an iterable of objects as JSONL or a JSON array.

The function writes one item at a time, so data can be any iterable, including a generator. For normal file suffixes such as .jsonl, each item is serialized as one JSON object per line. If file ends with .json, the same stream of items is written as a valid JSON array instead.

bytes values are converted to UTF-8 strings by the default custom JSON encoder.

Parameters:

• data – Iterable of JSON-serializable objects to save.

• file – Destination file path. A .json suffix switches the output format from JSONL to a JSON array.

• **kwargs – Extra keyword arguments forwarded to json.dumps().

Raises:

ValueError – If JSONL output would contain a newline inside one line, for example when passing pretty-print options such as indent=2.

dl_utils.data.save_and_load.iter_jsonl(file, max_samples: int | None = None)

Create an iterable view over a JSONL file.

Empty lines are skipped. The returned object supports iteration, len(), and non-negative integer indexing. Length and indexing are backed by a lazy offset index built on first use.

Parameters:

• file – Source JSONL file path.

• max_samples – Optional maximum number of non-empty JSONL records to expose.

Returns:

An iterable object yielding one decoded JSON object per non-empty line.

dl_utils.data.save_and_load.load_jsonl(file, max_samples: int | None = None)

Load a JSONL file into a list.

Parameters:

• file – Source JSONL file path.

• max_samples – Optional maximum number of non-empty JSONL records to load.

Returns:

A list of decoded JSON objects.

dl_utils.data.save_and_load.concurrent_file_loader(file_paths: Iterable[str | Path], loader: Callable[[...], Any] | None = None, load_kwargs: Dict[str, Any] | None = None, concurrency_limit: int | None = None, chunk_size: int | None = None, **kwargs) → Iterable[Any]

Load many files concurrently.

This is a small IO-bound primitive used by higher-level save/load helpers. loader receives one file path and returns the loaded object. Results are yielded in the same order as file_paths when using joblib’s default ordered generator mode.

Parameters:

• file_paths – File paths to read.

• loader – Function used to load one file path. If None, load_bytes() is used.

• load_kwargs – Extra keyword arguments passed to loader.

• concurrency_limit – Alias for joblib n_jobs.

• chunk_size – Alias for joblib batch_size.

• **kwargs – Extra keyword arguments passed to joblib.Parallel.

Returns:

An iterable of loaded file contents.

dl_utils.data.save_and_load.iter_files(files_or_dir: str | Path | Iterable[str | Path], *, pattern: str | None = None, sort: bool = True, n_jobs: int | None = None, flatten: bool = False, loader: Callable[[...], Any] | None = None, load_kwargs: Dict[str, Any] | None = None, **parallel_kwargs) → Iterator[Any]

Iterate over objects loaded from a directory or file list.

loader is any callable that accepts a file path and returns the loaded object, such as load_json(), load_text(), load_pickle(), or a custom function. load_bytes() is used by default.

Parameters:

• files_or_dir – A directory, a single file, or an iterable of files and/or directories. If all inputs are explicit files, pattern can be omitted.

• pattern – Optional glob pattern used when expanding directories. Required if files_or_dir is a directory or contains directories. Common examples are "*.json" for direct JSON children, "*.jsonl" for direct JSONL children, "**/*.json" for recursive JSON matching with pathlib.Path.glob(), and "part-*.json" for prefixed shard files.

• sort – Whether directory expansion should be sorted. Explicit file lists keep caller-provided order.

• n_jobs – Number of parallel reader jobs. Defaults to all CPUs.

• flatten – If True and a loaded object is a list, yield each list item; otherwise yield one object per file.

• loader – Callable used to load one file path. Defaults to load_bytes().

• load_kwargs – Extra keyword arguments passed to loader.

• **parallel_kwargs – Extra keyword arguments passed to joblib.Parallel.

Yields:

Loaded objects, or items inside loaded lists when flatten=True.

Examples

Skip files that fail to load by wrapping the loader with try / except and filtering out None values:

def safe_load_json(path):
    try:
        return load_json(path)
    except Exception as exc:
        print(f"Failed to load {path}: {exc}")
        return None

items = (
    item
    for item in iter_files(folder, pattern="**/*.json", loader=safe_load_json)
    if item is not None
)

dl_utils.data.save_and_load.load_files(files_or_dir: str | Path | Iterable[str | Path], *, pattern: str | None = None, sort: bool = True, n_jobs: int | None = None, flatten: bool = False, loader: Callable[[...], Any] | None = None, load_kwargs: Dict[str, Any] | None = None, **parallel_kwargs) → List[Any]

Load many files into memory as a list.

This is the eager counterpart of iter_files().

Parameters:

• files_or_dir – A directory, a single file, or an iterable of files and/or directories. If all inputs are explicit files, pattern can be omitted.

• pattern – Optional glob pattern used when expanding directories. Required if files_or_dir is a directory or contains directories. Common examples are "*.json" for direct JSON children, "*.jsonl" for direct JSONL children, "**/*.json" for recursive JSON matching with pathlib.Path.glob(), and "part-*.json" for prefixed shard files.

• sort – Whether directory expansion should be sorted. Explicit file lists keep caller-provided order.

• n_jobs – Number of parallel reader jobs. Defaults to all CPUs.

• flatten – If True and a loaded object is a list, append each list item; otherwise append one object per file.

• loader – Callable used to load one file path. Defaults to load_bytes().

• load_kwargs – Extra keyword arguments passed to loader.

• **parallel_kwargs – Extra keyword arguments passed to joblib.Parallel.

Returns:

A list of loaded objects, or flattened list items when flatten=True.

LMDB

class dl_utils.data.lmdb.JsonLmdb(env: Environment, autogrow: bool)

Bases: Lmdb

https://pypi.org/project/lmdbm/

Sampling

dl_utils.data.sample.sample_evenly(input_data: List[Any] | ndarray | Tensor | Sequence[Any], n: int) → ndarray | List[Any] | Tensor

Evenly sample N elements from input_data. Supports list, numpy array, or torch tensor. The input_data can be empty, and n can be less than or equal to 0, in which case it will return empty data.

Parameters:

• input_data – List, numpy array, or torch tensor to sample from.

• n – Number of elements to sample.

Returns:

Sampled data in the same type as input_data.

dl_utils.data.sample.sample_randomly(input_data: List[Any] | ndarray | Tensor | Sequence[Any], n: int, ordered: bool = False, seed: int = None, put_back: bool = False) → ndarray | List[Any] | Tensor

Randomly sample N elements from input_data. Supports list, numpy array, or torch tensor.

Parameters:

• input_data – List, numpy array, or torch tensor to sample from.

• n – Number of elements to sample.

• ordered – Whether to return sampled elements in the original order.

• seed – Random seed for reproducibility.

• put_back – If True, sample with replacement.

Returns:

Sampled data in the same type as input_data.

dl_utils.data.sample.sample_contiguous(input_data: List[Any] | ndarray | Tensor | Sequence[Any], n: int, seed: int | None = None)

Randomly sample a contiguous slice of length n from input_data.

Parameters:

• input_data – List, numpy array, or torch tensor to sample from.

• n – Length of the contiguous sequence to sample.

• seed – Random seed for reproducibility.

Returns:

A contiguous slice of length n from input_data, in the same type as input_data.

Download

dl_utils.data.download.download(url: str, filepath: str | None = None, expected_sha256: str | None = None, cache_dir: str = '~/.cache/dl_utils')

Download file from URL to the given path in a multi-process safe way.

Internal logic: - Return early if file already downloaded - Verify using SHA256 if provided - File lock to avoid concurrent downloads - Temporary file + atomic rename to avoid partial files

Parameters:

• url – Download URL.

• filepath – Path to save the file. If not provided, the file is saved under cache_dir with filename os.path.basename(url).

• expected_sha256 – Expected SHA256 checksum.

• cache_dir – Cache directory if filepath is not specified. By default, it is ~/.cache/dl_utils.

Returns:

Path to the downloaded file.

Distributed

Basic

dl_utils.distributed.basic.get_global_rank() → int

Get the global rank, the global index of the GPU.

dl_utils.distributed.basic.get_local_rank() → int

Get the local rank, the local index of the GPU.

dl_utils.distributed.basic.get_world_size() → int

Get (global) world size, the total amount of GPUs.

dl_utils.distributed.basic.get_device() → device

Get current rank device.

dl_utils.distributed.basic.get_master_addr() → str | None

dl_utils.distributed.basic.get_master_port() → int | None

dl_utils.distributed.basic.barrier_if_distributed(*args, **kwargs)

Synchronizes all processes if under distributed context.

dl_utils.distributed.basic.rank0() → bool

Global rank 0

dl_utils.distributed.basic.local_rank0() → bool

Local rank 0 (of each node)

dl_utils.distributed.basic.rank0_print(*args, **kwargs)

Print only on rank 0

dl_utils.distributed.basic.rank0_log(fn, *args, **kwargs)

Log only on rank 0.

dl_utils.distributed.basic.rank0_wrapper(fn)

Wrap any function to only run on rank 0.

dl_utils.distributed.basic.dist_info(print_fn: ~typing.Callable[[str], ~typing.Any] = <built-in function print>, prefix: str = '')

Print torch distributed information for debugging.

Device

dl_utils.distributed.device.recursive_to(obj: Any, *args, **kwargs) → Any

Recursively move all torch.Tensor in obj to the given device/dtype following the same behavior as torch.Tensor.to(). Supports: Tensor, list, tuple, dict, set. Leaves other objects intact.

Parameters:

• obj – The object to move.

• *args – Arguments to pass to torch.Tensor.to().

• **kwargs – Keyword arguments to pass to torch.Tensor.to().

Returns:

The object with all torch.Tensor moved to the given device.

Note

If no device is specified, the current (gpu) device will be used. If no gpu is available, the cpu will be used.

Note

This operation is in-place. You should copy if you want to keep the original object.

Gather

dl_utils.distributed.gather.gather_objects(list_object: List[Any]) → List[Any]

gather a list of something from multiple GPU.

Breakpoint

dl_utils.distributed.breakpoint.dist_breakpoint(rank: int = 0)

Breakpoint for distributed training. Enter the breakpoint only if the current rank is rank, and block all other processes using distributed barrier.

File System

Dir

dl_utils.fs.dir.make_parent_dirs(path: str | PathLike)

List Files

dl_utils.fs.list_files.list_files(path: str, depth: int | None = None) → List[str]

List all files in a folder recursively.

Parameters:

• path – Root path to start the search.

• depth – Maximum depth to search. If None, there is no depth limit. If 0 or less, stop searching deeper.

Returns: A List of file paths found under the given path.

dl_utils.fs.list_files.list_files_multithread(directory, n_jobs=16, depth: int | None = None)

List all files in a directory recursively using multiple threads. Useful for list files on NFS.

Parameters:

• directory – The directory to search.

• n_jobs – Number of parallel jobs (threads) to use.

• depth – Maximum recursion depth. If None, no depth limit.

Returns: List of all file paths found under the directory.

Inference

LLM Utils

dl_utils.inference.llm_utils.extract_json(llm_output: str) → Any

Extract and parse JSON from an LLM output string. The output may contain text before or after the JSON, including markdown code fences. Returns the parsed Python object.

Raises:

ValueError – if no valid JSON can be found or parsed.

dl_utils.inference.llm_utils.format_prompt(prompt: str, variables: Mapping[str, Any] | None = None, *, style: str | None = 'format', **kwargs: Any) → str

Fill placeholders in a prompt string.

Parameters:

• prompt – Prompt text to format.

• variables – Placeholder values.

• style – Placeholder style. format/brace/python uses Python str.format syntax like {name} and escapes literal braces as {{text}}. bracket/square uses [name] and escapes literal brackets as [[text]]. none disables formatting.

• **kwargs – Extra placeholder values. Values here override variables.

Returns:

Formatted prompt text.

dl_utils.inference.llm_utils.load_prompt(prompt: str | PathLike | None = None, *, prompt_dir: str | PathLike | None = None, prompt_name: str | None = None, version: str | None = None, extensions: Sequence[str] = ('.txt', '.md'), encoding: str = 'utf-8', input_type: str = 'auto') → str

Load prompt text from a file path, a prompt directory, or return text directly.

Parameters:

• prompt – Prompt file path or already-loaded prompt text. In auto mode, existing paths are loaded from disk, strings that look like prompt paths raise FileNotFoundError if missing, and other strings are treated as text. Use input_type="text" to force a path-like string to be treated as literal prompt text.

• prompt_dir – Directory that stores prompt files.

• prompt_name – Prompt name under prompt_dir. If it has no suffix, extensions are tried in order. When version is set, common versioned names are tried, such as name_v1.txt, name.v1.txt and name/v1.txt.

• version – Optional prompt version saved as a separate file.

• extensions – Candidate prompt file extensions. Defaults to .txt and .md.

• encoding – File encoding used when reading prompt files.

• input_type – One of auto, path or text.

Returns:

Loaded prompt text.

dl_utils.inference.llm_utils.render_prompt(prompt: str | PathLike | None = None, variables: Mapping[str, Any] | None = None, *, prompt_dir: str | PathLike | None = None, prompt_name: str | None = None, version: str | None = None, extensions: Sequence[str] = ('.txt', '.md'), encoding: str = 'utf-8', input_type: str = 'auto', style: str | None = 'format', **kwargs: Any) → str

Load a prompt from path/text and fill placeholders in one call.

Parameters:

• prompt – Prompt file path or already-loaded prompt text.

• variables – Placeholder values.

• prompt_dir – Directory that stores prompt files.

• prompt_name – Prompt name under prompt_dir.

• version – Optional prompt version saved as a separate file.

• extensions – Candidate prompt file extensions.

• encoding – File encoding used when reading prompt files.

• input_type – One of auto, path or text.

• style – Placeholder style passed to format_prompt().

• **kwargs – Extra placeholder values. Values here override variables.

Returns:

Rendered prompt text.

QPS Control

class dl_utils.inference.qps_control.QPSLimiter(max_qps: int = 1, max_concurrent: int | None = None, init_tokens: int = 0)

Bases: object

QPSLimiter is a high performance QPS limiter designed for asyncio workloads. It schedules a background refill task in __init__, so instantiate it inside an async function (i.e., when an event loop is running).

Example

import asyncio
import time
from dl_utils.inference.qps_control import QPSLimiter

def call_api(x: int) -> dict:
    # Any blocking/sync function is OK; it runs in a thread pool.
    time.sleep(0.05)
    return {"x": x}

async def main():
    limiter = QPSLimiter(max_qps=20, max_concurrent=50, init_tokens=20)
    try:
        tasks = [limiter.run(call_api, i) for i in range(100)]
        results = await asyncio.gather(*tasks)
        print("real_qps=", limiter.real_qps())
    finally:
        await limiter.shutdown()

asyncio.run(main())

real_qps() → float

Return the real QPS based on the query count and time.

async shutdown()

Shutdown token refill

async acquire()

async run(func: Callable, *args, **kwargs) → dict[str, Any]

Ray Inference Utils

class dl_utils.inference.ray_inference_utils.RayActorScheduler(actors: List[ray.actor.ActorHandle], actor_fn: Callable[[ray.actor.ActorHandle, ...], ray.ObjectRef], queue_max_size: int = 2)

Bases: object

A high-performance load balanced scheduler for Ray actors. Balance the load of multiple Ray actors using token bucket.

Parameters:

• actors – A list of Ray actor handles to schedule work onto.

• actor_fn – A callable that takes an actor handle plus user-provided *args/**kwargs and submits a Ray task (typically an actor method call). It must return a ray.ObjectRef representing the submitted task.

• queue_max_size – Token bucket size per actor. Effectively caps the number of in-flight tasks allowed per actor to this value.

• scheduler_max_concurrency – Concurrency for internal Ray scheduler actor.

submit(*args, **kwargs)

Ollama Utils

dl_utils.inference.ollama_utils.is_ollama_running(host='127.0.0.1', port=11434) → bool

dl_utils.inference.ollama_utils.start_ollama_background(*args)

Start ollama serve detached from the parent process.

dl_utils.inference.ollama_utils.ensure_ollama_running(*args, max_retries: int = 16, wait_seconds: int = 1)

Ensure that Ollama server is running, start it if not.

Parameters:

• *args – Additional arguments to pass to ollama serve.

• max_retries – Maximum number of retries to check if Ollama is running.

• wait_seconds – Seconds to wait between retries.

Returns:

class dl_utils.inference.ollama_utils.OllamaModel(model_name: str)

Bases: object

chat(*args, **kwargs) → ollama.ChatResponse

async async_chat(*args, **kwargs)

Misc

Chunk

dl_utils.chunk_utils.chunk(data: List[Any], n_chunks: int | None = None, chunk_size: int | None = None, idx: int | None = None) → List[List[Any]] | List[Any]

Split a list into multiple smaller chunks.

You can specify either:

• n_chunks: the number of chunks to create, which will divide the list into approximately equal parts, or

• chunk_size: the number of elements each chunk should have.

Parameters:

• data – The list to chunk.

• n_chunks – The number of chunks to split the list into.

• chunk_size – The size of each chunk.

• idx – The index of the chunk to return. If None, return all chunks.

Returns:

A list of chunks, or a single chunk if idx is specified.

dl_utils.chunk_utils.sort_chunk(data: ~typing.List[~typing.Any], key: ~typing.Callable[[~typing.Any], ~typing.Any] = <function <lambda>>, reverse: bool = False, *args, **kwargs) → List[List[Any]] | List[Any]

Sort a list and then split it into chunks. Useful in distributed processing where input data may be unordered and must be sorted before splitting into chunks.

Parameters:

• data – The list to sort and chunk.

• key – A function to extract the sort key for each element.

• reverse – Whether to sort in descending order.

• args – Arguments passed to the chunk function (e.g., n_chunks or chunk_size).

• kwargs – Arguments passed to the chunk function (e.g., n_chunks or chunk_size).

Returns:

A list of chunks of the sorted data.

Decorators

dl_utils.decorators.log_on_entry(fn: Callable | None = None, print_fn: Callable[[str], Any] | None = None) → Callable

Functions with this decorator will log the function name at entry. When using multiple decorators, this must be applied innermost to properly capture the name.

dl_utils.decorators.barrier_on_entry(fn: Callable) → Callable

Functions with this decorator will start executing when all ranks are ready to enter.

dl_utils.decorators.flex_kwargs(fn: Callable) → Callable

Decorator that allows a function to ignore extra keyword arguments. Only the parameters that the function actually accepts will be passed in.

Cache

class dl_utils.global_cache_utils.GlobalCache(max_size=None)

Bases: MutableMapping

Create a singleton key-value cache.

Parameters:

max_size – Maximum size of the cache. Defaults to None.

Examples

Set and retrieve a cached value:

>>> cache = GlobalCache()
>>> cache["a"] = 1
>>> cache2 = GlobalCache()
>>> cache2["a"]
1

Retrieve a value with a function and arguments, automatically caching the result:

>>> cache = GlobalCache()
>>> cache.get("b", fn=lambda x: x+1, x=1)
2
>>> cache["b"]
2

get(key, fn: Callable | None = None, **kwargs) → Any

Retrieve a value from the cache, or compute and store it if missing. This method will call the provided function fn with kwargs to compute the value

Parameters:

• key – The key to look up in the cache.

• fn (Optional[Callable]) – A function to compute the value if the key is missing.

• **kwargs – Keyword arguments to pass to fn.

Returns:

The cached or newly computed value.

Raises:

KeyError – If the key is missing and no fn is provided.

save(file_path: str)

Save only cache content to file.

Parameters:

file_path – Path to the cache file using pickle.

classmethod load(file_path: str, max_size: int | None = None, update: bool = True) → GlobalCache

Load cache from file and optionally set new max_size, return a GlobalCache instance.

Parameters:

• file_path – Path to the cache file.

• max_size – Maximum size of the cache.

• update – If True (default), merge the loaded cache into the existing cache. If False, replace the current cache entirely with the loaded data.

to_bytes() → bytes

Serialize cache content to bytes using pickle.

Returns:

Serialized cache content in bytes.

classmethod from_bytes(data: bytes, max_size: int | None = None, update: bool = True) → GlobalCache

Deserialize bytes and optionally set new max_size, return a GlobalCache instance.

Parameters:

• data – Serialized cache content.

• max_size – Maximum size of the cache.

• update – If True (default), merge the loaded cache into the existing cache. If False, replace the current cache entirely with the loaded data.

Env

dl_utils.env_utils.get_env(name: str, allow_multiline: bool = False, cwd: str | Path | None = None) → str | list[str]

Get a secret or config value from environment or a dotfile in the working directory.

Lookup order (later overwrites earlier if both exist):

1. A file named {name} under the current working directory (or cwd if provided).

2. A file named .{name} under the current working directory (or cwd if provided).

3. Environment variable name.

By default, the returned string contains no line breaks. Set allow_multiline=True to preserve multiple lines (line endings normalized to \n).

Parameters:

• name – Environment variable name, and the filename stem for the .{name} file.

• allow_multiline – Whether to preserve multiple lines in the returned string.

• cwd – Optional directory to resolve the .{name} file from. Defaults to Path.cwd().

Returns:

The resolved secret/config value as a string.

Raises:

KeyError – If the value is not found in either the environment or the .{name} file.

ID

dl_utils.id_utils.list_ids(roots: str | list[str], include: str | Pattern | list[Pattern] = 'auto', exclude: str | Pattern | list[Pattern] | None = None, matching: str | Pattern = 'auto', simple: bool = True, return_filepath: bool = False, return_dict: bool = False) → list[str] | tuple[list[str], list[str]] | dict[str, str]

Extract sample IDs from file or folder names under given root paths.

Parameters:

• roots – Root path(s) to search for files or folders.

• include – Patterns to include files/folders in directory. If “auto”, include all files if there are more files than folders, else include all folders.

• exclude – Patterns to exclude files/folders in directory.

• matching – Patterns to extract IDs from file/folder names. If “auto”, use common patterns: for files use the full name without extension, for folders use the full folder name.

• simple – If True, use os.listdir and treat entries with a suffix as files (heuristic). This avoids filesystem stat calls and is faster on large NFS directories.

• return_filepath – Whether to return the full file/folder paths along with IDs.

• return_dict – Whether to return a dictionary mapping IDs to file/folder paths. Only effective when return_filepath is True.

Returns: A list of extracted IDs, or a tuple of (IDs, file/folder paths) if return_filepath is True.

dl_utils.id_utils.generate_id(seed: str | None = None) → str

Generate a sample ID.

Two modes: - seed is None: generate a random ID (UUIDv4). - seed is not None: generate a deterministic ID from the input string (hash).

Parameters:

seed – Seed string used for deterministic ID generation. If None, return a random ID.

Returns:

A 32-character lowercase hex string.

dl_utils.id_utils.index_by_id(entries: Iterable[Any] | None, key: Callable[[Any], Any] | None = None, ignore_duplicates: bool = False) → dict[str, Any]

Build a dict index from entries.

Parameters:

• entries – Entries to index.

• key – Function used to extract a key from each entry. If None, try common id-like keys from dict entries.

• ignore_duplicates – Whether to ignore duplicated keys. If False, raise a ValueError when duplicated keys are found. If True, skip duplicate checks and let later entries overwrite earlier ones.

Returns:

A dictionary mapping extracted keys to original entries.

class dl_utils.id_utils.IdSampleManager(ids: Iterable[Any], records: dict[str, dict[str, Any]] | None = None)

Bases: object

Manage a unique ID pool, named selections, and non-overlapping sampling.

This class only operates on unique IDs so it can be reused with arbitrary data structures outside the manager. Sampling results are recorded as named selections, and consumed selections are automatically excluded from later sampling operations.

Parameters:

• ids – Unique IDs representing the full pool.

• records – Existing named records to restore. Each record should contain ids, and can optionally include kind, consumed, and metadata.

property all_ids: list[str]

Return all managed IDs.

property records: dict[str, dict[str, Any]]

Return a copy of all named records.

has_record(name: str) → bool

Return whether a record with the given name exists.

list_records(kind: str | None = None, consumed: bool | None = None) → list[str]

List record names filtered by kind and/or consumed flag.

get_record(name: str) → dict[str, Any]

Return a copy of a named record.

get_record_ids(name: str) → list[str]

Return the IDs stored in a named record.

find_records_by_id(id_: Any, *, kind: str | None = None, consumed: bool | None = None) → list[str]

Find all record names containing the given ID.

add_record(name: str, ids: Iterable[Any], *, kind: str = 'manual', consumed: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Add a named ID record.

Parameters:

• name – Record name.

• ids – IDs to register. IDs must belong to the managed pool.

• kind – Record type, e.g. sample, manual, or derived.

• consumed – Whether these IDs should be excluded from future sampling.

• metadata – Optional metadata stored with the record.

• overwrite – Whether to overwrite an existing record with the same name.

Returns:

The normalized IDs stored in the record.

remove_record(name: str) → dict[str, Any]

Remove and return a named record.

get_selected_ids(consumed_only: bool = False) → list[str]

Return the union of selected IDs across records.

get_remaining_ids() → list[str]

Return IDs that have not been consumed by sampling records.

sample_record(n: int, name: str, *, seed: int | None = None, include_ids: Iterable[Any] | None = None, exclude_ids: Iterable[Any] | None = None, allow_partial: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Sample IDs without repeating previously consumed IDs.

Parameters:

• n – Number of IDs to sample.

• name – Record name used to store this sampling result.

• seed – Optional random seed for reproducibility.

• include_ids – Optional subset of pool IDs to sample from.

• exclude_ids – Optional extra IDs to exclude.

• allow_partial – If True, return all remaining candidates when n exceeds the available count.

• metadata – Optional metadata stored with the sampling record.

• overwrite – Whether to overwrite an existing record with the same name.

Returns:

The sampled IDs.

union(*names: str, name: str | None = None, consumed: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Return the union of multiple records; store it only when name is provided.

intersection(*names: str, name: str | None = None, consumed: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Return the intersection of multiple records; store it only when name is provided.

difference(base_name: str, *other_names: str, name: str | None = None, consumed: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Return IDs in base_name but not in the other records; store only when name is provided.

complement(*names: str, name: str | None = None, consumed: bool = False, metadata: dict[str, Any] | None = None, overwrite: bool = False) → list[str]

Return IDs in the full pool excluding the given records; store only when name is provided.

to_dict() → dict[str, Any]

Serialize the manager into a JSON-friendly dictionary.

classmethod from_dict(data: dict[str, Any]) → IdSampleManager

Restore a manager from a dictionary returned by to_dict().

save_json(path: str | Path, indent: int = 2)

Save manager state as JSON.

classmethod load_json(path: str | Path) → IdSampleManager

Load manager state from a JSON file.

Mask

dl_utils.mask_utils.binarize_mask(mask, threshold: int = 127)

Convert a mask to a 2D boolean NumPy array.

Parameters:

• mask – Input mask. Can be a NumPy array or a PIL image. Supported shapes: - (H, W) boolean or numeric - (H, W, 1) - (H, W, 3) or (H, W, 4) (treated as any-channel > threshold)

• threshold – Threshold used to binarize numeric masks. Pixels greater than this value are treated as True.

Returns:

A 2D boolean mask (dtype np.bool_) with shape (H, W).

Raises:

ValueError – If the input mask has an unsupported shape.

dl_utils.mask_utils.load_mask(path: str | Path, threshold: int = 127, invert: bool = False, as_bool: bool = True)

Load a mask image from disk.

The image is read via Pillow and converted to grayscale (mode "L") for consistent behavior.

Parameters:

• path – Path to the mask image.

• threshold – Threshold applied on the grayscale image (0-255). Pixels greater than this value are treated as True.

• invert – Whether to invert the mask after binarization.

• as_bool – If True, return a boolean mask. If False, return a uint8 mask in {0, 255}.

Returns:

a 2D boolean mask (dtype np.bool_). If as_bool=False: a 2D np.uint8 mask with values in {0, 255}.

Return type:

If as_bool=True (default)

dl_utils.mask_utils.save_mask(mask, path: str | Path, threshold: int = 127, invert: bool = False)

Save a mask to an image file.

Parameters:

• mask – Input mask. It will be converted to a 2D boolean mask using binarize_mask().

• path – Output image path.

• threshold – Threshold used when binarizing numeric masks.

• invert – Whether to invert the mask before saving.

Returns:

None. This function saves the mask to disk.

Notes

The output image is an 8-bit grayscale image (Pillow mode "L") with values in {0, 255}.

dl_utils.mask_utils.union_masks(*masks, threshold: int = 127)

Union (logical OR) of multiple masks.

Parameters:

• *masks – One or more input masks. Each mask is binarized with binarize_mask().

• threshold – Threshold used when binarizing numeric masks.

Returns:

A 2D boolean mask representing the union.

Raises:

ValueError – If no masks are provided, or if mask shapes do not match.

dl_utils.mask_utils.intersect_masks(*masks, threshold: int = 127)

Intersection (logical AND) of multiple masks.

Parameters:

• *masks – One or more input masks. Each mask is binarized with binarize_mask().

• threshold – Threshold used when binarizing numeric masks.

Returns:

A 2D boolean mask representing the intersection.

Raises:

ValueError – If no masks are provided, or if mask shapes do not match.

dl_utils.mask_utils.subtract_mask(a, b, threshold: int = 127)

Set difference: keep pixels in a but not in b (a AND (NOT b)).

Parameters:

• a – Input mask.

• b – Input mask.

• threshold – Threshold used when binarizing numeric masks.

Returns:

A 2D boolean mask.

Raises:

ValueError – If mask shapes do not match.

dl_utils.mask_utils.invert_mask(mask, threshold: int = 127)

Invert a mask (logical NOT).

Parameters:

• mask – Input mask. It is binarized with binarize_mask().

• threshold – Threshold used when binarizing numeric masks.

Returns:

A 2D boolean mask.

dl_utils.mask_utils.mask_iou(a, b, threshold: int = 127) → float

Compute IoU (Intersection over Union) between two masks.

Parameters:

• a – Input mask.

• b – Input mask.

• threshold – Threshold used when binarizing numeric masks.

Returns:

IoU value in [0, 1]. By convention, if both masks are empty (union == 0), IoU is 1.0.

Raises:

ValueError – If mask shapes do not match.

dl_utils.mask_utils.unbinarize_mask(mask, true_value: int = 255, false_value: int = 0, dtype=<class 'numpy.uint8'>)

Convert a 2D boolean mask to a numeric mask (typically uint8 0/255).

This is the reverse of binarize_mask() when the input mask is already boolean.

Parameters:

• mask – A 2D boolean mask (dtype must be np.bool_).

• true_value – Value to write where mask is True.

• false_value – Value to write where mask is False.

• dtype – Output dtype.

Returns:

A 2D array with values in {false_value, true_value}.

Raises:

ValueError – If the input is not a 2D boolean mask.

Inspect Data

dl_utils.inspect_data_utils.inspect_data(data: Any, max_items: int = 10, max_dict_items: int | None = None, max_list_items: int | None = None, max_depth: int = 2, name: str | None = None) → None

Recursively inspects and prints the structure of a data object using a rich Tree.

Parameters:

• data – The data object to inspect.

• max_items – Maximum number of items to display for each container (dict, list, tuple).

• max_dict_items – Maximum number of items to display for dictionaries.

• max_list_items – Maximum number of items to display for lists and tuples.

• max_depth – Maximum recursion depth.

• name – Optional name to label the variable.

Memory

class dl_utils.memory_utils.MemoryStats(tag: str, gpu_before_mb: float | None, gpu_after_mb: float | None, cpu_before_mb: float | None, cpu_after_mb: float | None, elapsed_s: float)

Bases: object

tag: str

gpu_before_mb: float | None

gpu_after_mb: float | None

cpu_before_mb: float | None

cpu_after_mb: float | None

elapsed_s: float

as_dict()

property gpu_delta_mb: float | None

property cpu_delta_mb: float | None

format_message(include_cpu: bool = True) → str

Return a formatted string summary, safe even if some values are None.

dl_utils.memory_utils.get_gpu_memory_state(device: device | None = None, sync: bool = True) → float | None

Get current GPU memory state.

Parameters:

• device – Device to measure memory usage. If None, the current device will be used.

• sync – Whether to synchronize the device before measuring memory usage.

Returns:

Current GPU memory usage in bytes. If no GPU is available, return None.

dl_utils.memory_utils.get_cpu_memory_state() → float | None

Get current CPU memory state in bytes. Usually this is not important, so we allow it to fail silently.

dl_utils.memory_utils.gc_and_empty_cache()

Simply a combination of gc.collect() and torch.cuda.empty_cache().

dl_utils.memory_utils.measure_memory(tag: str = '', device: device | None = None, sync: bool = True, verbose: bool = True, report_cpu: bool = True)

Context manager to measure GPU/CPU memory usage during a code block.

Parameters:

• tag – Tag to identify the memory usage.

• device – Device to measure memory usage. If None, the current device will be used.

• sync – Whether to synchronize the device before measuring memory usage.

• verbose – Whether to print the memory usage.

• report_cpu – Whether to report CPU memory usage.

Examples

>>> with measure_memory("forward") as m:
...     out = model(x)
>>> print(m.gpu_delta_mb)

Prefetcher

class dl_utils.prefetcher.CudaPreFetcher(data_loader: DataLoader)

Bases: object

preload()

static cuda(x: list | tuple | dict | Tensor)

Timer

dl_utils.timer.get_timestamp() → str

Return the current time in a format suitable for filenames.

Examples

>>> get_timestamp()
'2022-10-11T13-41-45W'

dl_utils.timer.get_readable_timestamp()

Return the current time in a readable format.

Examples

>>> get_readable_timestamp()
'2022-10-11 13:41:45'

dl_utils.timer.get_current_time_in_ms(precision: int) → float

Return the current time in milliseconds.

class dl_utils.timer.Timer(precision=3)

Bases: object

A simple timer for getting duration in milliseconds.

get_current_time_in_ms() → float

Return the current time in milliseconds.

reset() → float

get_duration() → float

get_duration_and_reset() → float

class dl_utils.timer.ExecutionTimer(history_size: int | None = None, precision: int = 2, start_prompt: str = None, end_prompt: str = None, log: bool = False, name: str | None = None)

Bases: Timer

A timer for tracking the execution time of sequential stages.

Parameters:

• history_size – Number of history to store for each stage. If None, all history will be stored.

• precision – Precision of the duration in milliseconds.

• start_prompt – Format string for the start prompt. If None, the default prompt will be used.

• end_prompt – Format string for the end prompt. If None, the default prompt will be used.

• log – Whether to print the start and end prompts at the beginning and end of each stage.

• name – Recognizable name of the timer to be used in outputs.

Examples

>>> timer = ExecutionTimer(log=True)
>>> timer.start_stage("stage_1")
(2025-10-21 16:17:32) => Starting stage: stage_1...
>>> time.sleep(0.1) # do something
>>> timer.start_stage("stage_2")
(2025-10-21 16:17:32) => Finished stage: stage_1 | Took 103.030 ms.
(2025-10-21 16:17:32) => Starting stage: stage_2...
>>> time.sleep(0.2) # do something
>>> timer.end_stage("stage_2")
(2025-10-21 16:17:32) => Finished stage: stage_2 | Took 203.100 ms.
>>> timer.print_table()
Stage      Total (ms)    Count    Min (ms)    Max (ms)    Avg (ms)
-------  ------------  -------  ----------  ----------  ----------
stage_1        103.03        1      103.03      103.03      103.03
stage_2        203.10        1      203.10      203.10      203.10
-------------------------
Total Time: 306.13 ms

start_stage(name: str | None = None)

Log the start of a stage. If there is a previous stage, the previous stage will be ended automatically by calling end_stage().

Parameters:

name – The name of the stage to start. If None, one must be specified in the next call to end_stage.

Note

Only for sequential use.

end_stage(name: str | None = None)

Log the end of a stage.

Parameters:

name – The name of the stage to end. If None, the name of the last call to start_stage must be specified.

Note

Only for sequential use.

summary()

Return a summary of all stages as a dictionary.

Each key is a stage name, and its value is a dictionary with:

• “total”: total time (ms)

• “avg”: average time (ms)

• “min”: minimum time (ms)

• “max”: maximum time (ms)

• “last”: last recorded time (ms)

• “count”: number of runs

print_table()

Print the summary of the timer in a table format.