Whisper worker as a service (Jobs don't work yet in airflow)

README.md

@@ -41,7 +41,7 @@ apt-get -y update && apt-get -y upgrade && apt-get install -y --no-install-recom
 ## Running the worker using a CUDA-compatible GPU
 
 To run the worker with a CUDA-compatible GPU instead of the CPU, either:
-- skip steps 3 & 4 from "Docker CPU run" 
+- skip steps 3 & 4 from "Docker CPU run"
 - skip step 3 from "Local run"
 
 **(OUTDATED BUT STILL MIGHT BE RELEVANT)** To run it using a GPU via Docker, check [the instructions from the dane-example-worker](https://github.com/beeldengeluid/dane-example-worker/wiki/Containerization#running-the-container-locally-using-cuda-compatible-gpu).
@@ -50,7 +50,7 @@ Make sure to replace `dane-example-worker` in the `docker run` command with `dan
 
 ## Expected run
 
-The expected run of this worker (whose pipeline is defined in `asr.py`) should 
+The expected run of this worker (whose pipeline is defined in `asr.py`) should
 
 1. download the input file if it isn't downloaded already in `/data/input/` via `download.py`
 
@@ -80,4 +80,4 @@ The pre-trained Whisper model version can be adjusted in the `.env` file by edit
 
 We recommend version `large-v2` as it performs better than `large-v3` in our [benchmarks](https://opensource-spraakherkenning-nl.github.io/ASR_NL_results/).
 
-You can also specify an S3 URI if you have your own custom model available via S3 (by modifying the `W_MODEL` parameter).
+You can also specify an S3 URI if you have your own custom model available via S3 (by modifying the `W_MODEL` parameter).

asr.py

@@ -1,8 +1,26 @@
 import logging
 import os
+import time
+import tomli
+
+from base_util import (
+    get_asset_info,
+    asr_output_dir,
+    save_provenance,
+    PROVENANCE_JSON_FILE,
+)
+from config import (
+    s3_endpoint_url,
+    s3_bucket,
+    s3_folder_in_bucket,
+    w_word_timestamps,
+    w_device,
+    w_model,
+    w_beam_size,
+    w_best_of,
+    w_vad,
+)
 
-from base_util import get_asset_info, asr_output_dir
-from config import s3_endpoint_url, s3_bucket, s3_folder_in_bucket
 from download import download_uri
 from whisper import run_asr, WHISPER_JSON_FILE
 from s3_util import S3Store
@@ -12,48 +30,115 @@
 logger = logging.getLogger(__name__)
 
 
-def run(input_uri: str, output_uri: str) -> bool:
+def _get_project_meta():
+    with open("pyproject.toml", mode="rb") as pyproject:
+        return tomli.load(pyproject)["tool"]["poetry"]
+
+
+pkg_meta = _get_project_meta()
+version = str(pkg_meta["version"])
+
+
+def run(input_uri: str, output_uri: str, model=None) -> bool:
     logger.info(f"Processing {input_uri} (save to --> {output_uri})")
+    start_time = time.time()
+    prov_steps = []  # track provenance
     # 1. download input
     result = download_uri(input_uri)
     logger.info(result)
     if not result:
         logger.error("Could not obtain input, quitting...")
         return False
 
+    prov_steps.append(result.provenance)
+
     input_path = result.file_path
     asset_id, extension = get_asset_info(input_path)
     output_path = asr_output_dir(input_path)
 
     # 2. check if the input file is suitable for processing any further
-    transcoded_file_path = try_transcode(input_path, asset_id, extension)
-    if not transcoded_file_path:
+    transcode_output = try_transcode(input_path, asset_id, extension)
+    if not transcode_output:
         logger.error("The transcode failed to yield a valid file to continue with")
         return False
     else:
-        input_path = transcoded_file_path
+        input_path = transcode_output.transcoded_file_path
+        prov_steps.append(transcode_output.provenance)
 
     # 3. run ASR
     if not asr_already_done(output_path):
         logger.info("No Whisper transcript found")
-        run_asr(input_path, output_path)
+        whisper_prov = run_asr(input_path, output_path, model)
+        if whisper_prov:
+            prov_steps.append(whisper_prov)
     else:
         logger.info(f"Whisper transcript already present in {output_path}")
+        provenance = {
+            "activity_name": "Whisper transcript already exists",
+            "activity_description": "",
+            "processing_time_ms": "",
+            "start_time_unix": "",
+            "parameters": [],
+            "software_version": "",
+            "input_data": "",
+            "output_data": "",
+            "steps": [],
+        }
+        prov_steps.append(provenance)
 
     # 4. generate JSON transcript
     if not daan_transcript_already_done(output_path):
         logger.info("No DAAN transcript found")
-        success = generate_daan_transcript(output_path)
-        if not success:
+        daan_prov = generate_daan_transcript(output_path)
+        if daan_prov:
+            prov_steps.append(daan_prov)
+        else:
             logger.warning("Could not generate DAAN transcript")
     else:
         logger.info(f"DAAN transcript already present in {output_path}")
+        provenance = {
+            "activity_name": "DAAN transcript already exists",
+            "activity_description": "",
+            "processing_time_ms": "",
+            "start_time_unix": "",
+            "parameters": [],
+            "software_version": "",
+            "input_data": "",
+            "output_data": "",
+            "steps": [],
+        }
+        prov_steps.append(provenance)
+
+    end_time = (time.time() - start_time) * 1000
+    final_prov = {
+        "activity_name": "Whisper ASR Worker",
+        "activity_description": "Worker that gets a video/audio file as input and outputs JSON transcripts in various formats",
+        "processing_time_ms": end_time,
+        "start_time_unix": start_time,
+        "parameters": {
+            "word_timestamps": w_word_timestamps,
+            "device": w_device,
+            "vad": w_vad,
+            "model": w_model,
+            "beam_size": w_beam_size,
+            "best_of": w_best_of,
+        },
+        "software_version": version,
+        "input_data": input_uri,
+        "output_data": output_uri if output_uri else output_path,
+        "steps": prov_steps,
+    }
+
+    prov_success = save_provenance(final_prov, output_path)
+    if not prov_success:
+        logger.warning("Could not save the provenance")
 
     # 5. transfer output
     if output_uri:
         transfer_asr_output(output_path, asset_id)
     else:
         logger.info("No output_uri specified, so all is done")
+
     return True
 
 
@@ -84,19 +169,20 @@ def transfer_asr_output(output_path: str, asset_id: str) -> bool:
         [
             os.path.join(output_path, DAAN_JSON_FILE),
             os.path.join(output_path, WHISPER_JSON_FILE),
+            os.path.join(output_path, PROVENANCE_JSON_FILE),
         ],
     )
 
 
 # check if there is a whisper-transcript.json
-def asr_already_done(output_dir):
+def asr_already_done(output_dir) -> bool:
     whisper_transcript = os.path.join(output_dir, WHISPER_JSON_FILE)
     logger.info(f"Checking existence of {whisper_transcript}")
     return os.path.exists(os.path.join(output_dir, WHISPER_JSON_FILE))
 
 
 # check if there is a daan-es-transcript.json
-def daan_transcript_already_done(output_dir):
+def daan_transcript_already_done(output_dir) -> bool:
     daan_transcript = os.path.join(output_dir, DAAN_JSON_FILE)
     logger.info(f"Checking existence of {daan_transcript}")
     return os.path.exists(os.path.join(output_dir, DAAN_JSON_FILE))

base_util.py

@@ -1,18 +1,19 @@
 import logging
-import ntpath
 import os
 import subprocess
+import json
 from typing import Tuple
 from config import data_base_dir
 
 
 LOG_FORMAT = "%(asctime)s|%(levelname)s|%(process)d|%(module)s|%(funcName)s|%(lineno)d|%(message)s"
+PROVENANCE_JSON_FILE = "provenance.json"
 logger = logging.getLogger(__name__)
 
 
 # the file name without extension is used as asset ID
 def get_asset_info(input_file: str) -> Tuple[str, str]:
-    file_name = ntpath.basename(input_file)
+    file_name = os.path.basename(input_file)
     asset_id, extension = os.path.splitext(file_name)
     logger.info(f"working with this asset ID {asset_id}")
     return asset_id, extension
@@ -56,3 +57,19 @@ def run_shell_command(cmd: str) -> bool:
     except Exception:
         logger.exception("Exception")
         return False
+
+
+def save_provenance(provenance: dict, asr_output_dir: str) -> bool:
+    logger.info(f"Saving provenance to: {asr_output_dir}")
+    try:
+        # write provenance.json
+        with open(
+            os.path.join(asr_output_dir, PROVENANCE_JSON_FILE), "w+", encoding="utf-8"
+        ) as f:
+            logger.info(provenance)
+            json.dump(provenance, f, ensure_ascii=False, indent=4)
+    except EnvironmentError as e:  # OSError or IOError...
+        logger.exception(os.strerror(e.errno))
+        return False
+
+    return True

daan_transcript.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import time
 from typing import TypedDict, List, Optional
 from whisper import WHISPER_JSON_FILE
 
@@ -19,12 +20,13 @@ class ParsedResult(TypedDict):
 
 
 # asr_output_dir e.g /data/output/whisper-test/
-def generate_daan_transcript(asr_output_dir: str) -> bool:
+def generate_daan_transcript(asr_output_dir: str) -> Optional[dict]:
     logger.info(f"Generating transcript from: {asr_output_dir}")
+    start_time = time.time()
     whisper_transcript = load_whisper_transcript(asr_output_dir)
     if not whisper_transcript:
         logger.error("No whisper_transcript.json found")
-        return False
+        return None
 
     transcript = parse_whisper_transcript(whisper_transcript)
 
@@ -37,9 +39,21 @@ def generate_daan_transcript(asr_output_dir: str) -> bool:
             json.dump(transcript, f, ensure_ascii=False, indent=4)
     except EnvironmentError as e:  # OSError or IOError...
         logger.exception(os.strerror(e.errno))
-        return False
+        return None
 
-    return True
+    end_time = (time.time() - start_time) * 1000
+    provenance = {
+        "activity_name": "Whisper transcript -> DAAN transcript",
+        "activity_description": "Converts the output of Whisper to the DAAN index format",
+        "processing_time_ms": end_time,
+        "start_time_unix": start_time,
+        "parameters": [],
+        "software_version": "",
+        "input_data": whisper_transcript,
+        "output_data": transcript,
+        "steps": [],
+    }
+    return provenance
 
 
 def load_whisper_transcript(asr_output_dir: str) -> Optional[dict]:

download.py

@@ -17,7 +17,8 @@
 @dataclass
 class DownloadResult:
     file_path: str  # target_file_path,  # TODO harmonize with dane-download-worker
-    mime_type: str  # download_data.get("mime_type", "unknown"),
+    mime_type: str
+    provenance: dict
     download_time: float = -1  # time (ms) taken to receive data after request
     content_length: int = -1  # download_data.get("content_length", -1),
 
@@ -47,11 +48,25 @@ def http_download(url: str) -> Optional[DownloadResult]:
             os.makedirs(input_file_dir)
         with open(input_file, "wb") as file:
             response = requests.get(url)
+            if response.status_code >= 400:
+                logger.error(f"Could not download url: {response.status_code}")
+                return None
             file.write(response.content)
             file.close()
     download_time = (time.time() - start_time) * 1000  # time in ms
+    provenance = {
+        "activity_name": "Input download",
+        "activity_description": "Downloads the input file from INPUT_URI",
+        "processing_time_ms": download_time,
+        "start_time_unix": start_time,
+        "parameters": [],
+        "software_version": "",
+        "input_data": url,
+        "output_data": input_file,
+        "steps": [],
+    }
     return DownloadResult(
-        input_file, mime_type, download_time  # TODO add content_length
+        input_file, mime_type, provenance, download_time  # TODO add content_length
     )
 
 
@@ -86,7 +101,23 @@ def s3_download(s3_uri: str) -> Optional[DownloadResult]:
         if not success:
             logger.error("Failed to download input data from S3")
             return None
-    download_time = (time.time() - start_time) * 1000  # time in ms
+
+        download_time = int((time.time() - start_time) * 1000)  # time in ms
+    else:
+        download_time = -1  # Report back?
+
+    provenance = {
+        "activity_name": "Input download",
+        "activity_description": "Downloads the input file from INPUT_URI",
+        "processing_time_ms": download_time,
+        "start_time_unix": start_time,
+        "parameters": [],
+        "software_version": "",
+        "input_data": s3_uri,
+        "output_data": input_file,
+        "steps": [],
+    }
+
     return DownloadResult(
-        input_file, mime_type, download_time  # TODO add content_length
+        input_file, mime_type, provenance, download_time  # TODO add content_length
     )