def main(kid): # Load connection info and init communications. cf = find_connection_file(kid) km = BlockingKernelClient(connection_file=cf) km.load_connection_file() km.start_channels() # Define a function that is useful from within the user's notebook: juneau_connect() can be # used to directly connect the notebook to the source database. Note that this includes the # full "root" credentials. # FIXME: allow for user-specific credentials on SQL tables. The DBMS may also not be at localhost. code = f""" from sqlalchemy import create_engine def juneau_connect(): engine = create_engine( "postgresql://{config.sql.name}:{config.sql.password}@{config.sql.host}/{config.sql.dbname}", connect_args={{ "options": "-csearch_path='{config.sql.dbs}'" }} ) return engine.connect() """ km.execute_interactive(code, timeout=TIMEOUT) km.stop_channels()
def main(kid, var): # Load connection info and init communications. cf = find_connection_file(kid) # str(port)) km = BlockingKernelClient(connection_file=cf) km.load_connection_file() km.start_channels() code = f""" import pandas as pd import numpy as np if type({var}) in [pd.DataFrame, np.ndarray, list]: print({var}.to_json(orient='split', index=False)) """ km.execute_interactive(code, timeout=TIMEOUT) km.stop_channels()
def JKConnect(self) -> None: runtime_dir = pathlib.Path(jupyter_core.paths.jupyter_runtime_dir()) connection_files = runtime_dir.glob("kernel-*.json") source = '\n'.join( connection_file.name.lstrip('kernel-').rstrip('.json') + ' ' + datetime.fromtimestamp(connection_file.stat().st_ctime).strftime( "%m/%d %H:%M") for connection_file in connection_files) proc = subprocess.run("fzf-tmux|awk '{print $1}'", input=source, stdout=PIPE, shell=True, text=True) connection_file = 'kernel-%s.json' % proc.stdout.strip() connection_file = runtime_dir.joinpath(connection_file).as_posix() kc = BlockingKernelClient() try: kc.load_connection_file(connection_file) kc.execute_interactive('', timeout=1) except (TimeoutError, FileNotFoundError): self.nvim.command("echoerr 'Selected connection is dead!'") else: self.executor = ExecutePreprocessor() self.executor.kc = kc self.nvim.command("echo 'Successfully connected!'")
class PyExecutor(pyexecutor_pb2_grpc.PyExecutorServicer): def __init__(self): self.manager = MultiKernelManager() kernel_id = self.manager.start_kernel() self.kernel = self.manager.get_kernel(kernel_id) self.client = BlockingKernelClient() self.client.load_connection_file(self.kernel.connection_file) def Execute(self, request, context): response = self.client.execute_interactive( code=request.command, user_expressions={'test': request.expression}) expression = response['content']['user_expressions']['test']['data'] result = expression[ 'text/html'] if 'text/html' in expression else expression[ 'text/plain'] return pyexecutor_pb2.ExecuteResponse(result=result)
def setup(): global client kernel = Popen([sys.executable, '-m', 'ipykernel'], stdout=PIPE, stderr=PIPE) connection_file = os.path.join( paths.jupyter_runtime_dir(), 'kernel-%i.json' % kernel.pid, ) sleep(1) client = BlockingKernelClient(connection_file=connection_file) client.load_connection_file() client.start_channels() client.wait_for_ready() loaded = client.execute_interactive(load_splonky) if loaded['content']['status'] == 'error': raise Exception("Could not load core Splonky libraries") os_process_id = re.findall('.*\/kernel-(\d+)\.json$', connection_file)[0] return os_process_id
def execute_from_command_line(): if sys.argv.count('--existing') != 1: raise ValueError(f'{sys.argv}\n' f'--existing argument must occur once only.') kernel_arg_index = sys.argv.index('--existing') try: kernel_name = sys.argv[kernel_arg_index + 1] except IndexError: # Following the command-line API of jupyter console, qtconsole etc, the --existing argument # can be used without a value, meaning use the kernel whose connection file has most # recently been accessed. We support that here when --existing is the last element of the # command line. Otherwise, the behavior of the no-argument-value form can be achieved with # --existing ''. kernel_name = None else: sys.argv.pop(kernel_arg_index + 1) sys.argv.pop(kernel_arg_index) if {'shell', 'shell_plus'} & set(sys.argv): # Special case: use `jupyter console` for management commands requesting a python shell. argv = [ 'jupyter', 'console', '--Completer.use_jedi=False', '--existing' ] if kernel_name: argv.append(kernel_name) os.execlp(argv[0], *argv) connection_file = find_connection_file( kernel_name) if kernel_name else find_connection_file() kernel_client = BlockingKernelClient(connection_file=connection_file) kernel_client.load_connection_file() response = kernel_client.execute_interactive(f""" from devkernel.kernel import execute_from_command_line execute_from_command_line('{json.dumps(sys.argv)}') """) exit_status = 0 if response['metadata']['status'] == 'ok' else 1 sys.exit(exit_status)
def main(kid, var, pid): # load connection info and init communication cf = find_connection_file(kid) # str(port)) km = BlockingKernelClient(connection_file=cf) km.load_connection_file() km.start_channels() # Step 0: get all the inputs load_input_code = f""" proc_id="{pid}" var={var} var_name="{var}" sql_name = "{cfg.sql_name}" sql_password = "******" sql_dbname = "{cfg.sql_dbname}" sql_schema_name = "{cfg.sql_schema_name}" sql_table_name = "{cfg.sql_table_name}" json_file_name = "/Users/peterchan/Desktop/GitHub/jupyter-extension/juneau_extension/data_file.json" """ # Step 1: access the table and convert it to JSON request_var_code = f""" import numpy as np import pandas as pd import json import copy if type(var) is pd.DataFrame or type(var) is np.ndarray or type(var) is list: df_json_string = var.to_json(orient='split', index=False) df_ls = json.loads(df_json_string)['data'] df_ls_copy = copy.deepcopy(df_ls) """ # Step 2: define the functions used to write to the JSON file json_lock_code = """ def initialize(): data = { "ownerID": "", "id123": "operating", "id124": "finish" } with open("./juneau_extension/data_file.json", "w") as file: json.dump(data, file, indent=4) def acquire_lock(pid): with open(json_file_name, "r+") as file: try: data = json.load(file) if data["ownerID"]: return False else: file.seek(0) file.truncate() data['ownerID'] = pid json.dump(data, file, indent=4) return True except Exception: return False def release_lock(pid): with open(json_file_name, "r+") as file: data = json.load(file) if data['ownerID'] == pid: file.seek(0) file.truncate() data['ownerID'] = "" json.dump(data, file, indent=4) # input: id of the process # remove from the file if the process is completed/ terminated/ timed out def update_exec_status(status, pid): done = False while not done: success = acquire_lock(pid) if success: try: with open(json_file_name, "r+") as file: data = json.load(file) if not data['ownerID'] == pid: continue file.seek(0) file.truncate() data[pid] = status json.dump(data, file, indent=4) release_lock(pid) done = True except Exception: continue return True """ # Step 3: connect to SQL and insert the table insert_code = """ from sqlalchemy import create_engine conn_string = f"postgresql://{sql_name}:{sql_password}@localhost/{sql_dbname}" table_string = f"{sql_schema_name}.{sql_table_name}" engine = create_engine(conn_string) with engine.connect() as connection: insertion_string = f'CREATE TABLE {sql_schema_name}.{var_name} ("A" int, "B" int, "C" int, "D" int);' for ls in df_ls_copy: insertion_string += f"INSERT INTO {sql_schema_name}.{var_name} VALUES ({ls[0]}, {ls[1]}, {ls[2]}, {ls[3]});" connection.execute(insertion_string) print(proc_id) update_exec_status("done", proc_id) rows = connection.execute(f"select * from {sql_schema_name}.{var_name} limit 5;") for row in rows: print(row) """ code = load_input_code + request_var_code + json_lock_code + insert_code km.execute_interactive(code, timeout=TIMEOUT) km.stop_channels()
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 DaisyWorkflow_client: def __init__(self, connection_file=None, executable=False): import os self.alg_keys = [] self.dat_keys = [] #super().__init__() self.kc = BlockingKernelClient() if connection_file == None: raise Exception( 'Please Specific Connection file to Remote IPython Kernel first' ) if os.access(connection_file, os.R_OK) == False: raise Exception('The connection file can no be read!') self.kc.load_connection_file(connection_file) try: self.kc.start_channels() except RuntimeError: raise Exception( 'Can not start channels, Please CHECK REMOTE KERNEL STATUS') self.executable = executable self.remote_name = None self.alg_clients = {} def initialize(self, class_name=None, workflow_name=None, workflow_cfgfile=None, algorithms_cfgfile=None): import os, json if class_name == None: raise Exception('Please Specific Workflow class name first') cmd = "from Daisy.Workflow import " + class_name self.kc.execute_interactive(cmd) if workflow_name == None: workflow_name = class_name self.remote_name = workflow_name cmd = self.remote_name + " = " + class_name + "('" + workflow_name + "')" self.kc.execute_interactive(cmd) if workflow_cfgfile == None: raise Exception('Please Specific Workflow Config file first') if os.access(workflow_cfgfile, os.R_OK) == False: raise Exception('The Workflow Config file can no be read!') with open(workflow_cfgfile, 'r') as json_file: string = json_file.read() wf_cfg = json.loads(string) temp_name = 'cfg_dict' + str(randint(1, 1000000)) cmd = temp_name + ' = ' + str(wf_cfg) self.kc.execute_interactive(cmd) cmd = self.remote_name + ".initialize(workflow_engine='PyWorkflowEngine', workflow_environment = " + temp_name + ")" self.kc.execute_interactive(cmd) self.kc.execute_interactive('del ' + temp_name) def setLogLevel(self, level): pass #Sniper.setLogLevel(level) #super().setLogLevel(level) def data_keys(self): cmd = self.remote_name + ".data_keys()" msg_id = self.kc.execute(cmd) exe_msg = self.execute_status(msg_id) dat_keys = [] if 'data' in exe_msg: msg = exe_msg['data'] msg = msg[msg.find("[") + 1:msg.rfind("]")] items = msg.split(',') #items = exe_msg['data'].split('\n') for i in items: begin = i.find("'") end = i.rfind("'") dat_keys.append(i[begin + 1:end]) self.dat_keys = dat_keys return self.dat_keys def algorithm_keys(self): cmd = self.remote_name + ".algorithm_keys()" msg_id = self.kc.execute(cmd) exe_msg = self.execute_status(msg_id) alg_keys = [] if 'data' in exe_msg: msg = exe_msg['data'] msg = msg[msg.find("[") + 1:msg.rfind("]")] items = msg.split(',') for i in items: begin = i.find("'") end = i.rfind("'") alg_keys.append(i[begin + 1:end]) self.alg_keys = alg_keys return self.alg_keys def get_data(self, data_name): raise Exception('Cannot get DataObject from Server') #return self.engine.datastore[data_name] def get_algorithm(self, algorithm_name): if algorithm_name in self.alg_clients.keys(): return self.alg_clients[algorithm_name] cmd = self.remote_name if algorithm_name in self.alg_keys: cmd = cmd + ".get_algorithm('" + algorithm_name + "')" elif algorithm_name in self.algorithm_keys(): cmd = cmd + ".get_algorithm('" + algorithm_name + "')" else: return False msg_id = self.kc.execute(cmd) exe_msg = self.execute_status(msg_id) self.alg_clients[algorithm_name] = DaisyAlgorithm_client( self.kc, exe_msg['code']) print(exe_msg['data']) return self.alg_clients[algorithm_name] def execute(self): pass #raise Exception('Must') def finalize(self): cmd = self.remote_name cmd = cmd + ".finalize()" self.kc.execute_interactive(cmd) #exe_msg = self.execute_status(msg_id) #print(exe_msg) #self.engine.finalize() def execute_status(self, msg_id): code_flag = False data_flag = False exe_msg = {'msg_id': msg_id} while True: try: kc_msg = self.kc.get_iopub_msg(timeout=5) if 'parent_header' in kc_msg and kc_msg['parent_header'][ 'msg_id'] != exe_msg['msg_id']: continue #_output_hook_default(kc_msg) msg_type = kc_msg['header']['msg_type'] msg_cont = kc_msg['content'] if msg_type == 'stream': exe_msg[msg_cont['name']] = msg_cont['text'] elif msg_type == 'error': exe_msg['error'] = msg_cont['traceback'] print(msg_cont['traceback']) break elif msg_type in ('display_data', 'execute_result'): if 'data' in msg_cont: data_flag = True exe_msg['data'] = msg_cont['data'].get( 'text/plain', '') if 'code' in msg_cont: code_flag = True exe_msg['code'] = msg_cont['code'] if code_flag and data_flag: break except: print('timeout kc.get_iopub_msg') break return exe_msg