Make coupled-dataset path handling work for cloud and local stores

scripts/data_process/create_coupled_datasets.py

@@ -18,6 +18,7 @@
     compute_coupled_ocean,
     compute_coupled_sea_ice,
 )
+from fs_utils import path_exists
 from get_stats import ClimateDataType, StatsConfig, get_stats
 from merge_stats import MergeStatsConfig, merge_stats
 from writer_utils import OutputWriterConfig
@@ -61,7 +62,7 @@ def _get_stats(
     # Check if stats directory already exists
     run_name = os.path.basename(input_zarr_path).replace(".zarr", "-stats")
     stats_dir = os.path.join(os.path.dirname(input_zarr_path), run_name)
-    if os.path.exists(stats_dir):
+    if path_exists(stats_dir):
         logging.info(
             f"Stats directory {stats_dir} already exists. Skipping stats computation."
         )
@@ -121,7 +122,7 @@ def _merge_stats(
 
     # uncoupled atmosphere training (coupled_sea_ice + uncoupled_atmos)
     output_dir = os.path.join(output_directory, "uncoupled_atmosphere")
-    if not os.path.exists(output_dir):
+    if not path_exists(output_dir):
         logging.info("Merging stats for uncoupled atmosphere training")
         logging.info(f"  Input: {sea_ice_dir}")
         logging.info(f"  Input: {uncoupled_atmos_dir}")
@@ -145,7 +146,7 @@ def _merge_stats_helper(
             return
 
         output_dir = os.path.join(output_directory, output_subdir)
-        if os.path.exists(output_dir):
+        if path_exists(output_dir):
             logging.info(
                 f"Merged {label} stats already exist at {output_dir}, skipping"
             )
@@ -212,7 +213,7 @@ def _combine_ensemble_stats(
         stats_roots = []
         for run_name in run_names:
             category_dir = os.path.join(ensemble_stats_dir, run_name, category)
-            if os.path.exists(category_dir):
+            if path_exists(category_dir):
                 # combine_stats uses string concatenation (root + filename),
                 # so paths must end with "/"
                 stats_roots.append(category_dir + "/")
@@ -483,9 +484,9 @@ def write_datasets_and_stats(
         logging.info("=" * 80)
 
         # Check if outputs already exist (unless in debug mode)
-        sea_ice_exists = not debug and os.path.exists(sea_ice_output_store)
-        ocean_exists = not debug and os.path.exists(ocean_output_store)
-        atmos_exists = not debug and os.path.exists(atmosphere_output_store)
+        sea_ice_exists = not debug and path_exists(sea_ice_output_store)
+        ocean_exists = not debug and path_exists(ocean_output_store)
+        atmos_exists = not debug and path_exists(atmosphere_output_store)
 
         compute_ocean = self.coupled_sea_surface is not None
         compute_atmos = self.coupled_ts is not None

scripts/data_process/fs_utils.py

@@ -0,0 +1,40 @@
+"""Filesystem helpers shared by the data processing scripts.
+
+The filesystem is resolved from the path's URL scheme via fsspec, so callers
+use a single code path for local paths and remote URLs (e.g. ``gs://``).
+"""
+
+import fsspec
+from fsspec.implementations.local import LocalFileSystem
+
+
+def _fs_and_path(path: str) -> tuple[fsspec.AbstractFileSystem, str]:
+    return fsspec.core.url_to_fs(path)
+
+
+def path_exists(path: str) -> bool:
+    """Return whether a file, directory, or object-store prefix exists."""
+    fs, fs_path = _fs_and_path(path)
+    return fs.exists(fs_path)
+
+
+def is_dir(path: str) -> bool:
+    """Return whether the path is a directory (or object-store prefix)."""
+    fs, fs_path = _fs_and_path(path)
+    return fs.isdir(fs_path)
+
+
+def makedirs(path: str) -> None:
+    """Create a directory if it doesn't exist; a no-op on object stores."""
+    fs, fs_path = _fs_and_path(path)
+    fs.makedirs(fs_path, exist_ok=True)
+
+
+def is_local(path: str) -> bool:
+    """Return whether the path resolves to the local filesystem.
+
+    isinstance is used because fsspec dispatches filesystem behavior by
+    class, and local-vs-remote is exactly that distinction.
+    """
+    fs, _ = _fs_and_path(path)
+    return isinstance(fs, LocalFileSystem)

scripts/data_process/get_stats.py

@@ -15,6 +15,7 @@
 import fsspec
 import xarray as xr
 import yaml
+from fs_utils import is_local, makedirs, path_exists
 
 # these are auxiliary variables that exist in dataset for convenience, e.g. to do
 # masking or to more easily compute vertical integrals. But they are not inputs
@@ -105,11 +106,7 @@ class Config:
 
 def _out_dir_exists(out_dir: str) -> bool:
     """Check if the stats output directory already has results."""
-    if out_dir.startswith("gs:"):
-        fs = fsspec.filesystem("gs")
-        return fs.exists(out_dir + "/centering.nc")
-    else:
-        return os.path.exists(os.path.join(out_dir, "centering.nc"))
+    return path_exists(os.path.join(out_dir, "centering.nc"))
 
 
 def get_stats(
@@ -192,15 +189,14 @@ def get_stats(
             f"Standard deviation computed over all variables: {all_var_stddev.values}"
         )
     else:
-        if out_dir.startswith("gs:"):
+        if is_local(out_dir):
+            makedirs(out_dir)
+            local_dir = out_dir
+            remote_dir: Optional[str] = None
+        else:
             temp_dir = tempfile.TemporaryDirectory()
             local_dir = temp_dir.name
-            remote_dir: Optional[str] = out_dir
-        else:
-            if not os.path.isdir(out_dir):
-                os.makedirs(out_dir)
-            local_dir = out_dir
-            remote_dir = None
+            remote_dir = out_dir
 
         centering.to_netcdf(os.path.join(local_dir, "centering.nc"))
         if remote_dir is not None:

scripts/data_process/merge_stats.py

@@ -8,6 +8,7 @@
 import fsspec
 import xarray as xr
 import yaml
+from fs_utils import makedirs
 from get_stats import copy
 
 STATS_NC_FILE_NAMES = [
@@ -66,8 +67,7 @@ def merge_stats(config: MergeStatsConfig):
     """
     stats_dir = config.output_directory
 
-    if not os.path.exists(stats_dir):
-        os.makedirs(stats_dir)
+    makedirs(stats_dir)
 
     for fname in STATS_NC_FILE_NAMES:
         logging.info(f"Combining {fname} stats datasets")

scripts/data_process/test_create_coupled_datasets.py

@@ -0,0 +1,199 @@
+"""End-to-end tests for create_coupled_datasets.write_datasets_and_stats on
+tiny synthetic atmosphere and ocean zarr stores in a temporary directory."""
+
+import numpy as np
+import pytest
+import xarray as xr
+
+pytest.importorskip("xpartition")
+
+import create_coupled_datasets  # noqa: E402
+from coupled_dataset_utils import (  # noqa: E402
+    CoupledSeaIceConfig,
+    CoupledSeaSurfaceConfig,
+    CoupledSurfaceTemperatureConfig,
+)
+from create_coupled_datasets import (  # noqa: E402
+    CoupledDatasetsConfig,
+    CoupledInputDatasetConfig,
+    CoupledStatsConfig,
+    CreateCoupledDatasetsConfig,
+    InputDatasetsConfig,
+    InputStatsConfig,
+)
+from create_window_avg_dataset import WindowAvgDatasetConfig  # noqa: E402
+from merge_stats import STATS_NC_FILE_NAMES  # noqa: E402
+from writer_utils import OutputWriterConfig  # noqa: E402
+
+NLAT = 4
+NLON = 8
+N_ATMOS_TIMES = 40  # 6-hourly steps covering ten days
+FIRST_WINDOW_END = "2000-01-06T00:00:00"
+
+
+def _times(start, periods, freq):
+    return xr.date_range(
+        start, periods=periods, freq=freq, calendar="noleap", use_cftime=True
+    )
+
+
+def _field(times, offset=0.0, seed=0):
+    rng = np.random.default_rng(seed)
+    data = offset + rng.uniform(0.0, 1.0, size=(len(times), NLAT, NLON))
+    return xr.DataArray(
+        data,
+        dims=["time", "lat", "lon"],
+        coords={
+            "time": times,
+            "lat": np.linspace(-80, 80, NLAT),
+            "lon": np.linspace(0, 315, NLON),
+        },
+    )
+
+
+def _write_input_zarrs(input_dir):
+    atmos_times = _times("2000-01-01T06:00:00", N_ATMOS_TIMES, "6h")
+    # 5-daily snapshots at the ends of the two 120h windows
+    ocean_times = _times(FIRST_WINDOW_END, 2, "120h")
+
+    atmos = xr.Dataset(
+        {
+            "surface_temperature": _field(atmos_times, offset=280.0, seed=1),
+            "sea_ice_fraction": _field(atmos_times, seed=2),
+            "ocean_fraction": _field(atmos_times, seed=3),
+            "latent_heat_flux": _field(atmos_times, offset=100.0, seed=4),
+        }
+    )
+    atmos["land_fraction"] = _field(atmos_times, seed=5).isel(time=0, drop=True)
+
+    ocean = xr.Dataset(
+        {
+            "sst": _field(ocean_times, offset=275.0, seed=6),
+            "hfds": _field(ocean_times, offset=10.0, seed=7),
+        }
+    )
+    ocean["sea_surface_fraction"] = _field(ocean_times, seed=8).isel(time=0, drop=True)
+
+    atmos_path = str(input_dir / "atmosphere.zarr")
+    ocean_path = str(input_dir / "ocean.zarr")
+    atmos.to_zarr(atmos_path)
+    ocean.to_zarr(ocean_path)
+    return atmos_path, ocean_path
+
+
+def _write_fake_stats_dir(stats_dir, var_name):
+    """Write the four stats netCDFs an uncoupled input's stats directory holds."""
+    stats_dir.mkdir(parents=True)
+    ds = xr.Dataset({var_name: xr.DataArray(1.0)})
+    ds.attrs["input_samples"] = N_ATMOS_TIMES
+    for fname in STATS_NC_FILE_NAMES:
+        ds.to_netcdf(str(stats_dir / fname))
+
+
+def _make_config(tmp_path):
+    input_dir = tmp_path / "inputs"
+    input_dir.mkdir()
+    atmos_path, ocean_path = _write_input_zarrs(input_dir)
+    _write_fake_stats_dir(input_dir / "atmosphere-stats", "uncoupled_atmos_var")
+    _write_fake_stats_dir(input_dir / "ocean-stats", "uncoupled_ocean_var")
+
+    output_dir = tmp_path / "outputs"
+    output_dir.mkdir()
+
+    return CreateCoupledDatasetsConfig(
+        version="v1",
+        family_name="synthetic",
+        output_directory=str(output_dir),
+        coupled_datasets=CoupledDatasetsConfig(
+            coupled_sea_ice=CoupledSeaIceConfig(
+                window_avg=WindowAvgDatasetConfig(
+                    window_timedelta="120h", first_timestamp=FIRST_WINDOW_END
+                ),
+            ),
+            coupled_ts=CoupledSurfaceTemperatureConfig(
+                how="threshold", ocean_fraction_threshold=0.9
+            ),
+            coupled_sea_surface=CoupledSeaSurfaceConfig(
+                surface_flux_window_avg=WindowAvgDatasetConfig(
+                    window_timedelta="120h",
+                    first_timestamp=FIRST_WINDOW_END,
+                    subset_names=["latent_heat_flux"],
+                ),
+                sst_threshold=275.5,
+            ),
+            output_writer=OutputWriterConfig(n_split=1),
+        ),
+        input_datasets=InputDatasetsConfig(
+            climate_data_type="CM4",
+            stats=InputStatsConfig(
+                atmosphere_dir=str(input_dir / "atmosphere-stats"),
+                ocean_dir=str(input_dir / "ocean-stats"),
+            ),
+            atmosphere=CoupledInputDatasetConfig(
+                zarr_path=atmos_path, time_chunk_size=20
+            ),
+            ocean=CoupledInputDatasetConfig(zarr_path=ocean_path, time_chunk_size=2),
+        ),
+        stats=CoupledStatsConfig(),
+    )
+
+
+class _DummyDaskClient:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def close(self):
+        pass
+
+
+@pytest.fixture(autouse=True)
+def no_distributed_client(monkeypatch):
+    """Stats computation on tiny data doesn't need a distributed cluster."""
+    distributed = pytest.importorskip("distributed")
+    monkeypatch.setattr(distributed, "Client", _DummyDaskClient)
+
+
+def _output_paths(config):
+    return [
+        config.sea_ice_output_store,
+        config.ocean_output_store,
+        config.atmosphere_output_store,
+    ]
+
+
+def test_write_datasets_and_stats_end_to_end_and_resume(tmp_path, monkeypatch):
+    config = _make_config(tmp_path)
+    config.write_coupled_datasets(debug=False, subsample=False)
+
+    for store in _output_paths(config):
+        assert (
+            xr.open_zarr(store).sizes["time"] > 0
+        ), f"expected non-empty zarr at {store}"
+
+    for scenario in ["uncoupled_atmosphere", "coupled_atmosphere", "ocean"]:
+        for fname in STATS_NC_FILE_NAMES:
+            merged = tmp_path / "outputs" / "v1-synthetic-stats" / scenario / fname
+            assert merged.exists(), f"missing merged stats file {merged}"
+
+    # A re-run must resume on the existing outputs rather than recompute:
+    # fail loudly if any compute stage or writer is invoked again.
+    def _fail(*args, **kwargs):
+        raise AssertionError("recomputed a stage that already has outputs")
+
+    for name in [
+        "compute_coupled_sea_ice",
+        "compute_coupled_ocean",
+        "compute_coupled_atmosphere",
+    ]:
+        monkeypatch.setattr(create_coupled_datasets, name, _fail)
+    monkeypatch.setattr(OutputWriterConfig, "write", _fail)
+
+    config.write_coupled_datasets(debug=False, subsample=False)
+
+
+def test_write_datasets_and_stats_debug_writes_nothing(tmp_path):
+    config = _make_config(tmp_path)
+    config.write_coupled_datasets(debug=True, subsample=False)
+
+    output_dir = tmp_path / "outputs"
+    assert list(output_dir.iterdir()) == []

scripts/data_process/test_fs_utils.py

@@ -0,0 +1,30 @@
+import os
+
+import fsspec
+from fs_utils import is_dir, is_local, makedirs, path_exists
+
+
+def test_path_helpers_local(tmp_path):
+    subdir = tmp_path / "some" / "nested" / "dir"
+    assert not path_exists(str(subdir))
+    makedirs(str(subdir))
+    assert path_exists(str(subdir))
+    assert is_dir(str(subdir))
+    makedirs(str(subdir))  # idempotent
+
+    file_path = subdir / "file.txt"
+    file_path.write_text("data")
+    assert path_exists(str(file_path))
+    assert not is_dir(str(file_path))
+    assert is_local(str(file_path))
+
+
+def test_path_helpers_resolve_filesystem_from_url_scheme():
+    path = "memory://bucket/prefix/store"
+    assert not is_local(path)
+    assert not path_exists(path)
+    makedirs(path)
+    with fsspec.open(os.path.join(path, "obj.txt"), "wb") as f:
+        f.write(b"data")
+    assert path_exists(os.path.join(path, "obj.txt"))
+    assert is_dir(path)