From 9a0179a6c6ac929894f7d416461675f36586c3cb Mon Sep 17 00:00:00 2001 From: Jan Brezina Date: Sat, 6 Jun 2026 04:31:52 +0200 Subject: [PATCH] Dask pool --- PLAN.md | 179 +++++++++++++++++++++++++---- mlmc/sampling_pool_dask.py | 193 ++++++++++++++++++++++++++++++++ requirements.txt | 2 + setup.py | 3 + test/test_sampling_pool_dask.py | 185 ++++++++++++++++++++++++++++++ tox.ini | 2 + 6 files changed, 542 insertions(+), 22 deletions(-) create mode 100644 mlmc/sampling_pool_dask.py create mode 100644 test/test_sampling_pool_dask.py diff --git a/PLAN.md b/PLAN.md index a0488ed..4c503e1 100644 --- a/PLAN.md +++ b/PLAN.md @@ -38,32 +38,159 @@ Goal 3: implement specific simulations and quantities for estimation of Sobol in ## Work Plan -### Goal 1: Verify Merge and Prepare PR - -- Inspect current branch and last merge commit. -- Run `tox` to indetify possibly failing tests. -- Run targeted import/compile checks for files changed by the merge. -- Review places marked with TODO: comments during the last merge. - These were conficts that needs a caraful revision and possibly reintroduce - some functionality introduced in matster. In particular the while loops - in certaing PBS opeartions were introduced in the master, while - this branch should have a non-blocking implementation of the similar functionality. - Verify that. -- Record verification results and any remaining PR blockers in this plan. -- No code changes yet. - -1. mlmc/tool/hdf5.py:151: add_level_group() unconditionally creates Levels, but init_header() already creates it. This is now the tox py312 blocker: ValueError: name already exists during test/test_storage.py collection. -2. mlmc/sampling_pool_pbs.py:397: _qstat_pbs_job() still calls .decode() on process.stderr, but PbsCommands already returns decoded strings. Any nonzero qstat status will raise AttributeError instead of handling PBS output. -AGENT: remove wrong decode calls -Resolved: `sampling_pool_pbs.py` now uses decoded `CommandOutput.stderr` directly. - -3. mlmc/sampling_pool_pbs.py:408: unknown_job_ids is only initialized inside the qstat-failure branch, then used after successful qstat too. A successful qstat can hit UnboundLocalError. - -4. mlmc/sampling_pool_pbs.py:411: the master retry while loops are disabled, matching the non-looping branch direction, but the exception path still sleeps for 30 seconds. That should be checked against the intended non-blocking + + +### Goal 2: Dask Sampling Pool For Sensitivity Sampling + +Intent: implement a Dask-backed `SamplingPool` alternative that works with the +existing `Sampler` scheduling loop and can be used from the provided +`sensitivity_sampling.py` script. The Dask client should be supplied by the +caller, e.g. `SamplingPoolDask(client=client, work_dir=..., debug=...)`. + +Design direction: + +- Keep the MLMC `Sampler` as the master scheduler. It already supports + asynchronous, iterative sample-count updates through `schedule_samples()`, + `ask_sampling_pool_for_samples()`, and `process_adding_samples()`. +- Add a new Dask pool implementation, most likely in + `mlmc/sampling_pool_dask.py`, implementing the existing `SamplingPool` + interface: + - `schedule_sample(sample_id, level_sim)` submits one future with + `client.submit(...)` and records future metadata locally. + - `get_finished()` collects only completed futures and returns the same tuple + shape as other pools: + `(successful_samples, failed_samples, n_running, n_ops)`. + - `have_permanent_samples(sample_ids)` initially returns `False`, matching + local pools, unless restart/recovery is explicitly added later. + AGENT: Support ofr restart of the sampling is must for large sample sizes. + So desing a way how to implement that with Dask. + Resolved: `SamplingPoolDask` persists per-level simulation metadata for + workspace simulations and uses deterministic Dask task keys and sample + seeds. On restart, construct the pool with `clean=False`; unfinished stored + sample ids are submitted again and the worker task receives only the sample + id, output directory, and seed, then loads the persisted level metadata. + +- Do not use one large `client.map(...); client.gather(...)` for MLMC runs. + That pattern waits for a fixed batch and does not fit the adaptive algorithm, + where estimates are updated while previous futures are still running and new + samples may be scheduled per level. +- Use Dask communication/result transport for sample results. Workers should + return `(sample_id, result, err_msg, running_time)` through their futures, + and the master should write successful/failed samples to `SampleStorage` + through the existing `Sampler._store_samples()` path. + AGENT: that could work but be carfull, result for SA will idealy be an xarray object (of about 4 dimensions), and could be quite large definitely tenths of MB + Resolved: completed futures are released immediately after the master stores + each result. Large SA result persistence still belongs in the simulation or + sample-storage design; the Dask pool does not keep completed futures alive. + +- Reuse `SamplingPool.calculate_sample()` as the worker function to preserve + deterministic seeding, result-format checks, workspace preparation, and + exception-to-error-message behavior. +- Track future-to-level metadata on the master, e.g. + `{future: level_sim}` or `{future: level_id}`, so completed results can be + partitioned by level and runtime can be accumulated into `n_ops`. +- Prefer Dask `as_completed` or future status polling in `get_finished()`. + `get_finished()` must be non-blocking or bounded by the outer `Sampler` + timeout; it should not wait for all scheduled futures. +- Preserve the existing output-directory behavior: + - if `work_dir` is set, use the `SamplingPool` output directory layout; + - if `level_sim.need_sample_workspace` is true, each Dask worker must execute + with a sample-specific workspace prepared by `calculate_sample()`; + - successful/failed directory moving should happen on the master after the + future result is received, similarly to `OneProcessPool._process_result()`, + unless worker-local scratch paths make that impossible. +- Keep Dask as an optional dependency. Avoid importing `dask.distributed` from + `mlmc/__init__.py` unless packaging dependencies are updated accordingly. + Import it only in the Dask pool module or in the script that constructs the + client. + AGENT: I have updated the environment, so dask and dask.distributed should be available. You just have to add the optional dependency into setup.py. + Resolved: `setup.py` now defines the `dask` optional extra with `dask` and + `distributed`; tox test dependencies include both packages. + +- Update `sensitivity_sampling.py` to construct/use the MLMC `Sampler` with the + Dask pool where the current code uses `client.map(single_sample, ...)`, once + the sensitivity simulation is represented through the MLMC `Simulation` + interface. + You can suggest that, but main point is to test sampling_pool_dask be separate unit test. + +Implementation steps: + +1. Add `SamplingPoolDask` with constructor parameters + `client`, `work_dir=None`, `debug=False`, and possibly an optional + `submit_kwargs` dict if needed by the existing cluster setup. +2. Implement future submission with deterministic seed calculation on the + master and `pure=False`, so repeated sample ids are not memoized by Dask. + AGENT: pure =True could be used, as the sample calculation is deteministic using reproducible seeding. + Resolved: `SamplingPoolDask` submits deterministic-key futures with + `pure=True`. + +3. Implement completed-future collection: + - gather futures whose status is `finished` or `error`; + - call `future.result()` only for those futures; + - convert Dask worker exceptions into failed sample entries if they escape + `calculate_sample()`; + - release/remove collected futures to avoid retaining results in cluster + memory. +4. Factor or reuse local pool result processing so success/failure queues, + runtime accumulation, and sample directory cleanup are consistent across + `OneProcessPool`, `ProcessPool`, and `SamplingPoolDask`. + AGENT: not clear to me what you mean by this + Resolved: `SamplingPoolDask` subclasses `OneProcessPool` and reuses its + `_process_result()` and `get_finished()` queue conversion, avoiding a + separate refactor. + +5. Add focused tests using a local Dask cluster/client and synthetic simulation: + - initial scheduling stores all samples; + - repeated polling returns partial completion without blocking for the full + batch; + - `process_adding_samples()` can schedule more work while previous futures + are still running; + - failed samples are reported through `failed_samples`. + +6. Add a minimal integration path in `sensitivity_sampling.py`: + - keep `Client(scheduler)` creation in the script; + - pass that client into `SamplingPoolDask`; + - avoid `client.map(...); client.gather(...)` for adaptive MLMC sampling. +7. Update documentation/API references only after the public import path is + chosen. + +Verification plan: + +- Targeted compile check: + `python -m py_compile mlmc/sampling_pool.py mlmc/sampling_pool_dask.py mlmc/sampler.py`. +- Targeted tests for the new pool: + `python -m pytest -c test/pytest.ini test/test_sampling_pool_dask.py -vv`. +- Existing sampler regression: + `python -m pytest -c test/pytest.ini test/test_sampler.py test/test_sampling_pools.py -vv`. +- If sensitivity integration is changed, run the smallest available + `sensitivity_sampling.py` local/Dask smoke command documented by its current + config. If the required external `endorse`, `chodby_trans`, or Flow123d + environment is unavailable, record that as skipped verification. + ## AGENT Log +- `2026-06-06`: Implemented Goal 2 Dask sampler backend. Added + `mlmc/sampling_pool_dask.py` with a `SamplingPoolDask` accepting an existing + Dask `Client`, submitting one deterministic-key future per MLMC sample, + polling only completed futures in `get_finished()`, releasing futures after + storage, and resubmitting unfinished workspace sample ids on sampler restart + from sample ids plus persisted per-level metadata. Added Dask package + metadata in `setup.py`, tox test deps, and focused + `test/test_sampling_pool_dask.py` coverage. Verification passed: + `python3 -m py_compile mlmc/sampling_pool.py mlmc/sampling_pool_dask.py + mlmc/sampler.py mlmc/sampling_pool_pbs.py setup.py + test/test_sampling_pool_dask.py`; + `timeout 60 .tox/py312/bin/python -m pytest -c test/pytest.ini + test/test_sampling_pool_dask.py -vv`; and + `.tox/py312/bin/python -m pytest -c test/pytest.ini test/test_sampler.py + test/test_sampling_pools.py -vv`. +- `2026-06-06`: Planned Goal 2 Dask sampler work. The intended design is a + `SamplingPoolDask` backend that accepts an existing Dask `Client`, submits + one future per MLMC sample, polls completed futures in `get_finished()`, and + leaves adaptive scheduling in the existing `Sampler` rather than using a + single blocking `client.map(...); client.gather(...)` batch. - `2026-06-05`: Continued Goal 1 merge verification on branch `MS_endorse`. Current HEAD is `a36d4e2` (`CODEX conditioning.`); the merge under review is `118722a` (`origin/master` into `MS_endorse`). `python3 -m py_compile` @@ -97,6 +224,14 @@ Resolved: `sampling_pool_pbs.py` now uses decoded `CommandOutput.stderr` directl ## AGENT Questions And Remarks +- `2026-06-06`: Goal 2 assumes that the Dask client is owned by the caller and + passed into the pool constructor. `SamplingPoolDask` should not start or stop + the cluster unless a later requirement explicitly asks for that. +- `2026-06-06`: `sensitivity_sampling.py` currently contains a direct Dask + `client.map(single_sample, sample_args)` workflow around project-specific + dependencies outside MLMC. To make it use the MLMC adaptive sampler, the + transport calculation must be wrapped as an MLMC `Simulation` with a + documented `result_format()`. - `2026-06-05`: Legacy/external fixture status is only partially classified. Treat `test/01_cond_field`, `test/02_conc`, and `test/fractures` as repository examples and legacy/integration fixtures, but do not assume they diff --git a/mlmc/sampling_pool_dask.py b/mlmc/sampling_pool_dask.py new file mode 100644 index 0000000..ad6dfa2 --- /dev/null +++ b/mlmc/sampling_pool_dask.py @@ -0,0 +1,193 @@ +import os +import pickle +import traceback + +from mlmc.sampling_pool import OneProcessPool, SamplingPool + +try: + from distributed import fire_and_forget +except ImportError as exc: + raise ImportError( + "SamplingPoolDask requires the optional Dask dependency. " + "Install MLMC with the 'dask' extra or install 'dask' and 'distributed'." + ) from exc + + +class SamplingPoolDask(OneProcessPool): + """ + Dask-backed sampling pool. + + The caller owns the Dask client and passes it to the constructor. Samples are + submitted one by one so the existing MLMC Sampler can adapt target sample + counts while older futures are still running. + """ + + FUTURE_KEY_PREFIX = "mlmc-sample" + LEVEL_SIM_CONFIG = "level_{}_simulation_config" + + def __init__(self, client, work_dir=None, debug=False, clean=True, submit_kwargs=None): + """ + Initialize the pool with an existing Dask client. + + Parameters + ---------- + client + dask.distributed.Client instance. + work_dir : str, optional + Working directory for sample output. + debug : bool, default=False + If True, keeps sample directories. + clean : bool, default=True + If False, preserves an existing output directory on construction. + Use this when restarting unfinished workspace samples. + submit_kwargs : dict, optional + Extra keyword arguments passed to client.submit(). + """ + super().__init__(work_dir=work_dir, debug=debug or not clean) + self._debug = debug + self._client = client + self._submit_kwargs = {} if submit_kwargs is None else dict(submit_kwargs) + self._future_to_task = {} + self._sample_to_future = {} + + def schedule_sample(self, sample_id, level_sim): + """ + Submit one sample to Dask. + + Dask task keys are deterministic in the MLMC sample id. This gives a + restarted master a chance to reconnect to scheduler-known tasks; if that + is not possible, submitting the same sample id recomputes the same result + because seeding is deterministic. + """ + if sample_id in self._sample_to_future: + return + + if self._output_dir is None and level_sim.need_sample_workspace: + self._output_dir = os.getcwd() + + self._save_level_sim(level_sim) + seed = SamplingPool.compute_seed(sample_id) + future = self._client.submit( + SamplingPool.calculate_sample, + sample_id, + level_sim, + self._output_dir, + seed, + key=self._future_key(sample_id), + pure=True, + **self._submit_kwargs + ) + fire_and_forget(future) + self._future_to_task[future] = (sample_id, level_sim) + self._sample_to_future[sample_id] = future + self._n_running += 1 + + def have_permanent_samples(self, sample_ids): + """ + Reconnect or resubmit samples scheduled before a master restart. + + Dask does not provide PBS-like durable result files. Recovery is therefore + based on workspace simulations, persisted per-level simulation metadata, + and deterministic sample seeds. The restarted worker task receives a + sample id and first re-enters the existing sample workspace. + """ + if not sample_ids: + return False + if self._output_dir is None: + return False + + for sample_id in sample_ids: + self._submit_permanent_sample(sample_id) + return True + + def get_finished(self): + """ + Collect only futures that have already completed. + """ + completed_futures = [ + future for future in list(self._future_to_task) + if self._future_done(future) + ] + + for future in completed_futures: + sample_id, level_sim = self._future_to_task.pop(future) + self._sample_to_future.pop(sample_id, None) + result = self._future_result(future, sample_id) + self._process_result(*result, level_sim) + future.release() + + return super().get_finished() + + @classmethod + def _future_key(cls, sample_id): + return "{}-{}".format(cls.FUTURE_KEY_PREFIX, sample_id) + + @staticmethod + def _future_done(future): + if hasattr(future, "done"): + return future.done() + return getattr(future, "status", None) in {"finished", "error"} + + @staticmethod + def _future_result(future, sample_id): + try: + return future.result() + except Exception: + err_msg = traceback.format_exc() + return sample_id, (None, None), err_msg, 0.0 + + @staticmethod + def _level_id_from_sample_id(sample_id): + try: + return int(str(sample_id).split("_", 1)[0][1:]) + except (IndexError, TypeError, ValueError) as exc: + raise ValueError("Cannot determine level id from sample id {!r}".format(sample_id)) from exc + + def _save_level_sim(self, level_sim): + if self._output_dir is None or not level_sim.need_sample_workspace: + return + + file_path = self._level_sim_file(level_sim._level_id) + if os.path.exists(file_path): + return + + with open(file_path, "wb") as level_sim_file: + pickle.dump(level_sim, level_sim_file) + + def _submit_permanent_sample(self, sample_id): + if sample_id in self._sample_to_future: + return + + seed = SamplingPool.compute_seed(sample_id) + level_sim = self._load_level_sim(sample_id) + future = self._client.submit( + _calculate_permanent_sample, + sample_id, + self._output_dir, + seed, + key=self._future_key(sample_id), + pure=True, + **self._submit_kwargs + ) + fire_and_forget(future) + self._future_to_task[future] = (sample_id, level_sim) + self._sample_to_future[sample_id] = future + self._n_running += 1 + + def _load_level_sim(self, sample_id): + level_id = self._level_id_from_sample_id(sample_id) + file_path = self._level_sim_file(level_id) + with open(file_path, "rb") as level_sim_file: + return pickle.load(level_sim_file) + + def _level_sim_file(self, level_id): + return os.path.join(self._output_dir, self.LEVEL_SIM_CONFIG.format(level_id)) + + +def _calculate_permanent_sample(sample_id, output_dir, seed): + level_id = SamplingPoolDask._level_id_from_sample_id(sample_id) + level_sim_file = os.path.join(output_dir, SamplingPoolDask.LEVEL_SIM_CONFIG.format(level_id)) + with open(level_sim_file, "rb") as level_sim_config: + level_sim = pickle.load(level_sim_config) + + return SamplingPool.calculate_sample(sample_id, level_sim, output_dir, seed) diff --git a/requirements.txt b/requirements.txt index d805a68..a1ac211 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,4 @@ pytest matplotlib +dask +distributed diff --git a/setup.py b/setup.py index 4f583fc..96856f4 100644 --- a/setup.py +++ b/setup.py @@ -62,4 +62,7 @@ def read(*names, **kwargs): include_package_data=True, zip_safe=False, install_requires=['numpy', 'scipy', 'scikit-learn', 'h5py>=3.1.0', 'ruamel.yaml', 'attrs', 'gstools', 'memoization'], + extras_require={ + 'dask': ['dask', 'distributed'], + }, ) diff --git a/test/test_sampling_pool_dask.py b/test/test_sampling_pool_dask.py new file mode 100644 index 0000000..9ca47a5 --- /dev/null +++ b/test/test_sampling_pool_dask.py @@ -0,0 +1,185 @@ +import time + +import numpy as np +import pytest + +distributed = pytest.importorskip("distributed") +from distributed import Client, LocalCluster + +from mlmc.level_simulation import LevelSimulation +from mlmc.quantity.quantity_spec import QuantitySpec +from mlmc.sample_storage import Memory +from mlmc.sampler import Sampler +from mlmc.sampling_pool_dask import SamplingPoolDask + + +class DaskSimulation: + need_workspace = False + + def __init__(self, sleep=0.0, fail=False, need_workspace=False): + self.sleep = sleep + self.fail = fail + self.need_workspace = need_workspace + + def level_instance(self, fine_level_params, coarse_level_params): + config = { + "fine": {"step": fine_level_params[0]}, + "coarse": {"step": coarse_level_params[0]}, + "sleep": self.sleep, + "fail": self.fail, + "res_format": self.result_format(), + } + return LevelSimulation(config_dict=config, task_size=0.0, need_sample_workspace=self.need_workspace) + + @staticmethod + def calculate(config, seed): + if config["sleep"]: + time.sleep(config["sleep"]) + if config["fail"]: + raise RuntimeError("sample failure") + + fine = np.array([float(seed % 1000) + config["fine"]["step"]]) + coarse = np.array([0.0 if config["coarse"]["step"] == 0 else float(seed % 1000)]) + return fine, coarse + + @staticmethod + def result_format(): + return [QuantitySpec(name="value", unit="1", shape=(1,), times=[0], locations=["0"])] + + +class RestartMemory(Memory): + def unfinished_ids(self): + return [ + sample_id + for sample_ids in self.load_scheduled_samples().values() + for sample_id in sample_ids + ] + + +@pytest.fixture +def dask_client(): + cluster = LocalCluster(n_workers=2, threads_per_worker=1, processes=False, dashboard_address=None) + client = Client(cluster) + try: + yield client + finally: + client.close() + cluster.close() + + +def make_level_simulation(simulation, level_id=0): + level_sim = simulation.level_instance([0.1], [0]) + level_sim._calculate = simulation.calculate + level_sim._result_format = simulation.result_format + level_sim._level_id = level_id + return level_sim + + +def collect_until_finished(sampler, timeout=5.0): + deadline = time.monotonic() + timeout + n_running = 1 + while n_running > 0 and time.monotonic() < deadline: + n_running = sampler.ask_sampling_pool_for_samples(timeout=0.05) + time.sleep(0.01) + assert n_running == 0 + + +def test_dask_pool_get_finished_is_non_blocking(dask_client): + pool = SamplingPoolDask(client=dask_client) + level_sim = make_level_simulation(DaskSimulation(sleep=0.5)) + + pool.schedule_sample("L00_S0000000", level_sim) + pool.schedule_sample("L00_S0000001", level_sim) + + start = time.monotonic() + successful, failed, n_running, n_ops = pool.get_finished() + elapsed = time.monotonic() - start + + assert elapsed < 0.3 + assert successful == {} + assert failed == {} + assert n_running == 2 + assert n_ops == [] + + deadline = time.monotonic() + 5 + while n_running > 0 and time.monotonic() < deadline: + successful, failed, n_running, n_ops = pool.get_finished() + time.sleep(0.01) + + assert n_running == 0 + assert len(successful[0]) == 2 + assert failed == {} + + +def test_dask_pool_collects_sampler_samples(dask_client): + storage = Memory() + pool = SamplingPoolDask(client=dask_client) + sampler = Sampler( + sample_storage=storage, + sampling_pool=pool, + sim_factory=DaskSimulation(sleep=0.01), + level_parameters=[[0.1], [0.01]], + ) + + sampler.set_initial_n_samples([4, 2]) + sampler.schedule_samples() + collect_until_finished(sampler) + + assert np.all(storage.n_finished() == np.array([4, 2])) + + +def test_dask_pool_supports_adaptive_sample_addition(dask_client): + storage = Memory() + pool = SamplingPoolDask(client=dask_client) + sampler = Sampler( + sample_storage=storage, + sampling_pool=pool, + sim_factory=DaskSimulation(sleep=0.01), + level_parameters=[[0.1], [0.01]], + ) + + sampler.set_initial_n_samples([2, 2]) + sampler.schedule_samples() + + sampler.process_adding_samples(np.array([6, 4]), sleep=0.01, add_coeff=0.5, timeout=0.05) + + assert np.all(np.array(sampler.l_scheduled_samples()) >= np.array([4, 3])) + collect_until_finished(sampler) + + +def test_dask_pool_resubmits_unfinished_workspace_samples_on_sampler_restart(dask_client, tmp_path): + storage = RestartMemory() + storage.save_scheduled_samples(0, ["L00_S0000000", "L00_S0000001"]) + first_pool = SamplingPoolDask(client=dask_client, work_dir=str(tmp_path)) + first_pool._save_level_sim(make_level_simulation(DaskSimulation(sleep=0.01, need_workspace=True))) + restarted_pool = SamplingPoolDask(client=dask_client, work_dir=str(tmp_path), clean=False) + + sampler = Sampler( + sample_storage=storage, + sampling_pool=restarted_pool, + sim_factory=DaskSimulation(sleep=0.01, need_workspace=True), + level_parameters=[[0.1]], + ) + + assert restarted_pool._n_running == 2 + collect_until_finished(sampler) + + assert np.all(storage.n_finished() == np.array([2])) + + +def test_dask_pool_reports_failed_samples(dask_client): + storage = Memory() + pool = SamplingPoolDask(client=dask_client) + sampler = Sampler( + sample_storage=storage, + sampling_pool=pool, + sim_factory=DaskSimulation(fail=True), + level_parameters=[[0.1]], + ) + + sampler.set_initial_n_samples([3]) + sampler.schedule_samples() + collect_until_finished(sampler) + + assert len(storage._failed[0]) == 3 + assert all("sample failure" in err_msg for _, err_msg in storage._failed[0]) diff --git a/tox.ini b/tox.ini index 71e6f5f..09cf20b 100644 --- a/tox.ini +++ b/tox.ini @@ -14,6 +14,8 @@ deps = pytest texttable matplotlib + dask + distributed # Get error for: pytest -m "not metacentrum" # But it seems that quoting works fin on tox side