def exec_code(kid, code):
"""
Executes arbitrary `code` in the kernel with id `kid`.
Returns:
- tuple: the output of the code and the error, if any.
"""
# Load connection info and init communications.
cf = find_connection_file(kid)
with jupyter_lock:
km = BlockingKernelClient(connection_file=cf)
km.load_connection_file()
km.start_channels()
msg_id = km.execute(code, store_history=False)
reply = km.get_shell_msg(msg_id, timeout=60)
output, error = None, None
while km.is_alive():
msg = km.get_iopub_msg(timeout=10)
if ("content" in msg and "name" in msg["content"]
and msg["content"]["name"] == "stdout"):
output = msg["content"]["text"]
break
km.stop_channels()
if reply["content"]["status"] != "ok":
logging.error(f"Status is {reply['content']['status']}")
logging.error(output)
error = output
output = None
return output, error
def exec_code(kid, var, code):
# load connection info and init communication
cf = find_connection_file(kid) # str(port))
global jupyter_lock
jupyter_lock.acquire()
try:
km = BlockingKernelClient(connection_file=cf)
km.load_connection_file()
km.start_channels()
# logging.debug('Executing:\n' + str(code))
msg_id = km.execute(code, store_history=False)
reply = km.get_shell_msg(msg_id, timeout=10)
# logging.info('Execution reply:\n' + str(reply))
state = 'busy'
output = None
idle_count = 0
try:
while km.is_alive():
try:
msg = km.get_iopub_msg(timeout=10)
# logging.debug('Read ' + str(msg))
if not 'content' in msg:
continue
if 'name' in msg['content'] and msg['content'][
'name'] == 'stdout':
# logging.debug('Got data '+ msg['content']['text'])
output = msg['content']['text']
break
if 'execution_state' in msg['content']:
# logging.debug('Got state')
state = msg['content']['execution_state']
if state == 'idle':
idle_count = idle_count + 1
except Empty:
pass
except KeyboardInterrupt:
logging.error('Keyboard interrupt')
pass
finally:
# logging.info('Kernel IO finished')
km.stop_channels()
# logging.info(str(output))
error = ''
if reply['content']['status'] != 'ok':
logging.error('Status is ' + reply['content']['status'])
logging.error(str(output))
error = output
output = None
finally:
jupyter_lock.release()
return output, error
client.load_connection_file(f)
print("prepared")
# msg_id = client.execute("print('hello')")
msg_id = client.execute("1 + 10")
# client.wait_for_ready()
res = client.get_shell_msg(msg_id, timeout=1)
print(res)
print("----------------------------------------")
msg = res["msg_id"]
for i in range(10):
if not client.is_alive():
print("not alived")
break
try:
res = client.get_iopub_msg(msg_id, timeout=1)
except Empty as e:
print("!", repr(e))
time.sleep(0.1)
continue
if res["msg_type"] != "status":
print("!!", res)
break
print(res)
client = BlockingKernelClient()
client.load_connection_file(f)
print("prepared")
# msg_id = client.execute("print('hello')")
msg_id = client.execute("1 + 10")
# client.wait_for_ready()
res = client.get_shell_msg(msg_id, timeout=1)
print(res)
print("----------------------------------------")
msg = res["msg_id"]
for i in range(10):
if not client.is_alive():
print("not alived")
break
try:
res = client.get_iopub_msg(msg_id, timeout=1)
except Empty as e:
print("!", repr(e))
time.sleep(0.1)
continue
if res["msg_type"] != "status":
print("!!", res)
break
print(res)
class SshKernel:
"""Remote ipykernel via SSH
Raises:
SshKernelException: "Could not execute remote command, connection died"
SshKernelException: "Connection failed"
SshKernelException: "Could not create kernel_info file"
Arguments:
host {str} -- host where the remote ipykernel should be started
connection_info {dict} -- Local ipykernel connection info as provided by Juypter lab
python_path {str} -- Remote python path to be used to start ipykernel
Keyword Arguments:
sudo {bool} -- Start ipykernel as root if necessary (default: {False})
timeout {int} -- SSH connection timeout (default: {5})
env {str} -- Environment variables passd to the ipykernel "VAR1=VAL1 VAR2=VAL2" (default: {""})
ssh_config {str} -- Path to the local SSH config file (default: {Path.home() / ".ssh" / "config"})
"""
def __init__(
self,
host,
connection_info,
python_path,
sudo=False,
timeout=5,
env="",
ssh_config=None,
quiet=True,
verbose=False,
msg_interval=30,
logger=None,
):
self.host = host
self.connection_info = connection_info
self.python_full_path = PurePosixPath(python_path) / "bin/python"
self.sudo = sudo
self.timeout = timeout
self.env = env
self.ssh_config = (Path.home() / ".ssh" /
"config" if ssh_config is None else ssh_config
) # OS specific path
self.quiet = quiet
self.verbose = verbose
self._connection = None
self.remote_ports = {}
self.uuid = str(uuid.uuid4())
self.fname = "/tmp/.ssh_ipykernel_%s.json" % self.uuid # POSIX path
if logger is None:
self._logger = setup_logging("SshKernel")
else:
self._logger = logger
self._logger.debug("Remote kernel info file: {0}".format(self.fname))
self._logger.debug(
"Local connection info: {0}".format(connection_info))
self.kernel_pid = 0
self.status = Status(connection_info, self._logger)
self.msg_interval = int(msg_interval / timeout)
self.msg_counter = 0
def _execute(self, cmd):
try:
result = subprocess.check_output(cmd)
return 0, result
except subprocess.CalledProcessError as e:
return e.returncode, e.args
def _ssh(self, cmd):
return self._execute([SSH, self.host, cmd])
def close(self):
"""Close pcssh connection
"""
if self._connection is not None: # and self._connection.isalive():
if self._connection.isalive():
self._connection.logout()
self._logger.debug("Ssh connection closed")
if self.kc.is_alive():
self.kc.stop_channels()
self._logger.debug("Kernel client channels stopped")
def create_remote_connection_info(self):
"""Create a remote ipykernel connection info file
Uses KERNEL_SCRIPT to execute jupyter_client.write_connection_file remotely to request remote ports.
The remote ports will be returned as json and stored to built the SSH tunnels later.
The pxssh connection will be closed at the end.
Raises:
SshKernelException: "Could not create kernel_info file"
"""
self._logger.info("Creating remote connection info")
script = KERNEL_SCRIPT.format(fname=self.fname, **self.connection_info)
cmd = "{python} -c '{command}'".format(python=self.python_full_path,
command="; ".join(
script.strip().split("\n")))
result = self._ssh(cmd)
self._logger.debug(result)
if result[0] == 0:
self.remote_ports = json.loads(result[1].decode("utf-8"))
self._logger.debug("Local ports = %s" % {
k: v
for k, v in self.connection_info.items() if "_port" in k
})
self._logger.debug("Remote ports = %s" % self.remote_ports)
else:
self.status.set_unreachable(self.kernel_pid, self.sudo)
raise SshKernelException("Could not create kernel_info file")
def kernel_client(self):
self.kc = BlockingKernelClient()
self.kc.load_connection_info(self.connection_info)
self.kc.start_channels()
def kernel_init(self):
done = False
if self.check_alive(show_pid=False):
i = 0
while not done:
try:
i += 1
self._logger.debug("Retrieving kernel pid, attempt %d" % i)
result = self.kc.execute_interactive(
"import os",
user_expressions={"pid": "os.getpid()"},
store_history=False,
silent=True,
timeout=2,
)
self._logger.debug("result = %s" % str(result["content"]))
self.kernel_pid = int(result["content"]["user_expressions"]
["pid"]["data"]["text/plain"])
self._logger.debug("Remote kernel pid %d" %
self.kernel_pid)
done = True
except Exception as ex:
msg = str(ex)
if msg == "Timeout waiting for output":
self._logger.warning("Warning: {}".format(msg))
if i > 5:
self._logger.error(
"Max attempts (5) reached, stopping")
raise SshKernelException(
"Could not initialize kernel")
break
else:
self._logger.error("Warning: {}".format(str(ex)))
return done
def kernel_customize(self):
pass
def check_alive(self, show_pid=True):
alive = self._connection.isalive() and self.kc.is_alive()
if show_pid:
msg = "Remote kernel ({}, pid = {}) is {}alive".format(
self.host, self.kernel_pid, "" if alive else "not ")
else:
msg = "Remote kernel is {}alive".format("" if alive else "not ")
if not alive or self.msg_counter % self.msg_interval == 0:
self.msg_counter = 0
self._logger.info(msg)
self.msg_counter += 1
return alive
def interrupt_kernel(self):
if self._connection.isalive():
if is_windows:
self._logger.warning(
'On Windows use "Interrupt remote kernel" button')
else:
self._logger.warning("Sending interrupt to remote kernel")
self._connection.sendintr() # send SIGINT
def start_kernel_and_tunnels(self):
"""Start Kernels and SSH tunnels
A new pxssh connection will be created that will
- set up the necessary ssh tunnels between remote kernel ports and local kernel ports
- start the ipykernel on the remote host
"""
self._logger.info("Setting up ssh tunnels")
ssh_tunnels = []
for port_name in self.remote_ports.keys():
ssh_tunnels += [
"-L",
"{local_port}:127.0.0.1:{remote_port}".format(
local_port=self.connection_info[port_name],
remote_port=self.remote_ports[port_name],
),
]
self._logger.info("Starting remote kernel")
# Build remote command
sudo = "sudo " if self.sudo else ""
if self.env is not None:
env = " ".join(self.env)
cmd = "{sudo} {env} {python} -m ipykernel_launcher -f {fname}".format(
sudo=sudo, env=env, python=self.python_full_path, fname=self.fname)
# Build ssh command with all flags and tunnels
if self.quiet:
args = ["-q"]
elif self.verbose:
args = ["-v"]
else:
args = []
args += ["-t", "-F", str(self.ssh_config)
] + ssh_tunnels + [self.host, cmd]
self._logger.debug("%s %s" % (SSH, " ".join(args)))
try:
# Start the child process
self._connection = expect.spawn(SSH,
args=args,
timeout=self.timeout,
**ENCODING)
# subprocess.check_output([SSH] + args)
#
# get blocking kernel client
self.kernel_client()
# initialize it
if self.kernel_init():
self.status.set_running(self.kernel_pid, self.sudo)
# run custom code if part of sub class
self.kernel_customize()
else:
self.status.set_connect_failed(sudo=self.sudo)
except Exception as e:
tb = sys.exc_info()[2]
self._logger.error(str(e.with_traceback(tb)))
self._logger.error("Cannot contiune, exiting")
sys.exit(1)
prompt = re.compile(r"\n")
while True:
try:
# Wait for prompt
self._connection.expect(prompt)
# print the outputs
self._logger.info(self._connection.before.strip("\r\n"))
except KeyboardInterrupt:
self.interrupt_kernel()
self.check_alive()
except expect.TIMEOUT:
self.check_alive()
except expect.EOF:
# The program has exited
self._logger.info("The program has exited.")
self.status.set_down(self.kernel_pid, self.sudo)
break
self.close()
self.status.close()
class SendToIPython(object):
def __init__(self, nvim):
self.nvim = nvim
self.client = None
self.kerneldir = Path(jupyter_runtime_dir())
@neovim.function('RunningKernels', sync=True)
def running_kernels(self, args):
l = self.kerneldir.glob('kernel-*.json')
l = sorted(l, reverse=True, key=lambda f: f.stat().st_ctime)
return [f.name for f in l]
@neovim.command('SendTo', complete='customlist,RunningKernels', nargs='?')
def send_to(self, args):
cfs = args or self.running_kernels(None)
if not cfs:
self.nvim.command('echom "No kernel found"')
return
if self.client is not None:
self.client.stop_channels()
cf = cfs[0]
self.client = BlockingKernelClient()
self.client.load_connection_file(self.kerneldir / cf)
self.client.start_channels()
# run function once to register it for the `funcref` function
self.nvim.command('call SendLinesToJupyter()')
self.nvim.command(
'let g:send_target = {"send": funcref("SendLinesToJupyter")}')
self.nvim.command('echom "Sending to %s"' % cf)
@neovim.function('SendLinesToJupyter')
def send_lines(self, args):
if args:
self.client.execute('\n'.join(args[0]))
@neovim.function('SendComplete', sync=True)
def complete(self, args):
findstart, base = args
if self.client is None:
return -3 # no client setup yet: cancel silently and leave completion mode
if findstart:
line = self.nvim.current.line
if not line:
return -2 # empty line: cancel silently but stay in completion mode
pos = self.nvim.current.window.cursor[1]
try:
reply = self.client.complete(line,
pos,
reply=True,
timeout=timeout)['content']
except TimeoutError:
return -2
self.completions = [{
'word': w,
'info': ' '
} for w in reply['matches']]
return reply['cursor_start']
else:
# TODO: use vim's complete_add/complete_check for async operation
get_info(self.client, self.completions)
return {'words': self.completions, 'refresh': 'always'}
@neovim.function('SendCanComplete', sync=True)
def can_complete(self, args):
return args[
0] != '' and self.client is not None and self.client.is_alive()
class ToreeClient:
def __init__(self, connectionFileLocation):
self.client = BlockingKernelClient(
connection_file=connectionFileLocation)
self.client.load_connection_file(
connection_file=connectionFileLocation)
def is_alive(self):
return self.client.is_alive()
def is_ready(self):
try:
result = self.eval('1')
if result == '1':
return True
else:
return False
except:
return False
def wait_for_ready(self, timeout=TIMEOUT):
# Wait for initialization, by receiving an 'idle' message
# Flush Shell channel
abs_timeout = time.time() + timeout
while True:
try:
msg = self.client.shell_channel.get_msg(block=True,
timeout=0.2)
except Empty:
break
# Check if current time is ready check time plus timeout
if time.time() > abs_timeout:
raise RuntimeError("Kernel didn't respond in %d seconds" %
timeout)
# Flush IOPub channel
while True:
try:
msg = self.client.iopub_channel.get_msg(block=True,
timeout=0.2)
except Empty:
break
# Check if current time is ready check time plus timeout
if time.time() > abs_timeout:
raise RuntimeError("Kernel didn't respond in %d seconds" %
timeout)
def eval(self, code, timeout=TIMEOUT):
# Validate that remote kernel is available before submitting request
if self.client.is_alive() == False:
raise Exception(
'Problem connecting to remote kernel: Kernel is NOT alive')
debug_print('-----------------------------------------')
debug_print('Executing: ')
debug_pprint(code)
# submit request and retrieve the message id for the execution
msg_id = self.client.execute(code=code, allow_stdin=False)
debug_print('Message id for code execution:' + msg_id)
# now the kernel should be 'busy' with [parent_header][msg_id] being the current message
try:
busy_msg = self.client.iopub_channel.get_msg(block=True,
timeout=timeout)
except:
raise Exception('Error: Timeout retrieving busy status message')
debug_print('Current kernel status (%s): %s' %
(busy_msg['parent_header']['msg_id'],
busy_msg['content']['execution_state']))
if busy_msg['content']['execution_state'] == 'busy':
debug_print('busy_message received as expected')
else:
debug_print('Error: did not receive busy message for request %s' %
msg_id)
debug_pprint(busy_msg)
# Check message reply status (ok / error)
debug_print('Waiting for status reply')
reply = self.client.get_shell_msg(block=True, timeout=timeout)
debug_print('message reply: %s' % reply['content']['status'])
debug_pprint(reply)
type = ''
results = []
while True:
try:
msg = self.client.get_iopub_msg(timeout=timeout)
except:
raise Exception("Error: Timeout executing request")
debug_print('message')
debug_pprint(msg)
# validate that the responses are still related to current request
if msg['parent_header']['msg_id'] != msg_id:
debug_print('Warning: Invalid message id received ' +
msg['parent_header']['msg_id'] + ' expected ' +
msg_id)
continue
# validate execute_inputs are from current code
elif msg['msg_type'] == 'execute_input':
debug_print('current message status: ' + msg['msg_type'])
debug_print('current message content code: ' +
msg['content']['code'])
if msg['content']['code'] == code:
continue
# Stream results are being returned, accumulate them to results
elif msg['msg_type'] == 'stream':
type = 'stream'
results.append(msg['content']['text'])
continue
# Execute_Results are being returned:
# They can be text/plain or text/html
# accumulate them to results
elif msg['msg_type'] == 'execute_result':
debug_print('Received results of type: %s ' %
msg['content']['data'])
if 'text/plain' in msg['content']['data']:
type = 'text'
results.append(msg['content']['data']['text/plain'])
elif 'text/html' in msg['content']['data']:
type = 'html'
results.append(msg['content']['data']['text/html'])
continue
# When idle, responses have all been processed/returned
elif msg['msg_type'] == 'status':
debug_print('current message status: ' +
msg['content']['execution_state'])
if msg['content']['execution_state'] == 'idle':
break
else:
debug_print('Message ignored: %s' % msg['msg_type'])
if reply['content']['status'] == 'ok':
debug_print('Returning sucessful invocation result')
if type == 'html':
html = ''.join(results)
htmlWrapper = HtmlOutput(html)
return htmlWrapper
else:
return ''.join(results)
else:
debug_print('Returning failed invocation exception')
error = ''
if 'ename' in reply['content']:
error = reply['content']['ename']
error_message = ''
if 'evalue' in reply['content']:
error_message = reply['content']['evalue']
raise Exception('Error: %s - %s' % (error, error_message))
