78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
|
|
"""
|
|
This file contains definitions for the LocalPythonConductor, in order to
|
|
execute Python jobs on the local resource.
|
|
|
|
Author(s): David Marchant
|
|
"""
|
|
import os
|
|
import shutil
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Tuple
|
|
|
|
from core.correctness.vars import JOB_TYPE_PYTHON, PYTHON_FUNC, JOB_STATUS, \
|
|
STATUS_RUNNING, JOB_START_TIME, PYTHON_EXECUTION_BASE, JOB_ID, META_FILE, \
|
|
STATUS_DONE, JOB_END_TIME, STATUS_FAILED, JOB_ERROR, PYTHON_OUTPUT_DIR, \
|
|
JOB_TYPE, JOB_TYPE_PAPERMILL
|
|
from core.correctness.validation import valid_job
|
|
from core.functionality import read_yaml, write_yaml
|
|
from core.meow import BaseConductor
|
|
|
|
|
|
# TODO add comments to me
|
|
class LocalPythonConductor(BaseConductor):
|
|
def __init__(self)->None:
|
|
super().__init__()
|
|
|
|
def valid_execute_criteria(self, job:dict[str,Any])->Tuple[bool,str]:
|
|
"""Function to determine given an job defintion, if this conductor can
|
|
process it or not. This conductor will accept any Python job type"""
|
|
try:
|
|
valid_job(job)
|
|
if job[JOB_TYPE] in [JOB_TYPE_PYTHON, JOB_TYPE_PAPERMILL]:
|
|
return True, ""
|
|
except Exception as e:
|
|
pass
|
|
return False, str(e)
|
|
|
|
def execute(self, job:dict[str,Any])->None:
|
|
valid_job(job)
|
|
|
|
job_dir = os.path.join(job[PYTHON_EXECUTION_BASE], job[JOB_ID])
|
|
meta_file = os.path.join(job_dir, META_FILE)
|
|
|
|
# update the status file with running status
|
|
job[JOB_STATUS] = STATUS_RUNNING
|
|
job[JOB_START_TIME] = datetime.now()
|
|
write_yaml(job, meta_file)
|
|
|
|
# execute the job
|
|
try:
|
|
job_function = job[PYTHON_FUNC]
|
|
job_function(job)
|
|
|
|
# get up to date job data
|
|
job = read_yaml(meta_file)
|
|
|
|
# Update the status file with the finalised status
|
|
job[JOB_STATUS] = STATUS_DONE
|
|
job[JOB_END_TIME] = datetime.now()
|
|
write_yaml(job, meta_file)
|
|
|
|
except Exception as e:
|
|
# get up to date job data
|
|
job = read_yaml(meta_file)
|
|
|
|
# Update the status file with the error status
|
|
job[JOB_STATUS] = STATUS_FAILED
|
|
job[JOB_END_TIME] = datetime.now()
|
|
msg = f"Job execution failed. {e}"
|
|
job[JOB_ERROR] = msg
|
|
write_yaml(job, meta_file)
|
|
|
|
# Move the contents of the execution directory to the final output
|
|
# directory.
|
|
job_output_dir = os.path.join(job[PYTHON_OUTPUT_DIR], job[JOB_ID])
|
|
shutil.move(job_dir, job_output_dir)
|