完成训练模块的转移

core/event.py

@@ -6,7 +6,7 @@
 # @IDE            : PyCharm
 # @desc           : 全局事件
 
-
+import torch
 from fastapi import FastAPI
 from motor.motor_asyncio import AsyncIOMotorClient
 from application.settings import REDIS_DB_URL, MONGO_DB_URL, MONGO_DB_NAME, EVENTS
@@ -68,6 +68,9 @@ async def connect_redis(app: FastAPI, status: bool):
             response = await rd.ping()
             if response:
                 print("Redis 连接成功")
+                # 数据初始化
+                is_gpu = torch.cuda.is_available()
+                rd.set('is_gpu', str(is_gpu))
             else:
                 print("Redis 连接失败")
         except AuthenticationError as e:

core/websocket_app.py

@@ -0,0 +1,30 @@
+from fastapi import WebSocket
+from fastapi import FastAPI
+from starlette.websockets import WebSocketState
+
+from utils.websocket_server import room_manager
+
+
+def websocket_config(app: FastAPI):
+
+    @app.websocket("/{room}")
+    async def websocket_room(websocket: WebSocket, room: str):
+        """
+        websocket 房间管理
+        :param websocket:
+        :param room:
+        :return:
+        """
+        await websocket.accept()
+        await room_manager.add_to_room(room, websocket)
+        try:
+            while True:
+                data = await websocket.receive_text()
+                await room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
+        except Exception as e:
+            if websocket.client_state != WebSocketState.DISCONNECTED:
+                await websocket.close(code=1000)
+        finally:
+            await room_manager.remove_from_room(room, websocket)
+            if websocket.client_state != WebSocketState.DISCONNECTED:
+                await websocket.close(code=1001)