Implement job_result_download for experiment service (#632)

* change job_result_download to experiment service * orm get jobs by experiment * add experiment methods * add service deps * fix service deps * fix dataset input * remove comment * fix schemas * Update lumigator/python/mzai/backend/backend/services/jobs.py Co-authored-by: Davide Eynard <[email protected]> Signed-off-by: Vicki Boykis <[email protected]> * Update lumigator/python/mzai/backend/backend/services/jobs.py Co-authored-by: Davide Eynard <[email protected]> Signed-off-by: Vicki Boykis <[email protected]> * fix experiment query * fetch jobs by experiment from experiment repo * remove schema change * add imports after rebase * clean up imports * refactor jobs call * add dataset service * change test and experiment response * testing * change types for routes * fix failing int tests * cleanup make * remove extra get by id method since covered in jobs * hotfix broken links --------- Signed-off-by: Vicki Boykis <[email protected]> Co-authored-by: Davide Eynard <[email protected]>
mozilla-ai · Jan 28, 2025 · 1c73e22 · 1c73e22
1 parent 2609a52
commit 1c73e22
Show file tree

Hide file tree

Showing 10 changed files with 262 additions and 203 deletions.
diff --git a/docs/source/conceptual-guides/endpoints.md b/docs/source/conceptual-guides/endpoints.md
@@ -18,11 +18,11 @@ the backend itself relies on:
   are used in the API which allows one to exactly know which kind of data has to be passed to it.
 
 * **Services** implement the actual functionalities and are called by the different methods exposed
-  in the API (see: [backend services](https://github.com/mozilla-ai/lumigator/tree/31d4cd3ecdafba6292a71e125b31f23498ddf9da/lumigator/python/mzai/backend/backend/services)).
+  in the API (see: [backend services](https://github.com/mozilla-ai/lumigator/tree/main/lumigator/python/mzai/backend/backend/services)).
 
 * **Repositories** implement the [repository pattern](https://www.cosmicpython.com/book/chapter_02_repository.html)
-  as an abstraction over the SQL database (see: [code for repositories](https://github.com/mozilla-ai/lumigator/tree/31d4cd3ecdafba6292a71e125b31f23498ddf9da/lumigator/python/mzai/backend/backend/repositories)).
-  They make use of [record classes](https://github.com/mozilla-ai/lumigator/tree/31d4cd3ecdafba6292a71e125b31f23498ddf9da/lumigator/python/mzai/backend/backend/records) to refer to actual records in the database.
+  as an abstraction over the SQL database (see: [code for repositories](https://github.com/mozilla-ai/lumigator/tree/main/lumigator/python/mzai/backend/backend/repositories)).
+  They make use of [record classes](https://github.com/mozilla-ai/lumigator/tree/main/lumigator/python/mzai/backend/backend/records) to refer to actual records in the database.
 
 ![Lumigator Backend](../../assets/lumigator-backend.svg)
 

diff --git a/lumigator/python/mzai/backend/backend/api/deps.py b/lumigator/python/mzai/backend/backend/api/deps.py
@@ -62,7 +62,9 @@ def get_job_service(session: DBSessionDep, dataset_service: DatasetServiceDep) -
 
 
 def get_experiment_service(
-    session: DBSessionDep, job_service: JobServiceDep, dataset_service: DatasetServiceDep
+    session: DBSessionDep,
+    job_service: JobServiceDep,
+    dataset_service: DatasetServiceDep,
 ) -> ExperimentService:
     job_repo = JobRepository(session)
     experiment_repo = ExperimentRepository(session)

diff --git a/lumigator/python/mzai/backend/backend/api/routes/experiments_new.py b/lumigator/python/mzai/backend/backend/api/routes/experiments_new.py
@@ -9,6 +9,7 @@
     ExperimentResultResponse,
 )
 from lumigator_schemas.extras import ListingResponse
+from lumigator_schemas.jobs import JobResponse
 
 from backend.api.deps import ExperimentServiceDep, JobServiceDep
 from backend.services.exceptions.base_exceptions import ServiceError
@@ -38,8 +39,10 @@ def get_experiment(service: ExperimentServiceDep, experiment_id: UUID) -> Experi
 @router.get("/{experiment_id}/jobs")
 def get_experiment_jobs(
     service: ExperimentServiceDep, experiment_id: UUID
-) -> ListingResponse[UUID]:
-    return service.get_all_owned_jobs(experiment_id)
+) -> ListingResponse[JobResponse]:
+    return ListingResponse[JobResponse].model_validate(
+        service._get_experiment_jobs(experiment_id).model_dump()
+    )
 
 
 @router.get("/")
@@ -66,10 +69,10 @@ def get_experiment_result(
 
 @router.get("/{experiment_id}/result/download")
 def get_experiment_result_download(
-    service: JobServiceDep,
+    service: ExperimentServiceDep,
     experiment_id: UUID,
 ) -> ExperimentResultDownloadResponse:
     """Return experiment results file URL for downloading."""
     return ExperimentResultDownloadResponse.model_validate(
-        service.get_job_result_download(experiment_id).model_dump()
+        service.get_experiment_result_download(experiment_id).model_dump()
     )
diff --git a/lumigator/python/mzai/backend/backend/repositories/jobs.py b/lumigator/python/mzai/backend/backend/repositories/jobs.py
@@ -10,7 +10,7 @@ class JobRepository(BaseRepository[JobRecord]):
     def __init__(self, session: Session):
         super().__init__(JobRecord, session)
 
-    def get_by_experiment_id(self, experiment_id: UUID) -> list[JobRecord]:
+    def get_by_experiment_id(self, experiment_id: UUID) -> list[JobRecord] | None:
         return self.session.query(JobRecord).where(JobRecord.experiment_id == experiment_id).all()
 
 

diff --git a/lumigator/python/mzai/backend/backend/services/experiments.py b/lumigator/python/mzai/backend/backend/services/experiments.py
@@ -1,21 +1,24 @@
+import json
 from uuid import UUID
 
 import loguru
 from fastapi import BackgroundTasks
-from lumigator_schemas.experiments import ExperimentCreate, ExperimentResponse
-from lumigator_schemas.extras import ListingResponse
-from lumigator_schemas.jobs import (
-    JobEvalLiteCreate,
-    JobInferenceCreate,
-    JobStatus,
+from lumigator_schemas.experiments import (
+    ExperimentCreate,
+    ExperimentResponse,
+    ExperimentResultDownloadResponse,
 )
+from lumigator_schemas.extras import ListingResponse
+from lumigator_schemas.jobs import JobEvalLiteCreate, JobInferenceCreate, JobResponse, JobStatus
+from s3fs import S3FileSystem
 
 from backend.records.jobs import JobRecord
 from backend.repositories.experiments import ExperimentRepository
 from backend.repositories.jobs import JobRepository
 from backend.services.datasets import DatasetService
 from backend.services.exceptions.experiment_exceptions import ExperimentNotFoundError
 from backend.services.jobs import JobService
+from backend.settings import settings
 
 
 class ExperimentService:
@@ -139,6 +142,51 @@ def get_experiment(self, experiment_id: UUID) -> ExperimentResponse:
 
         return ExperimentResponse.model_validate(record)
 
+    def _get_experiment_jobs(self, experiment_id: UUID) -> ListingResponse[JobResponse]:
+        records = self._job_repo.get_by_experiment_id(experiment_id)
+        return ListingResponse(
+            total=len(records),
+            items=[JobResponse.model_validate(x) for x in records],
+        )
+
+    def get_experiment_result_download(
+        self, experiment_id: UUID
+    ) -> ExperimentResultDownloadResponse:
+        """Return experiment results file URL for downloading."""
+        s3 = S3FileSystem()
+        # get jobs matching this experiment
+        # have a query returning a list of (two) jobs, one inference and one eval,
+        # matching the current experiment id. Note that each job has its own type baked in
+        # (per https://github.com/mozilla-ai/lumigator/pull/576)
+        jobs = self._get_experiment_jobs(experiment_id)
+
+        # iterate on jobs and depending on which job this is, import selected fields
+        results = {}
+
+        for job in jobs:
+            # Get the dataset from the S3 bucket
+            result_key = self._job_service._get_results_s3_key(job.id)
+            with s3.open(f"{settings.S3_BUCKET}/{result_key}", "r") as f:
+                job_results = json.loads(f.read())
+
+            # we just merge the two dictionaries for now
+            results.update(job_results)
+
+        loguru.logger.error(results)
+
+        # Generate presigned download URL for the object
+        result_key = self._job_service._get_results_s3_key(experiment_id)
+        download_url = self._dataset_service.s3_client.generate_presigned_url(
+            "get_object",
+            Params={
+                "Bucket": settings.S3_BUCKET,
+                "Key": result_key,
+            },
+            ExpiresIn=settings.S3_URL_EXPIRATION,
+        )
+
+        return ExperimentResultDownloadResponse(id=experiment_id, download_url=download_url)
+
     def list_experiments(
         self, skip: int = 0, limit: int = 100
     ) -> ListingResponse[ExperimentResponse]:

diff --git a/lumigator/python/mzai/backend/backend/services/jobs.py b/lumigator/python/mzai/backend/backend/services/jobs.py
@@ -291,7 +291,7 @@ def _get_config_template(self, job_type: str, model_name: str) -> str:
 
         return config_template
 
-    def _set_model_type(self, request: BaseModel) -> str:
+    def _set_model_type(self, request) -> str:
         """Sets model URL based on protocol address"""
         if request.model.startswith("oai://"):
             model_url = settings.OAI_API_URL

diff --git a/lumigator/python/mzai/backend/backend/tests/conftest.py b/lumigator/python/mzai/backend/backend/tests/conftest.py
@@ -15,6 +15,8 @@
 from fastapi.testclient import TestClient
 from fsspec.implementations.memory import MemoryFileSystem
 from loguru import logger
+
+from lumigator_schemas.experiments import ExperimentResponse
 from lumigator_schemas.jobs import (
     JobConfig,
     JobResponse,
@@ -95,7 +97,7 @@ def wait_for_experiment(client, experiment_id: UUID) -> bool:
     for _ in range(1, MAX_POLLS):
         get_experiment_response = client.get(f"/experiments_new/{experiment_id}")
         assert get_experiment_response.status_code == 200
-        get_experiment_response_model = JobResponse.model_validate(get_experiment_response.json())
+        get_experiment_response_model = ExperimentResponse.model_validate(get_experiment_response.json())
         if get_experiment_response_model.status == JobStatus.SUCCEEDED.value:
             succeeded = True
             timed_out = False