app.py

import os
import json
import random
import shutil
import hashlib
import datetime
from flask import Flask, request

def get_key_sha256(key):
    sha256 = hashlib.sha256()
    sha256.update(key.encode('utf-8'))
    return sha256.hexdigest()

match_sha256 = 'ffa007286cff75ab9c14e42282edd8e7e28b467a6e1bcd3d9b90dc40d375cae3'
data_folder = 'jobs'
stages = ['new', 'doing', 'done']
for stage in stages:
    os.makedirs(os.path.join(data_folder, stage), exist_ok=True)

app = Flask(__name__)

@app.route("/add", methods = ['POST'])
def put():
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    machine_type = d.get('machine_type', 'default')
    job = d['job']
    sha256 = get_key_sha256(job)
    folder = os.path.join(data_folder, 'new', machine_type)
    doing_folder = os.path.join(data_folder, 'doing', machine_type)
    done_folder = os.path.join(data_folder, 'done', machine_type)
    os.makedirs(folder, exist_ok=True)
    os.makedirs(doing_folder, exist_ok=True)
    os.makedirs(done_folder, exist_ok=True)
    file = os.path.join(folder, sha256)
    job_json = {'creat_time': str(datetime.datetime.utcnow()), 'job': job, 'machine_type': machine_type}
    open(file, 'w').write(json.dumps(job_json))
    return 'ok'


@app.route("/get", methods = ['GET'])
def get():
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    machine_type = d['machine_type']
    worker = d['worker']
    folder = os.path.join(data_folder, 'new', machine_type)
    if not os.path.isdir(folder):
        return ''
    doing_folder = os.path.join(data_folder, 'doing', machine_type)
    files = os.listdir(folder)
    if len(files) == 0:
        return ''
    file = random.choice(files)
    source_file = os.path.join(folder, file)
    target_file = os.path.join(doing_folder, file)
    txt = open(source_file).read()
    job_json = json.loads(txt)
    job = job_json['job']
    job_json['worker'] = worker
    json.dump(job_json, open(target_file, 'w'))
    try:
        os.remove(source_file)
    except:
        print('deleteerr')
    return json.dumps({'sha256': file, 'job': job, 'start_time': str(datetime.datetime.utcnow())})

@app.route("/done", methods = ['POST'])
def done():
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    sha256 = d['sha256']
    machine_type = d['machine_type']
    output = d['output']
    dt = d['dt']
    #start_time = d['start_time'] # could be found in doing file
    end_time = str(datetime.datetime.utcnow())
    src = os.path.join(data_folder, 'doing', machine_type, sha256)
    dst = os.path.join(data_folder, 'done', machine_type, sha256)
    job_json = json.load(open(src))
    job_json['dt'] = dt
    job_json['end_time'] = end_time
    dst_output = dst+'.output'
    json.dump(job_json, open(dst, 'w'))
    open(dst_output, 'w').write(output)
    try:
        os.remove(src)
    except:
        return 'derr'
    return 'ok'


def get_folder_list(folder, full=False):
    machine_types = os.listdir(folder)
    d = {}
    for machine_type in machine_types:
        fs = os.listdir(os.path.join(folder, machine_type))
        dd = {}
        for f in fs:
            if full:
                dd[f] = open(os.path.join(folder, machine_type, f)).read()
            else:
                dd[f] = open(os.path.join(folder, machine_type, f)).read(10)
        d[machine_type] = dd
    return d

@app.route("/stat", methods = ['GET'])
def stat():
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    full = int(d['full'])
    stat = {}
    stat['new'] = get_folder_list(os.path.join(data_folder, 'new'), full)
    stat['doing'] = get_folder_list(os.path.join(data_folder, 'doing'), full)
    stat['done'] = get_folder_list(os.path.join(data_folder, 'done'), full)
    return json.dumps(stat) 

@app.route("/mv", methods = ['POST'])
def mv():
    '''
    move from doing to new
    '''
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    machine_type = d['machine_type']
    sha256 = d['sha256']
    src = os.path.join(data_folder, 'doing', machine_type, sha256)
    dst = os.path.join(data_folder, 'new', machine_type, sha256)
    os.system(f'mv {src} {dst}')
    return 'ok'


@app.route("/clear", methods = ['POST'])
def clear():
    '''
    clear all of new, doing, or done
    '''
    d = request.json
    key = d['key']
    if not get_key_sha256(key) == match_sha256:
        return 'err'
    what = d['what']
    folder = os.path.join(data_folder, what)
    shutil.rmtree(folder)
    os.makedirs(folder)
    return 'ok'


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)