diff --git a/CHANGELOG.md b/CHANGELOG.md index 1043f64..e7f3690 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead. ## Unreleased +### Added + +- Added `ignore_failures` option to `Beaker.cluster.preempt_jobs()`. + ## [v1.26.1](https://github.com/allenai/beaker-py/releases/tag/v1.26.1) - 2024-02-28 ### Fixed diff --git a/beaker/exceptions.py b/beaker/exceptions.py index 5eca2ca..074b917 100644 --- a/beaker/exceptions.py +++ b/beaker/exceptions.py @@ -31,6 +31,7 @@ "ValidationError", "HTTPError", "RequestException", + "BeakerPermissionsError", "NotFoundError", "AccountNotFound", "OrganizationNotFound", @@ -76,6 +77,12 @@ class BeakerError(Exception): """ +class BeakerPermissionsError(BeakerError): + """ + Raised when a user doesn't have sufficient permissions to perform an action. + """ + + class NotFoundError(BeakerError): """ Base class for all "not found" error types. diff --git a/beaker/services/cluster.py b/beaker/services/cluster.py index 326bed0..3dd1f9a 100644 --- a/beaker/services/cluster.py +++ b/beaker/services/cluster.py @@ -363,11 +363,14 @@ def url(self, cluster: Union[str, Cluster]) -> str: cluster_name = self.resolve_cluster(cluster).full_name return f"{self.config.agent_address}/cl/{cluster_name}/details" - def preempt_jobs(self, cluster: Union[str, Cluster]) -> List[Job]: + def preempt_jobs( + self, cluster: Union[str, Cluster], ignore_failures: bool = False + ) -> List[Job]: """ Preempt all preemptible jobs on the cluster. :param cluster: The cluster ID, full name, or object. + :param ignore_failures: If ``True``, any jobs that fail to preempt will be ignored. :raises ClusterNotFound: If the cluster doesn't exist. :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur. @@ -387,7 +390,15 @@ def preempt_jobs(self, cluster: Union[str, Cluster]) -> List[Job]: continue if job.execution.spec.context.priority != Priority.preemptible: continue - preempted_jobs.append(self.beaker.job.preempt(job)) + try: + preempted_jobs.append(self.beaker.job.preempt(job)) + except BeakerPermissionsError: + if ignore_failures: + self.logger.warning( + "Failed to preempt job '%s': insufficient permissions", job.id + ) + else: + raise return preempted_jobs def _not_found_err_msg(self, cluster: Union[str, Cluster]) -> str: diff --git a/beaker/services/service_client.py b/beaker/services/service_client.py index 5f50c39..ffd3f65 100644 --- a/beaker/services/service_client.py +++ b/beaker/services/service_client.py @@ -125,7 +125,10 @@ def make_request(session: requests.Session) -> requests.Response: and 400 <= response.status_code < 500 ): # Raise a BeakerError if we're misusing the API (4xx error code). - raise BeakerError(msg) + if response.status_code == 403: + raise BeakerPermissionsError(msg) + else: + raise BeakerError(msg) elif msg is not None: raise HTTPError(msg, response=response) # type: ignore else: