Add budget support for dataset and image creation

CHANGELOG.md

@@ -9,6 +9,11 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added optional `budget` parameter to `Beaker.dataset.create()` method for specifying budget during dataset creation.
+- Added optional `budget` parameter to `Beaker.image.create()` method for specifying budget during image creation.
+
 ## [v1.37.0](https://github.com/allenai/beaker-py/releases/tag/v1.37.0) - 2025-09-15
 
 ### Added

beaker/data_model/dataset.py

@@ -275,6 +275,7 @@ class DatasetInfo(BaseModel):
 class DatasetSpec(BaseModel):
     workspace: Optional[str] = None
     description: Optional[str] = None
+    budget: Optional[str] = None
 
 
 class DatasetPatch(BaseModel):

beaker/data_model/image.py

@@ -115,6 +115,7 @@ class ImageSpec(BaseModel):
     image_id: Optional[str] = None
     image_tag: Optional[str] = None
     description: Optional[str] = None
+    budget: Optional[str] = None
 
 
 class ImagePatch(BaseModel):

beaker/services/dataset.py

@@ -92,6 +92,7 @@ def create(
         target: Optional[PathOrStr] = None,
         workspace: Optional[str] = None,
         description: Optional[str] = None,
+        budget: Optional[str] = None,
         force: bool = False,
         max_workers: Optional[int] = None,
         quiet: bool = False,
@@ -108,6 +109,7 @@ def create(
         :param workspace: The workspace to upload the dataset to. If not specified,
             :data:`Beaker.config.default_workspace <beaker.Config.default_workspace>` is used.
         :param description: Text description for the dataset.
+        :param budget: Budget to associate with the dataset. If not specified, uses workspace default if available.
         :param force: If ``True`` and a dataset by the given name already exists, it will be overwritten.
         :param max_workers: The maximum number of thread pool workers to use to upload files concurrently.
         :param quiet: If ``True``, progress won't be displayed.
@@ -142,7 +144,11 @@ def make_dataset() -> Dataset:
                     "datasets",
                     method="POST",
                     query={"name": name},
-                    data=DatasetSpec(workspace=workspace_id, description=description),
+                    data=DatasetSpec(
+                        workspace=workspace_id,
+                        description=description,
+                        budget=self.resolve_budget(budget) if budget is not None else None,
+                    ),
                     exceptions_for_status={409: DatasetConflict(name)},
                 ).json()
             )

beaker/services/image.py

@@ -58,6 +58,7 @@ def create(
         image_tag: str,
         workspace: Optional[Union[Workspace, str]] = None,
         description: Optional[str] = None,
+        budget: Optional[str] = None,
         quiet: bool = False,
         commit: bool = True,
     ) -> Image:
@@ -69,6 +70,7 @@ def create(
         :param workspace: The workspace to upload the image to. If not specified,
             :data:`Beaker.config.default_workspace <beaker.Config.default_workspace>` is used.
         :param description: Text description of the image.
+        :param budget: Budget to associate with the image. If not specified, uses workspace default.
         :param quiet: If ``True``, progress won't be displayed.
         :param commit: Whether to commit the image after successful upload.
 
@@ -96,6 +98,7 @@ def create(
                 image_id=image.id,
                 image_tag=image_tag,
                 description=description,
+                budget=self.resolve_budget(budget) if budget else None,
             ),
             query={"name": name},
             exceptions_for_status={409: ImageConflict(name)},

integration_tests/datasets_test.py

@@ -33,6 +33,7 @@ def test_dataset_write_error(self, client: Beaker, dataset_name: str):
             client.dataset.sync(dataset, self.file_b.name)
 
     def test_dataset_basics(self, client: Beaker, dataset_name: str, alternate_dataset_name: str):
+        # Test create dataset without budget parameter
         dataset = client.dataset.create(
             dataset_name,
             self.file_a.name,
@@ -59,6 +60,23 @@ def test_dataset_basics(self, client: Beaker, dataset_name: str, alternate_datas
         dataset = client.dataset.rename(dataset, alternate_dataset_name)
         assert dataset.name == alternate_dataset_name
 
+        # Test with dataset creation with specified budget parameter
+        dataset_with_budget = client.dataset.create(
+            alternate_dataset_name + "-budget",
+            self.file_a.name,
+            budget="ai2/compute",
+            commit=True,
+            description="Testing dataset with budget",
+        )
+        assert dataset_with_budget.name == alternate_dataset_name + "-budget"
+
+        contents = b"".join(
+            list(client.dataset.stream_file(dataset_with_budget, Path(self.file_a.name).name))
+        )
+        assert contents == self.file_a_contents
+
+        client.dataset.delete(dataset_with_budget)
+
 
 class TestLargeFileDataset:
     def setup_method(self):

integration_tests/images_test.py

@@ -24,6 +24,18 @@ def test_image_create_workflow(
     image = client.image.rename(image, alternate_beaker_image_name)
     assert image.name == alternate_beaker_image_name
 
+    # Test with budget parameter
+    print(f"Creating image with budget '{beaker_image_name}-budget'")
+    image_with_budget = client.image.create(
+        beaker_image_name + "-budget",
+        LOCAL_IMAGE_TAG,
+        budget="ai2/compute",
+        description="Test image with budget",
+    )
+    assert image_with_budget.name == beaker_image_name + "-budget"
+    assert image_with_budget.original_tag == LOCAL_IMAGE_TAG
+    client.image.delete(image_with_budget)
+
 
 def test_image_pull_workflow(client: Beaker, beaker_python_image_name: str):
     print(f"Pulling image '{beaker_python_image_name}' from Beaker")