Upload app.py with huggingface_hub

app.py

@@ -1,7 +1,9 @@
 import os
 import json
 import random
+import shutil
 import hashlib
+import datetime
 from flask import Flask, request
 
 def get_key_sha256(key):
@@ -10,7 +12,7 @@ def get_key_sha256(key):
     return sha256.hexdigest()
 
 match_sha256 = 'ffa007286cff75ab9c14e42282edd8e7e28b467a6e1bcd3d9b90dc40d375cae3'
-
+match_sha256 = 'ffa007286cff75ab9c14e42282edd8e7e28b467a6e1bcd3d9b90dc40d375cae3'
 data_folder = 'jobs'
 stages = ['new', 'doing', 'done']
 for stage in stages:
@@ -34,7 +36,8 @@ def put():
     os.makedirs(doing_folder, exist_ok=True)
     os.makedirs(done_folder, exist_ok=True)
     file = os.path.join(folder, sha256)
-    open(file, 'w').write(job)
+    job_json = {'creat_time': str(datetime.datetime.utcnow()), 'job': job, 'machine_type': machine_type}
+    open(file, 'w').write(json.dumps(job_json))
     return 'ok'
 
 
@@ -45,6 +48,7 @@ def get():
     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 ''
@@ -56,12 +60,15 @@ def get():
     source_file = os.path.join(folder, file)
     target_file = os.path.join(doing_folder, file)
     txt = open(source_file).read()
-    open(target_file, 'w').write(txt)
+    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': txt})
+    return json.dumps({'sha256': file, 'job': job, 'start_time': str(datetime.datetime.utcnow())})
 
 @app.route("/done", methods = ['POST'])
 def done():
@@ -72,9 +79,19 @@ def done():
     sha256 = d['sha256']
     machine_type = d['machine_type']
     output = d['output']
-    file = os.path.join(data_folder, 'doing', machine_type, sha256)
+    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(file)
+        os.remove(src)
     except:
         return 'derr'
     return 'ok'
@@ -104,8 +121,42 @@ def stat():
     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__":