✨Comp backend: Clusters-keeper can create a cluster with no wallet ID defined (#4746)

Author: sanderegg

packages/models-library/src/models_library/rpc_schemas_clusters_keeper/clusters.py

@@ -20,6 +20,6 @@ class OnDemandCluster(BaseModel):
     authentication: ClusterAuthentication
     state: ClusterState
     user_id: UserID
-    wallet_id: WalletID
+    wallet_id: WalletID | None
     gateway_ready: bool
     eta: datetime.timedelta

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

@@ -24,7 +24,7 @@
 
 
 async def create_cluster(
-    app: FastAPI, *, user_id: UserID, wallet_id: WalletID
+    app: FastAPI, *, user_id: UserID, wallet_id: WalletID | None
 ) -> list[EC2InstanceData]:
     ec2_client = get_ec2_client(app)
     app_settings = get_application_settings(app)
@@ -56,7 +56,7 @@ async def get_all_clusters(app: FastAPI) -> list[EC2InstanceData]:
 
 
 async def get_cluster(
-    app: FastAPI, *, user_id: UserID, wallet_id: WalletID
+    app: FastAPI, *, user_id: UserID, wallet_id: WalletID | None
 ) -> EC2InstanceData:
     app_settings = get_application_settings(app)
     assert app_settings.CLUSTERS_KEEPER_EC2_INSTANCES  # nosec
@@ -70,7 +70,7 @@ async def get_cluster(
 
 
 async def cluster_heartbeat(
-    app: FastAPI, *, user_id: UserID, wallet_id: WalletID
+    app: FastAPI, *, user_id: UserID, wallet_id: WalletID | None
 ) -> None:
     app_settings = get_application_settings(app)
     assert app_settings.CLUSTERS_KEEPER_EC2_INSTANCES  # nosec

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

@@ -17,7 +17,7 @@
 
 @router.expose()
 async def get_or_create_cluster(
-    app: FastAPI, *, user_id: UserID, wallet_id: WalletID
+    app: FastAPI, *, user_id: UserID, wallet_id: WalletID | None
 ) -> OnDemandCluster:
     """Get or create cluster for user_id and wallet_id
     This function will create a new instance on AWS if needed or return the already running one.

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

@@ -66,7 +66,7 @@ def _create_eta(
 def create_cluster_from_ec2_instance(
     instance: EC2InstanceData,
     user_id: UserID,
-    wallet_id: WalletID,
+    wallet_id: WalletID | None,
     gateway_password: SecretStr,
     *,
     gateway_ready: bool,

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

@@ -17,7 +17,7 @@
 
 
 def creation_ec2_tags(
-    app_settings: ApplicationSettings, *, user_id: UserID, wallet_id: WalletID
+    app_settings: ApplicationSettings, *, user_id: UserID, wallet_id: WalletID | None
 ) -> EC2Tags:
     assert app_settings.CLUSTERS_KEEPER_EC2_INSTANCES  # nosec
     return _DEFAULT_CLUSTERS_KEEPER_TAGS | {
@@ -38,7 +38,7 @@ def get_user_id_from_tags(tags: EC2Tags) -> UserID:
 
 
 def ec2_instances_for_user_wallet_filter(
-    user_id: UserID, wallet_id: WalletID
+    user_id: UserID, wallet_id: WalletID | None
 ) -> EC2Tags:
     return (
         _DEFAULT_CLUSTERS_KEEPER_TAGS

services/clusters-keeper/tests/unit/test_rpc_clusters.py

@@ -71,7 +71,7 @@ def _base_configuration(
 async def _assert_cluster_instance_created(
     ec2_client: EC2Client,
     user_id: UserID,
-    wallet_id: WalletID,
+    wallet_id: WalletID | None,
 ) -> None:
     instances = await ec2_client.describe_instances()
     assert len(instances["Reservations"]) == 1
@@ -92,11 +92,16 @@ async def _assert_cluster_instance_created(
     assert instances["Reservations"][0]["Instances"][0]["Tags"][1]["Key"] == "Name"
     assert "Value" in instances["Reservations"][0]["Instances"][0]["Tags"][1]
     instance_name = instances["Reservations"][0]["Instances"][0]["Tags"][1]["Value"]
-
-    parse_result = search("user_id:{user_id:d}-wallet_id:{wallet_id:d}", instance_name)
+    search_str = (
+        "user_id:{user_id:d}-wallet_id:{wallet_id:d}"
+        if wallet_id
+        else "user_id:{user_id:d}-wallet_id:None"
+    )
+    parse_result = search(search_str, instance_name)
     assert isinstance(parse_result, Result)
     assert parse_result["user_id"] == user_id
-    assert parse_result["wallet_id"] == wallet_id
+    if wallet_id:
+        assert parse_result["wallet_id"] == wallet_id
 
 
 async def _assert_cluster_heartbeat_on_instance(
@@ -136,26 +141,30 @@ def mocked_dask_ping_gateway(mocker: MockerFixture) -> MockedDaskModule:
     )
 
 
+@pytest.mark.parametrize("use_wallet_id", [True, False])
 async def test_get_or_create_cluster(
     _base_configuration: None,
     clusters_keeper_rabbitmq_rpc_client: RabbitMQRPCClient,
     ec2_client: EC2Client,
     user_id: UserID,
     wallet_id: WalletID,
+    use_wallet_id: bool,
     mocked_dask_ping_gateway: MockedDaskModule,
 ):
     # send rabbitmq rpc to create_cluster
     rpc_response = await clusters_keeper_rabbitmq_rpc_client.request(
         CLUSTERS_KEEPER_NAMESPACE,
         RPCMethodName("get_or_create_cluster"),
         user_id=user_id,
-        wallet_id=wallet_id,
+        wallet_id=wallet_id if use_wallet_id else None,
     )
     assert rpc_response
     assert isinstance(rpc_response, OnDemandCluster)
     created_cluster = rpc_response
     # check we do have a new machine in AWS
-    await _assert_cluster_instance_created(ec2_client, user_id, wallet_id)
+    await _assert_cluster_instance_created(
+        ec2_client, user_id, wallet_id if use_wallet_id else None
+    )
     # it is called once as moto server creates instances instantly
     mocked_dask_ping_gateway.ping_gateway.assert_called_once()
     mocked_dask_ping_gateway.ping_gateway.reset_mock()
@@ -165,7 +174,7 @@ async def test_get_or_create_cluster(
         CLUSTERS_KEEPER_NAMESPACE,
         RPCMethodName("get_or_create_cluster"),
         user_id=user_id,
-        wallet_id=wallet_id,
+        wallet_id=wallet_id if use_wallet_id else None,
     )
     assert rpc_response
     assert isinstance(rpc_response, OnDemandCluster)

services/director-v2/src/simcore_service_director_v2/modules/clusters_keeper.py

@@ -1,14 +1,12 @@
 import datetime
 import logging
-from typing import Final
 
 from models_library.clusters import BaseCluster, ClusterTypeInModel
 from models_library.rpc_schemas_clusters_keeper.clusters import (
     ClusterState,
     OnDemandCluster,
 )
 from models_library.users import UserID
-from models_library.wallets import WalletID
 from servicelib.rabbitmq import (
     RabbitMQRPCClient,
     RemoteMethodNotRegisteredError,
@@ -23,7 +21,6 @@
 
 _logger = logging.getLogger(__name__)
 
-_TEMPORARY_DEFAULT_WALLET_ID: Final[WalletID] = 43
 _TIME_FORMAT = "{:02d}:{:02d}"  # format for minutes:seconds
 
 
@@ -40,7 +37,7 @@ async def get_or_create_on_demand_cluster(
             RPCMethodName("get_or_create_cluster"),
             timeout_s=300,
             user_id=user_id,
-            wallet_id=_TEMPORARY_DEFAULT_WALLET_ID,
+            wallet_id=None,  # NOTE: --> MD this will need to be replaced by the real walletID
         )
         _logger.info("received cluster: %s", returned_cluster)
         if returned_cluster.state is not ClusterState.RUNNING: