Updating codec pipeline batch size for GPU codecs to enable parallelism

akshaysubr · akshaysubr · commit a8c0db34ef48 · 2025-02-27T16:24:21.000Z
diff --git a/src/zarr/core/config.py b/src/zarr/core/config.py
@@ -68,6 +68,10 @@ def enable_gpu(self) -> ConfigSet:
                 "buffer": "zarr.core.buffer.gpu.Buffer",
                 "ndbuffer": "zarr.core.buffer.gpu.NDBuffer",
                 "codecs": {"zstd": "zarr.codecs.gpu.NvcompZstdCodec"},
+                "codec_pipeline": {
+                    "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",
+                    "batch_size": 65536,
+                },
             }
         )
 

Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,10 @@ def enable_gpu(self) -> ConfigSet:`
`68`	`68`	`"buffer": "zarr.core.buffer.gpu.Buffer",`
`69`	`69`	`"ndbuffer": "zarr.core.buffer.gpu.NDBuffer",`
`70`	`70`	`"codecs": {"zstd": "zarr.codecs.gpu.NvcompZstdCodec"},`
	`71`	`+ "codec_pipeline": {`
	`72`	`+ "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",`
	`73`	`+ "batch_size": 65536,`
	`74`	`+ },`
`71`	`75`	`}`
`72`	`76`	`)`
`73`	`77`