yeq6x commited on
Commit
9757452
·
1 Parent(s): b4880bf
Files changed (3) hide show
  1. Dockerfile +0 -7
  2. app.py +22 -22
  3. requirements.txt +0 -1
Dockerfile CHANGED
@@ -46,12 +46,5 @@ RUN echo '#!/bin/bash\nredis-server --daemonize yes\npython app.py' > /app/start
46
  RUN useradd -m appuser && chown -R appuser:appuser /app
47
  USER appuser
48
 
49
- # 環境変数を設定
50
- ENV REDIS_HOST=redis
51
- ENV REDIS_PORT=6379
52
-
53
- # ポートを指定
54
- EXPOSE 7860
55
-
56
  # 起動コマンドを変更
57
  CMD ["/app/start.sh"]
 
46
  RUN useradd -m appuser && chown -R appuser:appuser /app
47
  USER appuser
48
 
 
 
 
 
 
 
 
49
  # 起動コマンドを変更
50
  CMD ["/app/start.sh"]
app.py CHANGED
@@ -4,7 +4,6 @@ from flask_cors import CORS
4
  from flask_limiter import Limiter
5
  from flask_limiter.util import get_remote_address
6
  import concurrent.futures
7
- import redis
8
 
9
  import io
10
  import os
@@ -27,9 +26,6 @@ app = Flask(__name__)
27
  CORS(app)
28
  socketio = SocketIO(app, cors_allowed_origins="*")
29
 
30
- # Redisクライアントの初期化(レート制限とキャッシュのため)
31
- redis_client = redis.Redis(host=os.environ.get('REDIS_HOST', 'localhost'), port=int(os.environ.get('REDIS_PORT', 6379)), db=0)
32
-
33
  # レート制限の設定
34
  limiter = Limiter(
35
  get_remote_address,
@@ -113,6 +109,23 @@ def worker():
113
  # ワーカースレッドの開始
114
  threading.Thread(target=worker, daemon=True).start()
115
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
116
  @app.route('/submit_task', methods=['POST'])
117
  @limiter.limit("10 per minute") # 1分間に10回までのリクエストに制限
118
  def submit_task():
@@ -123,7 +136,7 @@ def submit_task():
123
  client_ip = get_remote_address()
124
 
125
  # 同一IPからの同時タスク数を制限
126
- if redis_client.get(f'tasks:{client_ip}') and int(redis_client.get(f'tasks:{client_ip}')) >= 2:
127
  return jsonify({'error': 'Maximum number of concurrent tasks reached'}), 429
128
 
129
  task_id = str(uuid.uuid4())
@@ -145,8 +158,7 @@ def submit_task():
145
  active_tasks[task_id] = task
146
 
147
  # 同一IPからのタスク数をインクリメント
148
- redis_client.incr(f'tasks:{client_ip}')
149
- redis_client.expire(f'tasks:{client_ip}', 3600) # 1時間後に期限切れ
150
 
151
  update_queue_status(f'Task submitted: {task_id}')
152
 
@@ -169,7 +181,7 @@ def cancel_task(task_id):
169
  del task_futures[task_id]
170
  del active_tasks[task_id]
171
  # タスク数をデクリメント
172
- redis_client.decr(f'tasks:{client_ip}')
173
  update_queue_status('Task cancelled')
174
  return jsonify({'message': 'Task cancellation requested'})
175
  else:
@@ -177,7 +189,7 @@ def cancel_task(task_id):
177
  if task.task_id == task_id and task.client_ip == client_ip:
178
  task.cancel_flag = True
179
  # タスク数をデクリメント
180
- redis_client.decr(f'tasks:{client_ip}')
181
  return jsonify({'message': 'Task cancellation requested for queued task'})
182
  return jsonify({'error': 'Task not found'}), 404
183
 
@@ -190,18 +202,6 @@ def handle_get_task_order(task_id):
190
  task_order = get_active_task_order(task_id)
191
  return jsonify({'task_order': task_order})
192
 
193
- @socketio.on('connect')
194
- def handle_connect(auth):
195
- # クライアント接続数の制限
196
- if redis_client.get('connected_clients') and int(redis_client.get('connected_clients')) > 100:
197
- return False # 接続を拒否
198
- redis_client.incr('connected_clients')
199
- emit('queue_update', {'active_tasks': len(active_tasks), 'active_task_order': None})
200
-
201
- @socketio.on('disconnect')
202
- def handle_disconnect():
203
- redis_client.decr('connected_clients')
204
-
205
  # Flaskルート
206
  # ルートパスのGETリクエストに対するハンドラ
207
  @app.route('/', methods=['GET'])
@@ -263,7 +263,7 @@ if __name__ == '__main__':
263
  args = parser.parse_args()
264
 
265
  # initialize(args.use_local, args.use_gpu, args.use_dotenv)
 
266
  port = int(os.environ.get('PORT', 7860))
267
- print(f"Starting server on port {port}")
268
  server = pywsgi.WSGIServer(('0.0.0.0', port), app, handler_class=WebSocketHandler)
269
  server.serve_forever()
 
4
  from flask_limiter import Limiter
5
  from flask_limiter.util import get_remote_address
6
  import concurrent.futures
 
7
 
8
  import io
9
  import os
 
26
  CORS(app)
27
  socketio = SocketIO(app, cors_allowed_origins="*")
28
 
 
 
 
29
  # レート制限の設定
30
  limiter = Limiter(
31
  get_remote_address,
 
109
  # ワーカースレッドの開始
110
  threading.Thread(target=worker, daemon=True).start()
111
 
112
+ # グローバル変数を使用して接続数とタスク数を管理
113
+ connected_clients = 0
114
+ tasks_per_client = {}
115
+
116
+ @socketio.on('connect')
117
+ def handle_connect(auth):
118
+ global connected_clients
119
+ if connected_clients > 100:
120
+ return False # 接続を拒否
121
+ connected_clients += 1
122
+ emit('queue_update', {'active_tasks': len(active_tasks), 'active_task_order': None})
123
+
124
+ @socketio.on('disconnect')
125
+ def handle_disconnect():
126
+ global connected_clients
127
+ connected_clients -= 1
128
+
129
  @app.route('/submit_task', methods=['POST'])
130
  @limiter.limit("10 per minute") # 1分間に10回までのリクエストに制限
131
  def submit_task():
 
136
  client_ip = get_remote_address()
137
 
138
  # 同一IPからの同時タスク数を制限
139
+ if tasks_per_client.get(client_ip, 0) >= 2:
140
  return jsonify({'error': 'Maximum number of concurrent tasks reached'}), 429
141
 
142
  task_id = str(uuid.uuid4())
 
158
  active_tasks[task_id] = task
159
 
160
  # 同一IPからのタスク数をインクリメント
161
+ tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) + 1
 
162
 
163
  update_queue_status(f'Task submitted: {task_id}')
164
 
 
181
  del task_futures[task_id]
182
  del active_tasks[task_id]
183
  # タスク数をデクリメント
184
+ tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
185
  update_queue_status('Task cancelled')
186
  return jsonify({'message': 'Task cancellation requested'})
187
  else:
 
189
  if task.task_id == task_id and task.client_ip == client_ip:
190
  task.cancel_flag = True
191
  # タスク数をデクリメント
192
+ tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
193
  return jsonify({'message': 'Task cancellation requested for queued task'})
194
  return jsonify({'error': 'Task not found'}), 404
195
 
 
202
  task_order = get_active_task_order(task_id)
203
  return jsonify({'task_order': task_order})
204
 
 
 
 
 
 
 
 
 
 
 
 
 
205
  # Flaskルート
206
  # ルートパスのGETリクエストに対するハンドラ
207
  @app.route('/', methods=['GET'])
 
263
  args = parser.parse_args()
264
 
265
  # initialize(args.use_local, args.use_gpu, args.use_dotenv)
266
+
267
  port = int(os.environ.get('PORT', 7860))
 
268
  server = pywsgi.WSGIServer(('0.0.0.0', port), app, handler_class=WebSocketHandler)
269
  server.serve_forever()
requirements.txt CHANGED
@@ -20,6 +20,5 @@ pytorch_lightning
20
  python-dotenv
21
  peft==0.11.1
22
  flask_limiter==3.7.0
23
- redis==5.0.7
24
  gevent==24.2.1
25
  gevent-websocket==0.10.1
 
20
  python-dotenv
21
  peft==0.11.1
22
  flask_limiter==3.7.0
 
23
  gevent==24.2.1
24
  gevent-websocket==0.10.1