From 079daae952aeae8c1e6dd1a009ec3d950d0f4767 Mon Sep 17 00:00:00 2001
From: faresobeid <111092724+faresobeid@users.noreply.github.com>
Date: Sat, 3 Oct 2026 00:28:52 +0000
Subject: [PATCH] chore: remove legacy dead code (OAI patches, pathing
 re-export, no-op stubs)

---
 .../src/prime_rl/configs/orchestrator.py      |   4 +-
 src/prime_rl/eval/runner.py                   |   7 --
 .../inference/vllm/worker/filesystem.py       |   4 -
 src/prime_rl/orchestrator/concurrency.py      |   2 +-
 src/prime_rl/orchestrator/dispatcher.py       |   8 +-
 src/prime_rl/orchestrator/orchestrator.py     |  12 +-
 src/prime_rl/orchestrator/patches.py          | 115 ------------------
 src/prime_rl/trainer/ckpt.py                  |   2 +-
 .../models/glm4_moe/modeling_glm4_moe.py      |   2 -
 .../glm_moe_dsa/modeling_glm_moe_dsa.py       |   2 -
 .../trainer/models/llama/modeling_llama.py    |   2 -
 src/prime_rl/trainer/rl/train.py              |   3 +-
 src/prime_rl/utils/utils.py                   |  19 ---
 .../test_serving_chat_with_tokens.py          |   1 -
 14 files changed, 11 insertions(+), 172 deletions(-)
 delete mode 100644 src/prime_rl/orchestrator/patches.py
 delete mode 100644 tests/unit/inference/test_serving_chat_with_tokens.py

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
index c8ce83c5857..75c8e32f041 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -527,8 +527,8 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 
 class ConcurrencyConfig(BaseConfig):
     """Adaptive in-flight concurrency control. The orchestrator sizes the
-    in-flight episode cap from engine KV capacity and learned per-env episode
-    costs; these fields only bound and seed it."""
+    in-flight episode cap from engine KV pressure; these fields only bound and
+    seed it."""
 
     initial_inflight: int | None = Field(None, ge=1)
     """Optional initial in-flight episodes to start from. Set it when a good value is known to skip the initial ramp; otherwise auto-derive a pessimistic bound at runtime."""
diff --git a/src/prime_rl/eval/runner.py b/src/prime_rl/eval/runner.py
index 17752ffa798..6a51e04f31a 100644
--- a/src/prime_rl/eval/runner.py
+++ b/src/prime_rl/eval/runner.py
@@ -34,10 +34,6 @@
 from prime_rl.orchestrator.eval_source import EvalSource
 from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector
 from prime_rl.orchestrator.metrics import dispatch_failure_metrics
-from prime_rl.orchestrator.patches import (
-    monkey_patch_chat_completion_logprobs,
-    monkey_patch_oai_iterable_types,
-)
 from prime_rl.orchestrator.periodic_logger import PeriodicLogger
 from prime_rl.orchestrator.types import DispatchFailure, EvalBatch, GroupCancellation, Policy
 from prime_rl.orchestrator.utils import (
@@ -49,9 +45,6 @@
 from prime_rl.utils.logger import format_time, get_logger
 from prime_rl.utils.pathing import get_config_dir
 
-monkey_patch_oai_iterable_types()
-monkey_patch_chat_completion_logprobs()
-
 # How often ``run_epoch`` re-checks for a superseding checkpoint while it waits for episodes.
 POLL_INTERVAL_S = 2.0
 
diff --git a/src/prime_rl/inference/vllm/worker/filesystem.py b/src/prime_rl/inference/vllm/worker/filesystem.py
index 1049d17b0c0..53471969df4 100644
--- a/src/prime_rl/inference/vllm/worker/filesystem.py
+++ b/src/prime_rl/inference/vllm/worker/filesystem.py
@@ -17,10 +17,6 @@
 class FileSystemWeightUpdateWorker(Worker):
     """vLLM worker extension for updating weights in-place using shared filesystem."""
 
-    def init_broadcaster(self) -> None:
-        """Initialize the broadcaster."""
-        ...
-
     def liveness_probe(self) -> None:
         """No-op RPC used by the API server liveness endpoint."""
         return None
diff --git a/src/prime_rl/orchestrator/concurrency.py b/src/prime_rl/orchestrator/concurrency.py
index 8874064279a..f6056ee6739 100644
--- a/src/prime_rl/orchestrator/concurrency.py
+++ b/src/prime_rl/orchestrator/concurrency.py
@@ -163,7 +163,7 @@ def bind(
 
     # ── inbound hooks ────────────────────────────────────────────────────────
 
-    def record_episode(self, env_name: str, kind: str, tokens: int, duration: float) -> None:
+    def record_episode(self, tokens: int) -> None:
         """One completed unit (from the dispatcher). Advances the turnover
         clock and, while the last poll was green, grows the cap — so growth
         is paced by the pipeline itself at any concurrency scale."""
diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py
index e3cfcede58f..b9410915cea 100644
--- a/src/prime_rl/orchestrator/dispatcher.py
+++ b/src/prime_rl/orchestrator/dispatcher.py
@@ -143,7 +143,7 @@ def __init__(
         max_off_policy_steps: int,
         run_id: str,
         run_name: str | None,
-        on_episode_complete: Callable[[str, str, int, float], None] | None = None,
+        on_episode_complete: Callable[[int], None] | None = None,
     ) -> None:
         self.policy = policy
         self.progress = progress
@@ -157,7 +157,7 @@ def __init__(
         self.max_off_policy_steps = max_off_policy_steps
         self.run_id = run_id
         self.run_name = run_name
-        # ``(env_name, kind, total_tokens, duration_s)`` per completed episode
+        # Called with ``total_tokens`` per completed episode
         self.on_episode_complete = on_episode_complete
 
         # Starting value of the dynamic cap (the concurrency controller moves
@@ -694,9 +694,7 @@ async def handle_completed_request(self, task: asyncio.Task) -> None:
         if not episode.ok and not episode.traces:
             self.metrics.record_error(kind=meta.kind, env_name=meta.env_name)
         if self.on_episode_complete is not None and meta.started_at > 0:
-            self.on_episode_complete(
-                meta.env_name, meta.kind, episode.num_total_tokens, time.monotonic() - meta.started_at
-            )
+            self.on_episode_complete(episode.num_total_tokens)
         await self.emit_episode(meta, group, episode)
 
     def complete_group_member(self, meta: InflightEpisode, group: GroupState | None) -> int:
diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py
index b7c32bfb243..ecf3846adc1 100644
--- a/src/prime_rl/orchestrator/orchestrator.py
+++ b/src/prime_rl/orchestrator/orchestrator.py
@@ -48,10 +48,6 @@
 from prime_rl.orchestrator.inference_metrics import InferenceMetricsCollector
 from prime_rl.orchestrator.metrics import TrainEpisodes, dispatch_failure_metrics
 from prime_rl.orchestrator.packing import BatchPacker
-from prime_rl.orchestrator.patches import (
-    monkey_patch_chat_completion_logprobs,
-    monkey_patch_oai_iterable_types,
-)
 from prime_rl.orchestrator.periodic_logger import PeriodicLogger
 from prime_rl.orchestrator.train_sink import TrainSink
 from prime_rl.orchestrator.train_source import TrainSource
@@ -78,12 +74,8 @@
 from prime_rl.utils.async_utils import EventLoopLagMonitor, EventLoopLagStats, safe_cancel
 from prime_rl.utils.heartbeat import Heartbeat
 from prime_rl.utils.logger import format_time, get_logger, setup_logger
-from prime_rl.utils.pathing import get_broadcast_dir, get_config_dir
-from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step
-
-monkey_patch_oai_iterable_types()
-monkey_patch_chat_completion_logprobs()
-
+from prime_rl.utils.pathing import get_broadcast_dir, get_config_dir, resolve_latest_ckpt_step
+from prime_rl.utils.utils import clean_exit
 
 # Wall-clock budget for post-training cleanup; force-exit if graceful
 # shutdown wedges (env-server ZMQ recv, vLLM admin aclose, etc)
diff --git a/src/prime_rl/orchestrator/patches.py b/src/prime_rl/orchestrator/patches.py
deleted file mode 100644
index 9f0d0f7690a..00000000000
--- a/src/prime_rl/orchestrator/patches.py
+++ /dev/null
@@ -1,115 +0,0 @@
-from typing import Any, Optional, TypedDict, Union
-
-import openai.types.chat
-from openai.types.chat.chat_completion import ChatCompletion, Choice
-from openai.types.chat.chat_completion_assistant_message_param import (
-    Audio,
-    ContentArrayOfContentPart,
-)
-from openai.types.chat.chat_completion_content_part_param import ChatCompletionContentPartParam
-from openai.types.chat.chat_completion_content_part_text_param import ChatCompletionContentPartTextParam
-from openai.types.chat.chat_completion_developer_message_param import ChatCompletionDeveloperMessageParam
-from openai.types.chat.chat_completion_function_message_param import ChatCompletionFunctionMessageParam
-from openai.types.chat.chat_completion_message import FunctionCall
-from openai.types.chat.chat_completion_message_tool_call_union_param import ChatCompletionMessageToolCallUnionParam
-from openai.types.chat.chat_completion_system_message_param import ChatCompletionSystemMessageParam
-from openai.types.chat.chat_completion_user_message_param import ChatCompletionUserMessageParam
-from typing_extensions import Literal, Required
-
-
-def monkey_patch_oai_iterable_types():
-    """
-    This monkey patch is necessary to avoid Pydantic validating fields using
-    typing.Iterable (e.g. in multimodal or tool call messages) lazily which
-    leads to tokenization errors, for more info see
-    https://github.com/PrimeIntellect-ai/prime-rl/pull/1249
-    """
-
-    class ModdedChatCompletionDeveloperMessageParam(TypedDict, total=False):
-        """Same as openai.types.chat.chat_completion_developer_message_param.ChatCompletionDeveloperMessageParam, but replacing typing.Iterable with list to not mess up Pydantic."""
-
-        content: Required[Union[str, list[ChatCompletionContentPartTextParam]]]
-        role: Required[Literal["developer"]]
-        name: str
-
-    class ModdedChatCompletionSystemMessageParam(TypedDict, total=False):
-        """Same as openai.types.chat.chat_completion_system_message_param.ChatCompletionSystemMessageParam, but replacing typing.Iterable with list to not mess up Pydantic."""
-
-        content: Required[Union[str, list[ChatCompletionContentPartTextParam]]]
-        role: Required[Literal["system"]]
-        name: str
-
-    class ModdedChatCompletionUserMessageParam(TypedDict, total=False):
-        """Same as openai.types.chat.chat_completion_user_message_param.ChatCompletionUserMessageParam, but replacing typing.Iterable with list to not mess up Pydantic."""
-
-        content: Required[Union[str, list[ChatCompletionContentPartParam]]]
-        role: Required[Literal["user"]]
-        name: str
-
-    class ModdedChatCompletionAssistantMessageParam(TypedDict, total=False):
-        """Same as openai.types.chat.chat_completion_assistant_message_param.ChatCompletionAssistantMessageParam, but replacing typing.Iterable with list to not mess up Pydantic."""
-
-        role: Required[Literal["assistant"]]
-        audio: Optional[Audio]
-        content: Union[str, list[ContentArrayOfContentPart], None]
-        function_call: Optional[FunctionCall]
-        name: str
-        refusal: Optional[str]
-        tool_calls: list[ChatCompletionMessageToolCallUnionParam]
-
-    class ModdedChatCompletionToolMessageParam(TypedDict, total=False):
-        """Same as openai.types.chat.chat_completion_tool_message_param.ChatCompletionToolMessageParam, but replacing typing.Iterable with list to not mess up Pydantic."""
-
-        content: Required[Union[str, list[ChatCompletionContentPartTextParam]]]
-        role: Required[Literal["tool"]]
-        tool_call_id: Required[str]
-
-    # Patch OAI types
-    openai.types.chat.chat_completion_developer_message_param.ChatCompletionDeveloperMessageParam = (
-        ModdedChatCompletionDeveloperMessageParam
-    )
-    openai.types.chat.chat_completion_system_message_param.ChatCompletionSystemMessageParam = (
-        ModdedChatCompletionSystemMessageParam
-    )
-    openai.types.chat.chat_completion_user_message_param.ChatCompletionUserMessageParam = (
-        ModdedChatCompletionUserMessageParam
-    )
-    openai.types.chat.chat_completion_assistant_message_param.ChatCompletionAssistantMessageParam = (
-        ModdedChatCompletionAssistantMessageParam
-    )
-    openai.types.chat.chat_completion_tool_message_param.ChatCompletionToolMessageParam = (
-        ModdedChatCompletionToolMessageParam
-    )
-
-    openai.types.chat.chat_completion_message_param.ChatCompletionMessageParam = Union[
-        ChatCompletionDeveloperMessageParam,
-        ChatCompletionSystemMessageParam,
-        ChatCompletionUserMessageParam,
-        ModdedChatCompletionAssistantMessageParam,
-        ModdedChatCompletionToolMessageParam,
-        ChatCompletionFunctionMessageParam,
-    ]
-
-
-def monkey_patch_chat_completion_logprobs():
-    """
-    At large batch sizes and context, constructing OAI's Pydantic model
-    ChatCompletion with logprobs is causing heavy CPU overhead (~200ms per
-    object at 32K context, which translates to >10min overhead at 4K batch
-    size). This function monkey-patches the OAI type and verifiers'
-    post-processing utils to avoid validating the complex logprobs field.
-    """
-
-    class ChoiceAny(Choice):
-        """Same as openai.types.chat.chat_completion.Choice, but without type validation for logprobs field."""
-
-        logprobs: Optional[Any] = None
-
-    class ModdedChatCompletion(ChatCompletion):
-        """Same as openai.types.chat.chat_completion.ChatCompletion, but using ChoiceAny instead of Choice."""
-
-        choices: list[ChoiceAny]  # type: ignore
-
-    # Patch OAI types
-    openai.types.chat.chat_completion.Choice = ChoiceAny
-    openai.types.chat.chat_completion.ChatCompletion = ModdedChatCompletion
diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py
index 356a0bfae3b..ddda992a78f 100644
--- a/src/prime_rl/trainer/ckpt.py
+++ b/src/prime_rl/trainer/ckpt.py
@@ -27,7 +27,7 @@
 from prime_rl.trainer.optim import OffloadOptimizer, OptimizerLike
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.logger import format_time, get_logger
-from prime_rl.utils.utils import get_all_ckpt_steps, get_ckpt_dir, get_step_path
+from prime_rl.utils.pathing import get_all_ckpt_steps, get_ckpt_dir, get_step_path
 
 
 @dataclass
diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
index 3b2dfa20601..19cda1c6259 100644
--- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
+++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py
@@ -23,7 +23,6 @@
 from transformers.modeling_outputs import BaseModelOutputWithPast
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring
-from transformers.utils.deprecation import deprecate_kwarg
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.glm4_moe.configuration_glm4_moe import Glm4MoeConfig
@@ -89,7 +88,6 @@ def __init__(self, config: Glm4MoeConfig, layer_idx: int):
         self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
         self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
 
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
index 08717e20654..1278f73b369 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
@@ -10,7 +10,6 @@
 from transformers.modeling_outputs import BaseModelOutputWithPast
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring
-from transformers.utils.deprecation import deprecate_kwarg
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig, _index_cache_skip_topk
@@ -87,7 +86,6 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
         self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
         self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
 
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/models/llama/modeling_llama.py b/src/prime_rl/trainer/models/llama/modeling_llama.py
index e17d7464d25..e7e01b45cdd 100644
--- a/src/prime_rl/trainer/models/llama/modeling_llama.py
+++ b/src/prime_rl/trainer/models/llama/modeling_llama.py
@@ -32,7 +32,6 @@
 from transformers.models.llama.configuration_llama import LlamaConfig
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
-from transformers.utils.deprecation import deprecate_kwarg
 
 from prime_rl.trainer.models.base import PreTrainedModelPrimeRL
 from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig
@@ -71,7 +70,6 @@ def __init__(self, config: LlamaConfig, layer_idx: int):
         self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
         self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps))
 
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
     def forward(
         self,
         hidden_states: torch.Tensor,
diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py
index 740a70504a3..a7dbf046d1f 100644
--- a/src/prime_rl/trainer/rl/train.py
+++ b/src/prime_rl/trainer/rl/train.py
@@ -71,7 +71,8 @@
 from prime_rl import monitors
 from prime_rl.utils.config import cli
 from prime_rl.utils.process import set_proc_title
-from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step
+from prime_rl.utils.pathing import resolve_latest_ckpt_step
+from prime_rl.utils.utils import clean_exit
 
 
 @clean_exit
diff --git a/src/prime_rl/utils/utils.py b/src/prime_rl/utils/utils.py
index 4365dc327d7..5a70c78934b 100644
--- a/src/prime_rl/utils/utils.py
+++ b/src/prime_rl/utils/utils.py
@@ -2,33 +2,14 @@
 import functools
 import importlib
 import math
-import os
 import sys
-from collections import defaultdict
 from contextlib import contextmanager
-from pathlib import Path
 from typing import Any, Callable
 
 import wandb
 
 from prime_rl.utils.logger import get_logger
 
-# TODO: Change all imports to use utils.pathing
-# ruff: noqa: F401
-from prime_rl.utils.pathing import (
-    get_all_ckpt_steps,
-    get_batch_dir,
-    get_broadcast_dir,
-    get_ckpt_dir,
-    get_eval_dir,
-    get_file_monitor_dir,
-    get_log_dir,
-    get_step_path,
-    resolve_latest_ckpt_step,
-    sync_wait_for_path,
-    wait_for_path,
-)
-
 
 def import_object(dotted_path: str) -> Any:
     """Import an object from a dotted path like 'my_module.submodule.MyClass'."""
diff --git a/tests/unit/inference/test_serving_chat_with_tokens.py b/tests/unit/inference/test_serving_chat_with_tokens.py
deleted file mode 100644
index 8b137891791..00000000000
--- a/tests/unit/inference/test_serving_chat_with_tokens.py
+++ /dev/null
@@ -1 +0,0 @@
-