rework interaction with jq (#567)

* rework interaction with jq
---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: fynnos

.vscode/launch.json

@@ -11,7 +11,6 @@
       "cwd": "${workspaceFolder}/backend",
       "envFile": "${workspaceFolder}/backend/.env",
       "env": {
-        "BACKEND_TYPE": "api",
         "PYTHONPATH": "${workspaceFolder}/backend/src"
       }
     },
@@ -26,7 +25,6 @@
       "cwd": "${workspaceFolder}/backend",
       "envFile": "${workspaceFolder}/backend/.env",
       "env": {
-        "BACKEND_TYPE": "worker",
         "PYTHONPATH": "${workspaceFolder}/backend/src"
       }
     },

backend/src/backend_api_entrypoint.sh

@@ -3,8 +3,6 @@
 set -e
 source .venv/bin/activate
 
-export BACKEND_TYPE="api"
-
 LOG_LEVEL=${LOG_LEVEL:-debug}
 API_PORT=${API_PORT:-5500}
 API_WORKERS=${API_WORKERS:-10}

backend/src/celery_background_jobs_worker_entrypoint.sh

@@ -5,7 +5,6 @@ source .venv/bin/activate
 
 export OMP_NUM_THREADS=1
 export MKL_NUM_THREADS=1
-export BACKEND_TYPE="worker"
 
 LOG_LEVEL=${LOG_LEVEL:-debug}
 CELERY_BACKGROUND_JOBS_WORKER_CONCURRENCY=${CELERY_BACKGROUND_JOBS_WORKER_CONCURRENCY:-1}

backend/src/modules/concept_over_time_analysis/cota_endpoint.py

@@ -2,9 +2,7 @@
 from common.dependencies import get_current_user, get_db_session
 from core.auth.authz_user import AuthzUser
 from fastapi import APIRouter, Depends
-from modules.concept_over_time_analysis.cota_crud import (
-    crud_cota,
-)
+from modules.concept_over_time_analysis.cota_crud import crud_cota
 from modules.concept_over_time_analysis.cota_dto import (
     COTACreate,
     COTACreateIntern,
@@ -234,5 +232,5 @@ async def get_cota_job(
     authz_user: AuthzUser = Depends(),
 ) -> COTARefinementJobRead:
     job = js.get_job(job_id=cota_job_id)
-    authz_user.assert_in_project(job.meta["project_id"])
+    authz_user.assert_in_project(job.get_project_id())
     return COTARefinementJobRead.from_rq_job(job=job)

backend/src/modules/concept_over_time_analysis/cota_refinement_jobs.py

@@ -1,21 +1,13 @@
-from modules.concept_over_time_analysis.cota_crud import (
-    crud_cota,
-)
-from modules.concept_over_time_analysis.cota_dto import (
-    COTARead,
-    COTARefinementJobInput,
-)
+from modules.concept_over_time_analysis.cota_crud import crud_cota
+from modules.concept_over_time_analysis.cota_dto import COTARead, COTARefinementJobInput
 from modules.concept_over_time_analysis.refinement_steps.finetune_apply_compute import (
     finetune_apply_compute,
 )
 from modules.concept_over_time_analysis.refinement_steps.init_search_space import (
     init_search_space,
 )
-from modules.concept_over_time_analysis.refinement_steps.store_in_db import (
-    store_in_db,
-)
-from rq import get_current_job
-from systems.job_system.job_dto import EndpointGeneration, JobPriority
+from modules.concept_over_time_analysis.refinement_steps.store_in_db import store_in_db
+from systems.job_system.job_dto import EndpointGeneration, Job, JobPriority
 from systems.job_system.job_register_decorator import register_job
 
 
@@ -26,22 +18,18 @@
     priority=JobPriority.DEFAULT,
     generate_endpoints=EndpointGeneration.NONE,
 )
-def cota_refinement(
-    payload: COTARefinementJobInput,
-) -> None:
+def cota_refinement(payload: COTARefinementJobInput, job: Job) -> None:
     from repos.db.sql_repo import SQLRepo
 
-    job = get_current_job()
-    assert job is not None, "Job must be running in a worker context"
-
     # init steps / current_step
-    job.meta["steps"] = [
-        "Initialize search space",
-        "Finetune and apply compute",
-        "Store in DB",
-    ]
-    job.meta["current_step"] = 0
-    job.save_meta()
+    job.update(
+        steps=[
+            "Initialize search space",
+            "Finetune and apply compute",
+            "Store in DB",
+        ],
+        current_step=0,
+    )
 
     with SQLRepo().db_session() as db:
         # make sure the cota exists!
@@ -55,10 +43,8 @@ def cota_refinement(
         # Do the refinement in 3 steps:
         search_space = init_search_space(db=db, cota=cota)
 
-        job.meta["current_step"] = 1
-        job.save_meta()
+        job.update(current_step=1)
         search_space = finetune_apply_compute(cota=cota, search_space=search_space)
 
-        job.meta["current_step"] = 2
-        job.save_meta()
+        job.update(current_step=1)
         store_in_db(db=db, cota_id=cota.id, search_space=search_space)

backend/src/modules/concept_over_time_analysis/cota_service.py

@@ -7,9 +7,7 @@
 from core.metadata.project_metadata_crud import crud_project_meta
 from core.metadata.project_metadata_dto import ProjectMetadataRead
 from fastapi.encoders import jsonable_encoder
-from modules.concept_over_time_analysis.cota_crud import (
-    crud_cota,
-)
+from modules.concept_over_time_analysis.cota_crud import crud_cota
 from modules.concept_over_time_analysis.cota_dto import (
     COTACreateIntern,
     COTARead,
@@ -234,6 +232,6 @@ def start_refinement_job(
             db=db,
             id=payload.cota_id,
             update_dto=COTAUpdateIntern(
-                last_refinement_job_id=job.id,
+                last_refinement_job_id=job.get_id(),
             ),
         )

backend/src/modules/duplicate_finder/duplicate_finder_jobs.py

@@ -10,10 +10,9 @@
 )
 from modules.word_frequency.word_frequency_crud import crud_word_frequency
 from repos.db.sql_repo import SQLRepo
-from rq import get_current_job
 from scipy import sparse
 from sklearn.metrics.pairwise import manhattan_distances
-from systems.job_system.job_dto import EndpointGeneration, JobPriority
+from systems.job_system.job_dto import EndpointGeneration, Job, JobPriority
 from systems.job_system.job_register_decorator import register_job
 
 
@@ -26,12 +25,9 @@
 )
 def find_duplicates_job(
     payload: DuplicateFinderInput,
+    job: Job,
 ) -> DuplicateFinderOutput:
-    job = get_current_job()
-    assert job is not None, "Job must be running in a worker context"
-
-    job.meta["status_message"] = "Started duplicate finding"
-    job.save_meta()
+    job.update(status_message="Started duplicate finding")
 
     logger.info("Finding duplicate text sdocs")
     t0 = time.time()
@@ -40,8 +36,7 @@ def find_duplicates_job(
             db, project_id=payload.project_id, doctype=DocType.text
         )
     t1 = time.time()
-    job.meta["status_message"] = "Fetched word frequencies from database"
-    job.save_meta()
+    job.update(status_message="Fetched word frequencies from database")
     logger.info(f"query took: {t1 - t0}")
 
     t0 = time.time()
@@ -75,8 +70,7 @@ def find_duplicates_job(
         (values, (index, indices)), shape=(len(idx2sdoc_id), vocab_size)
     )
     t1 = time.time()
-    job.meta["status_message"] = "Created document word vectors"
-    job.save_meta()
+    job.update(status_message="Created document word vectors")
     logger.info(f"document vector creation took: {t1 - t0}")
     logger.info(f"vocab size: {vocab_size}")
     logger.info(f"document_vectors shape: {document_vectors.shape}")
@@ -85,8 +79,7 @@ def find_duplicates_job(
     t0 = time.time()
     word_dists = manhattan_distances(document_vectors, document_vectors)
     t1 = time.time()
-    job.meta["status_message"] = "Computed distances between documents"
-    job.save_meta()
+    job.update(status_message="Computed distances between documents")
     logger.info(f"manhatten distance took: {t1 - t0}")
 
     # mask out self distances and one half of the matrix
@@ -103,8 +96,7 @@ def find_duplicates_job(
         )
     ).tolist()
     t1 = time.time()
-    job.meta["status_message"] = "Identified duplicate pairs"
-    job.save_meta()
+    job.update(status_message="Identified duplicate pairs")
     logger.info(f"finding duplicates took: {t1 - t0}")
 
     # map back to sdoc_ids
@@ -120,8 +112,7 @@ def find_duplicates_job(
     G.to_undirected()
     subgraph_nodes = [list(subgraph) for subgraph in nx.connected_components(G)]
     t1 = time.time()
-    job.meta["status_message"] = "Finished finding duplicates!"
-    job.save_meta()
+    job.update(status_message="Finished finding duplicates!")
     logger.info(f"graph grouping took: {t1 - t0}")
 
     return DuplicateFinderOutput(duplicates=subgraph_nodes)

backend/src/modules/eximport/export_jobs.py

@@ -1,5 +1,5 @@
 from modules.eximport.export_job_dto import ExportJobInput, ExportJobOutput
-from systems.job_system.job_dto import EndpointGeneration, JobPriority
+from systems.job_system.job_dto import EndpointGeneration, Job, JobPriority
 from systems.job_system.job_register_decorator import register_job
 
 
@@ -10,9 +10,7 @@
     priority=JobPriority.DEFAULT,
     generate_endpoints=EndpointGeneration.MINIMAL,
 )
-def export_data(
-    payload: ExportJobInput,
-) -> ExportJobOutput:
+def export_data(payload: ExportJobInput, job: Job) -> ExportJobOutput:
     from modules.eximport.export_service import ExportService
 
     return ExportService().handle_export_job(

backend/src/modules/eximport/import_endpoint.py

@@ -3,11 +3,7 @@
 from common.dependencies import get_current_user
 from core.auth.authz_user import AuthzUser
 from fastapi import APIRouter, Depends, UploadFile
-from modules.eximport.import_job_dto import (
-    ImportJobInput,
-    ImportJobRead,
-    ImportJobType,
-)
+from modules.eximport.import_job_dto import ImportJobInput, ImportJobRead, ImportJobType
 from modules.eximport.import_service import ImportJobPreparationError, ImportService
 from repos.filesystem_repo import FilesystemRepo
 from systems.job_system.job_service import JobService
@@ -137,7 +133,7 @@ def get_import_job(
     *, import_job_id: str, authz_user: AuthzUser = Depends()
 ) -> ImportJobRead:
     job = js.get_job(import_job_id)
-    authz_user.assert_in_project(job.meta["project_id"])
+    authz_user.assert_in_project(job.get_project_id())
     return ImportJobRead.from_rq_job(job)
 
 
@@ -152,5 +148,5 @@ def get_all_import_jobs(
     authz_user.assert_in_project(project_id)
 
     jobs = js.get_jobs_by_project(job_type="import", project_id=project_id)
-    jobs.sort(key=lambda x: x.meta["created"], reverse=True)
+    jobs.sort(key=lambda x: x.get_created(), reverse=True)
     return [ImportJobRead.from_rq_job(job) for job in jobs]

backend/src/modules/eximport/import_jobs.py

@@ -1,5 +1,5 @@
 from modules.eximport.import_job_dto import ImportJobInput
-from systems.job_system.job_dto import EndpointGeneration, JobPriority
+from systems.job_system.job_dto import EndpointGeneration, Job, JobPriority
 from systems.job_system.job_register_decorator import register_job
 
 
@@ -10,9 +10,7 @@
     priority=JobPriority.DEFAULT,
     generate_endpoints=EndpointGeneration.NONE,
 )
-def import_data(
-    payload: ImportJobInput,
-) -> None:
+def import_data(payload: ImportJobInput, job: Job) -> None:
     from modules.eximport.import_service import ImportService
 
     return ImportService().handle_import_job(