🎨 Expose wallet_id to computational and dynamic services via env va…

…rs upon request (#7125) Co-authored-by: Andrei Neagu <[email protected]> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
ITISFoundation · Jan 29, 2025 · 1433e6e · 1433e6e
1 parent 1713021
commit 1433e6e
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 0 deletions.
diff --git a/services/director-v2/src/simcore_service_director_v2/modules/dask_client.py b/services/director-v2/src/simcore_service_director_v2/modules/dask_client.py
@@ -399,6 +399,7 @@ async def send_computation_tasks(
                     node_image=node_image,
                     metadata=metadata,
                     resource_tracking_run_id=resource_tracking_run_id,
+                    wallet_id=metadata.get("wallet_id"),
                 )
                 task_owner = dask_utils.compute_task_owner(
                     user_id, project_id, node_id, metadata.get("project_metadata", {})

diff --git a/...rector-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/docker_compose_specs.py b/...rector-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/docker_compose_specs.py
@@ -23,6 +23,7 @@
 from models_library.services_types import ServiceRunID
 from models_library.users import UserID
 from models_library.utils.docker_compose import replace_env_vars_in_compose_spec
+from models_library.wallets import WalletID
 from pydantic import ByteSize
 from servicelib.resources import CPU_RESOURCE_LIMIT_KEY, MEM_RESOURCE_LIMIT_KEY
 from settings_library.docker_registry import RegistrySettings
@@ -280,6 +281,7 @@ async def assemble_spec(  # pylint: disable=too-many-arguments # noqa: PLR0913
     simcore_user_agent: str,
     swarm_stack_name: str,
     service_run_id: ServiceRunID,
+    wallet_id: WalletID | None,
 ) -> str:
     """
     returns a docker-compose spec used by
@@ -353,6 +355,7 @@ async def assemble_spec(  # pylint: disable=too-many-arguments # noqa: PLR0913
             project_id=project_id,
             node_id=node_id,
             service_run_id=service_run_id,
+            wallet_id=wallet_id,
         )
 
         add_egress_configuration(
@@ -392,6 +395,7 @@ async def assemble_spec(  # pylint: disable=too-many-arguments # noqa: PLR0913
         project_id=project_id,
         node_id=node_id,
         service_run_id=service_run_id,
+        wallet_id=wallet_id,
     )
 
     stringified_service_spec: str = replace_env_vars_in_compose_spec(

diff --git a/...core_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_events_user_services.py b/...core_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_events_user_services.py
@@ -101,6 +101,9 @@ async def submit_compose_sepc(app: FastAPI, scheduler_data: SchedulerData) -> No
         simcore_user_agent=scheduler_data.request_simcore_user_agent,
         swarm_stack_name=dynamic_services_scheduler_settings.SWARM_STACK_NAME,
         service_run_id=scheduler_data.run_id,
+        wallet_id=(
+            scheduler_data.wallet_info.wallet_id if scheduler_data.wallet_info else None
+        ),
     )
 
     _logger.debug(

diff --git a/...ces/director-v2/src/simcore_service_director_v2/modules/osparc_variables/substitutions.py b/...ces/director-v2/src/simcore_service_director_v2/modules/osparc_variables/substitutions.py
@@ -19,6 +19,7 @@
 from models_library.services_types import ServiceRunID
 from models_library.users import UserID
 from models_library.utils.specs_substitution import SpecsSubstitutionsResolver
+from models_library.wallets import WalletID
 from pydantic import BaseModel
 from servicelib.fastapi.app_state import SingletonInAppStateMixin
 from servicelib.logging_utils import log_context
@@ -121,6 +122,7 @@ def create(cls, app: FastAPI):
             ("OSPARC_VARIABLE_PRODUCT_NAME", "product_name"),
             ("OSPARC_VARIABLE_STUDY_UUID", "project_id"),
             ("OSPARC_VARIABLE_SERVICE_RUN_ID", "run_id"),
+            ("OSPARC_VARIABLE_WALLET_ID", "wallet_id"),
             ("OSPARC_VARIABLE_USER_ID", "user_id"),
             ("OSPARC_VARIABLE_API_HOST", "api_server_base_url"),
         ]:
@@ -183,6 +185,7 @@ async def resolve_and_substitute_session_variables_in_model(
     project_id: ProjectID,
     node_id: NodeID,
     service_run_id: ServiceRunID,
+    wallet_id: WalletID | None,
 ) -> TBaseModel:
     result: TBaseModel = model
     try:
@@ -204,6 +207,7 @@ async def resolve_and_substitute_session_variables_in_model(
                 project_id=project_id,
                 node_id=node_id,
                 run_id=service_run_id,
+                wallet_id=wallet_id,
                 api_server_base_url=app_settings.DIRECTOR_V2_PUBLIC_API_BASE_URL,
             ),
         )
@@ -226,6 +230,7 @@ async def resolve_and_substitute_session_variables_in_specs(
     project_id: ProjectID,
     node_id: NodeID,
     service_run_id: ServiceRunID,
+    wallet_id: WalletID | None,
 ) -> dict[str, Any]:
     table = OsparcSessionVariablesTable.get_from_app_state(app)
     resolver = SpecsSubstitutionsResolver(specs, upgrade=False)
@@ -248,6 +253,7 @@ async def resolve_and_substitute_session_variables_in_specs(
                     project_id=project_id,
                     node_id=node_id,
                     run_id=service_run_id,
+                    wallet_id=wallet_id,
                     api_server_base_url=app_settings.DIRECTOR_V2_PUBLIC_API_BASE_URL,
                 ),
             )

diff --git a/services/director-v2/src/simcore_service_director_v2/utils/dask.py b/services/director-v2/src/simcore_service_director_v2/utils/dask.py
@@ -28,6 +28,7 @@
 from models_library.services import ServiceKey, ServiceVersion
 from models_library.services_types import ServiceRunID
 from models_library.users import UserID
+from models_library.wallets import WalletID
 from pydantic import AnyUrl, ByteSize, TypeAdapter, ValidationError
 from servicelib.logging_utils import log_catch, log_context
 from simcore_sdk import node_ports_v2
@@ -344,6 +345,7 @@ async def compute_task_envs(
     node_image: Image,
     metadata: RunMetadataDict,
     resource_tracking_run_id: ServiceRunID,
+    wallet_id: WalletID | None,
 ) -> ContainerEnvsDict:
     product_name = metadata.get("product_name", UNDEFINED_DOCKER_LABEL)
     task_envs = node_image.envs
@@ -363,6 +365,7 @@ async def compute_task_envs(
             project_id=project_id,
             node_id=node_id,
             service_run_id=resource_tracking_run_id,
+            wallet_id=wallet_id,
         )
         # NOTE: see https://github.com/ITISFoundation/osparc-simcore/issues/3638
         # we currently do not validate as we are using illegal docker key names with underscores

diff --git a/services/director-v2/tests/unit/test_modules_osparc_variables.py b/services/director-v2/tests/unit/test_modules_osparc_variables.py
@@ -22,6 +22,7 @@
 from models_library.users import UserID
 from models_library.utils.specs_substitution import SubstitutionValue
 from models_library.utils.string_substitution import OSPARC_IDENTIFIER_PREFIX
+from models_library.wallets import WalletID
 from pydantic import TypeAdapter
 from pytest_mock import MockerFixture
 from pytest_simcore.helpers.faker_compose_specs import generate_fake_docker_compose
@@ -159,11 +160,13 @@ async def fake_app(faker: Faker) -> AsyncIterable[FastAPI]:
         yield app
 
 
+@pytest.mark.parametrize("wallet_id", [None, 12])
 async def test_resolve_and_substitute_session_variables_in_specs(
     mock_user_repo: None,
     mock_osparc_variables_api_auth_rpc: None,
     fake_app: FastAPI,
     faker: Faker,
+    wallet_id: WalletID | None,
 ):
     specs = {
         "product_name": "${OSPARC_VARIABLE_PRODUCT_NAME}",
@@ -175,6 +178,7 @@ async def test_resolve_and_substitute_session_variables_in_specs(
         "api_key": "${OSPARC_VARIABLE_API_KEY}",
         "api_secret": "${OSPARC_VARIABLE_API_SECRET}",
         "service_run_id": "${OSPARC_VARIABLE_SERVICE_RUN_ID}",
+        "wallet_id": "${OSPARC_VARIABLE_WALLET_ID}",
     }
     print("SPECS\n", specs)
 
@@ -186,10 +190,12 @@ async def test_resolve_and_substitute_session_variables_in_specs(
         project_id=faker.uuid4(cast_to=None),
         node_id=faker.uuid4(cast_to=None),
         service_run_id=ServiceRunID("some_run_id"),
+        wallet_id=wallet_id,
     )
     print("REPLACED SPECS\n", replaced_specs)
 
     assert OSPARC_IDENTIFIER_PREFIX not in f"{replaced_specs}"
+    assert f"'wallet_id': '{wallet_id}'" in f"{replaced_specs}"
 
 
 @pytest.fixture

diff --git a/services/director-v2/tests/unit/with_dbs/test_utils_dask.py b/services/director-v2/tests/unit/with_dbs/test_utils_dask.py
@@ -661,5 +661,6 @@ async def test_compute_task_envs(
         node_image=sleeper_task.image,
         metadata=run_metadata,
         resource_tracking_run_id=resource_tracking_run_id,
+        wallet_id=run_metadata.get("wallet_id", None),
     )
     assert task_envs == expected_computed_task_envs