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))