feat: add notebook: cell type for multi-notebook composition.

[Ya-shh]

Backend and labextension for the composable-notebooks KEP (KEP-0812). Compose multiple Kale notebooks into one Kubeflow pipeline through the existing kale --nb path and the Kale panel, with no new command or flag.

Demo of the composition flow on YouTube: https://youtu.be/rJEnyU6t3S8?si=LTkmlLsmMPCjdKlh

What it does

A notebook: cell type lets a notebook reference other notebooks. Two shapes are supported:

• Sequence of notebooks (Story 1): the root notebook holds only notebook: references. Each referenced notebook becomes a KFP sub-pipeline (GraphComponent), so cell-level step visibility is preserved. Data flow between notebooks is inferred by matching variable names, the same AST and PyFlakes mechanism Kale already uses between cells, applied one level up.

• Mixed steps and notebooks (Story 2): the root notebook can also carry its own step:, imports and functions cells alongside the references. Each such step becomes a top-level component, and ordering and data cross the step/notebook boundary the same way they cross a step/step one. A root step runs after every unit above it and before every unit below it (reading order), while independent references between steps still run in parallel.

Whether a notebook compiles as a plain pipeline, a sequence, or a mixed composition is decided entirely by the cells it contains.

How it works (reworked after review)

The first version of this PR put composition in a separate processors/workflow.py behind a branch in cli.py. Two review findings retired that design:

• @StefanoFioravanzo pointed out the branch read as a codepath bolted onto the primary flow, and that the recursion belongs inside NotebookProcessor, the way a compiler recurses.
• @ederign found that the Kale panel's compile RPC never reached workflow.py at all, so composing from the UI silently produced a pipeline with the referenced notebooks missing.

Composition now lives inside NotebookProcessor, and workflow.py and the cli.py branch are deleted. Every notebook goes through one compile path, which is what fixes the UI case:

1. Parse. The processor reads cells in order. On a notebook: cell it recurses into a child NotebookProcessor and attaches the resulting pipeline as a SubPipeline node, a peer of Step in the same graph.
2. Link. dependencies_detection resolves only along edges that already exist, and a composition has no prev: tags to create them. So the edges are inferred first: data flow between units, plus reading order for the root notebook's own steps.
3. Resolve. The existing dependencies_detection then runs unchanged, so a variable crossing a notebook boundary is resolved by exactly the code that resolves one crossing a cell boundary.
4. Emit. Compiler walks the mixed graph and writes one DSL module per referenced notebook plus the orchestrator.

DSL file structure (per review feedback)

Each notebook produces its own independent DSL module (.kale/kale_notebook_<name>.py) and the orchestrating DSL file imports those modules as dependencies. This mirrors the recursive shape of the processing: each notebook is an independent module with no circular references, each owns its own configuration, and each is compilable standalone. In a mixed composition the root notebook's own steps live in the orchestrator, which is that notebook's own module. The kale_notebook_ prefix means a notebook named json.ipynb, class.ipynb or 2train.ipynb compiles rather than erroring.

What is in the PR

Backend

processors/nbprocessor.py: the notebook:<name> tag (referenced path in cell metadata) added to the tag language; recursive parsing of references into SubPipeline nodes; link_composition_units() for edge inference; propagate_subpipeline_boundaries() to hand each boundary variable to the inner step that produces or consumes it. A notebook: cell breaks the code merge chain (KEP-0812 caveat 5): untagged cells after a reference attach to the next step: cell, and orphaned code raises rather than being dropped.

step.py: SubPipeline(Step), a referenced notebook as a node of the same graph, carrying the same name, ins, outs and config so dependency detection and code generation treat a reference and a step uniformly.

compiler.py: generate_composition() and the per-notebook module emission. Each referenced notebook's module is generated by a compiler bound to that notebook's own pipeline, so its pipeline parameters and imports are the ones its steps are built with, exactly as when it is compiled on its own.

templates/subpipeline_template.jinja2 and templates/composition_template.jinja2: the per-notebook module (components, typed sub-pipeline, standalone compile hook) and the orchestrator that imports the modules, holds any root-step components, and wires everything.

common/kfputils.py: compile_pipeline carries the DSL script's sibling modules into its temp dir so the orchestrator's imports resolve.

Frontend (labextension)

A Notebook cell type in the Kale cell editor: pick it, enter the notebook path, and it writes the notebook: tag plus the notebook_path metadata. Reference cells render an inline notebook: <name> chip and are excluded from prev: dependency choices and the merge-chain notice.

Examples and tests

examples/composition/ holds main.ipynb (three references, dataset to model to prediction) and main_mixed.ipynb (references plus a root step: cell). test_composition.py covers reference-to-node conversion, inferred order and data, variables crossing the step/notebook boundary both ways, reading order, preserved parallelism, the merge rule, orphan code, missing paths, reference cycles, name collisions, per-module output, safe module naming, parameter survival, nested-reference rejection, reported step config, and two real kfp.compiler compilations.

The single-notebook path is untouched: a notebook with no notebook: cells compiles exactly as before, and no existing cell type or API changes.

Testing

• kale --nb examples/composition/main.ipynb compiles to a root DAG of three sub-DAGs, each holding its notebook's own cell-level steps; main_mixed.ipynb adds the root's step as a top-level component wired in reading order.
• The Kale panel's compile RPC produces the same result, which is the case that previously dropped references.
• Both submit and run green against a KFP cluster.
• 297 backend tests passing, ruff clean, labextension tsc, eslint and prettier clean.

Known limitations

• Nested references. A notebook: cell inside a referenced notebook raises rather than compiling to something that flattens a level of nesting. Nested composition is not part of the KEP stories, so this is a deliberate rejection rather than a partial implementation. Happy to change it if reviewers would prefer it supported.
• Per-step configuration. The composition templates do not yet carry over per-step limits, annotations, labels, retries or timeouts. Compiling reports exactly which step and which field is affected rather than dropping it quietly. Restoring this is the remaining step of the agreed rework plan.
• Explicit typed input/output declarations remain the v2 discussion from the KEP review.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Ya-shh]

You can watch the full demo of this feature on YouTube
https://youtu.be/rJEnyU6t3S8?si=LTkmlLsmMPCjdKlh

[Copilot]

Pull request overview

This PR adds Phase 2 (backend) support for composable notebooks by introducing a notebook: cell type that lets a “composition notebook” reference other notebooks and compile them into a single KFP pipeline while preserving per-cell step visibility via sub-pipelines.

Changes:

• Adds a multi-notebook composition compiler (kale.processors.workflow) that infers cross-notebook data flow and generates one DSL module per notebook plus an orchestrator DSL.
• Extends the notebook tag language to recognize notebook:<name> and step:notebook:<name> and updates kale --nb to route to composition when references are present.
• Adds unit tests and an examples/composition/ demo notebook set.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
kale/tests/unit_tests/test_workflow.py	Adds unit coverage for topo sort, cycle detection, boundary inference, name collisions, and wiring semantics.
kale/templates/workflow_template.jinja2	New orchestrator DSL template that imports per-notebook modules and wires sub-pipelines.
kale/templates/subpipeline_template.jinja2	New per-notebook DSL module template defining components + sub-pipeline + standalone compile hook.
kale/processors/workflow.py	Implements composition: reference extraction, dependency inference, boundary variable wiring, and DSL generation.
kale/processors/nbprocessor.py	Extends tag parsing with NOTEBOOK_TAG and records notebook reference metadata.
kale/common/kfputils.py	Updates pipeline compilation to copy sibling modules so orchestrator imports resolve in a temp dir.
kale/cli.py	Routes kale --nb to composition when notebook: cells exist; prints concise user-facing errors.
examples/composition/main.ipynb	Adds a composition notebook that references A/B/C as sub-pipelines.
examples/composition/notebook_a.ipynb	Adds example “producer” notebook for the demo.
examples/composition/notebook_b.ipynb	Adds example intermediate notebook for the demo.
examples/composition/notebook_c.ipynb	Adds example “consumer” notebook for the demo.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Ya-shh]

moved the kfp import above the sys.path mutation so no generated module can shadow it, and the composer now rejects notebook names that collide with kfp/kale/stdlib. Kept insert(0) rather than append deliberately: the generated modules must win resolution of their own names, otherwise a same named module elsewhere on sys.path would silently shadow them

[Ya-shh]

New updated DSL structure
https://youtu.be/AFvds4Jwdnw?si=YMoywpODG25WNArl

thank you @jesuino for the suggestion!

[StefanoFioravanzo]

@Ya-shh I went through the code. For of all, congrats for submitting this first version. It's a lot of work and I really appreciate the lengths you went to understand the code deeply and implement a sophisticated mechanism.

I didn't run it myself, but I saw the demo and I know the approach work, practically. So we are in a very good position, because you proved this can be done.

I added some comments inline, but to me the bottom line is that we should try to rethink a bit how you put together the whole logic inside workflow.py. When thinking about an elegant, fluid, and well architected solution, I imagine a fully recursive implementation inside NotebookProcessor, instead of having to build a separate dedicated logic inside a separate module (workflow.py). The function compose_notebooks_as_subpipelines is very long and a bit hard to follow as well.

I'd like to hear what @ederign thinkgs as well, because what I am proposing would mean changing this code A LOT, but at the same time would allows us to have code that anyone can understand and feel like it was built for multi-notebook orchestration from day 1. As it is now, it still feels like it was built afterwards, with bits and pieces conneting sometimes in an awkward way

Also - mind you - I know it's hard to build a fully recursive loop, but that's also why it will be so much more elegant and beautiful :)

[Ya-shh]

@StefanoFioravanzo @ederign step 3 of the rework is landed in 5bc0b6e

Composition now happens inside NotebookProcessor, and processors/workflow.py plus the _compile_composition branch in cli.py are deleted
A notebook: cell spawns a child processor and attaches the result as a SubPipeline node beside the Steps, so every notebook goes through one compile path
Edges are inferred before resolution, because dependencies_detection only resolves along edges that already exist and a composition has no prev: tags to create them
The resolution itself is unchanged, so a variable crossing a notebook boundary is handled by exactly the code that handles one crossing a cell boundary

R1 is fixed, The panel's compile RPC was the path that silently dropped references, It now produces a root DAG where each referenced notebook is a nested sub-DAG holding its own cell-level steps, and a mixed root adds its own steps as top-level components, verified against the reproducers in the gist

Two things I want to flag :

1. A referenced notebook's pipeline-parameters were being dropped, It compiled clean but would have failed at runtime with a NameError, since the component signature lost the name its body still used ,fixed in the same commit: each referenced notebook's module is now generated by a compiler bound to that notebook's own pipeline, which is also closer to the per-module design @ederign asked for
2. Per-step config the composition templates do not carry over yet (limits, annotations, labels, retries, timeouts) is now logged by step and field instead of disappearing quietly, carrying it over is step 4, the remaining piece

A notebook: cell inside a referenced notebook currently raises ,It used to flatten that level of nesting silently, which I did not think was acceptable, but nested composition is not in the four steps or in the KEP stories either
Is rejecting it the behaviour you want for now, or should it compile?

[Ya-shh]

Step 4 is landed too, in d1fa92b, so the four step plan is complete !!

The per-task config block is now a single shared macro imported by every emission path, so a step's limit:, label:, annotation: and cache: tags have the same effect whether its notebook is composed or compiled on its own
previously a referenced notebook's GPU limit was dropped, which meant a step tagged for a GPU silently ran on CPU once composed
The warning I mentioned above is removed, since the gap it reported is now closed

Single notebook output is unchanged, verified byte for byte against the golden DSL fixtures in test_e2e.py

One correction to my note above:
retry_* and timeout turned out not to be emitted by any Kale template, so they are unused StepConfig fields across the project rather than something composition drops
probably worth a separate look :)

[Ya-shh]

@StefanoFioravanzo

Template: done subpipeline_template and composition_template are both
deleted
root, referenced and composing notebooks all render from
pipeline_template now
Two checks that it is the right seam: the compiled KFP IR is byte identical before and after for five pipelines, compositions included, so nothing about what runs moved and it then survived merging #907, where PVC mounting slotted into the shared macro in four lines and composed notebooks got PVC mounts for free

What is genuinely sub-pipeline specific turned out to be four things, not the
three I said
I found the fourth by breaking it: KFP raises "Platform-specific
features can only be set on primitive components" if you set pod labels or a
security context on a nested pipeline task
so a sub-pipeline call takes caching and nothing else, and that is KFP's constraint rather than ours

SubPipeline I would still keep and that error is part of why, It is not a
pipeline, it is a use of one: reference the same notebook twice and you get two
nodes over two pipelines with their own ins, outs and config, and only one
of the two kinds of node can carry platform config at all
Definition vs call site, eager to know your view on this :)

separately, @ederign's second round of reproducers is at six of seven passing

[StefanoFioravanzo]

@Ya-shh congrats on this amazing work, and thanks for sticking with us through all the rounds of reviews!
I gave another coarse look at the code and architecture, we are definitely in a good spot now. I didn't manage to test the feature flag, I was in a hurry and I won't have the laptop with me for the next few days. But if @ederign or others cross check this and confirm that the feature flag works as expected, I am all on for merging this PR! This is a great milestone and paves the way to the second phase - creating a full fledges low-code notebook orchestration UI :D

[ederign]

@Ya-shh merging this. Congrats, it's a big one for Kale. It will open a sea of new possibilities.

The part I want to call out isn't the feature, it's that you took "restructure most of it" without starting over. Exactly as we proposed in the beginning.

Follow-ups, none blocking:

1. [backend] Root pipeline parameters do not reach steps inside a referenced notebook — compiles clean, NameError at runtime #945 - root pipeline-parameters don't reach steps inside a referenced notebook. Compiles clean, NameError at runtime. Same family as the four you fixed, and it's the mirror of the one you caught yourself in 5bc0b6e. This is the one I'd take first.
2. [backend] Generated sub-pipeline modules collide across notebooks on the kale --nb path #942 - the kale_notebook_* prefix keys on the pipeline name, and kale --nb defaults that to kale-pipeline, so two compositions in one
   dir still overwrite each other on the CLI path. Panel path is fine.
3. Document multi-notebook composition — no user-facing docs, no examples README, env var undocumented #944 - docs.

Thanks for taking the feedback well, and the error messages and comments explaining why rather than what is not a common thing to see.

Keep pushing!!

[ederign]

/approve
/lgtm

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: ederign

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

• OWNERS [ederign]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment