Knonik DataLoader

This guide explains how to use the Knonik multidataloader from the wheel-based Knonik package. The dataloader is primarily a Python API.

1. What The Dataloader Does

The Knonik multidataloader streams training batches from KHLP datasets produced by knonik pack.

It supports:

• ▸Single sharded datasets with ShardedDataLoader.
• ▸Multi-dataset training with MultiShardedDataLoader.
• ▸Local filesystems and S3.
• ▸NumPy output for JAX / TensorFlow / framework-neutral pipelines.
• ▸Torch output for PyTorch pipelines.
• ▸Ordered threaded/process augment execution.
• ▸Optional ordered batch prefetch.
• ▸Optional zero-copy batching on the fast paths.
• ▸Global sampling for better training diversity.

2. Install

Create or activate the training environment:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

Install the Knonik wheel:

python -m pip install /path/to/knonik-0.1.0-cp311-cp311-manylinux*.whl

Verify the install:

python -c "import knonik_multidataloader; print('dataloader import OK')"

Install PyTorch (Use the PyTorch build that matches your CUDA/runtime environment.)

python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

3. Log In

Log in once with an account entitled for multidataloader:

knonik login --product multidataloader

On headless Linux machines without an OS keyring:

knonik login --product multidataloader --allow-file-key-store

Check status:

knonik status --product multidataloader

Log out:

knonik logout --product multidataloader

4. Single-Dataset Python Example

Use ShardedDataLoader when training from one KHLP manifest.

from knonik_multidataloader.api.sharded_dataloader import ShardedDataLoader

loader = ShardedDataLoader(
    loader_cfg={
        "decode_mode": "random_access",
        "hlp_manifest": "/data/knonik_thread_velcro_hlp/hlp_manifest.json",
        "chunk_len": 10,
        "batch_size": 64,
        "framework": "numpy",
        "fs_type": "local",
        "fs_params": {"root": "/"},
        "shuffle_mode": "global_triplet",
        "num_decoders": 8,
        "augment_execution": "threaded",
        "precollate_batches": 2
    },
    stream_cfg={},          # empty means auto-discover streams from manifest
    dataset_name="thread_velcro",
)

try:
    for batch in loader:
        # train(batch)
        print(batch.keys())
        break
finally:
    loader.shutdown()

5. Multi-Dataset Python Example

Use MultiShardedDataLoader when mixing several KHLP datasets.

from knonik_multidataloader.api.multi_sharded_dataloader import MultiShardedDataLoader

datasets = [
    (
        "thread_velcro",
        {
            "loader": {
                "hlp_manifest": "/data/thread_velcro_hlp/hlp_manifest.json",
                "fs_type": "local",
                "fs_params": {"root": "/"}
            },
            "streams": {}
        },
    ),
    (
        "bimanual_transfer",
        {
            "loader": {
                "hlp_manifest": "/data/bimanual_transfer_hlp/hlp_manifest.json",
                "fs_type": "local",
                "fs_params": {"root": "/"}
            },
            "streams": {}
        },
    ),
]

loader = MultiShardedDataLoader(
    datasets=datasets,
    weights=[0.7, 0.3],
    loader_cfg={
        "decode_mode": "random_access",
        "chunk_len": 10,
        "batch_size": 64,
        "framework": "numpy",
        "shuffle_mode": "global_triplet",
        "num_decoders": 8,
        "augment_execution": "threaded",
        "precollate_batches": 2
    },
)

try:
    batch = next(iter(loader))
    print(batch["_sample_dataset_names"][:8])
finally:
    loader.shutdown()

Multi-Dataset Epoch & Mixing Keys

Alongside weights, MultiShardedDataLoader accepts two loader_cfg keys that control how the mix behaves over an epoch:

6. Batch Shape

Each batch is a dictionary.

Typical stream keys:

batch["rgb_cam_high"]    # shape: (B, T, H, W, 3)
batch["state"]           # shape: (B, T, state_dim)
batch["action"]          # shape: (B, T, action_dim)

How a batch is built

B is batch_size. The loader picks B sample anchors (the order/diversity of anchors is what shuffle_mode / the planner control), then gathers the per-stream frames for each anchor and stacks the B samples along axis 0. The time dimension T is set per stream by chunk_len and delta_timestamps:

Plain chunk_len. Every stream gets the same contiguous run of chunk_len frames starting at the anchor, so T = chunk_len for all streams:

# chunk_len=16  -> every stream has T=16
batch["rgb_cam_high"]  # (B, 16, H, W, 3)
batch["state"]         # (B, 16, state_dim)
batch["action"]        # (B, 16, action_dim)

With delta_timestamps. Each stream you list samples frames at the given offsets in seconds around the anchor, so its T is the number of offsets. Unlisted streams fall back to the contiguous chunk_len window — so streams can have different T in the same batch:

# chunk_len=16, fps=50
loader_cfg["delta_timestamps"] = {
    "state":  [-0.10, -0.05, 0.0, 0.05, 0.10],     # 5 offsets  -> T=5
    "action": [i / 50 for i in range(16)],         # 16 offsets -> T=16
}
# rgb_cam_high is not listed -> contiguous chunk_len window -> T=16

batch["state"]         # (B, 5,  state_dim)
batch["action"]        # (B, 16, action_dim)
batch["rgb_cam_high"]  # (B, 16, H, W, 3)

pad_missing=True clamps offsets that fall outside the episode to its boundaries (instead of dropping the sample); when padding happens the batch also carries a <stream>_valid_mask so you can tell real frames from padded ones. Within one anchor every stream is aligned to the same anchor frame, so element t of one stream and element t of another correspond to the offsets you requested.

Typical metadata keys:

batch["_dataset_name"]          # single-dataset name
batch["_dataset_names"]         # dataset names represented in the batch
batch["_sample_dataset_names"]  # one dataset name per sample
batch["_hlp_ids"]
batch["_episode_ids"]
batch["_shard_idxs"]
batch["_chunk_starts"]
batch["_stream_names"]

With framework="numpy", arrays are numpy.ndarray. With framework="torch", arrays are CPU torch tensors.

7. Recommended Framework Modes

JAX / TensorFlow / NumPy

Recommended:

{
  "framework": "numpy",
  "augment_execution": "threaded",
  "augment_workers": 4,
  "precollate_batches": 2
}

This path is:

decode numpy -> numpy augment -> numpy batch

This avoids numpy -> torch -> numpy conversion churn.

PyTorch

Recommended:

{
  "framework": "torch",
  "augment_execution": "threaded",
  "augment_workers": 4,
  "precollate_batches": 2
}

This path is:

decode numpy -> torch augment -> torch batch

The dataloader does not move batches to GPU. Move to GPU in your training loop.

8. Augment Functions

augment_fn runs per sample before collation.

NumPy example:

def augment_numpy(sample):
    image = sample["streams"]["rgb_cam_high"]
    sample = dict(sample)
    streams = dict(sample["streams"])
    streams["rgb_cam_high"] = image[..., ::-1, :] 
    sample["streams"] = streams
    return sample

loader_cfg = {
    "framework": "numpy",
    "augment_fn": augment_numpy,
    "augment_execution": "threaded",
    "augment_workers": 4
}

Torch example:

def augment_torch(sample):
    image = sample["streams"]["rgb_cam_high"]
    sample = dict(sample)
    streams = dict(sample["streams"])
    streams["rgb_cam_high"] = image.float()
    sample["streams"] = streams
    return sample

loader_cfg = {
    "framework": "torch",
    "augment_fn": augment_torch,
    "augment_execution": "threaded",
    "augment_workers": 4
}

Augment execution modes:

For JAX/OpenPI-style training, prefer threaded unless the transform is very CPU-heavy and safely pickleable.

9. Ordered Batch Prefetch (precollate_batches)

precollate_batches keeps fully-collated, ready-to-consume batches in a background queue so your training step never waits on batch assembly. It is enabled by default with a depth of 2; raise it when model steps are expensive, or set it to 0 to disable.

{
  "precollate_batches": 2,
  "batch_prefetch_timeout_s": 0.1,
  "batch_prefetch_daemon": true
}

Properties:

• ▸Preserves exact batch order.
• ▸Prefetches fully collated batches.
• ▸Works for single and multi sharded loaders.
• ▸Does not change planner or sampling behavior.
• ▸Default is 2; set precollate_batches: 0 to disable.

10. KHLP Storage-Aware Loader Modes

The packer chooses how RGB frames are stored. The loader chooses how to read those bytes. Good throughput comes from matching the two choices.

global_triplet_fast does not change the planner, seed behavior, sample order, or batch metadata relative to global_triplet. The only difference is how the bytes are read, via partial decode.

Training-Fast Loader Recipe

Use this for local or cloud training when loader throughput matters more than minimum storage size. The simple surface is just the shuffle mode:

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/datasets/my_dataset_kimg/hlp_manifest.json",
    "chunk_len": 4,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet_fast",
    "num_decoders": 8,            # raise if the GPU is input-starved
}

Training-Compressed Loader Recipe

Use this when storage size matters but sparse training access is still required. On local files, add the single toggle kdelta_direct_local_decode to read KDLT byte ranges straight from the KHLP file; on S3 leave it off (the default).

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/datasets/my_dataset_kdelta/hlp_manifest.json",
    "chunk_len": 1,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet_fast",
    "delta_timestamps": {
        "rgb_cam_high": [0.0, 0.1, 0.2, 0.3],
        "state": [0.0, 0.1, 0.2, 0.3],
        "action": [0.0, 0.1, 0.2, 0.3],
    },
    "pad_missing": True,
    "kdelta_direct_local_decode": True,   # local files only; omit for S3
}

That is the whole simple surface. Everything else the direct path needs: request bundling, per-worker slot sizing, byte-range merging, file-handle caching, orchestration threads, is auto-defaulted to sane values, so you do not set it. If profiling later shows a specific bottleneck, the individual override knobs live in Fine-grained control.

Delta-Timestamp Sampling With global_triplet_fast

delta_timestamps is the right interface when streams are sampled around an anchor step. For example, at 50 FPS, offsets [0, 5, 10, 15] frames are:

offsets = [0 / 50, 5 / 50, 10 / 50, 15 / 50]
loader_cfg["delta_timestamps"] = {
    "rgb_cam_high": offsets,
    "state": offsets,
    "action": offsets
}

The fast path resolves these offsets to concrete per-stream frame indices before decoding. It does not silently fall back to sequential frames unless your timestamps themselves describe adjacent frames.

11. Zero-Copy Batching

Zero-copy batching (zerocopy) is an opt-in fast-path optimization that can roughly double loader throughput (~2×) in exchange for a larger RAM footprint. It is off by default — every existing run behaves exactly as before unless you turn it on — and it requires shuffle_mode="global_triplet_fast" or shuffle_mode="global_triplet".

When you enable it, declare each RGB stream's frame shape up front with zerocopy_frame_shapes:

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_kimg_hlp/hlp_manifest.json",
    "chunk_len": 4,
    "batch_size": 64,
    "framework": "torch",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet_fast",
    "zerocopy": True,
    "zerocopy_frame_shapes": {"rgb_cam_high": [480, 640, 3]},   # [H, W, C] per RGB stream
}

Constraints:

• ▸augment_fn is not supported with zerocopy=True (it raises). Apply the transform to the batched tensor in your training loop instead.
• ▸pin_memory is ignored under zero-copy (it warns).

12. Useful Loader Config Keys

These are the keys a normal training job sets. Each loader_cfg key is optional except the four required ones; the rest have sensible defaults, and internal / auto-tuned knobs are kept out of this list (see Fine-grained control). The keys fall into a few groups, described before each table below: required (what to load and the sample shape), throughput (how fast to read and decode), sampling (which frames become samples and in what order), diversity & memory (the RAM-vs-shuffle-quality trade), and a single storage-aware toggle for the training_fast / training_compressed profiles.

Required / common — identify the dataset, the sample shape, the output array type, and where the files live. These four (decode_mode, hlp_manifest, chunk_len, batch_size) are the only mandatory keys.

Throughput — control how far ahead the loader reads and how many shards it decodes in parallel. Raise these when the GPU is waiting on input; lower them under CPU/RAM pressure. For most runs num_decoders is the only one you touch.

Sampling & shape — decide which frames become samples and in what order.

Diversity & memory — trade shuffle quality against RAM use.

Run control & output — cap epoch length, pin the output for faster GPU transfer, and preset the loader for training.

Zero-copy — see Zero-Copy Batching for zerocopy and its keys.

Storage-aware reads — a single toggle for datasets packed with the training_fast or training_compressed storage profiles.

Fine-grained Control (Advanced)

You should not need any of these for a normal run. They are auto-tuned or have sane fixed defaults; setting one overrides the automatic value. Reach for them only when profiling points to a specific bottleneck.

Concurrency & decode internals:

The loader pipeline has two work stages with separate parallelism:

• ▸Fetch / IO — threads that issue byte-range reads from the filesystem (or S3) to pull shard bytes. Sized by num_fetchers.
• ▸Decode — worker processes that turn those shard bytes into numpy arrays. Sized by decode_concurrency.

num_decoders is the single high-level knob that sets both at once. decode_concurrency exists to override just the decode-process count, decoupling it from the fetch side, useful when the two stages bottleneck differently:

• ▸Decode-bound on a CPU-rich box → raise decode_concurrency without inflating fetch threads.
• ▸S3/latency-bound with cheap decode, or limited CPU/RAM → keep num_decoders (hence num_fetchers) high but cap decode_concurrency lower.

In short: num_decoders = "overall decode parallelism" (drives fetch + decode); decode_concurrency = "decode-process count only" (overrides the decode half). Everything below auto-scales from these; override one only to decouple it.

Global-planner internals (under global_triplet):

Partial decode (Only with global_triplet_fast):

global_triplet_fast does partial decode: it figures out exactly which frame indices each planned batch needs and decodes only those, skipping the rest of the shard. This is what makes the training_fast and training_compressed profiles fast for sparse random access.

You normally just set shuffle_mode="global_triplet_fast" and leave the rest at their auto defaults. The knobs that tune the partial-decode path — partial_decode_coalesce_batches, partial_decode_window_prefetch — live in Fine-grained control;

in short:

• ▸partial_decode_coalesce_batches — how many upcoming planned batches are merged into one decode "window", so frames from nearby batches that land in the same shard are decoded together (fewer, larger decode jobs).
• ▸partial_decode_window_prefetch — how many such windows to work on ahead.

Precollate queue tuning:

13. Worked Examples

Each example is a loader_cfg you drop into the standard call from section 4:

loader = ShardedDataLoader(loader_cfg=<one of the dicts below>,
                           stream_cfg={}, dataset_name="my_dataset")

Copy one and adapt the paths. All assume you are logged in (knonik login --product multidataloader).

13.1 Minimal local run

Defaults for everything except the dataset and sample shape.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": 16,
    "batch_size": 8,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
}

13.2 EC2 / cloud training from S3

Reading shards from S3 on an EC2 box. Widen the I/O side to hide network latency — more event loops, more files in flight — and keep a decoded-shard RAM cache so epoch 2+ is fast. Needs pip install aiobotocore and valid AWS credentials (standard credential chain).

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "s3://my-bucket/datasets/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": 16,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "s3",
    "fs_params": {
        "bucket": "my-bucket",
        "prefix": "datasets/my_dataset_hlp",   # "" if manifest URIs already include it
        "region": "us-east-1",
        "max_pool_connections": 64,
    },
    # cloud-friendly I/O
    "io_loops": 4,
    "hlps_in_flight": 16,
    "num_decoders": 16,
    "ram_cache_bytes": 16 << 30,   # 16 GiB
}

13.3 Memory-constrained training

Small RAM footprint: no decoded cache, shallow prefetch, fewer workers, fewer files in flight, bounded working set.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": 16,
    "batch_size": 8,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet",
    "planner_locality": "hlp_bounded",
    "ram_cache_bytes": 0,
    "prefetch_shards": 4,
    "num_decoders": 4,
    "hlps_in_flight": 2,
    "low_memory": True,
}

13.4 High-throughput on a big box

Lots of RAM and cores — widen everything.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": 16,
    "batch_size": 32,
    "framework": "torch",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet",
    "planner_locality": "global",
    "num_decoders": 16,
    "hlps_in_flight": 8,
    "prefetch_shards": 16,
    "ram_cache_bytes": 32 << 30,
    "precollate_batches": 3,
}

13.5 Single-frame samples (chunk_len=1)

One frame per sample with a large batch — bound the active working set.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": 1,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet",
    "planner_locality": "low_mem",
    "low_memory": True,
    "ram_cache_bytes": 16 << 30,
}

13.6 LeRobot-style delta_timestamps

Sample state history and action lookahead at chosen offsets (seconds) around the anchor frame. Streams not listed fall back to a contiguous chunk_len window.

FPS = 50
CHUNK_LEN = 16
loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_hlp/hlp_manifest.json",
    "chunk_len": CHUNK_LEN,
    "batch_size": 8,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet",
    "delta_timestamps": {
        "state":  [-0.10, -0.05, 0.0, 0.05, 0.10],
        "action": [i / FPS for i in range(CHUNK_LEN)],
    },
    "pad_missing": True,
}

13.7 training_fast dataset

For a dataset packed with storage_profile="training_fast". Use the fast partial-decode path; each frame is an independent image so random access is cheap.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_kimg_hlp/hlp_manifest.json",
    "chunk_len": 4,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet_fast",
    "num_decoders": 8,
}

13.8 training_compressed dataset

For a dataset packed with storage_profile="training_compressed" on local disk. Add the single kdelta_direct_local_decode toggle — bundle sizing is auto-derived from the dataset.

loader_cfg = {
    "decode_mode": "random_access",
    "hlp_manifest": "/data/my_dataset_kdelta_hlp/hlp_manifest.json",
    "chunk_len": 1,
    "batch_size": 64,
    "framework": "numpy",
    "fs_type": "local",
    "fs_params": {"root": "/"},
    "shuffle_mode": "global_triplet_fast",
    "kdelta_direct_local_decode": True,    # local only; omit for S3
    "delta_timestamps": {
        "rgb_cam_high": [0.0, 0.1, 0.2, 0.3],
        "state": [0.0, 0.1, 0.2, 0.3],
        "action": [0.0, 0.1, 0.2, 0.3],
    },
    "pad_missing": True,
}

13.9 Multi-dataset mixing

Mix two datasets at fixed weights with MultiShardedDataLoader; the smaller one cycles so the ratio holds for the whole epoch.

from knonik_multidataloader.api.multi_sharded_dataloader import MultiShardedDataLoader

datasets = [
    ("dataset_a", {"loader": {"hlp_manifest": "/data/a_hlp/hlp_manifest.json",
                              "fs_type": "local", "fs_params": {"root": "/"}},
                   "streams": {}}),
    ("dataset_b", {"loader": {"hlp_manifest": "/data/b_hlp/hlp_manifest.json",
                              "fs_type": "local", "fs_params": {"root": "/"}},
                   "streams": {}}),
]

loader = MultiShardedDataLoader(
    datasets=datasets,
    weights=[0.7, 0.3],            # 70% A / 30% B, guaranteed per batch
    loader_cfg={
        "decode_mode": "random_access",
        "chunk_len": 16,
        "batch_size": 8,
        "framework": "numpy",
        "shuffle_mode": "global_triplet",
        "planner_locality": "hlp_bounded",
        "num_decoders": 8,
        "ram_cache_bytes": 16 << 30,
    },
)
try:
    for batch in loader:
        # batch["_sample_dataset_names"] -> which dataset each sample came from
        pass
finally:
    loader.shutdown()

14. Tuning By Symptom

Skip the reference; come here when something is wrong.

GPU is waiting on input (low utilization). In order: raise num_decoders (4 → 8 → 16) — fetchers, slots, and orchestration scale with it; raise prefetch_shards (12 → 16 or 24); for S3 raise io_loops (1 → 2 → 4); add or grow ram_cache_bytes.

Memory usage too high / RSS grows without bound. Lower ram_cache_bytes; lower prefetch_shards (12 → 4–8); set low_memory: true; tighten planner_locality (global → hlp_bounded → low_mem → ultra_low_mem); lower num_decoders if the pressure is from the decode pool.

Batches look temporally correlated. Loosen planner_locality (ultra_low_mem → low_mem → hlp_bounded → global); raise planner_recent_batches (4 → 16 → 32).

Too many cross-shard samples (still correct, just costlier). Increase --shard-size at pack time; reduce chunk_len; increase stride.

"Slow start" — first batch is slow. That is warmup. On large datasets it can take seconds to tens of seconds, paid once per loader. If it never arrives, lower num_decoders / num_fetchers.

Throughput swings wildly batch-to-batch. Add or grow ram_cache_bytes; increase prefetch_shards; check storage (variable I/O latency → variable batch latency).

slot capacity exceeded. Per-decode budget too small for your data. Set shard_max_size (bytes) slightly above the largest decoded shard (or, on the KDLT direct path, lower kdelta_direct_bundle_target_frames).

15. Shutdown Correctly

Always call shutdown() when the loader is no longer needed:

loader = ShardedDataLoader(...)
try:
    for batch in loader:
        train_step(batch)
finally:
    loader.shutdown()

This closes decode pools, augment executors, prefetch threads, caches, and metrics files.