Add transport-aware Skippy stage ordering

[i386]

Summary

Teach the Skippy topology planner to consider directed pipeline edge cost when ordering selected stages.

This is the transport-specific slice from the MLX comparison work. It does not change runtime networking, protocol bytes, or artifact planning. It only adds an optional planner input and a new package-aware planning entry point.

Before

Package-aware planning could score nodes by capacity, cache locality, missing artifact cost, and placement signals. That answers “which nodes should participate?”

For a staged pipeline, that is not enough. Decode and prefill activations move stage-to-stage, so the ordered edges matter:

flowchart LR
  A["stage 0: node A"] -->|"slow"| B["stage 1: node B"]
  B -->|"slow"| C["stage 2: node C"]
  C -->|"fast"| D["stage 3: node D"]

Loading

The same node set can be valid but unnecessarily slow if adjacent stages are ordered across expensive links.

How It Works

This PR adds StageEdgeSignal:

• source_node_id
• target_node_id
• optional directed rtt_ms
• optional large_frame_bytes_per_sec
• direct_prediction_return_supported

plan_package_aware_contiguous_with_transport(...) first applies the existing package-aware node ordering, then reorders the same candidate nodes by adjacent directed edge cost.

For small stage counts, it checks all permutations. For larger plans, it falls back to a greedy ordering so planning stays bounded.

flowchart LR
  A["stage 0: node A"] -->|"2 ms"| C["stage 1: node C"]
  C -->|"3 ms"| B["stage 2: node B"]
  B -->|"4 ms"| D["stage 3: node D"]

Loading

The accepted plan emits diagnostics for measured stage edges:

[info] pipeline edge node-a -> node-c after stage-0 has measured rtt 5 ms

Why This Is Good

Skippy split serving is chain-latency sensitive. A single slow adjacent stage edge can dominate token latency even when each node is individually capable.

MLX pipeline mode gets this discipline naturally from rank-neighbor communication. This PR brings the same idea to Skippy without copying MLX’s runtime: the topology planner can now prefer the cheapest activation chain when the selected nodes are otherwise equivalent.

Scope

• Adds planner-only StageEdgeSignal.
• Adds plan_package_aware_contiguous_with_transport(...).
• Keeps the existing plan_package_aware_contiguous_with_signals(...) API as a compatibility wrapper.
• Does not infer true peer-to-peer edge costs from coordinator-to-peer RTT.
• Does not change wire protocol or server transport.

Validation

• cargo fmt --all -- --check
• cargo test -p skippy-topology --lib
• cargo clippy -p skippy-topology --all-targets -- -D warnings

Summary by CodeRabbit

• New Features

  • Planner orders pipeline stages using measured per-edge network latency (RTT) when available.
  • New API accepts transport/edge signals so planning can prefer lower-latency stage transitions.
  • Ties in edge cost preserve original package order deterministically.
  • Plans include edge-related diagnostics showing per-transition network cost details.

• Tests

  • Added transport-aware tests validating ordering, tie behavior, and diagnostic messages.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 1e203ee6-9df2-4edc-b3af-ea4cf7c9b754

📥 Commits

Reviewing files that changed from the base of the PR and between c26f84f and a16644d.

📒 Files selected for processing (2)

• crates/skippy-topology/src/lib.rs
• crates/skippy-topology/src/tests.rs

🚧 Files skipped from review as they are similar to previous changes (2)

• crates/skippy-topology/src/lib.rs
• crates/skippy-topology/src/tests.rs

---

📝 Walkthrough

Walkthrough

Adds edge-signal-driven node ordering to topology planning: a StageEdgeSignal contract, exhaustive and greedy ordering strategies using RTT-derived costs, planning API integration that conditionally runs edge-aware ordering, diagnostics attachment for measured edges, and three unit tests validating ordering and tie-breaking.

Changes

Transport-Aware Node Ordering for Topology Planning

Layer / File(s)	Summary
Edge signal model and ordering entrypoint crates/skippy-topology/src/edge_order.rs	StageEdgeSignal carries per-edge RTT and capability metadata. order_pipeline_nodes dispatcher routes small node counts to exhaustive search and larger counts to greedy heuristic. append_edge_diagnostics enriches the plan with informational edge RTT diagnostics.
Ordering strategies and cost computation crates/skippy-topology/src/edge_order.rs	Exhaustive permutation search picks the ordering with minimum aggregate edge cost; greedy search selects the lowest-latency start node and repeatedly appends the next node with minimum edge cost. Helper functions compute edge cost (using measured RTT or per-node latency), order score with deterministic tie-breaking, and signal lookups.
Planning API and edge signal integration crates/skippy-topology/src/lib.rs, crates/skippy-topology/src/edge_order.rs	Module exposes StageEdgeSignal and new plan_package_aware_contiguous_with_transport function. Refactors existing wrapper to delegate to the new transport-aware entrypoint. Conditions fast-path fallback on empty edge signals. Replaces prior local sorting with order_pipeline_nodes and post-processes with append_edge_diagnostics.
Transport-aware planning tests crates/skippy-topology/src/tests.rs	Test helpers edge() and stage_layout() construct signals and extract orderings. Three tests validate that planning selects the minimum-cost edge path, prefers cheaper directed edges, and preserves package order (warm before cold) when costs tie.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

• michaelneale
• ndizazzo

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.05% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and concisely describes the main change: adding transport-aware stage ordering to the Skippy topology planner.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/skippy-transport-edge-ordering

---

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

[i386]

@michaelneale could you review this one when you get a chance?

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

crates/skippy-topology/src/lib.rs (1)

965-970: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Tie-breaking uses pre-sort indices, not package order.

You assign tuple indices before package-aware sorting. In exhaustive mode, equal-cost tie-break then prefers original request order, which can undo package-priority ordering.

Proposed fix

     let mut nodes = request
         .nodes
         .iter()
         .cloned()
         .enumerate()
         .collect::<Vec<_>>();
     nodes.sort_by(|(left_index, left), (right_index, right)| {
         let left_signal = placement_signal_for(placement_signals, &left.node_id);
         let right_signal = placement_signal_for(placement_signals, &right.node_id);
         node_package_score(right, right_signal)
             .cmp(&node_package_score(left, left_signal))
             .then_with(|| left_index.cmp(right_index))
             .then_with(|| left.node_id.cmp(&right.node_id))
     });
+    let nodes = nodes
+        .into_iter()
+        .map(|(_, node)| node)
+        .enumerate()
+        .collect::<Vec<_>>();
     let nodes = edge_order::order_pipeline_nodes(nodes, placement_signals, edge_signals);

Also applies to: 979-979

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/skippy-topology/src/lib.rs` around lines 965 - 970, Tie-breaking
currently uses pre-sort indices because you enumerate and assign tuple indices
to request.nodes before performing package-aware sorting, so in exhaustive mode
equal-cost ties prefer original request order and can override package-priority
ordering; to fix, defer the enumerate() / index assignment until after you sort
request.nodes by package priority (i.e. perform the package-aware sort on
request.nodes first, then do .iter().cloned().enumerate().collect::<Vec<_>>() to
build the nodes vector), and apply the same change to the other occurrence that
mirrors this pattern around the nodes handling at the later block.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/skippy-topology/src/edge_order.rs`:
- Line 27: The early-return in edge_order.rs currently skips transport-aware
ordering for two-stage plans because it checks `if nodes.len() < 3 ||
edge_signals.is_empty() { ... }`; change this condition so only trivial graphs
are skipped (e.g., `if nodes.len() < 2 || edge_signals.is_empty() { ... }` or
`if nodes.len() <= 1 || edge_signals.is_empty() { ... }`) so that when there are
exactly two nodes and `edge_signals` is present the directed edge cost logic
still runs; update the conditional guarding the transport-aware ordering (the
`nodes.len()` check involving `edge_signals`) accordingly.

In `@crates/skippy-topology/src/tests.rs`:
- Around line 91-92: The test
transport_aware_plan_preserves_package_order_when_edges_tie currently passes
edge_signals: &[] so it never exercises transport edge tie-breaking; update the
test to construct and pass a non-empty edge_signals (e.g. a Vec of
placement_signal calls or edge-specific signals for the transports under test)
to the planner invocation so the edge-aware ordering path is executed, ensuring
the test asserts ordering when transport costs tie by comparing results for the
tied-edge signals instead of using edge_signals: &[].

---

Outside diff comments:
In `@crates/skippy-topology/src/lib.rs`:
- Around line 965-970: Tie-breaking currently uses pre-sort indices because you
enumerate and assign tuple indices to request.nodes before performing
package-aware sorting, so in exhaustive mode equal-cost ties prefer original
request order and can override package-priority ordering; to fix, defer the
enumerate() / index assignment until after you sort request.nodes by package
priority (i.e. perform the package-aware sort on request.nodes first, then do
.iter().cloned().enumerate().collect::<Vec<_>>() to build the nodes vector), and
apply the same change to the other occurrence that mirrors this pattern around
the nodes handling at the later block.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 3c600581-2fb5-4fc1-b502-a8b8a6a7efea

📥 Commits

Reviewing files that changed from the base of the PR and between ef40600 and b5d6441.

📒 Files selected for processing (3)

• crates/skippy-topology/src/edge_order.rs
• crates/skippy-topology/src/lib.rs
• crates/skippy-topology/src/tests.rs

[ndizazzo]

minor fixup

[coderabbitai]

🧹 Nitpick comments (1)

crates/skippy-topology/src/tests.rs (1)

90-112: ⚡ Quick win

Consider verifying diagnostics for consistency.

The similar transport test at lines 54-88 asserts that NetworkPipelineCost diagnostics are emitted for measured edges. Adding a diagnostics check here would improve test consistency and coverage.

📋 Suggested diagnostics assertion

     assert_eq!(
         stage_layout(&plan),
         vec![("node-b", 0, 3), ("node-a", 3, 6)]
     );
+    assert!(plan.diagnostics.iter().any(|diagnostic| diagnostic.code
+        == PlanReasonCode::NetworkPipelineCost
+        && diagnostic.message.contains("node-b -> node-a")));
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/skippy-topology/src/tests.rs` around lines 90 - 112, The test
transport_aware_plan_orders_two_stages_by_edge_cost should also assert that
NetworkPipelineCost diagnostics are emitted for measured edges like the earlier
transport test: after calling
plan_package_aware_contiguous_with_transport(&request, &[], &[...] ) and before
or after checking stage_layout(&plan), inspect the plan's diagnostics (same
access pattern used in the other test) and assert that a NetworkPipelineCost
diagnostic exists for the measured edges (edge("node-a","node-b",200) and
edge("node-b","node-a",5))—use the same diagnostic keys/types and helper(s) from
the other test to locate and compare diagnostics so the test mirrors the
consistency checks present in the test at lines 54-88.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/skippy-topology/src/tests.rs`:
- Around line 90-112: The test
transport_aware_plan_orders_two_stages_by_edge_cost should also assert that
NetworkPipelineCost diagnostics are emitted for measured edges like the earlier
transport test: after calling
plan_package_aware_contiguous_with_transport(&request, &[], &[...] ) and before
or after checking stage_layout(&plan), inspect the plan's diagnostics (same
access pattern used in the other test) and assert that a NetworkPipelineCost
diagnostic exists for the measured edges (edge("node-a","node-b",200) and
edge("node-b","node-a",5))—use the same diagnostic keys/types and helper(s) from
the other test to locate and compare diagnostics so the test mirrors the
consistency checks present in the test at lines 54-88.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 033c9079-a2d6-4484-a4ab-641b41fdc598

📥 Commits

Reviewing files that changed from the base of the PR and between b5d6441 and c26f84f.

📒 Files selected for processing (3)

• crates/skippy-topology/src/edge_order.rs
• crates/skippy-topology/src/lib.rs
• crates/skippy-topology/src/tests.rs

🚧 Files skipped from review as they are similar to previous changes (2)

• crates/skippy-topology/src/lib.rs
• crates/skippy-topology/src/edge_order.rs