Update Jupyter notebook server execution example

This commit is contained in:
allegroai 2020-03-09 12:20:34 +02:00
parent 2bffc24347
commit 1043c22d0a

View File

@ -3,12 +3,12 @@ import sys
import subprocess import subprocess
from copy import deepcopy from copy import deepcopy
import socket import socket
import psutil
from tempfile import mkstemp from tempfile import mkstemp
# make sure we have jupyter in the auto requirements # make sure we have jupyter in the auto requirements
import jupyter import jupyter
from trains import Task from trains import Task
# set default docker image, with network configuration # set default docker image, with network configuration
os.environ['TRAINS_DOCKER_IMAGE'] = 'nvidia/cuda --network host' os.environ['TRAINS_DOCKER_IMAGE'] = 'nvidia/cuda --network host'
@ -29,15 +29,27 @@ for key in os.environ:
param = { param = {
'jupyter_server_base_directory': '~/', 'jupyter_server_base_directory': '~/',
'ssh_server': True, 'ssh_server': True,
'ssh_password': 'training'
} }
task.connect(param) task.connect(param)
try:
hostname = socket.gethostname()
hostnames = socket.gethostbyname(socket.gethostname())
except Exception:
def get_ip_addresses(family):
for interface, snics in psutil.net_if_addrs().items():
for snic in snics:
if snic.family == family:
yield snic.address
hostnames = list(get_ip_addresses(socket.AF_INET))
hostname = hostnames[0]
if param.get('ssh_server'): if param.get('ssh_server'):
print('Installing SSH Server on {} [{}]'.format(socket.gethostname(), print('Installing SSH Server on {} [{}]'.format(hostname, hostnames))
socket.gethostbyname(socket.gethostname()))) ssh_password = param.get('ssh_password', 'training')
ssh_password = 'training'
try: try:
import psutil
used_ports = [i.laddr.port for i in psutil.net_connections()] used_ports = [i.laddr.port for i in psutil.net_connections()]
port = [i for i in range(10022, 15000) if i not in used_ports][0] port = [i for i in range(10022, 15000) if i not in used_ports][0]
@ -55,7 +67,7 @@ if param.get('ssh_server'):
if result == 0: if result == 0:
print('\n#\n# SSH Server running on {} [{}] port {}\n# LOGIN u:root p:{}\n#\n'.format( print('\n#\n# SSH Server running on {} [{}] port {}\n# LOGIN u:root p:{}\n#\n'.format(
socket.gethostname(), socket.gethostbyname(socket.gethostname()), port, ssh_password)) hostname, hostnames, port, ssh_password))
else: else:
raise ValueError() raise ValueError()
except: except:
@ -65,9 +77,8 @@ if param.get('ssh_server'):
fd, local_filename = mkstemp() fd, local_filename = mkstemp()
cwd = os.path.expandvars(os.path.expanduser(param['jupyter_server_base_directory'])) \ cwd = os.path.expandvars(os.path.expanduser(param['jupyter_server_base_directory'])) \
if param['jupyter_server_base_directory'] else os.getcwd() if param['jupyter_server_base_directory'] else os.getcwd()
print('Running Jupyter Notebook Server on {} [{}] at {}'.format(socket.gethostname(), print('Running Jupyter Notebook Server on {} [{}] at {}'.format(hostname, hostnames, cwd))
socket.gethostbyname(socket.gethostname()), cwd)) process = subprocess.Popen([sys.executable, '-m', 'jupyter', 'notebook', '--no-browser', '--allow-root', '--ip', '0.0.0.0'],
process = subprocess.Popen([sys.executable, '-m', 'jupyter', 'notebook', '--no-browser', '--allow-root'],
env=env, stdout=fd, stderr=fd, cwd=cwd) env=env, stdout=fd, stderr=fd, cwd=cwd)
# print stdout/stderr # print stdout/stderr
@ -100,4 +111,4 @@ os.close(fd)
try: try:
os.unlink(local_filename) os.unlink(local_filename)
except: except:
pass pass