17.7. Extending The Pipeline#

Pipeline Steps describes the built-in extension path: a new operation is implemented in pycsamt.emtools, given a StepSpec entry in pycsamt/pipeline/_registry.py, and reviewed with tests and documentation before it ships in a pyCSAMT release. That path is intentional for anything that should become part of the shared scientific catalogue.

It is the wrong path for a step that only makes sense for one project, one instrument vendor, or one organisation’s internal correction. For that case, pyCSAMT accepts a second kind of step: a pipeline plugin, registered at runtime with register_step and never touching pyCSAMT’s own source tree. Once registered, a plugin step behaves exactly like a built-in one – it can be used from Step, Pipeline, presets, configuration files, and the CLI.

17.7.1. Two Extension Paths#

Built-in step

Plugin step

Where it lives

pycsamt/pipeline/_registry.py

Any installed package, or a project script

How it ships

Reviewed pull request, released with pyCSAMT

pip install of a third-party package, or a runtime call

Registered by

The literal StepSpec list in the registry module

register_step from Python, or a plugin entry-point group callable

StepSpec.origin

"builtin"

"plugin"

Good fit for

Operations useful across surveys and organisations

Site-specific corrections, vendor formats, internal QC rules

17.7.2. Registering A Step#

register_step takes a fully-built StepSpec and inserts it into the same step registry that the 47 built-in steps live in:

 1>>> from pycsamt.pipeline import StepSpec, register_step, lookup_step
 2>>> def scale_amplitude(sites, factor: float = 2.0):
 3...     """Multiply every impedance value in *sites* by *factor* (toy transform)."""
 4...     return sites
 5...
 6>>> spec = register_step(
 7...     StepSpec(
 8...         code="DEMO001",
 9...         name="scale_amplitude",
10...         label="Demo Amplitude Scale",
11...         category="demo",
12...         override_fn=scale_amplitude,
13...         defaults={"factor": 2.0},
14...     )
15... )
16>>> spec.origin
17'plugin'
18>>> lookup_step("DEMO001") is spec
19True

origin is always stamped "plugin" by register_step itself, regardless of what the caller passed – a plugin author cannot accidentally label a step "builtin". Once registered, the step is usable exactly like any other:

1>>> from pycsamt.pipeline import Step
2>>> Step("DEMO001", factor=3.0).params
3{'factor': 3.0}

Registration is a mutation of one process-wide registry, the same way pycsamt.forward.maxwell’s backend registry is process-wide. A second registration under the same code or name fails rather than silently swapping the implementation, because a silent swap could change numerical behaviour somewhere else in a long-running session:

 1>>> register_step(
 2...     StepSpec(
 3...         code="DEMO001",
 4...         name="scale_amplitude",
 5...         label="duplicate",
 6...         category="demo",
 7...         override_fn=scale_amplitude,
 8...     )
 9... )
10Traceback (most recent call last):
11...
12ValueError: Pipeline step code='DEMO001' or name='scale_amplitude' is already registered.  Pass replace_existing=True to overwrite it.

Passing replace_existing=True allows a deliberate overwrite, including of a built-in step – useful for a site that wants to patch one operation’s defaults without forking pyCSAMT. Formally, if \(R\) is the current registry mapping codes to specs and a caller registers spec \(s\) under code \(c\),

\[\begin{split}R' = \begin{cases} R \cup \{c: s\}, & c \notin R,\\ R \setminus \{c: R[c]\} \cup \{c: s\}, & c \in R \text{ and } \texttt{replace\_existing},\\ \text{raise } \texttt{ValueError}, & c \in R \text{ and not } \texttt{replace\_existing}. \end{cases}\end{split}\]

register_step also validates the spec before inserting it: it calls spec.get_fn() once, so a typo’d module path fails immediately at registration time rather than three steps into a pipeline run:

 1>>> register_step(
 2...     StepSpec(
 3...         code="DEMO_BAD",
 4...         name="demo_bad",
 5...         label="Bad",
 6...         category="demo",
 7...         mod="my_package.pipeline_steps",
 8...         fn_name="not_a_real_function",
 9...     )
10... )
11Traceback (most recent call last):
12...
13ModuleNotFoundError: No module named 'my_package'

The registry is left untouched when validation fails; DEMO_BAD never appears in lookup_step or list_steps.

17.7.3. Removing A Step#

unregister_step reverses a registration, by code or by name:

1>>> from pycsamt.pipeline import unregister_step
2>>> unregister_step("DEMO001")
3>>> lookup_step("DEMO001")
4Traceback (most recent call last):
5...
6KeyError: "No pipeline step found for 'DEMO001'.  Call list_steps() to see all available steps."

unregister_step("SOME_CODE", missing_ok=True) is a safe no-op when the step was never registered – the form to reach for in a test fixture’s teardown, so a test does not fail just because an earlier test already cleaned up the same code.

17.7.4. Packaging A Plugin#

A plugin package announces itself through the pycsamt.pipeline.steps plugin entry-point group. The package’s own pyproject.toml points at a zero-argument callable that performs the registration:

1[project.entry-points."pycsamt.pipeline.steps"]
2demo = "demo_pipe_plugin:register"

demo is the plugin’s display name – it is what shows up in pycsamt pipe plugins output. demo_pipe_plugin:register is module:callable, resolved with importlib.metadata the same way console_scripts entry points resolve for the pycsamt command itself. The referenced module calls register_step for whatever it contributes:

 1"""Toy pyCSAMT pipeline plugin used to smoke-test the plugin discovery API."""
 2
 3from pycsamt.pipeline import StepSpec, register_step
 4
 5
 6def scale_amplitude(sites, factor: float = 2.0):
 7    """Multiply every impedance value in *sites* by *factor* (toy transform)."""
 8    return sites
 9
10
11def register() -> None:
12    register_step(
13        StepSpec(
14            code="DEMO001",
15            name="scale_amplitude",
16            label="Demo Amplitude Scale",
17            category="demo",
18            mod="demo_pipe_plugin",
19            fn_name="scale_amplitude",
20            defaults={"factor": 2.0},
21        )
22    )

Installing the package with pip install demo-pipe-plugin (or pip install -e . during development) is enough for importlib.metadata to see the entry point. It is not, by itself, enough for the step to appear in STEP_REGISTRY – see the next section.

17.7.5. Explicit Discovery#

Installing a plugin package never runs its register function by itself. Plugin discovery – the act of scanning the entry-point group and calling every callable found there – only happens when pycsamt.pipeline.discover_plugins is called:

1>>> from pycsamt.pipeline import discover_plugins
2>>> discover_plugins()
3[PluginLoadResult(name='demo', ok=True, error=None)]

This is deliberate, for two reasons. First, running arbitrary third-party code merely because a package happens to be installed is a real security surface pyCSAMT does not want to open by default. Second, the scan itself is not free: importlib.metadata.entry_points() has to enumerate every installed distribution to find the ones that declare pycsamt.pipeline.steps, and on the Anaconda environment used to build this page – several hundred installed packages – a single scan measured several seconds. Paying that cost on every import pycsamt.pipeline, or on every single pycsamt pipe invocation, would make the common case (no plugins at all) noticeably slower for everyone in order to serve a case most users never hit.

A plugin that fails to load is reported, not fatal to the others – and this demo plugin’s own register() is itself a case in point. It calls register_step unconditionally with no replace_existing, so calling discover_plugins a second time in the same process reports it as failed, not because anything broke, but because DEMO001 is already registered from the first call:

1>>> results = discover_plugins()
2>>> [(r.name, r.ok) for r in results]
3[('demo', False)]

This is why pyCSAMT’s own CLI calls discover_plugins at most once per process – either inside pycsamt pipe plugins or from the --with-plugins group flag, never both (see CLI Plugin Discovery below). A plugin author who expects register() to be called more than once per process should make it tolerate that, either by checking spec.code not in STEP_REGISTRY first or by passing replace_existing=True. Pass on_error="raise" to stop at the first failure instead of collecting a PluginLoadResult for it – useful in a CI job that should fail loudly on a broken plugin rather than merely warn.

17.7.6. CLI Plugin Discovery#

pycsamt pipe plugins always discovers, because discovery is the entire point of that command:

1pycsamt pipe plugins
1Discovered 1 pipeline plugin(s):
2  demo                     ok
3
4Registered 1 plugin step(s):
5  DEMO001    scale_amplitude              [demo]  Demo Amplitude Scale

Every other pycsamt pipe subcommand leaves discovery off by default, for the same latency reason explained above. Trying to use a plugin step’s code without opting in fails with a hint toward the fix:

1pycsamt pipe run --steps DEMO001 --survey ./edis/ --dry-run
1Error: Invalid value for '--steps': Unknown step 'DEMO001'.  Run  pycsamt pipe steps  to see all available steps.  If this is a plugin step, pass  pipe --with-plugins  or run  pycsamt pipe plugins  first.

--with-plugins is a flag on the pipe group itself, not on individual subcommands such as run – it has to run before Click parses --steps, because --steps is validated against the registry at parse time. Setting the PYCSAMT_PIPELINE_LOAD_PLUGINS environment variable has the same effect without retyping the flag on every invocation. With discovery opted into, the same plugin step resolves in one shot:

1pycsamt pipe --with-plugins run --steps DEMO001 --survey ./edis/ --dry-run
1Pipeline  'cli_pipeline'  ───────────────────────────────────────────────  1 step
2  ( 1) scale_amplitude  [DEMO001]  Demo Amplitude Scale  factor=2.0
3────────────────────────────────────────────────────────────────────────────────
4
5Sites   : 25
6Steps   : 1
7Out dir : pipe_results (default)
8
9Dry run — no processing performed.

Captured against a real 25-site WILLY survey directory. pycsamt pipe steps --info (also under --with-plugins) formats a plugin step the same way it formats a built-in one:

1pycsamt pipe --with-plugins steps --info DEMO001
1DEMO001  Demo Amplitude Scale  [demo]
2  name     : scale_amplitude
3  function : demo_pipe_plugin.scale_amplitude
4  defaults : factor=2.0
5  qc plots : —
6  returns  : Sites (transform)

17.7.7. Plugin Origin#

Every StepSpec carries an origin field so built-in and plugin steps can be told apart programmatically:

1>>> from pycsamt.pipeline import STEP_REGISTRY
2>>> sorted(s.code for s in STEP_REGISTRY.values() if s.origin == "plugin")
3['DEMO001']
4>>> lookup_step("NR001").origin
5'builtin'

A script that wants to report only what a survey’s specific pipeline configuration relies on beyond the shared catalogue can filter on this field rather than maintaining a separate list of “steps I added.”

17.7.8. First-Party AI Steps#

pyCSAMT ships one opt-in step of its own, built on exactly the mechanism this page teaches for third-party plugins: AI001 / audit_survey, wrapping pycsamt.ai.domain_gap.audit.audit_survey() as a diagnostic pipeline step. It is not a built-in the way the other 50 steps are, and it is not discovered through the entry-point mechanism above either – it is registered directly by pycsamt.pipeline.register_ai_steps():

1>>> from pycsamt.pipeline import register_ai_steps, lookup_step
2>>> registered = register_ai_steps()
3>>> [s.code for s in registered]
4['AI001']
5>>> lookup_step("AI001").origin
6'plugin'

Why this needs its own opt-in, distinct from register_step collisions or entry-point scanning: resolving AI001 imports pycsamt.ai, and pycsamt.ai’s own package __init__ eagerly imports pycsamt.ai.nets.drcnn, which imports torch at module level (a deliberate choice there so its classes stay picklable for checkpointing – see that module’s own comment). That is a real, multi-second cost the CLI must not force on every pipeline user merely for pycsamt pipe run. The pipe group’s --with-ai-steps flag (or the PYCSAMT_PIPELINE_LOAD_AI_STEPS environment variable) opts in the same way --with-plugins does, and for the same structural reason: it is a group-level flag so it runs before --steps is parsed and validated:

1pycsamt pipe --with-ai-steps steps --info AI001
1AI001  AI Domain-Gap Survey Audit  [ai]
2  name     : audit_survey
3  function : pycsamt.pipeline.ai_steps.qc_audit_survey
4  defaults : —
5  qc plots : —
6  returns  : Sites unchanged (diagnostic)

Captured running the step for real against the 25-station WILLY L22 line:

1pycsamt pipe --with-ai-steps run --steps AI001 \
2    --survey data/AMT/WILLY_DATA/L22PLT --out results/line22_audit
 1Survey audit (generated 2026-08-14T08:05:48Z)
 2  Stations: 25 input, 25 included, 0 excluded
 3  Frequency grid: matched
 4  Impedance coverage: 100.0%
 5  Frequency range: 1.008-10400 Hz
 6  Declared error / |Z|: p05=0.0214, p50=0.0839, p95=0.3492
 7  Station spacing (m): min=59.9, median=100.0, max=120.5
 8  CRS declared: False
 9  Elevation coverage: 100.0%
10  Dimensionality: n=1325, 1D=9.5%, 2D=18.0%, 3D=72.5%
11  Strike (consensus): -29.0 deg (IQR 118.9 deg)
12  Static shift log10 sigma: 0.1219
13  Distortion sigma: gain(log10)=0.0000, twist_deg=14.63, shear=0.3680, anisotropy=0.1234

qc_audit_survey (pycsamt.pipeline.ai_steps) also accepts a report_path parameter that writes the full report as JSON via write_json(), for runs where the structured report should be kept alongside the rest of the pipeline’s output.

17.7.9. Testing Plugin Steps#

Plugin registration is process-wide mutable state, the same way the Maxwell backend registry is. A test that registers a plugin step must clean it up, or it leaks into unrelated tests that run later in the same process:

 1>>> import pytest
 2>>> from pycsamt.pipeline import StepSpec, register_step, unregister_step
 3>>>
 4>>> @pytest.fixture()
 5... def demo_plugin_step():
 6...     spec = register_step(
 7...         StepSpec(
 8...             code="DEMO001",
 9...             name="scale_amplitude",
10...             label="Demo Amplitude Scale",
11...             category="demo",
12...             override_fn=scale_amplitude,
13...         )
14...     )
15...     yield spec
16...     unregister_step("DEMO001", missing_ok=True)
17...
18>>> def test_plugin_step_usable_in_a_pipeline(demo_plugin_step, sites):
19...     from pycsamt.pipeline import Pipeline, Step
20...     pipe = Pipeline([("scale", Step("DEMO001"))])
21...     result = pipe.run(sites, outdir=None, save_plots=False)
22...     assert result.ok

missing_ok=True in the teardown matters: if the test body itself already called unregister_step, or failed before registration completed, a strict teardown would raise a second, unrelated error that masks the real failure.

17.7.10. Troubleshooting#

ValueError: ... is already registered

Another step already claims this code or name. Pass replace_existing=True if overwriting it is intentional, or choose a less generic code – a vendor or project prefix such as ACME_DEMO001 avoids colliding with both the built-in catalogue and other plugins.

A plugin step is “unknown” from the CLI

Discovery did not run. Pass pycsamt pipe --with-plugins before the subcommand, run pycsamt pipe plugins first, or set PYCSAMT_PIPELINE_LOAD_PLUGINS=1.

The first pycsamt pipe plugins call feels slow

That is importlib.metadata.entry_points() enumerating every installed package, not the plugin’s own code. It is a one-time cost per process, not per step.

A plugin’s register function raised

discover_plugins reports it as a failed PluginLoadResult and warns, rather than crashing the whole discovery pass. Check pycsamt pipe plugins output, or PluginLoadResult.error, for the underlying exception message.