{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Switching & Cancel\n",
    "\n",
    "This notebook shows how to:\n",
    "- Route orders to the right product using `product_id`\n",
    "- Switch models mid-training using `discard()`\n",
    "- Cancel individual orders with `cancel()`\n",
    "- Replace a supplier gracefully"
   ],
   "id": "cell-0"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "\n",
    "from softs import (\n",
    "    Broker, Supplier, Client, ShmMedium,\n",
    "    BatchConfig, TensorSpec, EndpointConfig,\n",
    ")"
   ],
   "id": "cell-1"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "One supplier serves 3 product ids. Each product produces different data\n",
    "so we can verify routing is correct."
   ],
   "id": "cell-2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = BatchConfig([TensorSpec(\"x\", (4,), \"float32\")])\n",
    "endpoints = EndpointConfig()\n",
    "\n",
    "broker = Broker(endpoints=endpoints)\n",
    "broker.start()\n",
    "time.sleep(0.2)\n",
    "\n",
    "# Supplier produces product-dependent data\n",
    "MODEL_VALUES = {\"layer_0\": 0.0, \"layer_1\": 1.0, \"layer_2\": 2.0}\n",
    "\n",
    "def generator(product_id: str) -> bytes:\n",
    "    val = MODEL_VALUES[product_id]\n",
    "    return config.encode(x=torch.full((4,), val))\n",
    "\n",
    "supplier = Supplier(\n",
    "    generator_fn=generator,\n",
    "    product_ids=list(MODEL_VALUES.keys()),\n",
    "    endpoint=endpoints.backend,\n",
    "    medium_cls=ShmMedium,\n",
    "    slot_size=config.nbytes(),\n",
    ")\n",
    "supplier.start()\n",
    "time.sleep(0.2)\n",
    "\n",
    "client = Client(\n",
    "    endpoint=endpoints.frontend,\n",
    "    medium_cls=ShmMedium,\n",
    "    slot_size=config.nbytes(),\n",
    "    num_slots=8,\n",
    ")\n",
    "client.hello()\n",
    "print(\"Broker started, supplier ready.\")"
   ],
   "id": "cell-3"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model-aware routing\n",
    "\n",
    "Orders are routed to the correct product. Each product returns its own value."
   ],
   "id": "cell-4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for product_id, expected_val in MODEL_VALUES.items():\n",
    "    slot = client.request_sample(product_id, timeout_ms=2000)\n",
    "    data = config.decode(client.medium.read(slot))\n",
    "    client.release_slot(slot)\n",
    "    assert data[\"x\"][0].item() == expected_val\n",
    "    print(f\"{product_id}: {data['x']}\")"
   ],
   "id": "cell-5"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model switching with `discard()`\n",
    "\n",
    "Simulate layer-by-layer training: train on `layer_0`, then switch to `layer_1`.\n",
    "`discard()` cancels all pending orders and drains stale completions."
   ],
   "id": "cell-6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for layer_idx in range(3):\n",
    "    product_id = f\"layer_{layer_idx}\"\n",
    "    print(f\"\\nTraining on {product_id}:\")\n",
    "\n",
    "    for step in range(3):\n",
    "        slot = client.request_sample(product_id, timeout_ms=2000)\n",
    "        data = config.decode(client.medium.read(slot))\n",
    "        client.release_slot(slot)\n",
    "        assert data[\"x\"][0].item() == float(layer_idx)\n",
    "        print(f\"  step {step}: x[0]={data['x'][0].item()}\")\n",
    "\n",
    "    # Switch to next layer: discard old work\n",
    "    cancelled = client.discard()\n",
    "    print(f\"  Discarded {cancelled} pending orders\")"
   ],
   "id": "cell-7"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cancel a specific order"
   ],
   "id": "cell-8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "order_id = client.request_slot(\"layer_0\")\n",
    "print(f\"Requested order: {order_id}\")\n",
    "time.sleep(0.1)\n",
    "\n",
    "ok = client.cancel(order_id)\n",
    "print(f\"Cancelled: {ok}\")\n",
    "print(f\"Pending slots after cancel: {len(client._pending_slots)}\")"
   ],
   "id": "cell-9"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supplier replacement\n",
    "\n",
    "Stop the current supplier, start a new one. The broker detects the GOODBYE\n",
    "immediately and routes to the replacement."
   ],
   "id": "cell-10"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supplier.stop()\n",
    "time.sleep(0.2)\n",
    "print(f\"Suppliers after stop: {broker.get_stats().connected_suppliers}\")\n",
    "\n",
    "# Start replacement\n",
    "supplier2 = Supplier(\n",
    "    generator_fn=generator,\n",
    "    product_ids=list(MODEL_VALUES.keys()),\n",
    "    endpoint=endpoints.backend,\n",
    "    medium_cls=ShmMedium,\n",
    "    slot_size=config.nbytes(),\n",
    ")\n",
    "supplier2.start()\n",
    "time.sleep(0.2)\n",
    "print(f\"Suppliers after new start: {broker.get_stats().connected_suppliers}\")\n",
    "\n",
    "slot = client.request_sample(\"layer_0\", timeout_ms=2000)\n",
    "data = config.decode(client.medium.read(slot))\n",
    "client.release_slot(slot)\n",
    "print(f\"New supplier serving: {data['x']}\")"
   ],
   "id": "cell-11"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graceful shutdown"
   ],
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "supplier2.stop()\n",
    "time.sleep(0.2)\n",
    "broker.stop()\n",
    "print(\"Shutdown complete.\")"
   ],
   "id": "cell-13"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}