feat(core): PymoniK rewrite

[AncientPatata]

Motivation

PymoniK 0.1 proved out the @task decorator model under time pressure but is fragile under the hood: a schema-less magic-string argument protocol, four overlapping ways to set task options, a ~50-line
hand-rolled option merge, double-pickling via LazyArgs, a worker that re-uses the client class via is_worker=True and rebuilds a client per task, polling instead of event streams, and no async,
retries, OTel, local-test story, or typed errors.

This PR is the v1 clean-break rewrite specified: decorator-first ergonomics, typed end-to-end, a single typed wire envelope, first-class local execution, etc.

Description

Greenfield src/pymonik/ package replacing the v1 pymonik/ + pymonik_worker/ + test_client/ trees. Highlights:

• Client/session split (client.py, session.py): PymonikClient owns the channel/credentials/OTel; Session owns the session id, default options, the multiplexed events stream (with a polling fallback),
  retries, and result resolution. Sync + async surfaces (session()/session_async(), .result()/await).
• Task model (task.py, options.py, composition.py): @task with ParamSpec/@overload signature preservation; .spawn/.map/.starmap/.with_options/.tail; TaskOpts as a frozen/slots/kw_only dataclass with a
  field-driven .merge(); gather/as_completed (+ sync siblings).
• Typed wire envelope (envelope.py): one msgspec.Struct with a version field and an embedded client Python-minor guard that raises instead of SIGSEGV'ing on a cloudpickle minor mismatch. Arg refs
  (_internal/refs.py) replace 0.1's magic strings; oversize inline args auto-spill to content-addressed blobs.
• Futures (future.py, multiresult.py): Future[T]/FutureList[T], lazy materialization (download only on .result()/await), MultiResult/MultiResultHandle/TailPromise for multi-output and tail-call sub-tasking.
• Worker runtime (worker.py, worker_session.py, _internal/task_runner.py, subprocess_dispatch.py): single boot, shared submission pipeline, in-process splice vs isolate=True subprocess for runtime deps.
• Cross-cutting: typed PymonikError hierarchy; structlog; optional [otel] trace propagation; public pymonik.hooks; fluent introspection DSL client.tasks.where(...).list(); LocalCluster in-process backend (testing/local.py); content-addressed runtime-dep envs and a structural-key result-reuse cache.
• Adds .docs/guides/, runnable examples/.py, tests/, worker-image/Dockerfile; removes 0.1 trees and automation.py.

Testing

multiple unit/local tests (tests/) against LocalCluster and stubbed armonik clients — no network. Strong coverage of TaskOpts.merge, env_id canonicalization, envelope round-trips, MultiResult AST extraction, map/starmap, lazy futures, multi-partition, hooks.

Impact

No direct impact on ArmoniK. Might need to make changes to the existing PymoniK partition.

Additional Information

None yet.

Checklist

• My code adheres to the coding and style guidelines of the project.
• I have performed a self-review of my code.
• I have commented my code, particularly in hard-to-understand areas.
• [~] I have made corresponding changes to the documentation.
• I have thoroughly tested my modifications and added tests when necessary.
• Tests pass locally and in the CI.
• I have assessed the performance impact of my modifications.

[CLAassistant]

All committers have signed the CLA.

[aneojgurhem]

Security review

One medium-confidence finding worth fixing before merge.

Credential leak in worker logs — src/pymonik/_internal/env_builder.py:164

ensure_env() logs the full index_url at INFO level:

log.info(
    "env build start",
    env_id=env_id,
    deps=list(canonical_deps(spec.deps)),
    index_url=spec.index_url or None,   # ← emits embedded token in plain text
)

The docs recommend embedding credentials in this URL (https://ci-user:ghp_TOKEN@private-pypi.example.com/simple/). When a user follows that pattern the token lands in every log sink — stdout, Seq, k8s pod logs, the polling-agent pipeline. Anyone with log-read access (a lower-privilege role than cluster admin) can extract the token.

Fix: redact the userinfo component before logging:

from urllib.parse import urlparse, urlunparse

def _redact_url(url: str) -> str:
    p = urlparse(url)
    if p.username or p.password:
        host = p.hostname + (f":{p.port}" if p.port else "")
        return urlunparse(p._replace(netloc=f"<redacted>@{host}"))
    return url

log.info(
    "env build start",
    env_id=env_id,
    deps=list(canonical_deps(spec.deps)),
    index_url=_redact_url(spec.index_url) if spec.index_url else None,
)

No other high-confidence issues found. Pickle deserialization, zip slip, and uv arg injection were all examined and ruled out as false positives or below threshold given the existing cloudpickle trust model.

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of lines (20,000). Try reducing the number of changed lines and requesting a review from Copilot again.

[aneojgurhem]

Automated review from 10-angle multi-agent analysis (15 confirmed findings across 3 critical / 8 high / 4 medium).

[aneojgurhem]

⚠️ High: _submit_tail omits ctx_param — context injection broken for tail-called tasks

This TaskEnvelope(...) call does not pass ctx_param, so it defaults to "". The worker dispatch code gates context injection on if envelope.ctx_param: — the condition is false, injection is skipped. A tail-called task that declares ctx: pymonik.Ctx is called without the argument → TypeError: <fn>() missing required argument: 'ctx', surfacing as TaskFailed on the parent future.

The main submit path in _internal/submit.py:280 correctly passes ctx_param=task.ctx_param or "".

Fix: add ctx_param=task.ctx_param or "", to the TaskEnvelope(...) call here.

[aneojgurhem]

⚠️ High: __exit__ sets _stop but never sets _result_events — deadlocks executor.shutdown(wait=True) when tasks have pending deps

Two-task DAG: task B depends on A's output. If an exception is raised in the with block before A finishes, __exit__ sets _stop and clears _pending, but does not set any _result_events. B's dispatcher thread is blocked at ev.wait() waiting for A's result event, which is never set. LocalCluster.__exit__ calls executor.shutdown(wait=True) → deadlock. Note: cancel() avoids this because it explicitly sets all events; a plain context-manager exception does not.

Fix: during teardown (before executor.shutdown), iterate _result_events and call .set() on each.

[aneojgurhem]

ℹ️ Medium: LocalCluster(cache=True) is a silent no-op — cache is never read or written

This method unconditionally returns ({}, list(range(len(normalised))), {}) — every call is a miss and the keys dict is always {}. Because fut._cache_key is never set, the put_bytes write path in _dispatch_result is also never triggered. The ExecCache object is constructed, passed in, and completely ignored. Running the same @task(cache=True) twice always re-executes.

Fix: implement this method using self._cache.get_bytes(key), mirroring Session._cache_classify.

[aneojgurhem]

⚠️ High: retry thread has no error handler — resubmit failure silently kills the thread and leaves the future permanently unresolved

If submit_many(...) raises (e.g., a partition validation error), this daemon thread dies silently. fut._done is never set. Any fut.result() or await fut hangs forever. The equivalent in session.py:385 wraps the call in try/except and resolves the future with an error on failure — that guard is absent here.

Fix: mirror session.py:385 — wrap the submit_many call in try/except Exception as e: and call fut._resolve_error(TaskFailed(fut.task_id, str(e))).

[aneojgurhem]

⚠️ Medium: sys.path check-insert-remove is not locked — concurrent tasks with the same venv dep race to ModuleNotFoundError

if site not in _sys.path: _sys.path.insert(0, site) and the finally: _sys.path.remove(spliced_path) are unprotected. LocalCluster dispatches on a ThreadPoolExecutor. Two concurrent tasks with the same env_spec both pass the check before either inserts, both insert, and the first to finish removes the path — the second task's imports then fail with ModuleNotFoundError mid-execution. self._lock exists and is used elsewhere in this class, but not here.

Fix: acquire self._lock around the entire check-insert-execute-remove sequence, or use a reference-count so the path is only removed when the last holder is done.

[aneojgurhem]

⚠️ High: LocalSession._submit_tail omits ctx_param — context injection broken for tail-called tasks (mirrors worker_session.py:205)

Same mechanism as worker_session.py:205: this TaskEnvelope(...) call does not pass ctx_param, defaulting it to "". The local dispatcher at line 1058 checks if envelope.ctx_param: — false, injection skipped. A tail-called task with ctx: Ctx raises TypeError during local test runs, masking the production bug.

Fix: add ctx_param=task.ctx_param or "", to the TaskEnvelope(...) call here.

[lemaitre-aneo]

I think the code could be refactored in many places to reduce the number of edge case handling.

[lemaitre-aneo]

Why don't you use the API type here?

[lemaitre-aneo]

Do you really need the lock to be held during the download?