Graphrl - #107
Open
JamesKrW wants to merge 5 commits into
Open
Conversation
Single-process asyncio rollout in run_eval_parallel becomes the bottleneck
when per-job CPU work is heavy (PIL image load + resize, base64 encoding,
JSON serialization, regex parsing in vision-language reasoning loops): the
GPU-side scheduler sits idle waiting for the orchestrator to dispatch
requests because Python is GIL-bound to one core.
Add an opt-in multi-process mode: when run.num_workers > 1, run_eval.main
partitions jobs round-robin across N subprocess workers. Each worker runs
its own asyncio event loop calling the unmodified run_eval_parallel against
the same shared backend. The dump_dir layout is bit-identical to the
single-process path because every rollout writes to its own
tag_<id>/<seed>/ subdir — concurrent writers never collide.
Default num_workers=1 → legacy single-process path, byte-identical behaviour.
max_concurrent_jobs is now per-worker (so existing configs without
num_workers are unchanged). Total in-flight = num_workers * max_concurrent_jobs.
Also adds run.normal_finish_reasons override so callers (e.g. the reasoning-
augmentation pipeline that monkey-patches NORMAL_FINISH_REASONS to {"done"})
can propagate that override across the spawn boundary into child workers.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughThis PR introduces multi-process evaluation for VAGEN by adding job partitioning, subprocess worker infrastructure, and configurable routing in the main evaluation entry point. It also adds best-validation checkpoint tracking and HuggingFace upload integration to save and optionally upload the best actor checkpoint. ChangesMulti-process Evaluation Feature
Best-Validation Checkpointing & HF Uploads
Sequence Diagram(s)sequenceDiagram
participant ParentProcess
participant ProcessPoolExecutor
participant WorkerProcess
ParentProcess->>ProcessPoolExecutor: submit(run_eval_chunk_subprocess, payload(chunk, concurrency, normal_finish_reasons))
ProcessPoolExecutor->>WorkerProcess: spawn worker (spawn start method)
WorkerProcess->>WorkerProcess: apply normal_finish_reasons override
WorkerProcess->>WorkerProcess: re-register builtin envs
WorkerProcess->>WorkerProcess: asyncio.run(run_eval_parallel(...)) -> returns rollout records
WorkerProcess-->>ProcessPoolExecutor: return rollout records
ProcessPoolExecutor-->>ParentProcess: aggregated results or worker crash
ParentProcess->>ParentProcess: convert crashed chunk -> per-job error records
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vagen/evaluate/run_eval.py`:
- Around line 422-437: The payloads currently pass global concurrency settings
unchanged (max_concurrent_jobs and backend_cfg["max_concurrency"]) to each
worker, which multiplies backend load by num_workers; modify the construction of
worker_payloads so these limits are converted to per-worker limits (e.g.,
per_worker_max_concurrency = max(1, floor_div(global_max_concurrency,
num_workers)) or use ceil as appropriate) and set "max_concurrent_jobs" and
backend_cfg["max_concurrency"] in each payload to those per-worker values before
appending, or alternatively rename and document the fields as per-worker limits
if that behavior is intended; locate this logic around the worker_payloads loop
and the variables max_concurrent_jobs and backend_cfg to implement the change.
- Around line 492-495: The module-level constant NORMAL_FINISH_REASONS (imported
by value earlier) must be updated when run_cfg provides an override; change the
override block so after setting _runner_mod.NORMAL_FINISH_REASONS you also
assign this module's NORMAL_FINISH_REASONS = set(nfr_override) so
_purge_error_rollouts() and _collect_completed_runs() (which use the local name)
see the new set; i.e., update the local binding in addition to the runner
module's binding.
- Around line 452-454: run_eval_parallel currently uses
ProcessPoolExecutor.pool.map which will abort the whole batch if any worker
raises (bootstrap/pickling/import errors); change to submit each worker_payload
with pool.submit(run_eval_chunk_subprocess, payload) and iterate futures via
concurrent.futures.as_completed, calling future.result() inside a try/except so
you can convert worker exceptions into the same structured per-job error records
returned by run_eval_parallel; include identifying context from the
corresponding worker_payload (e.g., payload id or dataset info) in the error
record so other successful chunk_results are still collected into results and a
failing worker doesn't abort the batch.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 27b2020d-f21b-4c5b-969d-b13b5eddf335
📒 Files selected for processing (2)
vagen/evaluate/run_eval.pyvagen/evaluate/runner.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vagen/evaluate/run_eval.py (1)
506-506:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winType annotation drifts from
_collect_completed_runsreturn type.
_collect_completed_runsreturnsDict[Tuple[str, int, Union[int, str]], str](tag_id may be a string), but the local annotation here pins tag_id toint. Tighten or unify with the helper's signature; otherwise type checkers will flag every assignment from_collect_completed_runsand any future str-keyed lookup againstcompleted_indexhere.🩹 Suggested fix
- completed_index: Dict[Tuple[str, int, int], str] = {} + completed_index: Dict[Tuple[str, int, Union[int, str]], str] = {}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/evaluate/run_eval.py` at line 506, The local variable completed_index is annotated too narrowly compared to the helper _collect_completed_runs (which returns Dict[Tuple[str,int,Union[int,str]], str]); update completed_index's annotation to match that return type (or reuse the helper's type alias/return annotation) so the tuple key allows tag_id to be int or str and avoids type-checker errors when assigning or looking up values; ensure you import typing.Union if needed and adjust any other references to completed_index accordingly.
🧹 Nitpick comments (1)
vagen/evaluate/run_eval.py (1)
229-243: 💤 Low valueInline
_pickto silence Ruff B023 and avoid per-iteration closure churn.
_pickcaptures the loop-scopemeta/metrics(Ruff B023). It's safe today because it's invoked immediately within the same iteration, but the pattern is brittle (any future hoist of the call breaks it) and creates a new function object every iteration.♻️ Suggested refactor
meta = _read_json(os.path.join(rollout.path, "meta.json")) or {} - # Prefer meta.json over metrics.json, but fall back only when meta's - # value is missing/None — NOT when it's a falsy-but-valid 0. - def _pick(key: str) -> Any: - v = meta.get(key) - return v if v is not None else metrics.get(key) - env_name, seed, tag_id = _pick("env_name"), _pick("seed"), _pick("tag_id") + # Prefer meta.json over metrics.json, but fall back only when meta's + # value is missing/None — NOT when it's a falsy-but-valid 0. + def _pick(src_meta: Dict[str, Any], src_metrics: Dict[str, Any], key: str) -> Any: + v = src_meta.get(key) + return v if v is not None else src_metrics.get(key) + env_name = _pick(meta, metrics, "env_name") + seed = _pick(meta, metrics, "seed") + tag_id = _pick(meta, metrics, "tag_id")Or define
_pickonce at module scope.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/evaluate/run_eval.py` around lines 229 - 243, The helper function _pick defined inside the loop captures loop-scoped variables meta/metrics and allocates a new function every iteration (Ruff B023); either inline its logic where it's used (replace calls to _pick("key") with the two-line lookup v = meta.get(key); use v if v is not None else metrics.get(key)) or hoist a single reusable helper out of the loop (define _pick once at module scope so it doesn't close over per-iteration state); update the uses that set env_name, seed, tag_id and leave completed[(str(env_name), int(seed), tag_id)] logic unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vagen/evaluate/run_eval.py`:
- Around line 377-454: The multi-process workers can concurrently write the same
tag summary file causing corruption; to fix, when _run_jobs_multiprocess sees
live_summary is True and num_workers > 1, automatically disable live_summary for
child workers and log a warning, i.e. before building payload_base set
live_summary=False (and keep the original live_summary for a single-process
final summary), update payload_base used by run_eval_chunk_subprocess so child
processes do not call write_rollouts_summary_from_dump concurrently;
alternatively (optional) implement partitioning by tag_id (e.g., replace
split_jobs_round_robin with a split_jobs_by_tag_id path when live_summary is
requested) so each worker owns disjoint tag directories.
---
Outside diff comments:
In `@vagen/evaluate/run_eval.py`:
- Line 506: The local variable completed_index is annotated too narrowly
compared to the helper _collect_completed_runs (which returns
Dict[Tuple[str,int,Union[int,str]], str]); update completed_index's annotation
to match that return type (or reuse the helper's type alias/return annotation)
so the tuple key allows tag_id to be int or str and avoids type-checker errors
when assigning or looking up values; ensure you import typing.Union if needed
and adjust any other references to completed_index accordingly.
---
Nitpick comments:
In `@vagen/evaluate/run_eval.py`:
- Around line 229-243: The helper function _pick defined inside the loop
captures loop-scoped variables meta/metrics and allocates a new function every
iteration (Ruff B023); either inline its logic where it's used (replace calls to
_pick("key") with the two-line lookup v = meta.get(key); use v if v is not None
else metrics.get(key)) or hoist a single reusable helper out of the loop (define
_pick once at module scope so it doesn't close over per-iteration state); update
the uses that set env_name, seed, tag_id and leave completed[(str(env_name),
int(seed), tag_id)] logic unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 37ec3ba0-12c6-40c7-a4a9-0ac21d9c5239
📒 Files selected for processing (2)
vagen/evaluate/run_eval.pyvagen/evaluate/runner.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
vagen/utils/upload_hugging_face.py (2)
84-86: 💤 Low valueNarrow the exception class on
upload.Static analysis flagged the bare
except Exception:(Ruff BLE001). The intent here is clearly "don't crash the training loop on HF transport errors," which is reasonable, but catchingExceptionwill also swallow programmer errors (e.g.,AttributeError,TypeError) silently. Consider narrowing to the HF-specific surface, e.g.huggingface_hub.errors.HfHubHTTPErrorplusOSError/requests.RequestException, and re-raisingKeyboardInterrupt/SystemExitimplicitly by being more specific.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/utils/upload_hugging_face.py` around lines 84 - 86, The bare except in the upload handler should be narrowed to HF/transport errors so programmer errors aren't swallowed: replace "except Exception as e:" in the upload method (the block referencing self.repo_id and label) with a specific except that catches huggingface_hub.errors.HfHubHTTPError plus network/IO exceptions (e.g., requests.RequestException and OSError), log the error the same way and return None; add the necessary imports (HfHubHTTPError and requests.RequestException) at the top of the module and do not catch KeyboardInterrupt/SystemExit so they propagate.
122-147: 💤 Low valueValidation logic only runs when
hf_save_freqis set, butupload_contentsis also used by best-val uploads.The
isinstancecheck on_upload_contentsand theactor_kwargs/repo_idsetup on lines 140–165 are guarded byif not self._hf_save_freq: return. That means when a user configuressave_best_val: Trueand provideshuggingface_hub.upload_contents/repo_idbut leaveshf_save_frequnset, the actor is never created, somaybe_upload_best_valwill silently no-op (becauseenabledisFalse).This is consistent with the docstring ("No-op if HF upload is disabled"), but it means best-val HF upload is implicitly gated on
hf_save_freqrather than its own flag. Worth either:
- documenting this coupling explicitly in
vagen_multiturn.yamlnearhf_save_freq, or- decoupling best-val upload by allowing the actor to be created whenever
repo_idis set.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/utils/upload_hugging_face.py` around lines 122 - 147, The init currently bails out early when self._hf_save_freq is falsy which skips validation of self._upload_contents and creation of the HF actor, so best-val uploads (triggered by save_best_val/repo_id) silently no-op; change the logic in __init__ so that only hf-save-frequency-specific setup is skipped when _hf_save_freq is falsy, but still validate upload_contents (the isinstance check for self._upload_contents) and perform the actor/actor_kwargs/repo_id setup whenever a repo_id is configured or save_best_val is enabled (so maybe_upload_best_val can work independently); update the guards around the blocks that reference actor_kwargs, repo_id, and actor creation to check for repo_id or save_best_val instead of requiring _hf_save_freq.vagen/utils/best_val.py (2)
99-115: ⚡ Quick winDestructive
rmtreebefore save risks losing the previous best on save failure.The order is:
flush() → rmtree(best_val_root) → save_checkpoint(...) → prune → write_metadata → upload. Ifsave_checkpointraises (OOM, FSDP rank failure, disk full), the previous best — both on disk and remotely if it was already uploaded locally — is gone andself.best_scorehas already been bumped (lines 92–93). The next validation would not save anything until a new high score appears.Safer pattern: save into a temporary sibling directory and only swap on success, or at least defer
self.best_score/self.best_stepupdates until aftersave_checkpointreturns.🛡️ Minimal: defer state update until after save succeeds
- prev = self.best_score - self.best_score = score - self.best_step = global_steps - print( - f"[BestVal] New best validation score {score:.6f} at step " - f"{global_steps} (previous: {prev})." - ) - best_val_root = os.path.join(self.default_local_dir, "best_val") actor_local_path = os.path.join(best_val_root, "actor") # Flush any in-flight HF upload before mutating best_val/. hf_upload_manager.flush() - # Wipe stale best_val from the previous best step so only the new - # HF model survives the prune step below. - if os.path.isdir(best_val_root): - shutil.rmtree(best_val_root, ignore_errors=True) - - actor_rollout_wg.save_checkpoint( - actor_local_path, None, global_steps, max_ckpt_to_keep=None - ) - _prune_to_huggingface(actor_local_path) - _write_metadata(best_val_root, global_steps, score) - hf_upload_manager.maybe_upload_best_val(global_steps, best_val_root) + staging_root = best_val_root + ".tmp" + if os.path.isdir(staging_root): + shutil.rmtree(staging_root, ignore_errors=True) + actor_rollout_wg.save_checkpoint( + os.path.join(staging_root, "actor"), None, global_steps, max_ckpt_to_keep=None + ) + _prune_to_huggingface(os.path.join(staging_root, "actor")) + _write_metadata(staging_root, global_steps, score) + if os.path.isdir(best_val_root): + shutil.rmtree(best_val_root, ignore_errors=True) + os.replace(staging_root, best_val_root) + + prev = self.best_score + self.best_score = score + self.best_step = global_steps + print( + f"[BestVal] New best validation score {score:.6f} at step " + f"{global_steps} (previous: {prev})." + ) + hf_upload_manager.maybe_upload_best_val(global_steps, best_val_root)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/utils/best_val.py` around lines 99 - 115, The current flow flush() → rmtree(best_val_root) → save_checkpoint(...) risks losing the previous best if actor_rollout_wg.save_checkpoint fails; change to save into a temporary sibling directory (e.g., best_val_tmp) or keep the existing directory until save_checkpoint completes, then atomically swap/rename the temp dir into best_val_root and only call shutil.rmtree on the old dir after success; additionally defer updating self.best_score/self.best_step (the updates around lines 92–93) until after save_checkpoint, _prune_to_huggingface, _write_metadata, and hf_upload_manager.maybe_upload_best_val complete successfully so state only advances on a confirmed persisted upload.
139-145: 💤 Low valueSilently swallowing
OSErroron metadata write hides a bug.If
best_val/best_val.jsonfails to write, future restarts cannot rehydratebest_score/best_step, which is the same problem flagged on the constructor. Logging the exception (still without raising) would make this debuggable.def _write_metadata(best_val_root: str, global_steps: int, score: float) -> None: """Stamp ``best_val/best_val.json`` so consumers know which step won.""" try: with open(os.path.join(best_val_root, "best_val.json"), "w") as f: json.dump({"step": global_steps, "score": score}, f, indent=2) - except OSError: - pass + except OSError as e: + print(f"[BestVal] Warning: failed to write best_val.json: {e}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vagen/utils/best_val.py` around lines 139 - 145, The current _write_metadata function swallows OSError silently when writing best_val.json; change the except OSError: block to except OSError as e: and log the failure (including the exception details) instead of ignoring it — e.g., get a module logger (logging.getLogger(__name__)) and call logger.error or logger.exception with context like f"Failed to write {os.path.join(best_val_root, 'best_val.json')} for step {global_steps}" and include exc_info so the error is recorded, but do not re-raise (keep current behavior of not crashing).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vagen/utils/best_val.py`:
- Around line 46-62: The current aggregate_val_core_score implementation should
ignore booleans and NaNs and optionally allow selecting a single metric name
instead of naïve averaging: update aggregate_val_core_score to accept an
optional best_val_metric (or read trainer.best_val_metric) and if set, try to
return that specific key's numeric value (only if found and not bool/NaN);
otherwise filter val_metrics items with k.startswith("val-core/") and include
only numbers where not isinstance(v, bool) and not math.isnan(float(v)), then
return None if no valid values or the mean of the remaining values; ensure
callers that compare to self.best_score get None for no-valid-metrics so NaNs
never propagate.
- Around line 37-41: In __init__ of the BestVal helper (the constructor shown),
load and parse the existing best_val.json from the trainer default_local_dir
(self.default_local_dir) when save_best_val is enabled, and set self.best_score
and self.best_step from that file instead of leaving them as None; handle the
file not existing or parse errors gracefully (leave None) and ensure the same
path/name used elsewhere for writing best_val is used when reading so resumed
runs correctly rehydrate prior state.
---
Nitpick comments:
In `@vagen/utils/best_val.py`:
- Around line 99-115: The current flow flush() → rmtree(best_val_root) →
save_checkpoint(...) risks losing the previous best if
actor_rollout_wg.save_checkpoint fails; change to save into a temporary sibling
directory (e.g., best_val_tmp) or keep the existing directory until
save_checkpoint completes, then atomically swap/rename the temp dir into
best_val_root and only call shutil.rmtree on the old dir after success;
additionally defer updating self.best_score/self.best_step (the updates around
lines 92–93) until after save_checkpoint, _prune_to_huggingface,
_write_metadata, and hf_upload_manager.maybe_upload_best_val complete
successfully so state only advances on a confirmed persisted upload.
- Around line 139-145: The current _write_metadata function swallows OSError
silently when writing best_val.json; change the except OSError: block to except
OSError as e: and log the failure (including the exception details) instead of
ignoring it — e.g., get a module logger (logging.getLogger(__name__)) and call
logger.error or logger.exception with context like f"Failed to write
{os.path.join(best_val_root, 'best_val.json')} for step {global_steps}" and
include exc_info so the error is recorded, but do not re-raise (keep current
behavior of not crashing).
In `@vagen/utils/upload_hugging_face.py`:
- Around line 84-86: The bare except in the upload handler should be narrowed to
HF/transport errors so programmer errors aren't swallowed: replace "except
Exception as e:" in the upload method (the block referencing self.repo_id and
label) with a specific except that catches huggingface_hub.errors.HfHubHTTPError
plus network/IO exceptions (e.g., requests.RequestException and OSError), log
the error the same way and return None; add the necessary imports
(HfHubHTTPError and requests.RequestException) at the top of the module and do
not catch KeyboardInterrupt/SystemExit so they propagate.
- Around line 122-147: The init currently bails out early when
self._hf_save_freq is falsy which skips validation of self._upload_contents and
creation of the HF actor, so best-val uploads (triggered by
save_best_val/repo_id) silently no-op; change the logic in __init__ so that only
hf-save-frequency-specific setup is skipped when _hf_save_freq is falsy, but
still validate upload_contents (the isinstance check for self._upload_contents)
and perform the actor/actor_kwargs/repo_id setup whenever a repo_id is
configured or save_best_val is enabled (so maybe_upload_best_val can work
independently); update the guards around the blocks that reference actor_kwargs,
repo_id, and actor creation to check for repo_id or save_best_val instead of
requiring _hf_save_freq.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 84ada0e3-e8a1-4346-8df0-e8ad6791d7e5
📒 Files selected for processing (5)
vagen/configs/vagen_multiturn.yamlvagen/evaluate/run_eval.pyvagen/ray_trainer.pyvagen/utils/best_val.pyvagen/utils/upload_hugging_face.py
Comment on lines
+37
to
+41
| def __init__(self, config): | ||
| self.enabled: bool = bool(config.trainer.get("save_best_val", False)) | ||
| self.default_local_dir: str = config.trainer.default_local_dir | ||
| self.best_score: Optional[float] = None | ||
| self.best_step: Optional[int] = None |
Contributor
There was a problem hiding this comment.
best_score/best_step are not persisted across runs.
After a training restart (resume from checkpoint), self.best_score is re-initialized to None, so the first post-resume validation will always be treated as a new best and overwrite the on-disk best_val/ — even if its score is worse than the previously saved best. The best_val.json already on disk contains exactly the state needed to recover this.
🛡️ Suggested fix: rehydrate from existing best_val.json on init
def __init__(self, config):
self.enabled: bool = bool(config.trainer.get("save_best_val", False))
self.default_local_dir: str = config.trainer.default_local_dir
self.best_score: Optional[float] = None
self.best_step: Optional[int] = None
+ # Rehydrate prior best across restarts so we don't regress to a worse checkpoint.
+ meta_path = os.path.join(self.default_local_dir, "best_val", "best_val.json")
+ if self.enabled and os.path.isfile(meta_path):
+ try:
+ with open(meta_path) as f:
+ meta = json.load(f)
+ self.best_score = float(meta.get("score"))
+ self.best_step = int(meta.get("step"))
+ print(
+ f"[BestVal] Rehydrated best score {self.best_score} from step "
+ f"{self.best_step} ({meta_path})."
+ )
+ except (OSError, ValueError, TypeError, KeyError) as e:
+ print(f"[BestVal] Could not load prior best from {meta_path}: {e}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def __init__(self, config): | |
| self.enabled: bool = bool(config.trainer.get("save_best_val", False)) | |
| self.default_local_dir: str = config.trainer.default_local_dir | |
| self.best_score: Optional[float] = None | |
| self.best_step: Optional[int] = None | |
| def __init__(self, config): | |
| self.enabled: bool = bool(config.trainer.get("save_best_val", False)) | |
| self.default_local_dir: str = config.trainer.default_local_dir | |
| self.best_score: Optional[float] = None | |
| self.best_step: Optional[int] = None | |
| # Rehydrate prior best across restarts so we don't regress to a worse checkpoint. | |
| meta_path = os.path.join(self.default_local_dir, "best_val", "best_val.json") | |
| if self.enabled and os.path.isfile(meta_path): | |
| try: | |
| with open(meta_path) as f: | |
| meta = json.load(f) | |
| self.best_score = float(meta.get("score")) | |
| self.best_step = int(meta.get("step")) | |
| print( | |
| f"[BestVal] Rehydrated best score {self.best_score} from step " | |
| f"{self.best_step} ({meta_path})." | |
| ) | |
| except (OSError, ValueError, TypeError, KeyError) as e: | |
| print(f"[BestVal] Could not load prior best from {meta_path}: {e}") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vagen/utils/best_val.py` around lines 37 - 41, In __init__ of the BestVal
helper (the constructor shown), load and parse the existing best_val.json from
the trainer default_local_dir (self.default_local_dir) when save_best_val is
enabled, and set self.best_score and self.best_step from that file instead of
leaving them as None; handle the file not existing or parse errors gracefully
(leave None) and ensure the same path/name used elsewhere for writing best_val
is used when reading so resumed runs correctly rehydrate prior state.
Comment on lines
+46
to
+62
| @staticmethod | ||
| def aggregate_val_core_score(val_metrics: dict) -> Optional[float]: | ||
| """Reduce the trainer's ``val-core/...`` metrics to one scalar. | ||
|
|
||
| ``_validate()`` emits keys shaped like | ||
| ``val-core/<data_source>/<var>/<metric>@<N>``. Higher = better. | ||
| We average all ``val-core`` numeric scalars so the score works | ||
| for single- and multi-dataset configs alike. Returns ``None`` | ||
| when no ``val-core`` scalars exist. | ||
| """ | ||
| scores = [ | ||
| float(v) for k, v in val_metrics.items() | ||
| if k.startswith("val-core/") and isinstance(v, (int, float)) | ||
| ] | ||
| if not scores: | ||
| return None | ||
| return sum(scores) / len(scores) |
Contributor
There was a problem hiding this comment.
Averaging all val-core/ scalars is opinionated and silently broken for mixed-direction or NaN metrics.
A few concerns about aggregate_val_core_score:
isinstance(v, (int, float))admitsbool(Pythonbool⊂int) andfloat('nan'). A single NaN poisons the mean andscore <= self.best_scorebecomesFalsefor every subsequent step (NaN comparisons are always False), causing every NaN run to overwrite the saved best.- Different
val-core/<data_source>/...metrics may be on different scales (e.g., reward vs. accuracy), so a naive unweighted mean can be dominated by the largest-magnitude metric and is not a stable "higher = better" proxy across configs. - There is no way to opt into a single named metric (e.g.,
val-core/<ds>/acc/best@N), which would usually be what users want.
Consider filtering out NaN/bool and optionally exposing a configurable key (e.g. trainer.best_val_metric) that, when set, picks one scalar instead of averaging.
🛡️ Minimal NaN/bool guard
- scores = [
- float(v) for k, v in val_metrics.items()
- if k.startswith("val-core/") and isinstance(v, (int, float))
- ]
+ import math
+ scores = [
+ float(v) for k, v in val_metrics.items()
+ if k.startswith("val-core/")
+ and isinstance(v, (int, float))
+ and not isinstance(v, bool)
+ and not math.isnan(float(v))
+ ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vagen/utils/best_val.py` around lines 46 - 62, The current
aggregate_val_core_score implementation should ignore booleans and NaNs and
optionally allow selecting a single metric name instead of naïve averaging:
update aggregate_val_core_score to accept an optional best_val_metric (or read
trainer.best_val_metric) and if set, try to return that specific key's numeric
value (only if found and not bool/NaN); otherwise filter val_metrics items with
k.startswith("val-core/") and include only numbers where not isinstance(v, bool)
and not math.isnan(float(v)), then return None if no valid values or the mean of
the remaining values; ensure callers that compare to self.best_score get None
for no-valid-metrics so NaNs never propagate.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes / Reliability