Jb trans sa

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

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)

requirements.txt

@@ -1,2 +1,4 @@
 pytest
 matplotlib
+dask
+distributed