feat(rl): TokenSpeed control endpoint — gRPC/ZMQ engines behind the RL control plane

bindings/python/src/smg/rl.py

@@ -8,9 +8,10 @@
 SGLang requires a JSON body on `pause_generation` and `continue_generation`
 (a bodyless POST is a 400), so bodyless routes are sent as `{}`.
 
-Only HTTP workers can be proxied. A gRPC or ZMQ worker matched by a selector is
-reported in `failed[]` as `unsupported_connection_mode`, which makes `fanout`
-raise `FanoutError` unless `allow_partial=True`.
+A worker with no control endpoint (a gRPC or ZMQ worker without an
+`rl.control_url` label) matched by a selector is reported in `failed[]` as
+`no_control_endpoint`, which makes `fanout` raise `FanoutError` unless
+`allow_partial=True`.
 """
 
 from __future__ import annotations
@@ -52,6 +53,7 @@ class Worker:
     role: str | None
     health: str
     weight_version: str | None
+    control_url: str | None = None
     labels: dict[str, str] = field(default_factory=dict)
     capabilities: dict[str, Any] = field(default_factory=dict)
 

bindings/python/src/smg/serve.py

@@ -8,14 +8,18 @@
 
 import argparse
 import atexit
+import json
 import logging
 import os
 import random
 import signal
 import socket
 import subprocess
 import sys
+import threading
 import time
+import urllib.error
+import urllib.request
 from abc import ABC, abstractmethod
 
 from smg.launch_router import launch_router
@@ -83,6 +87,19 @@ def _zmq_handshake_port(ipc_url: str) -> int:
     return _ZMQ_HANDSHAKE_PORT_BASE + (h % _ZMQ_HANDSHAKE_PORT_SPAN)
 
 
+def _rl_control_port(port: int) -> int:
+    """Port for a TokenSpeed worker's in-engine RL control app.
+
+    Offset from the worker port so co-located workers (consecutive ports) and
+    the engine's own +233 distributed store never collide, reflected below the
+    u16 ceiling and hopped past SMG's ZMQ handshake band like ``dist_port``.
+    """
+    p = port + 400 if port + 400 <= 65535 else port - 400
+    if _ZMQ_HANDSHAKE_PORT_BASE <= p < _ZMQ_HANDSHAKE_PORT_BASE + _ZMQ_HANDSHAKE_PORT_SPAN:
+        p += _ZMQ_HANDSHAKE_PORT_SPAN
+    return p
+
+
 def _reject_handshake_port_collisions(ports: list[int]) -> None:
     """Fail before launch if two workers derive the same ZMQ handshake port.
 
@@ -354,6 +371,10 @@ class TokenspeedWorkerLauncher(WorkerLauncher):
     def _get_tp_size(self, args: argparse.Namespace) -> int:
         return getattr(args, "tensor_parallel_size", 1) or 1
 
+    def control_url(self, port: int) -> str:
+        """URL of the in-engine RL control app this launcher started for ``port``."""
+        return f"http://127.0.0.1:{_rl_control_port(port)}"
+
     def build_command(
         self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
     ) -> list[str]:
@@ -420,6 +441,10 @@ def _build_zmq_command(
             str(rpc_port),
             "--zmq-engine-index",
             "0",
+            "--rl-control-host",
+            "127.0.0.1",
+            "--rl-control-port",
+            str(_rl_control_port(port)),
         ]
         cmd.extend(
             self._backend_arg_defaults(
@@ -449,6 +474,8 @@ def _build_zmq_command(
                     "--data-parallel-address",
                     "--data-parallel-rpc-port",
                     "--zmq-engine-index",
+                    "--rl-control-host",
+                    "--rl-control-port",
                 ],
             )
         )
@@ -547,8 +574,6 @@ def build_command(
 def _http_health_check(url: str, timeout: float) -> bool:
     """GET the URL and return True on HTTP 200."""
     try:
-        import urllib.request
-
         req = urllib.request.Request(url, method="GET")
         with urllib.request.urlopen(req, timeout=timeout) as resp:
             return resp.status == 200
@@ -868,6 +893,60 @@ def parse_serve_args(
 _WORKER_SHUTDOWN_TIMEOUT = 30
 
 
+def _stamp_rl_control_labels(
+    gateway_url: str,
+    api_key: str | None,
+    targets: list[tuple[str, str]],
+    deadline_s: float,
+) -> None:
+    """Label each ZMQ worker with its control endpoint once the gateway lists it.
+
+    ZMQ discovery yields no labels, so the launcher, which owns both ends,
+    stamps ``rl.control_url`` through the worker update route (labels merge).
+    """
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    pending = dict(targets)
+    stop_at = time.monotonic() + deadline_s
+    while pending and time.monotonic() < stop_at:
+        try:
+            req = urllib.request.Request(f"{gateway_url}/workers", headers=headers, method="GET")
+            with urllib.request.urlopen(req, timeout=5) as resp:
+                workers = json.loads(resp.read()).get("workers", [])
+        except Exception as e:  # noqa: BLE001 — the gateway may not be up yet
+            logger.debug("rl.control_url stamping: gateway not ready: %s", e)
+            time.sleep(1)
+            continue
+        for w in workers:
+            url = str(w.get("url", ""))
+            if url not in pending:
+                continue
+            body = json.dumps({"labels": {"rl.control_url": pending[url]}}).encode()
+            req = urllib.request.Request(
+                f"{gateway_url}/workers/{w['id']}", data=body, headers=headers, method="PATCH"
+            )
+            try:
+                with urllib.request.urlopen(req, timeout=5):
+                    pass
+                logger.info("stamped rl.control_url=%s on %s", pending[url], url)
+                del pending[url]
+            except urllib.error.HTTPError as e:
+                if 400 <= e.code < 500:
+                    logger.warning(
+                        "rl.control_url stamping got HTTP %s for %s; not retrying", e.code, url
+                    )
+                    del pending[url]
+                else:
+                    logger.warning("rl.control_url stamping failed for %s: %s", url, e)
+            except Exception as e:  # noqa: BLE001
+                logger.warning("rl.control_url stamping failed for %s: %s", url, e)
+        if pending:
+            time.sleep(1)
+    for url in pending:
+        logger.warning("rl.control_url never stamped on %s (gateway did not list it)", url)
+
+
 class ServeOrchestrator:
     """Coordinate worker launch, health checking, router startup, and shutdown."""
 
@@ -890,6 +969,27 @@ def run(self) -> None:
             self._launch_workers()
             self._wait_healthy()
             router_args = self._build_router_args()
+            if (
+                getattr(router_args, "enable_rl", False)
+                and self.backend == "tokenspeed"
+                and getattr(self.args, "connection_mode", "grpc") == "zmq"
+            ):
+                control = getattr(self.launcher, "control_url", None)
+                if callable(control):
+                    targets = [
+                        (
+                            self.launcher.worker_url(self.args, self.args.worker_host, port),
+                            control(port),
+                        )
+                        for _, port in self.workers
+                    ]
+                    gateway_url = f"http://127.0.0.1:{router_args.port}"
+                    threading.Thread(
+                        target=_stamp_rl_control_labels,
+                        args=(gateway_url, getattr(router_args, "api_key", None), targets, 300.0),
+                        name="smg-rl-control-labels",
+                        daemon=True,
+                    ).start()
             launch_router(router_args)
         finally:
             self._cleanup_workers()

bindings/python/tests/test_refit_from_trainer_layout.py

@@ -0,0 +1,124 @@
+"""The trainer-side refit example's rank layout (no gateway, no GPU, no engine).
+
+`rank_layout` decides which ranks of the trainer's NCCL group each engine
+occupies, which is the one piece of `examples/rl/refit_from_trainer.py` that is
+pure arithmetic and the one an off-by-one silently deadlocks.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+REPO = Path(__file__).resolve().parents[3]
+SCRIPT = REPO / "examples" / "rl" / "refit_from_trainer.py"
+SRC = str(REPO / "bindings" / "python" / "src")
+
+
+def _smg_modules() -> list[str]:
+    return [name for name in sys.modules if name == "smg" or name.startswith("smg.")]
+
+
+@contextmanager
+def _loaded_script() -> Iterator[Any]:
+    """Import the example by path, with this checkout's `smg.rl` importable.
+
+    The example imports `smg.rl`, torch and transformers at the top, the way a
+    trainer would, and `smg` is usually already imported from wherever the
+    package was installed -- a checkout that need not be this one, and whose
+    compiled `smg_rs` extension is not in this source tree. So point `smg` at
+    this checkout's sources for the import and put the interpreter back
+    afterwards: leaving either `sys.path` or `sys.modules` shifted breaks every
+    later test in the session that imports `smg.smg_rs`.
+    """
+    saved_path = list(sys.path)
+    saved_modules = {name: sys.modules[name] for name in _smg_modules()}
+    try:
+        sys.path.insert(0, SRC)
+        for name in list(saved_modules):
+            del sys.modules[name]
+        spec = importlib.util.spec_from_file_location("refit_from_trainer", SCRIPT)
+        assert spec and spec.loader
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        yield module
+    finally:
+        sys.path[:] = saved_path
+        for name in _smg_modules():
+            del sys.modules[name]
+        sys.modules.update(saved_modules)
+
+
+@pytest.fixture(scope="module")
+def script() -> Iterator[Any]:
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    with _loaded_script() as module:
+        yield module
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("tp_sizes", "expected"),
+    [
+        ([], (1, [])),
+        ([1], (2, [1])),
+        ([1, 1], (3, [1, 2])),
+        ([2, 4], (7, [1, 3])),
+    ],
+)
+def test_rank_layout(script, tp_sizes, expected):
+    assert script.rank_layout(tp_sizes) == expected
+
+
+@pytest.mark.unit
+def test_rank_layout_rejects_a_nonsense_tp_size(script):
+    with pytest.raises(ValueError, match="tp_size must be >= 1"):
+        script.rank_layout([1, 0])
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("host", "loopback"),
+    [
+        ("127.0.0.1", True),
+        ("127.1.2.3", True),
+        ("localhost", True),
+        ("::1", True),
+        ("[::1]", True),
+        ("0.0.0.0", True),
+        ("10.0.1.52", False),
+        ("trainer.internal", False),
+        (None, False),
+        ("", False),
+    ],
+)
+def test_is_loopback(script, host, loopback):
+    assert script._is_loopback(host) is loopback
+
+
+@pytest.mark.unit
+def test_loading_the_example_restores_the_interpreter():
+    """Importing the example must not leave `smg` pointed at this source tree.
+
+    Every other test file in this suite imports `smg.smg_rs` from inside a test
+    body, and this source tree has no compiled extension, so a leaked
+    `sys.path` entry or a purged `sys.modules` entry fails all of them.
+    """
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    before_path = list(sys.path)
+    before_modules = {name: sys.modules[name] for name in _smg_modules()}
+
+    with _loaded_script() as module:
+        assert module.rank_layout([1]) == (2, [1])
+        assert sys.path[0] == SRC
+
+    assert sys.path == before_path
+    assert {name: sys.modules[name] for name in _smg_modules()} == before_modules

bindings/python/tests/test_rl_client.py

@@ -19,6 +19,7 @@
     "model_id": "m",
     "worker_type": "regular",
     "connection_mode": "http",
+    "control_url": "http://e:1",
     "tp_size": 1,
     "dp_size": 1,
     "pp_size": 1,
@@ -89,6 +90,7 @@ def test_workers_and_worker(stub):
     ws = rl.workers()
     assert len(ws) == 1 and ws[0].id == "w1" and ws[0].engine == "sglang"
     assert ws[0].capabilities["pause_modes"] == ["abort"]
+    assert ws[0].control_url == "http://e:1"
     assert rl.worker("w1").weight_version == "7"
     assert _Stub.seen[0]["auth"] == "Bearer k"
 
@@ -155,11 +157,13 @@ def test_worker_from_json_defaults_missing_dicts():
     del d["labels"]
     del d["capabilities"]
     del d["role"]
+    del d["control_url"]
     d["future_field"] = 1
     w = Worker.from_json(d)
     assert w.labels == {}
     assert w.capabilities == {}
     assert w.role is None
+    assert w.control_url is None
 
 
 def test_call_raises_on_smg_error_envelope(stub):