Skip to content

Commit deac822

Browse files
alexzhang13alexzhang13
committed
fix format
1 parent 4689c90 commit deac822

7 files changed

Lines changed: 154 additions & 57 deletions

File tree

‎rlm/utils/prompts.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@
153153
"the concrete sequence of turns — what each turn computes and which sub-LLM call "
154154
"(if any) it issues — like a condensed trajectory, before you execute them. "
155155
"Then execute one turn at a time: after each step `print` a small sample of the "
156-
"result, verify it looks right, and only flip `answer[\"ready\"] = True` once you "
156+
'result, verify it looks right, and only flip `answer["ready"] = True` once you '
157157
"have actually printed the candidate answer. If you are running out of turns "
158158
"without a confirmed answer, submit your best inference rather than letting the "
159159
"rollout terminate unsubmitted."

‎training/environments/oolong/oolong/env.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -135,12 +135,14 @@ def _keep(ex):
135135
"context": context,
136136
"root_prompt": f"{_QUESTION_INSTRUCTION}\n\nQuestion: {question}",
137137
}
138-
rows.append({
139-
"example_id": i,
140-
"prompt": [{"role": "user", "content": question}],
141-
"answer": str(s.get("answer", "")),
142-
"info": json.dumps(meta),
143-
})
138+
rows.append(
139+
{
140+
"example_id": i,
141+
"prompt": [{"role": "user", "content": question}],
142+
"answer": str(s.get("answer", "")),
143+
"info": json.dumps(meta),
144+
}
145+
)
144146
return Dataset.from_list(rows)
145147

146148

‎training/src/rlm_train/env.py‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,15 @@ async def get_prompt_messages(self, state: State) -> Messages:
180180
try:
181181
result = await backend.execute(code)
182182
except Exception as e: # noqa: BLE001
183-
outputs.append({
184-
"code": code, "stdout": "", "stderr": f"Worker error: {e}",
185-
"final_answer": None, "locals_keys": [],
186-
})
183+
outputs.append(
184+
{
185+
"code": code,
186+
"stdout": "",
187+
"stderr": f"Worker error: {e}",
188+
"final_answer": None,
189+
"locals_keys": [],
190+
}
191+
)
187192
continue
188193
outputs.append(_pack_exec(code, result))
189194
state["rlm_repl_calls"] = int(state.get("rlm_repl_calls") or 0) + 1
@@ -250,7 +255,8 @@ def _normalize_for_api(msgs: list) -> list:
250255
out.append(m)
251256
continue
252257
reasoning = (
253-
m.get("reasoning_content") if isinstance(m, dict)
258+
m.get("reasoning_content")
259+
if isinstance(m, dict)
254260
else getattr(m, "reasoning_content", None)
255261
)
256262
if isinstance(m, dict):
@@ -323,5 +329,7 @@ def _format_one(o: dict[str, Any]) -> str:
323329
parts.append(f"REPL variables: {list(o['locals_keys'])}\n")
324330
body = "\n\n".join(parts) if parts else "No output"
325331
if len(body) > _MAX_REPL_OUTPUT_CHARS:
326-
body = body[:_MAX_REPL_OUTPUT_CHARS] + f"... + [{len(body) - _MAX_REPL_OUTPUT_CHARS} chars...]"
332+
body = (
333+
body[:_MAX_REPL_OUTPUT_CHARS] + f"... + [{len(body) - _MAX_REPL_OUTPUT_CHARS} chars...]"
334+
)
327335
return body

‎training/src/rlm_train/proxy.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,9 @@ async def run_one(p: str) -> str:
196196
text, meta = await self._completion(handle, p, model)
197197
if handle.record_call is not None:
198198
try:
199-
handle.record_call({"model": model, "prompt": p, "response": text, **meta})
199+
handle.record_call(
200+
{"model": model, "prompt": p, "response": text, **meta}
201+
)
200202
except Exception:
201203
logger.exception("record_call failed")
202204
return text
@@ -208,7 +210,10 @@ async def run_one(p: str) -> str:
208210
return web.json_response({"responses": results})
209211

210212
async def _completion(
211-
self, handle: ClientHandle, prompt: str | list, model: str,
213+
self,
214+
handle: ClientHandle,
215+
prompt: str | list,
216+
model: str,
212217
) -> tuple[str, dict[str, Any]]:
213218
if handle.fake_query is not None:
214219
prompt_text = _flatten_prompt(prompt)

‎training/src/rlm_train/repl/subprocess.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,10 @@ async def start(self, proxy_url: str, rollout_id: str, depth: int = 1) -> None:
4747
env.setdefault("PYTHONUNBUFFERED", "1")
4848

4949
self._proc = await asyncio.create_subprocess_exec(
50-
self._python, "-u", "-m", self._worker_module,
50+
self._python,
51+
"-u",
52+
"-m",
53+
self._worker_module,
5154
stdin=asyncio.subprocess.PIPE,
5255
stdout=asyncio.subprocess.PIPE,
5356
stderr=asyncio.subprocess.PIPE,

‎training/src/rlm_train/rubric.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,7 @@ async def _call_correctness(
5656
result = await result
5757
return float(result)
5858

59-
def _make_main_correctness(
60-
self, correctness: Callable[..., float]
61-
) -> Callable[..., Any]:
59+
def _make_main_correctness(self, correctness: Callable[..., float]) -> Callable[..., Any]:
6260
gated_mode = self._gate_reward
6361
min_rew = self._min_reward
6462

@@ -76,9 +74,7 @@ async def main(**kwargs: Any) -> float:
7674
main.__name__ = getattr(correctness, "__name__", "correctness")
7775
return main
7876

79-
def _make_gated_metric(
80-
self, correctness: Callable[..., float]
81-
) -> Callable[..., Any]:
77+
def _make_gated_metric(self, correctness: Callable[..., float]) -> Callable[..., Any]:
8278
min_rew = self._min_reward
8379

8480
async def gated_reward(**kwargs: Any) -> float:

‎training/src/rlm_train/worker.py‎

Lines changed: 118 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -17,38 +17,105 @@
1717
from typing import Any
1818

1919
_SAFE_BUILTINS = {
20-
"print": print, "len": len, "str": str, "int": int, "float": float,
21-
"list": list, "dict": dict, "set": set, "tuple": tuple, "bool": bool,
22-
"type": type, "isinstance": isinstance, "issubclass": issubclass,
23-
"enumerate": enumerate, "zip": zip, "map": map, "filter": filter,
24-
"sorted": sorted, "reversed": reversed, "range": range,
25-
"min": min, "max": max, "sum": sum, "abs": abs, "round": round,
26-
"any": any, "all": all, "pow": pow, "divmod": divmod,
27-
"chr": chr, "ord": ord, "hex": hex, "bin": bin, "oct": oct,
28-
"repr": repr, "ascii": ascii, "format": format, "hash": hash, "id": id,
29-
"iter": iter, "next": next, "slice": slice, "callable": callable,
30-
"hasattr": hasattr, "getattr": getattr, "setattr": setattr, "delattr": delattr,
31-
"dir": dir, "vars": vars,
32-
"bytes": bytes, "bytearray": bytearray, "memoryview": memoryview,
33-
"complex": complex, "object": object, "super": super, "property": property,
34-
"staticmethod": staticmethod, "classmethod": classmethod,
35-
"__import__": __import__, "open": open,
36-
"Exception": Exception, "BaseException": BaseException,
37-
"ValueError": ValueError, "TypeError": TypeError, "KeyError": KeyError,
38-
"IndexError": IndexError, "AttributeError": AttributeError,
39-
"FileNotFoundError": FileNotFoundError, "OSError": OSError, "IOError": IOError,
40-
"RuntimeError": RuntimeError, "NameError": NameError, "ImportError": ImportError,
41-
"StopIteration": StopIteration, "AssertionError": AssertionError,
42-
"NotImplementedError": NotImplementedError, "ArithmeticError": ArithmeticError,
43-
"LookupError": LookupError, "Warning": Warning,
44-
"input": None, "eval": None, "exec": None, "compile": None,
45-
"globals": None, "locals": None,
20+
"print": print,
21+
"len": len,
22+
"str": str,
23+
"int": int,
24+
"float": float,
25+
"list": list,
26+
"dict": dict,
27+
"set": set,
28+
"tuple": tuple,
29+
"bool": bool,
30+
"type": type,
31+
"isinstance": isinstance,
32+
"issubclass": issubclass,
33+
"enumerate": enumerate,
34+
"zip": zip,
35+
"map": map,
36+
"filter": filter,
37+
"sorted": sorted,
38+
"reversed": reversed,
39+
"range": range,
40+
"min": min,
41+
"max": max,
42+
"sum": sum,
43+
"abs": abs,
44+
"round": round,
45+
"any": any,
46+
"all": all,
47+
"pow": pow,
48+
"divmod": divmod,
49+
"chr": chr,
50+
"ord": ord,
51+
"hex": hex,
52+
"bin": bin,
53+
"oct": oct,
54+
"repr": repr,
55+
"ascii": ascii,
56+
"format": format,
57+
"hash": hash,
58+
"id": id,
59+
"iter": iter,
60+
"next": next,
61+
"slice": slice,
62+
"callable": callable,
63+
"hasattr": hasattr,
64+
"getattr": getattr,
65+
"setattr": setattr,
66+
"delattr": delattr,
67+
"dir": dir,
68+
"vars": vars,
69+
"bytes": bytes,
70+
"bytearray": bytearray,
71+
"memoryview": memoryview,
72+
"complex": complex,
73+
"object": object,
74+
"super": super,
75+
"property": property,
76+
"staticmethod": staticmethod,
77+
"classmethod": classmethod,
78+
"__import__": __import__,
79+
"open": open,
80+
"Exception": Exception,
81+
"BaseException": BaseException,
82+
"ValueError": ValueError,
83+
"TypeError": TypeError,
84+
"KeyError": KeyError,
85+
"IndexError": IndexError,
86+
"AttributeError": AttributeError,
87+
"FileNotFoundError": FileNotFoundError,
88+
"OSError": OSError,
89+
"IOError": IOError,
90+
"RuntimeError": RuntimeError,
91+
"NameError": NameError,
92+
"ImportError": ImportError,
93+
"StopIteration": StopIteration,
94+
"AssertionError": AssertionError,
95+
"NotImplementedError": NotImplementedError,
96+
"ArithmeticError": ArithmeticError,
97+
"LookupError": LookupError,
98+
"Warning": Warning,
99+
"input": None,
100+
"eval": None,
101+
"exec": None,
102+
"compile": None,
103+
"globals": None,
104+
"locals": None,
46105
}
47106

48-
RESERVED_TOOL_NAMES = frozenset({
49-
"llm_query", "llm_query_batched", "rlm_query", "rlm_query_batched",
50-
"SHOW_VARS", "answer", "context", "history",
51-
})
107+
RESERVED_TOOL_NAMES = frozenset(
108+
{
109+
"llm_query",
110+
"llm_query_batched",
111+
"rlm_query",
112+
"rlm_query_batched",
113+
"SHOW_VARS",
114+
"answer",
115+
"context",
116+
"history",
117+
}
118+
)
52119

53120

54121
class _AnswerDict(dict):
@@ -143,7 +210,10 @@ def _proxy_post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
143210
url = f"{self.proxy_url}/rollout/{self.rollout_id}/{path.lstrip('/')}"
144211
body = json.dumps(payload).encode("utf-8")
145212
req = urllib.request.Request(
146-
url, data=body, headers={"Content-Type": "application/json"}, method="POST",
213+
url,
214+
data=body,
215+
headers={"Content-Type": "application/json"},
216+
method="POST",
147217
)
148218
try:
149219
with urllib.request.urlopen(req, timeout=600) as resp:
@@ -153,7 +223,11 @@ def _proxy_post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
153223
detail = e.read().decode("utf-8")
154224
except Exception:
155225
detail = ""
156-
too_large = e.code == 413 or "Entity Too Large" in (detail or "") or "Entity Too Large" in (e.reason or "")
226+
too_large = (
227+
e.code == 413
228+
or "Entity Too Large" in (detail or "")
229+
or "Entity Too Large" in (e.reason or "")
230+
)
157231
return {"error": f"HTTP {e.code}: {detail or e.reason}", "too_large": too_large}
158232
except Exception as e:
159233
return {"error": f"Proxy request failed: {e}"}
@@ -253,7 +327,8 @@ def execute(self, code: str) -> dict[str, Any]:
253327
final_answer = self._last_final_answer
254328
self._last_final_answer = None
255329
simple_keys = [
256-
k for k, v in self.locals.items()
330+
k
331+
for k, v in self.locals.items()
257332
if not k.startswith("_")
258333
and k not in ("__builtins__", "__name__", "__doc__")
259334
and isinstance(v, (str, int, float, bool, list, dict, tuple))
@@ -305,7 +380,9 @@ def main() -> None:
305380
result = worker.execute(req.get("code", ""))
306381
_send({"id": rid, "ok": True, **result})
307382
except BaseException as e: # noqa: BLE001
308-
_send({"id": rid, "ok": False, "error": f"exec failed: {e}\n{traceback.format_exc()}"})
383+
_send(
384+
{"id": rid, "ok": False, "error": f"exec failed: {e}\n{traceback.format_exc()}"}
385+
)
309386
elif kind == "load_context":
310387
try:
311388
idx = worker.load_context(req.get("payload"), req.get("index"))
@@ -319,7 +396,13 @@ def main() -> None:
319396
exec(compile(code, "<bootstrap>", "exec"), worker.globals)
320397
_send({"id": rid, "ok": True})
321398
except BaseException as e: # noqa: BLE001
322-
_send({"id": rid, "ok": False, "error": f"bootstrap failed: {e}\n{traceback.format_exc()}"})
399+
_send(
400+
{
401+
"id": rid,
402+
"ok": False,
403+
"error": f"bootstrap failed: {e}\n{traceback.format_exc()}",
404+
}
405+
)
323406
elif kind == "shutdown":
324407
_send({"id": rid, "ok": True})
325408
return

0 commit comments

Comments
 (0)