Update app.py

app.py

@@ -6,7 +6,7 @@ import logging
 import psutil
 import subprocess
 from flask import Flask, request, jsonify, render_template, send_file
-import multiprocessing
+from mpi4py import MPI
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -15,9 +15,9 @@ logger = logging.getLogger(__name__)
 # Initialize Flask app
 app = Flask(__name__)
 
-connected_cpus = {"localhost": {"cpu_count": psutil.cpu_count(logical=False), "usage": 0.0}}
+connected_cpus = {"localhost": {"cpu_count": psutil.cpu_count(logical=False), "usage": psutil.cpu_percent(interval=1)}}
 
-# Define the target function for multiprocessing
+# Define the target function for MPI
 def target_function(script_path, folder_path):
     output_log = tempfile.TemporaryFile(mode='w+t')
     try:
@@ -55,7 +55,7 @@ def handle_upload():
         with open(script_path, 'w') as script_file:
             script_file.write(script_content)
         
-        # Run the script using multiprocessing
+        # Run the script using MPI
         log_output = run_script(script_path, folder_path)
         
         # Create a zip file of the entire folder
@@ -85,7 +85,8 @@ def get_cpu_info():
     try:
         info = []
         for host, data in connected_cpus.items():
-            info.append(f"{host}: {data['cpu_count']} CPUs, {data['usage']}% usage")
+            usage = psutil.cpu_percent(interval=1) if host == "localhost" else data['usage']
+            info.append(f"{host}: {data['cpu_count']} CPUs, {usage}% usage")
         return jsonify({"status": "success", "cpu_info": "\n".join(info)})
     except Exception as e:
         logger.error(f"Error in get_cpu_info: {e}")
@@ -154,14 +155,21 @@ def index():
     return render_template('index.html')
 
 def run_script(script_path, folder_path):
-    # Collect all available CPUs including the local host CPU
-    total_cpus = sum(cpu['cpu_count'] for cpu in connected_cpus.values()) + 1
-    
-    # Use multiprocessing to run the script
-    with multiprocessing.Pool(total_cpus) as pool:
-        log_outputs = pool.starmap(target_function, [(script_path, folder_path)] * total_cpus)
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
 
-    return '\n'.join(log_outputs)
+    if rank == 0:
+        # Master process
+        log_outputs = []
+        for i in range(1, size):
+            log_output = comm.recv(source=i, tag=11)
+            log_outputs.append(log_output)
+        return '\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__":
     app.run(host='0.0.0.0', port=7860, threaded=True)