First attempt at a parametrized JobCreate

lumigator/backend/backend/api/deps.py

@@ -2,7 +2,7 @@
 from typing import Annotated
 
 import boto3
-from fastapi import Depends
+from fastapi import BackgroundTasks, Depends
 from mypy_boto3_s3.client import S3Client
 from ray.job_submission import JobSubmissionClient
 from s3fs import S3FileSystem
@@ -59,11 +59,13 @@ def get_dataset_service(
 DatasetServiceDep = Annotated[DatasetService, Depends(get_dataset_service)]
 
 
-def get_job_service(session: DBSessionDep, dataset_service: DatasetServiceDep) -> JobService:
+def get_job_service(
+    session: DBSessionDep, dataset_service: DatasetServiceDep, background_tasks: BackgroundTasks
+) -> JobService:
     job_repo = JobRepository(session)
     result_repo = JobResultRepository(session)
     ray_client = JobSubmissionClient(settings.RAY_DASHBOARD_URL)
-    return JobService(job_repo, result_repo, ray_client, dataset_service)
+    return JobService(job_repo, result_repo, ray_client, dataset_service, background_tasks)
 
 
 JobServiceDep = Annotated[JobService, Depends(get_job_service)]
@@ -87,9 +89,12 @@ def get_workflow_service(
     tracking_client: TrackingClientDep,
     job_service: JobServiceDep,
     dataset_service: DatasetServiceDep,
+    background_tasks: BackgroundTasks,
 ) -> WorkflowService:
     job_repo = JobRepository(session)
-    return WorkflowService(job_repo, job_service, dataset_service, tracking_client=tracking_client)
+    return WorkflowService(
+        job_repo, job_service, dataset_service, background_tasks, tracking_client=tracking_client
+    )
 
 
 WorkflowServiceDep = Annotated[WorkflowService, Depends(get_workflow_service)]

lumigator/backend/backend/api/routes/experiments.py

@@ -5,13 +5,11 @@
 from lumigator_schemas.experiments import (
     ExperimentCreate,
     ExperimentIdCreate,
-    ExperimentIdResponse,
-    ExperimentResponse,
     GetExperimentResponse,
 )
 from lumigator_schemas.extras import ListingResponse
 from lumigator_schemas.jobs import (
-    JobEvalCreate,
+    JobCreate,
     JobResponse,
     JobResultDownloadResponse,
     JobResultResponse,
@@ -31,13 +29,16 @@ def experiment_exception_mappings() -> dict[type[ServiceError], HTTPStatus]:
 
 
 @router.post("/", status_code=status.HTTP_201_CREATED)
-def create_experiment(service: JobServiceDep, request: ExperimentCreate) -> JobResponse:
-    return service.create_job(JobEvalCreate.model_validate(request.model_dump()))
+def create_experiment(
+    service: JobServiceDep,
+    request: ExperimentCreate,
+) -> GetExperimentResponse:
+    return service.create_job(JobCreate.model_validate(request.model_dump()))
 
 
 @router.get("/{experiment_id}")
 def get_experiment(service: JobServiceDep, experiment_id: UUID) -> JobResponse:
-    return ExperimentResponse.model_validate(service.get_job(experiment_id).model_dump())
+    return GetExperimentResponse.model_validate(service.get_job(experiment_id).model_dump())
 
 
 @router.get("/")
@@ -64,9 +65,7 @@ def get_experiment_result_download(
     experiment_id: UUID,
 ) -> JobResultDownloadResponse:
     """Return experiment results file URL for downloading."""
-    return JobResultDownloadResponse.model_validate(
-        service.get_job_result_download(experiment_id).model_dump()
-    )
+    return JobResultDownloadResponse.model_validate(service.get_job_result_download(experiment_id).model_dump())
 
 
 ####################################################################################################
@@ -78,11 +77,10 @@ def get_experiment_result_download(
 # but right now it is a placeholder while we build up the Workflows routes
 # It's not included in the OpenAPI schema for now so it's not visible in the docs
 @router.post("/new", status_code=status.HTTP_201_CREATED, include_in_schema=True)
-def create_experiment_id(
-    service: ExperimentServiceDep, request: ExperimentIdCreate
-) -> ExperimentIdResponse:
+def create_experiment_id(service: ExperimentServiceDep, request: ExperimentIdCreate) -> GetExperimentResponse:
     """Create an experiment ID."""
-    return ExperimentIdResponse.model_validate(service.create_experiment(request).model_dump())
+    # FIXME Shouldn't the model set this
+    return GetExperimentResponse.model_validate(service.create_experiment(request).model_dump())
 
 
 # TODO: FIXME this should not need the /all suffix.
@@ -92,11 +90,9 @@ def list_experiments_new(
     service: ExperimentServiceDep,
     skip: int = 0,
     limit: int = 100,
-) -> ListingResponse[ExperimentResponse]:
+) -> ListingResponse[GetExperimentResponse]:
     """List all experiments."""
-    return ListingResponse[ExperimentResponse].model_validate(
-        service.list_experiments(skip, limit).model_dump()
-    )
+    return ListingResponse[GetExperimentResponse].model_validate(service.list_experiments(skip, limit).model_dump())
 
 
 @router.get("/new/{experiment_id}", include_in_schema=False)

lumigator/backend/backend/api/routes/jobs.py

@@ -6,19 +6,21 @@
 
 import loguru
 import requests
-from fastapi import APIRouter, BackgroundTasks, HTTPException, Query, status
+from fastapi import APIRouter, HTTPException, Query, status
 from lumigator_schemas.datasets import DatasetResponse
 from lumigator_schemas.extras import ListingResponse
 from lumigator_schemas.jobs import (
     Job,
     JobAnnotateCreate,
-    JobEvalCreate,
+    JobCreate,
     JobEvalLiteCreate,
+    JobEvaluateCreate,
     JobInferenceCreate,
     JobLogsResponse,
     JobResponse,
     JobResultDownloadResponse,
     JobResultResponse,
+    JobType,
 )
 from ray.job_submission import JobDetails as RayJobDetails
 from starlette.requests import Request
@@ -37,11 +39,9 @@
 router = APIRouter()
 
 
-def job_exception_mappings() -> (
-    dict[
-        type[JobNotFoundError] | type[JobTypeUnsupportedError] | type[JobUpstreamError] | type[JobValidationError], int
-    ]
-):
+def job_exception_mappings() -> dict[
+    type[JobNotFoundError] | type[JobTypeUnsupportedError] | type[JobUpstreamError] | type[JobValidationError], int
+]:
     return {
         JobNotFoundError: status.HTTP_404_NOT_FOUND,
         JobTypeUnsupportedError: status.HTTP_501_NOT_IMPLEMENTED,
@@ -56,12 +56,9 @@ def create_inference_job(
     job_create_request: JobInferenceCreate,
     request: Request,
     response: Response,
-    background_tasks: BackgroundTasks,
 ) -> JobResponse:
     job_response = service.create_job(job_create_request)
 
-    service.add_background_task(background_tasks, service.handle_inference_job, job_response.id, job_create_request)
-
     url = request.url_for(get_job.__name__, job_id=job_response.id)
     response.headers[HttpHeaders.LOCATION] = f"{url}"
 
@@ -74,26 +71,22 @@ def create_annotation_job(
     job_create_request: JobAnnotateCreate,
     request: Request,
     response: Response,
-    background_tasks: BackgroundTasks,
 ) -> JobResponse:
     """This uses a hardcoded model, that is, Lumigator's opinion on what
     reference model should be used to generate annotations.
     See more: https://blog.mozilla.ai/lets-build-an-app-for-evaluating-llms/
     """
-    inference_job_create_request = JobInferenceCreate(
-        **job_create_request.model_dump(),
-        model="hf://facebook/bart-large-cnn",
-        output_field="ground_truth",
-    )
-    inference_job_create_request.store_to_dataset = True
-    job_response = service.create_job(inference_job_create_request)
+    inference_job_create_config_dict = job_create_request.job_config.dict()
+    inference_job_create_config_dict["model"] = "hf://facebook/bart-large-cnn"
+    inference_job_create_config_dict["output_field"] = "ground_truth"
+    inference_job_create_config_dict["store_to_dataset"] = True
+    inference_job_create_config_dict["job_type"] = JobType.INFERENCE
 
-    service.add_background_task(
-        background_tasks,
-        service.handle_inference_job,
-        job_response.id,
-        inference_job_create_request,
-    )
+    inference_job_create_request_dict = job_create_request.model_dump()
+    inference_job_create_request_dict["job_config"] = inference_job_create_config_dict
+
+    inference_job_create_request = JobCreate(**inference_job_create_request_dict)
+    job_response = service.create_job(inference_job_create_request)
 
     url = request.url_for(get_job.__name__, job_id=job_response.id)
     response.headers[HttpHeaders.LOCATION] = f"{url}"
@@ -104,7 +97,7 @@ def create_annotation_job(
 @router.post("/evaluate/", status_code=status.HTTP_201_CREATED)
 def create_evaluation_job(
     service: JobServiceDep,
-    job_create_request: JobEvalCreate,
+    job_create_request: JobEvaluateCreate,
     request: Request,
     response: Response,
 ) -> JobResponse:

lumigator/backend/backend/api/routes/workflows.py

@@ -1,6 +1,6 @@
 from http import HTTPStatus
 
-from fastapi import APIRouter, BackgroundTasks, status
+from fastapi import APIRouter, status
 from lumigator_schemas.jobs import JobLogsResponse
 from lumigator_schemas.workflows import (
     WorkflowCreateRequest,
@@ -27,14 +27,14 @@ def workflow_exception_mappings() -> dict[type[ServiceError], HTTPStatus]:
 
 @router.post("/", status_code=status.HTTP_201_CREATED)
 async def create_workflow(
-    service: WorkflowServiceDep, request: WorkflowCreateRequest, background_tasks: BackgroundTasks
+    service: WorkflowServiceDep, request: WorkflowCreateRequest
 ) -> WorkflowResponse:
     """A workflow is a single execution for an experiment.
     A workflow is a collection of 1 or more jobs.
     It must be associated with an experiment id,
     which means you must already have created an experiment and have that ID in the request.
     """
-    return WorkflowResponse.model_validate(service.create_workflow(request, background_tasks))
+    return WorkflowResponse.model_validate(service.create_workflow(request))
 
 
 @router.get("/{workflow_id}")