Skip to content

API Reference

Run Ubunye tasks from Python instead of the CLI: in a notebook, a job, or a test. This page lists the public classes and functions, generated from their docstrings.

Which API do I use?

You want to Use It returns
Run one task, end to end ubunye.run_task(task_dir, ...) {output name: frame}: what the task wrote, as the backend's own frames (a pandas DataFrame on pandas, a Spark DataFrame on Spark)
Run several tasks of one pipeline, in order ubunye.run_pipeline(usecase_dir, usecase, pipeline, tasks, ...) {task name: {output name: frame}}, without the tasks resume=True skipped
Step through a task in a notebook: read, look, transform, look, write ubunye.notebook(task_dir, ...) A NotebookContext: read() gives {input name: frame}, transform() gives {output name: frame}, write() writes them, run() does all three
Drive the engine yourself from a parsed config ubunye.core.runtime.Engine Engine.run(cfg) gives {output name: frame} (None with dry_run=True)

Most people want run_task. All of them take mode (the profile: DEV, PROD, ...; default DEV) and backend ("pandas", "spark", ...). See Execution backends and Profiles.


Python API

The main use: Databricks notebooks and jobs, where a SparkSession already exists and the engine attaches to it.

import ubunye

# Run a single task
outputs = ubunye.run_task(
    task_dir="pipelines/fraud_detection/ingestion/claim_etl",
    mode="PROD",
    dt="202510",
)

# The same task with no Spark and no Java
outputs = ubunye.run_task("pipelines/sales/etl/daily", backend="pandas")

# Run multiple tasks in order
results = ubunye.run_pipeline(
    usecase_dir="pipelines",
    usecase="fraud_detection",
    pipeline="ingestion",
    tasks=["claim_etl", "feature_engineering"],
    mode="PROD",
    dt="202510",
)

ubunye.api.run_task

run_task(task_dir, *, mode=None, dt=None, dtf=None, spark=None, backend=None, variables=None, lineage=None, lineage_dir=None, profile=None, hooks=None, rerun=False, lineage_backend=None, run_records=None)

Run a single Ubunye task and return the outputs map.

Parameters

task_dir : str Path to the task directory containing config.yaml and transformations.py. mode : str, optional Run mode, used for Spark profile merging. If not given: the job parameter ubunye_mode, else $UBUNYE_MODE, else "DEV". dt : str, optional Data timestamp, injected as {{ dt }} in Jinja templates. "" counts as not given; then the job parameter dt, else $UBUNYE_DT. dtf : str, optional Data timestamp format, injected as {{ dtf }}. spark : SparkSession, optional Explicit SparkSession to reuse. If None, auto-detects an active session (Databricks) or creates a new one. backend : Backend or str, optional The backend to run on: a registered name ("pandas" for a run with no Spark and no Java, "spark", or any installed plugin) or an instance. If None: the spark session if given, else the platform's active session, else a new Spark session. variables : dict, optional Extra template variables for the config ({"region": "gauteng"} makes {{ region }} available), as --var does on the command line. May set dt or dtf; the same name with two values is refused. lineage : bool, optional Record lineage for this run. If None (the default), records are on when a run-record location is configured (run_records, the job parameter ubunye_run_records, $UBUNYE_RUN_RECORDS); False turns them off. lineage_dir : str, optional Where lineage records go with lineage=True. For the filesystem store, a directory under the use case directory (default .ubunye/lineage); for any other store, its location as given (a Delta table name or path). profile : str, optional Config profile for validation (passed to load_config). hooks : iterable of Hook, optional Hooks to add to the engine's own (the ubunye.hooks entry point, CONFIG.monitors and the run record). They never replace them. rerun : bool Replace a batch (same dt and variables) that a finished run already appended. Without it, such a run is refused, since it would append the batch twice. lineage_backend : str, optional The run record store, by name (ADR 011): "filesystem", "delta" or an installed plugin. If None: $UBUNYE_LINEAGE_BACKEND, else "filesystem". run_records : str, optional Where run records go, which also turns them on (ADR 016): a table name (main.ops.ubunye_runs) goes to the Delta store, a path to files. If None: the job parameter ubunye_run_records, else $UBUNYE_RUN_RECORDS.

Returns

Dict[str, Any] Mapping of output name → DataFrame.

ubunye.api.run_pipeline

run_pipeline(usecase_dir, usecase, pipeline=None, tasks=None, *, mode=None, dt=None, dtf=None, spark=None, backend=None, variables=None, lineage=None, lineage_dir=None, profile=None, hooks=None, rerun=False, lineage_backend=None, resume=False, package=None, run_records=None)

Run multiple tasks sequentially and return all outputs.

Parameters

usecase_dir : str Root directory for use cases (e.g. "./pipelines"). usecase : str Use case name. pipeline : str Pipeline name (the folder between the use case and the tasks). tasks : List[str] Task names to run in order. mode, dt, dtf, spark, backend, variables, lineage, lineage_dir, profile, hooks, rerun, lineage_backend, run_records Same as :func:run_task. One backend runs every task. resume : bool Finish a pipeline that stopped half way: a task that already finished this batch is skipped (with a warning, and is not in the result); the rest run. A task that appends knows it finished; one that only overwrites is known to have finished from its run records, so it is skipped only with lineage=True (its latest run of the batch succeeded). Without resume a finished append is refused (:class:~ubunye.core.runs.BatchFinished). package : str, optional Old name of pipeline. Still works, with a DeprecationWarning; removed in 1.0.

Returns

Dict[str, Dict[str, Any]] Mapping of task name → outputs map (without tasks skipped by resume).

Job health and upstream checks

Queries over run records (ADR 013).

ubunye.api.job_health

job_health(usecase_dir, *, usecase=None, pipeline=None, tasks=None, max_age=None, stuck_after='6h', lineage_dir='.ubunye/lineage', lineage_backend=None, package=None)

Job health per task, from run records (ADR 013); what ubunye runs health prints.

Name the tasks (tasks=["uc/pipeline/task"]) or a use case (usecase=, with pipeline= to narrow it; package= is its old name, deprecated). Returns one :class:~ubunye.lineage.health.TaskHealth per task; .healthy is False when the task is failing, has a run likely dead (running past stuck_after) or, with max_age, has not succeeded recently enough.

ubunye.api.wait_for_upstream

wait_for_upstream(usecase_dir, task, *, dt=None, variables=None, timeout=0, poll=60, lineage_dir='.ubunye/lineage', lineage_backend=None)

The run of task that succeeded for this batch, waiting up to timeout.

The batch is dt and variables; every one must match. Raises :class:~ubunye.lineage.upstream.UpstreamNotReady when there is none in time. What ubunye runs wait-for does; a task can ask the same in its config with CONFIG.requires.


Notebook API

A Databricks job notebook is three steps inside one with block:

import ubunye

with ubunye.notebook(task_dir, dt=dt) as run:
    sources = run.read()
    outputs = run.transform(sources)
    run.write(outputs)

The block is one run (ADR 016):

  • Run records are on when a location is configured: run_records=, else the job parameter ubunye_run_records, else $UBUNYE_RUN_RECORDS. A table name (main.ops.ubunye_runs) goes to Delta, a path to files. The first save creates the table; if the job's user may not, the warning says to create it once or grant CREATE TABLE. lineage=False turns records off.
  • Every failure is recorded with the stage it failed in, as failed_stage: setup, read, expectations, transform, write, notebook (your code between the steps), or a name you give with with run.step("score"):.
  • One line per step is printed (quiet=True stops it):
ubunye: read 1 input (claims) in 2.10s
ubunye: transform 1 output (bronze_claims) in 0.40s
ubunye: write 1 output (bronze_claims) in 5.02s
ubunye: run 3a1b7ed8 success in 7.60s; run record: delta main.ops.ubunye_runs
  • dt="" (an empty job parameter) counts as not given; then the job parameter dt, else $UBUNYE_DT. mode not given is the job parameter ubunye_mode, else $UBUNYE_MODE, else DEV.
  • A backup copy of a record the store could not save goes to the folder in ubunye_run_records_fallback (a Unity Catalog Volume), only when one is set; otherwise there is a warning and the job goes on.
  • {{ env.VAR }} values in config.yaml come from widgets and secrets, with no manual os.environ setup.

run_task and run_pipeline read the same settings (run_records=, the job parameter, the environment) and record the same stage names.

Without with, the notebook works as before, for stepping through a task in cells; each write() leaves a run record when records are on:

ctx = ubunye.notebook(
    "/Workspace/pipelines/claims/claim_etl",
    mode="PROD",
    dt="2026-01-01",
)

sources = ctx.read()
sources["raw_claims"].display()

outputs = ctx.transform(sources)
outputs["bronze_claims"].show()

ctx.write(outputs)

ubunye.notebook.notebook

notebook(task_dir, *, mode=None, dt=None, dtf=None, spark=None, backend=None, variables=None, env=None, secrets_scope=None, secrets_map=None, lineage=None, lineage_dir=None, profile=None, auto_env=True, hooks=None, lineage_backend=None, run_records=None, quiet=False)

A task run step by step from a notebook; use it with with for one recorded run.

::

with ubunye.notebook(task_dir, dt=dt) as run:
    sources = run.read()
    outputs = run.transform(sources)
    run.write(outputs)

All {{ env.VAR }} references in the task's config.yaml are auto-resolved from Databricks widgets and secrets: no manual os.environ setup needed.

Run records are on when a location is configured: run_records= (a table name goes to the Delta store, a path to files), else the job parameter ubunye_run_records, else $UBUNYE_RUN_RECORDS. lineage=True with lineage_dir= still works as before; lineage=False turns records off. dt and mode given as "" or None come from the job parameters dt / ubunye_mode ($UBUNYE_DT / $UBUNYE_MODE); mode defaults to DEV. quiet=True stops the one-line-per-step summary.

backend picks the engine as in :func:ubunye.run_task: a registered name ("pandas") or an instance; by default the notebook's session. lineage_backend picks the run record store as in :func:ubunye.run_task.

ubunye.notebook.NotebookContext

A task run step by step from a notebook; with with, one recorded run.

run_id property

run_id

This run's id, as its run record has it.

run_records property

run_records

Where this run's record goes ("delta main.ops.runs"), or None when off.

config property

config

The fully resolved config object.

config_dict property

config_dict

The config as a plain dict.

env_vars property

env_vars

The env vars that were auto-resolved and injected.

spark property

spark

Shortcut to the active SparkSession.

task_name property

task_name

The task directory name.

read

read()

Read all inputs defined in config.yaml. Returns {name: DataFrame}.

transform

transform(sources=None)

Apply the task's transform logic.

If sources is None, uses the result of the last :meth:read call (or reads automatically if :meth:read was never called).

write

write(outputs=None)

Write outputs to configured sinks, and record the run.

If outputs is None, uses the result of the last :meth:transform call. Inside with the write is part of the one run; without it, every write leaves a run record when records are on, exactly as ubunye.run_task does.

run

run()

Read, transform, and write in one call (same as ubunye.run_task).

step

step(name)

A step of the notebook's own (with run.step("score"):): timed, printed, and the stage a failure in it is recorded under.

close

close()

Restore env vars and remove task dir from sys.path.


Core Engine

ubunye.core.runtime.Engine

Executes a task by reading inputs, applying one or more transforms, and writing outputs.

Minimal required config structure

cfg[CONFIG][inputs] : mapping input_name -> reader cfg (must include 'format') cfg[CONFIG][outputs] : mapping output_name -> writer cfg (must include 'format') cfg[CONFIG][transform]: EITHER a single transform dict with 'type', OR a list of transform dicts to form a pipeline.

Observation (logging, metrics, tracing, user monitors) is delegated to :class:ubunye.core.hooks.Hook instances. Pass hooks= to override the default set.

backend property writable

backend

The backend running this engine's tasks, resolved on first use.

__init__

__init__(backend=None, registry=None, context=None, hooks=None, extra_hooks=None, manage_backend=True)
Parameters

hooks : iterable of Hook, optional Replace the default hook set entirely. extra_hooks : iterable of Hook, optional Append these hooks to the default set. Ignored when hooks is also given. manage_backend : bool, default True If True (default), the engine calls backend.start() and backend.stop(). Set to False when the caller owns the backend lifecycle (e.g. Python API running multiple tasks on one session).

apply_transforms

apply_transforms(sources, cfg)

Apply configured transforms to sources.

Returns a dict mapping output name to the backend's own frame type. Frames that did not come from :meth:read_inputs (a notebook's own sample, say) are checked against the input contracts and counted for a reconcile first, like read ones: a reconcile compares an output with the frames the transform actually got.

open_run

open_run(cfg)

Hold one run open: what :meth:read_inputs, :meth:apply_transforms, :meth:write_outputs and :meth:step do inside it is one run (ADR 016).

The hooks see one task, from entry to exit. A failure anywhere inside ends it as failed, with the stage it was in (state["stage"]: read, expectations, transform, write or a named step; the caller names any other); a clean exit ends it as a success. Yields the run's state, the one the hooks are given. This is how a notebook's with block becomes a recorded run.

read_inputs

read_inputs(cfg)

Read all inputs defined in CONFIG.inputs.

Returns a dict mapping input name to the backend's own frame type (a pandas.DataFrame on pandas) — suitable for interactive inspection before calling :meth:apply_transforms. The inputs' expectations (input contracts) are checked here, as in :meth:run, and the inputs a reconcile names are counted; both are kept for a later :meth:write_outputs.

run

run(cfg, *, dry_run=False)

Run a task using the provided config mapping.

Parameters

cfg : dict Parsed task configuration (already rendered/validated). dry_run : bool, default False If True, validates and returns without executing readers/writers.

Returns

Optional[Dict[str, Any]] Optionally returns the outputs map (None if dry_run).

step

step(name, cfg=None)

A step the caller names (run.step("score")): timed, seen by the hooks, and, inside :meth:open_run, the stage a failure in it is recorded under.

write_outputs

write_outputs(outputs, cfg, *, as_run=False, inputs=None)

Write outputs to the sinks defined in CONFIG.outputs.

With as_run=True the write is wrapped as a whole task for the hooks, so lineage and monitors record it exactly as they record run(). This is how a notebook that reads, transforms and writes step by step still leaves a run record. A reconcile compares with the inputs counted before the transform (by :meth:read_inputs or :meth:apply_transforms). inputs given here that were not counted then are counted now, after the transform: a transform that changed them in place hides its loss.


Interfaces (Protocols)

See the full Interfaces page for design principles, discovery, and auto-detection.

ubunye.interfaces

Formal interfaces (typing.Protocol) for Ubunye's pluggable seams.

Each protocol defines the structural contract that backends must satisfy. Implementations are discovered at runtime via entry-point groups and auto-detect classmethods.

Cross-boundary dataclasses accompany each protocol so that callers never depend on a concrete backend's internal types.

AuthBackend

Bases: Protocol

Structural interface for authentication backends.

Concrete implementations: TokenAuthBackend (built-in), ServicePrincipalAuthBackend (Phase 2).

resolve
resolve(workspace_url, **kwargs)

Resolve credentials for the given workspace.

Parameters

workspace_url: The Databricks (or other) workspace host URL. **kwargs: Backend-specific hints (e.g. token_env="MY_TOKEN").

Returns

Credentials Populated credentials ready for client construction.

Raises

ubunye.core.errors.AuthNotFoundError When credentials cannot be located (missing env var, etc.).

validate
validate(credentials)

Validate that credentials are accepted by the workspace.

Returns True on success.

Raises

ubunye.core.errors.AuthInvalidError When the workspace rejects the credentials.

Credentials dataclass

Resolved authentication credentials.

The auth_type field tells callers which fields are populated. metadata carries backend-specific extras (e.g. token expiry, scope, tenant ID) without polluting the core fields.

DeployAdapter

Bases: Protocol

Structural interface for deployment backends.

Each adapter is responsible for:

  1. Validating that the target dict contains the fields it needs.
  2. Preparing artifacts (notebooks, config bundles) for the target.
  3. Uploading artifacts and creating/updating the remote job.
  4. Returning a :class:DeployResult with job metadata.
name property
name

Short identifier for this adapter (e.g. "databricks").

deploy
deploy(ctx, target, *, dry_run=False)

Deploy a task to the target environment.

Parameters

ctx: Task metadata, config dict, and filesystem paths. target: Adapter-specific target configuration (validated by the adapter via :meth:validate_target). dry_run: If True, build and return the job spec without making any remote API calls.

Returns

DeployResult Job metadata including name, id, spec, and workspace path.

get_job_url
get_job_url(job_id, target)

Return a browser URL to the deployed job.

Parameters

job_id: The job identifier returned by :meth:deploy. target: The same target dict, which typically contains the host URL.

validate_target
validate_target(target)

Validate adapter-specific target config.

Raises

ubunye.core.errors.TargetNotFoundError When required keys are missing or values are invalid.

DeployContext dataclass

Structured metadata for a deploy operation.

Callers construct this from CLI flags and the parsed config; the adapter never derives task identity from filesystem paths.

DeployResult dataclass

Returned by :meth:DeployAdapter.deploy.

metadata carries adapter-specific extras (notebook source, cluster id, etc.) without polluting the core fields.

LineageBackend

Bases: Protocol

Structural interface for lineage storage backends.

Concrete implementations: FilesystemLineageBackend (built-in), DeltaLineageBackend (Phase 2).

compact
compact()

Optimize the underlying storage.

For filesystem backends this is a no-op. For Delta backends this triggers OPTIMIZE on the lineage table.

get_run
get_run(run_id)

Load a single run record by ID.

Raises

ubunye.core.errors.LineageRecordNotFoundError When no record exists for run_id.

record
record(record)

Persist a lineage record (non-blocking per Decision 1).

Creates or overwrites the record for the given run_id. Typically called twice per task: once at start (status="running") and once at end (status="success" or "error").

The actual write is dispatched to a background worker thread. On failure, the record is written to a fallback manifest (Decision 2).

Parameters

record: The lineage record to persist. Callers construct this with the current engine version and task metadata.

search
search(*, task=None, usecase=None, pipeline=None, status=None, since=None, limit=100)

Search across recorded runs with optional filters.

Parameters

task: Filter by task name. usecase: Filter by use case. pipeline: Filter by pipeline. status: Filter by status ("success", "error", "running"). since: ISO-8601 datetime; only return runs recorded at or after. limit: Maximum number of records to return.

LineageRecord dataclass

Cross-boundary lineage record (Decision 3 compliant).

Core columns are strict; everything else goes into metadata. Every record stamps engine_version so analysts can filter by writer version.

ModelVersionInfo dataclass

Cross-boundary model version record (Decision 3 compliant).

Core fields are strict; everything else goes into metadata. Every record stamps engine_version so downstream consumers can filter by writer version.

RegistryBackend

Bases: Protocol

Structural interface for model registries.

Concrete implementations: FilesystemRegistryBackend (built-in), MLflowModelRegistryBackend (mlflow, ADR 014).

delete
delete(*, use_case, model_name, version)

Remove a version and its artifacts.

demote
demote(*, use_case, model_name, version, to_stage)

Demote a version to a lower stage.

Never move a version up: refuse (RegistryError) a move to production, or to staging from below. Call :func:ubunye.models.lifecycle.check_demote first.

get
get(*, use_case, model_name, version=None, stage=None)

Retrieve a model version by exact version or active stage.

Returns

tuple of (artifact_path, ModelVersionInfo) artifact_path is an opaque string — it may be a local path, an MLflow artifact URI, or a cloud storage URL. Callers pass it to ModelClass.load(artifact_path).

Raises

ubunye.core.errors.VersionNotFoundError When no matching version exists.

list_versions
list_versions(*, use_case, model_name)

List all registered versions, newest first.

promote
promote(*, use_case, model_name, version, to_stage, gates=None, promoted_by=None)

Promote a version to a higher lifecycle stage.

If gates are provided, all checks must pass before the transition. On gate failure, raises :class:~ubunye.core.errors.PromotionBlockedError — this is NOT caught by the graceful degradation layer (auth and gate failures must propagate).

Metadata update is non-blocking per Decision 1.

register
register(*, use_case, model_name, version, model, metrics, metadata=None, lineage_run_id=None)

Register a trained model as a new version.

Artifact writes are synchronous — the method returns only after the model is safely persisted. Index/metadata updates are non-blocking per Decision 1.

Parameters

use_case: Logical namespace (e.g. "fraud_detection"). model_name: Model identifier within the use case. version: Explicit semver string, or None for auto-increment. model: A trained model instance whose save()/metadata() methods will be called by the backend. metrics: Training metrics dict. metadata: Optional flexible key-value pairs (Decision 3). lineage_run_id: Optional lineage run ID for cross-referencing.

RunRecordStore

Bases: Protocol

Where run records live (ADR 011). Four methods, no base class needed.

A task is named by its task_path, "usecase/package/task". Newest first means the record written last comes first. A store is built with one argument, its location (a directory, a table name, a URL), given unchanged.

list_runs
list_runs(task_path, n=20)

The n newest records of one task, newest first.

load
load(task_path, run_id)

One record, every field as saved.

Raises :class:ubunye.core.errors.LineageRecordNotFoundError when the task has no run with that id.

save
save(ctx)

Create or replace the record of ctx.run_id for its task.

search
search(task_path=None, status=None, since=None, limit=None)

Records newest first, filtered by task (or its first parts), status and started_at >= since, at most limit of them.


Config

ubunye.config.schema.UbunyeConfig

Bases: BaseModel

Top-level Ubunye task config (the full contents of a config.yaml).

MODEL and VERSION are optional; they default to etl and "0.0.0-dev" respectively. Set them explicitly in production pipelines where job type or version is load-bearing (lineage records, model registry, orchestrator metadata).

check_profile

check_profile(profile)

Raise ConfigProfileError when profile names no profile of this config.

A config with no ENGINE.profiles accepts any mode. run, validate, plan, doctor and the Python API all go through this one rule.

merged_spark_conf

merged_spark_conf(profile=None)

Return base spark_conf merged with the named profile's overrides.

resolved_catalog

resolved_catalog(profile=None)

Return the catalog name, with profile override if available.

resolved_schema

resolved_schema(profile=None)

Return the schema name, with profile override if available.

ubunye.config.schema.TaskConfig

Bases: BaseModel

The CONFIG section of a task: inputs, transform, outputs.

ubunye.config.schema.IOConfig

Bases: BaseModel

Input or output connector configuration.

Extra fields are allowed so plugin-specific keys (e.g. rest_api's auth, pagination, headers) pass through to the plugin unchanged via model_dump().

ubunye.config.schema.EngineConfig

Bases: BaseModel

Spark/compute settings with optional per-profile overrides.

ubunye.config.schema.RegistryConfig

Bases: BaseModel

Configuration for model registry integration within a transform.

ubunye.config.schema.ModelTransformParams

Bases: BaseModel

Typed params for transform.type: model — for documentation and validation.


Backends

Backends are plugins, chosen by name. See Execution backends for how to choose one and write your own.

ubunye.core.backends

Backends are plugins, found by name (ADR 001) and resolved in one order (ADR 003).

Every backend, the shipped ones included, registers in the ubunye.backends entry point group::

[project.entry-points."ubunye.backends"]
spark = "ubunye.backends.spark_backend:SparkBackend"
pandas = "ubunye.backends.pandas_backend:PandasBackend"

so a third party engine plugs in the same way, with no edit here. This module never imports an engine: it loads the class a name points at, only when asked.

Which backend runs a task, in order:

  1. the one asked for by name (--backend on the CLI, backend= in the API);
  2. the platform's: a backend that finds an active session it should attach to (on Databricks, the notebook's SparkSession);
  3. :data:DEFAULT_BACKEND.

The task's config.yaml never names its engine; the same folder runs on any backend that can do what it asks.

DEFAULT_BACKEND module-attribute

DEFAULT_BACKEND = 'spark'

resolve

resolve(name=None, *, app_name='ubunye', conf=None)

The backend a run should use: the named one, else the platform's, else the default.

create

create(name, *, app_name='ubunye', conf=None)

Build the backend registered as name, if what it needs is installed.

Looking a backend up (:func:load_class) needs nothing installed, so it can be listed and inspected anywhere. Using one does: a missing package is named here, with the install that fixes it, not on the first frame.

load_class

load_class(name)

The Backend class registered as name (case blind), or a clear error.

available

available()

Registered backend names, sorted. Loads nothing.

ubunye.core.capabilities.Capabilities dataclass

A backend's declaration of what it can do.

file_formats and write_modes of None mean "any": the backend passes them through to an engine that knows its own formats (Spark reads orc, avro, xml... with the right packages).

describe

describe()

A plain, JSON ready description (ubunye backends).

unknown classmethod

unknown()

For a backend that has not declared anything.

The shipped backends (spark, databricks, pandas): what each is for and how each reads and writes is on Execution backends.

ubunye.backends.spark_backend.SparkBackend

Bases: Backend

Creates and manages a SparkSession for a Ubunye run.

Parameters

app_name : str Spark application name (appears in Spark UI/history server). conf : Optional[Dict[str, str]] Spark configuration key-values (e.g., {"spark.master": "yarn"}).

Notes
  • pyspark is imported lazily inside start() to keep installation lightweight.
  • spark property is only valid after start() (or inside the context manager).

app_name property

app_name

Configured Spark app name.

conf_effective property

conf_effective

The effective Spark configuration currently applied to the session.

Returns an empty dict if the session hasn't been started yet.

conf_input property

conf_input

The configuration dict passed into this backend at construction time.

is_spark property

is_spark

Whether this backend is Spark-based (always True here).

spark property

spark

Return the active SparkSession (after start()).

Raises

RuntimeError If start() has not been called.

timezone property

timezone

The session time zone this backend works in, for the run record (ADR 007).

Known before the session starts: the task's setting, else a running session's (which this backend would reuse, unchanged), else UTC.

environment_problem classmethod

environment_problem()

Why this backend cannot start a session here, or None (ubunye doctor).

materialise

materialise(frame)

Compute an output once for its checks, write and hash (ADR 009).

start

start()

Create (or reuse) a SparkSession.

Safe to call multiple times; a second call is a no-op if a session already exists.

stop

stop()

Stop the SparkSession if this backend started it.

A session that was already running when :meth:start attached to it (the user's own, or a notebook's) is left running.

ubunye.backends.databricks_backend.DatabricksBackend

Bases: Backend

Backend that reuses the existing Databricks SparkSession.

Parameters

spark : Optional[SparkSession] An explicit session to wrap. If None (the default), the active session is retrieved via SparkSession.getActiveSession().

timezone property

timezone

The attached session's time zone, for the run record (ADR 007).

create classmethod

create(*, app_name='ubunye', conf=None)

Attach to the active session when started (there is no app name to set).

from_platform classmethod

from_platform(*, app_name='ubunye', conf=None)

Claim the run when a SparkSession is already active (a Databricks notebook).

Attaching, rather than creating, means the run uses the cluster the platform set up, and never stops a session it did not start.

materialise

materialise(frame)

Compute an output once for its checks, write and hash (ADR 009).

start

start()

Attach to the active SparkSession, and apply ENGINE.spark_conf to it.

stop

stop()

No-op — we don't own the session, so we never stop it.

ubunye.backends.pandas_backend.PandasBackend

Bases: Backend

Execute a Ubunye task with pandas, no Spark and no JVM.

timezone property

timezone

The zone timestamp text is read in (UTC unless set).

check_io classmethod

check_io(direction, cfg)

Options and schemas this backend cannot honour, found before a run.

A file format it does not handle is left to the capability check, which already reports it.

frame_from_records

frame_from_records(records, *, schema=None)

Records typed as Spark's createDataFrame types them (F-015).

iter_records

iter_records(frame)

Each row as a dict, as Spark's Row.asDict(recursive=True) gives it.

start

start()

No session to create; check pyarrow can handle timezones here.

pyarrow before 24 cannot find a timezone database on Windows, so every timestamp read or written fails deep inside it. Say so now, with the fix. The pandas extra already asks for 24 on Windows; this catches an environment that was assembled by hand.

stop

stop()

No session to stop.

to_native

to_native(frame)

A transform gets the plain pandas.DataFrame (ADR 004).

to_port

to_port(frame)

The engine gets the port, where count() means rows (ADR 004).


Models (UbunyeModel contract)

ubunye.models.base.UbunyeModel

Bases: ABC

Abstract contract every user model must implement.

The engine calls these methods during the pipeline lifecycle. Users implement them with whatever ML library they choose — the engine never imports or knows about the underlying library.

Minimal example::

from ubunye.models.base import UbunyeModel
import joblib

class FraudRiskModel(UbunyeModel):
    def __init__(self):
        self._clf = None

    def train(self, df):
        import pandas as pd
        from sklearn.ensemble import GradientBoostingClassifier
        pdf = df.toPandas()
        X, y = pdf.drop("label", axis=1), pdf["label"]
        self._clf = GradientBoostingClassifier().fit(X, y)
        return {"accuracy": self._clf.score(X, y)}

    def predict(self, df):
        pdf = df.toPandas()
        pdf["score"] = self._clf.predict_proba(pdf)[:, 1]
        return spark.createDataFrame(pdf)

    def save(self, path):
        import pathlib, joblib
        pathlib.Path(path).mkdir(parents=True, exist_ok=True)
        joblib.dump(self._clf, f"{path}/clf.joblib")

    @classmethod
    def load(cls, path):
        import joblib
        m = cls()
        m._clf = joblib.load(f"{path}/clf.joblib")
        return m

    def metadata(self):
        import sklearn
        return {
            "library": "sklearn",
            "library_version": sklearn.__version__,
            "features": list(self._clf.feature_names_in_),
            "params": self._clf.get_params(),
        }

load abstractmethod classmethod

load(path)

Deserialize a model from the given directory path.

Parameters:

Name Type Description Default
path str

Directory path where model files were saved.

required

Returns:

Type Description
'UbunyeModel'

An instance of the model ready for prediction.

metadata abstractmethod

metadata()

Return model metadata for the registry.

Should include at minimum: - library: str — e.g. "xgboost", "sklearn", "pytorch" - library_version: str — e.g. "2.0.3" - features: List[str] — feature names used during training - params: dict — model hyperparameters

Additional keys are allowed and will be stored in the registry.

predict abstractmethod

predict(df)

Generate predictions on the provided DataFrame.

Parameters:

Name Type Description Default
df Any

Spark DataFrame with feature columns.

required

Returns:

Type Description
Any

Spark DataFrame (or DataFrame-like) with prediction columns added.

save abstractmethod

save(path)

Serialize the model to the given directory path.

The user chooses the serialization format (pickle, joblib, ONNX, etc.). The engine only provides the path — it never inspects the saved files.

Parameters:

Name Type Description Default
path str

Directory path where model files should be written.

required

train abstractmethod

train(df)

Train the model on the provided DataFrame.

Parameters:

Name Type Description Default
df Any

Spark DataFrame (or any DataFrame-like) with training data.

required

Returns:

Type Description
Dict[str, Any]

Dict of metrics e.g. {"auc": 0.94, "f1": 0.87, "accuracy": 0.91}.

Dict[str, Any]

These metrics are stored in the registry alongside the model version.

validate

validate(df)

Optional: validate model on holdout data.

Returns:

Type Description
Dict[str, Any]

Metrics dict in the same format as :meth:train.

Raises:

Type Description
NotImplementedError

if the subclass has not overridden this method.

ubunye.models.loader.load_model_class

load_model_class(task_dir, class_name)

Dynamically load a UbunyeModel subclass.

The class_name is a dotted path like "model.FraudRiskModel" where the first segment is the Python module filename (without .py) and the rest is the class name.

Resolution order: 1. If task_dir is provided → load <module_segment>.py directly from that directory using importlib.util.spec_from_file_location. 2. If task_dir is None → attempt a standard importlib.import_module (the module must already be on sys.path; _run_single_task adds the task directory to sys.path before transforms are invoked, so this works automatically when called from the engine).

Parameters:

Name Type Description Default
task_dir Optional[str]

Absolute or relative path to the directory containing the model file (e.g. the pipeline task directory). Pass None to rely on sys.path.

required
class_name str

Dotted path to the class, e.g. "model.FraudRiskModel" or "models.risk.FraudRiskModel".

required

Returns:

Type Description
Type[UbunyeModel]

The model class (not an instance). Callers should do cls() or

Type[UbunyeModel]

cls.load(path) as needed.

Raises:

Type Description
FileNotFoundError

The module file does not exist in task_dir.

ImportError

The class name was not found in the module.

TypeError

The class does not subclass :class:ubunye.models.base.UbunyeModel.


Model Registry

ubunye.models.registry.ModelRegistry

Manages model versions and lifecycle transitions.

Parameters:

Name Type Description Default
store_path str

Where the models live. For the filesystem backend, the root directory (or s3://-style URL) for all artifacts and registry files. Any other backend gets it unchanged.

required
backend Any

The registry backend (ADR 010): a name registered under the ubunye.registry_backends entry point group, or a backend object. Not given, UBUNYE_REGISTRY_BACKEND decides, else filesystem. Every backend gets this whole API built on its six protocol methods (:mod:ubunye.models.lifecycle).

None

archive

archive(use_case, model_name, version)

Archive a model version.

compare_versions

compare_versions(use_case, model_name, version_a, version_b)

Compare metrics between two versions.

Returns:

Type Description
Dict[str, Dict[str, Any]]

Dict mapping metric name to {"a": val_a, "b": val_b, "delta": diff}.

Dict[str, Dict[str, Any]]

Only metrics present in at least one version are included.

demote

demote(use_case, model_name, version, to_stage)

Demote a model version to a lower stage.

Demote only moves a version down: to development or archived, or to staging from production. Moving up is a promotion; use :meth:promote, which checks the promotion gates.

Raises:

Type Description
RegistryError

if the move would take the version up a stage.

get_model

get_model(use_case, model_name, task_dir=None, version=None, stage=None)

Return the artifact path and version metadata for a model.

Specify either version (exact) or stage (active version in that stage). task_dir is accepted for older callers and not used.

Returns:

Type Description
str

(artifact_path, model_version). The caller calls

ModelVersion

ModelClass.load(artifact_path) to obtain the model instance.

Raises:

Type Description
VersionNotFoundError

No matching version found.

list_versions

list_versions(use_case, model_name)

List all registered versions for a model (newest first by registered_at).

Raises:

Type Description
RegistryNotFoundError

Model not found in registry.

promote

promote(use_case, model_name, version, to_stage, promoted_by=None, gates=None, force=False)

Promote a model version to a higher stage.

Every gate must pass first: gates if given, else the gates stored with the model when it was registered. force=True skips them, and the version records that it was forced and which gates it skipped. If promoting to production, the current production version is automatically archived.

Raises:

Type Description
VersionNotFoundError

Version not found.

PromotionBlockedError

One or more promotion gates failed.

Returns:

Name Type Description
Updated ModelVersion

class:ModelVersion.

promotion_gates

promotion_gates(use_case, model_name)

The gates stored with a model (empty if none were set).

register

register(use_case, model_name, version, model, metrics, lineage_run_id=None, registered_by=None, promotion_gates=None)

Register a trained model as a new version (stage=development).

Parameters:

Name Type Description Default
use_case str

Logical use-case grouping (e.g. "fraud_detection").

required
model_name str

Model identifier (e.g. "FraudRiskModel").

required
version Optional[str]

Semver string. Pass None to auto-generate.

required
model UbunyeModel

A trained :class:UbunyeModel instance.

required
metrics Dict[str, Any]

Metrics dict returned by :meth:UbunyeModel.train.

required
lineage_run_id Optional[str]

The run_id of the training run. When omitted inside a run (a task's transformations.py), it is that run's id, so the model points at its run record.

None
registered_by Optional[str]

Optional username for the audit trail.

None
promotion_gates Optional[Dict[str, Any]]

The gates every later promotion of this model must pass. Given, they replace the model's stored gates.

None

Returns:

Type Description
ModelVersion

The newly created :class:ModelVersion.

rollback

rollback(use_case, model_name, to_version)

Roll production back to to_version; the current production version is archived.

ubunye.models.registry.ModelVersion dataclass

Metadata for a single registered model version.

to_info

to_info(*, name='', artifact_path='')

Convert to the cross-boundary :class:ModelVersionInfo type.

The model's metadata keeps its JSON types, so it converts back to the same :class:ModelVersion (:func:ubunye.models.lifecycle.to_model_version).

ubunye.models.registry.ModelStage

Bases: str, Enum


Promotion Gates

ubunye.models.gates.PromotionGate

Evaluates a set of metric thresholds before allowing model promotion.

Parameters:

Name Type Description Default
gate_config Dict[str, Any]

Dict of gate rules. Supported keys: - min_<metric>: metric must be >= value - max_<metric>: metric must be <= value - require_drift_check: boolean; checks metadata["drift_check_passed"]

required

all_passed

all_passed(metrics, metadata=None)

Return True only if every configured gate passes.

Parameters:

Name Type Description Default
metrics Dict[str, Any]

Model metrics from :meth:UbunyeModel.train.

required
metadata Optional[Dict[str, Any]]

Optional model metadata from :meth:UbunyeModel.metadata.

None

evaluate

evaluate(metrics, metadata=None)

Evaluate all configured gates against the provided metrics.

Parameters:

Name Type Description Default
metrics Dict[str, Any]

Model metrics dict from :meth:UbunyeModel.train.

required
metadata Optional[Dict[str, Any]]

Optional model metadata dict from :meth:UbunyeModel.metadata.

None

Returns:

Type Description
List[GateResult]

List of :class:GateResult — one per configured gate.

List[GateResult]

Promotion is allowed only if all results have passed=True.

failed_gates

failed_gates(metrics, metadata=None)

Return only the gates that did not pass.

ubunye.models.gates.GateResult dataclass

Result of evaluating a single promotion gate.


Lineage

ubunye.lineage.recorder.LineageRecorder

Monitor plugin that persists run lineage as structured JSON.

Parameters

store: The old name for backend, kept with a DeprecationWarning ("s3" still gives the stub store, without a warning). base_dir: The store's location, passed to it unchanged: a directory for filesystem, a table name or path for delta. sample_fraction: Ignored since 0.7.0 and kept so existing configs still load: every row is hashed now (the rows-v1 content hash), so there is nothing to sample. backend: The run record store, by name (ADR 011): "filesystem", "delta" or any installed ubunye.lineage_backends plugin. If None, $UBUNYE_LINEAGE_BACKEND, else "filesystem". fallback: Where a record the store could not save is kept (ADR 016): "~" (the default) is ~/.ubunye/fallback, as before; a folder (a Unity Catalog Volume) keeps it there; None keeps it nowhere and only warns, for a job whose home folder dies with it (Databricks serverless).

reads_inputs property

reads_inputs

Whether inputs are hashed; if not, no source version is taken (F-046).

recorded

recorded(run_id)

Whether this recorder has a record of the run (it saw the run start).

shared_store

shared_store()

The run record store, unless it is the filesystem one.

Rerun safety (ADR 008) reads a dead run's status through it; with the filesystem store it reads the record files directly, as before.

task_end

task_end(*, context, config, outputs, status, duration_sec, inputs=None, expectations=None, timings=None, llm_calls=None, llm_budget=None, model_runs=None, hash_basis=None, source_versions=None, error=None, failed_stage=None)

Update the run record with final status, duration, and step hashes.

hash_basis is the engine's word, by output name, on whether each output frame is the rows that were written ("materialised") or will be computed again for the hash ("recomputed"). An output it does not name, and every input, gets the basis its frame type implies.

source_versions is each input's source version when it was read (:mod:ubunye.lineage.source_version). A recomputed input's version is taken again right after its hash; if it moved, the step says the digest is of a later state than the one read (F-046).

model_runs is each model the run trained or scored with: the version, the stage and, for training, the gate result (ADR 012).

error is why a failed run failed (the exception's type and message). It is kept in the record, with the values of secret-looking variables masked, so lineage show, OpenLineage and ubunye prove can say why (F-054).

failed_stage is the stage a failed run was in (ADR 016): setup, read, transform, expectations, write, notebook or the name of a step the notebook named (run.step("score")).

task_start

task_start(*, context, config)

Create a "running" lineage record and persist it immediately.

ubunye.lineage.context.RunContext dataclass

Full lineage record for one task execution.

ubunye.lineage.storage.FileSystemLineageStore

Bases: LineageStore

Stores lineage records as JSON files in a local directory tree.

Satisfies both the legacy :class:LineageStore ABC and the :class:~ubunye.interfaces.LineageBackend protocol.

Layout::

{base_dir}/{usecase}/{package}/{task_name}/{run_id}.json

compact

compact()

No-op for filesystem backend.

flush

flush(timeout=None)

Drain the background worker queue.

Returns the number of records that went to fallback manifests.

get_run

get_run(run_id)

Load a single run record by run ID.

Uses a run_id→path index for O(1) lookup instead of rglob.

load

load(task_path, run_id)

One run of a task, by its run id or the first characters of it (as ubunye lineage list prints them). A prefix that fits more than one run is refused, never guessed.

record

record(record)

Persist a lineage record (non-blocking per Decision 1).

Dispatches the write to a background worker thread. On failure, the record is written to a fallback manifest (Decision 2).

search_records

search_records(*, task=None, usecase=None, pipeline=None, status=None, since=None, limit=100)

Protocol-compatible search returning LineageRecord objects.


Plugins: Readers

ubunye.plugins.readers.rest_api.RestApiReader

Bases: Reader

Read records from a REST API endpoint into the backend's frame.

Handles pagination (offset, cursor, next_link), authentication (bearer, api_key, basic), rate limiting, and retry with exponential backoff. The HTTP side is the same on every backend; the backend builds the frame (frame_from_records), typed as Spark's createDataFrame types it.

read

read(cfg, backend)

Fetch all pages from the API and return the backend's frame.

Parameters

cfg : dict Reader configuration. Required key: url. See module docstring for full reference. backend : Backend Any backend with the records capability (spark, databricks, pandas).


Plugins: Writers

ubunye.plugins.writers.rest_api.RestApiWriter

Bases: Writer

Write a frame to a REST API endpoint in JSON batches.

Takes the rows from the backend (iter_records), groups them into batches of batch_size, and POSTs each batch as {"records": [...]}. Tracks and logs success/failure counts per batch.

write

write(df, cfg, backend)

POST the frame's rows to a REST endpoint in batches.

Parameters

df : frame The output, as the engine hands it to writers. cfg : dict Writer configuration. Required key: url. See module docstring for full reference. backend : Backend Any backend with the records capability (spark, databricks, pandas).


Plugins: Transforms

ubunye.plugins.transforms.model_transform.ModelTransform

Bases: Transform

Transform plugin for ML operations (train / predict).

Implements :class:ubunye.core.interfaces.Transform so it plugs into the existing engine dispatch with zero changes to the core runtime.

apply

apply(inputs, cfg, backend)

Dispatch to train or predict based on cfg["action"].

Parameters:

Name Type Description Default
inputs Dict[str, Any]

Named DataFrames from the read phase.

required
cfg dict

Transform params from config.yaml (the params sub-dict).

required
backend Backend

Active engine backend (provides Spark session if needed).

required

Returns:

Type Description
Dict[str, Any]

Dict of named outputs. For train: {"model_metrics": dict}.

Dict[str, Any]

For predict: {"predictions": DataFrame}.


Internal ML wrappers

Note

These are internal wrappers used by Ubunye's own sklearn/Spark ML adapters. User-defined models should implement UbunyeModel, not these classes.

ubunye.plugins.ml.base.BaseModel

Bases: HasSchema, ABC

Framework-agnostic estimator interface.

Subclasses must implement
  • _fit_core(X, y)
  • _predict_core(X) -> y_pred (and optionally probabilities)
  • _save_core(path)
  • _load_core(path)
  • params property (read-only view)

params property

params

Model hyperparameters (read-only view).

fit

fit(X, y=None)

Fit the model; X can be Spark DF, pandas DF, numpy array, etc.

load

load(path)

Load model and metadata from path.

metrics

metrics()

Optional: training/validation metrics (subclasses may override).

predict

predict(X, proba=False)

Predict on inputs; returns labels or (labels, probabilities) if proba=True.

save

save(path)

Persist model and metadata to a directory path.

ubunye.plugins.ml.sklearn.SklearnModel

Bases: BaseModel

Wrap any sklearn estimator with a unified Ubunye interface.

ubunye.plugins.ml.pysparkml.SparkMLModel

Bases: BaseModel

Wrap a Spark ML Estimator or loaded PipelineModel.