Fix ART ty failures in local backend and MoE conversion

Author: Kovbo

src/art/local/backend.py

@@ -162,6 +162,9 @@ def _allocated_gpu_count(self, model: Model) -> int:
     def __enter__(self) -> Self:
         return self
 
+    async def __aenter__(self) -> Self:
+        return self
+
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
@@ -170,11 +173,19 @@ def __exit__(
     ) -> None:
         self._close()
 
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        await self.close()
+
     async def close(self) -> None:
         """
         If running vLLM in a separate process, this will kill that process and close the communication threads.
         """
-        self._close()
+        await self._aclose()
 
     def _close(self) -> None:
         for _, service in self._services.items():
@@ -183,6 +194,17 @@ def _close(self) -> None:
                 close()
             close_proxy(service)
 
+    async def _aclose(self) -> None:
+        for _, service in self._services.items():
+            aclose = getattr(service, "aclose", None)
+            if aclose is not None:
+                await aclose()
+            else:
+                close = getattr(service, "close", None)
+                if close is not None:
+                    close()
+            close_proxy(service)
+
     async def register(
         self,
         model: Model,
@@ -505,6 +527,7 @@ async def train(  # type: ignore[override]
         *,
         # Core training parameters
         learning_rate: float = 5e-6,
+        loss_fn: Literal["cispo", "ppo"] | None = None,
         # KL-penalized advantage adjustment
         kl_penalty_coef: float = 0.0,
         kl_penalty_reference_step: int | None = None,
@@ -600,6 +623,8 @@ async def train(  # type: ignore[override]
             # await model.log(metrics=result.metrics, step=result.step)
         """
         groups_list = list(trajectory_groups)
+        if loss_fn is not None:
+            ppo = loss_fn == "ppo"
 
         resolved_kl_ref_adapter_path = kl_ref_adapter_path
         if (

src/art/utils/convert_moe_lora.py

@@ -12,13 +12,15 @@
   ...
 """
 
+import importlib
 import json
 import os
 import re
 
-import safetensors.torch
 import torch
 
+safetensors_torch = importlib.import_module("safetensors.torch")
+
 
 def _has_fused_moe_lora(tensors: dict[str, torch.Tensor]) -> bool:
     """Check if the adapter contains fused MoE LoRA tensors."""
@@ -152,7 +154,7 @@ def convert_checkpoint_if_needed(checkpoint_dir: str) -> None:
     if not os.path.exists(adapter_path) or not os.path.exists(config_path):
         return
 
-    tensors = safetensors.torch.load_file(adapter_path)
+    tensors = safetensors_torch.load_file(adapter_path)
     if not _has_fused_moe_lora(tensors):
         return
 
@@ -168,7 +170,7 @@ def convert_checkpoint_if_needed(checkpoint_dir: str) -> None:
     )
 
     # Overwrite the adapter with the converted tensors
-    safetensors.torch.save_file(new_tensors, adapter_path)
+    safetensors_torch.save_file(new_tensors, adapter_path)
 
     # Update adapter_config.json target_modules
     adapter_config["target_modules"] = [