propagate task creation all the way out

Author: bisgaard-itis

services/api-server/src/simcore_service_api_server/_service_function_jobs.py

@@ -34,9 +34,8 @@
 from ._service_jobs import JobService
 from .models.api_resources import JobLinks
 from .models.domain.functions import (
+    FunctionJobPatch,
     PreRegisteredFunctionJobData,
-    ProjectFunctionJobPatch,
-    SolverFunctionJobPatch,
 )
 from .models.schemas.jobs import JobInputs, JobPricingSpecification
 from .services_http.webserver import AuthSession
@@ -209,14 +208,13 @@ async def patch_registered_function_job(
         user_id: UserID,
         product_name: ProductName,
         patches: Annotated[
-            list[ProjectFunctionJobPatch] | list[SolverFunctionJobPatch],
+            list[FunctionJobPatch],
             Field(max_length=50, min_length=1),
         ],
     ) -> list[RegisteredFunctionJob]:
         patch_inputs = []
         for patch in patches:
             if patch.function_class == FunctionClass.PROJECT:
-                assert isinstance(patch, ProjectFunctionJobPatch)  # nosec
                 patch_inputs.append(
                     RegisteredProjectFunctionJobPatch(
                         title=None,
@@ -228,7 +226,6 @@ async def patch_registered_function_job(
                     )
                 )
             elif patch.function_class == FunctionClass.SOLVER:
-                assert isinstance(patch, SolverFunctionJobPatch)  # nosec
                 patch_inputs.append(
                     RegisteredSolverFunctionJobPatch(
                         title=None,
@@ -279,7 +276,8 @@ async def run_function(
                 user_id=self.user_id,
                 product_name=self.product_name,
                 patches=[
-                    ProjectFunctionJobPatch(
+                    FunctionJobPatch(
+                        function_class=FunctionClass.PROJECT,
                         function_job_id=pre_registered_function_job_data.function_job_id,
                         job_creation_task_id=None,
                         project_job_id=study_job.id,
@@ -309,7 +307,8 @@ async def run_function(
                 user_id=self.user_id,
                 product_name=self.product_name,
                 patches=[
-                    SolverFunctionJobPatch(
+                    FunctionJobPatch(
+                        function_class=FunctionClass.SOLVER,
                         function_job_id=pre_registered_function_job_data.function_job_id,
                         job_creation_task_id=None,
                         solver_job_id=solver_job.id,

services/api-server/src/simcore_service_api_server/_service_function_jobs_task_client.py

@@ -38,17 +38,16 @@
 )
 from sqlalchemy.ext.asyncio import AsyncEngine
 
+from ._meta import APP_NAME
 from ._service_function_jobs import FunctionJobService
 from ._service_functions import FunctionService
 from ._service_jobs import JobService
 from .api.dependencies.authentication import Identity
-from .exceptions.function_errors import (
-    FunctionJobCacheNotFoundError,
-)
 from .models.api_resources import JobLinks
 from .models.domain.celery_models import ApiServerOwnerMetadata
+from .models.domain.functions import FunctionJobPatch
 from .models.schemas.functions import FunctionJobCreationTaskStatus
-from .models.schemas.jobs import JobInputs, JobPricingSpecification
+from .models.schemas.jobs import JobPricingSpecification
 from .services_http.webserver import AuthSession
 from .services_rpc.storage import StorageService
 from .services_rpc.wb_api_server import WbApiRpcClient
@@ -238,29 +237,6 @@ async def inspect_function_job(
             check_write_permissions=False,
         )
 
-    async def get_cached_function_job(
-        self,
-        *,
-        function: RegisteredFunction,
-        job_inputs: JobInputs,
-    ) -> RegisteredFunctionJob:
-        """Raises FunctionJobCacheNotFoundError if no cached job is found"""
-        if cached_function_jobs := await self._web_rpc_client.find_cached_function_jobs(
-            function_id=function.uid,
-            inputs=job_inputs.values,
-            user_id=self.user_id,
-            product_name=self.product_name,
-        ):
-            for cached_function_job in cached_function_jobs:
-                job_status = await self.inspect_function_job(
-                    function=function,
-                    function_job=cached_function_job,
-                )
-                if job_status.status == RunningState.SUCCESS:
-                    return cached_function_job
-
-        raise FunctionJobCacheNotFoundError
-
     async def function_job_outputs(
         self,
         *,
@@ -319,7 +295,7 @@ async def function_job_outputs(
             check_write_permissions=False,
         )
 
-    async def create_function_job_creation_task(
+    async def create_function_job_creation_tasks(
         self,
         *,
         function: RegisteredFunction,
@@ -329,7 +305,7 @@ async def create_function_job_creation_task(
         job_links: JobLinks,
         parent_project_uuid: ProjectID | None = None,
         parent_node_id: NodeID | None = None,
-    ) -> RegisteredFunctionJob:
+    ) -> list[RegisteredFunctionJob]:
         inputs = [
             self._function_job_service.create_function_job_inputs(
                 function=function, function_inputs=input_
@@ -359,9 +335,10 @@ async def create_function_job_creation_task(
         )
 
         # run function in celery task
-
         owner_metadata = ApiServerOwnerMetadata(
-            user_id=user_identity.user_id, product_name=user_identity.product_name
+            user_id=user_identity.user_id,
+            product_name=user_identity.product_name,
+            owner=APP_NAME,
         )
         task_uuids = [
             await self._celery_task_manager.submit_task(
@@ -382,10 +359,22 @@ async def create_function_job_creation_task(
             for pre_registered_function_job_data in pre_registered_function_job_data_list
         ]
 
-        return await self._function_job_service.patch_registered_function_job(
+        patched_jobs = await self._function_job_service.patch_registered_function_job(
             user_id=user_identity.user_id,
             product_name=user_identity.product_name,
-            function_job_id=pre_registered_function_job_data_list.function_job_id,
-            function_class=function.function_class,
-            job_creation_task_id=TaskID(task_uuid),
+            patches=[
+                FunctionJobPatch(
+                    function_class=function.function_class,
+                    function_job_id=pre_registered_function_job_data.function_job_id,
+                    job_creation_task_id=TaskID(task_uuid),
+                    project_job_id=None,
+                    solver_job_id=None,
+                )
+                for task_uuid, pre_registered_function_job_data in zip(
+                    task_uuids, pre_registered_function_job_data_list
+                )
+            ],
         )
+        patched_jobs_iter = iter(patched_jobs)
+        _ = lambda job: job if job is not None else next(patched_jobs_iter)
+        return [_(job) for job in cached_jobs]

services/api-server/src/simcore_service_api_server/api/routes/functions_routes.py

@@ -16,13 +16,12 @@
     RegisteredFunctionJob,
     RegisteredFunctionJobCollection,
 )
-from models_library.functions import FunctionJobCollection, FunctionJobID
+from models_library.functions import FunctionJobCollection
 from models_library.products import ProductName
 from models_library.projects import ProjectID
 from models_library.projects_nodes_io import NodeID
 from models_library.users import UserID
 from servicelib.fastapi.dependencies import get_reverse_url_mapper
-from servicelib.utils import limited_gather
 
 from ..._service_function_jobs import FunctionJobService
 from ..._service_function_jobs_task_client import FunctionJobTaskClientService
@@ -350,15 +349,17 @@ async def run_function(
     )
     job_links = await function_service.get_function_job_links(to_run_function, url_for)
 
-    return await function_job_task_client_service.create_function_job_creation_task(
+    jobs = await function_job_task_client_service.create_function_job_creation_tasks(
         function=to_run_function,
-        function_inputs=function_inputs,
+        function_inputs=[function_inputs],
         user_identity=user_identity,
         pricing_spec=pricing_spec,
         job_links=job_links,
         parent_project_uuid=parent_project_uuid,
         parent_node_id=parent_node_id,
     )
+    assert len(jobs) == 1  # nosec
+    return jobs[0]
 
 
 @function_router.delete(
@@ -429,42 +430,23 @@ async def map_function(
     )
     job_links = await function_service.get_function_job_links(to_run_function, url_for)
 
-    async def _run_single_function(function_inputs: FunctionInputs) -> FunctionJobID:
-        result = (
-            await function_job_task_client_service.create_function_job_creation_task(
-                function=to_run_function,
-                function_inputs=function_inputs,
-                user_identity=user_identity,
-                pricing_spec=pricing_spec,
-                job_links=job_links,
-                parent_project_uuid=parent_project_uuid,
-                parent_node_id=parent_node_id,
-            )
-        )
-        return result.uid
-
-    # Run all tasks concurrently, allowing them to complete even if some fail
-    results = await limited_gather(
-        *[
-            _run_single_function(function_inputs)
-            for function_inputs in function_inputs_list
-        ],
-        reraise=False,
-        limit=1,
+    jobs = await function_job_task_client_service.create_function_job_creation_tasks(
+        function=to_run_function,
+        function_inputs=function_inputs_list,
+        user_identity=user_identity,
+        pricing_spec=pricing_spec,
+        job_links=job_links,
+        parent_project_uuid=parent_project_uuid,
+        parent_node_id=parent_node_id,
     )
 
-    # Check if any tasks raised exceptions and raise the first one found
-    for result in results:
-        if isinstance(result, BaseException):
-            raise result
-
     # At this point, all results are FunctionJobID since we've checked for exceptions
     function_job_collection_description = f"Function job collection of map of function {to_run_function.uid} with {len(function_inputs_list)} inputs"
     return await web_api_rpc_client.register_function_job_collection(
         function_job_collection=FunctionJobCollection(
             title="Function job collection of function map",
             description=function_job_collection_description,
-            job_ids=results,  # type: ignore
+            job_ids=jobs,  # type: ignore
         ),
         user_id=user_identity.user_id,
         product_name=user_identity.product_name,

services/api-server/src/simcore_service_api_server/models/domain/functions.py

@@ -1,5 +1,3 @@
-from typing import NamedTuple
-
 from models_library.functions import (
     FunctionClass,
     FunctionJobID,
@@ -9,7 +7,7 @@
     TaskID,
 )
 from models_library.projects import ProjectID
-from pydantic import BaseModel
+from pydantic import BaseModel, model_validator
 
 from ...models.pagination import Page
 from ...models.schemas.jobs import JobInputs
@@ -28,15 +26,28 @@ class PageRegisteredFunctionJobWithorWithoutStatus(
     pass
 
 
-class ProjectFunctionJobPatch(NamedTuple):
-    function_class = FunctionClass.PROJECT
-    function_job_id: FunctionJobID
-    job_creation_task_id: TaskID | None
-    project_job_id: ProjectID | None
-
-
-class SolverFunctionJobPatch(NamedTuple):
-    function_class = FunctionClass.SOLVER
+class FunctionJobPatch(BaseModel):
+    function_class: FunctionClass
     function_job_id: FunctionJobID
-    job_creation_task_id: TaskID | None
-    solver_job_id: SolverJobID | None
+    job_creation_task_id: TaskID | None = None
+    project_job_id: ProjectID | None = None
+    solver_job_id: SolverJobID | None = None
+
+    @model_validator(mode="after")
+    def validate_function_class_consistency(self) -> "FunctionJobPatch":
+        """Validate consistency between function_class and job IDs."""
+        if (
+            self.solver_job_id is not None
+            and self.function_class != FunctionClass.SOLVER
+        ):
+            msg = f"solver_job_id must be None when function_class is {self.function_class}, expected {FunctionClass.SOLVER}"
+            raise ValueError(msg)
+
+        if (
+            self.project_job_id is not None
+            and self.function_class != FunctionClass.PROJECT
+        ):
+            msg = f"project_job_id must be None when function_class is {self.function_class}, expected {FunctionClass.PROJECT}"
+            raise ValueError(msg)
+
+        return self