compute_pool / app.py
oscarwang2's picture
Update app.py
889f571 verified
raw
history blame
2.2 kB
from flask import Flask, request, jsonify
import os
import subprocess
import tempfile
import json
import multiprocessing
app = Flask(__name__)
# Dictionary to store donated CPU information
donated_cpus = {}
total_donated_cpus = 0
@app.route('/donate_cpu', methods=['POST'])
def donate_cpu():
global total_donated_cpus
data = request.get_json()
host = data['host']
cpu_count = data['cpu_count']
donated_cpus[host] = {"cpu_count": cpu_count, "usage": 0.0}
total_donated_cpus += cpu_count
return jsonify({"status": "success", "message": f"CPU donated by {host}"})
@app.route('/update_cpu_usage', methods=['POST'])
def update_cpu_usage():
data = request.get_json()
host = data['host']
usage = data['usage']
if host in donated_cpus:
donated_cpus[host]['usage'] = usage
return jsonify({"status": "success", "message": f"Updated CPU usage for {host}: {usage}%"})
return jsonify({"status": "error", "message": "Host not found"}), 404
@app.route('/execute_script', methods=['POST'])
def execute_script():
try:
data = request.get_json()
script_content = data['script_content']
# Save the script to a temporary file
script_file = tempfile.NamedTemporaryFile(delete=False, suffix='.py')
script_file.write(script_content.encode())
script_file.close()
# Run the script using the donated CPUs
cpu_count = total_donated_cpus or 1 # Default to 1 if no CPUs are donated
result = subprocess.run(['python', script_file.name], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=60)
os.remove(script_file.name)
return jsonify({
"stdout": result.stdout,
"stderr": result.stderr,
"returncode": result.returncode
})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route('/cpu_info', methods=['GET'])
def cpu_info():
info = [{"host": host, "cpu_count": data['cpu_count'], "usage": data['usage']} for host, data in donated_cpus.items()]
return jsonify({"cpu_info": info})
if __name__ == "__main__":
app.run(host='0.0.0.0', port=7860)