Interface to set inference options

[gmertes]

Is your feature request related to a problem? Please describe.

We are starting to have some options in the model that are set during inference, like the mapper chunking. Right now, we pass them via environment variable.

Describe the solution you'd like

An interface to pass a dictionary of inference options that the model can read from.

So we can do in inference something like:

inference_options = {
    "num_chunks": 2,
    "do_something": True
}

model = torch.load('checkpoint.ckpt')
model.set_inference_options(inference_options)

model.predict_step(...)

The downside of this is that this only allows for runtime options, not for any options that are needed in initialisation (we don't have any of those right now, but we might in the future?). To do the latter we could do:

from anemoi.models.inference_options import set_inference_options

set_inference_options({...}) 
model = torch.load('checkpoint.ckpt')
...

And in models we can do something like this, so it is accessible from everywhere in the module:

from anemoi.models.inference_options import inference_options

num_chunks = inference_options("num_chunks")

Describe alternatives you've considered

No response

Additional context

No response

Organisation

No response