annescribe/transcription/job_utils.py

import os
import shutil
from multiprocessing import Lock

# multiprocessing lock for thread safety
lock = Lock()

chunk_input_subdir = 'input_chunks'
chunk_output_subdir = 'output_chunks'


def get_job_data(job_name: str, config: dict) -> dict:
    """
    Get job data from the database.
    :param job_name: name of the job
    :param config: app configuration
    :return: job data
    """

    job_dir = os.path.join(config['jobs']['root_directory'], job_name)

    if not os.path.exists(job_dir):
        raise ValueError(f'Job {job_name} does not exist')

    job_data = dict()
    job_data['name'] = job_name

    output_file = os.path.join(job_dir, 'output.txt')
    if os.path.exists(output_file):
        with open(output_file, 'r') as f:
            job_data['output'] = f.read()
        job_data["completed"] = True
    else:
        job_data['output'] = None
        job_data["completed"] = False

    processing_file = os.path.join(job_dir, '_PROCESSING')
    if os.path.exists(processing_file):
        job_data["processing"] = True
    else:
        job_data["processing"] = False

    # get progress of task
    job_data["progress"] = get_progress(job_name, config)

    return job_data


def get_existing_jobs(config: dict) -> list:
    """
    Get existing jobs from the database.
    :param config: app configuration
    :return: list of jobs
    """

    job_root_dir = config['jobs']['root_directory']

    if not os.path.exists(job_root_dir):
        return []

    jobs = list()
    for folder in os.listdir(job_root_dir):
        job_data = get_job_data(folder, config)
        jobs.append(job_data)
    return jobs


def is_job(job_name: str, config: dict) -> bool:
    """
    Check if a job exists.
    :param job_name: name of the job
    :param config: app configuration
    :return: True if the job exists, False otherwise
    """

    try:
        get_job_data(job_name, config)
        return True
    except ValueError:
        return False


def create_job(job_name: str, config: dict) -> None:
    """
    Create a job in the file system.
    :param job_name: name of the job
    :param config: app configuration
    """
    # check, if job root directory exists
    if not os.path.exists(config['jobs']['root_directory']):
        os.makedirs(config['jobs']['root_directory'])

    job_dir = os.path.join(config['jobs']['root_directory'], job_name)

    if os.path.exists(job_dir):
        raise ValueError(f'Job {job_name} already exists')

    lock.acquire()

    os.mkdir(job_dir)

    lock.release()


def rename_job(job_name: str, new_job_name: str, config: dict) -> None:
    """
    Rename a job in the file system.
    :param job_name: name of the job
    :param new_job_name: new name of the job
    :param config: app configuration
    """
    job_dir = os.path.join(config['jobs']['root_directory'], job_name)
    new_job_dir = os.path.join(config['jobs']['root_directory'], new_job_name)

    if not os.path.exists(job_dir):
        raise ValueError(f'Job {job_name} does not exist')

    if os.path.exists(new_job_dir):
        raise ValueError(f'Job {new_job_name} already exists')

    lock.acquire()

    os.rename(job_dir, new_job_dir)

    lock.release()


def remove_job(job_name: str, config: dict) -> None:
    """
    Remove a job from the file system.
    :param job_name: name of the job
    :param config: app configuration
    """
    job_dir = os.path.join(config['jobs']['root_directory'], job_name)

    if not os.path.exists(job_dir):
        raise ValueError(f'Job {job_name} does not exist')

    lock.acquire()

    # remove the directory
    shutil.rmtree(job_dir)

    lock.release()


def get_processing(job_name: str, config: dict) -> bool:
    """
    Get the processing status of a job.
    :param job_name: name of the job
    :param config: app configuration
    :return: True if the job is processing, False otherwise
    """

    job_dir = os.path.join(config['jobs']['root_directory'], job_name)
    processing_file = os.path.join(job_dir, '_PROCESSING')

    return os.path.exists(processing_file)


def set_processing(job_name: str, config: dict, processing: bool) -> None:
    """
    Set the processing status of a job.
    :param job_name: name of the job
    :param config: app configuration
    :param processing: processing status
    """

    lock.acquire()

    job_dir = os.path.join(config['jobs']['root_directory'], job_name)
    processing_file = os.path.join(job_dir, '_PROCESSING')

    if processing:
        open(processing_file, 'w').close()
    else:
        os.remove(processing_file)

    lock.release()


def get_progress(job_name: str, config: dict) -> int:
    """
    Get the progress of a job.
    :param job_name: name of the job
    :param config: app configuration
    :return: progress of the job
    """

    input_dir = os.path.join(config['jobs']['root_directory'], job_name, chunk_input_subdir)
    output_dir = os.path.join(config['jobs']['root_directory'], job_name, chunk_output_subdir)

    if not os.path.exists(input_dir) or not os.path.exists(output_dir):
        return 0

    num_input_files = len(os.listdir(input_dir))
    num_output_files = len(os.listdir(output_dir))

    if num_input_files == 0 or num_output_files == 0:
        return 0

    if num_output_files >= num_input_files:
        return 100

    return int(num_output_files / num_input_files * 100)