diff --git a/src/access_moppy/executors/nc_anchor.py b/src/access_moppy/executors/nc_anchor.py new file mode 100644 index 00000000..41053303 --- /dev/null +++ b/src/access_moppy/executors/nc_anchor.py @@ -0,0 +1,58 @@ +"""Keep one netCDF file open in every Dask worker for the worker's lifetime. + +netCDF-C keeps every open file in a process-global table that has no lock, and +frees the whole table when its open-file count reaches zero +(``libdispatch/nclistmgr.c``: ``del_from_NCList`` calls ``free_NCList`` when +``numfiles == 0``; unchanged on netcdf-c ``main``). A Dask worker touches +netCDF-C from two threads: the compute thread opens files under xarray's +backend lock, while ``CachingFileManager.__del__`` closes them from whichever +thread drops the last reference -- without that lock. If such a close takes the +worker's count to zero while the compute thread is inside ``nc_open``, the +table is freed under it and the new ncid no longer resolves: +``RuntimeError: NetCDF: Not a valid ID``, raised from ``_get_format`` or +``_get_vars``. + +Batch workers reach zero often: each source partition opens a handful of +files, and they are closed as soon as the partition's tasks are released. Two +threads opening and closing with nothing else held crash within seconds in an +isolated netCDF4 test; holding one extra file removes the failures entirely. + +This is a workaround for an unsynchronised close, not a fix for it. It closes +the path where the count legitimately reaches zero; a lost update on the +non-atomic counter could still, in principle, drive it there. +""" + +from __future__ import annotations + +import os + +from distributed import WorkerPlugin + +ANCHOR_FILENAME = "moppy_nc_anchor.nc" + + +class NetCDFAnchorPlugin(WorkerPlugin): + """Hold a tiny read-only netCDF file open so netCDF-C never frees its table. + + Registered with ``client.register_plugin``; the scheduler re-applies it to + any worker the nanny restarts, so the anchor survives worker restarts. + """ + + name = "moppy-nc-anchor" + + def __init__(self) -> None: + self._anchor = None + + def setup(self, worker) -> None: + import netCDF4 + + path = os.path.join(worker.local_directory, ANCHOR_FILENAME) + if not os.path.exists(path): + with netCDF4.Dataset(path, "w", format="NETCDF4") as ds: + ds.createDimension("x", 1) + self._anchor = netCDF4.Dataset(path, "r") + + def teardown(self, worker) -> None: + if self._anchor is not None and self._anchor.isopen(): + self._anchor.close() + self._anchor = None diff --git a/src/access_moppy/templates/cmor_python_script.j2 b/src/access_moppy/templates/cmor_python_script.j2 index 4898f3b2..74af1295 100644 --- a/src/access_moppy/templates/cmor_python_script.j2 +++ b/src/access_moppy/templates/cmor_python_script.j2 @@ -18,6 +18,7 @@ from access_moppy.executors.dask_config import ( recommend_dask_config, record_measured_peak, ) +from access_moppy.executors.nc_anchor import NetCDFAnchorPlugin from access_moppy.task_status import TaskStatusFile @@ -265,6 +266,12 @@ def main(): silence_logs=logging.WARNING, **dask_config, ) + # Before any file is opened on a worker: holding one netCDF file open + # per worker stops netCDF-C freeing its file table under a concurrent + # nc_open ("NetCDF: Not a valid ID"). See executors/nc_anchor.py. + anchored = client.register_plugin(NetCDFAnchorPlugin()) + n_anchored = sum(r.get('status') == 'OK' for r in anchored.values()) + print(f'netCDF anchor held on {n_anchored}/{len(anchored)} workers') print(f'Dask dashboard: {client.dashboard_link}') # Run each source partition through an independent CMORiser so its diff --git a/tests/unit/test_batch_cmoriser.py b/tests/unit/test_batch_cmoriser.py index 1095a45f..3448da82 100644 --- a/tests/unit/test_batch_cmoriser.py +++ b/tests/unit/test_batch_cmoriser.py @@ -2996,6 +2996,11 @@ def _run_generated_main( mock_discover.side_effect = discover_error mock_cmoriser = Mock() mock_recommend_dask_config = Mock(return_value={}) + # register_plugin returns {worker address: {"status": ...}}; the worker + # script counts the "OK" entries to report how many workers hold the + # netCDF anchor. + mock_client = Mock() + mock_client.register_plugin.return_value = {"tcp://worker": {"status": "OK"}} with ( patch("access_moppy.file_discovery.discover_files", mock_discover), @@ -3004,7 +3009,7 @@ def _run_generated_main( "access_moppy.executors.dask_config.recommend_dask_config", mock_recommend_dask_config, ), - patch("dask.distributed.Client", Mock()), + patch("dask.distributed.Client", Mock(return_value=mock_client)), ): namespace = {"__name__": "generated"} exec(compile(script.read_text(), str(script), "exec"), namespace) diff --git a/tests/unit/test_nc_anchor.py b/tests/unit/test_nc_anchor.py new file mode 100644 index 00000000..93d1c8e2 --- /dev/null +++ b/tests/unit/test_nc_anchor.py @@ -0,0 +1,74 @@ +"""Tests for the netCDF anchor worker plugin.""" + +import os +from types import SimpleNamespace + +import pytest + +from access_moppy.executors.nc_anchor import ANCHOR_FILENAME, NetCDFAnchorPlugin + + +@pytest.mark.unit +def test_setup_holds_file_open_and_teardown_closes_it(tmp_path): + worker = SimpleNamespace(local_directory=str(tmp_path)) + plugin = NetCDFAnchorPlugin() + + plugin.setup(worker) + + assert os.path.exists(tmp_path / ANCHOR_FILENAME) + assert plugin._anchor.isopen() + + plugin.teardown(worker) + + assert plugin._anchor is None + + +@pytest.mark.unit +def test_setup_reuses_an_existing_anchor_file(tmp_path): + # A restarted worker can come back in the same local directory. + worker = SimpleNamespace(local_directory=str(tmp_path)) + first = NetCDFAnchorPlugin() + first.setup(worker) + first.teardown(worker) + + second = NetCDFAnchorPlugin() + second.setup(worker) + + assert second._anchor.isopen() + second.teardown(worker) + + +@pytest.mark.unit +def test_teardown_without_setup_is_a_no_op(tmp_path): + NetCDFAnchorPlugin().teardown(SimpleNamespace(local_directory=str(tmp_path))) + + +def _anchor_is_open(dask_worker): + plugin = dask_worker.plugins.get(NetCDFAnchorPlugin.name) + return plugin is not None and plugin._anchor is not None and plugin._anchor.isopen() + + +@pytest.mark.integration +def test_every_worker_holds_an_anchor_and_keeps_it_after_restart(): + from distributed import Client, LocalCluster + + # Same shape as the batch worker: separate processes, one thread each. + with ( + LocalCluster( + n_workers=2, + threads_per_worker=1, + processes=True, + dashboard_address=None, + ) as cluster, + Client(cluster) as client, + ): + client.register_plugin(NetCDFAnchorPlugin()) + assert all(client.run(_anchor_is_open).values()) + + # The nanny starts fresh worker processes; the scheduler must + # re-apply the plugin to them. + client.restart() + client.wait_for_workers(2) + result = client.run(_anchor_is_open) + assert len(result) == 2 + assert all(result.values()) diff --git a/tests/unit/test_templates.py b/tests/unit/test_templates.py index de0acd8e..437c787e 100644 --- a/tests/unit/test_templates.py +++ b/tests/unit/test_templates.py @@ -68,6 +68,14 @@ def test_partitioned_worker_template_compiles(self): assert "supports only direct mappings" not in rendered assert "calculation.get('partitionable') is not True" in rendered assert "only monthly and" in rendered + + # The anchor must be in place before any worker opens a file, i.e. + # right after the client starts and before the first CMORiser. + client_at = rendered.index("client = dd.Client(") + anchor_at = rendered.index("client.register_plugin(NetCDFAnchorPlugin())") + assert "netCDF anchor held on" in rendered + cmoriser_at = rendered.index("cmoriser = ACCESS_ESM_CMORiser(") + assert client_at < anchor_at < cmoriser_at assert "daily variables support source partitioning" in rendered assert "self-contained " in rendered assert "mappings need no input files" in rendered