refactor to use new RunID

Author: Andrei Neagu

packages/models-library/src/models_library/rabbitmq_messages.py

@@ -15,6 +15,7 @@
 from .projects_state import RunningState
 from .services import ServiceKey, ServiceType, ServiceVersion
 from .services_resources import ServiceResourcesDict
+from .services_types import RunID
 from .users import UserID
 from .utils.enums import StrAutoEnum
 from .wallets import WalletID
@@ -178,7 +179,7 @@ class RabbitResourceTrackingMessageType(StrAutoEnum):
 class RabbitResourceTrackingBaseMessage(RabbitMessageBase):
     channel_name: Literal["io.simcore.service.tracking"] = "io.simcore.service.tracking"
 
-    service_run_id: str = Field(
+    service_run_id: RunID = Field(
         ..., description="uniquely identitifies the service run"
     )
     created_at: datetime.datetime = Field(

packages/models-library/src/models_library/services_types.py

@@ -1,18 +1,28 @@
-from typing import Annotated, Any, TypeAlias
+from typing import TYPE_CHECKING, Annotated, Any, Self, TypeAlias
 from uuid import uuid4
 
 import arrow
-from pydantic import GetCoreSchemaHandler, StringConstraints, ValidationInfo
+from pydantic import (
+    GetCoreSchemaHandler,
+    PositiveInt,
+    StringConstraints,
+    ValidationInfo,
+)
 from pydantic_core import CoreSchema, core_schema
 
 from .basic_regex import PROPERTY_KEY_RE, SIMPLE_VERSION_RE
+from .projects_nodes_io import NodeID
 from .services_regex import (
     COMPUTATIONAL_SERVICE_KEY_RE,
     DYNAMIC_SERVICE_KEY_RE,
     FILENAME_RE,
     SERVICE_ENCODED_KEY_RE,
     SERVICE_KEY_RE,
 )
+from .users import UserID
+
+if TYPE_CHECKING:
+    from .projects import ProjectID
 
 ServicePortKey: TypeAlias = Annotated[str, StringConstraints(pattern=PROPERTY_KEY_RE)]
 
@@ -44,12 +54,15 @@ class RunID(str):
     and old volumes for different runs.
     Avoids overwriting data that left dropped on the node (due to an error)
     and gives the osparc-agent an opportunity to back it up.
+    The resource-usage-tracker tracker uses these RunIDs to keep track of
+    resource usage from comp and dynamic services.
     """
 
     __slots__ = ()
 
     @classmethod
-    def create(cls) -> "RunID":
+    def create_for_dynamic_sidecar(cls) -> Self:
+        """used for dynamic services"""
         # NOTE: there was a legacy version of this RunID
         # legacy version:
         #   '0ac3ed64-665b-42d2-95f7-e59e0db34242'
@@ -59,6 +72,17 @@ def create(cls) -> "RunID":
         run_id_format = f"{utc_int_timestamp}_{uuid4()}"
         return cls(run_id_format)
 
+    @classmethod
+    def get_resource_tracking_run_id(
+        cls,
+        user_id: UserID,
+        project_id: "ProjectID",
+        node_id: NodeID,
+        iteration: PositiveInt,
+    ) -> Self:
+        """used by computational services"""
+        return cls(f"comp_{user_id}_{project_id}_{node_id}_{iteration}")
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls,

packages/models-library/tests/test_services_types.py

@@ -0,0 +1,36 @@
+import pytest
+from models_library.projects import ProjectID
+from models_library.projects_nodes import NodeID
+from models_library.services_types import RunID
+from models_library.users import UserID
+from pydantic import PositiveInt
+
+
+@pytest.mark.parametrize(
+    "user_id, project_id, node_id, iteration, expected_result",
+    [
+        (
+            2,
+            ProjectID("e08356e4-eb74-49e9-b769-2c26e34c61d9"),
+            NodeID("a08356e4-eb74-49e9-b769-2c26e34c61d1"),
+            5,
+            "comp_2_e08356e4-eb74-49e9-b769-2c26e34c61d9_a08356e4-eb74-49e9-b769-2c26e34c61d1_5",
+        )
+    ],
+)
+def test_run_id_get_resource_tracking_run_id(
+    user_id: UserID,
+    project_id: ProjectID,
+    node_id: NodeID,
+    iteration: PositiveInt,
+    expected_result: str,
+):
+    resource_tracking_run_id = RunID.get_resource_tracking_run_id(
+        user_id, project_id, node_id, iteration
+    )
+    assert isinstance(resource_tracking_run_id, RunID)
+    assert resource_tracking_run_id == expected_result
+
+
+def test_run_id_create_for_dynamic_sidecar():
+    assert isinstance(RunID.create_for_dynamic_sidecar(), RunID)

services/agent/tests/unit/conftest.py

@@ -57,7 +57,7 @@ def test_client(initialized_app: FastAPI) -> TestClient:
 
 @pytest.fixture
 def run_id() -> RunID:
-    return RunID.create()
+    return RunID.create_for_dynamic_sidecar()
 
 
 @pytest.fixture

services/director-v2/src/simcore_service_director_v2/models/dynamic_services_scheduler.py

@@ -377,7 +377,7 @@ class SchedulerData(CommonServiceDetails, DynamicSidecarServiceLabels):
         description="Name of the current dynamic-sidecar being observed",
     )
     run_id: RunID = Field(
-        default_factory=RunID.create,
+        default_factory=RunID.create_for_dynamic_sidecar,
         description=(
             "Uniquely identify the dynamic sidecar session (a.k.a. 2 "
             "subsequent exact same services will have a different run_id)"

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_base.py

@@ -26,6 +26,7 @@
 from models_library.projects_nodes_io import NodeID, NodeIDStr
 from models_library.projects_state import RunningState
 from models_library.services import ServiceType
+from models_library.services_types import RunID
 from models_library.users import UserID
 from networkx.classes.reportviews import InDegreeView
 from servicelib.common_headers import UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE
@@ -66,7 +67,6 @@
     TASK_TO_START_STATES,
     WAITING_FOR_START_STATES,
     create_service_resources_from_task,
-    get_resource_tracking_run_id,
 )
 
 _logger = logging.getLogger(__name__)
@@ -295,7 +295,7 @@ def _need_heartbeat(task: CompTaskAtDB) -> bool:
                 *(
                     publish_service_resource_tracking_heartbeat(
                         self.rabbitmq_client,
-                        get_resource_tracking_run_id(
+                        RunID.get_resource_tracking_run_id(
                             user_id, t.project_id, t.node_id, iteration
                         ),
                     )
@@ -348,7 +348,7 @@ async def _process_started_tasks(
             *(
                 publish_service_resource_tracking_started(
                     self.rabbitmq_client,
-                    service_run_id=get_resource_tracking_run_id(
+                    service_run_id=RunID.get_resource_tracking_run_id(
                         user_id, t.project_id, t.node_id, iteration
                     ),
                     wallet_id=run_metadata.get("wallet_id"),

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_dask.py

@@ -18,6 +18,7 @@
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
 from models_library.rabbitmq_messages import SimcorePlatformStatus
+from models_library.services_types import RunID
 from models_library.users import UserID
 from servicelib.common_headers import UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE
 from servicelib.logging_utils import log_catch
@@ -48,7 +49,6 @@
 from ..db.repositories.comp_runs import CompRunsRepository
 from ..db.repositories.comp_tasks import CompTasksRepository
 from ._scheduler_base import BaseCompScheduler
-from ._utils import get_resource_tracking_run_id
 
 _logger = logging.getLogger(__name__)
 
@@ -129,7 +129,7 @@ async def _start_tasks(
                         hardware_info=task.hardware_info,
                         callback=wake_up_callback,
                         metadata=comp_run.metadata,
-                        resource_tracking_run_id=get_resource_tracking_run_id(
+                        resource_tracking_run_id=RunID.get_resource_tracking_run_id(
                             user_id, project_id, node_id, comp_run.iteration
                         ),
                     )
@@ -322,7 +322,9 @@ async def _process_task_result(
             # resource tracking
             await publish_service_resource_tracking_stopped(
                 self.rabbitmq_client,
-                get_resource_tracking_run_id(user_id, project_id, node_id, iteration),
+                RunID.get_resource_tracking_run_id(
+                    user_id, project_id, node_id, iteration
+                ),
                 simcore_platform_status=simcore_platform_status,
             )
             # instrumentation

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_utils.py

@@ -2,19 +2,15 @@
 
 from fastapi import FastAPI
 from models_library.docker import DockerGenericTag
-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 (
     ResourceValue,
     ServiceResourcesDict,
     ServiceResourcesDictHelpers,
 )
-from models_library.users import UserID
 from servicelib.redis import RedisClientSDK
 from settings_library.redis import RedisDatabase
 
-from ...models.comp_runs import Iteration
 from ...models.comp_tasks import CompTaskAtDB
 from ..redis import get_redis_client_manager
 
@@ -55,12 +51,6 @@
 }
 
 
-def get_resource_tracking_run_id(
-    user_id: UserID, project_id: ProjectID, node_id: NodeID, iteration: Iteration
-) -> str:
-    return f"comp_{user_id}_{project_id}_{node_id}_{iteration}"
-
-
 def create_service_resources_from_task(task: CompTaskAtDB) -> ServiceResourcesDict:
     assert task.image.node_requirements  # nosec
     return ServiceResourcesDictHelpers.create_from_single_service(

services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_event_create_sidecars.py

@@ -237,7 +237,7 @@ async def action(cls, app: FastAPI, scheduler_data: SchedulerData) -> None:
 
         # Each time a new dynamic-sidecar service is created
         # generate a new `run_id` to avoid resource collisions
-        scheduler_data.run_id = RunID.create()
+        scheduler_data.run_id = RunID.create_for_dynamic_sidecar()
 
         rpc_client: RabbitMQRPCClient = app.state.rabbitmq_rpc_client
 

services/director-v2/src/simcore_service_director_v2/utils/rabbitmq.py

@@ -15,6 +15,7 @@
 )
 from models_library.services import ServiceKey, ServiceType, ServiceVersion
 from models_library.services_resources import ServiceResourcesDict
+from models_library.services_types import RunID
 from models_library.users import UserID
 from models_library.wallets import WalletID
 from pydantic import NonNegativeFloat
@@ -70,7 +71,7 @@ async def publish_service_stopped_metrics(
 
 async def publish_service_resource_tracking_started(  # pylint: disable=too-many-arguments # noqa: PLR0913
     rabbitmq_client: RabbitMQClient,
-    service_run_id: str,
+    service_run_id: RunID,
     *,
     wallet_id: WalletID | None,
     wallet_name: str | None,
@@ -127,7 +128,7 @@ async def publish_service_resource_tracking_started(  # pylint: disable=too-many
 
 async def publish_service_resource_tracking_stopped(
     rabbitmq_client: RabbitMQClient,
-    service_run_id: str,
+    service_run_id: RunID,
     *,
     simcore_platform_status: SimcorePlatformStatus,
 ) -> None:
@@ -138,7 +139,7 @@ async def publish_service_resource_tracking_stopped(
 
 
 async def publish_service_resource_tracking_heartbeat(
-    rabbitmq_client: RabbitMQClient, service_run_id: str
+    rabbitmq_client: RabbitMQClient, service_run_id: RunID
 ) -> None:
     message = RabbitResourceTrackingHeartbeatMessage(service_run_id=service_run_id)
     await rabbitmq_client.publish(message.channel_name, message)