linting fixes

Author: alexzhang13

training/environments/oolong/oolong/env.py

@@ -9,12 +9,10 @@
 from typing import Any
 
 import dateutil.parser
+import rlm_train
 import verifiers as vf
 from datasets import Dataset, load_dataset
 
-import rlm_train
-
-
 COMPARISON_PHRASES = ("more common than", "less common than", "same frequency as")
 
 

training/src/rlm_train/env.py

@@ -5,7 +5,8 @@
 import asyncio
 import logging
 import uuid
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import verifiers as vf
 from verifiers.types import Messages, State
@@ -17,7 +18,6 @@
     build_rlm_system_prompt,
     build_user_prompt,
 )
-
 from rlm_train.proxy import ClientHandle, SubLLMProxy
 from rlm_train.repl.base import ExecResult, ReplBackend
 from rlm_train.repl.subprocess import SubprocessReplBackend

training/src/rlm_train/proxy.py

@@ -5,16 +5,17 @@
 import asyncio
 import inspect
 import logging
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
-from typing import Any, Awaitable, Callable, Union
+from typing import Any
 
 from aiohttp import web
 
 logger = logging.getLogger(__name__)
 
 
-FakeQuery = Callable[[str, "Any"], Union[str, Awaitable[str]]]
-FakeQueryBatched = Callable[[list[str], "Any"], Union[list[str], Awaitable[list[str]]]]
+FakeQuery = Callable[[str, "Any"], str | Awaitable[str]]
+FakeQueryBatched = Callable[[list[str], "Any"], list[str] | Awaitable[list[str]]]
 
 
 @dataclass
@@ -178,7 +179,7 @@ async def _handle_batched(self, request: web.Request) -> web.Response:
                 if not isinstance(responses, list) or len(responses) != len(prompts):
                     return web.json_response({"error": "fake_query_batched returned wrong shape"})
                 if handle.record_call is not None:
-                    for p, r in zip(prompts, responses):
+                    for p, r in zip(prompts, responses, strict=True):
                         try:
                             handle.record_call({"model": model, "prompt": p, "response": r})
                         except Exception:

training/src/rlm_train/repl/subprocess.py

@@ -57,7 +57,7 @@ async def start(self, proxy_url: str, rollout_id: str, depth: int = 1) -> None:
         self._stderr_task = asyncio.create_task(self._drain_stderr())
         try:
             init = await asyncio.wait_for(self._read_line(), timeout=self._startup_timeout)
-        except asyncio.TimeoutError as e:
+        except TimeoutError as e:
             await self.stop()
             raise WorkerStartupError(
                 f"Worker did not produce init line in {self._startup_timeout}s; "
@@ -75,7 +75,7 @@ async def stop(self) -> None:
                 try:
                     await self._send({"id": "_shutdown", "type": "shutdown"})
                     await asyncio.wait_for(self._read_line(), timeout=2.0)
-                except (BrokenPipeError, ConnectionResetError, asyncio.TimeoutError, WorkerProtocolError):
+                except (TimeoutError, BrokenPipeError, ConnectionResetError, WorkerProtocolError):
                     pass
         finally:
             if self._proc.returncode is None:
@@ -85,7 +85,7 @@ async def stop(self) -> None:
                     pass
                 try:
                     await asyncio.wait_for(self._proc.wait(), timeout=2.0)
-                except asyncio.TimeoutError:
+                except TimeoutError:
                     self._proc.kill()
                     await self._proc.wait()
             if self._stderr_task and not self._stderr_task.done():
@@ -126,7 +126,7 @@ async def _request(self, payload: dict[str, Any]) -> dict[str, Any]:
             await self._send(payload)
             try:
                 resp = await asyncio.wait_for(self._read_line(), timeout=self._request_timeout)
-            except asyncio.TimeoutError as e:
+            except TimeoutError as e:
                 await self._kill_worker()
                 raise WorkerProtocolError(
                     f"Request {payload.get('type')} exceeded parent watchdog "
@@ -147,7 +147,7 @@ async def _kill_worker(self) -> None:
             return
         try:
             await asyncio.wait_for(self._proc.wait(), timeout=2.0)
-        except asyncio.TimeoutError:
+        except TimeoutError:
             pass
 
     async def _send(self, payload: dict[str, Any]) -> None:

training/src/rlm_train/rubric.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import verifiers as vf
 from verifiers.types import State