Function test_initial_hint

dali/test/python/test_external_source_parallel.py:1182–1215

()

def test_initial_hint():
    sample_size = 32 * 1024 * 1024
    shapes = [(32, 1024, 1024)]
    # make the initial size still smaller than necessary to check if reallocation works
    bytes_per_sample_hint = 4 * 1024 * 1024
    batch_size = 7
    num_workers = 3
    min_samples_in_mini_batch = batch_size // num_workers
    max_samples_in_mini_batch = (batch_size + num_workers - 1) // num_workers
    initial_shm_size = max_samples_in_mini_batch * bytes_per_sample_hint
    expected_min_chunk_size = min_samples_in_mini_batch * sample_size

    pipe = per_iter_shape_pipeline(
        shapes,
        bytes_per_sample_hint=bytes_per_sample_hint,
        batch_size=batch_size,
        py_num_workers=num_workers,
    )
    sizes = pipe.external_source_shm_statistics()["capacities"]
    assert len(sizes) > 0
    for size in sizes:
        assert (
            size == initial_shm_size
        ), f"Expected initial size to be {initial_shm_size}, got {size}."

    for _ in range(5):
        pipe.run()

    sizes = pipe.external_source_shm_statistics()["capacities"]
    assert len(sizes) > 0
    for size in sizes:
        assert (
            size >= expected_min_chunk_size
        ), f"Expected the size to be at least {expected_min_chunk_size}, got {size}."


def test_variable_sample_size():