Reading sharded array with multiprocessing produces wrong result

[ziw-liu]

Zarr version

v3.0.10

Numcodecs version

0.16.1

Python Version

3.13

Operating System

macOS

Installation

pip in venv

Description

Sharding doesn't work with multiprocessing.

Steps to reproduce

With ziw-liu#1, run pytest -v -k multiprocessing.

Additional output

No response

[d-v-b]

i think sharding invites race conditions if multiple processes attempt to write to the same shard file at the same time. so we need a mechanism to ensure that sharded writes in a multiprocessing context are per-shard, and not per-chunk. I wonder if this is related to #3217

[ziw-liu]

i think sharding invites race conditions if multiple processes attempt to write to the same shard file at the same time.

The test linked in OP is only reading in parallel but still fails.

[ziw-liu]

Here's a script version of the test in case reproducing the dev environment is too much trouble:

import multiprocessing as mp

import numpy as np
import zarr


def _index_array(arr, index):
    return arr[index]


if __name__ == "__main__":
    data = np.arange(100)
    arr = zarr.create_array(
        store=zarr.storage.MemoryStore(), data=data, shards=(20,), chunks=(10,)
    )
    ctx = mp.get_context("forkserver")
    with ctx.Pool() as pool:
        results = pool.starmap(_index_array, [(arr, slice(len(data)))])
    assert all(np.array_equal(r, data) for r in results)

On Linux it doesn't work when shards=(20,), but passes with shards=(10,) (i.e. equal to chunks).

[d-v-b]

thanks for the reproducer. hopefully this is a simple bug to fix, since it's on the reading side

[ziw-liu]

It seems like it's only reading the last chunk. results[0] with chunk size fixed to 10:

# shard size 20
[ 0  0  0  0  0  0  0  0  0  0 10 11 12 13 14 15 16 17 18 19  0  0  0  0
  0  0  0  0  0  0 30 31 32 33 34 35 36 37 38 39  0  0  0  0  0  0  0  0
  0  0 50 51 52 53 54 55 56 57 58 59  0  0  0  0  0  0  0  0  0  0 70 71
 72 73 74 75 76 77 78 79  0  0  0  0  0  0  0  0  0  0 90 91 92 93 94 95
 96 97 98 99]

# shard size 30
[ 0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0 20 21 22 23
 24 25 26 27 28 29  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
  0  0 50 51 52 53 54 55 56 57 58 59  0  0  0  0  0  0  0  0  0  0  0  0
  0  0  0  0  0  0  0  0 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
 96 97 98 99]

# shard size 90
[ 0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
  0  0  0  0  0  0  0  0 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
 96 97 98 99]

And all the workers read the same results independent of the number of workers.