api: Add tracking_id query param to API endpoints

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

app/api/routers/evaluation.py

@@ -4,7 +4,7 @@
 import uuid
 import tempfile
 
-from typing import List
+from typing import List, Union
 from starlette.status import HTTP_202_ACCEPTED, HTTP_503_SERVICE_UNAVAILABLE
 from typing_extensions import Annotated
 from fastapi import APIRouter, Query, Depends, UploadFile, Request, File
@@ -34,6 +34,7 @@
              description="Evaluate the model being served with a trainer export")
 async def get_evaluation_with_trainer_export(request: Request,
                                              trainer_export: Annotated[List[UploadFile], File(description="One or more trainer export files to be uploaded")],
+                                             tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the evaluation task")] = None,
                                              model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> JSONResponse:
     files = []
     file_names = []
@@ -54,7 +55,7 @@ async def get_evaluation_with_trainer_export(request: Request,
     json.dump(concatenated, data_file)
     data_file.flush()
     data_file.seek(0)
-    evaluation_id = str(uuid.uuid4())
+    evaluation_id = tracking_id or str(uuid.uuid4())
     evaluation_accepted = model_service.train_supervised(data_file, 0, sys.maxsize, evaluation_id, ",".join(file_names))
     if evaluation_accepted:
         return JSONResponse(content={"message": "Your evaluation started successfully.", "evaluation_id": evaluation_id}, status_code=HTTP_202_ACCEPTED)
@@ -69,6 +70,7 @@ async def get_evaluation_with_trainer_export(request: Request,
              description="Sanity check the model being served with a trainer export")
 def get_sanity_check_with_trainer_export(request: Request,
                                          trainer_export: Annotated[List[UploadFile], File(description="One or more trainer export files to be uploaded")],
+                                         tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the sanity check task")] = None,
                                          model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> StreamingResponse:
     files = []
     file_names = []
@@ -88,8 +90,9 @@ def get_sanity_check_with_trainer_export(request: Request,
     metrics = sanity_check_model_with_trainer_export(concatenated, model_service, return_df=True, include_anchors=False)
     stream = io.StringIO()
     metrics.to_csv(stream, index=False)
+    tracking_id = tracking_id or str(uuid.uuid4())
     response = StreamingResponse(iter([stream.getvalue()]), media_type="text/csv")
-    response.headers["Content-Disposition"] = f'attachment ; filename="sanity_check_{str(uuid.uuid4())}.csv"'
+    response.headers["Content-Disposition"] = f'attachment ; filename="sanity_check_{tracking_id}.csv"'
     return response
 
 
@@ -102,7 +105,8 @@ def get_inter_annotator_agreement_scores(request: Request,
                                          trainer_export: Annotated[List[UploadFile], File(description="A list of trainer export files to be uploaded")],
                                          annotator_a_project_id: Annotated[int, Query(description="The project ID from one annotator")],
                                          annotator_b_project_id: Annotated[int, Query(description="The project ID from another annotator")],
-                                         scope: Annotated[str, Query(enum=[s.value for s in Scope], description="The scope for which the score will be calculated, e.g., per_concept, per_document or per_span")]) -> StreamingResponse:
+                                         scope: Annotated[str, Query(enum=[s.value for s in Scope], description="The scope for which the score will be calculated, e.g., per_concept, per_document or per_span")],
+                                         tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the IAA task")] = None) -> StreamingResponse:
     files = []
     for te in trainer_export:
         temp_te = tempfile.NamedTemporaryFile()
@@ -126,8 +130,9 @@ def get_inter_annotator_agreement_scores(request: Request,
             raise AnnotationException(f'Unknown scope: "{scope}"')
         stream = io.StringIO()
         iaa_scores.to_csv(stream, index=False)
+        tracking_id = tracking_id or str(uuid.uuid4())
         response = StreamingResponse(iter([stream.getvalue()]), media_type="text/csv")
-        response.headers["Content-Disposition"] = f'attachment ; filename="iaa_{str(uuid.uuid4())}.csv"'
+        response.headers["Content-Disposition"] = f'attachment ; filename="iaa_{tracking_id}.csv"'
         return response
 
 
@@ -137,7 +142,8 @@ def get_inter_annotator_agreement_scores(request: Request,
              dependencies=[Depends(cms_globals.props.current_active_user)],
              description="Concatenate multiple trainer export files into a single file for download")
 def get_concatenated_trainer_exports(request: Request,
-                                     trainer_export: Annotated[List[UploadFile], File(description="A list of trainer export files to be uploaded")]) -> JSONResponse:
+                                     trainer_export: Annotated[List[UploadFile], File(description="A list of trainer export files to be uploaded")],
+                                     tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the concatenation task")] = None) -> JSONResponse:
     files = []
     for te in trainer_export:
         temp_te = tempfile.NamedTemporaryFile()
@@ -148,8 +154,9 @@ def get_concatenated_trainer_exports(request: Request,
     concatenated = concat_trainer_exports([file.name for file in files], allow_recurring_doc_ids=False)
     for file in files:
         file.close()
+    tracking_id = tracking_id or str(uuid.uuid4())
     response = JSONResponse(concatenated, media_type="application/json; charset=utf-8")
-    response.headers["Content-Disposition"] = f'attachment ; filename="concatenated_{str(uuid.uuid4())}.json"'
+    response.headers["Content-Disposition"] = f'attachment ; filename="concatenated_{tracking_id}.json"'
     return response
 
 
@@ -159,7 +166,8 @@ def get_concatenated_trainer_exports(request: Request,
              dependencies=[Depends(cms_globals.props.current_active_user)],
              description="Get annotation stats of trainer export files")
 def get_annotation_stats(request: Request,
-                         trainer_export: Annotated[List[UploadFile], File(description="One or more trainer export files to be uploaded")]) -> StreamingResponse:
+                         trainer_export: Annotated[List[UploadFile], File(description="One or more trainer export files to be uploaded")],
+                         tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the annotation stats task")] = None) -> StreamingResponse:
     files = []
     file_names = []
     for te in trainer_export:
@@ -177,6 +185,7 @@ def get_annotation_stats(request: Request,
     stats = get_stats_from_trainer_export(concatenated, return_df=True)
     stream = io.StringIO()
     stats.to_csv(stream, index=False)
+    tracking_id = tracking_id or str(uuid.uuid4())
     response = StreamingResponse(iter([stream.getvalue()]), media_type="text/csv")
-    response.headers["Content-Disposition"] = f'attachment ; filename="stats_{str(uuid.uuid4())}.csv"'
+    response.headers["Content-Disposition"] = f'attachment ; filename="stats_{tracking_id}.csv"'
     return response

app/api/routers/invocation.py

@@ -132,6 +132,7 @@ def get_entities_from_multiple_texts(request: Request,
              description="Upload a file containing a list of plain text and extract the NER entities in JSON")
 def extract_entities_from_multi_text_file(request: Request,
                                           multi_text_file: Annotated[UploadFile, File(description="A file containing a list of plain texts, in the format of [\"text_1\", \"text_2\", ..., \"text_n\"]")],
+                                          tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the bulk processing task")] = None,
                                           model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> StreamingResponse:
     with tempfile.NamedTemporaryFile() as data_file:
         for line in multi_text_file.file:
@@ -160,8 +161,9 @@ def extract_entities_from_multi_text_file(request: Request,
         output = json.dumps(body)
         logger.debug(output)
         json_file = BytesIO(output.encode())
+        tracking_id = tracking_id or str(uuid.uuid4())
         response = StreamingResponse(json_file, media_type="application/json")
-        response.headers["Content-Disposition"] = f'attachment ; filename="concatenated_{str(uuid.uuid4())}.json"'
+        response.headers["Content-Disposition"] = f'attachment ; filename="concatenated_{tracking_id}.json"'
         return response
 
 

app/api/routers/metacat_training.py

@@ -29,6 +29,7 @@ async def train_metacat(request: Request,
                         epochs: Annotated[int, Query(description="The number of training epochs", ge=0)] = 1,
                         log_frequency: Annotated[int, Query(description="The number of processed documents after which training metrics will be logged", ge=1)] = 1,
                         description: Annotated[Union[str, None], Query(description="The description on the training or change logs")] = None,
+                        tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the training task")] = None,
                         model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> JSONResponse:
     files = []
     file_names = []
@@ -49,7 +50,7 @@ async def train_metacat(request: Request,
     json.dump(concatenated, data_file)
     data_file.flush()
     data_file.seek(0)
-    training_id = str(uuid.uuid4())
+    training_id = tracking_id or str(uuid.uuid4())
     try:
         training_accepted = model_service.train_metacat(data_file,
                                                         epochs,

app/api/routers/preview.py

@@ -27,14 +27,16 @@
              description="Extract the NER entities in HTML for preview")
 async def get_rendered_entities_from_text(request: Request,
                                           text: Annotated[str, Body(description="The text to be sent to the model for NER", media_type="text/plain")],
+                                          tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the preview task")] = None,
                                           model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> StreamingResponse:
     annotations = model_service.annotate(text)
     entities = annotations_to_entities(annotations, model_service.model_name)
     logger.debug("Entities extracted for previewing %s", entities)
     ent_input = Doc(text=text, ents=entities)
     data = displacy.render(ent_input.dict(), style="ent", manual=True)
+    tracking_id = tracking_id or str(uuid.uuid4())
     response = StreamingResponse(BytesIO(data.encode()), media_type="application/octet-stream")
-    response.headers["Content-Disposition"] = f'attachment ; filename="preview_{str(uuid.uuid4())}.html"'
+    response.headers["Content-Disposition"] = f'attachment ; filename="preview_{tracking_id}.html"'
     return response
 
 
@@ -47,7 +49,8 @@ def get_rendered_entities_from_trainer_export(request: Request,
                                               trainer_export: Annotated[List[UploadFile], File(description="One or more trainer export files to be uploaded")] = [],
                                               trainer_export_str: Annotated[str, Form(description="The trainer export raw JSON string")] = "{\"projects\": []}",
                                               project_id: Annotated[Union[int, None], Query(description="The target project ID, and if not provided, all projects will be included")] = None,
-                                              document_id: Annotated[Union[int, None], Query(description="The target document ID, and if not provided, all documents of the target project(s) will be included")] = None) -> Response:
+                                              document_id: Annotated[Union[int, None], Query(description="The target document ID, and if not provided, all documents of the target project(s) will be included")] = None,
+                                              tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the trainer export preview task")] = None) -> Response:
     data: Dict = {"projects": []}
     if trainer_export is not None:
         files = []
@@ -88,8 +91,9 @@ def get_rendered_entities_from_trainer_export(request: Request,
             doc = Doc(text=document["text"], ents=entities, title=f"P{project['id']}/D{document['id']}")
             htmls.append(displacy.render(doc.dict(), style="ent", manual=True))
     if htmls:
+        tracking_id = tracking_id or str(uuid.uuid4())
         response = StreamingResponse(BytesIO("<br/>".join(htmls).encode()), media_type="application/octet-stream")
-        response.headers["Content-Disposition"] = f'attachment ; filename="preview_{str(uuid.uuid4())}.html"'
+        response.headers["Content-Disposition"] = f'attachment ; filename="preview_{tracking_id}.html"'
     else:
         logger.debug("Cannot find any matching documents to preview")
         return JSONResponse(content={"message": "Cannot find any matching documents to preview"}, status_code=HTTP_404_NOT_FOUND)

app/api/routers/supervised_training.py

@@ -32,6 +32,7 @@ async def train_supervised(request: Request,
                            test_size: Annotated[Union[float, None], Query(description="The override of the test size in percentage. (For a 'huggingface-ner' model, a negative value can be used to apply the train-validation-test split if implicitly defined in trainer export: 'projects[0]' is used for training, 'projects[1]' for validation, and 'projects[2]' for testing)")] = 0.2,
                            log_frequency: Annotated[int, Query(description="The number of processed documents after which training metrics will be logged", ge=1)] = 1,
                            description: Annotated[Union[str, None], Form(description="The description of the training or change logs")] = None,
+                           tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the training task")] = None,
                            model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> JSONResponse:
     files = []
     file_names = []
@@ -51,7 +52,7 @@ async def train_supervised(request: Request,
     json.dump(concatenated, data_file)
     data_file.flush()
     data_file.seek(0)
-    training_id = str(uuid.uuid4())
+    training_id = tracking_id or str(uuid.uuid4())
     try:
         training_accepted = model_service.train_supervised(data_file,
                                                            epochs,

app/api/routers/unsupervised_training.py

@@ -33,6 +33,7 @@ async def train_unsupervised(request: Request,
                              test_size: Annotated[Union[float, None], Query(description="The override of the test size in percentage", ge=0.0)] = 0.2,
                              log_frequency: Annotated[int, Query(description="The number of processed documents after which training metrics will be logged", ge=1)] = 1000,
                              description: Annotated[Union[str, None], Query(description="The description of the training or change logs")] = None,
+                             tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the training task")] = None,
                              model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> JSONResponse:
     """
     Upload one or more plain text files and trigger the unsupervised training
@@ -61,7 +62,7 @@ async def train_unsupervised(request: Request,
     logger.debug("Training data concatenated")
     data_file.flush()
     data_file.seek(0)
-    training_id = str(uuid.uuid4())
+    training_id = tracking_id or str(uuid.uuid4())
     try:
         training_accepted = model_service.train_unsupervised(data_file,
                                                              epochs,
@@ -96,6 +97,7 @@ async def train_unsupervised_with_hf_dataset(request: Request,
                                              test_size: Annotated[Union[float, None], Query(description="The override of the test size in percentage will only take effect if the dataset does not have predefined validation or test splits", ge=0.0)] = 0.2,
                                              log_frequency: Annotated[int, Query(description="The number of processed documents after which training metrics will be logged", ge=1)] = 1000,
                                              description: Annotated[Union[str, None], Query(description="The description of the training or change logs")] = None,
+                                             tracking_id: Annotated[Union[str, None], Query(description="The tracking ID of the training task")] = None,
                                              model_service: AbstractModelService = Depends(cms_globals.model_service_dep)) -> JSONResponse:
     """
     Trigger the unsupervised training with a dataset from Hugging Face Hub
@@ -129,7 +131,7 @@ async def train_unsupervised_with_hf_dataset(request: Request,
         logger.debug("Training dataset downloaded and transformed")
         hf_dataset.save_to_disk(data_dir.name)
 
-    training_id = str(uuid.uuid4())
+    training_id = tracking_id or str(uuid.uuid4())
     training_accepted = model_service.train_unsupervised(data_dir,
                                                          epochs,
                                                          log_frequency,