Zarr/OMX skim precedence changes with Sharrow shared-memory loading

[zneedell]

Hello! I ran into a source-precedence behavior that seems worth clarifying, and may be a small compatibility bug.

Describe the bug

When both an OMX file and a Zarr skim store are configured for taz_skims, it is possible for both files to be valid but contain different skim values. In that situation, which on-disk version of the skims "wins" appears to depend on the internal Sharrow/shared-memory loading path, rather than on an explicit user-facing rule.

In particular:

• With Sharrow disabled, the legacy skim path uses OMX; the configured Zarr store is not used.
• With Sharrow enabled but shared-memory skims disabled, the loader uses Zarr values when the store is valid. (this doesn't allow for multiprocessing)
• With Sharrow and shared-memory skims enabled, the normal aligned path creates the shared memory backing but fills it from OMX rather than Zarr. A Zarr realignment path can behave differently and populate shared memory from Zarr.

Also in the current implementation, configured OMX paths are also resolved before Zarr selection, so a valid Zarr store does not make a configured OMX file optional (this is what led me to look into this to begin with).

To Reproduce

Here's verification against the prototype_mtc example

import tempfile
from pathlib import Path

import numpy as np
import openmatrix
import sharrow as sh
import zarr

import activitysim.abm
from activitysim.core import skim_dataset, workflow


EXAMPLE = Path(
    "activitysim_examples/prototype_mtc" # Path to prototype_mtc configs and data directory
)


def load_dist(cache_dir, output_dir, *, store_skims_in_shm, multiprocess):
    state = workflow.State.make_default(
        configs_dir=(EXAMPLE / "configs_mp", EXAMPLE / "configs"),
        data_dir=EXAMPLE / "data",
        cache_dir=cache_dir,
        output_dir=output_dir,
        settings={
            "sharrow": True,
            "multiprocess": multiprocess,
            "store_skims_in_shm": store_skims_in_shm,
        },
    )
    skims = skim_dataset.load_skim_dataset_to_shared_memory(state)
    try:
        return float(skims["DIST"].isel(otaz=0, dtaz=1).values)
    finally:
        if skims.shm.is_shared_memory:
            skims.shm.release_shared_memory()
        else:
            skims.close()


with tempfile.TemporaryDirectory() as temp_dir:
    temp_dir = Path(temp_dir)
    cache_dir = temp_dir / "cache"
    output_dir = temp_dir / "output"
    cache_dir.mkdir()
    output_dir.mkdir()
    zarr_path = cache_dir / "skims.zarr"

    # 1. Create valid Zarr skims from the actual 25-zone OMX data.
    with openmatrix.open_file(EXAMPLE / "data" / "skims.omx", mode="r") as omx:
        print("OMX DIST[0, 1]:", float(omx["DIST"][0, 1]))
        sh.dataset.from_omx_3d(
            [omx],
            index_names=("otaz", "dtaz", "time_period"),
            time_periods=["EA", "AM", "MD", "PM", "EV"],
        ).to_zarr_with_attr(zarr_path)

    # 2. Change only Zarr's DIST values.
    zarr.open_group(str(zarr_path), mode="a")["DIST"][:] = np.float32(22)
    print("Zarr DIST[0, 1]: 22.0")

    # 3. Process-local Zarr loading: Zarr values are returned.
    print(
        "process-local Zarr:",
        load_dist(cache_dir, output_dir, store_skims_in_shm=False, multiprocess=False),
    )

    # 4. Aligned shared-memory loading: code reloads from OMX.
    print(
        "aligned shared memory:",
        load_dist(cache_dir, output_dir, store_skims_in_shm=True, multiprocess=True),
    )

gives

OMX DIST[0, 1]: 0.24
Zarr DIST[0, 1]: 22.0
process-local Zarr: 22.0
aligned shared memory: 0.23999999463558197

Expected behavior
I'd expect that a valid enabled Zarr store should supply the effective skim values regardless of whether ActivitySim keeps the dataset process-local or materializes it into shared memory. OMX would remain the fallback when Zarr isn't used. Alternately, the zarr store could be more clearly labeled as an internal convenience cache only, not to be modified.

Additional context
This seems to be related but distinct from #1035: both involve the no-realignment reload_from_omx_3d shared-memory path, but this concerns which configured source supplies values.

I have a narrow patch in a fork that populates shared memory from Zarr rather than overwriting it from OMX. I would be happy to turn that into a PR if this matches the intended direction. It could also be possible to update sharrow to build a version of reload_from_omx_3d targeted at zarr skims to avoid the Dask dependency, and I'd be happy to take a look at that or any other simpler precedence rules if you all would prefer that.