def setup_server_listener(conn_filename, parent_pid, lower_port, upper_port,
response_addr, kernel_id, public_key):
ip = "0.0.0.0"
key = str(uuid.uuid4()).encode() # convert to bytes
ports = _select_ports(5, lower_port, upper_port)
write_connection_file(fname=conn_filename,
ip=ip,
key=key,
shell_port=ports[0],
iopub_port=ports[1],
stdin_port=ports[2],
hb_port=ports[3],
control_port=ports[4])
if response_addr:
comm_socket = return_connection_info(conn_filename, response_addr,
int(lower_port), int(upper_port),
kernel_id, public_key,
int(parent_pid))
if comm_socket: # socket in use, start server listener thread
server_listener_thread = Thread(target=server_listener,
args=(
comm_socket,
int(parent_pid),
))
server_listener_thread.start()
return
def setup_gateway_listener(conn_filename, parent_pid, lower_port, upper_port,
response_addr, kernel_id, public_key):
ip = "0.0.0.0"
key = str_to_bytes(str(uuid.uuid4()))
ports = _select_ports(5, lower_port, upper_port)
write_connection_file(fname=conn_filename,
ip=ip,
key=key,
shell_port=ports[0],
iopub_port=ports[1],
stdin_port=ports[2],
hb_port=ports[3],
control_port=ports[4])
if response_addr:
gateway_socket = return_connection_info(conn_filename, response_addr,
int(lower_port),
int(upper_port), kernel_id,
public_key, int(parent_pid))
if gateway_socket: # socket in use, start gateway listener thread
gateway_listener_thread = Thread(target=gateway_listener,
args=(
gateway_socket,
int(parent_pid),
))
gateway_listener_thread.start()
return
def test_write_connection_file():
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
assert os.path.exists(cf)
with open(cf, 'r') as f:
info = json.load(f)
info['key'] = str_to_bytes(info['key'])
assert info == sample_info
def test_write_connection_file():
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
nt.assert_true(os.path.exists(cf))
with open(cf, 'r') as f:
info = json.load(f)
info['key'] = str_to_bytes(info['key'])
nt.assert_equal(info, sample_info)
def test_write_connection_file():
with TemporaryDirectory() as d:
cf = os.path.join(d, "kernel.json")
connect.write_connection_file(cf, **sample_info)
assert os.path.exists(cf)
with open(cf, "r") as f:
info = json.load(f)
info["key"] = info["key"].encode()
assert info == sample_info
def test_app_load_connection_file():
"""test `ipython console --existing` loads a connection file"""
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
app = DummyConsoleApp(connection_file=cf)
app.initialize(argv=[])
for attr, expected in sample_info.items():
if attr in ('key', 'signature_scheme'):
continue
value = getattr(app, attr)
assert value == expected, "app.%s = %s != %s" % (attr, value, expected)
def test_app_load_connection_file():
"""test `ipython console --existing` loads a connection file"""
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
app = DummyConsoleApp(connection_file=cf)
app.initialize(argv=[])
for attr, expected in sample_info.items():
if attr in ('key', 'signature_scheme'):
continue
value = getattr(app, attr)
nt.assert_equal(value, expected, "app.%s = %s != %s" % (attr, value, expected))
def test_get_connection_info():
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
json_info = connect.get_connection_info(cf)
info = connect.get_connection_info(cf, unpack=True)
nt.assert_equal(type(json_info), type(""))
nt.assert_equal(info, sample_info)
info2 = json.loads(json_info)
info2['key'] = str_to_bytes(info2['key'])
nt.assert_equal(info2, sample_info)
def test_get_connection_info():
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
json_info = connect.get_connection_info(cf)
info = connect.get_connection_info(cf, unpack=True)
nt.assert_equal(type(json_info), type(""))
nt.assert_equal(info, sample_info)
info2 = json.loads(json_info)
info2['key'] = str_to_bytes(info2['key'])
nt.assert_equal(info2, sample_info)
def test_load_connection_file_session_with_kn():
"""test load_connection_file() after """
session = Session()
app = DummyConsoleApp(session=Session())
app.initialize(argv=[])
session = app.session
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info_kn)
app.connection_file = cf
app.load_connection_file()
assert session.key == sample_info_kn['key']
assert session.signature_scheme == sample_info_kn['signature_scheme']
def test_load_connection_file_session():
"""test load_connection_file() after """
session = Session()
app = DummyConsoleApp(session=Session())
app.initialize(argv=[])
session = app.session
with TemporaryDirectory() as d:
cf = os.path.join(d, 'kernel.json')
connect.write_connection_file(cf, **sample_info)
app.connection_file = cf
app.load_connection_file()
nt.assert_equal(session.key, sample_info['key'])
nt.assert_equal(session.signature_scheme, sample_info['signature_scheme'])
def setup_gateway_listener(fname, parent_pid, lower_port, upper_port):
ip = "0.0.0.0"
key = str_to_bytes(str(uuid.uuid4()))
gateway_socket = prepare_gateway_socket(lower_port, upper_port)
gateway_listener_thread = Thread(target=gateway_listener,
args=(
gateway_socket,
parent_pid,
))
gateway_listener_thread.start()
basename = os.path.splitext(os.path.basename(fname))[0]
fd, conn_file = tempfile.mkstemp(suffix=".json", prefix=basename + "_")
os.close(fd)
ports = _select_ports(5, lower_port, upper_port)
conn_file, config = write_connection_file(fname=conn_file,
ip=ip,
key=key,
shell_port=ports[0],
iopub_port=ports[1],
stdin_port=ports[2],
hb_port=ports[3],
control_port=ports[4])
try:
os.remove(conn_file)
except:
pass
# Add in the gateway_socket and parent_pid fields...
config['comm_port'] = gateway_socket.getsockname()[1]
config['pid'] = parent_pid
with open(fname, 'w') as f:
f.write(json.dumps(config, indent=2))
return
if create_spark_context:
spark = WaitingForSparkSessionToBeInitialized(global_variable_name='spark')
sc = WaitingForSparkSessionToBeInitialized(global_variable_name='sc')
sqlContext = WaitingForSparkSessionToBeInitialized(global_variable_name='sqlContext')
sqlCtx = WaitingForSparkSessionToBeInitialized(global_variable_name='sqlCtx')
thread_to_initialize_spark_session = Thread(target=initialize_spark_session)
# If the connection file doesn't exist, then create it.
if not os.path.isfile(connection_file):
key = str_to_bytes(str(uuid.uuid4()))
connection_file = determine_connection_file(connection_file)
ports = _select_ports(5, lower_port, upper_port)
write_connection_file(fname=connection_file, ip=ip, key=key, shell_port=ports[0], iopub_port=ports[1],
stdin_port=ports[2], hb_port=ports[3], control_port=ports[4])
if response_addr:
gateway_socket = return_connection_info(connection_file, response_addr, disable_gateway_socket,
lower_port, upper_port)
if gateway_socket: # socket in use, start gateway listener thread
gateway_listener_thread = Thread(target=gateway_listener, args=(gateway_socket,))
gateway_listener_thread.start()
# start to initialize the Spark session in the background
if create_spark_context:
thread_to_initialize_spark_session.start()
# launch the IPython kernel instance
embed_kernel(connection_file=connection_file, ip=ip)
try:
dest='disable_gateway_socket',
action='store_true',
help='Disable use of gateway socket for extended communications',
default=False)
arguments = vars(parser.parse_args())
connection_file = arguments['connection_file']
response_addr = arguments['response_address']
disable_gateway_socket = arguments['disable_gateway_socket']
ip = "0.0.0.0"
# If the connection file doesn't exist, then create it.
if not os.path.isfile(connection_file):
key = str_to_bytes(str(uuid.uuid4()))
connection_file = determine_connection_file(connection_file)
write_connection_file(fname=connection_file, ip=ip, key=key)
if response_addr:
gateway_socket = return_connection_info(connection_file, ip,
response_addr,
disable_gateway_socket)
if gateway_socket: # socket in use, start gateway listener thread
gateway_listener_thread = Thread(target=gateway_listener,
args=(gateway_socket, ))
gateway_listener_thread.start()
# start to initialize the Spark session in the background
thread_to_initialize_spark_session.start()
# launch the IPython kernel instance
embed_kernel(connection_file=connection_file, ip=ip)
raise RuntimeError(
"At least one of the parameters: 'connection_file' or "
"'--RemoteProcessProxy.kernel-id' must be provided!")
# If the connection file doesn't exist, then create it.
if (connection_file
and not os.path.isfile(connection_file)) or kernel_id is not None:
key = str_to_bytes(str(uuid.uuid4()))
connection_file = determine_connection_file(connection_file, kernel_id)
ports = _select_ports(5, lower_port, upper_port)
write_connection_file(fname=connection_file,
ip=ip,
key=key,
shell_port=ports[0],
iopub_port=ports[1],
stdin_port=ports[2],
hb_port=ports[3],
control_port=ports[4])
if response_addr:
gateway_socket = return_connection_info(connection_file,
response_addr, lower_port,
upper_port)
if gateway_socket: # socket in use, start gateway listener thread
gateway_listener_process = Process(target=gateway_listener,
args=(
gateway_socket,
os.getpid(),
))
gateway_listener_process.start()
def initialize(self):
relay_host = self.remote_connection_file['relayHost']
relay_port = self.remote_connection_file['relayPort']
signaling_port = self.remote_connection_file['signalPort']
# generate the local file
from jupyter_client.connect import write_connection_file
local_connection_file_name, local_connection_file = write_connection_file()
local_connection_file['remoteHost'] = socket.gethostname()
local_connection_file['ip'] = '0.0.0.0'
local_connection_file['key'] = self.remote_connection_file['key'] # otherwise no auth
logging.info('Use connection file %s' % json.dumps(local_connection_file))
with open(local_connection_file_name, 'w') as f:
f.write(json.dumps(local_connection_file))
# notify the launcher
context = zmq.Context()
self.signaling_socket = context.socket(zmq.SUB)
self.signaling_socket.setsockopt(zmq.SUBSCRIBE, b'')
self.signaling_socket.connect('tcp://%s:%s' % (relay_host, signaling_port))
self.callback_socket = context.socket(zmq.REQ)
self.callback_socket.connect('tcp://%s:%s' % (relay_host, relay_port))
self.callback_socket.send(json.dumps(local_connection_file).encode('utf8'))
message = self.callback_socket.recv()
logging.info("Got %s" % message)
# proceed with heartbeating
heart_failure = False
hb_thread = threading.Thread(name="forwarder-watcher", target=self.hb_forwarder)
hb_thread.daemon = True
hb_thread.start()
signal_thread = threading.Thread(name="signaling-watcher", target=self.signaling_handler)
signal_thread.daemon = True
signal_thread.start()
# start relaying the sockets in the connection file
port_pairs = []
for port_type in ['shell_port', 'iopub_port', 'stdin_port', 'control_port', 'hb_port']:
local_port = local_connection_file.get(port_type, None)
remote_port = self.remote_connection_file.get(port_type, None)
if local_port is None or local_port == 0:
continue
if remote_port is None or remote_port == 0:
continue
port_pairs.append([local_port, remote_port, port_type[:-5]])
def printout(x):
logging.info(x)
# bind on 127.0.0.1 for the jupyter-server-facing side and on all interfaces for the kernel-facing side
def forward_ROUTER_DEALER(local_port, remote_port, port_type):
return ROUTER_DEALER_Forwarder(relay_host, remote_port, '127.0.0.1', local_port, port_type, printout, False, False)
def forward_PUB_SUB(local_port, remote_port, port_type):
return PUB_SUB_Forwarder(relay_host, remote_port, '127.0.0.1', local_port, port_type, printout, False, False)
def forward_REP_REQ(local_port, remote_port, port_type):
return REQ_REP_Forwarder(relay_host, remote_port, '127.0.0.1', local_port, port_type, printout, False, False)
socket_forwarders = {'hb' : forward_REP_REQ, 'shell' : forward_ROUTER_DEALER, 'iopub' : forward_PUB_SUB, 'stdin' : forward_ROUTER_DEALER, 'control': forward_ROUTER_DEALER}
for port_pair in port_pairs:
local_port = port_pair[0]
remote_port = port_pair[1]
port_type = port_pair[2]
logging.info("Relay port %s to %s on type %s" % (local_port, remote_port, port_type))
socket_forwarders[port_type](local_port, remote_port, port_type)
return local_connection_file_name
