call rfd3 with process.Popen to allow cancellation

app.py

@@ -13,6 +13,7 @@ from utils.pipelines import *
 #from gradio_molecule3d import Molecule3D
 from utils.handle_events import *
 from utils.handle_files import *
+import signal
 
 download_weights()
 
@@ -66,6 +67,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
 
     gen_directory = gr.State(None) # the directory where generation results are saved, used to trigger the download of results as zip file
     gen_results = gr.State(None) # the results of the generation, which is a list of dicts where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".
+    process = gr.State(None) # to store the subprocess.Popen object for the generation process, used to terminate the process when the stop button is clicked.
 
     # inputs from user
     with gr.Row():
@@ -131,19 +133,13 @@ with gr.Blocks(title="RFD3 Test") as demo:
     display_state = gr.Textbox(label="Selected Batch and Design", visible=True)
     display_state.value = "Please Select a Batch and Design number to show sequence"
     
-    def generate(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration):
-        if config_upload is None:
-            return gr.update(), None, None
-        else:
-            textbox_update, gen_directory, gen_results = generation_with_input_config(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration)
-            print(textbox_update)
-            return textbox_update, gen_directory, gen_results
+
         
     generation_event = run_btn.click(
         lambda: (gr.update(visible=False), gr.update(visible=True)),
         outputs=[run_btn, stop_btn]
     ).then(
-        generate, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_directory, gen_results]
+        generation_with_input_config, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_directory, gen_results, process]
     ).then(
         update_batch_choices,
         inputs=gen_results,
@@ -156,9 +152,19 @@ with gr.Blocks(title="RFD3 Test") as demo:
         lambda: (gr.update(visible=True), gr.update(visible=False)),
         outputs=[run_btn, stop_btn]
     )
+
+    def stop_generation(proc_state):
+        #if proc_state is not None and proc_state.poll() is None:
+        #    try:
+        #        # Kill whole process group so child processes die too
+        #        os.killpg(os.getpgid(proc_state.pid), signal.SIGTERM)
+        #    except Exception:
+        #        proc_state.terminate()
+        return gr.update(value="Generation cancelled by user.") 
     
     stop_btn.click(
-        lambda: gr.update(value="Generation cancelled by user."),
+        stop_generation,
+        inputs=process,
         outputs=runtextbox,
         cancels=[generation_event]
     ).then(
@@ -183,5 +189,6 @@ with gr.Blocks(title="RFD3 Test") as demo:
     #visualize_btn.click(load_viewer, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=viewer)
 
 if __name__ == "__main__":
+
     demo.queue()
     demo.launch()

utils/pipelines.py

@@ -42,7 +42,7 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
 
 
     if input_file is None:        
-        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json"), gr.update(), gr.update()
+        return gr.update(value="Please ensure you have uploaded a configuration file: .yaml or .json"), gr.update(), gr.update(), gr.update()
     elif pdb_file is None:
         status_update = f"Running generation for {num_batches} batches of {num_designs_per_batch}\n job configuration uploaded from file {os.path.basename(input_file)}\n no scaffold/target provided"
     else:
@@ -53,6 +53,7 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
     directory = f"./outputs/generation_with_input_config/session_{session_hash}_{time_stamp}"
     os.makedirs(directory, exist_ok=False)
 
+    process = None
     try: 
         if pdb_file is not None:
             # I need to do this because uploading files to a HF space stores each file in a separate temp directory so I need to copy them again to the same place.
@@ -70,11 +71,37 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
         print(f"Running command: {command}")
         status_update += f"\nRunning command: {command}."
         start = perf_counter()
-        res = subprocess.run(command, shell=True, check=True, text=True, capture_output=True)
-        print("Command took", perf_counter() - start, "seconds to run.")
-        status_update += f"\nGeneration successful! Command took {perf_counter() - start:.2f} seconds to run."
-
+        
+        # Start the process with Popen to allow cancellation
+        process = subprocess.Popen(
+            command,
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            preexec_fn=os.setsid,  # For Unix-like systems, start a new process group
+        )
+        
+        # Immediately yield to make process available in State
+        yield gr.update(value=status_update), gr.update(), gr.update(), process
+        
+        # Poll the process - yield regularly to allow Gradio to stop calling the generator
+        while process.poll() is None:
+            time.sleep(2)  # Wait before checking again
+            yield gr.update(value=status_update), gr.update(), gr.update(), process
+        
+        # Get the output after process completes
+        stdout, stderr = process.communicate()
+        
+        # Check if process succeeded
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(process.returncode, command, stdout, stderr)
+        
+        elapsed_time = perf_counter() - start
+        print(f"Command took {elapsed_time:.2f} seconds to run.")
+        status_update += f"\nGeneration successful! Command took {elapsed_time:.2f} seconds to run."
 
+        # Collect results from output directory
         results = []
         for file_name in os.listdir(directory):
             if file_name.endswith(".cif.gz"):
@@ -88,11 +115,28 @@ def generation_with_input_config(input_file, pdb_file, num_batches, num_designs_
                 results.append({"batch": batch, "design": design, "cif_path": cif_path, "pdb_path": pdb_path})
             
         print(results)
-        return gr.update(value=status_update), directory, results
+        return gr.update(value=status_update), directory, results, process
     
     except subprocess.CalledProcessError as e:
         print("subprocess threw an error", e.stderr)
-        return gr.update(value=f"Generation failed:\n{e.stderr}"), None, None
+        return gr.update(value=f"Generation failed:\n{e.stderr}"), None, None, None
+    
+    except Exception as e:
+        # Handle cancellation or other errors - terminate the subprocess if it's still running
+        if process and process.poll() is None:
+            print(f"Terminating subprocess due to: {e}")
+            try:
+                os.killpg(os.getpgid(process.pid), 15)  # SIGTERM to process group
+            except Exception:
+                process.terminate()
+            try:
+                process.wait(timeout=5)  # Wait up to 5 seconds for graceful termination
+            except subprocess.TimeoutExpired:
+                try:
+                    os.killpg(os.getpgid(process.pid), 9)  # SIGKILL to process group
+                except Exception:
+                    process.kill()  # Force kill if killpg fails
+        raise  # Re-raise the exception so Gradio knows the event was cancelled
 
 
 #def generation_with_input_config_factory(max_duration):