allenai/beaker-py

A pure-Python Beaker client

beaker-py

A lightweight, standalone, pure Python client for Beaker

---

Features

🪶 Lightweight

• Minimal dependencies.
• Only pure-Python dependencies.
• Communicates directly with the Beaker server via HTTP requests (Beaker CLI not required).

💪 Robust

• Automatically retries failed HTTP requests with exponential backoff.
• Runtime data validation.
• High test coverage.

📓 Exhaustively-typed and documented

• Thorough data model for all input / output types.
• Every expected HTTP error from the Beaker server is translated into a specific exception type.

Quick links

• Documentation
• PyPI package
• Contributing
• License

See also 👇

• Beaker (CLI)
• Beaker Gantry
• Beaker-relevant GitHub Actions
  • setup-beaker
  • beaker-command
  • beaker-run

Installing

Installing with pip

beaker-py is available on PyPI. Just run

pip install beaker-py

Installing from source

To install beaker-py from source, first clone the repository:

git clone https://github.com/allenai/beaker-py.git
cd beaker-py

Then run

pip install -e .

Quick start

If you've already configured the Beaker command-line client, beaker-py will find and use the existing configuration file (usually located at $HOME/.beaker/config.yml). Otherwise just set the environment variable BEAKER_TOKEN to your Beaker user token.

Either way, you should then instantiate the Beaker client with .from_env():

from beaker import Beaker

beaker = Beaker.from_env(default_workspace="my_org/my_workspace")

The API of beaker-py is meant to mirror - as closely as possible - the API of the Beaker CLI. For example, when you do this with the CLI:

beaker dataset create --name foo .

The beaker-py equivalent would be:

beaker.dataset.create("foo", ".")

See the API Docs to learn about the Beaker client's methods.

beaker-py's Issues

Add dataset search method

Want to be able to search based on name or description.

Add method to list nodes with free gpus

@codeviking already wrote a script for this: https://gist.github.com/codeviking/602dd302a1d698805efc8dbb104d8821. So we'll translate those beaker CLI calls into Python.

Create hyperparameter tuning example

JSON validation error when calling beaker.cluster.utilization

🐛 Describe the bug

Repro:

from beaker import Beaker

beaker = Beaker.from_env(default_workspace="ai2/rbg_dev")
beaker.cluster.utilization("ai2/prior-cirrascale")

Results in:

---------------------------------------------------------------------------
ValidationError                           Traceback (most recent call last)
Cell In[9], line 1
----> 1 beaker.cluster.utilization("ai2/prior-cirrascale")

File ~/miniconda3/envs/base-dev/lib/python3.11/site-packages/beaker/services/cluster.py:221, in ClusterClient.utilization(self, cluster)
    210 running_preemptible_jobs = 0
    211 node_to_util: Dict[str, Dict[str, Union[int, float]]] = {
    212     node.id: {
    213         "running_jobs": 0,
   (...)
    218     for node in nodes
    219 }
--> 221 for job in self.beaker.job.list(cluster=cluster, finalized=False):
    222     if job.status.current in (CurrentJobStatus.running, CurrentJobStatus.idle):
    223         if job.node not in node_to_util:

File ~/miniconda3/envs/base-dev/lib/python3.11/site-packages/beaker/services/job.py:100, in JobClient.list(self, cluster, experiment, finalized, kind, node)
     98 # Gather jobs, page by page.
     99 while True:
--> 100     page = Jobs.from_json(self.request("jobs", method="GET", query=request_opts).json())
    101     if page.data:
    102         jobs.extend(page.data)

File ~/miniconda3/envs/base-dev/lib/python3.11/site-packages/beaker/data_model/base.py:68, in BaseModel.from_json(cls, json_data)
     65 @classmethod
     66 def from_json(cls: Type[T], json_data: Dict[str, Any]) -> T:
     67     try:
---> 68         return cls(**json_data)
     69     except ValidationError:
     70         logger.error("Error validating raw JSON data for %s: %s", cls.__name__, json_data)

File ~/miniconda3/envs/base-dev/lib/python3.11/site-packages/pydantic/main.py:171, in BaseModel.__init__(self, **data)
    169 # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks
    170 __tracebackhide__ = True
--> 171 self.__pydantic_validator__.validate_python(data, self_instance=self)

ValidationError: 8 validation errors for Jobs
data.21.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.22.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.23.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.24.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.25.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.26.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.27.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]
data.28.status.canceled_code
  Input should be 0, 1, 2, 3 or 4 [type=enum, input_value=6, input_type=int]

Versions

Python 3.11.7
annotated-types==0.6.0
appnope @ file:///home/conda/feedstock_root/build_artifacts/appnope_1707233003401/work
asttokens @ file:///home/conda/feedstock_root/build_artifacts/asttokens_1698341106958/work
attrs==23.2.0
backcall==0.2.0
beaker-gantry==0.21.0
beaker-py==1.26.3
beautifulsoup4==4.12.3
bleach==6.1.0
certifi==2024.2.2
charset-normalizer==3.3.2
click==8.1.7
click-help-colors==0.9.4
comm @ file:///home/conda/feedstock_root/build_artifacts/comm_1704278392174/work
contourpy==1.2.0
cycler==0.12.1
debugpy @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_563_nwtkoc/croot/debugpy_1690905063850/work
decorator @ file:///home/conda/feedstock_root/build_artifacts/decorator_1641555617451/work
defusedxml==0.7.1
docker==6.1.3
docopt==0.6.2
exceptiongroup @ file:///home/conda/feedstock_root/build_artifacts/exceptiongroup_1704921103267/work
executing @ file:///home/conda/feedstock_root/build_artifacts/executing_1698579936712/work
fastjsonschema==2.19.1
filelock==3.13.1
fonttools==4.48.1
fsspec==2024.2.0
gitdb==4.0.11
GitPython==3.1.41
idna==3.6
importlib-metadata @ file:///home/conda/feedstock_root/build_artifacts/importlib-metadata_1703269254275/work
ipykernel @ file:///Users/runner/miniforge3/conda-bld/ipykernel_1707326353119/work
ipython==8.12.3
jedi @ file:///home/conda/feedstock_root/build_artifacts/jedi_1696326070614/work
Jinja2==3.1.3
jsonschema==4.21.1
jsonschema-specifications==2023.12.1
jupyter-console @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_62liw5pns2/croot/jupyter_console_1679999641189/work
jupyter_client @ file:///home/conda/feedstock_root/build_artifacts/jupyter_client_1699283905679/work
jupyter_core @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_782yoyc_98/croot/jupyter_core_1698937318631/work
jupyterlab_pygments==0.3.0
kiwisolver==1.4.5
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.8.2
matplotlib-inline @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-inline_1660814786464/work
mdurl==0.1.2
mistune==3.0.2
mpmath==1.3.0
nbclient==0.9.0
nbconvert==7.16.1
nbformat==5.9.2
nest_asyncio @ file:///home/conda/feedstock_root/build_artifacts/nest-asyncio_1705850609492/work
networkx==3.2.1
numpy==1.26.4
packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1696202382185/work
pandas==2.2.0
pandocfilters==1.5.1
parso @ file:///home/conda/feedstock_root/build_artifacts/parso_1638334955874/work
petname==2.6
pexpect @ file:///home/conda/feedstock_root/build_artifacts/pexpect_1706113125309/work
pickleshare @ file:///home/conda/feedstock_root/build_artifacts/pickleshare_1602536217715/work
pillow==10.2.0
pipreqs==0.5.0
platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1706713388748/work
prompt-toolkit @ file:///home/conda/feedstock_root/build_artifacts/prompt-toolkit_1702399386289/work
psutil @ file:///Users/cbousseau/work/recipes/ci_py311_2/psutil_1678995687212/work
ptyprocess @ file:///home/conda/feedstock_root/build_artifacts/ptyprocess_1609419310487/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
pure-eval @ file:///home/conda/feedstock_root/build_artifacts/pure_eval_1642875951954/work
pydantic==2.6.1
pydantic_core==2.16.2
Pygments @ file:///home/conda/feedstock_root/build_artifacts/pygments_1700607939962/work
pyparsing==3.1.1
python-dateutil @ file:///home/conda/feedstock_root/build_artifacts/python-dateutil_1626286286081/work
pytz==2024.1
PyYAML==6.0.1
pyzmq @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_43pxpbos3z/croot/pyzmq_1705605108344/work
referencing==0.33.0
requests==2.31.0
rich==13.7.0
rpds-py==0.18.0
seaborn==0.13.2
six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work
smmap==5.0.1
soupsieve==2.5
stack-data @ file:///home/conda/feedstock_root/build_artifacts/stack_data_1669632077133/work
sympy==1.12
tinycss2==1.2.1
torch==2.2.0
torchaudio==2.2.0
torchvision==0.17.0
tornado @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_3a5nrn2jeh/croot/tornado_1696936974091/work
traitlets @ file:///home/conda/feedstock_root/build_artifacts/traitlets_1704212992681/work
typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1702176139754/work
tzdata==2024.1
urllib3==2.2.0
wcwidth @ file:///home/conda/feedstock_root/build_artifacts/wcwidth_1704731205417/work
webencodings==0.5.1
websocket-client==1.7.0
yarg==0.1.9
zipp @ file:///home/conda/feedstock_root/build_artifacts/zipp_1695255097490/work

Add `since` option to `Beaker.job.logs()`

See https://github.com/allenai/beaker-service/pull/2120

Make sure quiet progress works in threads

Option to raise error from `.wait_on` when job fails

Add strict: bool argument to Beaker.(experiment|job).(wait_on|as_completed) that - when True - causes an error to be raised if a job fails / finalizing with non-zero exit code.

Add `description` parameter to `Beaker.*.create()` methods

tolerating unexpected fields in config file

Looks like beaker-py uses yaml.load() which does not allow unexpected fields in the config file.

Minimal repro of usage:

import yaml

from dataclasses import asdict, dataclass, fields

@dataclass
class Config:
    foo: str

    @classmethod
    def loadit(cls):
        with open("somefile.yaml") as config_file:
            return cls(**yaml.load(config_file, Loader=yaml.SafeLoader))

Config.loadit()

with file somefile.yaml:

a: 42
foo: "Asdf"

yields:

TypeError: __init__() got an unexpected keyword argument 'a'

As an aside:

The PyYAML documentation at https://pyyaml.org/wiki/PyYAMLDocumentation suggests not to use yaml.load() from untrusted sources:

Warning: It is not safe to call yaml.load with any data received from an untrusted source! yaml.load is as powerful as pickle.load and so may call any Python function. Check the yaml.safe_load function though.

Is that better to use than providing Loader=yaml.SafeLoader?

`Beaker.experiment.follow()` shouldn't have to make so many calls to `experiments/tasks/` API

Complete remaining components of Beaker API

Datasets

• beaker.dataset.ls (52952ab)
• beaker.dataset.rename (2a4aebd)
• beaker.dataset.size (a1e765b)
• beaker.dataset.stream_file (af423e6)
• beaker.dataset.sync (b07b5c8)

Experiments

• beaker.experiment.delete
• beaker.experiment.rename (#58)
• beaker.experiment.results (#59)
• beaker.experiment.resume (2cb9f6a)
• beaker.experiment.spec (#59)
• beaker.experiment.stop (3eec1bb)
• beaker.experiment.tasks (#59)

Jobs

• beaker.job.get (#55)
• beaker.job.finalize (#70)
• beaker.job.list (#55)
• beaker.job.results (d9f267a)
• beaker.job.stop (#70)

Workspaces

• beaker.workspace.archive (#71)
• beaker.workspace.datasets (#60)
• beaker.workspace.experiments (#60)
• beaker.workspace.images (#60)
• beaker.workspace.secrets (#61)
• beaker.workspace.list (#60)
• beaker.workspace.move (#71)
• beaker.workspace.permissions (#73)
• beaker.workspace.rename (#71)
• beaker.workspace.unarchive (#71)

Nodes

• beaker.node.get (#55)

Secrets

Added in #61

• beaker.secret.read
• beaker.secret.write
• beaker.secret.delete

Groups

Added in #76

• beaker.group.add_experiments
• beaker.group.create
• beaker.group.delete
• beaker.group.list_experiments
• beaker.group.get
• beaker.group.remove_experiments
• beaker.group.rename

Support directory sources for `Beaker.dataset.create()`

Remove all deprecated methods / parameters

Add `Beaker.experiment.await()` method

Improve the `BufferedReaderWithProgress` class

Should implement all required methods.

Follow what rich does here: https://rich.readthedocs.io/en/stable/_modules/rich/progress.html#Progress with the _Reader class.

Method to get the URL for an experiment

beaker.experiment.url()

Fetching experiments from a workspace results in exception

🐛 Describe the bug

I am trying to get all experiments from a workspace:

exps = beaker_client.workspace.experiments(workspace="ai2/nlpgym2.0")
print(len(exps))

But this results in following error:

File "/Users/rramamurth/opt/miniconda3/envs/nlpgym2/lib/python3.7/site-packages/beaker/data_model/base.py", line 51, in from_json return cls(**json_data) File "pydantic/main.py", line 341, in pydantic.main.BaseModel.__init__ pydantic.error_wrappers.ValidationError: 1 validation error for ExperimentsPage data -> 95 -> jobs -> 2 -> execution -> spec -> image -> __root__ Exactly one of 'beaker' or 'docker' must be specified for image source (type=value_error)

Versions

Python 3.7.13
absl-py==1.0.0
aiohttp==3.8.1
aiosignal==1.2.0
astunparse==1.6.3
async-timeout==4.0.2
asynctest==0.13.0
attrs==21.4.0
autopep8==1.6.0
beaker-py==1.6.4
bert-score==0.3.11
BLEURT @ git+https://github.com/google-research/bleurt.git@c6f2375c7c178e1480840cf27cb9e2af851394f9
cached-property==1.5.2
cachetools==5.1.0
certifi==2022.5.18.1
charset-normalizer==2.0.12
click==8.1.3
cloudpickle==1.6.0
colorama==0.4.4
commonmark==0.9.1
cycler==0.11.0
datasets==2.0.0
dill==0.3.5.1
diskcache==5.4.0
docker==5.0.0
docker-pycreds==0.4.0
filelock==3.7.0
flake8==4.0.1
flatbuffers==1.12
fonttools==4.33.3
frozenlist==1.3.0
fsspec==2022.5.0
gast==0.4.0
gem-metrics @ git+https://github.com/GEM-benchmark/GEM-metrics.git@431a8174bd6b3637e8d6118bfad2983e39e99733
gitdb==4.0.9
GitPython==3.1.27
google-auth==2.6.6
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
grpcio==1.46.3
gym==0.19.0
h5py==3.6.0
huggingface-hub==0.6.0
idna==3.3
importlib-metadata==4.11.4
iniconfig==1.1.1
joblib==1.1.0
jsonlines==3.0.0
keras==2.9.0
Keras-Preprocessing==1.1.2
kiwisolver==1.4.2
libclang==14.0.1
logzero==1.7.0
Markdown==3.3.7
matplotlib==3.5.2
mccabe==0.6.1
multidict==6.0.2
multiprocess==0.70.12.2
-e git+ssh://[email protected]/allenai/nlp-gym2.git@b40a06a44c342fe0cb49a6b2f6911224d56902a5#egg=nlp_gym
nltk==3.7
numpy==1.21.6
oauthlib==3.2.0
opt-einsum==3.3.0
overrides==3.1.0
packaging==21.3
pandas==1.3.5
pathtools==0.1.2
Pillow==9.1.1
pluggy==0.13.1
portalocker==2.4.0
promise==2.3
protobuf==3.20.1
psutil==5.9.1
py==1.11.0
pyarrow==8.0.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pycodestyle==2.8.0
pycountry==22.3.5
pydantic==1.9.1
pyflakes==2.4.0
Pygments==2.12.0
pyparsing==3.0.9
pytest==6.2.4
python-dateutil==2.8.2
pytz==2022.1
PyYAML==6.0
regex==2022.4.24
repro==0.1.4
requests==2.27.1
requests-oauthlib==1.3.1
responses==0.18.0
rich==12.4.4
rouge-score==0.0.4
rsa==4.8
sacrebleu==2.1.0
sacremoses==0.0.53
scipy==1.7.3
seaborn==0.11.2
sentencepiece==0.1.96
sentry-sdk==1.5.12
setproctitle==1.2.3
shortuuid==1.0.9
six==1.16.0
smmap==5.0.0
stable-baselines3==1.4.0
tabulate==0.8.9
tensorboard==2.9.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.9.0
tensorflow-estimator==2.9.0
tensorflow-io-gcs-filesystem==0.26.0
termcolor==1.1.0
tf-slim==1.1.0
tokenizers==0.12.1
toml==0.10.2
torch==1.11.0
torchvision==0.12.0
tqdm==4.64.0
transformers==4.18.0
typing_extensions==4.2.0
urllib3==1.26.9
wandb==0.12.15
websocket-client==1.3.2
Werkzeug==2.1.2
wrapt==1.14.1
xxhash==3.0.0
yarl==1.7.2
zipp==3.8.0

Beaker added a "allow_preemptible_restriction_exceptions" field

🐛 Describe the bug

...and it needs to be added to Beaker-py. I received a user warning today about it

RuntimeWarning: Found unknown field 'allow_preemptible_restriction_exceptions: True' for data model 'XXX'. This may be a newly added field that hasn't been defined in beaker-py yet. Please submit an issue report about this here:
https://github.com/allenai/beaker-py/issues/new?assignees=&labels=bug&template=bug_report.yml

Versions

Python 3.11.3

Add group CSV export method

See https://github.com/allenai/beaker-service/pull/2110

`beaker.experiment.create` should allow for an empty (or None) name

🚀 The feature, motivation and pitch

The beaker command line, such as beaker experiment create spec.yaml, does not require that you specify a name. If no name is provided it defaults to the experiment id. This is not the same behavior as in beaker.experiment.create which requires a non empty string name.

I am running many experiments using beaker.experiment.create, some of which fail (at least during the testing process). Thus coming up with unique, meaningful names for each experiment is challenging (recall that beaker requires experiment names to be unique, unless they have been manually deleted). I can get around this by using a hash or something of that sort but it would be more elegant to simply allow the experiment id to be the name. Thanks for considering!

Alternatives

No response

Additional context

No response

Document / improve experience for running tests locally

Deprecated Field Notice: DatasetFIleInfo.URL

See https://github.com/allenai/beaker/pull/2937#discussion_r1046679574.

In the next Beaker release we'll be removing DatasetInfo.URL. It was always "", so it shouldn't have much effect.

I believe it's modeled in allenai/beaker-py here:

beaker-py/beaker/data_model/dataset.py

Lines 129 to 142 in c68a51d

	class FileInfo(BaseModel, arbitrary_types_allowed=True):
	path: str
	digest: Digest
	updated: datetime
	size: Optional[int] = None
	"""
	The size of the file, if known.
	"""
	url: Optional[str] = None
	"""
	A URL that can be used to directly download the file.
	"""

...and because it's Optional, nothing should fail, I think!