From 76e1c2b82686576742d1d32b3d1200014946de72 Mon Sep 17 00:00:00 2001
From: Bohdan <bkovalevski@coreweave.com>
Date: Wed, 21 Jan 2026 13:57:19 -0800
Subject: [PATCH] Fix model registration

---
 src/art/dev/openai_server.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/art/dev/openai_server.py b/src/art/dev/openai_server.py
index 2f28560f..e6f400d1 100644
--- a/src/art/dev/openai_server.py
+++ b/src/art/dev/openai_server.py
@@ -18,18 +18,19 @@ def get_openai_server_config(
         config = OpenAIServerConfig()
     log_file = config.get("log_file", log_file)
 
-    # Extract step from lora_path for multi-checkpoint support
-    # lora_path format is: {output_dir}/checkpoints/{step:04d}
-    lora_name = model_name
+    # Build LoRA modules list for multi-checkpoint support
+    # Register under both model_name (for "current" model) and model_name@step (for specific checkpoint)
+    lora_modules: list[str] | None = None
     if lora_path:
         step = int(os.path.basename(lora_path))
-        lora_name = f"{model_name}@{step}"
+        lora_modules = [
+            f'{{"name": "{model_name}", "path": "{lora_path}"}}',
+            f'{{"name": "{model_name}@{step}", "path": "{lora_path}"}}',
+        ]
 
     server_args = ServerArgs(
         api_key="default",
-        lora_modules=(
-            [f'{{"name": "{lora_name}", "path": "{lora_path}"}}'] if lora_path else None
-        ),
+        lora_modules=lora_modules,
         return_tokens_as_token_ids=True,
         enable_auto_tool_choice=True,
         tool_call_parser="hermes",
@@ -37,7 +38,7 @@ def get_openai_server_config(
     server_args.update(config.get("server_args", {}))
     engine_args = EngineArgs(
         model=base_model,
-        served_model_name=base_model if lora_path else model_name,
+        served_model_name=model_name,
         generation_config="vllm",
     )
     engine_args.update(config.get("engine_args", {}))