File size: 3,865 Bytes
d97f2ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import gradio as gr
import os
import subprocess
import tempfile
import shutil
from zipfile import ZipFile
import logging
import json
import threading
import psutil

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

connected_cpus = {}

# Function to donate CPU
def donate_cpu(data):
    host = data['host']
    cpu_count = data['cpu_count']
    connected_cpus[host] = {"cpu_count": cpu_count, "usage": 0.0}
    logger.info(f"CPU donated by {host} with {cpu_count} CPUs.")
    return {"status": "success", "message": f"CPU donated by {host}"}

# Function to update CPU usage
def update_cpu_usage(data):
    host = data['host']
    usage = data['usage']
    if host in connected_cpus:
        connected_cpus[host]['usage'] = usage
        logger.info(f"Updated CPU usage for {host}: {usage}%")
    return {"status": "success"}

# Function to run the provided Python script using MPI
def run_script(script_name, folder_path):
    output_log = tempfile.TemporaryFile(mode='w+t')
    try:
        # Collect all available CPUs
        total_cpus = sum(cpu['cpu_count'] for cpu in connected_cpus.values())
        
        # Run the script using MPI
        result = subprocess.run(['mpiexec', '-n', str(total_cpus), 'python', script_name], 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

# Function to handle file uploads and script execution
def handle_upload(folder, script_name):
    # Create a temporary directory to store uploaded files
    temp_dir = tempfile.mkdtemp()
    
    # Save the uploaded folder contents to the temporary directory
    folder_path = os.path.join(temp_dir, 'uploaded_folder')
    os.makedirs(folder_path, exist_ok=True)
    for file_name, file_obj in folder.items():
        with open(os.path.join(folder_path, file_name), 'wb') as f:
            f.write(file_obj.read())

    # Run the script
    log_output = run_script(script_name, folder_path)
    
    # Create a zip file of the entire folder (including any new files created by the script)
    zip_path = os.path.join(temp_dir, 'output_folder.zip')
    with ZipFile(zip_path, 'w') as zipf:
        for root, _, files in os.walk(folder_path):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
    
    return log_output, zip_path

# Function to get connected CPUs information
def get_cpu_info():
    info = []
    for host, data in connected_cpus.items():
        info.append(f"{host}: {data['cpu_count']} CPUs, {data['usage']}% usage")
    return "\n".join(info)

# Gradio interface
def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Python Script Executor with Distributed Computing")
        
        with gr.Row():
            folder = gr.File(label="Upload Folder", file_count="multiple", file_types=['file'])
            script_name = gr.Textbox(label="Python Script Name")
        
        log_output = gr.Textbox(label="Log Output", interactive=False)
        output_folder = gr.File(label="Download Output Folder")
        cpu_info = gr.Textbox(label="Connected CPUs Info", interactive=False)
        
        run_button = gr.Button("Run Script")
        refresh_button = gr.Button("Refresh CPU Info")
        
        run_button.click(fn=handle_upload, inputs=[folder, script_name], outputs=[log_output, output_folder])
        refresh_button.click(fn=get_cpu_info, inputs=[], outputs=[cpu_info])
        
        # Define the donate CPU endpoint
        demo.api(donate_cpu, inputs=gr.JSON(), outputs=gr.JSON(), name="donate_cpu")
    
    demo.launch()

if __name__ == "__main__":
    gradio_interface()