From 383887009fc272726232136120af310215aefd42 Mon Sep 17 00:00:00 2001 From: luke-lombardi <33990301+luke-lombardi@users.noreply.github.com> Date: Mon, 29 Apr 2024 16:38:31 -0400 Subject: [PATCH] add volumes to taskqueues and endpoints (#166) - Add volumes to taskqueues - Add volumes to endpoints --------- Co-authored-by: Luke Lombardi --- sdk/pyproject.toml | 2 +- sdk/src/beta9/abstractions/endpoint.py | 5 ++++- sdk/src/beta9/abstractions/function.py | 15 ++++++++++++--- sdk/src/beta9/abstractions/taskqueue.py | 7 ++++++- 4 files changed, 23 insertions(+), 6 deletions(-) diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml index 76bbdc2f4..521d8018c 100644 --- a/sdk/pyproject.toml +++ b/sdk/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "beta9" -version = "0.1.4" +version = "0.1.6" description = "" authors = ["beam.cloud "] packages = [ diff --git a/sdk/src/beta9/abstractions/endpoint.py b/sdk/src/beta9/abstractions/endpoint.py index 93a9b2c95..921df387b 100644 --- a/sdk/src/beta9/abstractions/endpoint.py +++ b/sdk/src/beta9/abstractions/endpoint.py @@ -1,5 +1,5 @@ import os -from typing import Any, Callable, Optional, Union +from typing import Any, Callable, List, Optional, Union from .. import terminal from ..abstractions.base.runner import ( @@ -8,6 +8,7 @@ RunnerAbstraction, ) from ..abstractions.image import Image +from ..abstractions.volume import Volume from ..clients.endpoint import ( EndpointServiceStub, StartEndpointServeRequest, @@ -31,6 +32,7 @@ def __init__( keep_warm_seconds: int = 300, max_pending_tasks: int = 100, on_start: Optional[Callable] = None, + volumes: Optional[List[Volume]] = None, ): super().__init__( cpu=cpu, @@ -44,6 +46,7 @@ def __init__( keep_warm_seconds=keep_warm_seconds, max_pending_tasks=max_pending_tasks, on_start=on_start, + volumes=volumes, ) self.endpoint_stub: EndpointServiceStub = EndpointServiceStub(self.channel) diff --git a/sdk/src/beta9/abstractions/function.py b/sdk/src/beta9/abstractions/function.py index 516cdd6a5..d472944e8 100644 --- a/sdk/src/beta9/abstractions/function.py +++ b/sdk/src/beta9/abstractions/function.py @@ -37,6 +37,15 @@ class Function(RunnerAbstraction): applicable or no GPU required, leave it empty. Default is [GpuType.NoGPU](#gputype). image (Union[Image, dict]): The container image used for the task execution. Default is [Image](#image). + timeout (Optional[int]): + The maximum number of seconds a task can run before it times out. + Default is 3600. Set it to -1 to disable the timeout. + retries (Optional[int]): + The maximum number of times a task will be retried if the container crashes. Default is 3. + callback_url (Optional[str]): + An optional URL to send a callback to when a task is completed, timed out, or cancelled. + volumes (Optional[List[Volume]]): + A list of storage volumes to be associated with the function. Default is []. Example: ```python from beta9 import function, Image @@ -63,21 +72,21 @@ def __init__( cpu: Union[int, float, str] = 1.0, memory: int = 128, gpu: str = "", + image: Image = Image(), timeout: int = 3600, retries: int = 3, - image: Image = Image(), - volumes: Optional[List[Volume]] = None, callback_url: Optional[str] = "", + volumes: Optional[List[Volume]] = None, ) -> None: super().__init__( cpu=cpu, memory=memory, gpu=gpu, image=image, - volumes=volumes, timeout=timeout, retries=retries, callback_url=callback_url, + volumes=volumes, ) self.function_stub: FunctionServiceStub = FunctionServiceStub(self.channel) diff --git a/sdk/src/beta9/abstractions/taskqueue.py b/sdk/src/beta9/abstractions/taskqueue.py index 422eedf73..2ed3fa78e 100644 --- a/sdk/src/beta9/abstractions/taskqueue.py +++ b/sdk/src/beta9/abstractions/taskqueue.py @@ -1,6 +1,6 @@ import json import os -from typing import Any, Callable, Optional, Union +from typing import Any, Callable, List, Optional, Union from .. import terminal from ..abstractions.base.runner import ( @@ -10,6 +10,7 @@ RunnerAbstraction, ) from ..abstractions.image import Image +from ..abstractions.volume import Volume from ..clients.gateway import DeployStubRequest, DeployStubResponse from ..clients.taskqueue import ( StartTaskQueueServeRequest, @@ -68,6 +69,8 @@ class TaskQueue(RunnerAbstraction): loading models, or anything else computationally expensive. callback_url (Optional[str]): An optional URL to send a callback to when a task is completed, timed out, or cancelled. + volumes (Optional[List[Volume]]): + A list of storage volumes to be associated with the taskqueue. Default is []. Example: ```python from beta9 import task_queue, Image @@ -96,6 +99,7 @@ def __init__( max_pending_tasks: int = 100, on_start: Optional[Callable] = None, callback_url: Optional[str] = None, + volumes: Optional[List[Volume]] = None, ) -> None: super().__init__( cpu=cpu, @@ -110,6 +114,7 @@ def __init__( max_pending_tasks=max_pending_tasks, on_start=on_start, callback_url=callback_url, + volumes=volumes, ) self.taskqueue_stub: TaskQueueServiceStub = TaskQueueServiceStub(self.channel)