Model Switching & Cancel

This notebook shows how to:

• Route orders to the right product using product_id

• Switch models mid-training using discard()

• Cancel individual orders with cancel()

• Replace a supplier gracefully

import time
import torch

from softs import (
    Broker, Supplier, Client, ShmMedium,
    BatchConfig, TensorSpec, EndpointConfig,
)

Setup

One supplier serves 3 product ids. Each product produces different data so we can verify routing is correct.

config = BatchConfig([TensorSpec("x", (4,), "float32")])
endpoints = EndpointConfig()

broker = Broker(endpoints=endpoints)
broker.start()
time.sleep(0.2)

# Supplier produces product-dependent data
MODEL_VALUES = {"layer_0": 0.0, "layer_1": 1.0, "layer_2": 2.0}

def generator(product_id: str) -> bytes:
    val = MODEL_VALUES[product_id]
    return config.encode(x=torch.full((4,), val))

supplier = Supplier(
    generator_fn=generator,
    product_ids=list(MODEL_VALUES.keys()),
    endpoint=endpoints.backend,
    medium_cls=ShmMedium,
    slot_size=config.nbytes(),
)
supplier.start()
time.sleep(0.2)

client = Client(
    endpoint=endpoints.frontend,
    medium_cls=ShmMedium,
    slot_size=config.nbytes(),
    num_slots=8,
)
client.hello()
print("Broker started, supplier ready.")

Model-aware routing

Orders are routed to the correct product. Each product returns its own value.

for product_id, expected_val in MODEL_VALUES.items():
    slot = client.request_sample(product_id, timeout_ms=2000)
    data = config.decode(client.medium.read(slot))
    client.release_slot(slot)
    assert data["x"][0].item() == expected_val
    print(f"{product_id}: {data['x']}")

Model switching with discard()

Simulate layer-by-layer training: train on layer_0, then switch to layer_1. discard() cancels all pending orders and drains stale completions.

for layer_idx in range(3):
    product_id = f"layer_{layer_idx}"
    print(f"\nTraining on {product_id}:")

    for step in range(3):
        slot = client.request_sample(product_id, timeout_ms=2000)
        data = config.decode(client.medium.read(slot))
        client.release_slot(slot)
        assert data["x"][0].item() == float(layer_idx)
        print(f"  step {step}: x[0]={data['x'][0].item()}")

    # Switch to next layer: discard old work
    cancelled = client.discard()
    print(f"  Discarded {cancelled} pending orders")

Cancel a specific order

order_id = client.request_slot("layer_0")
print(f"Requested order: {order_id}")
time.sleep(0.1)

ok = client.cancel(order_id)
print(f"Cancelled: {ok}")
print(f"Pending slots after cancel: {len(client._pending_slots)}")

Supplier replacement

Stop the current supplier, start a new one. The broker detects the GOODBYE immediately and routes to the replacement.

supplier.stop()
time.sleep(0.2)
print(f"Suppliers after stop: {broker.get_stats().connected_suppliers}")

# Start replacement
supplier2 = Supplier(
    generator_fn=generator,
    product_ids=list(MODEL_VALUES.keys()),
    endpoint=endpoints.backend,
    medium_cls=ShmMedium,
    slot_size=config.nbytes(),
)
supplier2.start()
time.sleep(0.2)
print(f"Suppliers after new start: {broker.get_stats().connected_suppliers}")

slot = client.request_sample("layer_0", timeout_ms=2000)
data = config.decode(client.medium.read(slot))
client.release_slot(slot)
print(f"New supplier serving: {data['x']}")

Graceful shutdown

client.close()
supplier2.stop()
time.sleep(0.2)
broker.stop()
print("Shutdown complete.")