add student queue

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.

python/jittor_utils/student_queue.py

@@ -1,12 +1,67 @@
+# ***************************************************************
+# Copyright (c) 2021 Jittor. All Rights Reserved. 
+# Maintainers:
+#   Dun Liang <randonlang@gmail.com>.
+# 
+# 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)
+    skey = req.recv(1024).decode()
+    with lock:
+        if skey not in key_queue:
+            key_queue[skey] = []
+        queue = key_queue[skey]
+        queue.append(req)
+
+        req.send(str(len(queue)-1).encode())
     while True:
-        buf = req.recv(2048)
-        
+        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
         
-server = ThreadingTCPServer(("127.0.0.1", 8900), handle_connect)
-server.serve_forever()
+
+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()

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