fixed mypy

Author: sanderegg

services/clusters-keeper/src/simcore_service_clusters_keeper/constants.py

@@ -1,9 +1,15 @@
 from typing import Final
 
-from aws_library.ec2._models import AWSTagKey
+from aws_library.ec2._models import AWSTagKey, AWSTagValue
 from pydantic import parse_obj_as
 
 DOCKER_STACK_DEPLOY_COMMAND_NAME: Final[str] = "private cluster docker deploy"
 DOCKER_STACK_DEPLOY_COMMAND_EC2_TAG_KEY: Final[AWSTagKey] = parse_obj_as(
     AWSTagKey, "io.simcore.clusters-keeper.private_cluster_docker_deploy"
 )
+
+USER_ID_TAG_KEY: Final[AWSTagKey] = parse_obj_as(AWSTagKey, "user_id")
+WALLET_ID_TAG_KEY: Final[AWSTagKey] = parse_obj_as(AWSTagKey, "wallet_id")
+ROLE_TAG_KEY: Final[AWSTagKey] = parse_obj_as(AWSTagKey, "role")
+WORKER_ROLE_TAG_VALUE: Final[AWSTagValue] = parse_obj_as(AWSTagValue, "worker")
+MANAGER_ROLE_TAG_VALUE: Final[AWSTagValue] = parse_obj_as(AWSTagValue, "manager")

services/clusters-keeper/src/simcore_service_clusters_keeper/core/settings.py

@@ -53,13 +53,15 @@ class Config(EC2Settings.Config):
 class ClustersKeeperSSMSettings(SSMSettings):
     class Config(SSMSettings.Config):
         env_prefix = CLUSTERS_KEEPER_ENV_PREFIX
-        prefixed_examples: ClassVar[list[dict[str, Any]]] = [
-            {f"{CLUSTERS_KEEPER_ENV_PREFIX}{key}": var for key, var in example.items()}
-            for example in SSMSettings.Config.schema_extra["examples"]
-        ]
 
-        schema_extra: ClassVar[dict[str, Any]] = {
-            "examples": prefixed_examples,
+        schema_extra: ClassVar[dict[str, Any]] = {  # type: ignore[misc]
+            "examples": [
+                {
+                    f"{CLUSTERS_KEEPER_ENV_PREFIX}{key}": var
+                    for key, var in example.items()
+                }
+                for example in SSMSettings.Config.schema_extra["examples"]
+            ],
         }
 
 

services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_core.py

@@ -16,6 +16,10 @@
 from ..constants import (
     DOCKER_STACK_DEPLOY_COMMAND_EC2_TAG_KEY,
     DOCKER_STACK_DEPLOY_COMMAND_NAME,
+    ROLE_TAG_KEY,
+    USER_ID_TAG_KEY,
+    WALLET_ID_TAG_KEY,
+    WORKER_ROLE_TAG_VALUE,
 )
 from ..core.settings import get_application_settings
 from ..modules.clusters import (
@@ -26,7 +30,12 @@
 )
 from ..utils.clusters import create_deploy_cluster_stack_script
 from ..utils.dask import get_scheduler_auth, get_scheduler_url
-from ..utils.ec2 import HEARTBEAT_TAG_KEY, get_cluster_name
+from ..utils.ec2 import (
+    HEARTBEAT_TAG_KEY,
+    get_cluster_name,
+    user_id_from_instance_tags,
+    wallet_id_from_instance_tags,
+)
 from .dask import is_scheduler_busy, ping_scheduler
 from .ec2 import get_ec2_client
 from .ssm import get_ssm_client
@@ -196,24 +205,26 @@ async def check_clusters(app: FastAPI) -> None:
         ]
         started_instances_ready_for_command = ec2_connected_to_ssm_server
         if started_instances_ready_for_command:
-            ssm_command = await ssm_client.send_command(
-                [i.id for i in started_instances_ready_for_command],
-                command=create_deploy_cluster_stack_script(
-                    app_settings,
-                    cluster_machines_name_prefix=get_cluster_name(
+            # we need to send 1 command per machine here, as the user_id/wallet_id changes
+            for i in started_instances_ready_for_command:
+                ssm_command = await ssm_client.send_command(
+                    [i.id],
+                    command=create_deploy_cluster_stack_script(
                         app_settings,
-                        user_id=user_id,
-                        wallet_id=wallet_id,
-                        is_manager=False,
+                        cluster_machines_name_prefix=get_cluster_name(
+                            app_settings,
+                            user_id=user_id_from_instance_tags(i.tags),
+                            wallet_id=wallet_id_from_instance_tags(i.tags),
+                            is_manager=False,
+                        ),
+                        additional_custom_tags={
+                            USER_ID_TAG_KEY: i.tags[USER_ID_TAG_KEY],
+                            WALLET_ID_TAG_KEY: i.tags[WALLET_ID_TAG_KEY],
+                            ROLE_TAG_KEY: WORKER_ROLE_TAG_VALUE,
+                        },
                     ),
-                    additional_custom_tags={
-                        AWSTagKey("user_id"): AWSTagValue(f"{user_id}"),
-                        AWSTagKey("wallet_id"): AWSTagValue(f"{wallet_id}"),
-                        AWSTagKey("role"): AWSTagValue("worker"),
-                    },
-                ),
-                command_name=DOCKER_STACK_DEPLOY_COMMAND_NAME,
-            )
+                    command_name=DOCKER_STACK_DEPLOY_COMMAND_NAME,
+                )
             await ec2_client.set_instances_tags(
                 started_instances_ready_for_command,
                 tags={

services/clusters-keeper/src/simcore_service_clusters_keeper/modules/ssm.py

@@ -21,7 +21,7 @@ async def on_startup() -> None:
         app.state.ssm_client = None
         settings: SSMSettings | None = get_application_settings(
             app
-        ).AUTOSCALING_SSM_ACCESS
+        ).CLUSTERS_KEEPER_SSM_ACCESS
 
         if not settings:
             _logger.warning("SSM client is de-activated in the settings")

services/clusters-keeper/src/simcore_service_clusters_keeper/utils/ec2.py

@@ -7,6 +7,12 @@
 from pydantic import parse_obj_as
 
 from .._meta import VERSION
+from ..constants import (
+    MANAGER_ROLE_TAG_VALUE,
+    ROLE_TAG_KEY,
+    USER_ID_TAG_KEY,
+    WALLET_ID_TAG_KEY,
+)
 from ..core.settings import ApplicationSettings
 
 _APPLICATION_TAG_KEY: Final[str] = "io.simcore.clusters-keeper"
@@ -50,9 +56,9 @@ def creation_ec2_tags(
                     app_settings, user_id=user_id, wallet_id=wallet_id, is_manager=True
                 )
             ),
-            AWSTagKey("user_id"): AWSTagValue(f"{user_id}"),
-            AWSTagKey("wallet_id"): AWSTagValue(f"{wallet_id}"),
-            AWSTagKey("role"): AWSTagValue("manager"),
+            USER_ID_TAG_KEY: AWSTagValue(f"{user_id}"),
+            WALLET_ID_TAG_KEY: AWSTagValue(f"{wallet_id}"),
+            ROLE_TAG_KEY: MANAGER_ROLE_TAG_VALUE,
         }
         | app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_CUSTOM_TAGS
     )
@@ -67,8 +73,8 @@ def ec2_instances_for_user_wallet_filter(
 ) -> EC2Tags:
     return (
         _minimal_identification_tag(app_settings)
-        | {AWSTagKey("user_id"): AWSTagValue(f"{user_id}")}
-        | {AWSTagKey("wallet_id"): AWSTagValue(f"{wallet_id}")}
+        | {USER_ID_TAG_KEY: AWSTagValue(f"{user_id}")}
+        | {WALLET_ID_TAG_KEY: AWSTagValue(f"{wallet_id}")}
     )
 
 
@@ -81,3 +87,14 @@ def compose_user_data(bash_command: str) -> str:
 echo "completed user data bash script"
 """
     )
+
+
+def wallet_id_from_instance_tags(tags: EC2Tags) -> WalletID | None:
+    wallet_id_str = tags[WALLET_ID_TAG_KEY]
+    if wallet_id_str == "None":
+        return None
+    return WalletID(wallet_id_str)
+
+
+def user_id_from_instance_tags(tags: EC2Tags) -> UserID:
+    return UserID(tags[USER_ID_TAG_KEY])