implemented compute cluster total resources

Author: sanderegg

services/autoscaling/src/simcore_service_autoscaling/modules/cluster_scaling/_provider_computational.py

@@ -155,7 +155,9 @@ async def compute_cluster_total_resources(
         assert self  # nosec
         try:
             return await dask.compute_cluster_total_resources(
-                _scheduler_url(app), _scheduler_auth(app), instances
+                _scheduler_url(app),
+                _scheduler_auth(app),
+                [i.ec2_instance for i in instances],
             )
         except DaskNoWorkersError:
             return Resources.create_as_empty()

services/autoscaling/src/simcore_service_autoscaling/modules/dask.py

@@ -21,7 +21,7 @@
     DaskWorkerNotFoundError,
 )
 from ..core.settings import DaskMonitoringSettings
-from ..models import AssociatedInstance, DaskTask, DaskTaskId
+from ..models import DaskTask, DaskTaskId
 from ..utils.utils_ec2 import (
     node_host_name_from_ec2_private_dns,
     node_ip_from_ec2_private_dns,
@@ -306,23 +306,31 @@ def _list_processing_tasks_on_worker(
 async def compute_cluster_total_resources(
     scheduler_url: AnyUrl,
     authentication: ClusterAuthentication,
-    instances: list[AssociatedInstance],
+    instances: list[EC2InstanceData],
 ) -> Resources:
     if not instances:
         return Resources.create_as_empty()
     async with _scheduler_client(scheduler_url, authentication) as client:
-        instance_hosts = (
-            node_ip_from_ec2_private_dns(i.ec2_instance) for i in instances
-        )
+        instance_host_resources_map = {
+            node_ip_from_ec2_private_dns(i): i.resources for i in instances
+        }
         scheduler_info = client.scheduler_info()
         if "workers" not in scheduler_info or not scheduler_info["workers"]:
             raise DaskNoWorkersError(url=scheduler_url)
         workers: dict[str, Any] = scheduler_info["workers"]
+        cluster_resources = Resources.create_as_empty()
         for worker_details in workers.values():
-            if worker_details["host"] not in instance_hosts:
+            if worker_details["host"] not in instance_host_resources_map:
                 continue
+            worker_ram = worker_details["memory_limit"]
+            worker_threads = worker_details["nthreads"]
+            cluster_resources += Resources(
+                cpus=instance_host_resources_map[worker_details["host"]].cpus,
+                ram=TypeAdapter(ByteSize).validate_python(worker_ram),
+                generic_resources={DASK_WORKER_THREAD_RESOURCE_NAME: worker_threads},
+            )
 
-        return Resources.create_as_empty()
+        return cluster_resources
 
 
 async def try_retire_nodes(

services/autoscaling/tests/unit/test_modules_dask.py

@@ -380,8 +380,11 @@ def _add_fct(x: int, y: int) -> int:
 
 
 async def test_compute_cluster_total_resources(
+    dask_spec_local_cluster: distributed.SpecCluster,
     scheduler_url: AnyUrl,
     scheduler_authentication: ClusterAuthentication,
+    fake_ec2_instance_data: Callable[..., EC2InstanceData],
+    fake_localhost_ec2_instance_data: EC2InstanceData,
 ):
     # asking for resources of empty cluster returns empty resources
     assert (
@@ -390,6 +393,26 @@ async def test_compute_cluster_total_resources(
         )
         == Resources.create_as_empty()
     )
+    ec2_instance_data = fake_ec2_instance_data()
+    assert ec2_instance_data.resources.cpus > 0
+    assert ec2_instance_data.resources.ram > 0
+    assert ec2_instance_data.resources.generic_resources == {}
+    assert (
+        await compute_cluster_total_resources(
+            scheduler_url, scheduler_authentication, [ec2_instance_data]
+        )
+        == Resources.create_as_empty()
+    ), "this instance is not connected and should not be accounted for"
+
+    cluster_total_resources = await compute_cluster_total_resources(
+        scheduler_url, scheduler_authentication, [fake_localhost_ec2_instance_data]
+    )
+    assert cluster_total_resources.cpus > 0
+    assert cluster_total_resources.ram > 0
+    assert DASK_WORKER_THREAD_RESOURCE_NAME in cluster_total_resources.generic_resources
+    assert (
+        cluster_total_resources.generic_resources[DASK_WORKER_THREAD_RESOURCE_NAME] == 2
+    )
 
 
 @pytest.mark.parametrize(
@@ -454,7 +477,6 @@ async def test_is_worker_connected(
 
 
 async def test_is_worker_retired(
-    dask_spec_local_cluster: distributed.SpecCluster,
     scheduler_url: AnyUrl,
     scheduler_authentication: ClusterAuthentication,
     fake_ec2_instance_data: Callable[..., EC2InstanceData],