Update app.py

app.py

@@ -241,21 +241,6 @@ PandasTools.drawOptions.singleColourWedgeBonds = True
 PandasTools.drawOptions.useCDKAtomPalette()
 PandasTools.molSize = (100, 64)
 
-session = requests.Session()
-ADAPTER = HTTPAdapter(max_retries=Retry(total=5, backoff_factor=0.1, status_forcelist=[500, 502, 503, 504]))
-session.mount('http://', ADAPTER)
-session.mount('https://', ADAPTER)
-
-db = TinyDB(f'{SERVER_DATA_DIR}/db.json')
-# Set all RUNNING jobs to FAILED at TinyDB initialization
-Job = Query()
-jobs = db.all()
-for job in jobs:
-    if job['status'] == 'RUNNING':
-        db.update({'status': 'FAILED'}, Job.id == job['id'])
-
-scheduler = BackgroundScheduler()
-
 
 def remove_job_record(job_id):
     # Delete the job from the database
@@ -637,6 +622,7 @@ def ts_to_str(timestamp, timezone):
 def lookup_job(job_id):
     gr.Info('Start querying the job database...')
     stop = False
+    retry = 0
     while not stop:
         try:
             Job = Query()
@@ -655,11 +641,11 @@ def lookup_job(job_id):
                     sleep(5)
                     yield {
                         pred_lookup_status: f'''
-Your **{job_type}** job (ID: {job_id}) started at 
+Your **{job_type}** job (ID: **{job_id}**) started at 
 **{start_time}** and is **RUNNING...**
 
 It might take a few minutes up to a few hours depending on the prediction dataset, the model, and the queue status.
-You may keep the page open and wait for the completion, or close the page and revisit later to look up the job status 
+You may keep the page open and wait for job completion, or close the page and revisit later to look up the job status 
 using the job id. You will also receive an email notification once the job is done.
 ''',
                         pred_lookup_btn: gr.Button(visible=False),
@@ -693,9 +679,11 @@ using the job id. You will also receive an email notification once the job is do
                         tabs: gr.Tabs(selected='Prediction Status Lookup'),
                     }
             else:
-                stop = True
-                msg = f'Job ID {job_id} not found.'
+                stop = (retry > 3)
+                msg = f'Job ID {job_id} not found. Retrying... ({retry})'
                 gr.Info(msg)
+                retry += 1
+                sleep(5)
                 yield {
                     pred_lookup_status: msg,
                     pred_lookup_btn: gr.Button(visible=True),
@@ -707,9 +695,11 @@ using the job id. You will also receive an email notification once the job is do
             raise gr.Error(f'Failed to retrieve job status due to error: {str(e)}')
 
 
-def submit_predict(predict_filepath, task, preset, target_family, opts, state):
-    job_id = state['id']
-    status = "RUNNING"
+def submit_predict(predict_filepath, task, preset, target_family, opts, job_info):
+    job_id = job_info['id']
+    status = job_info['status']
+    send_email(job_info)
+    db.insert(job_info)
     error = None
     task_file_abbr = {'Compound-Protein Interaction': 'CPI', 'Compound-Protein Binding Affinity': 'CPA'}
     predictions_file = None
@@ -768,7 +758,7 @@ def submit_predict(predict_filepath, task, preset, target_family, opts, state):
             task_value = TASK_MAP[task]
             preset_value = PRESET_MAP[preset]
             predictions_file = (f'{SERVER_DATA_DIR}/'
-                                f'{job_id}_{task_file_abbr[task]}_{preset_value}_{target_family_value}_predictions.csv')
+                                f'{job_id}_{task_file_abbr[task]}_{preset}_{target_family_value}_predictions.csv')
 
             cfg = hydra.compose(
                 config_name="webserver_inference",
@@ -884,7 +874,7 @@ def submit_predict(predict_filepath, task, preset, target_family, opts, state):
                    'error': error,
                    'input_file': predict_filepath,
                    'output_file': predictions_file},
-                  job_query)
+                   job_query)
         if job_info := db.search(job_query)[0]:
             if job_info.get('email'):
                 send_email(job_info)
@@ -2156,7 +2146,7 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
 
     def common_job_initiate(job_id, job_type, email, request, task):
         gr.Info('Finished processing inputs. Initiating the prediction job... '
-                'You will be redirected to Prediction Status Lookup after the job is submitted.')
+                'You will be redirected to Prediction Status Lookup once the job has been submitted.')
         job_info = {'id': job_id,
                     'type': job_type,
                     'task': task,
@@ -2168,7 +2158,7 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
                     'end_time': None,
                     'expiry_time': None,
                     'error': None}
-        db.insert(job_info)
+        # db.insert(job_info)
         return job_info
 
 
@@ -2314,10 +2304,10 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
         concurrency_limit=100,
     )
 
-    drug_screen_click.success(
-        fn=send_email,
-        inputs=[run_state]
-    )
+    # drug_screen_click.success(
+    #     fn=send_email,
+    #     inputs=[run_state]
+    # )
 
     drug_screen_click.success(
         fn=submit_predict,
@@ -2404,10 +2394,10 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
         concurrency_limit=100
     )
 
-    target_identify_click.success(
-        fn=send_email,
-        inputs=[run_state]
-    )
+    # target_identify_click.success(
+    #     fn=send_email,
+    #     inputs=[run_state]
+    # )
 
     target_identify_click.success(
         fn=submit_predict,
@@ -2436,10 +2426,10 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
         concurrency_limit=100
     )
 
-    pair_infer_click.success(
-        fn=send_email,
-        inputs=[run_state]
-    )
+    # pair_infer_click.success(
+    #     fn=send_email,
+    #     inputs=[run_state]
+    # )
 
     pair_infer_click.success(
         fn=submit_predict,
@@ -2568,7 +2558,24 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
 
 if __name__ == "__main__":
     pandarallel.initialize()
+    
     hydra.initialize(version_base="1.3", config_path="configs", job_name="webserver_inference")
-    demo.queue(default_concurrency_limit=None, max_size=10).launch(show_api=False)
+    
+    session = requests.Session()
+    ADAPTER = HTTPAdapter(max_retries=Retry(total=5, backoff_factor=0.1, status_forcelist=[500, 502, 503, 504]))
+    session.mount('http://', ADAPTER)
+    session.mount('https://', ADAPTER)
+    
+    db = TinyDB(f'{SERVER_DATA_DIR}/db.json')
+    # Set all RUNNING jobs to FAILED at TinyDB initialization
+    Job = Query()
+    jobs = db.all()
+    for job in jobs:
+        if job['status'] == 'RUNNING':
+            db.update({'status': 'FAILED'}, Job.id == job['id'])
+    
+    scheduler = BackgroundScheduler()
     scheduler.add_job(check_expiry, 'interval', hours=1)
     scheduler.start()
+    
+    demo.queue(default_concurrency_limit=None, max_size=10).launch(show_api=False)