pass running mode to dask backend

ITISFoundation · Mar 20, 2023 · 1cdf4d8 · 1cdf4d8
1 parent 01e31f3
commit 1cdf4d8
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/services/director-v2/src/simcore_service_director_v2/models/domains/comp_tasks.py b/services/director-v2/src/simcore_service_director_v2/models/domains/comp_tasks.py
@@ -13,6 +13,7 @@
     ServiceOutput,
     ServicePortKey,
 )
+from models_library.services_resources import BootMode
 from pydantic import BaseModel, ByteSize, Extra, Field, parse_obj_as, validator
 from pydantic.types import PositiveInt
 from simcore_postgres_database.models.comp_tasks import NodeClass, StateType
@@ -34,6 +35,7 @@ class Image(BaseModel):
     node_requirements: Optional[NodeRequirements] = Field(
         None, description="the requirements for the service to run on a node"
     )
+    boot_mode: BootMode = BootMode.CPU
     command: list[str] = Field(
         default=[
             "run",

diff --git a/services/director-v2/src/simcore_service_director_v2/modules/dask_client.py b/services/director-v2/src/simcore_service_director_v2/modules/dask_client.py
@@ -31,6 +31,7 @@
 from models_library.projects import ProjectID
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
+from models_library.services_resources import BootMode
 from models_library.users import UserID
 from pydantic import parse_obj_as
 from pydantic.networks import AnyUrl
@@ -99,6 +100,7 @@
         LogFileUploadURL,
         Commands,
         Optional[S3Settings],
+        BootMode,
     ],
     TaskOutputData,
 ]
@@ -206,6 +208,7 @@ def _comp_sidecar_fct(
             log_file_url: AnyUrl,
             command: list[str],
             s3_settings: Optional[S3Settings],
+            boot_mode: BootMode,
         ) -> TaskOutputData:
             """This function is serialized by the Dask client and sent over to the Dask sidecar(s)
             Therefore, (screaming here) DO NOT MOVE THAT IMPORT ANYWHERE ELSE EVER!!"""
@@ -220,6 +223,7 @@ def _comp_sidecar_fct(
                 log_file_url,
                 command,
                 s3_settings,
+                boot_mode,
             )
 
         if remote_fct is None:
@@ -314,6 +318,7 @@ def _comp_sidecar_fct(
                     log_file_url=log_file_url,
                     command=node_image.command,
                     s3_settings=s3_settings,
+                    boot_mode=node_image.boot_mode,
                     key=job_id,
                     resources=dask_resources,
                     retries=0,

diff --git a/services/director-v2/src/simcore_service_director_v2/modules/db/repositories/comp_tasks.py b/services/director-v2/src/simcore_service_director_v2/modules/db/repositories/comp_tasks.py
@@ -10,6 +10,7 @@
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
 from models_library.services import ServiceDockerData
+from models_library.services_resources import BootMode
 from models_library.users import UserID
 from sqlalchemy import literal_column
 from sqlalchemy.dialects.postgresql import insert
@@ -63,6 +64,12 @@ def _compute_node_requirements(node_resources: dict[str, Any]) -> NodeRequiremen
     return NodeRequirements.parse_obj(node_defined_resources)
 
 
+def _compute_node_boot_mode(node_resources: dict[str, Any]) -> BootMode:
+    for image_data in node_resources.values():
+        return BootMode(image_data.get("boot_modes")[0])
+    raise RuntimeError("No BootMode")
+
+
 async def _generate_tasks_list_from_project(
     project: ProjectAtDB,
     catalog_client: CatalogClient,
@@ -100,6 +107,7 @@ async def _generate_tasks_list_from_project(
 
         if node_resources:
             data.update(node_requirements=_compute_node_requirements(node_resources))
+            data["boot_mode"] = _compute_node_boot_mode(node_resources)
         if node_extras and node_extras.container_spec:
             data.update(command=node_extras.container_spec.command)
         image = Image.parse_obj(data)