Reconstruct drives its own write loop, so it cannot resume interrupted work #574

Activity

1. ieivanov commented on Aug 5, 2026

   

   ieivanov

   on Aug 5, 2026

   ContributorAuthor

   More actions

   I went through this against the current code in waveorder, iohub feat/interrupt-safe-writes, and biahub feat/resume-interrupted-positions. The conclusion is that Option C is the right call and is cheaper than the issue estimates — the headline risk is already retired and item 3 turns out to be a no-op. But two of the sketches would change behaviour if taken literally, and one of them silently. Details, then a proposed path for deleting apply_inverse_to_zyx_and_save.

   Caveat up front: this is all from reading code, not from running a reconstruction. The verification checklist at the end is the part that needs an actual run.

   The fork → spawn risk is already paid

   The issue calls this "the main risk … the part most likely to bite", on the premise that reconstruct uses torch.multiprocessing.Pool and forks. That was true when the issue was written but is not true on main: 507a8af ("Use ProcessPoolExecutor with spawn for apply-inv-tf workers", #558) landed after #555, and apply_inverse_transfer_function_single_position now uses

   context = mp.get_context("spawn")
   with ProcessPoolExecutor(max_workers=num_processes, mp_context=context) as p:

   with the same tensorstore-fork-safety and BrokenProcessPool reasoning iohub carries. So reconstruct already runs under exactly the executor and start method process_single_position uses.

   The premise does not hold for the pre-507a8af code either. multiprocessing.Pool.starmap serializes its task payload through _ForkingPickler in Pool._handle_tasks regardless of start method, so input_xa and the whole transfer-function xr.Dataset were already being pickled once per task under fork. Fork was sharing the parent's imports and pre-pool state, not the arguments.

   And the transfer function specifically cannot fail to pickle. _load_transfer_function_dataset materialises every array eagerly — np.array(transfer_function_dataset[key][idx]) — and transfer_function_dataset.close() runs before the pool is created, so there is no lazily-opened handle anywhere in the payload. Nothing in apply_inverse_args is a lambda, a closure, a bound method, or a torch tensor; tensors are created inside the worker by _to_tensor's torch.from_numpy(...). There is no GPU state in the payload.

   What spawn actually costs is a full torch import per worker and the loss of copy-on-write. That is startup overhead, not a correctness risk, and it is already being paid today.

   Adopting iohub's loop does not change resume granularity

   Worth knowing before weighing the benefits, because it cuts one of them.

   iohub's unit is a shard-aligned batch of timepoints for one channel group, batched on output_shards[0] in _indices_to_shard_aligned_batches. Reconstruct's output plate goes through create_empty_plate with no chunks and no shards_ratio — get_reconstruction_output_metadata returns only shape, scale, channel names, dtype, plate metadata and version — so it gets _default_shards_ratio, whose docstring is "Shards ratio yielding a shard of (1, 1, Z, Y, X)". The shard extent along T is 1, so each unit is a single timepoint, which is already the granularity of p.submit(partial, t_idx).

   Two consequences. The resume this buys is unit-for-unit what instrumenting the existing loop would give, so the win here is genuinely the maintenance one — one loop instead of two — rather than a finer resume. And "iohub batches timepoints that share a shard into one write" is a no-op at this geometry; it only starts paying if something later passes an explicit shards_ratio with a T extent above 1.

   Item 3 (per-write metadata) is a no-op, not a risk

   The issue flags this as "the one item worth checking carefully before switching, since it is the only place where behaviour could silently change". Reading write_xarray against the way reconstruct calls it, nothing persistent is lost:

   • The translation transform is never written. transforms is only passed to create_zeros, inside if image not in self:. Reconstruct pre-creates the plate, so that branch is skipped and the derived translation is discarded. The output keeps only the scale transform create_empty_plate wrote.
   • The axes/units are never written. self.axes = [...] shadows a cached_property; write_xarray does not call dump_meta(), and close() only closes the store. The coord-attrs-derived units never reach disk (and match _DEFAULT_AXES anyway).
   • iohub.xarray_attrs is never written, because data_array.attrs is always empty. _build_output_xarray constructs the output without attrs=, and nothing in waveorder/api/ ever assigns .attrs to a DataArray.

   So there is nothing to move to plate creation. An index-based write loses nothing.

   Item 4 is more than a confirmation — it may be a live bug

   The issue says "worth confirming this is equivalent when time_indices is a subset, and when the input's t coords are offset or non-uniform." It isn't equivalent. write_xarray recovers the index as

   t_indices = np.round((t_coords - translation[0]) / scale[0]).astype(int)

   where t_coords come from the input store (to_xarray builds them as arange(size) * scale[i] + translation[i]) while translation is read from the output array — which create_empty_plate gave no translation transform at all, since it only ever writes a scale. That works out to

   t_index = t_idx + round(input_translation_t / output_scale_t)
   

   Zero in the mantis pipeline, where every store comes from create_empty_plate. But standalone waveorder reconstruct on a raw acquisition store is the ordinary single-user path, and if any of those carry a nonzero T translation, every timepoint is written to a shifted or out-of-bounds index, silently. This is worth checking against a real acquisition store independently of this issue. Addressing by index fixes it by construction, which I would promote to a stated motivation for the switch rather than a box to tick.

   The interchange: Option A works, and the drift argument dissolves

   The case made for changing iohub is that Option A would have waveorder rebuild coords that iohub already knows, "which is exactly where a subtle bug would live: coords drifting from the store they came from." Checking what the model functions actually read from the labeled array, that risk is smaller than it looks:

   • No apply_inverse_transfer_function in any of the four api modules reads coords at all. They touch the input only through .values and .shape (torch.tensor(czyx_data.values, ...), or .values[0] for the phase/fluorescence single-channel paths). All pixel sizes come from the pydantic settings.transfer_function, never from coords.
   • coords["c"] appears only in compute_transfer_function (birefringence.py:184, birefringence_and_phase.py:274), as a fallback for counting polarization states when input_channel_names is not passed — and reconstruct always passes it.
   • The models under waveorder/models/ never see an xarray object at all.
   • The only coord consumption is _build_output_xarray (waveorder/api/_utils.py) copying z/y/x coord values and attrs onto the output DataArray — whose coords the write then discards, per the section above.
   • The transfer-function xr.Dataset is used purely as a string-keyed bag of arrays; _named_dataarray deliberately gives it meaningless synthetic dim names.

   So the coords the adapter has to supply exist only to keep _build_output_xarray from raising, and their values have no effect on stored data. More to the point, the adapter does not have to reconstruct them from input_dataset.scale — it can take them from to_xarray() once in the parent and pass them down as plain (dims, values, attrs) tuples. They then come from the same code that reads the data, which is the property Option B was there to guarantee.

   That removes the reason to add a labeled= / array_type= switch to iohub. I would rather not put a public two-path flag into process_single_position that every future feature in that loop has to honour, in exchange for six lines in one caller.

   One trap in the sketch, and one real behaviour change

   Output channel indices must be resolved by name. The sketch has

   output_channel_indices = [list(range(len(output_channel_names)))]

   write_xarray currently resolves them through get_channel_index(ch). Those differ whenever the reconstructed channels do not start at index 0 — create_empty_plate appends channels to an existing plate, and #547 namespaced the settings by output channel names precisely so several configs can write different channel groups into one plate. Keep the name lookup against output_dataset.channel_names or this silently overwrites the wrong channels. The same applies on the input side: build the index list in config order, not sorted, since oindex respects the order given and that is what .sel(c=input_channel_names) guarantees today.

   The all-zero/NaN predicate changes. This is the one genuine behaviour change and it isn't in the issue's list. waveorder's check is over the whole CZYX block:

   return np.all(np.isnan(input_array)) or np.all(input_array == 0)

   iohub's, for a 4-D block, is per channel with any:

   return any(_zyx_is_all_zero_or_nan(input_array[c]) for c in range(input_array.shape[0]))

   iohub skips a strict superset. A timepoint where one polarization state is blank but the others carry signal is reconstructed today and would be skipped after the switch, leaving zeros in the output. For reconstruction that is arguably the better semantics — a missing state makes the result meaningless — but it needs to be a deliberate decision and a changelog note, not a side effect. If waveorder's laxer semantics have to be preserved, that is an argument for making the predicate injectable in iohub, since it has to run before func to be worth anything.

   Proposed path for deleting apply_inverse_to_zyx_and_save

   Step 1 — a module-level numpy adapter. It must be module-level, not a closure: process_single_position submits the partial to a spawn pool, so a locally-defined function would not pickle. This also matches biahub's _<verb>_czyx() convention.

   # waveorder/cli/utils.py
   def apply_inverse_czyx(
       czyx: NDArray,
       model_function,
       czyx_coords: dict,
       **kwargs,
   ) -> NDArray:
       """CZYX numpy in/out wrapper around a ``waveorder.api`` apply-inverse function."""
       xa = xr.DataArray(czyx, dims=("c", "z", "y", "x"), coords=czyx_coords)
       return model_function(xa, **kwargs).values

   czyx_coords is built once in the parent from the coords iohub already produced, so nothing is recomputed from scale:

   input_xa = input_dataset.to_xarray()
   czyx_coords = {
       "c": ("c", settings.input_channel_names),
       **{
           d: (d, input_xa.coords[d].values, dict(input_xa.coords[d].attrs))
           for d in ("z", "y", "x")
       },
   }

   model_function stays one of the four module-level waveorder.api.* functions, so the payload is still picklable and waveorder.api.* signatures are untouched.

   Step 2 — resolve channels to indices, in config order, as a single group.

   input_channel_indices = [[input_dataset.channel_names.index(n) for n in settings.input_channel_names]]
   output_channel_indices = [[output_dataset.channel_names.index(n) for n in settings.output_channel_names]]

   The nesting matters: a flat list makes iohub treat each channel as its own task, which is the biahub concatenate case. One group keeps the model receiving one CZYX.

   Step 3 — address timepoints by index. The output plate is allocated with the full input T, so input and output indices coincide even for a subset:

   time_indices = resolve_time_indices(settings.time_indices, T)

   passed as both input_time_indices and output_time_indices.

   Step 4 — replace the pool.

   process_single_position(
       func=apply_inverse_czyx,
       input_position_path=input_position_dirpath,
       output_position_path=output_position_dirpath,
       input_channel_indices=input_channel_indices,
       output_channel_indices=output_channel_indices,
       input_time_indices=time_indices,
       output_time_indices=time_indices,
       num_workers=num_processes,
       resume=resume,
       resume_token=reconstruction_fingerprint(settings, transfer_function_dataset),
       model_function=apply_inverse_model_function,
       czyx_coords=czyx_coords,
       **apply_inverse_args,
   )

   Everything after resume_token rides through **kwargs to the adapter. num_workers is additionally clamped by iohub's _available_cpus(), which honours SLURM_CPUS_PER_TASK — a small improvement over the current mp.cpu_count() bound in processes_option. use_threads=True is available if #566/#552 lands, so this composes with either outcome.

   Step 5 — the resume token must cover the transfer function. The issue says "ideally"; I would make it required. biahub's settings_fingerprint hashes only the resolved settings, so recomputing a transfer function and re-running apply-inverse into the same output store yields the same token and silently reuses timepoints computed with the old transfer function. That is a data-correctness bug, not wasted compute. The store already records what is needed — compute_transfer_function writes output_dataset.zattrs["settings"] = settings.model_dump() — so:

   def reconstruction_fingerprint(settings, transfer_function_dataset) -> str:
       payload = json.dumps(
           {
               "settings": settings.model_dump(mode="json"),
               "transfer_function": dict(transfer_function_dataset.zattrs.get("settings", {})),
           },
           sort_keys=True,
           default=str,
       )
       return hashlib.sha256(payload.encode()).hexdigest()[:16]

   Note this has to be computed before transfer_function_dataset.close().

   Step 6 — delete. apply_inverse_to_zyx_and_save and _check_nan_n_zeros from waveorder/cli/utils.py, and the partial plus executor block from apply_inverse_transfer_function.py. input_dataset.to_xarray() shrinks to a coords-only read. Kept as is: the settings and channel validation, resolve_time_indices, the transfer-function load, and the post-pool zattrs["waveorder"] write.

   Sequencing. Two waveorder PRs — the loop swap, then --resume — is worth it for reviewability, but both should target the same iohub, because the released process_single_position takes num_processes/num_threads while feat/interrupt-safe-writes renamed them to num_workers/use_threads. Splitting across versions means writing the call twice. Pin the branch for review and repoint at the release when it merges, which is what biahub#299 already does. Then the biahub side: pass --resume from biahub/apply_inverse_transfer_function.py, add it to the reconstruct Nextflow module (which has none today, unlike deskew and flat-field), and document @resume() in biahub's CLI-consistency canon, which is still one step behind that branch.

   Verification checklist

   1. Byte-identical output on a dataset with no blank channels, against main, for each of the four reconstruction modes. This is the main gate on step 1–4 being a pure refactor.
   2. The all-zero/NaN divergence, deliberately: a timepoint with one blank input channel and signal in the others. Confirm the new skip is what you want, and note it in the changelog.
   3. Channel offsetting: two configs writing different output channel groups into one plate, to confirm the name lookup in step 2 (this is the case range(len(...)) would break).
   4. A time_indices subset, and a store with a nonzero T translation if you can find or fabricate one — the second is the case the current coord round-trip gets wrong.
   5. Resume, by killing a position mid-run and confirming the retry recomputes only the unfinished timepoints, and that a changed config or a recomputed transfer function invalidates the records.
   6. A v0.4 input, since get_reconstruction_output_metadata propagates the input plate's version and tracking_available is False for Zarr v2 — resume degrades to a warning plus full recompute there, which is fine but worth echoing from the CLI rather than only warning from inside iohub.

   One thing this path does not cover, worth noting since it argues for keeping iohub's write-unit helpers importable rather than internal: virtual_stain and track also drive their own per-timepoint writes (biahub/virtual_stain.py, biahub/track.py) and cannot adopt process_single_position — virtual_stain runs a torch model with its own batching. iohub#455 already lists virtual_stain as getting neither repair nor resume. Unifying reconstruct fixes reconstruct; those two would still need plan_write_unit / unit_is_complete / progress_dir_for to be public before they can record anything.
2. ieivanov commented on Aug 6, 2026

   

   ieivanov

   on Aug 6, 2026

   ContributorAuthor

   More actions

   @talonchandler @srivarra Sorry for the wall of text, I'll try to go through it and decide what actually makes sense. Let me know if you are generally opposed to the changes suggested here - getting rid of apply_inverse_to_zyx_and_save and relying on iohub utils instead.
3. talonchandler commented on Aug 6, 2026

   

   talonchandler

   on Aug 6, 2026

   Collaborator

   More actions

   getting rid of apply_inverse_to_zyx_and_save and relying on iohub utils instead.

   Makes sense to me! Thanks @ieivanov.
4. ieivanov

   linked a pull request that will close this issuerefactor(cli)!: run apply-inv-tf through iohub's process_single_position, with --resume and per-reconstruction provenance keys #583on Sep 28, 2026
5. ieivanov

   added a commit that references this issue on Oct 16, 2026

   refactor(cli)!: run apply-inv-tf through iohub's process_single_posit…

   ...

   8d79650
6. ieivanov

   closed this as completedin #583on Oct 16, 2026