swactor/examples/pipeline-parallel-inference/TEST_SPEC.md

Test Spec: Pipeline-Parallel Inference, Two Nodes

This spec enumerates every component surface in SPEC.md and names the behavioral test that proves it correct. No implementations here — just the name and a one-line statement of what passing the test proves.

Test Philosophy

• Fail fast, cheap first. Pure-Rust unit tests run on every save. Local multi-process integration runs on every commit. Anything touching real vast.ai is gated and rare.
• Behavior over structure. Tests describe observable outcomes (a response arrives, a token id is valid, a cluster converges). They do not assert internal call sequences or echo back the implementation.
• One killer correctness test. The "non-empty response" bar is too low for sliced inference — silently wrong outputs would pass. Equivalence against single-node inference (§9) is the load-bearing test.

Test Tiers

Tier	What runs	When	Approx wall-clock
1. Pure unit (Rust)	codec, layer-range math, message types	cargo test	< 1 s
2. Python worker contract	stdin/stdout protocol, stub mode	pytest	a few seconds
3. Stage actor (Rust, in-process)	one stage actor + its worker subprocess	cargo test	seconds
4. Local cluster integration	3 swactor nodes in one process, real iroh, stub worker	cargo test	seconds
5. Local cluster + real tinygrad	same as 4, but real GGUF on CPU	cargo test -- --ignored	minutes
6. Binary E2E (localhost)	full pp-smoke-run --seed driving real pp-gpu-node processes	cargo test -- --ignored	minutes
7. vast.ai REST (mocked)	2-instance create/destroy with mocked HTTP	cargo test	< 1 s
8. vast.ai full run	actual smoke run on rented hardware	manual only	~10 min

---

1. Message Codec — tests/t_codec.rs

Surface: InferenceRequest, InferenceResponse, StageActivation, NextToken defined in §3 of SPEC.

Test name	Proves
stage_activation_roundtrips_through_codec_and_registry	StageActivation with non-trivial binary hidden payload encodes and decodes to identity.
next_token_roundtrips_through_codec_and_registry	NextToken encodes and decodes to identity including done flag.
inference_request_roundtrips_with_max_tokens_field	New max_tokens field on InferenceRequest roundtrips (forked from single-GPU).
inference_response_roundtrips_through_codec_and_registry	Same as single-GPU, kept for parity.
corrupted_bytes_produce_error_for_stage_activation	Random byte flips in encoded StageActivation decode to Err, never panic.
truncated_stage_activation_produces_error	Truncating the hidden blob is rejected with Err.
wrong_message_type_produces_error	Decoding bytes of one type with the codec of another returns Err.
empty_bytes_produce_error	Empty input decodes to Err for every message type.

Reused-as-is codec tests from single-gpu-inference/tests/t_codec.rs are not duplicated.

---

2. Layer-Range and Sampling Math — tests/test_worker.py::TestLayerMath

Surface: the layer-slice computation and argmax sampling described in §4 of SPEC.

Test name	Proves
test_stage_0_layer_range_is_lower_half	STAGE=0 NUM_STAGES=2 yields (0, total // 2).
test_stage_1_layer_range_covers_remainder	STAGE=1 NUM_STAGES=2 yields (total // 2, total), including any odd block.
test_layer_range_partition_is_total_coverage	The union of all stage ranges equals [0, total) with no gaps or overlap, for NUM_STAGES ∈ {2, 3, 4}. (Generalizes early, costs nothing.)
test_argmax_sampling_is_deterministic	Same logits tensor → same token_id, repeated calls.

---

3. Python Worker Contract — tests/test_worker.py

Surface: pp_tinygrad_worker.py stdin/stdout protocol described in §4 of SPEC. Run against a --stub mode that skips real GGUF load.

TestWorkerStartup

Test name	Proves
test_worker_emits_ready_with_pid_and_stage	Startup writes {"status": "ready", "pid": <int>, "stage": <int>} to stdout.
test_worker_rejects_invalid_stage_env	STAGE outside [0, NUM_STAGES) causes immediate exit with non-zero code.

TestStage0Operations (stub mode)

Test name	Proves
test_embed_and_forward_returns_hidden_for_prompt	op=embed_and_forward with a tokenized prompt returns a base64 hidden payload whose decoded byte length matches seq_len * hidden_dim * 2 (bf16).
test_decode_step_returns_hidden_for_single_token	op=decode_step with one token returns hidden bytes of length 1 * hidden_dim * 2.
test_kv_cache_grows_across_successive_decode_steps	Successive decode_step calls at increasing position succeed; the worker does not error or reset state between them.
test_stage_0_rejects_stage_1_ops	Sending op=forward_and_sample to a stage-0 worker returns {"error": ...}, worker stays alive.

TestStage1Operations (stub mode)

Test name	Proves
test_forward_and_sample_returns_valid_token_id	op=forward_and_sample returns a token_id in [0, vocab_size).
test_forward_and_sample_is_deterministic_for_same_input	Same hidden + position → same token_id.
test_stage_1_rejects_stage_0_ops	Sending op=embed_and_forward to a stage-1 worker returns {"error": ...}, worker stays alive.

TestWorkerMalformedInput

Test name	Proves
test_malformed_json_returns_error_and_continues	One garbage line yields an {"error": ...} reply and the worker accepts a valid request immediately after.
test_missing_op_field_returns_error	Valid JSON without op returns {"error": ...}.
test_oversized_hidden_payload_returns_error	hidden whose decoded length disagrees with declared seq_len is rejected.

TestWorkerEOFShutdown

Test name	Proves
test_eof_causes_clean_exit	Closing stdin makes the worker exit 0 within a few seconds.

---

4. Stage Actor — tests/t_actor.rs

Surface: the two new actor types (Stage0Actor, Stage1Actor) and their ProcessBridge to the Python worker. Uses the stub worker.

Stage0Actor

Test name	Proves
stage_0_actor_spawns_worker_and_reports_ready	After spawn, the actor emits WorkerReady { pid } once the stub worker prints its ready line.
inference_request_produces_stage_activation_to_next_address	Sending InferenceRequest causes the actor to emit a StageActivation to its configured next-stage address, with is_prefill=true and seq_len == prompt_token_count.
next_token_produces_stage_activation_for_decode_step	Sending NextToken { done: false } causes the actor to emit StageActivation with is_prefill=false, seq_len=1.
next_token_with_done_produces_no_further_activations	After NextToken { done: true }, the actor emits nothing more for that request.
stage_0_worker_crash_is_handled_without_poisoning_runtime	If the worker process dies, the actor reports the failure and the runtime continues serving other actors.
stopping_stage_0_actor_kills_child_worker	Dropping/stopping the actor cleans up the worker pid.

Stage1Actor

Test name	Proves
stage_1_actor_spawns_worker_and_reports_ready	Mirror of stage 0.
stage_activation_produces_next_token_to_prev_address	Receiving StageActivation causes the actor to emit NextToken to its configured prev-stage address.
stage_1_accumulates_tokens_and_emits_inference_response_on_eos	After enough StageActivations to produce the configured EOS token (stub-controlled), the actor emits InferenceResponse { text } with detokenized output.
stage_1_emits_inference_response_when_max_tokens_reached	After max_tokens activations, the actor emits InferenceResponse even without EOS.
stage_1_worker_crash_is_handled_without_poisoning_runtime	Mirror of stage 0.
stopping_stage_1_actor_kills_child_worker	Mirror of stage 0.

---

5. Name Resolution and Topology — tests/t_topology.rs

Surface: SWIM name registration described in §3 of SPEC (pp-entry, pp-exit, pp-stage-{i}) and per-stage computation of next/prev neighbor names.

Test name	Proves
stage_0_registers_pp_entry_and_pp_stage_0	Stage 0 boot registers both names; both resolve to the same ActorAddress.
stage_1_registers_pp_exit_and_pp_stage_1	Stage 1 boot registers both names; both resolve to the same ActorAddress.
stage_resolves_next_neighbor_after_cluster_join	Stage 0 resolves pp-stage-1 to a non-empty address after the 3-node cluster converges.
stage_resolves_prev_neighbor_after_cluster_join	Stage 1 resolves pp-stage-0 to a non-empty address after the 3-node cluster converges.
next_stage_name_computed_from_env_alone	A stage given STAGE=0, NUM_STAGES=2 computes its outbound target name without any external config.
last_stage_has_no_next_neighbor	Stage NUM_STAGES - 1 does not attempt to resolve a next neighbor.

---

6. Cluster Transport — tests/t_cluster.rs

Surface: 3-node swactor cluster (orchestrator + 2 stages) over real iroh, single-process, separate DistributedNodes.

Test name	Proves
three_node_cluster_converges_via_iroh_seed_join	All three nodes report each other alive within the SWIM convergence window.
stage_activation_roundtrips_stage_0_to_stage_1	A StageActivation sent stage-0 → stage-1 arrives intact (bytes + scalars).
next_token_roundtrips_stage_1_to_stage_0	A NextToken sent stage-1 → stage-0 arrives intact.
inference_response_roundtrips_stage_1_to_orchestrator	An InferenceResponse sent stage-1 → orchestrator arrives intact.
node_death_detected_via_swim_after_stage_shutdown	If a stage actor's host node shuts down, the other two mark it dead within the SWIM timeout.

---

7. vast.ai Client Extensions — tests/t_vastai.rs

Surface: changes needed in the vast.ai client to rent and destroy two instances. Mocks HTTP, does not hit the real API.

Test name	Proves
create_two_instances_sends_distinct_stage_env_vars	Two consecutive create_instance calls send STAGE=0 and STAGE=1 in the env payload respectively.
create_two_instances_both_receive_same_seed_addr	Both creation calls carry the same SEED_ADDR.
failure_to_create_second_instance_triggers_destroy_of_first	If the second create_instance errors, the orchestrator destroys the first instance before exiting.
destroy_two_instances_sends_two_delete_requests	Cleanup issues exactly two DELETEs, one per instance id.
destroy_continues_when_one_delete_fails	A failure to destroy one instance does not prevent the attempt to destroy the other.

Reused-as-is single-instance vast.ai tests in single-gpu-inference/tests/t_vastai.rs are not duplicated.

---

8. Local Integration (in-process, stub worker) — tests/t_integration.rs

Surface: orchestrator + stage 0 + stage 1, each on a separate DistributedNode inside one test process, real iroh, stub worker. Drives the autoregressive loop end-to-end without real model compute.

Test name	Proves
distributed_pipeline_through_stub_workers_returns_response	pp-smoke-run's in-process equivalent sends InferenceRequest { prompt: "Say hello", max_tokens: 4 }, receives non-empty InferenceResponse from stage 1.
decode_loop_terminates_on_stub_eos	Stub worker configured to emit EOS at token N causes the response to arrive with exactly N tokens of accumulated text.
decode_loop_terminates_on_max_tokens	Stub worker that never emits EOS, with max_tokens=4, produces a response of length 4.
stage_failure_mid_decode_surfaces_as_error_to_orchestrator	Killing one stage's worker mid-loop causes the orchestrator's request future to resolve to an error within the SWIM detection window, not hang.

---

9. Sliced-vs-Full Equivalence — tests/t_integration.rs (gated)

The load-bearing correctness test. Surface: that the 2-stage sliced forward pass produces the same outputs as a single-process full forward pass on the same prompt and model. Without this, every other test could pass with silently wrong inference.

Test name	Proves
sliced_two_stage_inference_matches_single_node_full_inference_for_say_hello	Given the same llama3.2:1b GGUF, same argmax sampling, and the prompt "Say hello", the token sequence produced by the 2-stage pipeline equals the token sequence produced by tinygrad's single-node Transformer.generate() up to max_tokens=8.
sliced_two_stage_inference_matches_single_node_for_three_diverse_prompts	Same equivalence holds for three short prompts of different shape (one word, one sentence, one with punctuation). Catches off-by-one errors in position handling, KV cache misalignment, and layer-boundary tensor shape mismatches.

Marked #[ignore]. Runs on CPU with real GGUF; takes minutes.

---

10. Binary E2E (localhost) — tests/t_binary.rs

Surface: the actual pp-smoke-run --seed and pp-gpu-node binaries, spawned as child processes from the test, no vast.ai.

Test name	Proves
binary_e2e_two_stub_workers_returns_hello_response	Running pp-smoke-run --seed against two pp-gpu-node child processes (STAGE=0/STAGE=1, stub worker) prints a non-empty response and exits 0.
binary_e2e_two_stub_workers_cleans_up_on_failure	If one pp-gpu-node is killed mid-run, pp-smoke-run exits non-zero and leaves no orphaned child processes.
binary_e2e_two_tinygrad_workers_returns_response	Same as the first but with the real tinygrad worker on CPU. #[ignore].

---

11. Manual Smoke Run — not a test, documented for completeness

Surface: actual rental on vast.ai. Not run by CI. Documented as a runbook command and an expected-output assertion.

Procedure	Pass criterion
VAST_API_KEY=... pp-smoke-run --vastai --gpu RTX_4090	Exits 0 within 15 minutes; prints a non-empty response; both rented instances appear destroyed in the vast.ai dashboard.

---

Out of Scope (Tests We Are Not Writing)

• Performance / latency / throughput tests. Spec §2 declares no perf targets.
• Fault-tolerance recovery: no retry, no rescheduling, so no test for them.
• N > 2 stage configurations beyond the layer-range-math sanity check in §2. The actor and topology tests are 2-stage only.
• Streaming token tests. No streaming API in the MVP.
• Sampling correctness beyond argmax determinism. No temperature, top-p, etc.
• Concurrent in-flight requests. Single-request only.
• Tests of reused-as-is components from single-gpu-inference (iroh transport, swactor process bridge mechanics, single-instance vast.ai operations). They are tested in that example.