feat: add disk-mode weight update flow to gateway (#1237)

* feat: add disk-mode weight update flow to gateway

Support disk-based weight transfer for weight-update pairs, including save-to-disk on train workers and load/update on inference workers. Ensure inference pause/continue is wrapped in a context manager so generation resumes even when updates fail.

Key changes:

- add disk mode fields to connect request and pair metadata

- implement disk transfer path with save/update_weights_from_disk and LoRA loading

- add controller and gateway tests for disk mode behavior

* fix: validate disk-mode gateway connect inputs

Fail fast on disk-mode misconfiguration so weight updates do not silently write to local relative paths instead of shared storage.

Key changes:
- reject empty disk-mode save_path values
- require absolute save_path values for disk mode
- require lora_name when disk-mode LoRA is enabled
- add gateway tests for rejected connect requests

---------

Co-authored-by: Wentai Zhang <rchardx@gmail.com>

Author: garrett4wade

areal/api/io_struct.py

@@ -165,7 +165,7 @@ def get_versioned_lora_name(lora_name: str, version: int) -> str:
 
 @dataclass
 class WeightUpdateMeta:
-    type: Literal["disk", "xccl"]
+    type: Literal["disk", "xccl", "awex"]
     path: str | None = None
     gen_allocation: ModelAllocation | None = None
 
@@ -267,6 +267,22 @@ def from_fsdp_xccl(
             base_model_name=base_model_name,
         )
 
+    @classmethod
+    def from_awex(
+        cls,
+        use_lora: bool = False,
+        lora_name: str = "",
+        lora_int_id: int = 1,
+        base_model_name: str = "",
+    ):
+        return cls(
+            type="awex",
+            use_lora=use_lora,
+            lora_name=lora_name,
+            lora_int_id=lora_int_id,
+            base_model_name=base_model_name,
+        )
+
 
 @dataclass
 class HttpRequest:

areal/experimental/weight_update/controller/controller.py

@@ -108,6 +108,10 @@ def connect(
         pair_name: str,
         train_worker_urls: list[str],
         inference_worker_urls: list[str],
+        mode: str = "awex",
+        save_path: str = "",
+        use_lora: bool = False,
+        lora_name: str = "",
     ) -> None:
         self._pair_name = pair_name
         resp = self._http.post(
@@ -116,11 +120,15 @@ def connect(
                 "pair_name": pair_name,
                 "train_worker_urls": train_worker_urls,
                 "inference_worker_urls": inference_worker_urls,
+                "mode": mode,
+                "save_path": save_path,
+                "use_lora": use_lora,
+                "lora_name": lora_name,
             },
             timeout=self.config.request_timeout,
         )
         resp.raise_for_status()
-        logger.info("Connected pair '%s'", pair_name)
+        logger.info("Connected pair '%s' (mode=%s)", pair_name, mode)
 
     def update_weights(self, version: int) -> WeightUpdateResult:
         if self._pair_name is None:

areal/experimental/weight_update/gateway/app.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import socket
 import time
 from contextlib import asynccontextmanager
@@ -29,6 +30,10 @@ class ConnectRequest(BaseModel):
     pair_name: str
     train_worker_urls: list[str]
     inference_worker_urls: list[str]
+    mode: str = "awex"  # "awex" or "disk"
+    save_path: str = ""
+    use_lora: bool = False
+    lora_name: str = ""
 
 
 class UpdateWeightsRequest(BaseModel):
@@ -191,6 +196,39 @@ async def connect(request: Request, body: ConnectRequest) -> ConnectResponse:
         train_urls = body.train_worker_urls
         inference_urls = body.inference_worker_urls
 
+        if body.mode == "disk":
+            if not body.save_path:
+                return JSONResponse(
+                    status_code=400,
+                    content={"error": "save_path is required when mode='disk'"},
+                )
+            if not os.path.isabs(body.save_path):
+                return JSONResponse(
+                    status_code=400,
+                    content={
+                        "error": "save_path must be an absolute path when mode='disk'"
+                    },
+                )
+            if body.use_lora and not body.lora_name:
+                return JSONResponse(
+                    status_code=400,
+                    content={"error": "lora_name is required when use_lora=True"},
+                )
+            pair_info = PairInfo(
+                pair_name=pair_name,
+                train_worker_urls=train_urls,
+                inference_worker_urls=inference_urls,
+                mode="disk",
+                save_path=body.save_path,
+                use_lora=body.use_lora,
+                lora_name=body.lora_name,
+            )
+            registry.register(pair_info)
+            logger.info(
+                "Connected disk pair '%s' (save_path=%s)", pair_name, body.save_path
+            )
+            return ConnectResponse(pair_name=pair_name)
+
         session = request.app.state.http_session
         init_timeout_s = config.init_timeout_s
 
@@ -325,6 +363,114 @@ async def connect(request: Request, body: ConnectRequest) -> ConnectResponse:
         logger.info("Connected pair '%s'", pair_name)
         return ConnectResponse(pair_name=pair_name)
 
+    @asynccontextmanager
+    async def _inference_paused(
+        session: aiohttp.ClientSession,
+        inference_urls: list[str],
+        timeout_s: float,
+        pair_name: str,
+    ):
+        await asyncio.gather(
+            *[
+                _post(session, f"{url}/pause_generation", timeout_s, json_data={})
+                for url in inference_urls
+            ]
+        )
+        try:
+            yield
+        finally:
+            try:
+                await asyncio.gather(
+                    *[
+                        _post(
+                            session,
+                            f"{url}/continue_generation",
+                            timeout_s,
+                            json_data={},
+                        )
+                        for url in inference_urls
+                    ]
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to resume inference for pair '%s'",
+                    pair_name,
+                    exc_info=True,
+                )
+
+    async def _awex_transfer_weights(
+        pair_info: PairInfo,
+        version: int,
+        session: aiohttp.ClientSession,
+        timeout_s: float,
+    ) -> None:
+        await asyncio.gather(
+            *[
+                _post(
+                    session,
+                    f"{url}/awex/update_weights",
+                    timeout_s,
+                    json_data={"version": version},
+                )
+                for url in pair_info.train_worker_urls + pair_info.inference_worker_urls
+            ]
+        )
+
+    async def _disk_transfer_weights(
+        pair_info: PairInfo,
+        version: int,
+        session: aiohttp.ClientSession,
+        timeout_s: float,
+    ) -> None:
+        from areal.api.io_struct import SaveLoadMeta, get_versioned_lora_name
+        from areal.infra.rpc.serialization import serialize_value
+
+        save_path = os.path.join(pair_info.save_path, f"weight_update_v{version}")
+
+        save_meta = SaveLoadMeta(path=save_path, weight_format="hf", with_optim=False)
+        save_payload = {
+            "args": serialize_value([save_meta]),
+            "kwargs": serialize_value({}),
+        }
+        await asyncio.gather(
+            *[
+                _post_json(session, f"{url}/save", timeout_s, json_data=save_payload)
+                for url in pair_info.train_worker_urls
+            ]
+        )
+
+        if pair_info.use_lora:
+            lora_name = get_versioned_lora_name(pair_info.lora_name, version)
+            await asyncio.gather(
+                *[
+                    _post_json(
+                        session,
+                        f"{url}/load_lora_adapter",
+                        timeout_s,
+                        json_data={
+                            "lora_name": lora_name,
+                            "lora_path": save_path,
+                        },
+                    )
+                    for url in pair_info.inference_worker_urls
+                ]
+            )
+        else:
+            await asyncio.gather(
+                *[
+                    _post_json(
+                        session,
+                        f"{url}/update_weights_from_disk",
+                        timeout_s,
+                        json_data={
+                            "model_path": save_path,
+                            "abort_all_requests": True,
+                        },
+                    )
+                    for url in pair_info.inference_worker_urls
+                ]
+            )
+
     @app.post("/update_weights")
     async def update_weights(
         request: Request, body: UpdateWeightsRequest
@@ -334,27 +480,51 @@ async def update_weights(
         pair_info = registry.get_by_name(body.pair_name)
         if pair_info is None:
             return JSONResponse(
-                status_code=404, content={"error": f"Pair '{body.pair_name}' not found"}
+                status_code=404,
+                content={"error": f"Pair '{body.pair_name}' not found"},
             )
 
         session = request.app.state.http_session
-        update_timeout_s = config.update_timeout_s
-
+        timeout_s = config.update_timeout_s
         start = time.monotonic()
-        tasks = [
-            _post(
+
+        try:
+            async with _inference_paused(
                 session,
-                f"{url}/awex/update_weights",
-                update_timeout_s,
-                json_data={"version": body.version},
+                pair_info.inference_worker_urls,
+                timeout_s,
+                pair_info.pair_name,
+            ):
+                if pair_info.mode == "disk":
+                    await _disk_transfer_weights(
+                        pair_info, body.version, session, timeout_s
+                    )
+                else:
+                    await _awex_transfer_weights(
+                        pair_info, body.version, session, timeout_s
+                    )
+        except Exception as e:
+            duration_ms = (time.monotonic() - start) * 1000
+            logger.error(
+                "Weight update failed for pair '%s': %s",
+                pair_info.pair_name,
+                e,
+            )
+            return WeightUpdateResult(
+                status="error",
+                version=body.version,
+                duration_ms=duration_ms,
+                error=str(e),
             )
-            for url in pair_info.train_worker_urls + pair_info.inference_worker_urls
-        ]
-        await asyncio.gather(*tasks)
 
         duration_ms = (time.monotonic() - start) * 1000
         pair_info.last_version = body.version
-
+        logger.info(
+            "Weight update completed for pair '%s' v%d (%.1fms)",
+            pair_info.pair_name,
+            body.version,
+            duration_ms,
+        )
         return WeightUpdateResult(
             status="ok", version=body.version, duration_ms=duration_ms
         )

areal/experimental/weight_update/gateway/config.py

@@ -54,6 +54,14 @@ class PairInfo:
     master_port: int = 0
     last_version: int = 0
 
+    # Disk-mode fields (used when mode="disk")
+    mode: str = "awex"  # "awex" or "disk"
+    save_path: str = ""
+    use_lora: bool = False
+    lora_name: str = ""
+
     def __post_init__(self):
         if not self.pair_name:
             raise ValueError("pair_name must not be empty")
+        if self.mode not in ("awex", "disk"):
+            raise ValueError(f"mode must be 'awex' or 'disk', got '{self.mode}'")