added ssm client, use ssm client, on the way to deploying stack via ssm

Author: sanderegg

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

@@ -0,0 +1,9 @@
+from typing import Final
+
+from aws_library.ec2._models import AWSTagKey
+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"
+)

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

@@ -18,6 +18,7 @@
 from ..modules.ec2 import setup as setup_ec2
 from ..modules.rabbitmq import setup as setup_rabbitmq
 from ..modules.redis import setup as setup_redis
+from ..modules.ssm import setup as setup_ssm
 from ..rpc.rpc_routes import setup_rpc_routes
 from .settings import ApplicationSettings
 
@@ -48,6 +49,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     setup_rabbitmq(app)
     setup_rpc_routes(app)
     setup_ec2(app)
+    setup_ssm(app)
     setup_redis(app)
     setup_clusters_management(app)
 

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

@@ -25,6 +25,7 @@
 from settings_library.ec2 import EC2Settings
 from settings_library.rabbit import RabbitSettings
 from settings_library.redis import RedisSettings
+from settings_library.ssm import SSMSettings
 from settings_library.utils_logging import MixinLoggingSettings
 from types_aiobotocore_ec2.literals import InstanceTypeType
 
@@ -49,6 +50,19 @@ 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,
+        }
+
+
 class WorkersEC2InstancesSettings(BaseCustomSettings):
     WORKERS_EC2_INSTANCES_ALLOWED_TYPES: dict[str, EC2InstanceBootSpecific] = Field(
         ...,
@@ -182,6 +196,12 @@ class PrimaryEC2InstancesSettings(BaseCustomSettings):
         "that take longer than this time will be terminated as sometimes it happens that EC2 machine fail on start.",
     )
 
+    PRIMARY_EC2_INSTANCES_DOCKER_DEFAULT_ADDRESS_POOL: str = Field(
+        default="172.20.0.0/14",
+        description="defines the docker swarm default address pool in CIDR format "
+        "(see https://docs.docker.com/reference/cli/docker/swarm/init/)",
+    )
+
     @validator("PRIMARY_EC2_INSTANCES_ALLOWED_TYPES")
     @classmethod
     def check_valid_instance_names(
@@ -249,6 +269,10 @@ class ApplicationSettings(BaseCustomSettings, MixinLoggingSettings):
         auto_default_from_env=True
     )
 
+    CLUSTERS_KEEPER_SSM_ACCESS: ClustersKeeperSSMSettings | None = Field(
+        auto_default_from_env=True
+    )
+
     CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES: PrimaryEC2InstancesSettings | None = Field(
         auto_default_from_env=True
     )

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

@@ -72,15 +72,7 @@ async def create_cluster(
         tags=creation_ec2_tags(app_settings, user_id=user_id, wallet_id=wallet_id),
         startup_script=create_startup_script(
             app_settings,
-            cluster_machines_name_prefix=get_cluster_name(
-                app_settings, user_id=user_id, wallet_id=wallet_id, is_manager=False
-            ),
             ec2_boot_specific=ec2_instance_boot_specs,
-            additional_custom_tags={
-                AWSTagKey("user_id"): AWSTagValue(f"{user_id}"),
-                AWSTagKey("wallet_id"): AWSTagValue(f"{wallet_id}"),
-                AWSTagKey("role"): AWSTagValue("worker"),
-            },
         ),
         ami_id=ec2_instance_boot_specs.ami_id,
         key_name=app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_KEY_NAME,

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

@@ -5,22 +5,31 @@
 
 import arrow
 from aws_library.ec2 import AWSTagKey, EC2InstanceData
+from aws_library.ec2._models import AWSTagValue
 from fastapi import FastAPI
 from models_library.users import UserID
 from models_library.wallets import WalletID
 from pydantic import parse_obj_as
 from servicelib.logging_utils import log_catch
+from servicelib.utils import limited_gather
 
+from ..constants import (
+    DOCKER_STACK_DEPLOY_COMMAND_EC2_TAG_KEY,
+    DOCKER_STACK_DEPLOY_COMMAND_NAME,
+)
 from ..core.settings import get_application_settings
 from ..modules.clusters import (
     delete_clusters,
     get_all_clusters,
     get_cluster_workers,
     set_instance_heartbeat,
 )
+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
+from ..utils.ec2 import HEARTBEAT_TAG_KEY, get_cluster_name
 from .dask import is_scheduler_busy, ping_scheduler
+from .ec2 import get_ec2_client
+from .ssm import get_ssm_client
 
 _logger = logging.getLogger(__name__)
 
@@ -157,6 +166,62 @@ async def check_clusters(app: FastAPI) -> None:
     # we send a command that contain:
     # the docker-compose file in binary,
     # the call to init the docker swarm and the call to deploy the stack
+    instances_in_need_of_deployment = {
+        i
+        for i in starting_instances - terminateable_instances
+        if DOCKER_STACK_DEPLOY_COMMAND_EC2_TAG_KEY not in i.tags
+    }
+
+    if instances_in_need_of_deployment:
+        app_settings = get_application_settings(app)
+        ssm_client = get_ssm_client(app)
+        ec2_client = get_ec2_client(app)
+        instances_in_need_of_deployment_ssm_connection_state = await limited_gather(
+            *[
+                ssm_client.is_instance_connected_to_ssm_server(i.id)
+                for i in instances_in_need_of_deployment
+            ],
+            reraise=False,
+            log=_logger,
+            limit=20,
+        )
+        ec2_connected_to_ssm_server = [
+            i
+            for i, c in zip(
+                instances_in_need_of_deployment,
+                instances_in_need_of_deployment_ssm_connection_state,
+                strict=True,
+            )
+            if c is True
+        ]
+        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(
+                        app_settings,
+                        user_id=user_id,
+                        wallet_id=wallet_id,
+                        is_manager=False,
+                    ),
+                    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,
+            )
+            await ec2_client.set_instances_tags(
+                started_instances_ready_for_command,
+                tags={
+                    DOCKER_STACK_DEPLOY_COMMAND_EC2_TAG_KEY: AWSTagValue(
+                        ssm_command.command_id
+                    ),
+                },
+            )
 
     # the remaining instances are broken (they were at some point connected but now not anymore)
     broken_instances = disconnected_instances - starting_instances

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

@@ -47,6 +47,7 @@ def setup(app: FastAPI):
         for s in [
             app_settings.CLUSTERS_KEEPER_EC2_ACCESS,
             app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES,
+            app_settings.CLUSTERS_KEEPER_SSM_ACCESS,
         ]
     ):
         logger.warning(

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

@@ -0,0 +1,56 @@
+import logging
+from typing import cast
+
+from aws_library.ssm import SimcoreSSMAPI
+from aws_library.ssm._errors import SSMNotConnectedError
+from fastapi import FastAPI
+from settings_library.ssm import SSMSettings
+from tenacity.asyncio import AsyncRetrying
+from tenacity.before_sleep import before_sleep_log
+from tenacity.stop import stop_after_delay
+from tenacity.wait import wait_random_exponential
+
+from ..core.errors import ConfigurationError
+from ..core.settings import get_application_settings
+
+_logger = logging.getLogger(__name__)
+
+
+def setup(app: FastAPI) -> None:
+    async def on_startup() -> None:
+        app.state.ssm_client = None
+        settings: SSMSettings | None = get_application_settings(
+            app
+        ).AUTOSCALING_SSM_ACCESS
+
+        if not settings:
+            _logger.warning("SSM client is de-activated in the settings")
+            return
+
+        app.state.ssm_client = client = await SimcoreSSMAPI.create(settings)
+
+        async for attempt in AsyncRetrying(
+            reraise=True,
+            stop=stop_after_delay(120),
+            wait=wait_random_exponential(max=30),
+            before_sleep=before_sleep_log(_logger, logging.WARNING),
+        ):
+            with attempt:
+                connected = await client.ping()
+                if not connected:
+                    raise SSMNotConnectedError  # pragma: no cover
+
+    async def on_shutdown() -> None:
+        if app.state.ssm_client:
+            await cast(SimcoreSSMAPI, app.state.ssm_client).close()
+
+    app.add_event_handler("startup", on_startup)
+    app.add_event_handler("shutdown", on_shutdown)
+
+
+def get_ssm_client(app: FastAPI) -> SimcoreSSMAPI:
+    if not app.state.ssm_client:
+        raise ConfigurationError(
+            msg="SSM client is not available. Please check the configuration."
+        )
+    return cast(SimcoreSSMAPI, app.state.ssm_client)

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

@@ -8,6 +8,7 @@
 import arrow
 import yaml
 from aws_library.ec2 import EC2InstanceBootSpecific, EC2InstanceData, EC2Tags
+from aws_library.ec2._models import CommandStr
 from fastapi.encoders import jsonable_encoder
 from models_library.api_schemas_clusters_keeper.clusters import (
     ClusterState,
@@ -107,47 +108,55 @@ def _convert_to_env_dict(entries: dict[str, Any]) -> str:
 def create_startup_script(
     app_settings: ApplicationSettings,
     *,
-    cluster_machines_name_prefix: str,
     ec2_boot_specific: EC2InstanceBootSpecific,
-    additional_custom_tags: EC2Tags,
 ) -> str:
     assert app_settings.CLUSTERS_KEEPER_EC2_ACCESS  # nosec
     assert app_settings.CLUSTERS_KEEPER_WORKERS_EC2_INSTANCES  # nosec
 
-    environment_variables = _prepare_environment_variables(
-        app_settings,
-        cluster_machines_name_prefix=cluster_machines_name_prefix,
-        additional_custom_tags=additional_custom_tags,
-    )
-
     startup_commands = ec2_boot_specific.custom_boot_scripts.copy()
+    return "\n".join(startup_commands)
+
+
+def create_deploy_cluster_stack_script(
+    app_settings: ApplicationSettings,
+    *,
+    cluster_machines_name_prefix: str,
+    additional_custom_tags: EC2Tags,
+) -> str:
+    deploy_script: list[CommandStr] = []
     assert app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES  # nosec
     if isinstance(
         app_settings.CLUSTERS_KEEPER_COMPUTATIONAL_BACKEND_DEFAULT_CLUSTER_AUTH,
         TLSAuthentication,
     ):
-
+        # get the dask certificates
         download_certificates_commands = [
             f"mkdir --parents {_HOST_CERTIFICATES_BASE_PATH}",
             f'aws ssm get-parameter --name "{app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_SSM_TLS_DASK_CA}" --region us-east-1 --with-decryption --query "Parameter.Value" --output text > {_HOST_TLS_CA_FILE_PATH}',
             f'aws ssm get-parameter --name "{app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_SSM_TLS_DASK_CERT}" --region us-east-1 --with-decryption --query "Parameter.Value" --output text > {_HOST_TLS_CERT_FILE_PATH}',
             f'aws ssm get-parameter --name "{app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_SSM_TLS_DASK_KEY}" --region us-east-1 --with-decryption --query "Parameter.Value" --output text > {_HOST_TLS_KEY_FILE_PATH}',
         ]
-        startup_commands.extend(download_certificates_commands)
+        deploy_script.extend(download_certificates_commands)
+
+    environment_variables = _prepare_environment_variables(
+        app_settings,
+        cluster_machines_name_prefix=cluster_machines_name_prefix,
+        additional_custom_tags=additional_custom_tags,
+    )
 
-    startup_commands.extend(
+    deploy_script.extend(
         [
             # NOTE: https://stackoverflow.com/questions/41203492/solving-redis-warnings-on-overcommit-memory-and-transparent-huge-pages-for-ubunt
             "sysctl vm.overcommit_memory=1",
             f"echo '{_docker_compose_yml_base64_encoded()}' | base64 -d > {_HOST_DOCKER_COMPOSE_PATH}",
             f"echo '{_prometheus_yml_base64_encoded()}' | base64 -d > {_HOST_PROMETHEUS_PATH}",
             f"echo '{_prometheus_basic_auth_yml_base64_encoded(app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_PROMETHEUS_USERNAME, app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_PROMETHEUS_PASSWORD.get_secret_value())}' | base64 -d > {_HOST_PROMETHEUS_WEB_PATH}",
             # NOTE: --default-addr-pool is necessary in order to prevent conflicts with AWS node IPs
-            "docker swarm init --default-addr-pool 172.20.0.0/14",
+            f"docker swarm init --default-addr-pool {app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_DOCKER_DEFAULT_ADDRESS_POOL}",
             f"{' '.join(environment_variables)} docker stack deploy --with-registry-auth --compose-file={_HOST_DOCKER_COMPOSE_PATH} dask_stack",
         ]
     )
-    return "\n".join(startup_commands)
+    return "\n".join(deploy_script)
 
 
 def _convert_ec2_state_to_cluster_state(

services/docker-compose.yml

@@ -198,6 +198,8 @@ services:
       PRIMARY_EC2_INSTANCES_SSM_TLS_DASK_KEY: ${PRIMARY_EC2_INSTANCES_SSM_TLS_DASK_KEY}
       PRIMARY_EC2_INSTANCES_PROMETHEUS_USERNAME: ${PRIMARY_EC2_INSTANCES_PROMETHEUS_USERNAME}
       PRIMARY_EC2_INSTANCES_PROMETHEUS_PASSWORD: ${PRIMARY_EC2_INSTANCES_PROMETHEUS_PASSWORD}
+      PRIMARY_EC2_INSTANCES_MAX_START_TIME: ${PRIMARY_EC2_INSTANCES_MAX_START_TIME}
+      PRIMARY_EC2_INSTANCES_DOCKER_DEFAULT_ADDRESS_POOL: ${PRIMARY_EC2_INSTANCES_DOCKER_DEFAULT_ADDRESS_POOL}
       RABBIT_HOST: ${RABBIT_HOST}
       RABBIT_PASSWORD: ${RABBIT_PASSWORD}
       RABBIT_PORT: ${RABBIT_PORT}