Update app.py

app.py

@@ -81,8 +81,10 @@ def predict_structure(prefix, feature_dict, model_runners, random_seed=0):
             f.write(protein.to_pdb(unrelaxed_protein))
     return plddts
 
-
+@ray.remote
 def run_protgpt2(startsequence, length, repetitionPenalty, top_k_poolsize, max_seqs):
+    print("running protgpt2"
+    print(gpu_usage())
     protgpt2 = pl("text-generation", model="nferruz/ProtGPT2")
     sequences = protgpt2(
         startsequence,
@@ -94,14 +96,14 @@ def run_protgpt2(startsequence, length, repetitionPenalty, top_k_poolsize, max_s
         eos_token_id=0,
     )
     print("Cleaning up after protGPT2")
-    print(gpu_usage())
+    #print(gpu_usage())
     #torch.cuda.empty_cache()
-    device = cuda.get_current_device()
-    device.reset()
-    print(gpu_usage())
+    #device = cuda.get_current_device()
+    #device.reset()
+    #print(gpu_usage())
     return sequences
 
-
+@ray.remote
 def run_alphafold(startsequence):
     print(gpu_usage())
     model_runners = {}
@@ -124,8 +126,7 @@ def run_alphafold(startsequence):
         **mk_mock_template(query_sequence),
     }
     plddts = predict_structure("test", feature_dict, model_runners)
-    print("Cleaning up after AF2")
-    print(gpu_usage())
+    print("AF2 done")
     #backend = jax.lib.xla_bridge.get_backend()
     #for buf in backend.live_buffers(): buf.delete()
     #device = cuda.get_current_device()
@@ -137,7 +138,7 @@ def run_alphafold(startsequence):
 def update_protGPT2(inp, length,repetitionPenalty, top_k_poolsize, max_seqs):
     startsequence = inp
     seqlen = length
-    generated_seqs = run_protgpt2(startsequence, seqlen, repetitionPenalty, top_k_poolsize, max_seqs)
+    generated_seqs = ray.get(run_protgpt2.remote(startsequence, seqlen, repetitionPenalty, top_k_poolsize, max_seqs))
     gen_seqs = [x["generated_text"] for x in generated_seqs]
     print(gen_seqs)
     sequencestxt = ""
@@ -151,7 +152,8 @@ def update_protGPT2(inp, length,repetitionPenalty, top_k_poolsize, max_seqs):
 def update(inp):
     print("Running AF on", inp)
     startsequence = inp
-    plddts = run_alphafold(startsequence)
+    # run alphafold using ray
+    plddts = ray.get(run_alphafold.remote(startsequence))
     print(plddts)
     x = np.arange(10)
     #plt.style.use(["seaborn-ticks", "seaborn-talk"])