codec registration silently fails

[d-v-b]

When trying to register a codec with the same name as a numcodecs codec (e.g., to override the behavior of that codec), the new codec cannot be retrieved from the registry.

here's a reproducer. The expected behavior is that, after registering my codec under a key, I should be able to pull it out under the same key. That fails, and the (broken) numcodecs codec is retrieved instead, which kind of defeats the entire purpose of a dynamic registry.

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "zarr@git+https://github.com/zarr-developers/zarr-python.git@main",
# ]
# ///
#
# This script automatically imports the development branch of zarr to check for issues


from numcodecs.zarr3 import Delta
import numpy as np
import zarr
from zarr.core.array_spec import ArraySpec
from zarr.dtype import parse_data_type
from zarr.registry import register_codec, get_codec_class
from dataclasses import replace

class PatchedDelta(Delta):
    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
        if astype := self.codec_config.get("astype"):
            return replace(chunk_spec, dtype=parse_data_type(astype, zarr_format=3))  # type: ignore[call-overload]
        return chunk_spec

key = 'numcodecs.delta'
value = PatchedDelta
register_codec(key, value)
print(get_codec_class(key))
# <class 'numcodecs.zarr3.Delta'>
assert get_codec_class(key) == value

The underlying problem stems from the design of the codec registry: the same codec name can be associated with multiple codecs, and in this case, a codec is retrieved in FILO order? The end result is that a numcodecs codec always wins, and the user-defined codec can never be retrieved:

zarr-python/src/zarr/registry.py

Lines 156 to 164 in 9ae874c

	config_entry = config.get("codecs", {}).get(key)
	if config_entry is None:
	if len(codec_classes) == 1:
	return next(iter(codec_classes.values()))
	warnings.warn(
	f"Codec '{key}' not configured in config. Selecting any implementation.",
	stacklevel=2,
	)
	return list(codec_classes.values())[-1]

I think the logic here was to support e.g. CPU and GPU implementations for the same codec, but our current solution is bad. The registry should be a 1:1 mapping between identifiers and codecs. If we need to support CPU codecs and GPU codecs, then we have 2 registries, not a single broken registry.

[TomAugspurger]

Discussed a bit on the call today.

The underlying problem stems from the design of the codec registry: the same codec name can be associated with multiple codecs, and in this case, a codec is retrieved in FILO order?

Having multiple implementations floating around for a single codec name seems problematic and not very useful. As far as we could tell:

1. the zarr JSON metadata includes a codec name + configuration
2. the zarr-python configuration should include a 1:1 mapping from configuration name to codec implementation (a python class)

so we shouldn't ever need a list of codec classes floating around. Just a way to build an instance of the currently active codec class.

[d-v-b]

i'm looking more into the design of the codec registry and I'm confused on a few points. my ultra-naive idea is that the codec registry would be structured something like this:

{"gzip": zarr.codecs.GZipCodec, ...}

i.e., for each codec identifier, there is 1 codec class.

instead, we have something like this:

{"gzip": {"zarr.codecs.GZipCodec": zarr.codecs.GZipCodec, "mypackage.GZip": mypackage.GZip}}

for each codec identifier there is a mapping from fully-qualified names to codec classes.

This makes the name "gzip" ambiguous. The registry doesn't associate "gzip" with a single codec class, which means extra information is needed to choose a codec class. There are two options: rely on insertion order (bad, and what we are currently doing) or require additional information when requesting a codec beyond the codec name. Given the structure of this registry, a resolvable codec request would have to take the form ("gzip", "zarr.codecs.GZipCodec") -> zarr.codecs.GZipCodec which... I don't understand at all 😄

@normanrz @brokkoli71 can you explain how this design is supposed to work?

[normanrz]

Having multiple codecs with the same identifier is necessary, if you want to have multiple implementations (e.g. for gpu).

[d-v-b]

Having multiple codecs with the same identifier is necessary, if you want to have multiple implementations (e.g. for gpu).

that makes sense but this would require a design where codecs can be addressed by (name, implementation). Instead, the codecs are addressed by (name, class name), and so it's possible to get multiple CPU codecs registered without any mechanism for picking one. So I think if the goal is to support selecting codecs by implementation, the current design doesn't work.

[normanrz]

The selection is based on the runtime config.

[d-v-b]

correct, which fails silently if a user registers a codec but doesn't also alter the config (hence this issue). I'm pretty sure we need a better process for this.