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 parameterubunye_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=Falseturns 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 withwith run.step("score"):. - One line per step is printed (
quiet=Truestops 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 parameterdt, else$UBUNYE_DT.modenot given is the job parameterubunye_mode, else$UBUNYE_MODE, elseDEV.- 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 manualos.environsetup.
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_records
property
¶
Where this run's record goes ("delta main.ops.runs"), or None when off.
transform ¶
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 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.
step ¶
A step of the notebook's own (with run.step("score"):): timed, printed,
and the stage a failure in it is recorded under.
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.
__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 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 ¶
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 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 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 ¶
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 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 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.).
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:
- Validating that the target dict contains the fields it needs.
- Preparing artifacts (notebooks, config bundles) for the target.
- Uploading artifacts and creating/updating the remote job.
- Returning a :class:
DeployResultwith job metadata.
deploy ¶
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.
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 ¶
Optimize the underlying storage.
For filesystem backends this is a no-op. For Delta backends
this triggers OPTIMIZE on the lineage table.
get_run ¶
Load a single run record by ID.
Raises¶
ubunye.core.errors.LineageRecordNotFoundError
When no record exists for run_id.
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 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).
demote ¶
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 ¶
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.
promote ¶
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 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.
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 ¶
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 ¶
Return base spark_conf merged with the named profile's overrides.
resolved_catalog ¶
Return the catalog name, with profile override if available.
resolved_schema ¶
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:
- the one asked for by name (
--backendon the CLI,backend=in the API); - the platform's: a backend that finds an active session it should attach to (on Databricks, the notebook's SparkSession);
- :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.
resolve ¶
The backend a run should use: the named one, else the platform's, else the default.
create ¶
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.
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).
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. sparkproperty is only valid afterstart()(or inside the context manager).
conf_effective
property
¶
The effective Spark configuration currently applied to the session.
Returns an empty dict if the session hasn't been started yet.
conf_input
property
¶
The configuration dict passed into this backend at construction time.
spark
property
¶
Return the active SparkSession (after start()).
Raises¶
RuntimeError
If start() has not been called.
timezone
property
¶
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
¶
Why this backend cannot start a session here, or None (ubunye doctor).
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 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().
create
classmethod
¶
Attach to the active session when started (there is no app name to set).
from_platform
classmethod
¶
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.
ubunye.backends.pandas_backend.PandasBackend ¶
Bases: Backend
Execute a Ubunye task with pandas, no Spark and no JVM.
check_io
classmethod
¶
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 ¶
Records typed as Spark's createDataFrame types them (F-015).
iter_records ¶
Each row as a dict, as Spark's Row.asDict(recursive=True) gives it.
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.
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
¶
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
¶
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
¶
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
¶
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 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. |
Dict[str, Any]
|
These metrics are stored in the registry alongside the model version. |
validate ¶
Optional: validate model on holdout data.
Returns:
| Type | Description |
|---|---|
Dict[str, Any]
|
Metrics dict in the same format as :meth: |
Raises:
| Type | Description |
|---|---|
NotImplementedError
|
if the subclass has not overridden this method. |
ubunye.models.loader.load_model_class ¶
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 |
required |
class_name
|
str
|
Dotted path to the class, e.g. |
required |
Returns:
| Type | Description |
|---|---|
Type[UbunyeModel]
|
The model class (not an instance). Callers should do |
Type[UbunyeModel]
|
|
Raises:
| Type | Description |
|---|---|
FileNotFoundError
|
The module file does not exist in |
ImportError
|
The class name was not found in the module. |
TypeError
|
The class does not subclass :class: |
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 |
required |
backend
|
Any
|
The registry backend (ADR 010): a name registered under the
|
None
|
compare_versions ¶
Compare metrics between two versions.
Returns:
| Type | Description |
|---|---|
Dict[str, Dict[str, Any]]
|
Dict mapping metric name to |
Dict[str, Dict[str, Any]]
|
Only metrics present in at least one version are included. |
demote ¶
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 ¶
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
|
|
ModelVersion
|
|
Raises:
| Type | Description |
|---|---|
VersionNotFoundError
|
No matching version found. |
list_versions ¶
List all registered versions for a model (newest first by registered_at).
Raises:
| Type | Description |
|---|---|
RegistryNotFoundError
|
Model not found in registry. |
promote ¶
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: |
promotion_gates ¶
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. |
required |
model_name
|
str
|
Model identifier (e.g. |
required |
version
|
Optional[str]
|
Semver string. Pass |
required |
model
|
UbunyeModel
|
A trained :class: |
required |
metrics
|
Dict[str, Any]
|
Metrics dict returned by :meth: |
required |
lineage_run_id
|
Optional[str]
|
The |
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: |
rollback ¶
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 ¶
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:
- |
required |
all_passed ¶
Return True only if every configured gate passes.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics
|
Dict[str, Any]
|
Model metrics from :meth: |
required |
metadata
|
Optional[Dict[str, Any]]
|
Optional model metadata from :meth: |
None
|
evaluate ¶
Evaluate all configured gates against the provided metrics.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics
|
Dict[str, Any]
|
Model metrics dict from :meth: |
required |
metadata
|
Optional[Dict[str, Any]]
|
Optional model metadata dict from :meth: |
None
|
Returns:
| Type | Description |
|---|---|
List[GateResult]
|
List of :class: |
List[GateResult]
|
Promotion is allowed only if all results have |
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
¶
Whether inputs are hashed; if not, no source version is taken (F-046).
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 ¶
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
flush ¶
Drain the background worker queue.
Returns the number of records that went to fallback manifests.
get_run ¶
Load a single run record by run ID.
Uses a run_id→path index for O(1) lookup instead of rglob.
load ¶
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 ¶
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 ¶
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.
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 ¶
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 ¶
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 |
required |
backend
|
Backend
|
Active engine backend (provides Spark session if needed). |
required |
Returns:
| Type | Description |
|---|---|
Dict[str, Any]
|
Dict of named outputs. For |
Dict[str, Any]
|
For |
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)
predict ¶
Predict on inputs; returns labels or (labels, probabilities) if proba=True.