PML-325 MerLinProcessor Cleaning

[ben9871]

Summary

Four-ticket internal refactor of MerlinProcessor executed as one stacked branch (9 commits, one logical step per commit). The processor's execution path goes from a single ~2,200-line class threading an anonymous state dict and monkey-patched futures, to a thin coordinator over typed, independently-testable units:

• PML-303 — the mutable state dict threaded through forward_async(), chunk orchestration, and job polling is replaced by a typed CallState object with named mutation helpers and documented per-field thread ownership.
• PML-305 — chunk orchestration and remote job execution are extracted into BatchChunker and RemoteJobRunner (merlin/core/execution.py), dependency-injected and testable with no-cloud fakes. Threading semantics unchanged; no concurrency-model redesign.
• PML-304 — the runtime monkey-patching of cancel_remote / status / job_ids onto plain torch.futures.Future instances is replaced by MerlinFuture, a Future subclass declaring the async contract in code and types. Exported at package root.
• PML-306 — all direct access to Perceval internals (RPC handler token/URL attributes, sampler command objects, job status fields, RemoteConfig, local experiment private state) is localized in a PercevalAdapter facade (merlin/core/perceval_adapter.py). Scope was deliberately extended to cover the local processor path (experiment _in_ports / _postselect / _circuit_changed handling), so a Perceval version bump now requires touching only the adapter.

User-visible behavior is unchanged. The PML-300 behavior-lock and PML-301 backend-parity suites pass throughout; the only test edits across the whole branch are mechanical (7 accessor lines for the CallState fixture, 10 Sampler patch-target retargets to the module that now owns the symbol).

Related Issue

Related Jira: PML-303, PML-304, PML-305, PML-306

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor / Cleanup
• Performance improvement
• CI / Build / Tooling
• Breaking change (requires migration notes)

Proposed changes

PML-303 — typed CallState (80c96602, 025b0f12)

• New CallState (+ frozen JobStatus): request_cancel(), record_job_id() (deduplicating, under lock), set_current_status(), add_planned_chunks(), mark_chunk_started()/finished(), status_snapshot().
• future.job_ids still aliases the live CallState list; future.status() payload is byte-identical; chunk counters move from the processor-wide lock to CallState's own lock.

PML-305 — execution units (bfb1607f)

• BatchChunker: batch splitting, bounded-concurrency chunk pooling, job labeling, output stitching, deadline → cancel_all → TimeoutError.
• RemoteJobRunner: fresh RemoteProcessor per attempt, iteration prep, probs-vs-sampling command selection, job-name capping, retries/backoff, polling, result mapping via injected callable.
• MerlinProcessor._run_chunks_pooled / _run_chunk / _submit_job / _poll_job become thin delegates with unchanged signatures; dependencies are bound methods and deferred lambdas so live attribute mutation and monkeypatched methods behave exactly as before.
• Known pre-existing quirk pinned by test (not changed): failures in fresh-processor construction are not retried; only submit/poll failures are.

PML-304 — typed async handle (d176114f)

• MerlinFuture(torch.futures.Future) declares job_ids (live property), status(), cancel_remote(). forward_async() is annotated -> MerlinFuture.
• isinstance(fut, torch.futures.Future) contract preserved (subclassing was mandated by the parity suite). Exported from merlin/__init__.py.

PML-306 — Perceval adapter (a41620ab, 805fb59e, 8075c3fe, 929ea1f5)

• PercevalAdapter (stateless static facade): token extraction (handler attr variants, Bearer header, RemoteConfig fallback), URL, RP cloning, session build, capability snapshot, processor configuration, sampler creation/dispatch (async + sync), Perceval 1.2 iterator-serialization shim, JobStatusSnapshot normalization, results/cancel, local processor rebuild with LocalExperimentSnapshot capture/restore, shot estimation.
• New Merlin-specific exceptions preserving historical contracts: TokenExtractionError(ValueError), RemoteJobFailedError(RuntimeError) — same messages, same catchability.
• _extract_rp_token, _clone_remote_processor, _create_fresh_local_processor remain as processor entry points (delegates), so existing patch targets are unaffected.

API changes

• Added (public): merlin.MerlinFuture.
• Added (internal but importable): merlin.core.merlin_processor.CallState / JobStatus; merlin.core.execution.BatchChunker / RemoteJobRunner; merlin.core.perceval_adapter.PercevalAdapter / JobStatusSnapshot / LocalExperimentSnapshot / TokenExtractionError / RemoteJobFailedError.
• Removed (private): MerlinProcessor._snapshot_local_experiment_metadata, _restore_local_experiment_metadata, _copy_circuit_for_execution, _ensure_serializable_sampler_iterator, _LocalExperimentMetadata (all relocated into the adapter).
• No public signature or behavior changes to MerlinProcessor.

Docs & quality (4f9d5fba)

• MerlinFuture documented in the merlin.core.merlin_processor API reference; the forward_async section no longer describes the removed monkey-patched contract.
• NumPy docstrings completed on CallState / PercevalAdapter per AGENTS.md; test files aligned with the suite's one-docstring-per-test pattern; 4 latent ruff S106 findings fixed.

How to test / How to run

1. Full processor-related suites (no cloud credentials needed; credentialed cloud tests skip):

pytest tests/core/test_call_state.py tests/core/test_execution_units.py tests/core/test_merlin_future.py tests/core/test_perceval_adapter.py tests/core/test_merlin_processor_unit.py tests/core/test_merlin_processor_gradient.py tests/core/cloud -q

Expected: 275 passed, 77 skipped (skips = live-backend integration tests).

2. Lint / format:

ruff check merlin/ tests/core/ && ruff format --check merlin/ tests/core/

3. Docs:

SPHINXOPTS="-W --keep-going -n" make -C docs clean html

Verified clean on this branch (via python -m sphinx -W --keep-going -n -b html docs/source <out>, exit 0).

4. Review guidance: the branch is designed to be reviewed commit-by-commit — each commit is independently green, and each "add unit" commit (80c96602, a41620ab) is a pure addition reviewable before its wiring commit.

Screenshots / Logs (optional)

n/a — behavior-preserving refactor; parity is demonstrated by the locked suites passing with only mechanical edits.

Performance considerations (optional)

Negligible. Per-call construction of the small runner/chunker/adapter objects adds a few allocations per chunk; polling cadence, thread model, retry/backoff, and remote payloads are unchanged.

Documentation

• User docs updated (Sphinx)
• Examples / notebooks updated
• Docstrings updated
• Updated the API

Checklist

• PR title includes Jira issue key (e.g., PML-126)

• "Related Jira ticket" section includes the Jira issue key (no URL)

• Code formatted (ruff format)

• Lint passes (ruff)

• Static typing passes (mypy) if applicable — not configured for this repo; not run

• Unit tests added/updated (pytest) — 104 new no-cloud tests across 4 new files

• Tests pass locally (pytest) — 275 passed, 77 skipped (credentialed cloud tests)

• Tests pass on GPU (pytest) — not run; no CUDA-path changes (inputs already moved to CPU for backend execution, unchanged)

• Test coverage not decreased significantly — coverage increased: token extraction, job-status guards, and chunk concurrency previously had no direct coverage

• Docs build locally if affected (sphinx)

• With this command:

    > SPHINXOPTS="-W --keep-going -n" make -C docs clean html
  

  the docs are built without any warning or errors.

• New public classes/methods/packages are added in the API following the methodology presented in other files.

• Dependencies updated (if needed) and pinned appropriately — no dependency changes

• PR description explains what changed and how to validate it

[CassNot]

@claude review

[claude]

Claude finished @CassNot's task in 4m 26s —— View job

---

Review — PML-325 MerlinProcessor Cleaning

• Gather context (AGENTS.md, diff vs bb2821b, base file comparison)
• Review merlin/core/perceval_adapter.py
• Review merlin/core/execution.py
• Review CallState / MerlinFuture in merlin_processor.py
• Review tests & docs
• Post review

---

Overall this is a genuinely good refactor. The commit-per-step structure holds up, the units are cleanly injected, and the extractions are faithful line-for-line against the base — I diffed _run_chunks_pooled, _run_chunk, _poll_job, _run_chunk_local, and the local-experiment snapshot path against bb2821b and the control flow, retry/backoff, sleep cadence, and error messages are all preserved. CallState is a real improvement over the anonymous dict, and 98 new no-cloud tests with meaningful assertions is the right kind of coverage to add here.

Below are the things I'd want addressed, roughly in order of importance.

Note on verification: I was not able to run pytest or ruff in this environment (command approval denied), so this is a static review only. The parity argument rests on the locked suites you ran locally.

---

1. _submit_job / _poll_job are now orphaned wrappers — overrides no longer reach the execution path

MerlinProcessor._submit_job (merlin/core/merlin_processor.py:1485) and _poll_job (:1532) each build a fresh RemoteJobRunner and call into it. But the real remote path never goes through them: RemoteJobRunner.run_chunk calls self.submit_job(...) and self.poll_job(...) on itself (merlin/core/execution.py:266, :271).

Consequence: _make_job_runner's docstring says dependencies are wired "so that … monkeypatched processor methods keep working exactly as before the extraction," and the PR description repeats that claim. It's true for _run_chunk (injected as a bound method into BatchChunker) but false for _submit_job / _poll_job — a subclass or patch.object override of either is silently bypassed in production.

The ~25 tests calling proc._submit_job(...) / proc._poll_job(...) still pass because the wrapper reaches the same code, but they now exercise a path production doesn't use. Pick one: inject submit_job/poll_job into the runner the way run_chunk is injected into the chunker, or delete the processor wrappers and retarget those tests at RemoteJobRunner. Keeping both unwired is the worst of the two options.

Fix this →

2. Two sources of truth for input-parameter ordering

merlin/core/execution.py:229 reads list(config.input_param_order) directly, while the local path (merlin_processor.py:1394) still goes through self._extract_input_params(config). Behaviour is identical today (_extract_input_params is exactly list(config.input_param_order)), but tests/core/cloud/test_parammapping.py:8 documents _extract_input_params as the param-routing seam in the pipeline, and it's now bypassed on the remote path. Inject it alongside the runner's other dependencies so both paths share one hook.

Fix this →

3. The snapshot handoff is hidden state — exactly what PML-306 was for

PercevalAdapter.rebuild_local_processor (perceval_adapter.py:575) stashes the LocalExperimentSnapshot on the returned processor via setattr(..., LOCAL_EXPERIMENT_SNAPSHOT_ATTR, ...), and the caller must then remember to getattr it, isinstance-check it, and call restore_experiment (merlin_processor.py:1410-1414). AGENTS.md is explicit: "Avoid hidden state, implicit side effects, and surprising mutation."

This is inherited from the base, but PML-306 is the commit that gets to fix it — returning tuple[AProcessor, LocalExperimentSnapshot] makes the contract impossible to get wrong. As written, a caller that forgets the restore step silently loses heralds/postselection rather than failing, and the isinstance guard at :1411 turns a snapshot-shape regression into silent metadata loss rather than an error.

Fix this →

4. Docstring completeness is inconsistent with the PR's own claim

The PR says "NumPy docstrings completed on CallState / PercevalAdapter per AGENTS.md" — and those two are indeed thorough. But the new public methods in execution.py are not, despite AGENTS.md requiring Parameters/Returns/Raises on all public methods:

Method	State
BatchChunker.split_batch (execution.py:73)	one-line
BatchChunker.run_chunks (:83)	one-line
RemoteJobRunner.run_chunk (:208)	one-line
RemoteJobRunner.poll_job (:359)	4 lines, no sections
RemoteJobRunner.submit_job (:305)	full — but see #5

The net effect is that the complete _poll_job documentation now lives on the thin wrapper in merlin_processor.py:1499-1531 while the implementation carries a stub.

Fix this →

5. RemoteJobRunner.submit_job: untyped signature + a parameter that is always the same value

def submit_job(self, sampler, nsample, job_base_label, capped_name):

merlin/core/execution.py:305 — no annotations (AGENTS.md: "a typed function signature"), and capped_name is always self._capped_name at the only real call site (:267). The parameter exists solely to preserve MerlinProcessor._submit_job's 4-arg signature for the existing tests. Type it, and consider dropping the parameter (let the wrapper absorb the arity difference) so the runner doesn't carry a vestigial injection point.

Fix this →

6. Smaller items

• chunk_concurrency is now snapshotted, not re-read. BatchChunker.__init__ (execution.py:69) captures it once; the base loop re-read self.chunk_concurrency on every iteration (bb2821b merlin_processor.py:1055). Mid-flight mutation no longer takes effect. Marginal in practice, but it is a divergence from "behavior unchanged" and worth either fixing (accept a callable, like the other deps) or noting. The max(1, int(...)) clamp is also now applied twice — here and at merlin_processor.py:873.
• MerlinFuture.job_ids is read-only. Previously fut.job_ids, fut.status, fut.cancel_remote were plain instance attributes, so assignment worked; now job_ids is a property with no setter and assignment raises AttributeError. This is a net improvement, but MerlinFuture is exported at package root — worth a line in the release notes, since the PR currently states "no public signature or behavior changes."
• Inconsistent guarding of get_rpc_handler(). extract_token (perceval_adapter.py:129) wraps it in try/except Exception; get_url (:175) calls it bare, and clone_remote_processor (:206) routes through get_url. Pick one policy — AGENTS.md's failure policy argues for the unguarded version in both.
• TokenExtractionError lives in the adapter but is only raised outside it (merlin_processor.py:857); PercevalAdapter.extract_token returns None. The docstring at perceval_adapter.py:39 ("Raised when no auth token can be resolved for a RemoteProcessor") describes behaviour that isn't in this module — either move the raise into extract_token or reword.
• configure_processor(processor, circuit, None) at merlin_processor.py:1407 calls the combined helper purely to set the circuit, relying on set_input's falsy check to no-op. A dedicated PercevalAdapter.set_circuit would read better than the None sentinel.
• BatchChunker.split_batch loops forever if microbatch_size <= 0. MerlinProcessor.__init__:864 assigns microbatch_size without validation. Pre-existing, but this is now a public static method with a documented contract — a cheap if microbatch_size <= 0: raise ValueError(...) closes it.
• Function-local from concurrent.futures import CancelledError at execution.py:219, :375, and merlin_processor.py:1384. Carried over from the base, but execution.py is a new file with no circular-import constraint — move it to the top.

---

What's good

• The CallState thread-ownership docblock (merlin_processor.py:76-90) is the kind of documentation that actually prevents future bugs. Moving the chunk counters off the processor-wide lock onto CallState's own lock is the right call.
• MerlinFuture subclassing torch.futures.Future preserves the isinstance contract while making the API discoverable and typed — and status()/job_ids don't collide with anything on the torch base class.
• Dependency injection in RemoteJobRunner is done with deferred lambdas (_make_job_runner, merlin_processor.py:1251-1265) specifically so live attribute mutation keeps working. That's careful work, and it's why the parity suites pass. Item Update issue templates #1 is the one place the same discipline wasn't applied.
• RemoteJobFailedError(RuntimeError) / TokenExtractionError(ValueError) as subclasses of the historical types is the right way to add specificity without breaking catchers.
• Test fakes in tests/core/test_execution_units.py are well-shaped — scripted result_events, recorded execute_kwargs, real threads for the concurrency assertions. Not mock-theatre.
  · branch PML_303_new