Merge branch 'main' into lora

Author: MikaStars39

areal/engine/vllm_ext/areal_vllm_server.py

@@ -152,12 +152,33 @@ async def update_weight_xccl(raw_request: Request):
 
 
 @router.post("/areal_update_weights_lora_xccl")
-async def update_weight_lora_xccl(raw_request: Request):
+async def update_weight_lora_xccl(
+    request: UpdateWeightsFromXcclRequestLora, raw_request: Request
+):
     logger.info("API server starts update_weight_lora via XCCL")
     llm = raw_request.app.state.engine_client
     ret_list = await llm.engine_core.call_utility_async(
         "areal_injected_update_weight_lora_xccl",
     )
+    # Only touch the registry after weights are actually updated
+    models_obj = raw_request.app.state.openai_serving_models
+    new_name = request.lora_name
+    lora_id = request.lora_int_id
+    for old_name, req in list(models_obj.lora_requests.items()):
+        if req.lora_int_id == lora_id:
+            del models_obj.lora_requests[old_name]
+            req.lora_name = new_name
+            models_obj.lora_requests[new_name] = req
+            logger.info(
+                f"Updated LoRA name of openai_serving_models "
+                f"from {old_name} -> {new_name}"
+            )
+            break
+    else:
+        logger.warning(
+            f"LoRA adapter with int_id={lora_id} not found in "
+            f"openai_serving_models.lora_requests"
+        )
     return build_response(ret_list)
 
 

areal/engine/vllm_remote.py

@@ -185,7 +185,7 @@ def build_distributed_weight_update_requests(
                 ),
                 HttpRequest(
                     endpoint=update_endpoint,
-                    payload={},
+                    payload={} if not meta.use_lora else payload,
                 ),
             ]
         )

areal/infra/rpc/rpc_server.py

@@ -431,15 +431,15 @@ def configure():
     try:
         data = request.get_json()
         if data is None:
-            return jsonify({"detail": "Invalid JSON in request body"}), 400
+            return jsonify({"error": "Invalid JSON in request body"}), 400
 
         config = data.get("config")
         if config is None:
-            return jsonify({"detail": "Missing 'config' field in request"}), 400
+            return jsonify({"error": "Missing 'config' field in request"}), 400
 
         rank = data.get("rank")
         if rank is None:
-            return jsonify({"detail": "Missing 'rank' field in request"}), 400
+            return jsonify({"error": "Missing 'rank' field in request"}), 400
 
         config = deserialize_value(config)
         config: BaseExperimentConfig

areal/infra/scheduler/local.py

@@ -866,10 +866,10 @@ def _configure_worker(self, worker_info: WorkerInfo, worker_rank: int):
                 logger.info(f"Configuration successfully on worker '{worker_id}'")
                 return
             elif response.status_code == 400:
-                error_detail = response.json().get("detail", "Unknown error")
+                error_detail = response.json().get("error", "Unknown error")
                 raise WorkerConfigurationError(worker_id, error_detail, str(400))
             elif response.status_code == 500:
-                error_detail = response.json().get("detail", "Unknown error")
+                error_detail = response.json().get("error", "Unknown error")
                 raise WorkerConfigurationError(worker_id, error_detail, str(500))
             else:
                 raise WorkerConfigurationError(
@@ -1118,7 +1118,7 @@ async def create_engine(
                     elif response.status == 400:
                         # Import error or bad request
                         error_detail = (await response.json()).get(
-                            "detail", "Unknown error"
+                            "error", "Unknown error"
                         )
                         if "Failed to import" in error_detail:
                             raise EngineImportError(engine, error_detail)
@@ -1127,7 +1127,7 @@ async def create_engine(
                     elif response.status == 500:
                         # Engine initialization failed
                         error_detail = (await response.json()).get(
-                            "detail", "Unknown error"
+                            "error", "Unknown error"
                         )
                         raise EngineCreationError(worker_id, error_detail, 500)
                     else:
@@ -1399,15 +1399,15 @@ async def async_call_engine(
                         elif response.status == 400:
                             # Bad request (e.g., method doesn't exist) - don't retry
                             error_detail = (await response.json()).get(
-                                "detail", "Unknown error"
+                                "error", "Unknown error"
                             )
                             raise EngineCallError(
                                 worker_id, method, error_detail, attempt
                             )
                         elif response.status == 500:
                             # Engine method failed - don't retry
                             error_detail = (await response.json()).get(
-                                "detail", "Unknown error"
+                                "error", "Unknown error"
                             )
                             raise EngineCallError(
                                 worker_id, method, error_detail, attempt
@@ -1492,11 +1492,11 @@ def _handle_call_response(
             return deserialized_result, False, None
         elif response.status_code == 400:
             # Bad request (e.g., method doesn't exist) - don't retry
-            error_detail = response.json().get("detail", "Unknown error")
+            error_detail = response.json().get("error", "Unknown error")
             raise EngineCallError(worker_id, method, error_detail, attempt)
         elif response.status_code == 500:
             # Engine method failed - don't retry
-            error_detail = response.json().get("detail", "Unknown error")
+            error_detail = response.json().get("error", "Unknown error")
             raise EngineCallError(worker_id, method, error_detail, attempt)
         elif response.status_code == 503:
             # Service unavailable - retry

areal/infra/scheduler/slurm.py

@@ -309,10 +309,10 @@ def _configure_worker(self, worker_info: SlurmWorkerInfo, worker_rank: int) -> N
                 logger.info(f"Configuration successful on worker '{worker_id}'")
                 return
             elif response.status_code == 400:
-                error_detail = response.json().get("detail", "Unknown error")
+                error_detail = response.json().get("error", "Unknown error")
                 raise WorkerConfigurationError(worker_id, error_detail, str(400))
             elif response.status_code == 500:
-                error_detail = response.json().get("detail", "Unknown error")
+                error_detail = response.json().get("error", "Unknown error")
                 raise WorkerConfigurationError(worker_id, error_detail, str(500))
             else:
                 raise WorkerConfigurationError(
@@ -1329,15 +1329,15 @@ async def create_engine(
                         return result.get("result")
                     elif response.status == 400:
                         error_detail = (await response.json()).get(
-                            "detail", "Unknown error"
+                            "error", "Unknown error"
                         )
                         if "Failed to import" in error_detail:
                             raise EngineImportError(engine, error_detail)
                         else:
                             raise EngineCreationError(worker_id, error_detail, 400)
                     elif response.status == 500:
                         error_detail = (await response.json()).get(
-                            "detail", "Unknown error"
+                            "error", "Unknown error"
                         )
                         raise EngineCreationError(worker_id, error_detail, 500)
                     else:
@@ -1439,7 +1439,7 @@ def call_engine(
                     result = response.json()
                     return deserialize_value(result.get("result"))
                 elif response.status_code == 500:
-                    error_detail = response.json().get("detail", "Unknown error")
+                    error_detail = response.json().get("error", "Unknown error")
                     # Check if retryable
                     if attempt < max_retries and "timeout" in error_detail.lower():
                         last_error = f"Engine method timeout: {error_detail}"
@@ -1457,7 +1457,7 @@ def call_engine(
                         f"Worker temporarily unavailable, retry {attempt}/{max_retries}"
                     )
                 else:
-                    error_detail = response.json().get("detail", "Unknown error")
+                    error_detail = response.json().get("error", "Unknown error")
                     raise EngineCallError(
                         worker_id,
                         method,
@@ -1580,7 +1580,7 @@ async def async_call_engine(
                             return deserialize_value(result.get("result"))
                         elif response.status == 500:
                             error_detail = (await response.json()).get(
-                                "detail", "Unknown error"
+                                "error", "Unknown error"
                             )
                             if (
                                 attempt < max_retries
@@ -1601,7 +1601,7 @@ async def async_call_engine(
                             )
                         else:
                             error_detail = (await response.json()).get(
-                                "detail", "Unknown error"
+                                "error", "Unknown error"
                             )
                             raise EngineCallError(
                                 worker_id,

tests/test_local_scheduler.py

@@ -1315,7 +1315,7 @@ def test_create_engine_import_error(self, scheduler, tmp_path):
         mock_response = AsyncMock()
         mock_response.status = 400
         mock_response.json = AsyncMock(
-            return_value={"detail": "Failed to import 'nonexistent.Engine'"}
+            return_value={"error": "Failed to import 'nonexistent.Engine'"}
         )
         mock_response.__aenter__.return_value = mock_response
         mock_response.__aexit__.return_value = None
@@ -1343,7 +1343,7 @@ def test_create_engine_initialization_error(self, scheduler, tmp_path):
         mock_response = AsyncMock()
         mock_response.status = 500
         mock_response.json = AsyncMock(
-            return_value={"detail": "Engine initialization failed: out of memory"}
+            return_value={"error": "Engine initialization failed: out of memory"}
         )
         mock_response.__aenter__.return_value = mock_response
         mock_response.__aexit__.return_value = None
@@ -1486,7 +1486,7 @@ def test_call_engine_method_error(self, scheduler, tmp_path):
         scheduler._workers["test"] = [worker]
 
         mock_response = create_mock_http_response(
-            status_code=400, json_data={"detail": "Method 'nonexistent' not found"}
+            status_code=400, json_data={"error": "Method 'nonexistent' not found"}
         )
 
         with patch.object(requests, "post", return_value=mock_response):