add volumes to taskqueues and endpoints (#166)

- Add volumes to taskqueues
- Add volumes to endpoints

---------

Co-authored-by: Luke Lombardi <[email protected]>

sdk/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beta9"
-version = "0.1.4"
+version = "0.1.6"
 description = ""
 authors = ["beam.cloud <[email protected]>"]
 packages = [

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)

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)

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)