Skip to content

Fix batch_func validation for invalid output coords#890

Open
fallintoplace wants to merge 1 commit into
NVIDIA:mainfrom
fallintoplace:fix/batch-guard-output-coords
Open

Fix batch_func validation for invalid output coords#890
fallintoplace wants to merge 1 commit into
NVIDIA:mainfrom
fallintoplace:fix/batch-guard-output-coords

Conversation

@fallintoplace
Copy link
Copy Markdown

Closes #889

Description

  • reject invalid output_coords earlier in batch_func._compress_batch
  • add a regression test for models whose output coordinates do not start with batch

Why

The existing guard only raised when both input_coords and output_coords were invalid. If input_coords was valid but output_coords was not, the call got past the compatibility check and later failed in _decompress_batch with KeyError: 'batch' instead of the intended early ValueError.\n\n## Impact\nThis makes the batching path fail fast with the compatibility error for misconfigured models and keeps the failure mode predictable.\n\n## Validation\n- uv run pytest 'test/models/test_batch.py::test_invalid_output_coords_raise' 'test/models/test_batch.py::test_invalid_ordering_raises[cpu]'\n- uv run ruff check earth2studio/models/batch.py test/models/test_batch.py

@fallintoplace fallintoplace marked this pull request as ready for review May 29, 2026 23:53
@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps Bot commented May 29, 2026

Greptile Summary

Fixes a silent pass-through in batch_func._compress_batch where a model with valid input_coords but invalid output_coords (missing the leading "batch" key) would bypass the early guard and later crash with an opaque KeyError in _decompress_batch.

  • The validation condition is corrected from and to or, so a ValueError is raised immediately if either coordinate system is missing "batch" as its first key.
  • A new regression test (test_invalid_output_coords_raise) constructs a model that deletes "batch" from its output coords and asserts the expected ValueError is raised at call time.

Confidence Score: 5/5

Safe to merge — the one-character operator change is logically correct and directly addresses the described failure mode.

The change is minimal and well-reasoned: or correctly flags either coordinate system being misconfigured, while and only caught the case where both were wrong. The new test directly reproduces the previously-silent failure path and confirms the fix holds.

No files require special attention.

Important Files Changed

Filename Overview
earth2studio/models/batch.py Single-character logic fix: andor in the batch-dimension guard, so a valid input_coords paired with an invalid output_coords now correctly raises ValueError instead of falling through to a KeyError later.
test/models/test_batch.py Adds test_invalid_output_coords_raise — a focused regression test for the exact bug scenario where input_coords starts with "batch" but output_coords removes it, ensuring the ValueError is raised at compress time.

Reviews (1): Last reviewed commit: "Fix batch guard for invalid output coord..." | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

batch_func lets invalid output coords through and then crashes later

1 participant