Golem infrastructure properties in CreateTaskReply (#56)

Author: mfranciszkiewicz

golem_task_api/proto/golem_task_api.proto

@@ -22,6 +22,9 @@ service ProviderApp {
   rpc Shutdown (ShutdownRequest) returns (ShutdownReply) {}
 }
 
+message Infrastructure {
+  float min_memory_mib = 1;
+}
 
 message CreateTaskRequest {
   string task_id = 1;
@@ -32,6 +35,7 @@ message CreateTaskRequest {
 message CreateTaskReply {
   string env_id = 1;
   string prerequisites_json = 2;
+  Infrastructure inf_requirements = 3;
 }
 
 message NextSubtaskRequest {

python/golem_task_api/client.py

@@ -37,7 +37,7 @@
     ProviderAppStub,
     RequestorAppStub,
 )
-from golem_task_api.structs import Subtask, Task
+from golem_task_api.structs import Subtask, Task, Infrastructure
 
 CONNECTION_TIMEOUT = 5.0  # seconds
 
@@ -190,10 +190,16 @@ async def create_task(
         request.max_subtasks_count = max_subtasks_count
         request.task_params_json = json.dumps(task_params)
         reply = await self._golem_app.CreateTask(request)
-        return Task(
+
+        inf_requirements = Infrastructure()
+        inf_requirements.min_memory_mib = reply.inf_requirements.min_memory_mib
+
+        task = Task(
             env_id=reply.env_id,
-            prerequisites=json.loads(reply.prerequisites_json)
-        )
+            prerequisites=json.loads(reply.prerequisites_json),
+            inf_requirements=inf_requirements)
+
+        return task
 
     async def next_subtask(
             self,

python/golem_task_api/envs/__init__.py

@@ -2,13 +2,14 @@
 
 from golem_task_api.envs.constants import DOCKER_CPU_ENV_ID, DOCKER_GPU_ENV_ID
 from golem_task_api.proto.envs_pb2 import DockerPrerequisites
-from golem_task_api.structs import Task
+from golem_task_api.structs import Task, Infrastructure
 
 
 def _create_docker_task(
         env_id: str,
         image: str,
-        tag: str
+        tag: str,
+        inf: Infrastructure,
 ) -> Task:
     prerequisites = DockerPrerequisites()
     prerequisites.image = image
@@ -19,13 +20,13 @@ def _create_docker_task(
 
     return Task(
         env_id=env_id,
-        prerequisites=prerequisites_dict
-    )
+        prerequisites=prerequisites_dict,
+        inf_requirements=inf)
 
 
-def create_docker_cpu_task(image: str, tag: str) -> Task:
-    return _create_docker_task(DOCKER_CPU_ENV_ID, image, tag)
+def create_docker_cpu_task(image: str, tag: str, inf: Infrastructure) -> Task:
+    return _create_docker_task(DOCKER_CPU_ENV_ID, image, tag, inf)
 
 
-def create_docker_gpu_task(image: str, tag: str) -> Task:
-    return _create_docker_task(DOCKER_GPU_ENV_ID, image, tag)
+def create_docker_gpu_task(image: str, tag: str, inf: Infrastructure) -> Task:
+    return _create_docker_task(DOCKER_GPU_ENV_ID, image, tag, inf)

python/golem_task_api/messages.py

@@ -5,6 +5,7 @@
     AbortTaskReply,
     CreateTaskRequest,
     CreateTaskReply,
+    Infrastructure,
     NextSubtaskRequest,
     NextSubtaskReply,
     SubtaskReply,