allenai / allenai/OLMo-core

BeakerCallback fails without the [beaker] extra: olmo_core.launch.beaker imports gantry at module top

Open
#850 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.