BeakerCallback fails without the [beaker] extra: olmo_core.launch.beaker imports gantry at module top
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 315
- Avg merge
- 1d 9h
- Merged PRs (30d)
- 11
Description
**Rev:** observed on `vision` @ `c54d228`, but the import structure looks the same on `main`.
`BeakerCallback.post_attach` does `from olmo_core.launch.beaker import is_running_in_beaker_batch_job`, and `olmo_core/launch/beaker.py` imports `from gantry.api import GitRepoState` at module top. When ai2-olmo-core is installed without the `[beaker]` extra (which declares `beaker-gantry`), attaching the callback in a Beaker batch job crashes:
```
ModuleNotFoundError: No module named 'gantry'
File "olmo_core/train/callbacks/beaker.py", line 49, in post_attach
from olmo_core.launch.beaker import is_running_in_beaker_batch_job
File "olmo_core/launch/beaker.py", line 25, in
from gantry.api import GitRepoState
```
The failure only manifests at trainer attach time inside an actual Beaker job, so it passes local tests and only dies in production.
**Suggested fixes:** guard the gantry import (the callback's description updates via `gantry.api.update_workload_description` could raise a clear "install ai2-olmo-core[beaker]" error or degrade to a no-op), or move `is_running_in_beaker_batch_job` somewhere gantry-free.
Found while reusing OLMo-core-native callbacks in allenai/open-instruct (open-instruct#1856).
Contributor guide
Assessment
This issue has not been assessed yet.