✨Autoscaling monitoring tool: make compatible with non-billable deploys (#6285)

Author: sanderegg

.pre-commit-config.yaml

@@ -28,6 +28,7 @@ repos:
         args:
           - "--py311-plus"
         name: upgrade code
+        exclude: ^scripts/maintenance/computational-clusters/autoscaled_monitor/cli\.py$ # Optional get replaced and typer does not like it
   - repo: https://github.com/hadialqattan/pycln
     rev: v2.1.4
     hooks:

scripts/maintenance/computational-clusters/autoscaled_monitor/cli.py

@@ -1,6 +1,6 @@
 import asyncio
 from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Optional
 
 import parse
 import rich
@@ -10,15 +10,22 @@
 from . import core as api
 from .constants import (
     DEFAULT_COMPUTATIONAL_EC2_FORMAT,
+    DEFAULT_COMPUTATIONAL_EC2_FORMAT_WORKERS,
     DEFAULT_DYNAMIC_EC2_FORMAT,
     DEPLOY_SSH_KEY_PARSER,
+    wallet_id_spec,
 )
 from .ec2 import autoscaling_ec2_client, cluster_keeper_ec2_client
 from .models import AppState
 
 state: AppState = AppState(
     dynamic_parser=parse.compile(DEFAULT_DYNAMIC_EC2_FORMAT),
-    computational_parser=parse.compile(DEFAULT_COMPUTATIONAL_EC2_FORMAT),
+    computational_parser_primary=parse.compile(
+        DEFAULT_COMPUTATIONAL_EC2_FORMAT, {"wallet_id_spec": wallet_id_spec}
+    ),
+    computational_parser_workers=parse.compile(
+        DEFAULT_COMPUTATIONAL_EC2_FORMAT_WORKERS, {"wallet_id_spec": wallet_id_spec}
+    ),
 )
 
 app = typer.Typer()
@@ -68,13 +75,24 @@ def main(
     state.ec2_resource_clusters_keeper = cluster_keeper_ec2_client(state)
 
     assert state.environment["EC2_INSTANCES_KEY_NAME"]
-    state.dynamic_parser = parse.compile(
-        f"{state.environment['EC2_INSTANCES_NAME_PREFIX']}-{{key_name}}"
-    )
+    dynamic_pattern = f"{state.environment['EC2_INSTANCES_NAME_PREFIX']}-{{key_name}}"
+    state.dynamic_parser = parse.compile(dynamic_pattern)
+    rich.print(f"using dynamic-naming-regex: {dynamic_pattern}")
     if state.environment["CLUSTERS_KEEPER_EC2_INSTANCES_PREFIX"]:
-        state.computational_parser = parse.compile(
-            f"{state.environment['CLUSTERS_KEEPER_EC2_INSTANCES_PREFIX']}-{DEFAULT_COMPUTATIONAL_EC2_FORMAT}"
+        state.computational_parser_primary = parse.compile(
+            f"{state.environment['CLUSTERS_KEEPER_EC2_INSTANCES_PREFIX'].strip('-')}-{DEFAULT_COMPUTATIONAL_EC2_FORMAT}",
+            {"wallet_id_spec", wallet_id_spec},
+        )
+        state.computational_parser_workers = parse.compile(
+            f"{state.environment['CLUSTERS_KEEPER_EC2_INSTANCES_PREFIX'].strip('-')}-{DEFAULT_COMPUTATIONAL_EC2_FORMAT_WORKERS}",
+            {"wallet_id_spec", wallet_id_spec},
         )
+    rich.print(
+        f"compuational-primary-naming-regex: {state.computational_parser_primary._expression}"  # noqa: SLF001
+    )
+    rich.print(
+        f"compuational-workers-naming-regex: {state.computational_parser_workers._expression}"  # noqa: SLF001
+    )
 
     # locate ssh key path
     for file_path in deploy_config.glob("**/*.pem"):
@@ -117,7 +135,10 @@ def summary(
 @app.command()
 def cancel_jobs(
     user_id: Annotated[int, typer.Option(help="the user ID")],
-    wallet_id: Annotated[int, typer.Option(help="the wallet ID")],
+    wallet_id: Annotated[
+        Optional[int | None],  # noqa: UP007 # typer does not understand | syntax
+        typer.Option(help="the wallet ID"),
+    ] = None,
     *,
     force: Annotated[
         bool,

scripts/maintenance/computational-clusters/autoscaled_monitor/constants.py

@@ -4,9 +4,20 @@
 import parse
 from pydantic import ByteSize
 
+
+@parse.with_pattern(r"None|\d+")
+def wallet_id_spec(text) -> None | int:
+    if text == "None":
+        return None
+    return int(text)
+
+
 DEFAULT_COMPUTATIONAL_EC2_FORMAT: Final[
     str
-] = r"osparc-computational-cluster-{role}-{swarm_stack_name}-user_id:{user_id:d}-wallet_id:{wallet_id:d}"
+] = r"osparc-computational-cluster-{role}-{swarm_stack_name}-user_id:{user_id:d}-wallet_id:{wallet_id:wallet_id_spec}"
+DEFAULT_COMPUTATIONAL_EC2_FORMAT_WORKERS: Final[
+    str
+] = r"osparc-computational-cluster-{role}-{swarm_stack_name}-user_id:{user_id:d}-wallet_id:{wallet_id:wallet_id_spec}-{key_name}"
 DEFAULT_DYNAMIC_EC2_FORMAT: Final[str] = r"osparc-dynamic-autoscaled-worker-{key_name}"
 DEPLOY_SSH_KEY_PARSER: Final[parse.Parser] = parse.compile(r"osparc-{random_name}.pem")
 

scripts/maintenance/computational-clusters/autoscaled_monitor/core.py

@@ -38,8 +38,14 @@ def _parse_computational(
     state: AppState, instance: Instance
 ) -> ComputationalInstance | None:
     name = utils.get_instance_name(instance)
-    if result := state.computational_parser.search(name):
+    if result := (
+        state.computational_parser_workers.parse(name)
+        or state.computational_parser_primary.parse(name)
+    ):
         assert isinstance(result, parse.Result)
+        # special handling for optional wallet
+        rich.print(result.named)
+
         last_heartbeat = utils.get_last_heartbeat(instance)
         return ComputationalInstance(
             role=InstanceRole(result["role"]),
@@ -538,12 +544,13 @@ async def cancel_jobs(  # noqa: C901, PLR0912
     rich.print(the_cluster.datasets)
     try:
         if response := typer.prompt(
-            "[yellow]Which dataset to cancel? all for all of them.[/yellow]",
+            "Which dataset to cancel? (all: will cancel everything, 1-5: will cancel jobs 1-5, or 4: will cancel job #4)",
             default="none",
         ):
             if response == "none":
                 rich.print("[yellow]not cancelling anything[/yellow]")
             elif response == "all":
+                rich.print("cancelling all tasks")
                 for comp_task, dask_task in task_to_dask_job:
                     if dask_task is not None and dask_task.state != "unknown":
                         await dask.trigger_job_cancellation_in_scheduler(
@@ -565,25 +572,43 @@ async def cancel_jobs(  # noqa: C901, PLR0912
 
                 rich.print("cancelled all tasks")
             else:
-                selected_index = TypeAdapter(int).validate_python(response)
-                comp_task, dask_task = task_to_dask_job[selected_index]
-                if dask_task is not None and dask_task.state != "unknown":
-                    await dask.trigger_job_cancellation_in_scheduler(
-                        state, the_cluster, dask_task.job_id
-                    )
-                    if comp_task is None:
-                        # we need to clear it of the cluster
-                        await dask.remove_job_from_scheduler(
-                            state, the_cluster, dask_task.job_id
-                        )
+                try:
+                    # Split the response and handle ranges
+                    indices = response.split("-")
+                    if len(indices) == 2:
+                        start_index, end_index = map(int, indices)
+                        selected_indices = range(start_index, end_index + 1)
+                    else:
+                        selected_indices = [int(indices[0])]
+
+                    for selected_index in selected_indices:
+                        comp_task, dask_task = task_to_dask_job[selected_index]
+                        if dask_task is not None and dask_task.state != "unknown":
+                            await dask.trigger_job_cancellation_in_scheduler(
+                                state, the_cluster, dask_task.job_id
+                            )
+                            if comp_task is None:
+                                # we need to clear it of the cluster
+                                await dask.remove_job_from_scheduler(
+                                    state, the_cluster, dask_task.job_id
+                                )
+
+                        if comp_task is not None and force:
+                            await db.abort_job_in_db(
+                                state, comp_task.project_id, comp_task.node_id
+                            )
+                    rich.print(f"Cancelled selected tasks: {response}")
 
-                if comp_task is not None and force:
-                    await db.abort_job_in_db(
-                        state, comp_task.project_id, comp_task.node_id
+                except ValidationError:
+                    rich.print(
+                        "[yellow]wrong index format, not cancelling anything[/yellow]"
+                    )
+                except IndexError:
+                    rich.print(
+                        "[yellow]index out of range, not cancelling anything[/yellow]"
                     )
-
     except ValidationError:
-        rich.print("[yellow]wrong index, not cancelling anything[/yellow]")
+        rich.print("[yellow]wrong input, not cancelling anything[/yellow]")
 
 
 async def trigger_cluster_termination(

scripts/maintenance/computational-clusters/autoscaled_monitor/models.py

@@ -18,7 +18,8 @@ class AppState:
     ec2_resource_autoscaling: EC2ServiceResource | None = None
     ec2_resource_clusters_keeper: EC2ServiceResource | None = None
     dynamic_parser: parse.Parser
-    computational_parser: parse.Parser
+    computational_parser_primary: parse.Parser
+    computational_parser_workers: parse.Parser
     deploy_config: Path | None = None
     ssh_key_path: Path | None = None