diff --git a/.flake8 b/.flake8 index 7bdc653..717710d 100644 --- a/.flake8 +++ b/.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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c3bdcf..dad71a7 100644 --- a/.github/workflows/ci.yml +++ b/.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/* diff --git a/CHANGELOG.md b/CHANGELOG.md index ba8e016..731d3a6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index c044e6d..f28c51c 100644 --- a/README.md +++ b/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() ``` diff --git a/beaker/__init__.py b/beaker/__init__.py index 6a8bab2..7dac329 100644 --- a/beaker/__init__.py +++ b/beaker/__init__.py @@ -1,3 +1,3 @@ -__all__ = ["Beaker"] - from .client import Beaker +from .config import Config +from .exceptions import * diff --git a/beaker/client.py b/beaker/client.py index 5826bff..e3ff86e 100644 --- a/beaker/client.py +++ b/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,11 +25,17 @@ 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": @@ -32,10 +43,7 @@ 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)}, + ) diff --git a/beaker/config.py b/beaker/config.py new file mode 100644 index 0000000..37ecce5 --- /dev/null +++ b/beaker/config.py @@ -0,0 +1,97 @@ +import os +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import ClassVar, Optional + +import yaml + +from .exceptions import ConfigurationError + + +@dataclass +class Config: + user_token: str + """ + Beaker user token that can be obtained from + `beaker.org `_. + """ + + agent_address: str = "https://beaker.org" + """ + The address of the Beaker server. + """ + + default_org: Optional[str] = None + """ + Default Beaker organization to use. + """ + + default_workspace: Optional[str] = None + """ + Default Beaker workspace to use. + """ + + DEFAULT_CONFIG_LOCATION: ClassVar[Path] = Path.home() / ".beaker" / "config.yml" + ADDRESS_KEY: ClassVar[str] = "BEAKER_ADDR" + CONFIG_PATH_KEY: ClassVar[str] = "BEAKER_CONFIG" + TOKEN_KEY: ClassVar[str] = "BEAKER_TOKEN" + + @classmethod + def from_env(cls) -> "Config": + """ + Initialize a config from environment variables or a local config file if one + can be found. + + .. note:: + Environment variables take precedence over values in the config file. + + """ + config: Config + + path = cls.find_config() + if path is not None: + config = cls.from_path(path) + if cls.TOKEN_KEY in os.environ: + config.user_token = os.environ[cls.TOKEN_KEY] + elif cls.TOKEN_KEY in os.environ: + config = cls( + user_token=os.environ[cls.TOKEN_KEY], + ) + else: + raise ConfigurationError( + f"Missing config file or environment variable '{cls.TOKEN_KEY}'" + ) + + # Override with environment variables. + if cls.ADDRESS_KEY in os.environ: + config.agent_address = os.environ[cls.ADDRESS_KEY] + + return config + + @classmethod + def from_path(cls, path: Path) -> "Config": + """ + Initialize a config from a local config file. + """ + with open(path) as config_file: + return cls(**yaml.load(config_file, Loader=yaml.SafeLoader)) + + def save(self, path: Optional[Path]): + """ + Save the config to the given path. + """ + path = path or self.DEFAULT_CONFIG_LOCATION + with open(path, "w") as config_file: + yaml.dump(asdict(self), config_file) + + @classmethod + def find_config(cls) -> Optional[Path]: + if cls.CONFIG_PATH_KEY in os.environ: + path = Path(os.environ[cls.CONFIG_PATH_KEY]) + if path.is_file(): + return path + + if cls.DEFAULT_CONFIG_LOCATION.is_file(): + return cls.DEFAULT_CONFIG_LOCATION + + return None diff --git a/beaker/exceptions.py b/beaker/exceptions.py new file mode 100644 index 0000000..418abad --- /dev/null +++ b/beaker/exceptions.py @@ -0,0 +1,42 @@ +from requests.exceptions import ( # noqa: F401, re-imported here for convenience + HTTPError, +) + + +class BeakerError(Exception): + """ + Base class for all Beaker errors other than :exc:`HTTPError` which is re-exported + from :exc:`requests.exceptions.HTTPError`. + """ + + +class ConfigurationError(BeakerError): + pass + + +class ImageNotFound(BeakerError): + pass + + +class ImageConflict(BeakerError): + pass + + +class WorkspaceNotFound(BeakerError): + pass + + +class ExperimentNotFound(BeakerError): + pass + + +class ExperimentConflict(BeakerError): + pass + + +class DatasetNotFound(BeakerError): + pass + + +class JobNotFound(BeakerError): + pass diff --git a/docs/source/api.rst b/docs/source/api.rst index 42c26b1..e960c67 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -1,5 +1,22 @@ API === +Client +------ + .. autoclass:: beaker.Beaker :members: + +Config +------ + +.. autoclass:: beaker.config.Config + :members: + +Exceptions +---------- + +.. autoexception:: beaker.exceptions.HTTPError + +.. automodule:: beaker.exceptions + :members: diff --git a/requirements.txt b/requirements.txt index ed416ec..01c97cc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ requests tqdm +PyYAML docker>=5.0,<6.0 diff --git a/scripts/hello_world.py b/scripts/hello_world.py new file mode 100644 index 0000000..1ce1885 --- /dev/null +++ b/scripts/hello_world.py @@ -0,0 +1,49 @@ +import sys + +from beaker import Beaker, ImageNotFound + + +def main( + experiment: str = "hello-world", + image: str = "hello-world", + cluster: str = "AI2/petew-cpu", + org: str = "AI2", +): + # Setup. + beaker = Beaker.from_env() + beaker.config.default_workspace = f"{org}/{beaker.user}" + + # Check if image exists on Beaker and create it if it doesn't. + beaker_image = image.replace(":", "-") + try: + image_data = beaker.get_image(f"{beaker.user}/{beaker_image}") + except ImageNotFound: + image_data = beaker.create_image( + name=beaker_image, + image_tag=image, + ) + + # Submit experiment. + experiment_data = beaker.create_experiment( + experiment, + { + "version": "v2-alpha", + "tasks": [ + { + "name": "main", + "image": {"beaker": image_data["id"]}, + "context": {"cluster": cluster}, + "result": {"path": "/unused"}, # required even if the task produces no output. + }, + ], + }, + ) + experiment_id = experiment_data["id"] + print( + f"Experiment {experiment_id} submitted.\n" + f"See progress at https://beaker.org/ex/{experiment_id}" + ) + + +if __name__ == "__main__": + main(*sys.argv[1:]) diff --git a/tests/client_test.py b/tests/client_test.py index 65ef6f0..8ae7d64 100644 --- a/tests/client_test.py +++ b/tests/client_test.py @@ -5,11 +5,14 @@ import pytest from beaker.client import Beaker +from beaker.config import Config @pytest.fixture(scope="module") def client(): - return Beaker.from_env(workspace="ai2/beaker-py") + config = Config.from_env() + config.default_workspace = "ai2/beaker-py" + return Beaker(config) @pytest.fixture(scope="module") @@ -27,5 +30,4 @@ def test_images(client, docker_client): beaker_image_name = petname.generate() + "-" + str(random.randint(0, 99)) push_result = client.create_image(beaker_image_name, "hello-world") assert push_result["originalTag"] == "hello-world" - client.delete_image(push_result["id"]) diff --git a/tests/hello_test.py b/tests/hello_test.py deleted file mode 100644 index 23b9253..0000000 --- a/tests/hello_test.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_hello(): - print("Hello, World!")