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.

1. Create an independent package

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"]

[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 and this package on local runners, Airflow scheduler/workers and Ray workers. Discovery uses configured roots such as:

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.

2. Define the main manifest

normalize/normalize.drtml:

drtml: 4
step: normalize
schema: normalize_v1

include:
  - fragments/params.drtml
  - fragments/ui.drtml
  - fragments/metrics.drtml

run:
  type: partition
  processor:
    batch: acme_pipeline_steps.normalize.processor:process_batch
    init: acme_pipeline_steps.normalize.processor:init_processor_ctx
  derive:
    token_registry:
      from: documents.tokens
      operation: count_values
      count: count

storage:
  resume: overwrite
  default: artifacts
  stores:
    artifacts:
      backend: s3
      root: pipeline-artifacts
      profile: minio_app
    warehouse:
      backend: postgres
      connection: postgres_app
      schema: pipeline

inputs:
  documents:
    run: mapping.upstream_runs
    schema:
      document_id: string
      language: string
      text: string
    feed:
      batch: 1024
      select: [document_id, language, text]
      group_by: [language]
      group_scope: shard

  metadata:
    as: metadata
    run: mapping.metadata_run
    schema:
      document_id: string
      source: string
    indexes:
      document_id:
        kind: routing
        columns: [document_id]
    attach:
      to: documents
      on: [document_id]
      route_by: document_id
      scope: primary_unit

outputs:
  normalized_documents:
    from: documents
    schema:
      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:
    store: warehouse
    kind: registry
    key: token
    schema:
      run_id: {type: string, not_null: true}
      token: {type: string, not_null: true}
      count: {type: int64, not_null: true}
    physical:
      table: drt_tbl_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
    measurements:
      vocabulary_size:
        method: count_distinct
        column: token
        unit_column: document_id
        cadence: checkpoint
        accuracy: exact

orchestration:
  dag_id: normalize_v1
  connections: [minio_s3, postgres_app]
  artifacts_root: env:DRTOLLER_ARTIFACTS_ROOT
  xcom:
    - key: normalized_outputs
      from: outputs.storage_manifests

observability:
  grafana:
    uid: normalize-v1

Key separation:

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

3. Declare defaults

fragments/params.drtml:

params:
  domain:
    lowercase: true
    collapse_whitespace: true
    minimum_token_length: 2
  runtime:
    metrics_enabled: true
    resume_enabled: true
    document_batch_size: 32
    corpus_batch_max_chars: 0
    max_docs: 0
    progress_log_every_n_docs: 100
    finalize_flush_every_n_docs: 0
    pipeline:
      read_q_max: 4
      write_q_max: 32
      lane_q_max: 32
    parallelism:
      backend: inprocess
      workers: 1
      ray_address: ""
  storage:
    parquet_compression: zstd
  mapping:
    upstream_runs: ""
    metadata_run: ""

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

4. Generate the Streamlit form

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.

5. Declare metrics and dashboard

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

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

processing:
  outputs:
    documents: normalized_documents
  registry_flush:
    vocab: 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.

7. Implement processor initialization

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.

8. Implement the batch processor

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 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.

9. Add the thin hook

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.

10. Unit-test domain code

from acme_pipeline_steps.normalize.processor import process_batch
from drtoller.framework.storage.feed 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.

11. Contract-test the manifest

Tests should load and compile the real packaged DRTML:

from pathlib import Path
from drtoller.framework.drtml import compile_step_manifest, load_manifest_path


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.

12. Integration-test storage and dispatch

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.

13. Run locally

Configure manifest discovery:

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

Then call the local API:

from drtoller.framework.integration.local 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.

14. Expose in Streamlit and 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.

15. Choose parallelism deliberately

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.

For uneven shard cost (long documents, mixed files), set runtime.parallelism.schedule: pull so idle workers take the next shard. On Ray, deflate_idle: true releases pods as the queue drains.

16. Production checklist

• 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.