Merge pull request #13 from allenai/api-improvements

API improvements

.flake8

@@ -7,6 +7,10 @@ ignore =
     E203
     # line break before binary operator
     W503
+    # star import
+    F403
+    # import may be defined in a star import
+    F405
 
 exclude =
     .venv

.github/workflows/ci.yml

@@ -167,6 +167,6 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:
           body_path: ${{ github.workspace }}-RELEASE_NOTES.md
-          prerelease: ${{ contains(env.TAG, '-rc') }}
+          prerelease: ${{ contains(env.TAG, 'rc') }}
           files: |
             dist/*

CHANGELOG.md

@@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added many new methods.
+- Added `beaker.exceptions` module.
+
+### Changed
+
+- Several breaking changes to `Beaker` method names.
+
 ## [v0.1.0](https://github.com/allenai/beaker-py/releases/tag/v0.1.0) - 2021-11-19
 
 ### Added

README.md

@@ -41,15 +41,14 @@ pip install -e .
 Create a Beaker client with your Beaker [user token](https://beaker.org/user):
 
 ```python
-from beaker import Beaker
+from beaker import Beaker, Config
 
-beaker = Beaker("my beaker token", workspace="my_org/my_workspace")
+beaker = Beaker(Config(user_token="my beaker token", workspace="my_org/my_workspace"))
 ```
 
-You can also create your client from environment variables with:
+You can also create your client from a beaker config file or environment variables with:
 
 ```python
-# Assumes your user token is set as the environment variable `BEAKER_TOKEN`.
 beaker = Beaker.from_env()
 ```
 

beaker/__init__.py

@@ -1,3 +1,3 @@
-__all__ = ["Beaker"]
-
 from .client import Beaker
+from .config import Config
+from .exceptions import *

beaker/client.py

@@ -10,6 +10,11 @@
 from requests.packages.urllib3.util.retry import Retry
 from tqdm import tqdm
 
+from .config import Config
+from .exceptions import *
+
+__all__ = ["Beaker"]
+
 
 class Beaker:
     """
@@ -20,22 +25,25 @@ class Beaker:
     MAX_RETRIES = 5
     API_VERSION = "v3"
 
-    def __init__(self, token: str, workspace: Optional[str] = None):
-        self.base_url = f"https://beaker.org/api/{self.API_VERSION}"
-        self.token = token
+    def __init__(self, config: Config):
+        self.config = config
+        self.base_url = f"{self.config.agent_address}/api/{self.API_VERSION}"
         self.docker = docker.from_env()
-        self.workspace = workspace
+
+    @property
+    def user(self) -> str:
+        """
+        The username associated with this account.
+        """
+        return self.whoami()["name"]
 
     @classmethod
     def from_env(cls, **kwargs) -> "Beaker":
         """
         Initialize client from environment variables. Expects the beaker auth token
         to be set as the ``BEAKER_TOKEN`` environment variable.
         """
-        import os
-
-        token = os.environ["BEAKER_TOKEN"]
-        return cls(token, **kwargs)
+        return cls(Config.from_env())
 
     @contextmanager
     def _session_with_backoff(self) -> requests.Session:
@@ -54,6 +62,7 @@ def request(
         method: str = "GET",
         query: Optional[Dict[str, Any]] = None,
         data: Optional[Dict[str, Any]] = None,
+        exceptions_for_status: Optional[Dict[int, BeakerError]] = None,
     ) -> requests.Response:
         with self._session_with_backoff() as session:
             url = f"{self.base_url}/{resource}"
@@ -62,11 +71,13 @@ def request(
             response = getattr(session, method.lower())(
                 url,
                 headers={
-                    "Authorization": f"Bearer {self.token}",
+                    "Authorization": f"Bearer {self.config.user_token}",
                     "Content-Type": "application/json",
                 },
                 data=None if data is None else json.dumps(data),
             )
+            if exceptions_for_status is not None and response.status_code in exceptions_for_status:
+                raise exceptions_for_status[response.status_code]
             response.raise_for_status()
             return response
 
@@ -76,23 +87,88 @@ def whoami(self) -> Dict[str, Any]:
         """
         return self.request("user").json()
 
-    def experiment(self, exp_id: str) -> Dict[str, Any]:
+    def get_workspace(self, workspace: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Get information about the workspace.
+
+        Raises
+        ------
+        WorkspaceNotFound
+        HTTPError
+
+        """
+        workspace_name = workspace or self.config.default_workspace
+        if workspace_name is None:
+            raise ValueError("'workspace' argument required")
+        return self.request(
+            f"workspaces/{urllib.parse.quote(workspace_name, safe='')}",
+            exceptions_for_status={404: WorkspaceNotFound(workspace_name)},
+        ).json()
+
+    def get_experiment(self, exp_id: str) -> Dict[str, Any]:
         """
         Get info about an experiment.
+
+        Raises
+        ------
+        ExperimentNotFound
+        HTTPError
+
         """
-        return self.request(f"experiments/{exp_id}").json()
+        return self.request(
+            f"experiments/{exp_id}", exceptions_for_status={404: ExperimentNotFound(exp_id)}
+        ).json()
+
+    def create_experiment(
+        self, name: str, spec: Dict[str, Any], workspace: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Create a new Beaker experiment with the given ``spec``.
+
+        Raises
+        ------
+        ExperimentConflict
+        HTTPError
+
+        """
+        workspace_name = workspace or self.config.default_workspace
+        if workspace_name is None:
+            raise ValueError("'workspace' argument required")
+        return self.request(
+            f"workspaces/{urllib.parse.quote(workspace_name, safe='')}/experiments",
+            method="POST",
+            query={"name": name},
+            data=spec,
+            exceptions_for_status={409: ExperimentConflict(name)},
+        ).json()
 
-    def dataset(self, dataset_id: str) -> Dict[str, Any]:
+    def get_dataset(self, dataset_id: str) -> Dict[str, Any]:
         """
         Get info about a dataset.
+
+        Raises
+        ------
+        DatasetNotFound
+        HTTPError
+
         """
-        return self.request(f"datasets/{dataset_id}").json()
+        return self.request(
+            f"datasets/{dataset_id}", exceptions_for_status={404: DatasetNotFound(dataset_id)}
+        ).json()
 
-    def logs(self, job_id: str) -> Generator[bytes, None, None]:
+    def get_logs(self, job_id: str) -> Generator[bytes, None, None]:
         """
         Download the logs for a job.
+
+        Raises
+        ------
+        JobNotFound
+        HTTPError
+
         """
-        response = self.request(f"jobs/{job_id}/logs")
+        response = self.request(
+            f"jobs/{job_id}/logs", exceptions_for_status={404: JobNotFound(job_id)}
+        )
         content_length = response.headers.get("Content-Length")
         total = int(content_length) if content_length is not None else None
         progress = tqdm(
@@ -103,28 +179,54 @@ def logs(self, job_id: str) -> Generator[bytes, None, None]:
                 progress.update(len(chunk))
                 yield chunk
 
-    def logs_for_experiment(
+    def get_logs_for_experiment(
         self, exp_id: str, job_id: Optional[str] = None
     ) -> Generator[bytes, None, None]:
         """
         Download the logs for an experiment.
+
+        Raises
+        ------
+        ExperimentNotFound
+        JobNotFound
         """
-        exp = self.experiment(exp_id)
+        exp = self.get_experiment(exp_id)
         if job_id is None:
             if len(exp["jobs"]) > 1:
                 raise ValueError(
                     f"Experiment {exp_id} has more than 1 job. You need to specify the 'job_id'."
                 )
             job_id = exp["jobs"][0]["id"]
-        return self.logs(job_id)
+        return self.get_logs(job_id)
+
+    def get_image(self, image: str) -> Dict[str, Any]:
+        """
+        Get info about an image.
+
+        Raises
+        ------
+        ImageNotFound
+        HTTPError
+
+        """
+        return self.request(
+            f"images/{urllib.parse.quote(image, safe='')}",
+            exceptions_for_status={404: ImageNotFound(image)},
+        ).json()
 
     def create_image(
         self, name: str, image_tag: str, workspace: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         Upload a Docker image to Beaker.
+
+        Raises
+        ------
+        ImageConflict
+        HTTPError
+
         """
-        workspace = workspace or self.workspace
+        workspace = workspace or self.config.default_workspace
 
         # Get local Docker image object.
         image = self.docker.images.get(image_tag)
@@ -135,6 +237,7 @@ def create_image(
             method="POST",
             data={"Workspace": workspace, "ImageID": image.id, "ImageTag": image_tag},
             query={"name": name},
+            exceptions_for_status={409: ImageConflict(name)},
         ).json()
 
         # Get the repo data for the Beaker image.
@@ -169,7 +272,20 @@ def create_image(
         self.request(f"images/{image_data['id']}", method="PATCH", data={"Commit": True})
 
         # Return info about the Beaker image.
-        return self.request(f"images/{image_data['id']}").json()
+        return self.get_image(image_data["id"])
 
     def delete_image(self, image_id: str):
-        self.request(f"images/{image_id}", method="DELETE")
+        """
+        Delete an image.
+
+        Raises
+        ------
+        ImageNotFound
+        HTTPError
+
+        """
+        self.request(
+            f"images/{image_id}",
+            method="DELETE",
+            exceptions_for_status={404: ImageNotFound(image_id)},
+        )