beaker-py 2.0.0.dev0__tar.gz → 2.0.1__tar.gz

beaker_py-2.0.1/PKG-INFO

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: beaker-py
+Version: 2.0.1
+Summary: A Python Beaker client
+Author-email: Allen Institute for Artificial Intelligence <contact@allenai.org>, Pete Walsh <petew@allenai.org>
+Project-URL: Homepage, https://github.com/allenai/beaker
+Project-URL: Repository, https://github.com/allenai/beaker
+Project-URL: Changelog, https://github.com/allenai/beaker/blob/main/bindings/python/CHANGELOG.md
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: requests
+Requires-Dist: packaging
+Requires-Dist: PyYAML
+Requires-Dist: grpcio>=1.70.0
+Requires-Dist: protobuf<6.0,>=5.0
+Requires-Dist: google-crc32c
+Provides-Extra: dev
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mypy<1.6,>=1.0; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+Requires-Dist: types-cachetools; extra == "dev"
+Requires-Dist: types-PyYAML; extra == "dev"
+Requires-Dist: types-protobuf; extra == "dev"
+Requires-Dist: black<24.0,>=23.0; extra == "dev"
+Requires-Dist: isort<5.13,>=5.12; extra == "dev"
+Requires-Dist: pytest<8.0; extra == "dev"
+Requires-Dist: twine>=1.11.0; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Requires-Dist: petname==2.6; extra == "dev"
+Requires-Dist: grpcio-tools; extra == "dev"
+
+# Beaker-py
+
+A lightweight pure-Python client for Beaker.
+
+## Installing
+
+### Installing with `pip`
+
+**beaker-py** is available [on PyPI](https://pypi.org/project/beaker-py/). Just run
+
+```bash
+pip install beaker-py
+```
+
+### Installing from source
+
+To install **beaker-py** from source, first clone [the repository](https://github.com/allenai/beaker):
+
+```bash
+git clone https://github.com/allenai/beaker.git
+```
+
+Then create or activate a Python virtual environment, and run:
+
+```bash
+cd beaker/bindings/python
+make dev-install
+```
+
+## Quick start
+
+If you've already configured the [Beaker command-line client](https://github.com/allenai/beaker/),
+**beaker-py** will find and use the existing configuration file (usually located at `$HOME/.beaker/config.yml`) or `BEAKER_TOKEN` environment variable.
+
+Then you can instantiate the Beaker client with the `.from_env()` class method:
+
+```python
+from beaker import Beaker
+
+with Beaker.from_env() as beaker:
+    ...
+```
+
+With the Python client, you can:
+- Query [**Clusters**](https://beaker-docs.apps.allenai.org/concept/clusters.html) with `beaker.cluster.*` methods, e.g. `beaker.cluster.get("ai2/jupiter-cirrascale-2")`.
+- Manage [**Datasets**](https://beaker-docs.apps.allenai.org/concept/datasets.html) with `beaker.dataset.*` methods, e.g. `beaker.dataset.create(dataset_name, source_dir)`.
+- Submit, track, and find [**Experiments**](https://beaker-docs.apps.allenai.org/concept/experiments.html) with `beaker.experiment.*`, `beaker.workload.*`, and `beaker.job.*` methods, e.g. `beaker.experiment.create(spec=spec, name=name)`.
+- Manage [**Workspaces**](https://beaker-docs.apps.allenai.org/concept/workspaces.html) with `beaker.workspace.*` methods, e.g. `beaker.workspace.create("ai2/new_workspace")`.
+- Manage [**Secrets**](https://beaker-docs.apps.allenai.org/concept/secrets.html) with `beaker.secret.*` methods, e.g. `beaker.secret.write(name, value)`.
+
+### Example workflow
+
+Launch and follow an experiment like [beaker-gantry](https://github.com/allenai/beaker-gantry) does:
+
+```python
+import time
+from beaker import Beaker, BeakerExperimentSpec, BeakerJobPriority
+
+
+with Beaker.from_env() as beaker:
+    # Build experiment spec...
+    spec = BeakerExperimentSpec.new(
+        description="beaker-py test run",
+        beaker_image="petew/hello-world",
+        priority=BeakerJobPriority.low,
+        preemptible=True,
+    )
+
+    # Create experiment workload...
+    workload = beaker.experiment.create(spec=spec)
+
+    # Wait for job to be created...
+    while (job := beaker.workload.get_latest_job(workload)) is None:
+        print("waiting for job to start...")
+        time.sleep(1.0)
+
+    # Follow logs...
+    print("Job logs:")
+    for job_log in beaker.job.logs(job, follow=True):
+        print(job_log.message.decode())
+```
+
+## Development
+
+After [installing from source](#installing-from-source), you can run checks and tests locally with:
+
+```bash
+make checks
+```
+
+### Releases
+
+At the moment releases need to be published manually by following these steps:
+
+1. Ensure you've authenticated with [PyPI](https://pypi.org/) through a `~/.pypirc` file and have write permissions to the [beaker-py project](https://pypi.org/project/beaker-py/).
+2. Ensure the target release version defined in `src/beaker/version.py` is correct, or change the version on the fly by adding the `Make` argument `BEAKER_PY_VERSION=X.X.X` to the command in the next step.
+3. Run `make publish` for a stable release or `make publish-nightly` for a nightly pre-release.

beaker_py-2.0.1/README.md

@@ -0,0 +1,97 @@
+# Beaker-py
+
+A lightweight pure-Python client for Beaker.
+
+## Installing
+
+### Installing with `pip`
+
+**beaker-py** is available [on PyPI](https://pypi.org/project/beaker-py/). Just run
+
+```bash
+pip install beaker-py
+```
+
+### Installing from source
+
+To install **beaker-py** from source, first clone [the repository](https://github.com/allenai/beaker):
+
+```bash
+git clone https://github.com/allenai/beaker.git
+```
+
+Then create or activate a Python virtual environment, and run:
+
+```bash
+cd beaker/bindings/python
+make dev-install
+```
+
+## Quick start
+
+If you've already configured the [Beaker command-line client](https://github.com/allenai/beaker/),
+**beaker-py** will find and use the existing configuration file (usually located at `$HOME/.beaker/config.yml`) or `BEAKER_TOKEN` environment variable.
+
+Then you can instantiate the Beaker client with the `.from_env()` class method:
+
+```python
+from beaker import Beaker
+
+with Beaker.from_env() as beaker:
+    ...
+```
+
+With the Python client, you can:
+- Query [**Clusters**](https://beaker-docs.apps.allenai.org/concept/clusters.html) with `beaker.cluster.*` methods, e.g. `beaker.cluster.get("ai2/jupiter-cirrascale-2")`.
+- Manage [**Datasets**](https://beaker-docs.apps.allenai.org/concept/datasets.html) with `beaker.dataset.*` methods, e.g. `beaker.dataset.create(dataset_name, source_dir)`.
+- Submit, track, and find [**Experiments**](https://beaker-docs.apps.allenai.org/concept/experiments.html) with `beaker.experiment.*`, `beaker.workload.*`, and `beaker.job.*` methods, e.g. `beaker.experiment.create(spec=spec, name=name)`.
+- Manage [**Workspaces**](https://beaker-docs.apps.allenai.org/concept/workspaces.html) with `beaker.workspace.*` methods, e.g. `beaker.workspace.create("ai2/new_workspace")`.
+- Manage [**Secrets**](https://beaker-docs.apps.allenai.org/concept/secrets.html) with `beaker.secret.*` methods, e.g. `beaker.secret.write(name, value)`.
+
+### Example workflow
+
+Launch and follow an experiment like [beaker-gantry](https://github.com/allenai/beaker-gantry) does:
+
+```python
+import time
+from beaker import Beaker, BeakerExperimentSpec, BeakerJobPriority
+
+
+with Beaker.from_env() as beaker:
+    # Build experiment spec...
+    spec = BeakerExperimentSpec.new(
+        description="beaker-py test run",
+        beaker_image="petew/hello-world",
+        priority=BeakerJobPriority.low,
+        preemptible=True,
+    )
+
+    # Create experiment workload...
+    workload = beaker.experiment.create(spec=spec)
+
+    # Wait for job to be created...
+    while (job := beaker.workload.get_latest_job(workload)) is None:
+        print("waiting for job to start...")
+        time.sleep(1.0)
+
+    # Follow logs...
+    print("Job logs:")
+    for job_log in beaker.job.logs(job, follow=True):
+        print(job_log.message.decode())
+```
+
+## Development
+
+After [installing from source](#installing-from-source), you can run checks and tests locally with:
+
+```bash
+make checks
+```
+
+### Releases
+
+At the moment releases need to be published manually by following these steps:
+
+1. Ensure you've authenticated with [PyPI](https://pypi.org/) through a `~/.pypirc` file and have write permissions to the [beaker-py project](https://pypi.org/project/beaker-py/).
+2. Ensure the target release version defined in `src/beaker/version.py` is correct, or change the version on the fly by adding the `Make` argument `BEAKER_PY_VERSION=X.X.X` to the command in the next step.
+3. Run `make publish` for a stable release or `make publish-nightly` for a nightly pre-release.

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/pyproject.toml

@@ -9,7 +9,6 @@ readme = "README.md"
 description = "A Python Beaker client"
 classifiers = [
     "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
@@ -22,10 +21,10 @@ requires-python = ">=3.10"
 dependencies = [
     "requests",
     "packaging",
-    "rich>=12.3,<14.0",
     "PyYAML",
     "grpcio>=1.70.0",  # NOTE: when updating min version, make sure to also update it in the Makefile
-    "protobuf>=5.0",
+    "protobuf>=5.0,<6.0",
+    "google-crc32c",
 ]
 
 [project.optional-dependencies]
@@ -50,8 +49,8 @@ dev = [
 [project.urls]
 Homepage = "https://github.com/allenai/beaker"
 Repository = "https://github.com/allenai/beaker"
-Changelog = "https://github.com/allenai/beaker-py/blob/main/CHANGELOG.md"
-Documentation = "https://beaker-py.readthedocs.io/"
+Changelog = "https://github.com/allenai/beaker/blob/main/bindings/python/CHANGELOG.md"
+#Documentation = "https://beaker-py.readthedocs.io/"
 
 [tool.setuptools]
 include-package-data = true

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/__init__.py

@@ -1,22 +1,39 @@
 from .client import Beaker
 from .types import (
     BeakerAuthRole,
+    BeakerCancelationCode,
+    BeakerCluster,
     BeakerConstraints,
     BeakerDataMount,
+    BeakerDataset,
+    BeakerDatasetFileAlgorithmType,
     BeakerDataSource,
+    BeakerEnvironment,
     BeakerEnvVar,
+    BeakerExperiment,
     BeakerExperimentSpec,
+    BeakerGpuType,
+    BeakerGroup,
+    BeakerImage,
     BeakerImageSource,
+    BeakerJob,
     BeakerJobPriority,
+    BeakerNode,
+    BeakerOrganization,
     BeakerResultSpec,
     BeakerRetrySpec,
+    BeakerSecret,
     BeakerSortOrder,
     BeakerSpecVersion,
+    BeakerTask,
     BeakerTaskContext,
     BeakerTaskResources,
     BeakerTaskSpec,
+    BeakerUser,
+    BeakerWorkload,
     BeakerWorkloadStatus,
     BeakerWorkloadType,
+    BeakerWorkspace,
 )
 
 __all__ = [
@@ -26,6 +43,9 @@ __all__ = [
     "BeakerJobPriority",
     "BeakerWorkloadType",
     "BeakerWorkloadStatus",
+    "BeakerDatasetFileAlgorithmType",
+    "BeakerCancelationCode",
+    "BeakerGpuType",
     "BeakerImageSource",
     "BeakerEnvVar",
     "BeakerDataSource",
@@ -38,4 +58,18 @@ __all__ = [
     "BeakerRetrySpec",
     "BeakerExperimentSpec",
     "BeakerConstraints",
+    "BeakerJob",
+    "BeakerWorkload",
+    "BeakerExperiment",
+    "BeakerTask",
+    "BeakerImage",
+    "BeakerCluster",
+    "BeakerNode",
+    "BeakerDataset",
+    "BeakerGroup",
+    "BeakerSecret",
+    "BeakerOrganization",
+    "BeakerUser",
+    "BeakerEnvironment",
+    "BeakerWorkspace",
 ]

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_cluster.py

@@ -61,7 +61,7 @@ class ClusterClient(ServiceClient):
                 count += 1
                 yield cluster
                 if limit is not None and count >= limit:
-                    break
+                    return
 
     def url(self, cluster: pb2.Cluster) -> str:
         return f"{self.config.agent_address}/orgs/{self.beaker.org_name}/clusters/{cluster.name}"

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_dataset.py

@@ -3,7 +3,7 @@ import os
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import Iterable, Literal
+from typing import Generator, Iterable, Literal
 from urllib.parse import urlparse
 
 import grpc
@@ -55,13 +55,20 @@ class DatasetClient(ServiceClient):
     REQUEST_SIZE_LIMIT = 32 * 1024 * 1024
     DOWNLOAD_CHUNK_SIZE = 10 * 1024
 
-    def get(self, dataset_id: str) -> pb2.Dataset:
+    def get(self, dataset: str) -> pb2.Dataset:
         return self.rpc_request(
             RpcMethod[pb2.GetDatasetResponse](self.service.GetDataset),
-            pb2.GetDatasetRequest(dataset_id=self.resolve_dataset_id(dataset_id)),
-            exceptions_for_status={grpc.StatusCode.NOT_FOUND: BeakerDatasetNotFound(dataset_id)},
+            pb2.GetDatasetRequest(dataset_id=self.resolve_dataset_id(dataset)),
+            exceptions_for_status={grpc.StatusCode.NOT_FOUND: BeakerDatasetNotFound(dataset)},
         ).dataset
 
+    def _get_storage(self, dataset: pb2.Dataset) -> _DatasetStorage:
+        dataset_info = self.http_request(
+            f"datasets/{self._url_quote(dataset.id)}",
+            exceptions_for_status={404: BeakerDatasetNotFound(dataset.id)},
+        ).json()
+        return _DatasetStorage(**dataset_info["storage"])
+
     def create(
         self,
         name: str,
@@ -136,7 +143,7 @@ class DatasetClient(ServiceClient):
             return dataset
 
     def commit(self, dataset: pb2.Dataset) -> pb2.Dataset:
-        if dataset.committed is not None and dataset.committed.ByteSize() > 0:
+        if dataset.HasField("committed"):
             return dataset
 
         @self._retriable()
@@ -164,8 +171,8 @@ class DatasetClient(ServiceClient):
         max_workers: int | None = None,
         strip_paths: bool = False,
     ) -> int:
-        if dataset.committed is not None and dataset.committed.ByteSize() > 0:
-            raise BeakerDatasetWriteError(dataset.id)
+        if dataset.HasField("committed"):
+            raise BeakerDatasetWriteError(f"Dataset '{dataset.id}' has already been committed")
 
         total_bytes = 0
         # map source path to (target_path, size)
@@ -224,6 +231,23 @@ class DatasetClient(ServiceClient):
 
         return total_bytes
 
+    def upload(self, dataset: pb2.Dataset, source: PathOrStr | bytes, target: PathOrStr) -> int:
+        """
+        Upload a file to a dataset.
+
+        :param dataset: The dataset to upload to (must be uncommitted).
+        :param source: Path to the local source file or the contents as bytes.
+        :param target: The path within the dataset to upload the file to.
+
+        :returns: The number of bytes uploaded.
+        """
+        if dataset.HasField("committed"):
+            raise BeakerDatasetWriteError(f"Dataset '{dataset.id}' has already been committed")
+
+        size = len(source) if isinstance(source, bytes) else Path(source).stat().st_size
+        storage = self._get_storage(dataset)
+        return self._upload_file(dataset, storage=storage, size=size, source=source, target=target)
+
     def _upload_file(
         self,
         dataset: pb2.Dataset,
@@ -325,6 +349,91 @@ class DatasetClient(ServiceClient):
         finally:
             source_file.close()
 
+    def stream_file(
+        self,
+        dataset: pb2.Dataset,
+        file_path: str,
+        *,
+        offset: int = 0,
+        length: int = -1,
+        chunk_size: int | None = None,
+        validate_checksum: bool = True,
+    ) -> Generator[bytes, None, None]:
+        prefix = os.path.dirname(file_path)
+        file: pb2.DatasetFile | None = None
+        for f in self.list_files(dataset, prefix=prefix):
+            if f.path == file_path:
+                file = f
+                break
+        else:
+            raise FileNotFoundError(file_path)
+
+        yield from self._stream_file(
+            dataset,
+            file,
+            offset=offset,
+            length=length,
+            chunk_size=chunk_size,
+            validate_checksum=validate_checksum,
+        )
+
+    def _stream_file(
+        self,
+        dataset: pb2.Dataset,
+        file: pb2.DatasetFile,
+        chunk_size: int | None = None,
+        offset: int = 0,
+        length: int = -1,
+        validate_checksum: bool = True,
+    ) -> Generator[bytes, None, None]:
+        def stream_file() -> Generator[bytes, None, None]:
+            headers = {}
+            if offset > 0 and length > 0:
+                headers["Range"] = f"bytes={offset}-{offset + length - 1}"
+            elif offset > 0:
+                headers["Range"] = f"bytes={offset}-"
+            response = self.http_request(
+                f"datasets/{dataset.id}/files/{self._url_quote(file.path)}",
+                method="GET",
+                stream=True,
+                headers=headers,
+                exceptions_for_status={404: FileNotFoundError(file.path)},
+            )
+            for chunk in response.iter_content(chunk_size=chunk_size or self.DOWNLOAD_CHUNK_SIZE):
+                yield chunk
+
+        contents_hash = None
+        if offset == 0 and validate_checksum and file.HasField("digest"):
+            contents_hash = BeakerDatasetFileAlgorithmType(file.digest.algorithm).hasher()
+
+        retries = 0
+        while True:
+            try:
+                for chunk in stream_file():
+                    offset += len(chunk)
+                    if contents_hash is not None:
+                        contents_hash.update(chunk)
+                    yield chunk
+                break
+            except RequestException as err:
+                if retries < self.beaker.MAX_RETRIES:
+                    self._log_and_wait(retries, err)
+                    retries += 1
+                else:
+                    raise
+
+        # Validate digest.
+        if file.HasField("digest") and contents_hash is not None:
+            import binascii
+
+            actual_digest = binascii.hexlify(contents_hash.digest()).decode()
+            expected_digest = binascii.hexlify(file.digest.value).decode()
+            if actual_digest != expected_digest:
+                raise BeakerChecksumFailedError(
+                    f"Checksum for '{file.path}' failed. "
+                    f"Expected '{expected_digest}', got '{actual_digest}'."
+                )
+
     def list_files(
         self, dataset: pb2.Dataset, *, prefix: str | None = None
     ) -> Iterable[pb2.DatasetFile]:
@@ -426,7 +535,7 @@ class DatasetClient(ServiceClient):
                 count += 1
                 yield dataset
                 if limit is not None and count >= limit:
-                    break
+                    return
 
     def url(self, dataset: pb2.Dataset) -> str:
         dataset_id = self.resolve_dataset_id(dataset)

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_experiment.py

@@ -1,4 +1,5 @@
 import grpc
+import yaml
 
 from . import beaker_pb2 as pb2
 from ._service_client import RpcMethod, ServiceClient
@@ -33,19 +34,21 @@ class ExperimentClient(ServiceClient):
 
         return self.beaker.workload.get(data["id"])
 
-    def get_spec(self, experiment: str | pb2.Experiment) -> str:
-        return self.rpc_request(
+    def get_spec(self, experiment: pb2.Experiment | pb2.Workload) -> BeakerExperimentSpec:
+        yaml_str = self.rpc_request(
             RpcMethod[pb2.GetExperimentYamlSpecResponse](self.service.GetExperimentYamlSpec),
             pb2.GetExperimentYamlSpecRequest(experiment_id=self.resolve_experiment_id(experiment)),
             exceptions_for_status={grpc.StatusCode.NOT_FOUND: BeakerExperimentNotFound(experiment)},
         ).experiment_spec
+        json_dict = yaml.safe_load(yaml_str)
+        return BeakerExperimentSpec.from_json(json_dict)
 
-    def restart_tasks(self, experiment: str | pb2.Experiment) -> pb2.Workload:
+    def restart_tasks(self, experiment: pb2.Experiment | pb2.Workload) -> pb2.Workload:
         return self.rpc_request(
             RpcMethod[pb2.RestartExperimentTasksResponse](self.service.RestartExperimentTasks),
             pb2.RestartExperimentTasksRequest(experiment_id=self.resolve_experiment_id(experiment)),
             exceptions_for_status={grpc.StatusCode.NOT_FOUND: BeakerExperimentNotFound(experiment)},
         ).workload
 
-    def url(self, experiment: pb2.Experiment) -> str:
-        return f"{self.config.agent_address}/ex/{self._url_quote(experiment.id)}"
+    def url(self, experiment: pb2.Experiment | pb2.Workload) -> str:
+        return f"{self.config.agent_address}/ex/{self._url_quote(self.resolve_experiment_id(experiment))}"

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_group.py

@@ -106,7 +106,7 @@ class GroupClient(ServiceClient):
                 count += 1
                 yield group
                 if limit is not None and count >= limit:
-                    break
+                    return
 
     def url(self, group: pb2.Group) -> str:
         group_id = self.resolve_group_id(group)

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_image.py

@@ -72,7 +72,7 @@ class ImageClient(ServiceClient):
                 count += 1
                 yield image
                 if limit is not None and count >= limit:
-                    break
+                    return
 
     def url(self, image: pb2.Image) -> str:
         image_id = self.resolve_image_id(image)

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_job.py

@@ -25,16 +25,26 @@ class JobClient(ServiceClient):
         follow: bool | None = None,
         since: datetime | None = None,
     ) -> Iterable[pb2.JobLog]:
-        yield from self.rpc_streaming_request(
+        request = pb2.StreamJobLogsRequest(
+            job_id=job.id,
+            tail_lines=tail_lines,
+            follow=follow,  # type: ignore
+            since=since,  # type: ignore
+        )
+
+        def on_failure():
+            nonlocal request, since
+            request.MergeFrom(pb2.StreamJobLogsRequest(since=since))  # type: ignore[arg-type]
+
+        for job_log in self.rpc_streaming_request(
             RpcStreamingMethod[pb2.JobLog](self.service.StreamJobLogs),
-            pb2.StreamJobLogsRequest(
-                job_id=job.id,
-                tail_lines=tail_lines,
-                follow=follow,  # type: ignore
-                since=since,  # type: ignore
-            ),
+            request,
             exceptions_for_status={grpc.StatusCode.NOT_FOUND: BeakerJobNotFound(job)},
-        )
+            on_failure=on_failure,
+            retriable=True,
+        ):
+            since = job_log.timestamp.ToDatetime()
+            yield job_log
 
     def list_summarized_events(
         self,
@@ -61,6 +71,7 @@ class JobClient(ServiceClient):
         *,
         org: pb2.Organization | None = None,
         task: pb2.Task | None = None,
+        environment: pb2.Environment | None = None,
         finalized: bool | None = None,
         elegible_for_cluster: pb2.Cluster | None = None,
         scheduled_on_node: pb2.Node | None = None,
@@ -85,6 +96,7 @@ class JobClient(ServiceClient):
                     ),
                     organization_id=self.resolve_org_id(org),
                     task_id=None if task is None else task.id,
+                    environment_id=None if environment is None else environment.id,
                     finalized=finalized,  # type: ignore
                     eligible_for_cluster_id=None
                     if elegible_for_cluster is None
@@ -106,7 +118,7 @@ class JobClient(ServiceClient):
                 count += 1
                 yield job
                 if limit is not None and count >= limit:
-                    break
+                    return
 
     def url(self, job: pb2.Job) -> str:
         job_id = job.id

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_node.py

@@ -65,4 +65,4 @@ class NodeClient(ServiceClient):
                 count += 1
                 yield node
                 if limit is not None and count >= limit:
-                    break
+                    return

{beaker_py-2.0.0.dev0 → beaker_py-2.0.1}/src/beaker/_secret.py

@@ -68,4 +68,4 @@ class SecretClient(ServiceClient):
                 count += 1
                 yield secret
                 if limit is not None and count >= limit:
-                    break
+                    return