Skip to content

Build a custom step

This guide builds a complete partition_loop step as an independent Python package. It does not assume the step repository, website, and framework live in one source checkout.

The example reads documents, normalizes text, joins optional metadata, and writes normalized rows plus a token-count registry.

acme-pipeline-steps/
pyproject.toml
src/
acme_pipeline_steps/
__init__.py
normalize/
__init__.py
normalize.drtml
hooks.py
processor.py
fragments/
params.drtml
ui.drtml
metrics.drtml
tests/
test_normalize_processor.py
test_normalize_manifest.py

Minimal packaging:

[project]
name = "acme-pipeline-steps"
version = "0.1.0"
dependencies = ["drtoller-framework"]
[build-system]
requires = ["setuptools>=70"]
build-backend = "setuptools.build_meta"
[tool.setuptools.package-data]
acme_pipeline_steps = ["**/*.drtml"]

Package every *.drtml fragment into the wheel and verify the installed tree contains them. Editable installs hide missing package data.

The deployment image installs both drtoller-framework and this package on local runners, Airflow scheduler/workers and Ray workers. Discovery uses configured roots such as:

Terminal window
export DRTOLLER_STEPS_ROOTS=/opt/acme-pipeline-steps/src/acme_pipeline_steps

Point the root at the directory that directly contains step folders (normalize/, …), not at an unrelated repository layout.

normalize/normalize.drtml:

drtml_version: 3
step_id: normalize
schema_id: normalize_v1
include:
- {path: fragments/params.drtml, merge: params_defaults}
- {path: fragments/ui.drtml, merge: ui}
- {path: fragments/metrics.drtml, merge: metrics}
execution:
mode: partition_loop
entry: acme_pipeline_steps.normalize.hooks:run
storage:
resume: overwrite
bindings:
artifacts:
backend: s3
root: pipeline-artifacts
profile: minio_app
datasets:
documents:
role: input
binding: artifacts
run_id_param: mapping.upstream_runs
columns:
document_id: string
language: string
text: string
feed:
loop: primary
read:
mode: batch
batch_rows: 1024
columns: [document_id, language, text]
partition:
source: file
group_by: [language]
metadata:
role: input
binding: artifacts
run_id_param: mapping.metadata_run
columns:
document_id: string
source: string
indexes:
document_id:
kind: routing
columns: [document_id]
feed:
loop: attach
read: {mode: scan}
attach:
to: documents
join_on: [document_id]
route_by: document_id
scope: primary_unit
normalized_documents:
role: output
binding: artifacts
backend: parquet
columns:
run_id: string
document_id: string
language: string
source: string
normalized_text: string
tokens: list<string>
flush:
mode: rows
value: 25000
max_rows_per_part: 100000
boundary_columns: [language]
checkpoint:
column: document_id
every_n: 500
token_registry:
role: output
backend: postgres
kind: registry
uniq: token
group_by: []
columns:
run_id: {type: string, not_null: true}
token: {type: string, not_null: true}
count: {type: int64, not_null: true}
postgres:
connection: postgres_app
schema: pipeline
table: normalize_token_registry_v1
schema_version: "2026.07.1"
indexes:
run_token:
columns: [run_id, token]
unique: true
on_conflict:
target: [run_id, token]
action: do_update
update_additive: [count]
checkpoint:
column: document_id
every_n: 500
metric_defs:
vocabulary_size:
method: count_distinct
column: token
unit_column: document_id
cadence: checkpoint
accuracy: exact
processing:
init: acme_pipeline_steps.normalize.processor:init_processor_ctx
processor: acme_pipeline_steps.normalize.processor:process_batch
batch_processor: acme_pipeline_steps.normalize.processor:process_batch
outputs:
documents: normalized_documents
registry_flush:
tokens: token_registry
registry_ingest:
from_logical: documents
mode: list_count
uniq_column: tokens
count_column: count
feed_attach:
metadata: metadata
orchestrator:
dag_id: normalize_v1
connections: [minio_s3, postgres_app]
artifacts_root: env:DRTOLLER_ARTIFACTS_ROOT
xcom:
- key: normalized_outputs
from: outputs.output_manifests_json
grafana:
uid: normalize-v1

Key separation:

  • DRTML names physical inputs/outputs, columns, storage and orchestration.
  • Processor returns logical output key documents; processing.outputs maps it to normalized_documents.
  • Framework ingests vocabulary state from logical processor rows and drains it to token_registry.
  • The attach dataset reaches the processor as logical key metadata.

fragments/params.drtml:

params_defaults:
domain.lowercase: true
domain.collapse_whitespace: true
domain.minimum_token_length: 2
runtime.metrics_enabled: true
runtime.resume_enabled: true
runtime.document_batch_size: 32
runtime.max_docs: 0
runtime.progress_log_every_n_docs: 100
runtime.finalize_flush_every_n_docs: 0
runtime.parallelism.backend: inprocess
runtime.parallelism.workers: 1
runtime.parallelism.ray_address: ""
storage.parquet_compression: zstd
mapping.upstream_runs: ""
mapping.metadata_run: ""

Every behavioral/resource number belongs here. Domain code reads domain.*; framework reads runtime.*, storage.*, mapping.*, and flush.*.

fragments/ui.drtml:

ui:
streamlit:
enabled: true
fields:
upstream_runs:
type: runs_picker
label: Document run IDs
param: mapping.upstream_runs
required: true
metadata_run:
type: text
label: Metadata run ID
param: mapping.metadata_run
lowercase:
type: checkbox
label: Lowercase text
param: domain.lowercase
minimum_token_length:
type: int
label: Minimum token length
param: domain.minimum_token_length
min: 1
max_docs:
type: int
label: Maximum documents (0 = all)
param: runtime.max_docs
min: 0
io_controls: {enabled: true}
runtime_controls: {enabled: true}
parallelism_controls:
enabled: true
backend_options: [inprocess, processes, ray]
flush_controls:
from_dataset_output: true

No Streamlit code is needed in the package. The generated overlay is merged with params_defaults before dispatch.

fragments/metrics.drtml:

metrics:
progress:
unit: documents
stage: run
total:
from: scan
cap_param: runtime.max_docs
prometheus:
- name: normalize_documents_total
kind: counter
labels: [run_id, step_id, stage]
source: counter_tick
- name: normalize_tokens_total
kind: counter
labels: [run_id, step_id, stage]
source: processor_stat
stat: tokens_total
- name: normalize_vocabulary_size
kind: gauge
labels: [run_id, step_id, stage, dataset]
source: dataset_metric
dataset: token_registry
metric: vocabulary_size
views:
panels:
- id: throughput
type: timeseries
title: Documents processed
metric: normalize_documents_total
stage: run
- id: vocabulary
type: timeseries
title: Vocabulary size
metric: normalize_vocabulary_size
stage: run
- id: cpu_stat
type: stat
title: CPU
metric: step_cpu_percent
stage: run
- id: ram_stat
type: stat
title: RAM
metric: step_ram_used_mib
stage: run

Processor stats are data returned in ProcessResult.stats; telemetry emits them according to DRTML. The processor does not import Prometheus.

6. Prefer registry_ingest for vocabularies

Section titled “6. Prefer registry_ingest for vocabularies”

Canonical path: return logical rows from the processor and let framework ingest them.

processing:
outputs:
documents: normalized_documents
registry_flush:
tokens: token_registry
registry_ingest:
from_logical: documents
mode: list_count # or row
uniq_column: tokens # list field on the logical rows
count_column: count

list_count reads a list column from each logical row and counts uniq values. row mode ingests one uniq key per row and may carry payload_columns / group_by_column. Every key in ProcessResult.rows must also appear in processing.outputs.

Framework drains on checkpoint, writes the flush dataset, and resumes through loader APIs. Domain code does not open PostgreSQL or parquet. Use a custom registry[] plugin only when ingest cannot express the state.

processor.py:

from __future__ import annotations
from typing import Any, Mapping
def init_processor_ctx(*, ctx, params: Mapping[str, Any], resume: bool) -> dict[str, Any]:
del ctx, resume
# Construct reusable in-memory objects here (regexes, NLP model, lookup tables).
return {
"lowercase": bool(params["domain.lowercase"]),
"collapse_whitespace": bool(params["domain.collapse_whitespace"]),
}

init_processor_ctx returns in-memory state. It must not load manifests, create sessions, write checkpoints, or emit metrics.

Framework calls the exact signature:

def process_batch(
raws: list[Any],
*,
ctx,
params: Mapping[str, Any],
processor_ctx: Mapping[str, Any],
partition_id: str,
) -> list[ProcessResult]:
...

With a feed, each item in raws is a ProcessorBatch:

from __future__ import annotations
import re
from typing import Any, Mapping
from drtoller.framework.processing.process_result import ProcessResult
from drtoller.framework.storage.feed.batch import ProcessorBatch
_WS = re.compile(r"\s+")
def process_batch(
raws: list[Any],
*,
ctx,
params: Mapping[str, Any],
processor_ctx: Mapping[str, Any],
partition_id: str,
) -> list[ProcessResult]:
del partition_id
min_len = int(params["domain.minimum_token_length"])
results: list[ProcessResult] = []
for raw in raws:
if not isinstance(raw, ProcessorBatch):
raise TypeError(f"expected ProcessorBatch, got {type(raw)!r}")
metadata_by_id = {
str(row["document_id"]): row
for row in raw.attach.get("metadata", ())
}
output_rows: list[dict] = []
token_count = 0
for row in raw.primary_rows:
text = str(row.get("text") or "")
if processor_ctx["collapse_whitespace"]:
text = _WS.sub(" ", text).strip()
if processor_ctx["lowercase"]:
text = text.lower()
tokens = [tok for tok in text.split(" ") if len(tok) >= min_len]
token_count += len(tokens)
document_id = str(row["document_id"])
meta = metadata_by_id.get(document_id, {})
output_rows.append({
"run_id": ctx.run_id,
"document_id": document_id,
"language": str(row.get("language") or ""),
"source": str(meta.get("source") or ""),
"normalized_text": text,
"tokens": tokens, # list_count registry_ingest
})
results.append(ProcessResult(
ok=True,
rows={"documents": output_rows},
stats={"tokens_total": token_count},
shard_id=raw.shard_id,
))
return results

ProcessResult fields:

  • ok: document/unit success flag.
  • rows: logical output key → logical row list.
  • stats: integer counters merged by framework and available to processor_stat metrics.
  • shard_id: optional source identity; framework can fill partition identity during merge.

Do not use the outdated form ProcessResult(dataset=..., rows=[...]); the contract uses a rows mapping.

hooks.py:

from pathlib import Path
from drtoller.framework.drtml.resolve import manifest_path_adjacent
from drtoller.framework.processing.partition_loop import run_partition_loop
def run(ctx, *, manifest_path=None):
path = Path(manifest_path) if manifest_path else manifest_path_adjacent(__file__)
return run_partition_loop(ctx, manifest_path=path)

The hook resolves the contract and delegates. It does not parse DRTML or manage a session.

from acme_pipeline_steps.normalize.processor import process_batch
from drtoller.framework.storage.feed.batch import ProcessorBatch
def test_normalizes_and_attaches_metadata():
batch = ProcessorBatch(
partition_id="en",
source_run_id="source-run",
primary_rows=({"document_id": "d1", "language": "en", "text": " HELLO World "},),
attach={"metadata": ({"document_id": "d1", "source": "web"},)},
shard_id="part-0",
)
ctx = type("Ctx", (), {"run_id": "run-1"})()
result = process_batch(
[batch],
ctx=ctx,
params={"domain.minimum_token_length": 2},
processor_ctx={"lowercase": True, "collapse_whitespace": True},
partition_id="en",
)[0]
assert result.ok
assert result.rows["documents"][0]["normalized_text"] == "hello world"
assert result.rows["documents"][0]["source"] == "web"
assert result.stats == {"tokens_total": 2}

No filesystem, parquet, database, Prometheus, Ray, or Airflow is needed for this test.

Tests should load and compile the real packaged DRTML:

from pathlib import Path
from drtoller.framework.drtml.loader import load_manifest_path
from drtoller.framework.drtml.compile import compile_step_manifest
def test_manifest_compiles():
path = Path(__file__).parents[1] / "src/acme_pipeline_steps/normalize/normalize.drtml"
manifest = load_manifest_path(path)
compiled = compile_step_manifest(manifest, merged_params=dict(manifest.params_defaults))
assert manifest.step_id == "normalize"
assert compiled.storage.dataset_output["normalized_documents"].boundary_columns == ("language",)
assert compiled.feed.primary_dataset_id == "documents"

Also assert:

  • every processing.outputs target exists;
  • every field param has a default/reader;
  • attach routing index exists;
  • metric references resolve;
  • CPU/RAM panels exist;
  • no forbidden I/O/import patterns in domain modules.

Use a temporary local binding and seed input via StorageSession.write; never write parquet directly with pyarrow. Then run run_from_manifest or run_step_local and inspect logical output rows.

For PostgreSQL/Qdrant, keep transport tests separate and gated by service fixtures. A processor unit test must stay backend-independent.

Configure manifest discovery:

Terminal window
export DRTOLLER_STEPS_ROOTS=/opt/acme-pipeline-steps/src/acme_pipeline_steps

Then call the local API:

from drtoller.framework.integration.local.run import run_step_local
result = run_step_local(
"normalize",
"normalize-dev-001",
user_params={
"mapping.upstream_runs": "source-run-001",
"runtime.max_docs": 1000,
},
)

The local runner applies the same default merge, params hook, env fallbacks, connections, and dispatch path as Airflow.

Streamlit discovers the same manifest and generates its form from ui.streamlit. Airflow creates/binds a per-step DAG, takes dag_run.conf as a user overlay, applies connections/env fallbacks, calls dispatch, and pushes standard plus declared XCom outputs.

No integration-specific code is added to the step.

Start with:

runtime.parallelism.backend: inprocess
runtime.parallelism.workers: 1

Move to processes or ray only when:

  • the feed mode is supported;
  • processor context and worker spec are picklable/reconstructible;
  • no checkpoint/resume is required;
  • run-global mutable state has been redesigned as partial rows + downstream reduction.
  • The step package installs independently of the framework source tree.
  • DRTML and fragments are included as package data.
  • No step id or domain column is hardcoded in framework code.
  • Domain processor is pure and backend-independent.
  • All numeric knobs are in params_defaults.
  • Generated form fields are live params.
  • One dataset has one backend.
  • PostgreSQL DDL is generated/applied at deploy, not runtime.
  • Prometheus and dashboards come from DRTML.
  • Generated Grafana JSON is not committed.
  • Layer and domain-boundary tests pass.