♻️✨🐛Dask-Sidecar: add RabbitMQ dependency and remove usage of deprecated Pub/Sub for logs 🚨🚨🚨

packages/dask-task-models-library/src/dask_task_models_library/container_tasks/events.py

@@ -1,6 +1,4 @@
-import logging
 from abc import ABC, abstractmethod
-from typing import TypeAlias
 
 import dask.typing
 from distributed.worker import get_worker
@@ -85,48 +83,3 @@ def ensure_between_0_1(cls, v):
         if 0 <= v <= 1:
             return v
         return min(max(0, v), 1)
-
-
-LogMessageStr: TypeAlias = str
-LogLevelInt: TypeAlias = int
-
-
-class TaskLogEvent(BaseTaskEvent):
-    log: LogMessageStr
-    log_level: LogLevelInt
-
-    @staticmethod
-    def topic_name() -> str:
-        return "task_logs"
-
-    @classmethod
-    def from_dask_worker(
-        cls, log: str, log_level: LogLevelInt, *, task_owner: TaskOwner
-    ) -> "TaskLogEvent":
-        worker = get_worker()
-        job_id = worker.get_current_task()
-        return cls(
-            job_id=_dask_key_to_dask_task_id(job_id),
-            log=log,
-            log_level=log_level,
-            task_owner=task_owner,
-        )
-
-    model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "job_id": "simcore/services/comp/sleeper:1.1.0:projectid_ec7e595a-63ee-46a1-a04a-901b11b649f8:nodeid_39467d89-b659-4914-9359-c40b1b6d1d6d:uuid_5ee5c655-450d-4711-a3ec-32ffe16bc580",
-                    "log": "some logs",
-                    "log_level": logging.INFO,
-                    "task_owner": {
-                        "user_id": 32,
-                        "project_id": "ec7e595a-63ee-46a1-a04a-901b11b649f8",
-                        "node_id": "39467d89-b659-4914-9359-c40b1b6d1d6d",
-                        "parent_project_id": None,
-                        "parent_node_id": None,
-                    },
-                },
-            ]
-        }
-    )

packages/dask-task-models-library/src/dask_task_models_library/container_tasks/io.py

@@ -17,46 +17,59 @@
     StrictInt,
     StrictStr,
 )
+from pydantic.config import JsonDict
 
 TaskCancelEventName = "cancel_event_{}"
 
 
 class PortSchema(BaseModel):
     required: bool
 
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "required": True,
+                    },
+                    {
+                        "required": False,
+                    },
+                ]
+            }
+        )
+
     model_config = ConfigDict(
         extra="forbid",
-        json_schema_extra={
-            "examples": [
-                {
-                    "required": True,
-                },
-                {
-                    "required": False,
-                },
-            ]
-        },
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
 class FilePortSchema(PortSchema):
     mapping: str | None = None
     url: AnyUrl
 
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "mapping": "some_filename.txt",
+                        "url": "sftp://some_file_url",
+                        "required": True,
+                    },
+                    {
+                        "required": False,
+                        "url": "s3://another_file_url",
+                    },
+                ]
+            }
+        )
+
     model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "mapping": "some_filename.txt",
-                    "url": "sftp://some_file_url",
-                    "required": True,
-                },
-                {
-                    "required": False,
-                    "url": "s3://another_file_url",
-                },
-            ]
-        }
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
@@ -70,18 +83,27 @@ class FileUrl(BaseModel):
         default=None, description="the file MIME type", pattern=MIME_TYPE_RE
     )
 
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "url": "https://some_file_url",
+                        "file_mime_type": "application/json",
+                    },
+                    {
+                        "url": "https://some_file_url",
+                        "file_mapping": "some_file_name.txt",
+                        "file_mime_type": "application/json",
+                    },
+                ]
+            }
+        )
+
     model_config = ConfigDict(
         extra="forbid",
-        json_schema_extra={
-            "examples": [
-                {"url": "https://some_file_url", "file_mime_type": "application/json"},
-                {
-                    "url": "https://some_file_url",
-                    "file_mapping": "some_file_name.txt",
-                    "file_mime_type": "application/json",
-                },
-            ]
-        },
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
@@ -99,18 +121,24 @@ class FileUrl(BaseModel):
 
 
 class TaskInputData(DictModel[ServicePortKey, PortValue]):
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "boolean_input": False,
+                        "int_input": -45,
+                        "float_input": 4564.45,
+                        "string_input": "nobody thinks like a string",
+                        "file_input": {"url": "s3://thatis_file_url"},
+                    },
+                ]
+            }
+        )
+
     model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "boolean_input": False,
-                    "int_input": -45,
-                    "float_input": 4564.45,
-                    "string_input": "nobody thinks like a string",
-                    "file_input": {"url": "s3://thatis_file_url"},
-                },
-            ]
-        }
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
@@ -126,26 +154,32 @@ class TaskOutputDataSchema(DictModel[ServicePortKey, PortSchemaValue]):
     # does not work well in that case. For that reason, the schema is
     # sent as a json-schema instead of with a dynamically-created model class
     #
-    model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "boolean_output": {"required": False},
-                    "int_output": {"required": True},
-                    "float_output": {"required": True},
-                    "string_output": {"required": False},
-                    "file_output": {
-                        "required": True,
-                        "url": "https://some_file_url",
-                        "mapping": "the_output_filename",
-                    },
-                    "optional_file_output": {
-                        "required": False,
-                        "url": "s3://one_file_url",
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "boolean_output": {"required": False},
+                        "int_output": {"required": True},
+                        "float_output": {"required": True},
+                        "string_output": {"required": False},
+                        "file_output": {
+                            "required": True,
+                            "url": "https://some_file_url",
+                            "mapping": "the_output_filename",
+                        },
+                        "optional_file_output": {
+                            "required": False,
+                            "url": "s3://one_file_url",
+                        },
                     },
-                },
-            ]
-        }
+                ]
+            }
+        )
+
+    model_config = ConfigDict(
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
@@ -181,16 +215,20 @@ def from_task_output(
 
         return cls.model_validate(data)
 
-    model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "boolean_output": False,
-                    "int_output": -45,
-                    "float_output": 4564.45,
-                    "string_output": "nobody thinks like a string",
-                    "file_output": {"url": "s3://yet_another_file_url"},
-                },
-            ]
-        }
-    )
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "boolean_output": False,
+                        "int_output": -45,
+                        "float_output": 4564.45,
+                        "string_output": "nobody thinks like a string",
+                        "file_output": {"url": "s3://yet_another_file_url"},
+                    },
+                ]
+            }
+        )
+
+    model_config = ConfigDict(json_schema_extra=_update_json_schema_extra)

packages/dask-task-models-library/src/dask_task_models_library/container_tasks/protocol.py

@@ -7,6 +7,7 @@
 from models_library.services_resources import BootMode
 from models_library.users import UserID
 from pydantic import AnyUrl, BaseModel, ConfigDict, model_validator
+from pydantic.config import JsonDict
 from settings_library.s3 import S3Settings
 
 from .docker import DockerBasicAuth
@@ -44,25 +45,31 @@ def check_parent_valid(cls, values: dict[str, Any]) -> dict[str, Any]:
             raise ValueError(msg)
         return values
 
+    @staticmethod
+    def _update_json_schema_extra(schema: JsonDict) -> None:
+        schema.update(
+            {
+                "examples": [
+                    {
+                        "user_id": 32,
+                        "project_id": "ec7e595a-63ee-46a1-a04a-901b11b649f8",
+                        "node_id": "39467d89-b659-4914-9359-c40b1b6d1d6d",
+                        "parent_project_id": None,
+                        "parent_node_id": None,
+                    },
+                    {
+                        "user_id": 32,
+                        "project_id": "ec7e595a-63ee-46a1-a04a-901b11b649f8",
+                        "node_id": "39467d89-b659-4914-9359-c40b1b6d1d6d",
+                        "parent_project_id": "887e595a-63ee-46a1-a04a-901b11b649f8",
+                        "parent_node_id": "aa467d89-b659-4914-9359-c40b1b6d1d6d",
+                    },
+                ]
+            }
+        )
+
     model_config = ConfigDict(
-        json_schema_extra={
-            "examples": [
-                {
-                    "user_id": 32,
-                    "project_id": "ec7e595a-63ee-46a1-a04a-901b11b649f8",
-                    "node_id": "39467d89-b659-4914-9359-c40b1b6d1d6d",
-                    "parent_project_id": None,
-                    "parent_node_id": None,
-                },
-                {
-                    "user_id": 32,
-                    "project_id": "ec7e595a-63ee-46a1-a04a-901b11b649f8",
-                    "node_id": "39467d89-b659-4914-9359-c40b1b6d1d6d",
-                    "parent_project_id": "887e595a-63ee-46a1-a04a-901b11b649f8",
-                    "parent_node_id": "aa467d89-b659-4914-9359-c40b1b6d1d6d",
-                },
-            ]
-        }
+        json_schema_extra=_update_json_schema_extra,
     )
 
 
@@ -83,13 +90,15 @@ class ContainerTaskParameters(BaseModel):
                 {
                     "image": "ubuntu",
                     "tag": "latest",
-                    "input_data": TaskInputData.model_config["json_schema_extra"]["examples"][0],  # type: ignore[index]
-                    "output_data_keys": TaskOutputDataSchema.model_config["json_schema_extra"]["examples"][0],  # type: ignore[index]
+                    "input_data": TaskInputData.model_json_schema()["examples"][0],
+                    "output_data_keys": TaskOutputDataSchema.model_json_schema()[
+                        "examples"
+                    ][0],
                     "command": ["sleep 10", "echo hello"],
                     "envs": {"MYENV": "is an env"},
                     "labels": {"io.simcore.thelabel": "is amazing"},
                     "boot_mode": BootMode.CPU.value,
-                    "task_owner": TaskOwner.model_config["json_schema_extra"]["examples"][0],  # type: ignore[index]
+                    "task_owner": TaskOwner.model_json_schema()["examples"][0],
                 },
             ]
         }
@@ -104,5 +113,4 @@ def __call__(
         docker_auth: DockerBasicAuth,
         log_file_url: LogFileUploadURL,
         s3_settings: S3Settings | None,
-    ) -> TaskOutputData:
-        ...
+    ) -> TaskOutputData: ...