0.13

Author: divideconcept

torchstudio/modeltrain.py

@@ -184,7 +184,7 @@ def deepcopy_cpu(value):
 
             sshclient = paramiko.SSHClient()
             sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-            sshclient.connect(hostname=sshaddress, port=int(sshport), username=username, password=password, pkey=pkey, timeout=5)
+            sshclient.connect(hostname=sshaddress, port=int(sshport), username=username, password=password, pkey=pkey, timeout=10)
 
             reverse_tunnel = sshtunnel.Tunnel(sshclient, sshtunnel.ReverseTunnel, 'localhost', 0, 'localhost', int(address[1]))
             address[1]=str(reverse_tunnel.lport)
@@ -319,6 +319,7 @@ def deepcopy_cpu(value):
 
         buffer=io.BytesIO()
         torch.save(deepcopy_cpu(model.state_dict()), buffer)
+        print("Training... epoch "+str(scheduler.last_epoch-1)+' | (save)\n\n', file=sys.stderr)
         tc.send_msg(app_socket, 'ModelState', buffer.getvalue())
 
         buffer=io.BytesIO()

torchstudio/pythoninstall.cmd

@@ -227,7 +227,7 @@ if EXIST "%pythonpath%" (
     )
 )
 
-set PATH=%PATH%;%pythonpath%;%pythonpath%\Library\mingw-w64\bin;%pythonpath%\Library\bin
+set PATH=%PATH%;%pythonpath%;%pythonpath%\Library\mingw-w64\bin;%pythonpath%\Library\bin;%pythonpath%\bin
 "%pythonpath%\python" -u -B -X utf8 -m torchstudio.pythoninstall --channel %channel% %cuda% %packages%
 if ERRORLEVEL 1 (
     echo. 1>&2

torchstudio/pythoninstall.py

@@ -38,12 +38,12 @@
             highest_cuda_string='.'.join([str(value) for value in highest_cuda_version])
             print("Using CUDA "+highest_cuda_string)
             print("")
-            conda_install+=" fastchan::cudatoolkit="+highest_cuda_string
+            conda_install+=f" {args.channel}::pytorch-cuda="+highest_cuda_string+" -c nvidia"
         else:
             conda_install+=f" {args.channel}::cpuonly"
     print(f"Downloading and installing {args.channel} packages...")
     print("")
-    conda_install+=" -k" #adding fastchan for cudatoolkit, faster than the conda-forge channel
+    conda_install+=" -k" #allow insecure ssl connections
     # https://stackoverflow.com/questions/41767340/using-conda-install-within-a-python-script
     (stdout_str, stderr_str, return_code_int) = Conda.run_command(Conda.Commands.INSTALL,conda_install.split(),use_exception_handler=True,stdout=sys.stdout,stderr=sys.stderr)
     if return_code_int!=0:

torchstudio/sshtunnel.py

@@ -258,7 +258,7 @@ def finish(self):
 
     print("Connecting to remote server...", file=sys.stderr)
     try:
-        sshclient.connect(hostname=args.sshaddress, port=args.sshport, username=args.username, password=args.password, pkey=paramiko.RSAKey.from_private_key_file(args.keyfile) if args.keyfile else None, timeout=5)
+        sshclient.connect(hostname=args.sshaddress, port=args.sshport, username=args.username, password=args.password, pkey=paramiko.RSAKey.from_private_key_file(args.keyfile) if args.keyfile else None, timeout=10)
     except:
         print("Error: could not connect to remote server", file=sys.stderr)
         exit(1)

torchstudio/tcpcodec.py

@@ -42,15 +42,25 @@ def connect(server_address=None, timeout=0):
         sock.settimeout(timeout)
     return sock
 
-def send_msg(sock, type, data = bytearray()):
-    type_bytes=bytes(type, 'utf-8')
+def send_msg(sock, data_type, data = bytearray()):
+    def sendall(sock, data):
+        while len(data) >0:
+            try:
+                ret = sock.send(data[:1048576]) #1MB chunks
+            except:
+                print("Lost connection (send timeout)", file=sys.stderr)
+                exit()
+            if ret == 0:
+                print("Lost connection (send null)", file=sys.stderr)
+                exit()
+            else:
+                data=data[ret:]
+
+    type_bytes=bytes(data_type, 'utf-8')
     type_size=len(type_bytes)
     msg = struct.pack(f'<B{type_size}sI', type_size, type_bytes, len(data)) + data
-    try:
-        sock.sendall(msg)
-    except:
-        print("Lost connection", file=sys.stderr)
-        exit()
+
+    sendall(sock, msg)
 
 def recv_msg(sock):
     def recvall(sock, n):
@@ -59,17 +69,17 @@ def recvall(sock, n):
             try:
                 packet = sock.recv(n - len(data))
             except:
-                print("Lost connection", file=sys.stderr)
+                print("Lost connection (receive timeout)", file=sys.stderr)
                 exit()
             if len(packet)==0:
-                print("Lost connection", file=sys.stderr)
+                print("Lost connection (receive null)", file=sys.stderr)
                 exit()
             data.extend(packet)
         return data
     type_size = struct.unpack('<B', recvall(sock, 1))[0]
-    type = struct.unpack(f'<{type_size}s', recvall(sock, type_size))[0]
+    data_type = struct.unpack(f'<{type_size}s', recvall(sock, type_size))[0]
     datalen = struct.unpack('<I', recvall(sock, 4))[0]
-    return str(type, 'utf-8'), recvall(sock, datalen)
+    return str(data_type, 'utf-8'), recvall(sock, datalen)