Skip to content

Commit

Permalink
Add author filter to Beaker.job.list()
Browse files Browse the repository at this point in the history
And add `Beaker.experiment.list()` method.
  • Loading branch information
epwalsh committed Jun 13, 2024
1 parent 7b2d35a commit b71e9a9
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 6 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@ use patch releases for compatibility fixes instead.

## Unreleased

### Added

- Added `author` filter to `Beaker.job.list()`.
- Added `Beaker.experiment.list()` method.

## [v1.28.0](https://github.com/allenai/beaker-py/releases/tag/v1.28.0) - 2024-06-13

### Added
Expand Down
39 changes: 39 additions & 0 deletions beaker/services/experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,45 @@ def _parse_create_args(
assert spec is not None
return spec, name, workspace

def list(
self,
*,
author: Optional[Union[str, Account]] = None,
cluster: Optional[Union[str, Cluster]] = None,
finalized: bool = False,
node: Optional[Union[str, Node]] = None,
) -> List[Experiment]:
"""
List experiments.
:param author: List only experiments by particular author.
:param cluster: List experiments on a particular cluster.
:param finalized: List only finalized or non-finalized experiments.
:param node: List experiments on a particular node.
.. important::
Either ``cluster``, ``author``, ``experiment``, or ``node`` must be specified.
If ``node`` is specified, neither ``cluster`` nor ``experiment`` can be
specified.
:raises ValueError: If the arguments are invalid, e.g. both ``node`` and
``cluster`` are specified.
:raises AccountNotFound: If the specified author doesn't exist.
:raises ClusterNotFound: If the specified cluster doesn't exist.
:raises NodeNotFound: If the specified node doesn't exist.
"""
jobs = self.beaker.job.list(
author=author, cluster=cluster, finalized=finalized, kind=JobKind.execution, node=node
)
experiments: List[Experiment] = []
exp_ids: Set[str] = set()
for job in jobs:
assert job.execution is not None
if (exp_id := job.execution.experiment) not in exp_ids:
exp_ids.add(exp_id)
experiments.append(self.get(exp_id))
return experiments

def create(
self,
*args,
Expand Down
21 changes: 15 additions & 6 deletions beaker/services/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ def get(self, job_id: str) -> Job:

def list(
self,
*,
author: Optional[Union[str, Account]] = None,
cluster: Optional[Union[str, Cluster]] = None,
experiment: Optional[Union[str, Experiment]] = None,
finalized: bool = False,
Expand All @@ -45,37 +47,44 @@ def list(
"""
List jobs.
:param cluster: List jobs on a cluster.
:param author: List only jobs by particular author.
:param cluster: List jobs on a particular cluster.
:param experiment: List jobs in an experiment.
:param finalized: List only finalized or non-finalized jobs.
:param kind: List jobs of a certain kind.
:param node: List jobs on a node.
:param node: List jobs on a particular node.
.. important::
Either ``cluster``, ``experiment``, or ``node`` must be specified.
Either ``cluster``, ``author``, ``experiment``, or ``node`` must be specified.
If ``node`` is specified, neither ``cluster`` nor ``experiment`` can be
specified.
:raises ValueError: If the arguments are invalid, e.g. both ``node`` and
``cluster`` are specified.
:raises AccountNotFound: If the specified author doesn't exist.
:raises ClusterNotFound: If the specified cluster doesn't exist.
:raises ExperimentNotFound: If the specified experiment doesn't exist.
:raises NodeNotFound: If the specified node doesn't exist.
"""
if node is None and cluster is None and experiment is None and author is None:
raise ValueError("You must specify one of 'node', 'cluster', 'experiment', or 'author'")

# Validate arguments.
if node is not None:
if cluster is not None:
raise ValueError("You cannot specify both 'node' and 'cluster'")
if experiment is not None:
raise ValueError("You cannot specify both 'node' and 'experiment'")
else:
if cluster is None and experiment is None:
raise ValueError("You must specify one of 'node', 'experiment', or 'cluster'")

jobs: List[Job] = []

# Build request options.
request_opts: Dict[str, Any] = {}
if author is not None:
author_id = (
author.id if isinstance(author, Account) else self.beaker.account.get(author).id
)
request_opts["author"] = author_id
if cluster is not None:
cluster_id = (
cluster.id if isinstance(cluster, Cluster) else self.beaker.cluster.get(cluster).id
Expand Down

0 comments on commit b71e9a9

Please sign in to comment.