pass the wallet ID

ITISFoundation · Oct 26, 2023 · 2d6e0fb · 2d6e0fb
1 parent ebed946
commit 2d6e0fb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/services/director-v2/src/simcore_service_director_v2/modules/clusters_keeper.py b/services/director-v2/src/simcore_service_director_v2/modules/clusters_keeper.py
@@ -6,6 +6,7 @@
     OnDemandCluster,
 )
 from models_library.users import UserID
+from models_library.wallets import WalletID
 from servicelib.rabbitmq import (
     RabbitMQRPCClient,
     RemoteMethodNotRegisteredError,
@@ -24,15 +25,18 @@
 
 
 async def get_or_create_on_demand_cluster(
-    user_id: UserID, rabbitmq_rpc_client: RabbitMQRPCClient
+    rabbitmq_rpc_client: RabbitMQRPCClient,
+    *,
+    user_id: UserID,
+    wallet_id: WalletID | None,
 ) -> BaseCluster:
     try:
         returned_cluster: OnDemandCluster = await rabbitmq_rpc_client.request(
             RPCNamespace("clusters-keeper"),
             RPCMethodName("get_or_create_cluster"),
             timeout_s=300,
             user_id=user_id,
-            wallet_id=None,  # NOTE: --> MD this will need to be replaced by the real walletID
+            wallet_id=wallet_id,
         )
         _logger.info("received cluster: %s", returned_cluster)
         if returned_cluster.state is not ClusterState.RUNNING:

diff --git a/...ices/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/dask_scheduler.py b/...ices/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/dask_scheduler.py
@@ -75,7 +75,9 @@ async def _cluster_dask_client(
     cluster: BaseCluster = scheduler.settings.default_cluster
     if pipeline_params.use_on_demand_clusters:
         cluster = await get_or_create_on_demand_cluster(
-            user_id, scheduler.rabbitmq_rpc_client
+            scheduler.rabbitmq_rpc_client,
+            user_id=user_id,
+            wallet_id=pipeline_params.run_metadata.get("wallet_id"),
         )
     if pipeline_params.cluster_id != DEFAULT_CLUSTER_ID:
         clusters_repo = ClustersRepository.instance(scheduler.db_engine)