diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c163fc..76d0151 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead. ## Unreleased +### Added + +- Added `cursor`, `sort_by` and ordering options to `Beaker.workspace.*` search methods. + ## [v1.12.1](https://github.com/allenai/beaker-py/releases/tag/v1.12.1) - 2022-12-08 ### Fixed diff --git a/beaker/config.py b/beaker/config.py index 0a3005a..8940417 100644 --- a/beaker/config.py +++ b/beaker/config.py @@ -35,7 +35,7 @@ class Config: The address of the Beaker server. """ - default_org: Optional[str] = None + default_org: Optional[str] = "ai2" """ Default Beaker organization to use. """ diff --git a/beaker/data_model/base.py b/beaker/data_model/base.py index 14b044c..6607b40 100644 --- a/beaker/data_model/base.py +++ b/beaker/data_model/base.py @@ -24,7 +24,7 @@ logger = logging.getLogger("beaker") -__all__ = ["BaseModel", "MappedSequence"] +__all__ = ["BaseModel", "MappedSequence", "StrEnum"] BUG_REPORT_URL = ( @@ -144,3 +144,8 @@ def keys(self): def values(self): return self._mapping.values() + + +class StrEnum(str, Enum): + def __str__(self) -> str: + return self.value diff --git a/beaker/data_model/cluster.py b/beaker/data_model/cluster.py index 03453da..c7d6c94 100644 --- a/beaker/data_model/cluster.py +++ b/beaker/data_model/cluster.py @@ -1,16 +1,15 @@ from datetime import datetime -from enum import Enum from typing import Optional, Tuple from pydantic import validator -from .base import BaseModel +from .base import BaseModel, StrEnum from .node import NodeResources, NodeUtilization __all__ = ["ClusterStatus", "Cluster", "ClusterUtilization", "ClusterSpec", "ClusterPatch"] -class ClusterStatus(str, Enum): +class ClusterStatus(StrEnum): """ Current status of a cluster. """ diff --git a/beaker/data_model/dataset.py b/beaker/data_model/dataset.py index 3a42da2..ce8f0c9 100644 --- a/beaker/data_model/dataset.py +++ b/beaker/data_model/dataset.py @@ -4,7 +4,7 @@ from pydantic import validator from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum from .workspace import WorkspaceRef __all__ = [ @@ -18,6 +18,7 @@ "DatasetsPage", "DatasetSpec", "DatasetPatch", + "DatasetSort", ] @@ -168,3 +169,10 @@ class DatasetPatch(BaseModel): name: Optional[str] = None description: Optional[str] = None commit: Optional[bool] = None + + +class DatasetSort(StrEnum): + created = "created" + author = "author" + dataset_name = "name" + dataset_name_or_description = "nameOrDescription" diff --git a/beaker/data_model/experiment.py b/beaker/data_model/experiment.py index a046d2c..90d9342 100644 --- a/beaker/data_model/experiment.py +++ b/beaker/data_model/experiment.py @@ -4,11 +4,11 @@ from pydantic import Field from .account import Account -from .base import BaseModel, MappedSequence +from .base import BaseModel, MappedSequence, StrEnum from .job import Job from .workspace import WorkspaceRef -__all__ = ["Experiment", "Task", "Tasks", "ExperimentsPage", "ExperimentPatch"] +__all__ = ["Experiment", "Task", "Tasks", "ExperimentsPage", "ExperimentPatch", "ExperimentSort"] class Experiment(BaseModel): @@ -70,3 +70,10 @@ class ExperimentsPage(BaseModel): class ExperimentPatch(BaseModel): name: Optional[str] = None description: Optional[str] = None + + +class ExperimentSort(StrEnum): + created = "created" + author = "author" + experiment_name = "name" + experiment_name_or_description = "nameOrDescription" diff --git a/beaker/data_model/experiment_spec.py b/beaker/data_model/experiment_spec.py index f261ee7..288a786 100644 --- a/beaker/data_model/experiment_spec.py +++ b/beaker/data_model/experiment_spec.py @@ -1,11 +1,10 @@ -from enum import Enum from typing import Any, Dict, List, Optional from pydantic import Field, root_validator, validator from ..aliases import PathOrStr from ..exceptions import * -from .base import BaseModel +from .base import BaseModel, StrEnum __all__ = [ "ImageSource", @@ -238,7 +237,7 @@ class TaskResources(BaseModel, frozen=False): """ -class Priority(str, Enum): +class Priority(StrEnum): urgent = "urgent" high = "high" normal = "normal" @@ -305,7 +304,7 @@ class TaskSpec(BaseModel, frozen=False): Context describes how and where this task should run. """ - constraints: Optional[Dict[str, List[str]]] + constraints: Optional[Dict[str, List[str]]] = None """ Each task can have many constraints. And each constraint can have many values. Constraints are rules that change where a task is executed, @@ -598,7 +597,7 @@ def with_constraint(self, **kwargs: List[str]) -> "TaskSpec": ) -class SpecVersion(str, Enum): +class SpecVersion(StrEnum): v2 = "v2" v2_alpha = "v2-alpha" diff --git a/beaker/data_model/group.py b/beaker/data_model/group.py index 246a23f..41ce461 100644 --- a/beaker/data_model/group.py +++ b/beaker/data_model/group.py @@ -1,12 +1,19 @@ from datetime import datetime -from enum import Enum -from typing import List, Optional +from typing import List, Optional, Tuple from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum from .workspace import WorkspaceRef -__all__ = ["Group", "GroupSpec", "GroupParameterType", "GroupParameter", "GroupPatch"] +__all__ = [ + "Group", + "GroupSpec", + "GroupParameterType", + "GroupParameter", + "GroupPatch", + "GroupsPage", + "GroupSort", +] class Group(BaseModel): @@ -31,7 +38,7 @@ class GroupSpec(BaseModel): experiments: Optional[List[str]] = None -class GroupParameterType(str, Enum): +class GroupParameterType(StrEnum): metric = "metric" env = "env" @@ -47,3 +54,17 @@ class GroupPatch(BaseModel): add_experiments: Optional[List[str]] = None remove_experiments: Optional[List[str]] = None parameters: Optional[List[GroupParameter]] = None + + +class GroupsPage(BaseModel): + data: Tuple[Group, ...] + next_cursor: Optional[str] = None + next: Optional[str] = None + + +class GroupSort(StrEnum): + created = "created" + modified = "modified" + author = "author" + group_name = "name" + group_name_or_description = "nameOrDescription" diff --git a/beaker/data_model/image.py b/beaker/data_model/image.py index 6d88ba2..0741309 100644 --- a/beaker/data_model/image.py +++ b/beaker/data_model/image.py @@ -1,11 +1,10 @@ from datetime import datetime -from enum import Enum from typing import Optional, Tuple from pydantic import validator from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum from .workspace import WorkspaceRef __all__ = [ @@ -20,6 +19,7 @@ "DockerLayerDownloadState", "ImageSpec", "ImagePatch", + "ImageSort", ] @@ -73,7 +73,7 @@ class DockerLayerProgress(BaseModel): total: Optional[int] = None -class DockerLayerUploadStatus(str, Enum): +class DockerLayerUploadStatus(StrEnum): preparing = "preparing" waiting = "waiting" pushing = "pushing" @@ -81,7 +81,7 @@ class DockerLayerUploadStatus(str, Enum): already_exists = "layer already exists" -class DockerLayerDownloadStatus(str, Enum): +class DockerLayerDownloadStatus(StrEnum): waiting = "waiting" downloading = "downloading" download_complete = "download complete" @@ -122,3 +122,9 @@ class ImagePatch(BaseModel): name: Optional[str] = None description: Optional[str] = None commit: Optional[bool] = None + + +class ImageSort(StrEnum): + created = "created" + author = "author" + image_name = "name" diff --git a/beaker/data_model/job.py b/beaker/data_model/job.py index 3a5fe42..bcb7bb6 100644 --- a/beaker/data_model/job.py +++ b/beaker/data_model/job.py @@ -5,7 +5,7 @@ from pydantic import Field, validator from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum from .experiment_spec import DataMount, EnvVar, ImageSource, Priority, TaskSpec __all__ = [ @@ -25,7 +25,7 @@ ] -class CurrentJobStatus(str, Enum): +class CurrentJobStatus(StrEnum): """ The status of a job. """ @@ -41,7 +41,7 @@ class CurrentJobStatus(str, Enum): preempted = "preempted" -class CanceledCode(int, Enum): +class CanceledCode(Enum): not_set = 0 system_preemption = 1 user_preemption = 2 @@ -125,7 +125,7 @@ class JobExecution(BaseModel): workspace: Optional[str] = None -class JobKind(str, Enum): +class JobKind(StrEnum): """ The kind of job. """ diff --git a/beaker/data_model/organization.py b/beaker/data_model/organization.py index 304213d..c249ba3 100644 --- a/beaker/data_model/organization.py +++ b/beaker/data_model/organization.py @@ -1,9 +1,8 @@ from datetime import datetime -from enum import Enum from typing import Optional from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum __all__ = ["Organization", "OrganizationRole", "OrganizationMember"] @@ -17,7 +16,7 @@ class Organization(BaseModel): pronouns: Optional[str] = None -class OrganizationRole(str, Enum): +class OrganizationRole(StrEnum): admin = "admin" member = "member" diff --git a/beaker/data_model/workspace.py b/beaker/data_model/workspace.py index 59fc363..50f73fa 100644 --- a/beaker/data_model/workspace.py +++ b/beaker/data_model/workspace.py @@ -1,9 +1,8 @@ from datetime import datetime -from enum import Enum from typing import Dict, List, Optional, Tuple from .account import Account -from .base import BaseModel +from .base import BaseModel, StrEnum __all__ = [ "WorkspaceSize", @@ -17,6 +16,7 @@ "WorkspacePatch", "WorkspacePermissionsPatch", "WorkspaceClearResult", + "WorkspaceSort", ] @@ -64,7 +64,7 @@ class WorkspaceTransferSpec(BaseModel): ids: List[str] -class Permission(str, Enum): +class Permission(StrEnum): """ Workspace permission levels. """ @@ -101,3 +101,9 @@ class WorkspaceClearResult(BaseModel): images_deleted: int = 0 datasets_deleted: int = 0 secrets_deleted: int = 0 + + +class WorkspaceSort(StrEnum): + created = "created" + modified = "modified" + workspace_name = "name" diff --git a/beaker/exceptions.py b/beaker/exceptions.py index 2be6ad1..e07e1df 100644 --- a/beaker/exceptions.py +++ b/beaker/exceptions.py @@ -55,6 +55,7 @@ "TaskStoppedError", "JobFailedError", "JobTimeoutError", + "ExperimentSpecError", ] diff --git a/beaker/services/__init__.py b/beaker/services/__init__.py index facb900..8e345ad 100644 --- a/beaker/services/__init__.py +++ b/beaker/services/__init__.py @@ -10,3 +10,18 @@ from .secret import SecretClient from .service_client import ServiceClient from .workspace import WorkspaceClient + +__all__ = [ + "AccountClient", + "ClusterClient", + "DatasetClient", + "ExperimentClient", + "GroupClient", + "ImageClient", + "JobClient", + "NodeClient", + "OrganizationClient", + "SecretClient", + "ServiceClient", + "WorkspaceClient", +] diff --git a/beaker/services/dataset.py b/beaker/services/dataset.py index 0cddc70..7858d2d 100644 --- a/beaker/services/dataset.py +++ b/beaker/services/dataset.py @@ -2,7 +2,16 @@ import os from datetime import datetime from pathlib import Path -from typing import TYPE_CHECKING, Deque, Dict, Generator, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + ClassVar, + Deque, + Dict, + Generator, + Optional, + Tuple, + Union, +) from ..aliases import PathOrStr from ..data_model import * @@ -27,7 +36,7 @@ class DatasetClient(ServiceClient): HEADER_LAST_MODIFIED = "Last-Modified" HEADER_CONTENT_LENGTH = "Content-Length" - REQUEST_SIZE_LIMIT = 32 * 1024 * 1024 + REQUEST_SIZE_LIMIT: ClassVar[int] = 32 * 1024 * 1024 def get(self, dataset: str) -> Dataset: """ diff --git a/beaker/services/workspace.py b/beaker/services/workspace.py index 5998e7a..d8c2838 100644 --- a/beaker/services/workspace.py +++ b/beaker/services/workspace.py @@ -4,6 +4,7 @@ from ..data_model import * from ..exceptions import * +from ..util import format_cursor from .service_client import ServiceClient @@ -52,7 +53,7 @@ def _get(id: str) -> Workspace: raise def create( - self, workspace: str, description: Optional[str] = None, public: bool = False + self, workspace: str, *, description: Optional[str] = None, public: bool = False ) -> Workspace: """ Create a workspace. @@ -210,36 +211,25 @@ def move( }, ) - def list( + def iter( self, org: Optional[Union[str, Organization]] = None, + *, author: Optional[Union[str, Account]] = None, match: Optional[str] = None, archived: Optional[bool] = None, limit: Optional[int] = None, - ) -> List[Workspace]: - """ - List workspaces belonging to an organization. - - :param org: The organization name or object. If not specified, - :data:`Beaker.config.default_org ` is used. - :param author: Only list workspaces authored by this account. - :param match: Only include workspaces matching the text. - :param archived: Only include/exclude archived workspaces. - :param limit: Limit the number of workspaces returned. - - :raises OrganizationNotFound: If the organization doesn't exist. - :raises OrganizationNotSet: If neither ``org`` nor - :data:`Beaker.config.default_org ` are set. - :raises AccountNotFound: If the author account doesn't exist. - :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur. - :raises RequestException: Any other exception that can occur when contacting the - Beaker server. - """ + sort_by: WorkspaceSort = WorkspaceSort.created, + descending: bool = True, + cursor: int = 0, + ) -> Generator[Workspace, None, None]: org = self.resolve_org(org) - workspaces: List[Workspace] = [] - cursor: Optional[str] = None - query: Dict[str, str] = {"org": org.id} + query: Dict[str, str] = { + "org": org.id, + "field": str(sort_by), + "order": "descending" if descending else "ascending", + "cursor": format_cursor(cursor), + } if author is not None: query["author"] = ( author.name if isinstance(author, Account) else self.beaker.account.get(author).name @@ -251,8 +241,8 @@ def list( if limit: query["limit"] = str(limit) + count = 0 while True: - query["cursor"] = cursor or "" page = WorkspacePage.from_json( self.request( "workspaces", @@ -260,21 +250,71 @@ def list( query=query, ).json() ) - workspaces.extend(page.data) - cursor = page.next_cursor or page.next - if not cursor: - break - if limit and len(workspaces) >= limit: - workspaces = workspaces[:limit] + for workspace in page.data: + count += 1 + yield workspace + if limit is not None and count >= limit: + return + + query["cursor"] = page.next_cursor or page.next # type: ignore + if not query["cursor"]: break - return workspaces + def list( + self, + org: Optional[Union[str, Organization]] = None, + *, + author: Optional[Union[str, Account]] = None, + match: Optional[str] = None, + archived: Optional[bool] = None, + limit: Optional[int] = None, + sort_by: WorkspaceSort = WorkspaceSort.created, + descending: bool = True, + cursor: int = 0, + ) -> List[Workspace]: + """ + List workspaces belonging to an organization. + + :param org: The organization name or object. If not specified, + :data:`Beaker.config.default_org ` is used. + :param author: Only list workspaces authored by this account. + :param match: Only include workspaces matching the text. + :param archived: Only include/exclude archived workspaces. + :param limit: Limit the number of workspaces returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. + + :raises OrganizationNotFound: If the organization doesn't exist. + :raises OrganizationNotSet: If neither ``org`` nor + :data:`Beaker.config.default_org ` are set. + :raises AccountNotFound: If the author account doesn't exist. + :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur. + :raises RequestException: Any other exception that can occur when contacting the + Beaker server. + """ + return list( + self.iter( + org=org, + author=author, + match=match, + archived=archived, + limit=limit, + sort_by=sort_by, + descending=descending, + cursor=cursor, + ) + ) def iter_images( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, limit: Optional[int] = None, + sort_by: ImageSort = ImageSort.created, + descending: bool = True, + cursor: int = 0, ) -> Generator[Image, None, None]: """ Iterate over the images in a workspace. @@ -283,6 +323,9 @@ def iter_images( :data:`Beaker.config.default_workspace ` is used. :param match: Only include images matching the text. :param limit: Limit the number of images returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -292,14 +335,18 @@ def iter_images( Beaker server. """ workspace_name = self.resolve_workspace(workspace, read_only_ok=True).full_name - cursor: Optional[str] = None - query: Dict[str, str] = {} + query: Dict[str, str] = { + "field": str(sort_by), + "order": "descending" if descending else "ascending", + "cursor": format_cursor(cursor), + } if match is not None: query["q"] = match + if limit: + query["limit"] = str(limit) count = 0 while True: - query["cursor"] = cursor or "" page = ImagesPage.from_json( self.request( f"workspaces/{self.url_quote(workspace_name)}/images", @@ -317,15 +364,19 @@ def iter_images( if limit is not None and count >= limit: return - cursor = page.next_cursor or page.next - if not cursor: + query["cursor"] = page.next_cursor or page.next # type: ignore + if not query["cursor"]: break def images( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, limit: Optional[int] = None, + sort_by: ImageSort = ImageSort.created, + descending: bool = True, + cursor: int = 0, ) -> List[Image]: """ List the images in a workspace. @@ -334,6 +385,9 @@ def images( :data:`Beaker.config.default_workspace ` is used. :param match: Only include images matching the text. :param limit: Limit the number of images returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -342,13 +396,26 @@ def images( :raises RequestException: Any other exception that can occur when contacting the Beaker server. """ - return list(self.iter_images(workspace=workspace, match=match, limit=limit)) + return list( + self.iter_images( + workspace=workspace, + match=match, + limit=limit, + sort_by=sort_by, + descending=descending, + cursor=cursor, + ) + ) def iter_experiments( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, limit: Optional[int] = None, + sort_by: ExperimentSort = ExperimentSort.created, + descending: bool = True, + cursor: int = 0, ) -> Generator[Experiment, None, None]: """ Iterate over the experiments in a workspace. @@ -357,6 +424,9 @@ def iter_experiments( :data:`Beaker.config.default_workspace ` is used. :param match: Only include experiments matching the text. :param limit: Limit the number of experiments returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -366,14 +436,18 @@ def iter_experiments( Beaker server. """ workspace_name = self.resolve_workspace(workspace, read_only_ok=True).full_name - cursor: Optional[str] = None - query: Dict[str, str] = {} + query: Dict[str, str] = { + "field": str(sort_by), + "order": "descending" if descending else "ascending", + "cursor": format_cursor(cursor), + } if match is not None: query["q"] = match + if limit: + query["limit"] = str(limit) count = 0 while True: - query["cursor"] = cursor or "" page = ExperimentsPage.from_json( self.request( f"workspaces/{self.url_quote(workspace_name)}/experiments", @@ -391,15 +465,19 @@ def iter_experiments( if limit is not None and count >= limit: return - cursor = page.next_cursor or page.next - if not cursor: + query["cursor"] = page.next_cursor or page.next # type: ignore + if not query["cursor"]: break def experiments( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, limit: Optional[int] = None, + sort_by: ExperimentSort = ExperimentSort.created, + descending: bool = True, + cursor: int = 0, ) -> List[Experiment]: """ List the experiments in a workspace. @@ -408,6 +486,9 @@ def experiments( :data:`Beaker.config.default_workspace ` is used. :param match: Only include experiments matching the text. :param limit: Limit the number of experiments returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -416,15 +497,28 @@ def experiments( :raises RequestException: Any other exception that can occur when contacting the Beaker server. """ - return list(self.iter_experiments(workspace=workspace, match=match, limit=limit)) + return list( + self.iter_experiments( + workspace=workspace, + match=match, + limit=limit, + sort_by=sort_by, + descending=descending, + cursor=cursor, + ) + ) def iter_datasets( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, results: Optional[bool] = None, uncommitted: Optional[bool] = None, limit: Optional[int] = None, + sort_by: DatasetSort = DatasetSort.created, + descending: bool = True, + cursor: int = 0, ) -> Generator[Dataset, None, None]: """ Iterate over the datasets in a workspace. @@ -435,6 +529,9 @@ def iter_datasets( :param results: Only include/exclude experiment result datasets. :param uncommitted: Only include/exclude uncommitted datasets. :param limit: Limit the number of datasets returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -444,18 +541,22 @@ def iter_datasets( Beaker server. """ workspace_name = self.resolve_workspace(workspace, read_only_ok=True).full_name - cursor: Optional[str] = None - query: Dict[str, str] = {} + query: Dict[str, str] = { + "field": str(sort_by), + "order": "descending" if descending else "ascending", + "cursor": format_cursor(cursor), + } if match is not None: query["q"] = match if results is not None: query["results"] = str(results).lower() if uncommitted is not None: query["committed"] = str(not uncommitted).lower() + if limit: + query["limit"] = str(limit) count = 0 while True: - query["cursor"] = cursor or "" page = DatasetsPage.from_json( self.request( f"workspaces/{self.url_quote(workspace_name)}/datasets", @@ -473,17 +574,21 @@ def iter_datasets( if limit is not None and count >= limit: return - cursor = page.next_cursor or page.next - if not cursor: + query["cursor"] = page.next_cursor or page.next # type: ignore + if not query["cursor"]: break def datasets( self, workspace: Optional[Union[str, Workspace]] = None, + *, match: Optional[str] = None, results: Optional[bool] = None, uncommitted: Optional[bool] = None, limit: Optional[int] = None, + sort_by: DatasetSort = DatasetSort.created, + descending: bool = True, + cursor: int = 0, ) -> List[Dataset]: """ List the datasets in a workspace. @@ -494,6 +599,9 @@ def datasets( :param results: Only include/exclude experiment result datasets. :param uncommitted: Only include/exclude uncommitted datasets. :param limit: Limit the number of datasets returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -509,6 +617,9 @@ def datasets( results=results, uncommitted=uncommitted, limit=limit, + sort_by=sort_by, + descending=descending, + cursor=cursor, ) ) @@ -538,12 +649,26 @@ def secrets(self, workspace: Optional[Union[str, Workspace]] = None) -> List[Sec ).json()["data"] ] - def groups(self, workspace: Optional[Union[str, Workspace]] = None) -> List[Group]: + def iter_groups( + self, + workspace: Optional[Union[str, Workspace]] = None, + *, + match: Optional[str] = None, + limit: Optional[int] = None, + sort_by: GroupSort = GroupSort.created, + descending: bool = True, + cursor: int = 0, + ) -> Generator[Group, None, None]: """ - List groups in a workspace. + Iterate over groups in a workspace. :param workspace: The Beaker workspace name, or object. If not specified, :data:`Beaker.config.default_workspace ` is used. + :param match: Only include groups matching the text. + :param limit: Limit the number of groups returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. :raises WorkspaceNotFound: If the workspace doesn't exist. :raises WorkspaceNotSet: If neither ``workspace`` nor @@ -553,16 +678,77 @@ def groups(self, workspace: Optional[Union[str, Workspace]] = None) -> List[Grou Beaker server. """ workspace_name = self.resolve_workspace(workspace, read_only_ok=True).full_name - return [ - Group.from_json(d) - for d in self.request( - f"workspaces/{self.url_quote(workspace_name)}/groups", - method="GET", - exceptions_for_status={ - 404: WorkspaceNotFound(self._not_found_err_msg(workspace_name)) - }, - ).json()["data"] - ] + query: Dict[str, str] = { + "field": str(sort_by), + "order": "descending" if descending else "ascending", + "cursor": format_cursor(cursor), + } + if match is not None: + query["q"] = match + if limit: + query["limit"] = str(limit) + + count = 0 + while True: + page = GroupsPage.from_json( + self.request( + f"workspaces/{self.url_quote(workspace_name)}/groups", + method="GET", + query=query, + exceptions_for_status={ + 404: WorkspaceNotFound(self._not_found_err_msg(workspace_name)) + }, + ).json() + ) + + for group in page.data: + count += 1 + yield group + if limit is not None and count >= limit: + return + + query["cursor"] = page.next_cursor or page.next # type: ignore + if not query["cursor"]: + break + + def groups( + self, + workspace: Optional[Union[str, Workspace]] = None, + *, + match: Optional[str] = None, + limit: Optional[int] = None, + sort_by: GroupSort = GroupSort.created, + descending: bool = True, + cursor: int = 0, + ) -> List[Group]: + """ + List groups in a workspace. + + :param workspace: The Beaker workspace name, or object. If not specified, + :data:`Beaker.config.default_workspace ` is used. + :param match: Only include groups matching the text. + :param limit: Limit the number of groups returned. + :param sort_by: The field to sort the results by. + :param descending: Order the results in descending order according to the ``sort_by`` field. + :param cursor: Set the starting cursor for the query. You can use this to paginate the results. + + :raises WorkspaceNotFound: If the workspace doesn't exist. + :raises WorkspaceNotSet: If neither ``workspace`` nor + :data:`Beaker.config.default_workspace ` are set. + :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur. + :raises RequestException: Any other exception that can occur when contacting the + Beaker server. + """ + return list( + self.iter_groups( + workspace=workspace, + match=match, + limit=limit, + sort_by=sort_by, + descending=descending, + cursor=cursor, + ) + ) def get_permissions( self, workspace: Optional[Union[str, Workspace]] = None @@ -706,6 +892,7 @@ def url(self, workspace: Optional[Union[str, Workspace]] = None) -> str: def clear( self, workspace: Optional[Union[str, Workspace]] = None, + *, groups: bool = True, experiments: bool = True, images: bool = True, diff --git a/beaker/util.py b/beaker/util.py index 984abd9..aeb6baf 100644 --- a/beaker/util.py +++ b/beaker/util.py @@ -1,3 +1,4 @@ +import base64 import re import time from collections import OrderedDict @@ -139,3 +140,10 @@ def retriable_method(*args, **kwargs) -> T: return retriable_method return parametrize_decorator + + +def format_cursor(cursor: int) -> str: + if cursor < 0: + raise ValueError("cursor must be >= 0") + + return base64.urlsafe_b64encode(cursor.to_bytes(8, "little")).decode() diff --git a/conftest.py b/conftest.py index f41b5bb..086f2d1 100644 --- a/conftest.py +++ b/conftest.py @@ -14,7 +14,7 @@ def unique_name() -> str: - return petname.generate() + "-" + str(uuid.uuid4())[:8] + return petname.generate() + "-" + str(uuid.uuid4())[:8] # type: ignore def beaker_object_fixture(client: Beaker, service: str): diff --git a/examples/sweep/run.py b/examples/sweep/run.py index 65d5ccf..682f175 100644 --- a/examples/sweep/run.py +++ b/examples/sweep/run.py @@ -8,7 +8,7 @@ def unique_name() -> str: - return petname.generate() + "-" + str(uuid.uuid4())[:8] + return petname.generate() + "-" + str(uuid.uuid4())[:8] # type: ignore def main(image: str, workspace: str, cluster: str): diff --git a/tests/data_model_test.py b/tests/data_model_test.py index f296359..aa1193c 100644 --- a/tests/data_model_test.py +++ b/tests/data_model_test.py @@ -14,7 +14,7 @@ def test_data_source_validation(): DataSource(beaker="foo", host_path="bar") with pytest.raises(ValidationError, match="Exactly one"): - DataSource(beaker="foo", hostPath="bar") + DataSource(beaker="foo", hostPath="bar") # type: ignore assert DataSource(host_path="bar").host_path == "bar" @@ -66,15 +66,15 @@ def test_experiment_spec_validation(): tasks=[ TaskSpec( name="main", - image={"docker": "hello-world"}, - context={"cluster": "foo"}, - result={"path": "/unused"}, + image={"docker": "hello-world"}, # type: ignore + context={"cluster": "foo"}, # type: ignore + result={"path": "/unused"}, # type: ignore ), TaskSpec( name="main", - image={"docker": "hello-world"}, - context={"cluster": "bar"}, - result={"path": "/unused"}, + image={"docker": "hello-world"}, # type: ignore + context={"cluster": "bar"}, # type: ignore + result={"path": "/unused"}, # type: ignore ), ] ) diff --git a/tests/experiment_test.py b/tests/experiment_test.py index 0c08b55..a102c54 100644 --- a/tests/experiment_test.py +++ b/tests/experiment_test.py @@ -24,8 +24,10 @@ def test_experiment_get(client: Beaker, hello_world_experiment_id: str): assert exp.jobs assert exp.jobs[0].status.current == CurrentJobStatus.finalized # Get with name. + assert exp.name is not None client.experiment.get(exp.name) # Get with full name. + assert exp.full_name is not None client.experiment.get(exp.full_name) diff --git a/tests/image_test.py b/tests/image_test.py index 5f5e535..89408c8 100644 --- a/tests/image_test.py +++ b/tests/image_test.py @@ -7,6 +7,7 @@ def test_image_get(client: Beaker, hello_world_image_name: str): # Get by ID. client.image.get(image.id) # Get by name. + assert image.name is not None client.image.get(image.name) diff --git a/tests/node_test.py b/tests/node_test.py index 696718d..f6e8c76 100644 --- a/tests/node_test.py +++ b/tests/node_test.py @@ -2,4 +2,6 @@ def test_node_get(client: Beaker, beaker_node_id: str): - assert client.node.get(beaker_node_id).limits.gpu_count > 0 + gpu_count = client.node.get(beaker_node_id).limits.gpu_count + assert gpu_count is not None + assert gpu_count > 0 diff --git a/tests/util_test.py b/tests/util_test.py index 48f3e55..5359077 100644 --- a/tests/util_test.py +++ b/tests/util_test.py @@ -1,3 +1,4 @@ +import base64 import time import pytest @@ -41,3 +42,9 @@ def x(self) -> int: client.config.default_workspace = alternate_workspace_name assert service_client.x == 3 + + +def test_format_cursor(): + cursor = 100 + formatted = format_cursor(100) + assert int.from_bytes(base64.urlsafe_b64decode(formatted), "little") == cursor