From e63186e209bf27b16551dcc43cb6661cf8fec806 Mon Sep 17 00:00:00 2001 From: Dun Liang Date: Tue, 4 May 2021 22:30:37 +0800 Subject: [PATCH] add student queue --- python/jittor/__init__.py | 4 +- python/jittor_utils/student_queue.py | 67 +++++++++++++++++++++++++--- script/converter_server.sh | 2 +- 3 files changed, 65 insertions(+), 8 deletions(-) diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py index 4c9f678d..7e2854e8 100644 --- a/python/jittor/__init__.py +++ b/python/jittor/__init__.py @@ -8,7 +8,7 @@ # This file is subject to the terms and conditions defined in # file 'LICENSE.txt', which is part of this source code package. # *************************************************************** -__version__ = '1.2.2.64' +__version__ = '1.2.2.65' from . import lock with lock.lock_scope(): ori_int = int @@ -39,6 +39,8 @@ import hashlib import sys, os import traceback +if "SKEY" in os.environ: + import jittor_utils.student_queue def safepickle(obj, path): # Protocol version 4 was added in Python 3.4. It adds support for very large objects, pickling more kinds of objects, and some data format optimizations. diff --git a/python/jittor_utils/student_queue.py b/python/jittor_utils/student_queue.py index f41b8cfd..b6f5c7df 100644 --- a/python/jittor_utils/student_queue.py +++ b/python/jittor_utils/student_queue.py @@ -1,12 +1,67 @@ +# *************************************************************** +# Copyright (c) 2021 Jittor. All Rights Reserved. +# Maintainers: +# Dun Liang . +# +# This file is subject to the terms and conditions defined in +# file 'LICENSE.txt', which is part of this source code package. +# *************************************************************** from socketserver import ThreadingTCPServer import socket +import os +import sys +import threading + +key_queue = {} def handle_connect(req:socket.socket, c_addr, server): print("get connect", c_addr, req) - while True: - buf = req.recv(2048) - - print(buf) + skey = req.recv(1024).decode() + with lock: + if skey not in key_queue: + key_queue[skey] = [] + queue = key_queue[skey] + queue.append(req) -server = ThreadingTCPServer(("127.0.0.1", 8900), handle_connect) -server.serve_forever() \ No newline at end of file + req.send(str(len(queue)-1).encode()) + while True: + buf = req.recv(1024).decode() + print(buf) + with lock: + if len(buf) == 0: + for i,r in enumerate(queue): + if r is req: + for j in range(i+1, len(queue)): + queue[j].send(str(j-1).encode()) + del queue[i] + print("queue size", len(queue)) + break + break + + +def wait_queue(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.connect(("127.0.0.1", 8900)) + s.send(skey.encode()) + while True: + buf = s.recv(1024).decode() + if len(buf) == 0: + print("Cannot connect to queue server, please report this issue to admin.") + sys.exit(1) + if buf == '0': + print("Begin") + os.system(f"sleep {os.environ.get('SWAIT', '60')} && bash -c ' if kill -9 {os.getpid()} 2>/dev/null; then echo Timeout; fi; ' &") + os.system("sleep 100") + break + else: + print("Pending", buf) + + + +if "SKEY" in os.environ: + skey = os.environ["SKEY"] + wait_queue() +else: + lock = threading.Lock() + server = ThreadingTCPServer(("127.0.0.1", 8900), handle_connect) + server.serve_forever() diff --git a/script/converter_server.sh b/script/converter_server.sh index 444ed8b7..7768d32e 100644 --- a/script/converter_server.sh +++ b/script/converter_server.sh @@ -9,6 +9,6 @@ docker build --tag jittor/converter_server -f /tmp/converter_server.dockerfile . # docker run --rm -it -m 16g --cpus=8 -p 0.0.0.0:5000:5000 jittor/converter_server bash -c "python3.7 -m pip install -U git+https://github.com/Jittor/jittor.git && python3.7 -m jittor.utils.converter_server" while true; do - timeout --foreground 24h docker run --rm -it -m 16g --cpus=8 -p 0.0.0.0:5000:5000 -v /etc/letsencrypt/:/https jittor/converter_server bash -c "python3.7 -m pip install -U jittor && python3.7 -m jittor.test.test_core && FLASK_APP=/usr/local/lib/python3.7/dist-packages/jittor/utils/converter_server python3.7 -m flask run --cert=/https/live/randonl.me/fullchain.pem --key=/https/live/randonl.me/privkey.pem --host=0.0.0.0" + timeout --foreground 24h docker run --rm -it -m 16g --cpus=8 -p 0.0.0.0:58187:5000 -v /etc/letsencrypt/:/https jittor/converter_server bash -c "python3.7 -m pip install -U jittor && python3.7 -m jittor.test.test_core && FLASK_APP=/usr/local/lib/python3.7/dist-packages/jittor/utils/converter_server python3.7 -m flask run --cert=/https/live/randonl.me/fullchain.pem --key=/https/live/randonl.me/privkey.pem --host=0.0.0.0" sleep 10 done \ No newline at end of file