Fix: Program vcpus and timeout could not be specified

aleph-im · Apr 29, 2022 · 1cf989a · 1cf989a
1 parent a26d011
commit 1cf989a
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 2 deletions.
diff --git a/src/aleph_client/__main__.py b/src/aleph_client/__main__.py
@@ -268,6 +268,8 @@ def program(
     entrypoint: str,
     channel: str = settings.DEFAULT_CHANNEL,
     memory: int = settings.DEFAULT_VM_MEMORY,
+    vcpus: int = settings.DEFAULT_VM_VCPUS,
+    timeout_seconds: float = settings.DEFAULT_VM_TIMEOUT,
     private_key: Optional[str] = settings.PRIVATE_KEY_STRING,
     private_key_file: Optional[str] = settings.PRIVATE_KEY_FILE,
     print_messages: bool = False,
@@ -364,6 +366,8 @@ def program(
             storage_engine=StorageEnum.storage,
             channel=channel,
             memory=memory,
+            vcpus=vcpus,
+            timeout_seconds=timeout_seconds,
             encoding=encoding,
             volumes=volumes,
             subscriptions=subscriptions,

diff --git a/src/aleph_client/asynchronous.py b/src/aleph_client/asynchronous.py
@@ -270,6 +270,8 @@ async def create_program(
     session: Optional[ClientSession] = None,
     api_server: str = settings.API_HOST,
     memory: int = settings.DEFAULT_VM_MEMORY,
+    vcpus: int = settings.DEFAULT_VM_VCPUS,
+    timeout_seconds: float = settings.DEFAULT_VM_TIMEOUT,
     encoding: Encoding = Encoding.zip,
     volumes: List[Dict] = None,
     subscriptions: Optional[List[Dict]] = None,
@@ -306,9 +308,9 @@ async def create_program(
                 "aleph_api": True,
             },
             "resources": {
-                "vcpus": 1,
+                "vcpus": vcpus,
                 "memory": memory,
-                "seconds": 30,
+                "seconds": timeout_seconds,
             },
             "runtime": {
                 "ref": runtime,

diff --git a/src/aleph_client/conf.py b/src/aleph_client/conf.py
@@ -18,6 +18,8 @@ class Settings(BaseSettings):
         "bd79839bf96e595a06da5ac0b6ba51dea6f7e2591bb913deccded04d831d29f4"
     )
     DEFAULT_VM_MEMORY: int = 128
+    DEFAULT_VM_VCPUS: int = 1
+    DEFAULT_VM_TIMEOUT: float = 30.
 
     CODE_USES_SQUASHFS: bool = which("mksquashfs") is not None  # True if command exists
 

diff --git a/src/aleph_client/synchronous.py b/src/aleph_client/synchronous.py
@@ -55,6 +55,8 @@ def create_program(
     session: Optional[ClientSession] = None,
     api_server: str = settings.API_HOST,
     memory: int = settings.DEFAULT_VM_MEMORY,
+    vcpus: int = settings.DEFAULT_VM_VCPUS,
+    timeout_seconds: float = settings.DEFAULT_VM_TIMEOUT,
     encoding: Encoding = Encoding.zip,
     volumes: List[Dict] = None,
     subscriptions: Optional[List[Dict]] = None,
@@ -70,6 +72,8 @@ def create_program(
         session=session,
         api_server=api_server,
         memory=memory,
+        vcpus=vcpus,
+        timeout_seconds=timeout_seconds,
         encoding=encoding,
         volumes=volumes,
         subscriptions=subscriptions,