[Vedaant-J]

Series: kumo-engine MLOps series: every snapshot, training table, model bundle and prediction run records where it came from; snapshots and models move between machines and registries; fits and prediction runs are tracked in MLflow or Snowflake; plain Python job scripts run on a scheduler through a small job helper; prediction runs are monitored for drift and realized performance.
Stack: Scheduling: #70 → #71
Merge order: one linear stack; merge bottom to top (#53 first)
This PR: Adds kumo.job(), a context manager that turns a plain Python script into a scheduled job: one JSON result line on stdout however the script ends, a 0 / non-zero exit code with the failure's retry class in the summary, and the scheduler's job ids recorded in every manifest and tracked run the script writes.

Summary

This PR adds kumo.job(), a context manager that turns a plain Python script into a scheduled job: one JSON result line on stdout however the script ends, a 0 / non-zero exit code with the failure's retry class in the summary, and the scheduler's job ids recorded in every manifest and tracked run the script writes.

Problem. Teams putting Kumo on a scheduler (Databricks Jobs, SageMaker, Snowflake ML Jobs, Airflow) need answers like these:

• Did last night's run fail because a source was briefly unavailable (retry), or because a table is missing (page someone)?
• Which bundle, snapshot and outputs did it produce, without scraping logs?
• How do the scheduler's job and run ids end up in the snapshot, training-table and bundle records and on the tracked run?

Cause. A job script uses the normal Python API. Progress lines, Lightning output, native-library output and warnings all share stdout, so a scheduler cannot tell results from log noise. A failure is a traceback with exit code 1, whatever its cause. Scheduler ids reach the manifests only if the script exports KUMO_JOB_ID/KUMO_JOB_RUN_ID itself.

Change.

from kumo.engine import PredictiveQuery, Trainer, job

# graph and plan come from your job script.
with job() as run:                                  # or job(job_id=..., job_run_id=...)
    snapshot = graph.snapshot()
    query = PredictiveQuery(graph=graph, query="PREDICT users.churned FOR EACH users.user_id")
    result = Trainer(plan).fit(query.generate_training_table(snapshot), snapshot=snapshot)
    run.record(bundle_id=result.bundle_id, snapshot_id=snapshot.id,
               outputs={"bundle": str(result.bundle_dir)})

• job(job_id=None, job_run_id=None) sets KUMO_JOB_ID/KUMO_JOB_RUN_ID when given; otherwise the existing variables or the detected platform (Record provenance for every artifact under one MLOps artifact root #50's scheduler context) apply. Every manifest the body writes (snapshot, training table, provenance.json) and every tracked fit run (kumo.scheduler, kumo.job_id, kumo.job_run_id, Log, register and serve fitted models through MLflow pyfunc #65) carries them.
• Logging goes to stderr (logging.basicConfig(level=INFO, stream=stderr)).
• run.record(bundle_id=, snapshot_id=, run_id=, outputs=) adds to the summary; None leaves a field as it was.
• On exit, however the body ends (success, exception, Ctrl-C, sys.exit(), SIGTERM), the environment and warnings.showwarning are restored and one JSON line goes to stdout (both in a finally); then a failed body exits the process.
• Warnings still reach the previous showwarning. When another library wraps kumo.job()'s handler after it is installed (MLflow autologging's warnings controller does this inside a Databricks trial), the re-entered call prints the warning directly instead of recursing; the test simulates that wrapper.

How the body ends	Exit code	error.kind
succeeded	0	none
PermanentSourceError in the cause chain (e.g. a missing table or no permission)	1	permanent: do not retry
TransientSourceError, ConnectionError or TimeoutError in the cause chain	1	transient: worth a retry
any other exception	1	unknown
SIGTERM	143	unknown
Ctrl-C	re-raised (130 in the summary)	unknown
sys.exit(n) in the body	n passes through	when n is non-zero

Databricks, SageMaker and Snowflake retry policies only see zero vs non-zero, so the retry class lives in the summary line rather than in the exit code.

Real stdout of a job script that builds and exports a snapshot of a synthetic users table, then of the same script failing on a missing table, with KUMO_JOB_ID=job-7 KUMO_JOB_RUN_ID=run-7 (hashes shortened; stderr held the logs):

$ python job.py ok; echo exit=$?
{"event": "kumo.engine.job", "status": "succeeded", "exit_code": 0, "scheduler": {"platform": "other", "job_id": "job-7", "job_run_id": "run-7"}, "bundle_id": null, "snapshot_id": "190b5d6c…@db819f3c5a154468", "run_id": null, "outputs": {"snapshot": "<store>/snapshots/190b5d6c…/db819f3c5a154468"}, "warnings": []}
exit=0

$ python job.py missing; echo exit=$?
{"event": "kumo.engine.job", "status": "failed", "exit_code": 1, "scheduler": {"platform": "other", "job_id": "job-7", "job_run_id": "run-7"}, "bundle_id": null, "snapshot_id": null, "run_id": null, "outputs": {}, "warnings": [], "error": {"type": "SourceNotFoundError", "kind": "permanent", "message": "no table users in the source"}}
exit=1

Summary fields: event (always kumo.engine.job), status, exit_code, scheduler (platform, job_id, job_run_id, or null when there is none), bundle_id, snapshot_id, run_id, outputs (what the body recorded; record full URIs or table names, since a scheduler cannot resolve connector-relative paths), warnings (the first 20 distinct warnings raised in the body, 300 characters each; deprecation and future warnings, mostly from third-party libraries, are left out and still go to stderr), and on failure error (type, kind, message cut to 500 characters; the traceback goes to stderr).

This replaces the kumo-engine console script and its JSON config that earlier revisions of this PR proposed: a job script already expresses graphs, connectors, tracking and registry options in the normal API, so a second config schema was only another surface to keep in sync. The job templates in #71 are plain Python scripts under kumo.job().

Back-compat

New API only.

Situation	Behaviour
Scripts that do not use job()	Unchanged.
KUMO_JOB_ID/KUMO_JOB_RUN_ID already set	Kept unless job() is given ids.
Body raises	One summary line, traceback on stderr, SystemExit(1).
Ctrl-C or sys.exit() in the body	The summary line is still printed and the environment restored; the exception propagates unchanged.
job() entered off the main thread	Works; SIGTERM is not intercepted there.
Previous SIGTERM handler and warnings.showwarning	Restored on exit.

Known limitations

• Inside a Databricks notebook job task the job and run ids are not visible to the process; pass them with job(job_id=..., job_run_id=...) from task parameters, or only platform: databricks is recorded.

Testing Done

• test_job.py (CPU): the summary line with recorded ids, outputs and the scheduler context, and the same scheduler block in a manifest written inside the body; warnings deduplicated, capped at 20 and 300 characters, deprecations left out; each failure class gives exit 1 with error.kind transient (in the cause chain, connection reset), permanent (missing source) or unknown, and a long message cut to 500 characters; Ctrl-C, sys.exit(3) and sys.exit(0) still print the summary with the right code and restore the job-id environment; SIGTERM prints the summary, exits 143 and restores the previous handler.
• Re-run on the current head: test_job.py + test_optional_imports.py: 56 passed. flake8 and isort clean.
• The output above is from an earlier head, with the exit code and error.kind updated to what this head prints.

Checklist

• No breaking changes (new API).
• README updated ("Scheduled jobs").
• Unit tests added.

🤖 Generated with Claude Code

[Vedaant-J]

Change map

File	Function / scope	What changed	Why
packages/kumo-engine/README.md	"Scheduled jobs"	new section: kumo.job(), summary fields, 0/1/143 exit codes and error.kind, warning cap, where scheduler ids are recorded	user docs
packages/kumo-engine/kumo/engine/__init__.py	exports	job, JobRun	public API
packages/kumo-engine/kumo/engine/_job.py	new: job, JobRun.record, _print_summary, _exit_code, _error_kind, Terminated, _sigterm_raises	sets KUMO_JOB_ID/KUMO_JOB_RUN_ID, logs to stderr, collects up to 20 distinct non-deprecation warnings (300 chars each); in finally restores the environment and prints one JSON summary line however the body ends; exits 0 / 1 / 143 (SIGTERM), a body's own sys.exit code passes through; error.kind transient/permanent/unknown from the cause chain	schedulers get a result line on every exit and a retry hint
packages/kumo-engine/test/test_job.py	new tests	summary with ids, outputs, capped warnings and scheduler (also in a manifest); each error kind; Ctrl-C and sys.exit still print the summary; SIGTERM	coverage

Requirement: a scheduled script gives the scheduler a machine-readable result on every exit, a 0/non-zero exit code and a retry hint, and its job ids reach every record. Existing: the Python API, #50's scheduler context read from KUMO_JOB_*, #65's scheduler tags on tracked runs. Smallest: one context manager over the existing API, no config schema. Stopping condition: success, each failure class, Ctrl-C and sys.exit produce one summary line and the right exit code; ids appear in manifests.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/kumo-engine
• Branch: vedaant/mlops-mlflow-model
• File: .coderabbit.yaml
• Commit: 5255ad9
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: c8cff3f5-787a-4b28-8b98-03e1f3fdb024

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[Vedaant-J]

Closing this PR: its pieces either moved or were dropped.

• Settings (artifact root, tracking, registry) that kumo.job() read from the environment now live in kumo.MLOps(...), turned on with with kumo.mlops(cfg): or mlops=cfg per call (Record provenance for every artifact under one MLOps artifact root #50). There are no environment settings.
• Example job scripts are plain Python scripts in Add example MLOps job scripts for Databricks, SageMaker and Snowflake #71. They take CLI flags (--artifact-root, --source, --registry, --model-name, --model-path, --predictions-table) and run inside with kumo.mlops(...). Add example MLOps job scripts for Databricks, SageMaker and Snowflake #71 now sits on Log, register and serve fitted models through MLflow pyfunc #65.
• Dropped: job and run ids, the JSON summary line, warning capture, error kinds, the exit-code contract and SIGTERM handling. Databricks Jobs, SageMaker Pipelines and Snowflake Tasks already track their own runs, exit status and logs, so the engine records no scheduler identity.