bug naosita

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, render_template, send_file, jsonify, send_from_directory, session
+from flask import Flask, request, render_template, send_file, jsonify, send_from_directory, session, copy_current_request_context
 from flask_socketio import SocketIO, join_room, leave_room, close_room, rooms, disconnect
 from flask_cors import CORS
 from flask_limiter import Limiter
@@ -23,9 +23,10 @@ from gevent import pywsgi
 from geventwebsocket.handler import WebSocketHandler
 
 app = Flask(__name__)
-app.secret_key = 'user'
 CORS(app)
 socketio = SocketIO(app, cors_allowed_origins="*")
+# クライアントIDとルーム情報を保存するグローバル辞書
+client_rooms = {}
 
 # レート制限の設定
 limiter = Limiter(
@@ -41,10 +42,12 @@ active_tasks = {}
 task_futures = {}
 
 # ThreadPoolExecutorの作成
-executor = concurrent.futures.ThreadPoolExecutor(max_workers=8) # 8vCPUのインスタンスを使用
+executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) # 8vCPUのインスタンスを使用
+
+gpu_lock = threading.Lock()
 
 class Task:
-    def __init__(self, task_id, mode, weight1, weight2, file_data, client_ip):
+    def __init__(self, task_id, mode, weight1, weight2, file_data, client_ip, client_id):
         self.task_id = task_id
         self.mode = mode
         self.weight1 = weight1
@@ -53,16 +56,14 @@ class Task:
         self.cancel_flag = False
         self.client_ip = client_ip
         self.is_processing = False
+        self.client_id = client_id
 
 # キューの状態を通知
 def update_queue_status(message):
-    socketio.emit('queue_update', {'active_tasks': len(active_tasks), 'message': message})
+    socketio.emit('queue_update', {'active_tasks': len(active_tasks), 'message': message}, namespace='/demo')
 
 def process_task(task):
     try:
-        update_queue_status("sid")
-        update_queue_status(str("sid" in session))
-        client_id = session.get("sid")
         task.is_processing = True
         # ファイルデータをPIL Imageに変換
         image = Image.open(io.BytesIO(task.file_data))
@@ -73,40 +74,49 @@ def process_task(task):
             return
 
         # 画像処理ロジックを呼び出す
-        sotai_image, sketch_image = process_image_as_base64(image, task.mode, task.weight1, task.weight2)
+        # GPU処理部分
+        with gpu_lock:
+            sotai_image, sketch_image = process_image_as_base64(image, task.mode, task.weight1, task.weight2)
 
         # キャンセルチェック
         if task.cancel_flag:
             return
-
-        socketio.emit('task_complete', {
-            'task_id': task.task_id, 
-            'sotai_image': sotai_image, 
-            'sketch_image': sketch_image
-        }, to=client_id)
+        
+        # クライアントIDをリクエストヘッダーから取得（クライアント側で設定する必要があります）
+        client_id = task.client_id
+        if client_id and client_id in client_rooms:
+            room = client_rooms[client_id]
+            # ルームにメッセージをemit
+            socketio.emit('task_complete', {
+                'task_id': task.task_id, 
+                'sotai_image': sotai_image, 
+                'sketch_image': sketch_image
+            }, to=room, namespace='/demo')
+        
     except Exception as e:
         print(f"Task error: {str(e)}")
         if not task.cancel_flag:
-            socketio.emit('task_error', {'task_id': task.task_id, 'error': str(e)}, to=client_id)
+            client_id = task.client_id
+            room = client_rooms[client_id]
+            socketio.emit('task_error', {'task_id': task.task_id, 'error': str(e)}, to=room, namespace='/demo')
     finally:
-        task.is_processing = False
-        if task.task_id in active_tasks:
-            del active_tasks[task.task_id]
-        if task.task_id in task_futures:
-            del task_futures[task.task_id]
-            
         # タスク数をデクリメント
         client_ip = task.client_ip
         tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
-        
+        print(f'Task {task.task_id} completed')
+        task.is_processing = False
+        if task.task_id in active_tasks.keys():
+            del active_tasks[task.task_id]
+        if task.task_id in task_futures.keys():
+            del task_futures[task.task_id]
+
         update_queue_status('Task completed or cancelled')
 
 def worker():
     while True:
         try:
             task = task_queue.get()
-            if task.task_id in active_tasks:
-                print(f"Processing task {task.task_id}")
+            if task.task_id in active_tasks.keys():
                 future = executor.submit(process_task, task)
                 task_futures[task.task_id] = future
         except Exception as e:
@@ -121,23 +131,31 @@ threading.Thread(target=worker, daemon=True).start()
 # グローバル変数を使用して接続数とタスク数を管理
 connected_clients = 0
 tasks_per_client = {}
-@socketio.on('connect', namespace='/')
+@socketio.on('connect', namespace='/demo')
 def handle_connect(auth):
-    session["sid"] = request.sid
-    client_id = request.sid  # クライアントIDを取得
-    join_room(client_id)  # クライアントを自身のルームに入れる
+    client_id = request.sid
+    room = f"room_{client_id}"  # クライアントごとに一意のルーム名を生成
+    join_room(room)
+    client_rooms[client_id] = room
+    print(f"Client {client_id} connected and joined room {room}")
+    
     global connected_clients
     connected_clients += 1
 
-@socketio.on('disconnect' , namespace='/')
+@socketio.on('disconnect' )
 def handle_disconnect():
-    client_id = request.sid  # クライアントIDを取得
-    leave_room(client_id)  # クライアントをルームから出す
+    client_id = request.sid
+    if client_id in client_rooms:
+        room = client_rooms[client_id]
+        leave_room(room)
+        del client_rooms[client_id]
+        print(f"Client {client_id} disconnected and removed from room {room}")
+
     global connected_clients
     connected_clients -= 1
     # キャンセル処理：接続が切断された場合、そのクライアントに関連するタスクをキャンセル。ただし、1番目で処理中のタスクはキャンセルしない
     client_ip = get_remote_address()
-    for task_id, task in list(active_tasks.items()):
+    for task_id, task in active_tasks.items():
         if task.client_ip == client_ip and not task.is_processing:
             task.cancel_flag = True
             if task_id in task_futures:
@@ -154,8 +172,6 @@ def submit_task():
 
     # クライアントIPアドレスを取得
     client_ip = get_remote_address()
-    # client_id = session.get("sid")
-    
     # 同一IPからの同時タスク数を制限
     if tasks_per_client.get(client_ip, 0) >= 2:
         return jsonify({'error': 'Maximum number of concurrent tasks reached'}), 429
@@ -174,30 +190,32 @@ def submit_task():
     # ファイルデータをバイト列として保存
     file_data = file.read()
     
-    task = Task(task_id, mode, weight1, weight2, file_data, client_ip)
+    client_id = request.headers.get('X-Client-ID')
+    task = Task(task_id, mode, weight1, weight2, file_data, client_ip, client_id)
     task_queue.put(task)
     active_tasks[task_id] = task
     
     # 同一IPからのタスク数をインクリメント
     tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) + 1
 
-    update_queue_status('Task submitted')
+    update_queue_status('Task submitted') # すべてに通知
     
     queue_size = task_queue.qsize()
-    return jsonify({'task_id': task_id, 'queue_size': queue_size})
+    task_order = get_active_task_order(task_id)
+    return jsonify({'task_id': task_id, 'task_order': task_order, 'queue_size': queue_size}) 
 
 @app.route('/cancel_task/<task_id>', methods=['POST'])
 def cancel_task(task_id):
     # クライアントIPアドレスを取得
     client_ip = get_remote_address()
 
-    if task_id in active_tasks:
+    if task_id in active_tasks.keys():
         task = active_tasks[task_id]
         # タスクの所有者を確認（IPアドレスで簡易的に判断）
         if task.client_ip != client_ip:
             return jsonify({'error': 'Unauthorized to cancel this task'}), 403
         task.cancel_flag = True
-        if task_id in task_futures:
+        if task_id in task_futures.keys():
             task_futures[task_id].cancel()
             del task_futures[task_id]
         del active_tasks[task_id]
@@ -216,21 +234,27 @@ def cancel_task(task_id):
 
 @app.route('/task_status/<task_id>', methods=['GET'])
 def task_status(task_id):
-    task = active_tasks.get(task_id, None)
-    if task:
-        return jsonify({'task_id': task_id, 'is_processing': task.is_processing})
-    return jsonify({'error': 'Task not found'}), 404
+    try:
+        if task_id in active_tasks.keys():
+            task = active_tasks[task_id]
+            return jsonify({'task_id': task_id, 'is_processing': task.is_processing})
+        else:
+            return jsonify({'task_id': task_id, 'is_processing': False})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
 
 def get_active_task_order(task_id):
-    non_processing_tasks = [tid for tid, task in active_tasks.items() if not task.is_processing]
-    return non_processing_tasks.index(task_id) if task_id in non_processing_tasks else 0
+    non_processing_task_ids = [tid for tid, task in active_tasks.items() if not task.is_processing]
+    return non_processing_task_ids.index(task_id) if task_id in non_processing_task_ids else 0
 
 # get_task_orderイベントハンドラー
 @app.route('/get_task_order/<task_id>', methods=['GET'])
 def handle_get_task_order(task_id):
-    if task_id not in active_tasks:
-        return jsonify({'error': 'Task not found'}), 404
-    return jsonify({'task_order': get_active_task_order(task_id)})
+    print(f'Active tasks order: {task_id}, Active tasks: {active_tasks.keys()}')
+    if task_id  in active_tasks.keys():
+        return jsonify({'task_order': get_active_task_order(task_id)})
+    else:
+        return jsonify({'task_order': 0})
 
 # Flaskルート
 # ルートパスのGETリクエストに対するハンドラ