From 80aa4171b2fea4ebd5a37a84b90cfc47b1df0746 Mon Sep 17 00:00:00 2001 From: uermel Date: Mon, 19 Oct 2026 21:13:33 -0700 Subject: [PATCH 1/2] feat: --xla, XLA-compiled inference for .h5 and .scnm models `copick inference easymode --xla` runs each model's forward pass as one XLA-compiled function. It is off by default. An `XLAModel` wrapper is called like the Keras model and passes every other attribute through, so easymode's tiling of an .h5 network and the .scnm inference use it unchanged. Each worker compiles at the first call of each input shape and logs the time. If a compiled call raises, the worker warns once and continues with the plain model. The report records `"xla"`, and each worker lists in `xla_fallbacks` any model that gave up XLA, with the reason. On three 10521 tomograms at TTA 4 on one L40S with TensorFlow 2.21: the microtubule .h5 network went from 310 s to 107 s per tomogram, 2.9x, and the actin .scnm slab model from 87 s to 67 s, 1.3x. The masks matched the plain model's to an IoU of 0.9998 or better. Co-Authored-By: Claude Opus 5.5 --- README.md | 23 +++++++-- src/copick_easymode/cli/inference.py | 18 ++++++- src/copick_easymode/core/dispatch.py | 9 ++++ src/copick_easymode/core/inference.py | 46 ++++++++++++++--- src/copick_easymode/core/xla.py | 71 ++++++++++++++++++++++++++ tests/test_cli.py | 13 +++++ tests/test_dispatch.py | 41 +++++++++++++++ tests/test_xla.py | 72 +++++++++++++++++++++++++++ 8 files changed, 280 insertions(+), 13 deletions(-) create mode 100644 src/copick_easymode/core/xla.py create mode 100644 tests/test_xla.py diff --git a/README.md b/README.md index 4b487c4..82f29ab 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,9 @@ copick inference easymode -c config.json -m ribosome -t wbp@10.0 --cpu # High quality with test-time augmentation copick inference easymode -c config.json -m ribosome -t wbp@10.0 --tta 16 +# XLA-compiled inference, about 3x faster on a GPU +copick inference easymode -c config.json -m microtubule -t wbp@10.0 --xla + # Don't add object definitions to config copick inference easymode -c config.json -m ribosome -t wbp@10.0 --no-add-objects @@ -108,6 +111,15 @@ in the command's own process. cannot collide. - **Config writes.** With `--add-objects`, only the parent process adds object definitions and saves the config, before the workers start; workers re-open the project from the config file and never write it. +- **XLA.** `--xla` runs each model's forward pass as one XLA-compiled function (`tf.function(jit_compile=True)`), + which fuses the memory-bound work between the convolutions. Each worker compiles at the first call of each input + shape, which takes tens of seconds and is logged as `XLA compiled for input (...)`, and reuses it for every + later tile of that shape. On 10521 tomograms of 320 x 772 x 548 voxels, at TTA 4 on one L40S with TensorFlow 2.21, + it took the microtubule `.h5` network from 310 s to 107 s per tomogram, 2.9x, and the actin `.scnm` slab model from + 87 s to 67 s, 1.3x. The masks matched the plain model's to an IoU of 0.9998 or better. If compiling fails, the + worker logs one warning, continues with the plain model and lists the model with the reason in its `xla_fallbacks` + in the report. It is off by default because a missing or too old `ptxas` can abort TensorFlow instead of raising an + error. - **Logs** of each worker start with `[worker k gpu X]`. ### Exit status and report @@ -124,7 +136,7 @@ the end (`"complete"` or `"failed"`), whatever the outcome: "tool": "copick-easymode", "version": "0.3.0", "status": "complete", "config": "/data/copick_config.json", "tomogram": "wbp@10.0", "runs": ["TS_01", "TS_02", "TS_03"], "user_id": "copick", "session_id": "1", "tta": 4, "threshold": 0.5, "batch_size": 1, "overwrite": false, - "models_requested": ["microtubule"], + "xla": false, "models_requested": ["microtubule"], "models": [{"feature": "microtubule", "object": "microtubule", "weights": "/shared/easymode/models/microtubule_sv3-3d.h5", "kind": "h5", "tag": "sv3-3d", "timestamp": "20260916223602", "bytes": 515770304, "apix": 10.0}], "missing": {}, "model_directory": "/shared/easymode", "online": true, "writable": true, "easymode_version": "1.2.5", @@ -133,9 +145,10 @@ the end (`"complete"` or `"failed"`), whatever the outcome: {"index": 0, "gpu": "0", "runs": ["TS_01", "TS_03"], "exitcode": 0, "seconds": 412.3, "processed": 2, "skipped": 0, "errors": [], "visible_devices": "0", "threads": 8, "items": [{"run": "TS_01", "feature": "microtubule", "status": "processed", - "seconds": {"read": 2.1, "prepare": 1.9, "infer": 198.0, "finish": 1.2, "write": 0.8}}]}, + "seconds": {"read": 2.1, "prepare": 1.9, "infer": 198.0, "finish": 1.2, "write": 0.8}}], + "xla_fallbacks": []}, {"index": 1, "gpu": "1", "runs": ["TS_02"], "exitcode": 0, "seconds": 205.7, "processed": 1, "skipped": 0, - "errors": [], "visible_devices": "1", "threads": 8, "items": []} + "errors": [], "visible_devices": "1", "threads": 8, "items": [], "xla_fallbacks": []} ], "processed": 3, "skipped": 0, "errors": [], "started_utc": "2026-10-08T20:00:00Z", "finished_utc": "2026-10-08T20:06:55Z", "seconds": 414.9 @@ -144,7 +157,8 @@ the end (`"complete"` or `"failed"`), whatever the outcome: `missing` names each model that could not be resolved, with easymode's reason. A worker's `gpu` is the `CUDA_VISIBLE_DEVICES` value it ran with (`null` with `--cpu`), and its `exitcode` is its process exit status (0 when -all of its runs were handled without an error). +all of its runs were handled without an error). With `--xla`, a worker's `xla_fallbacks` lists each model that fell +back to plain inference, as `{"feature": ..., "reason": ...}`; its runs still count as processed. ### Python API @@ -187,6 +201,7 @@ IDs, so a rerun finds them. | `--tta` | Test-time augmentation level 1-16 (8 at most for an `.scnm` model). Higher = better but slower. Default: 4 | | `--batch-size` | Batch size for inference. Default: 1 | | `--threshold` | Probability threshold for binarizing a segmentation. Default: 0.5 | +| `--xla/--no-xla` | Run each model as one XLA-compiled function: up to about 3x faster on a GPU, most for the 3D `.h5` networks. Default: disabled | | `--add-objects/--no-add-objects` | Add object definitions to config if missing. Default: enabled | | `--overwrite/--no-overwrite` | Overwrite existing segmentations. Default: disabled | | `--user-id` | User ID for created segmentations. Default: copick | diff --git a/src/copick_easymode/cli/inference.py b/src/copick_easymode/cli/inference.py index a5d7c81..1036c3e 100644 --- a/src/copick_easymode/cli/inference.py +++ b/src/copick_easymode/cli/inference.py @@ -13,7 +13,7 @@ def add_easymode_inference_options(func: click.Command) -> click.Command: """ - Add easymode inference options: --tta, --batch-size, --threshold. + Add easymode inference options: --tta, --batch-size, --threshold, --xla. Args: func (click.Command): The Click command to which the options will be added. @@ -46,6 +46,15 @@ def add_easymode_inference_options(func: click.Command) -> click.Command: show_default=True, help="Probability threshold for binarizing segmentation (0.0-1.0).", ), + click.option( + "--xla/--no-xla", + is_flag=True, + default=False, + show_default=True, + help="Run each model as one XLA-compiled function: up to about 3x faster on a GPU, most for the 3D .h5 " + "networks, after a compile of seconds to tens of seconds per worker. A model that fails to compile falls " + "back to plain inference, but a missing ptxas can abort TensorFlow.", + ), ] for opt in opts: @@ -229,6 +238,7 @@ def easymode( tta: int, batch_size: int, threshold: float, + xla: bool, gpus: str, cpu: bool, max_workers: int, @@ -277,6 +287,9 @@ def easymode( # High quality inference with TTA copick inference easymode -c config.json -m ribosome -t wbp@10.0 --tta 16 --batch-size 2 + # XLA-compiled inference, about 3x faster on a GPU + copick inference easymode -c config.json -m microtubule -t wbp@10.0 --xla + # Skip adding object definitions to config copick inference easymode -c config.json -m ribosome -t wbp@10.0 --no-add-objects @@ -368,7 +381,7 @@ def easymode( logger.info(f"Models: {model_list}") logger.info(f"Tomogram: {tomo_type}@{voxel_size}") logger.info(f"Runs: {run_list if run_list else 'all'}") - logger.info(f"TTA: {tta}, Batch size: {batch_size}, Threshold: {threshold}") + logger.info(f"TTA: {tta}, Batch size: {batch_size}, Threshold: {threshold}, XLA: {xla}") logger.info(f"GPUs: {'none (--cpu)' if cpu else gpus if gpus else 'all of the allocation'}") logger.info(f"Add objects: {add_objects}, Overwrite: {overwrite}") @@ -385,6 +398,7 @@ def easymode( batch_size=batch_size, threshold=threshold, overwrite=overwrite, + xla=xla, gpus=gpus, cpu=cpu, max_workers=max_workers, diff --git a/src/copick_easymode/core/dispatch.py b/src/copick_easymode/core/dispatch.py index 6fa798e..b48be51 100644 --- a/src/copick_easymode/core/dispatch.py +++ b/src/copick_easymode/core/dispatch.py @@ -93,6 +93,7 @@ def _worker_entry(spec: WorkerSpec, **updates) -> dict: "visible_devices": None, "threads": spec.threads, "items": [], + "xla_fallbacks": [], } entry.update(updates) return entry @@ -165,6 +166,8 @@ def run_worker(spec: WorkerSpec, logger=None, prefixed: bool = False) -> dict: visible = os.environ.get(ENV_VISIBLE) stats = new_stats() log.info(f"Starting: {len(spec.runs)} run(s), {spec.threads} thread(s), {ENV_VISIBLE}={visible!r}") + if spec.settings.xla: + log.info("XLA on: each model runs as one XLA-compiled function, compiled at the first call of each input shape") try: distribution = import_easymode(spec.model_dir, lock_path=spec.lock_path, logger=log) go_offline(distribution) @@ -202,6 +205,7 @@ def run_worker(spec: WorkerSpec, logger=None, prefixed: bool = False) -> dict: errors=list(stats["errors"]), visible_devices=visible, items=stats["items"], + xla_fallbacks=list(stats["xla_fallbacks"]), ) @@ -324,6 +328,7 @@ def dispatch_easymode_inference( batch_size: int = 1, threshold: float = 0.5, overwrite: bool = False, + xla: bool = False, gpus: Optional[str] = None, cpu: bool = False, max_workers: Optional[int] = None, @@ -351,6 +356,8 @@ def dispatch_easymode_inference( batch_size: Batch size for inference. threshold: Probability threshold for binarizing a segmentation. overwrite: Overwrite existing segmentations instead of skipping their runs. + xla: Run each model as one XLA-compiled function (see :mod:`copick_easymode.core.xla`); a worker whose model + gave up XLA lists it, with the reason, in its ``xla_fallbacks``. gpus: GPUs to use (see :func:`copick_easymode.core.devices.visible_gpus`); None means all of the allocation. cpu: One worker without a GPU. @@ -389,6 +396,7 @@ def dispatch_easymode_inference( "threshold": threshold, "batch_size": batch_size, "overwrite": overwrite, + "xla": xla, "models_requested": list(models), "models": [], "missing": {}, @@ -435,6 +443,7 @@ def tally() -> None: batch_size=batch_size, threshold=threshold, overwrite=overwrite, + xla=xla, logger=log, ) report.update(user_id=settings.user_id, session_id=settings.session_id) diff --git a/src/copick_easymode/core/inference.py b/src/copick_easymode/core/inference.py index 20380f5..55cd08a 100644 --- a/src/copick_easymode/core/inference.py +++ b/src/copick_easymode/core/inference.py @@ -13,6 +13,7 @@ """ import contextlib +import dataclasses import gc import os import time @@ -179,6 +180,8 @@ class Stages: prepare: Callable[[np.ndarray], Any] infer: Callable[[Any], np.ndarray] finish: Callable[[np.ndarray, Any], np.ndarray] + #: The XLA-compiled model ``infer`` calls (``settings.xla``), whose ``fallback`` says whether it was given up. + xla: Any = None def copick_identity(value: str, what: str = "name") -> str: @@ -207,6 +210,8 @@ class InferenceSettings: batch_size: int = 1 threshold: float = 0.5 overwrite: bool = False + #: Run each model as one XLA-compiled function (:class:`copick_easymode.core.xla.XLAModel`). + xla: bool = False def __post_init__(self): if self.user_id is not None: @@ -258,8 +263,8 @@ def end(self, status: str, message: str = "") -> "_Item": def new_stats() -> dict: - """Processing statistics: counts, error messages and one record per run and model.""" - return {"processed": 0, "skipped": 0, "errors": [], "items": []} + """Processing statistics: counts, error messages, one record per run and model, and each model that gave up XLA.""" + return {"processed": 0, "skipped": 0, "errors": [], "items": [], "xla_fallbacks": []} class _Quiet: @@ -419,8 +424,20 @@ def _enable_memory_growth(tf) -> None: def load_stages(model: "ResolvedModel", settings: InferenceSettings, logger=None) -> Stages: - """Load a resolved model's weights and bind its three stages to ``settings``.""" + """Load a resolved model's weights and bind its three stages to ``settings``. + + With ``settings.xla`` the network is wrapped in an :class:`~copick_easymode.core.xla.XLAModel`, which the stages + call exactly as they call the plain network, and which ``Stages.xla`` keeps for its fallback. + """ log = logger if logger is not None else _Quiet() + + def compiled(network): + if not settings.xla: + return None + from copick_easymode.core.xla import XLAModel + + return XLAModel(network, logger=logger, name=model.feature) + if model.kind == "scnm": from copick_easymode.core.scnm import MAX_TTA, finish_scnm, infer_scnm, load_scnm, prepare_scnm @@ -433,10 +450,14 @@ def load_stages(model: "ResolvedModel", settings: InferenceSettings, logger=None ) if model_tta != settings.tta: log.info(f"TTA {settings.tta} reduced to {model_tta}, the most an .scnm model supports") + xla = compiled(scnm.model) + if xla is not None: + scnm = dataclasses.replace(scnm, model=xla) return Stages( prepare=lambda volume: prepare_scnm(scnm, volume, settings.voxel_size), infer=lambda prepared: infer_scnm(scnm, prepared, model_tta), finish=lambda raw, prepared: finish_scnm(raw, prepared, model_tta), + xla=xla, ) # Imported here rather than first inside a stage thread (it configures TensorFlow at import). @@ -444,10 +465,14 @@ def load_stages(model: "ResolvedModel", settings: InferenceSettings, logger=None model_apix = model.apix if model.apix is not None else 10.0 network = easymode_env.load_network(model.weights) log.info(f"Model loaded from {model.weights}, inference at {model_apix} A/px") + xla = compiled(network) + if xla is not None: + network = xla return Stages( prepare=lambda volume: prepare_tomogram(volume, settings.voxel_size, model_apix), infer=lambda prepared: infer_tomogram(network, prepared, settings.tta, settings.batch_size), finish=lambda raw, prepared: finish_tomogram(raw, prepared, settings.tta), + xla=xla, ) @@ -470,7 +495,8 @@ def segment_with_models( ) -> dict: """Run every resolved model over ``runs``, one model at a time, freeing each before the next loads. - A model that fails to load is recorded as an error and the next one still runs. + A model that fails to load is recorded as an error and the next one still runs. A model that gave up XLA is + recorded in ``stats["xla_fallbacks"]`` with the reason; its runs still count as processed. """ import tensorflow as tf @@ -499,6 +525,8 @@ def segment_with_models( logger=logger, ) finally: + if stages is not None and stages.xla is not None and stages.xla.fallback: + stats.setdefault("xla_fallbacks", []).append({"feature": model.feature, "reason": stages.xla.fallback}) # Free the model's GPU memory after every attempt, a failed load included. stages = None tf.keras.backend.clear_session() @@ -539,6 +567,7 @@ def run_easymode_inference( config_path: Optional[str] = None, logger=None, model_dir: Optional[str] = None, + xla: bool = False, ) -> dict: """ Run easymode inference on copick tomograms in this process. @@ -565,11 +594,13 @@ def run_easymode_inference( logger: Logger instance for output messages. model_dir: easymode model directory to use instead of easymode's setting (in memory only; default: ``$COPICK_EASYMODE_MODEL_DIR``, else the setting). + xla: Run each model as one XLA-compiled function (see :mod:`copick_easymode.core.xla`). Returns: - Dictionary with processing statistics: processed, skipped, errors, and items (one record per run - and model, with its stage timings). A model that is not found, whose object cannot be added or that - fails to load is recorded as an error and the next model still runs. + Dictionary with processing statistics: processed, skipped, errors, items (one record per run and + model, with its stage timings) and xla_fallbacks (each model that gave up XLA, with the reason). A model + that is not found, whose object cannot be added or that fails to load is recorded as an error and the next + model still runs. """ stats = new_stats() settings = copick_settings( @@ -582,6 +613,7 @@ def run_easymode_inference( batch_size=batch_size, threshold=threshold, overwrite=overwrite, + xla=xla, logger=logger, ) diff --git a/src/copick_easymode/core/xla.py b/src/copick_easymode/core/xla.py new file mode 100644 index 0000000..7d408fa --- /dev/null +++ b/src/copick_easymode/core/xla.py @@ -0,0 +1,71 @@ +""" +XLA-compiled inference (``--xla``): a model whose forward pass runs as one XLA-compiled function. + +XLA fuses the memory-bound work between a network's convolutions (normalization, activations, residual adds), which +dominates a 3D UNet's time on large tiles. On real tomograms at TTA 4 this took the easymode microtubule ``.h5`` network +2.9x faster and the actin ``.scnm`` slab model 1.3x faster, with masks that agree with the plain model's to an IoU of +0.9998 or better. It is opt-in: compiling can fail, which falls back to the plain model, but a missing or too old +``ptxas`` can also abort the process inside TensorFlow, which no fallback can catch. +""" + +import logging +import time +from typing import Optional + + +def _shape_label(shapes) -> str: + return ", ".join(f"({', '.join('?' if d is None else str(d) for d in shape)})" for shape in shapes) + + +class XLAModel: + """A Keras model whose forward pass runs as one XLA-compiled function. + + Called like the model, ``m(x, training=False)``, so it drops in wherever the model is called: easymode's tiling of + an ``.h5`` network and the ``.scnm`` inference alike. Every other attribute (``input``, ``input_shape``, + ``output_shape``, ...) is the wrapped model's. Inference only: the compiled function always runs with + ``training=False``. + + The first call for each input shape compiles (tens of seconds for a large network) and is logged with its time; + later calls of that shape reuse it. If a call raises, e.g. because XLA cannot find libdevice, one warning is + logged, the reason is kept in :attr:`fallback`, and this call and every later one run the plain model. + """ + + def __init__(self, model, logger=None, name: str = "model"): + import tensorflow as tf + + self.model = model + self.name = name + #: Why the compiled function was given up (``": "``), or None while it is in use. + self.fallback: Optional[str] = None + self._log = logger if logger is not None else logging.getLogger(__name__) + self._tf = tf + self._compiled = tf.function(lambda x: model(x, training=False), jit_compile=True) + self._shapes = set() + + def __getattr__(self, name): + # Reached only for what the wrapper itself lacks. "model" missing means __init__ has not run (unpickling). + if name == "model": + raise AttributeError(name) + return getattr(self.model, name) + + def __call__(self, x, training=False): + if self.fallback is None: + tf = self._tf + inputs = tf.nest.map_structure(tf.convert_to_tensor, x) + shapes = tuple(tuple(t.shape) for t in tf.nest.flatten(inputs)) + started = time.perf_counter() + try: + output = self._compiled(inputs) + except Exception as exc: + reason = (str(exc).strip().splitlines() or [""])[0][:200] + self.fallback = f"{type(exc).__name__}: {reason}" + self._log.warning(f"XLA failed for {self.name} ({self.fallback}); continuing without XLA") + else: + if shapes not in self._shapes: + self._shapes.add(shapes) + self._log.info( + f"XLA compiled {self.name} for input {_shape_label(shapes)}: first call " + f"{time.perf_counter() - started:.1f} s", + ) + return output + return self.model(x, training=False) diff --git a/tests/test_cli.py b/tests/test_cli.py index 1300cea..fd3db39 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -78,6 +78,7 @@ def test_command_parses_lists_and_delegates(cli_env, tmp_path): "2", "--threshold", "0.25", + "--xla", "--gpus", "0,1", "--max-workers", @@ -107,6 +108,7 @@ def test_command_parses_lists_and_delegates(cli_env, tmp_path): batch_size=2, threshold=0.25, overwrite=True, + xla=True, gpus="0,1", cpu=False, max_workers=2, @@ -120,6 +122,17 @@ def test_command_parses_lists_and_delegates(cli_env, tmp_path): ) +def test_xla_is_off_unless_asked_for(cli_env): + config_path, dispatch = cli_env + dispatch.side_effect = None + dispatch.return_value = _report() + + for flag, expected in (([], False), (["--no-xla"], False), (["--xla"], True)): + result = CliRunner().invoke(easymode, ["-c", str(config_path), "-m", "ribosome", "-t", "wbp@10", *flag]) + assert result.exit_code == 0, result.output + assert dispatch.call_args.kwargs["xla"] is expected + + def test_a_failed_report_exits_nonzero_and_lists_the_errors(cli_env): config_path, dispatch = cli_env dispatch.side_effect = None diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index 46559e5..4eac8aa 100644 --- a/tests/test_dispatch.py +++ b/tests/test_dispatch.py @@ -78,6 +78,47 @@ def test_a_cpu_run_reports_what_it_did_and_exits_zero(tmp_path, isolated_env): assert json.loads(report_path.read_text())["skipped"] == 3 +def test_xla_reaches_the_worker_and_the_report_records_a_fallback(tmp_path, isolated_env, monkeypatch): + from copick_easymode.core import xla + + config, _ = make_project(tmp_path, RUNS) + model_dir = make_model_dir(tmp_path) + report_path = tmp_path / "easymode.json" + + plain = _invoke(config, model_dir, report_path, "-m", "tiny", "--cpu") + assert plain.exit_code == 0, plain.output + report = json.loads(report_path.read_text()) + assert report["xla"] is False and report["workers"][0]["xla_fallbacks"] == [] + + # XLA on the CPU: the tiny model compiles, and every run is segmented through it. + result = _invoke(config, model_dir, report_path, "-m", "tiny", "--cpu", "--xla", "--overwrite") + assert result.exit_code == 0, result.output + report = json.loads(report_path.read_text()) + assert report["xla"] is True and report["processed"] == 3 + assert report["workers"][0]["xla_fallbacks"] == [] + + # A compiled call that raises: the worker continues with the plain model and its report entry says why. + created = [] + init = xla.XLAModel.__init__ + + def broken(x): + raise RuntimeError("ptxas too old") + + def failing_init(self, *args, **kwargs): + init(self, *args, **kwargs) + self._compiled = broken + created.append(self) + + monkeypatch.setattr(xla.XLAModel, "__init__", failing_init) + result = _invoke(config, model_dir, report_path, "-m", "tiny", "--cpu", "--xla", "--overwrite") + assert result.exit_code == 0, result.output + report = json.loads(report_path.read_text()) + assert report["status"] == "complete" and report["processed"] == 3 and len(created) == 1 + [worker] = report["workers"] + assert worker["xla_fallbacks"] == [{"feature": "tiny", "reason": "RuntimeError: ptxas too old"}] + assert all(_segmented(config, run) for run in RUNS) + + def test_a_missing_model_fails_before_any_worker(tmp_path, isolated_env): config, _ = make_project(tmp_path, RUNS) model_dir = make_model_dir(tmp_path) diff --git a/tests/test_xla.py b/tests/test_xla.py new file mode 100644 index 0000000..c1790a3 --- /dev/null +++ b/tests/test_xla.py @@ -0,0 +1,72 @@ +"""The XLA wrapper: called like the model, attributes passed through, the plain model after a failed call.""" + +import dataclasses +import logging + +import numpy as np +import pytest + +keras = pytest.importorskip("keras") +pytest.importorskip("tensorflow") + +from conftest import make_scnm # noqa: E402 + +from copick_easymode.core.scnm import load_scnm # noqa: E402 +from copick_easymode.core.xla import XLAModel # noqa: E402 + + +def _slab_model(): + inp = keras.Input(shape=(None, None, 4, 1)) + x = keras.layers.Conv3D(2, 3, padding="same")(inp) + return keras.Model(inp, keras.layers.Conv3D(1, 1, activation="sigmoid")(x)) + + +def test_the_wrapper_passes_attributes_through_to_the_model(tmp_path): + model = _slab_model() + wrapped = XLAModel(model, name="slab") + assert wrapped.input is model.input + assert (wrapped.input_shape, wrapped.output_shape) == (model.input_shape, model.output_shape) + assert wrapped.layers == model.layers + with pytest.raises(AttributeError): + wrapped.no_such_attribute # noqa: B018 + + # An .scnm model reads its kind from the network's shapes, through the wrapper alike. + scnm = load_scnm(make_scnm(tmp_path / "models")) + compiled = dataclasses.replace(scnm, model=XLAModel(scnm.model)) + assert (compiled.is_slab, compiled.dimensionality) == (scnm.is_slab, scnm.dimensionality) == (False, 2.5) + + +def test_the_compiled_call_returns_what_the_model_does(caplog): + model = _slab_model() + wrapped = XLAModel(model, logger=logging.getLogger("xla-test"), name="slab") + tile = np.random.default_rng(0).normal(size=(1, 16, 24, 4, 1)).astype(np.float32) + + with caplog.at_level(logging.INFO, logger="xla-test"): + first = wrapped(tile, training=False).numpy() + wrapped(tile, training=False) + np.testing.assert_allclose(first, model(tile, training=False).numpy(), atol=1e-5) + assert wrapped.fallback is None + # One compile per input shape, logged once. + assert [r.getMessage().split(":")[0] for r in caplog.records] == ["XLA compiled slab for input (1, 16, 24, 4, 1)"] + + +def test_a_failing_compiled_call_falls_back_to_the_model_for_good(caplog): + model = _slab_model() + wrapped = XLAModel(model, logger=logging.getLogger("xla-test"), name="slab") + calls = [] + + def broken(x): + calls.append(x) + raise RuntimeError("libdevice not found at ./libdevice.10.bc\nmore detail") + + wrapped._compiled = broken + tile = np.random.default_rng(1).normal(size=(1, 16, 16, 4, 1)).astype(np.float32) + + with caplog.at_level(logging.INFO, logger="xla-test"): + outputs = [wrapped(tile, training=False).numpy() for _ in range(3)] + expected = model(tile, training=False).numpy() + assert all(np.allclose(out, expected) for out in outputs) + assert len(calls) == 1 # given up after the first failure + assert wrapped.fallback == "RuntimeError: libdevice not found at ./libdevice.10.bc" + [warning] = [r for r in caplog.records if r.levelno == logging.WARNING] + assert "continuing without XLA" in warning.getMessage() and "libdevice" in warning.getMessage() From 7054da9e7385b0f03b63844064488cd5a7149b86 Mon Sep 17 00:00:00 2001 From: uermel Date: Wed, 21 Oct 2026 10:43:44 -0700 Subject: [PATCH 2/2] feat: support easymode 1.2.6 easymode 1.2.6 and ais-cryoet 1.2.40 release the --xla options this package mirrors (mgflast/easymode#37, mgflast/Ais#16). The .h5 path uses only _pad_volume, _segment_tomogram_instance and load_model, which are unchanged, so the bump is the supported version and commit, the tests and the docs. The .scnm port follows ais-cryoet 1.2.40: its slab --stride keeps the depth // 2 default, which easymode does not override, and under XLA a short last 2.5D batch is padded with zeros to the batch size so the compiled function keeps one input shape. The docs also drop the claim that easymode is not on PyPI, and name its new metadata pins. Co-Authored-By: Claude Opus 5.5 --- README.md | 28 ++++++++++++------------ pyproject.toml | 6 ++--- src/copick_easymode/core/easymode_env.py | 21 +++++++++--------- src/copick_easymode/core/scnm.py | 15 ++++++++++--- src/copick_easymode/core/xla.py | 5 +++++ tests/test_easymode_env.py | 6 ++--- tests/test_xla.py | 14 +++++++++++- 7 files changed, 61 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index 82f29ab..cb93aa2 100644 --- a/README.md +++ b/README.md @@ -10,12 +10,11 @@ If you use this plugin, please cite the easymode preprint (see [Citation](#citat ## Installation -easymode is not published to PyPI (only older `0.0.x` releases are), so it is installed -from GitHub. Its packaging metadata also (incorrectly) pins `numpy<2` / `tensorflow<2.12`, which -conflicts with copick's `numpy>=2` — even though easymode runs fine on `numpy>=2`. To avoid that -conflict, install copick-easymode and its dependencies **first** (this brings in copick and -`numpy>=2`), then the supported easymode release (1.2.5) from its source commit with `--no-deps`, -so its bad pins are ignored: +easymode's packaging metadata pins `numpy<2` and `tensorflow<2.16` (and `ais-cryoet`, which needs +Keras 2), which conflicts with copick's `numpy>=2` — even though easymode runs fine on `numpy>=2` and +Keras 3. To avoid that conflict, install copick-easymode and its dependencies **first** (this brings +in copick and `numpy>=2`), then the supported easymode release (1.2.6) from its source commit with +`--no-deps`, so its pins are ignored: ```bash git clone https://github.com/copick/copick-easymode.git @@ -24,17 +23,17 @@ cd copick-easymode # 1. Install copick-easymode + dependencies (copick, numpy>=2, tensorflow>=2.16, easymode's runtime deps) pip install -e . -# 2. Install easymode 1.2.5 WITHOUT dependency resolution. +# 2. Install easymode 1.2.6 WITHOUT dependency resolution. # --no-deps keeps your numpy>=2 stack intact, and also upgrades over any older easymode -# (e.g. a 0.0.x already installed from PyPI). -pip install --no-deps "easymode @ git+https://github.com/mgflast/easymode.git@544572cfc071869113eb0567ef607791946a81c1" +# already installed. +pip install --no-deps "easymode @ git+https://github.com/mgflast/easymode.git@00ab9035b19ad30dca46fbe97bc77e6f11a2d981" ``` Verify the install: ```bash python -c "import numpy, easymode, importlib.metadata as m; print('numpy', numpy.__version__, '| easymode', m.version('easymode'))" -# expected: numpy 2.x | easymode 1.2.5 +# expected: numpy 2.x | easymode 1.2.6 ``` copick-easymode imports easymode in one place, `copick_easymode.core.easymode_env`. A missing easymode, or one @@ -112,7 +111,8 @@ in the command's own process. - **Config writes.** With `--add-objects`, only the parent process adds object definitions and saves the config, before the workers start; workers re-open the project from the config file and never write it. - **XLA.** `--xla` runs each model's forward pass as one XLA-compiled function (`tf.function(jit_compile=True)`), - which fuses the memory-bound work between the convolutions. Each worker compiles at the first call of each input + which fuses the memory-bound work between the convolutions, like `easymode segment --xla` (easymode 1.2.6) and + `ais segment --xla` (ais-cryoet 1.2.40). Each worker compiles at the first call of each input shape, which takes tens of seconds and is logged as `XLA compiled for input (...)`, and reuses it for every later tile of that shape. On 10521 tomograms of 320 x 772 x 548 voxels, at TTA 4 on one L40S with TensorFlow 2.21, it took the microtubule `.h5` network from 310 s to 107 s per tomogram, 2.9x, and the actin `.scnm` slab model from @@ -139,7 +139,7 @@ the end (`"complete"` or `"failed"`), whatever the outcome: "xla": false, "models_requested": ["microtubule"], "models": [{"feature": "microtubule", "object": "microtubule", "weights": "/shared/easymode/models/microtubule_sv3-3d.h5", "kind": "h5", "tag": "sv3-3d", "timestamp": "20260916223602", "bytes": 515770304, "apix": 10.0}], - "missing": {}, "model_directory": "/shared/easymode", "online": true, "writable": true, "easymode_version": "1.2.5", + "missing": {}, "model_directory": "/shared/easymode", "online": true, "writable": true, "easymode_version": "1.2.6", "cpu": false, "devices": ["0", "1"], "threads_per_worker": 8, "workers": [ {"index": 0, "gpu": "0", "runs": ["TS_01", "TS_03"], "exitcode": 0, "seconds": 412.3, "processed": 2, "skipped": 0, @@ -175,7 +175,7 @@ easymode publishes them in two formats, and both run in the same process and on - **`.h5` (3D)**: `ribosome`, `microtubule` and `tric`, loaded through `easymode.core.distribution.load_model`. - **`.scnm` (Ais 2D-engine models, 2.5D slice or 3D slab)**: every other feature, e.g. `actin`, `membrane`, `proteasome`, `atp_synthase`, `cytoplasm`, `nucleus`. `copick_easymode.core.scnm` loads them in Keras 3 and runs - a port of `ais segment`'s inference (Ais 1.2.38), so neither Ais nor its TensorFlow 2.11 is needed. `--tta` above 8 + a port of `ais segment`'s inference (Ais 1.2.40), so neither Ais nor its TensorFlow 2.11 is needed. `--tta` above 8 is reduced to 8 for these, the most Ais supports. copick does not allow underscores in object names, so a feature such as `atp_synthase` is written to copick as @@ -236,7 +236,7 @@ ready, so a run whose inference fails leaves it untouched. - copick >= 1.24.1 - numpy >= 2.0.2 - TensorFlow >= 2.16 -- easymode 1.2.5, commit `544572cfc071869113eb0567ef607791946a81c1` (installed separately; see [Installation](#installation)) +- easymode 1.2.6, commit `00ab9035b19ad30dca46fbe97bc77e6f11a2d981` (installed separately; see [Installation](#installation)) ## Citation diff --git a/pyproject.toml b/pyproject.toml index f1f6c94..7c142f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,9 +16,9 @@ dependencies = [ "numpy>=2.0.2", "scipy>=1.10.0", "tensorflow>=2.16", - # easymode itself is installed separately from GitHub with `--no-deps` (see README): it is - # not on PyPI, and its published metadata wrongly pins numpy<2 / tensorflow<2.12, which - # conflicts with copick (numpy>=2). easymode runs fine on numpy>=2, so we install copick + + # easymode itself is installed separately from GitHub with `--no-deps` (see README): its + # published metadata pins numpy<2 / tensorflow<2.16 / ais-cryoet (Keras 2), which conflicts + # with copick (numpy>=2). easymode runs fine on numpy>=2, so we install copick + # numpy>=2 first and declare easymode's runtime deps here so the env is complete. "mrcfile", "huggingface_hub", diff --git a/src/copick_easymode/core/easymode_env.py b/src/copick_easymode/core/easymode_env.py index 4d112a4..0f74764 100644 --- a/src/copick_easymode/core/easymode_env.py +++ b/src/copick_easymode/core/easymode_env.py @@ -2,12 +2,12 @@ The one place easymode is imported: its configuration, imported safely by concurrent processes, its models resolved once, and the private inference functions this package runs. -**The compatibility boundary.** easymode is not on PyPI, and copick-easymode uses parts of it that are not a -public API: the registry-based model distribution (``MODEL_CACHE_DIR``, ``REGISTRY_CACHE``, the ``_online`` -flag) and the private ``_pad_volume`` and ``_segment_tomogram_instance`` of ``easymode.segmentation.inference``. -No other module imports easymode. What is used is checked when it is imported, and a missing or incompatible -easymode raises :class:`EasymodeCompatibilityError`, which names the supported release and the command that -installs it (:data:`EASYMODE_INSTALL`). +**The compatibility boundary.** copick-easymode uses parts of easymode that are not a public API: the registry-based +model distribution (``MODEL_CACHE_DIR``, ``REGISTRY_CACHE``, the ``_online`` flag) and the private ``_pad_volume`` +and ``_segment_tomogram_instance`` of ``easymode.segmentation.inference``. No other module imports easymode. What is +used is checked when it is imported, and a missing or incompatible easymode raises +:class:`EasymodeCompatibilityError`, which names the supported release and the command that installs it +(:data:`EASYMODE_INSTALL`). **The import race.** ``easymode/core/config.py`` runs ``parse_settings()`` at import, and that function *rewrites* ``~/easymode/settings.txt`` (opening for writing truncates it, then the JSON is dumped). Processes @@ -45,10 +45,11 @@ except ImportError: # pragma: no cover - not POSIX: no cross-process lock fcntl = None -#: The easymode release this package is tested against, and its source commit (easymode is not on PyPI). -EASYMODE_VERSION = "1.2.5" -EASYMODE_COMMIT = "544572cfc071869113eb0567ef607791946a81c1" -#: easymode's own metadata pins numpy<2 and tensorflow<2.12, so it is installed without its dependencies. +#: The easymode release this package is tested against, and its source commit (tag ``v1.2.6``). +EASYMODE_VERSION = "1.2.6" +EASYMODE_COMMIT = "00ab9035b19ad30dca46fbe97bc77e6f11a2d981" +#: easymode's own metadata pins numpy<2, tensorflow<2.16 and ais-cryoet (Keras 2), so it is installed without its +#: dependencies. EASYMODE_INSTALL = ( f"python -m pip install --no-deps 'easymode @ git+https://github.com/mgflast/easymode.git@{EASYMODE_COMMIT}'" ) diff --git a/src/copick_easymode/core/scnm.py b/src/copick_easymode/core/scnm.py index e5248c0..724a627 100644 --- a/src/copick_easymode/core/scnm.py +++ b/src/copick_easymode/core/scnm.py @@ -9,9 +9,10 @@ model loads in Keras 3 once the ``groups`` argument Keras 2 wrote into its transposed convolutions is dropped. So this module loads it here and ports Ais's inference: ``_preprocess_tomo``, ``_segmentation_thread`` and ``_infer_slab`` of -``Ais/core/cli_fn.py`` (ais-cryoet 1.2.38, Mart G. F. So-Last, GPL-3.0 like this package), +``Ais/core/cli_fn.py`` (ais-cryoet 1.2.40, Mart G. F. So-Last, GPL-3.0 like this package), with the defaults easymode's ``ais segment`` call leaves in place (whole Z range, batch 1, -512 px XY tiles with a 64 px discarded border, no post-blur). +a slab Z stride of half the slab depth, 512 px XY tiles with a 64 px discarded border, no +post-blur). The result is the same 0-1 probability map ``segment_tomogram_from_array`` returns for an ``.h5`` model, so thresholding and everything downstream is shared. Like Ais, inference is split @@ -29,6 +30,8 @@ import numpy as np +from copick_easymode.core.xla import XLAModel + SCNM_SUFFIX = ".scnm" #: Ais's test-time augmentation: four 90-degree rotations, then the same four flipped. MAX_TTA = 8 @@ -190,9 +193,15 @@ def _slice_input(vi, j, depth, dimensionality): def _infer_slices(scnm: ScnmModel, vi, batch_size): si = np.zeros_like(vi, dtype=np.float32) n = vi.shape[0] + # Under XLA a short last batch is padded with zeros to batch_size, as Ais does, so the compiled function keeps one + # input shape; the padded rows are never read back. At batch 1, which .scnm inference runs at, no batch is short. + pad = isinstance(scnm.model, XLAModel) for start in range(0, n, batch_size): js = list(range(start, min(start + batch_size, n))) - out = _call(scnm.model, np.stack([_slice_input(vi, j, scnm.depth, scnm.dimensionality) for j in js])) + batch = np.stack([_slice_input(vi, j, scnm.depth, scnm.dimensionality) for j in js]) + if pad and len(js) < batch_size: + batch = np.concatenate([batch, np.zeros((batch_size - len(js),) + batch.shape[1:], dtype=batch.dtype)]) + out = _call(scnm.model, batch) for i, j in enumerate(js): si[j] = np.squeeze(out[i]) return si diff --git a/src/copick_easymode/core/xla.py b/src/copick_easymode/core/xla.py index 7d408fa..13fc972 100644 --- a/src/copick_easymode/core/xla.py +++ b/src/copick_easymode/core/xla.py @@ -6,6 +6,11 @@ 2.9x faster and the actin ``.scnm`` slab model 1.3x faster, with masks that agree with the plain model's to an IoU of 0.9998 or better. It is opt-in: compiling can fail, which falls back to the plain model, but a missing or too old ``ptxas`` can also abort the process inside TensorFlow, which no fallback can catch. + +It is the option easymode 1.2.6 has as ``easymode segment --xla`` and ais-cryoet 1.2.40 as ``ais segment --xla``. This +package keeps its own wrapper rather than easymode's ``XLAModel``: the ``.scnm`` inference reads the network's shapes +through it, each compiled shape is logged, and a fallback's reason goes into the report. Ais's ``_XLAModel`` is never +reached, since ``.scnm`` models run here rather than through ``ais segment``. """ import logging diff --git a/tests/test_easymode_env.py b/tests/test_easymode_env.py index 79db900..7bbe1e1 100644 --- a/tests/test_easymode_env.py +++ b/tests/test_easymode_env.py @@ -24,7 +24,7 @@ def _module(name, **attributes): def _distribution(**overrides): - """What easymode 1.2.5's ``easymode.core.distribution`` offers this package.""" + """What easymode 1.2.6's ``easymode.core.distribution`` offers this package.""" attributes = { "MODEL_CACHE_DIR": "/models", "REGISTRY_CACHE": "/models/registry.json", @@ -43,7 +43,7 @@ def fake_easymode(monkeypatch, isolated_env): """Install a fake easymode whose distribution is ``install(**overrides)``.""" for name in [name for name in sys.modules if name == "easymode" or name.startswith("easymode.")]: monkeypatch.delitem(sys.modules, name) - config = _module("easymode.core.config", settings={"MODEL_DIRECTORY": "/models"}, version="1.2.5") + config = _module("easymode.core.config", settings={"MODEL_DIRECTORY": "/models"}, version="1.2.6") def install(**overrides): distribution = _distribution(**overrides) @@ -116,7 +116,7 @@ def test_a_compatible_easymode_is_used_through_the_boundary(fake_easymode): distribution = fake_easymode() assert easymode_env.import_easymode() is distribution - assert easymode_env.easymode_version() == "1.2.5" + assert easymode_env.easymode_version() == "1.2.6" assert easymode_env.load_network("/models/ribosome.h5") == {"weights": "/models/ribosome.h5"} easymode_env.go_offline(distribution) assert distribution._online is False diff --git a/tests/test_xla.py b/tests/test_xla.py index c1790a3..c94abc3 100644 --- a/tests/test_xla.py +++ b/tests/test_xla.py @@ -11,7 +11,7 @@ from conftest import make_scnm # noqa: E402 -from copick_easymode.core.scnm import load_scnm # noqa: E402 +from copick_easymode.core.scnm import _infer_slices, load_scnm # noqa: E402 from copick_easymode.core.xla import XLAModel # noqa: E402 @@ -70,3 +70,15 @@ def broken(x): assert wrapped.fallback == "RuntimeError: libdevice not found at ./libdevice.10.bc" [warning] = [r for r in caplog.records if r.levelno == logging.WARNING] assert "continuing without XLA" in warning.getMessage() and "libdevice" in warning.getMessage() + + +def test_a_short_last_batch_is_padded_so_one_shape_compiles(tmp_path, caplog): + scnm = load_scnm(make_scnm(tmp_path / "models")) + compiled = dataclasses.replace(scnm, model=XLAModel(scnm.model, logger=logging.getLogger("xla-test"), name="tiny")) + vi = np.random.default_rng(2).normal(size=(3, 16, 24)).astype(np.float32) + + with caplog.at_level(logging.INFO, logger="xla-test"): + padded = _infer_slices(compiled, vi, batch_size=2) + np.testing.assert_allclose(padded, _infer_slices(scnm, vi, batch_size=2), atol=1e-5) + # Slices 0-1, then slice 2 padded to a batch of two: one input shape, one compile. + assert [r.getMessage().split(":")[0] for r in caplog.records] == ["XLA compiled tiny for input (2, 16, 24, 4)"]