Update app.py

app.py

@@ -56,7 +56,7 @@ def handle_upload():
             script_file.write(script_content)
         
         # Run the script using MPI
-        log_output = run_script(script_path, folder_path)
+        log_output = run_script_with_mpi(script_path, folder_path)
         
         # Create a zip file of the entire folder
         zip_path = os.path.join(temp_dir, 'output_folder.zip')
@@ -154,6 +154,31 @@ def update_cpu_usage():
 def index():
     return render_template('index.html')
 
+def run_script_with_mpi(script_path, folder_path):
+    # Create a temporary directory for MPI processes
+    mpi_temp_dir = tempfile.mkdtemp()
+    mpi_script_path = os.path.join(mpi_temp_dir, 'mpi_script.py')
+    
+    # Write the MPI script to the temporary directory
+    with open(mpi_script_path, 'w') as mpi_script_file:
+        mpi_script_file.write(f"""
+import os
+import tempfile
+import subprocess
+from mpi4py import MPI
+
+def target_function(script_path, folder_path):
+    output_log = tempfile.TemporaryFile(mode='w+t')
+    try:
+        result = subprocess.run(['python', script_path], cwd=folder_path, stdout=output_log, stderr=subprocess.STDOUT)
+        output_log.seek(0)
+        log_output = output_log.read()
+    except Exception as e:
+        log_output = str(e)
+    finally:
+        output_log.close()
+    return log_output
+
 def run_script(script_path, folder_path):
     comm = MPI.COMM_WORLD
     rank = comm.Get_rank()
@@ -165,11 +190,29 @@ def run_script(script_path, folder_path):
         for i in range(1, size):
             log_output = comm.recv(source=i, tag=11)
             log_outputs.append(log_output)
-        return '\n'.join(log_outputs)
+        with open(os.path.join(folder_path, 'mpi_log_output.txt'), 'w') as log_file:
+            log_file.write('\\n'.join(log_outputs))
     else:
         # Worker process
         log_output = target_function(script_path, folder_path)
         comm.send(log_output, dest=0, tag=11)
 
+if __name__ == "__main__":
+    run_script('{script_path}', '{folder_path}')
+""")
+    
+    # Run the MPI script using subprocess
+    result = subprocess.run(['mpiexec', '-n', str(psutil.cpu_count(logical=False)), 'python', mpi_script_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    
+    # Read the log output from the file
+    log_output_path = os.path.join(folder_path, 'mpi_log_output.txt')
+    with open(log_output_path, 'r') as log_file:
+        log_output = log_file.read()
+    
+    # Clean up the temporary directory
+    shutil.rmtree(mpi_temp_dir)
+    
+    return log_output
+
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860, threaded=True)