def setup_kernel(cmd): """start an embedded kernel in a subprocess, and wait for it to be ready Returns ------- kernel_manager: connected KernelManager instance """ kernel = Popen([sys.executable, "-c", cmd], stdout=PIPE, stderr=PIPE, env=env) connection_file = os.path.join(IPYTHONDIR, "profile_default", "security", "kernel-%i.json" % kernel.pid) # wait for connection file to exist, timeout after 5s tic = time.time() while not os.path.exists(connection_file) and kernel.poll() is None and time.time() < tic + SETUP_TIMEOUT: time.sleep(0.1) if kernel.poll() is not None: o, e = kernel.communicate() e = py3compat.cast_unicode(e) raise IOError("Kernel failed to start:\n%s" % e) if not os.path.exists(connection_file): if kernel.poll() is None: kernel.terminate() raise IOError("Connection file %r never arrived" % connection_file) client = BlockingKernelClient(connection_file=connection_file) client.load_connection_file() client.start_channels() client.wait_for_ready() try: yield client finally: client.stop_channels() kernel.terminate()
def connect(connection_file): if connection_file not in clients: print "[nyroglancer] connecting to: " + connection_file kernel_client = BlockingKernelClient(connection_file=connection_file) kernel_client.load_connection_file() kernel_client.start_channels() clients[connection_file] = kernel_client return kernel_client return clients[connection_file]
def __init__(self, active_dir): # kernel config is stored in a dot file with the active directory config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) # right now we're spawning a child process for IPython. we can # probably work directly with the IPython kernel API, but the docs # don't really explain how to do it. args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): if (os.path.isfile(config)): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, it wasn't # up and running. 1.5 seconds was arbitrarily chosen (but seems to work) while not os.path.isfile(config): time.sleep(1.5) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch)
def __init__(self, active_dir): # kernel config is stored in a temp file config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, it wasn't up # and running. 1.5 seconds was arbitrarily chosen (but seems to work) time.sleep(1.5) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch)
def register_worker_magic(connection_info, magic_name='worker'): """Register a %worker magic, given connection_info. Both a line and cell magic are registered, which run the given cell in a remote kernel. """ ip = get_ipython() kc = BlockingKernelClient(**connection_info) kc.session.key = connection_info['key'] kc.start_channels() def remote(line, cell=None): """Run the current cell on a remote IPython kernel""" if cell is None: # both line and cell magic cell = line run_cell_remote(ip, kc, cell) ip.register_magic_function(remote, magic_kind='line', magic_name=magic_name) ip.register_magic_function(remote, magic_kind='cell', magic_name=magic_name)
def test_start_ipython_workers(loop, zmq_ctx): from jupyter_client import BlockingKernelClient with cluster(1) as (s, [a]): with Client(('127.0.0.1', s['port']), loop=loop) as e: info_dict = e.start_ipython_workers() info = first(info_dict.values()) key = info.pop('key') kc = BlockingKernelClient(**info) kc.session.key = key kc.start_channels() kc.wait_for_ready(timeout=10) msg_id = kc.execute("worker") reply = kc.get_shell_msg(timeout=10) assert reply['parent_header']['msg_id'] == msg_id assert reply['content']['status'] == 'ok' kc.stop_channels()
def remote_magic(line, cell=None): """A magic for running code on a specified remote worker The connection_info dict of the worker will be looked up as the first positional arg to the magic. The rest of the line (or the entire cell for a %%cell magic) will be passed to the remote kernel. Usage: info = e.start_ipython(worker)[worker] %remote info print(worker.data) """ # get connection info from IPython's user namespace ip = get_ipython() split_line = line.split(None, 1) info_name = split_line[0] if info_name not in ip.user_ns: raise NameError(info_name) connection_info = dict(ip.user_ns[info_name]) if not cell: # line magic, use the rest of the line if len(split_line) == 1: raise ValueError("I need some code to run!") cell = split_line[1] # turn info dict to hashable str for use as lookup key in _clients cache key = ','.join(map(str, sorted(connection_info.items()))) session_key = connection_info.pop('key') if key in remote_magic._clients: kc = remote_magic._clients[key] else: kc = BlockingKernelClient(**connection_info) kc.session.key = session_key kc.start_channels() kc.wait_for_ready(timeout=10) remote_magic._clients[key] = kc # actually run the code run_cell_remote(ip, kc, cell)
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 register_worker_magic(connection_info, magic_name='worker'): """Register a %worker magic, given connection_info. Both a line and cell magic are registered, which run the given cell in a remote kernel. """ ip = get_ipython() info = dict(connection_info) # copy key = info.pop('key') kc = BlockingKernelClient(**connection_info) kc.session.key = key kc.start_channels() def remote(line, cell=None): """Run the current cell on a remote IPython kernel""" if cell is None: # both line and cell magic cell = line run_cell_remote(ip, kc, cell) remote.client = kc # preserve reference on kc, largely for mocking ip.register_magic_function(remote, magic_kind='line', magic_name=magic_name) ip.register_magic_function(remote, magic_kind='cell', magic_name=magic_name)
def register_worker_magic(connection_info, magic_name="worker"): """Register a %worker magic, given connection_info. Both a line and cell magic are registered, which run the given cell in a remote kernel. """ ip = get_ipython() kc = BlockingKernelClient() kc.load_connection_info(connection_info) kc.start_channels() def remote(line, cell=None): """Run the current cell on a remote IPython kernel""" if cell is None: # both line and cell magic cell = line run_cell_remote(ip, kc, cell) remote.client = kc # preserve reference on kc, largely for mocking ip.register_magic_function(remote, magic_kind="line", magic_name=magic_name) ip.register_magic_function(remote, magic_kind="cell", magic_name=magic_name)
def setup_kernel(cmd): """start an embedded kernel in a subprocess, and wait for it to be ready Returns ------- kernel_manager: connected KernelManager instance """ kernel = Popen([sys.executable, '-c', cmd], stdout=PIPE, stderr=PIPE, env=env) connection_file = os.path.join(IPYTHONDIR, 'profile_default', 'security', 'kernel-%i.json' % kernel.pid) # wait for connection file to exist, timeout after 5s tic = time.time() while not os.path.exists(connection_file) \ and kernel.poll() is None \ and time.time() < tic + SETUP_TIMEOUT: time.sleep(0.1) if kernel.poll() is not None: o, e = kernel.communicate() e = py3compat.cast_unicode(e) raise IOError("Kernel failed to start:\n%s" % e) if not os.path.exists(connection_file): if kernel.poll() is None: kernel.terminate() raise IOError("Connection file %r never arrived" % connection_file) client = BlockingKernelClient(connection_file=connection_file) client.load_connection_file() client.start_channels() client.wait_for_ready() try: yield client finally: client.stop_channels() kernel.terminate()
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 run(self): cf = find_connection_file() client = BlockingKernelClient(connection_file=cf) client.load_connection_file() client.start_channels(shell=False, iopub=True, stdin=False, control=True, hb=False) while True: try: msg = client.get_iopub_msg(TIMEOUT) self.pub_q.put(msg) except Empty: pass if self.cmd_q.qsize(): cmd = self.cmd_q.get() if cmd is None: print('Client thread closing') break client.execute(cmd) self.ctrl_q.put(client.get_shell_msg())
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 create_kernel_client(self, ci): kernel_client = BlockingKernelClient() kernel_client.load_connection_info(ci) kernel_client.start_channels(shell=True, iopub=False, stdin=False, hb=False) return kernel_client
def run(self): cf = find_connection_file() client = BlockingKernelClient(connection_file=cf) client.load_connection_file() client.start_channels(shell=False, iopub=True, stdin=False, control=False, hb=False) while True: msg = client.get_iopub_msg() self.received.emit(msg)
def test_start_ipython_scheduler(loop, zmq_ctx): from jupyter_client import BlockingKernelClient with cluster(1) as (s, [a]): with Client(('127.0.0.1', s['port']), loop=loop) as e: info = e.start_ipython_scheduler() key = info.pop('key') kc = BlockingKernelClient(**info) kc.session.key = key kc.start_channels() msg_id = kc.execute("scheduler") reply = kc.get_shell_msg(timeout=10) kc.stop_channels()
def __init__(self, active_dir, pyspark): # kernel config is stored in a dot file with the active directory config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) # right now we're spawning a child process for IPython. we can # probably work directly with the IPython kernel API, but the docs # don't really explain how to do it. log_file = None if pyspark: os.environ["IPYTHON_OPTS"] = "kernel -f %s" % config pyspark = os.path.join(os.environ.get("SPARK_HOME"), "bin/pyspark") spark_log = os.environ.get("SPARK_LOG", None) if spark_log: log_file = open(spark_log, "w") spark_opts = os.environ.get("SPARK_OPTS", "") args = [pyspark] + spark_opts.split() # $SPARK_HOME/bin/pyspark <SPARK_OPTS> p = subprocess.Popen(args, stdout=log_file, stderr=log_file) else: args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): if os.path.isfile(config): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, so we'll # wait until the config file exists before moving on while os.path.isfile(config)==False: time.sleep(0.1) def close_file(): if log_file: log_file.close() atexit.register(close_file) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch)
def send_to(self, args): if args and args[0].endswith('(newest)'): args[0] = args[0][:-len('(newest)')] cf = find_connection_file(*args) if cf not in self.clients: client = BlockingKernelClient() client.load_connection_file(cf) client.start_channels() self.clients[cf] = client return cf
def remote_magic(line, cell=None): """A magic for running code on a specified remote worker The connection_info dict of the worker will be looked up as the first positional arg to the magic. The rest of the line (or the entire cell for a %%cell magic) will be passed to the remote kernel. Usage: info = e.start_ipython(worker)[worker] %remote info print(worker.data) """ # get connection info from IPython's user namespace ip = get_ipython() split_line = line.split(None, 1) info_name = split_line[0] if info_name not in ip.user_ns: raise NameError(info_name) connection_info = dict(ip.user_ns[info_name]) if not cell: # line magic, use the rest of the line if len(split_line) == 1: raise ValueError("I need some code to run!") cell = split_line[1] # turn info dict to hashable str for use as lookup key in _clients cache key = ",".join(map(str, sorted(connection_info.items()))) if key in remote_magic._clients: kc = remote_magic._clients[key] else: kc = BlockingKernelClient() kc.load_connection_info(connection_info) kc.start_channels() kc.wait_for_ready(timeout=10) remote_magic._clients[key] = kc # actually run the code run_cell_remote(ip, kc, cell)
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 Kernel(object): # kernel config is stored in a dot file with the active directory def __init__(self, config, active_dir, pyspark): # right now we're spawning a child process for IPython. we can # probably work directly with the IPython kernel API, but the docs # don't really explain how to do it. log_file = None if pyspark: os.environ["IPYTHON_OPTS"] = "kernel -f %s" % config pyspark = os.path.join(os.environ.get("SPARK_HOME"), "bin/pyspark") spark_log = os.environ.get("SPARK_LOG", None) if spark_log: log_file = open(spark_log, "w") spark_opts = os.environ.get("SPARK_OPTS", "") args = [pyspark] + spark_opts.split() # $SPARK_HOME/bin/pyspark <SPARK_OPTS> p = subprocess.Popen(args, stdout=log_file, stderr=log_file) else: args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): if os.path.isfile(config): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, so we'll # wait until the config file exists before moving on while os.path.isfile(config)==False: time.sleep(0.1) def close_file(): if log_file: log_file.close() atexit.register(close_file) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") python_patch_file = os.path.join(dirname, "langs", "python-patch.py") self.client.execute("%run " + python_patch_file) def _run_code(self, execution_id, code, timeout=0.1): # this function executes some code and waits for it to completely finish # before returning. i don't think that this is neccessarily the best # way to do this, but the IPython documentation isn't very helpful for # this particular topic. # # 1) execute code and grab the ID for that execution thread # 2) look for messages coming from the "iopub" channel (this is just a # stream of output) # 3) when we get a message that is one of the following, save relevant # data to `data`: # - execute_result - content from repr # - stream - content from stdout # - error - ansii encoded stacktrace # the final piece is that we check for when the message indicates that # the kernel is idle and the message's parent is the original execution # ID (msg_id) that's associated with our executing code. if this is the # case, we'll return the data and the msg_id and exit msg_id = self.client.execute(code, allow_stdin=False) request = { "id": execution_id, "msg_id": msg_id, "code": code, "status": "started" } sys.stdout.write(json.dumps(request) + '\n') sys.stdout.flush() output = { "id": execution_id, "msg_id": msg_id, "output": "", "stream": None, "image": None, "error": None } while True: try: reply = self.client.get_iopub_msg(timeout=timeout) except Empty: continue if "execution_state" in reply['content']: if reply['content']['execution_state']=="idle" and reply['parent_header']['msg_id']==msg_id: if reply['parent_header']['msg_type']=="execute_request": request["status"] = "complete" sys.stdout.write(json.dumps(request) + '\n') sys.stdout.flush() return elif reply['header']['msg_type']=="execute_result": output['output'] = reply['content']['data'].get('text/plain', '') output['stream'] = reply['content']['data'].get('text/plain', '') elif reply['header']['msg_type']=="display_data": if 'image/png' in reply['content']['data']: output['image'] = reply['content']['data']['image/png'] elif 'text/html' in reply['content']['data']: output['html'] = reply['content']['data']['text/html'] elif reply['header']['msg_type']=="stream": output['output'] += reply['content'].get('text', '') output['stream'] = reply['content'].get('text', '') elif reply['header']['msg_type']=="error": output['error'] = "\n".join(reply['content']['traceback']) # TODO: if we have something non-trivial to send back... sys.stdout.write(json.dumps(output) + '\n') sys.stdout.flush() # TODO: should probably get rid of all this output['stream'] = None output['image'] = None output['html'] = None def _complete(self, execution_id, code, timeout=0.5): # Call ipython kernel complete, wait for response with the correct msg_id, # and construct appropriate UI payload. # See below for an example response from ipython kernel completion for 'el' # # { # 'parent_header': # {u'username': u'ubuntu', u'version': u'5.0', u'msg_type': u'complete_request', # u'msg_id': u'5222d158-ada8-474e-88d8-8907eb7cc74c', u'session': u'cda4a03d-a8a1-4e6c-acd0-de62d169772e', # u'date': datetime.datetime(2015, 5, 7, 15, 25, 8, 796886)}, # 'msg_type': u'complete_reply', # 'msg_id': u'a3a957d6-5865-4c6f-a0b2-9aa8da718b0d', # 'content': # {u'matches': [u'elif', u'else'], u'status': u'ok', u'cursor_start': 0, u'cursor_end': 2, u'metadata': {}}, # 'header': # {u'username': u'ubuntu', u'version': u'5.0', u'msg_type': u'complete_reply', # u'msg_id': u'a3a957d6-5865-4c6f-a0b2-9aa8da718b0d', u'session': u'f1491112-7234-4782-8601-b4fb2697a2f6', # u'date': datetime.datetime(2015, 5, 7, 15, 25, 8, 803470)}, # 'buffers': [], # 'metadata': {} # } # msg_id = self.client.complete(code) request = { "id": execution_id, "msg_id": msg_id, "code": code, "status": "started" } sys.stdout.write(json.dumps(request) + '\n') sys.stdout.flush() output = { "id": execution_id, "msg_id": msg_id, "output": None, "image": None, "error": None } while True: try: reply = self.client.get_shell_msg(timeout=timeout) except Empty: continue if "matches" in reply['content'] and reply['msg_type']=="complete_reply" and reply['parent_header']['msg_id']==msg_id: results = [] for completion in reply['content']['matches']: result = { "value": completion, "dtype": "---" } if "." in code: # result['text'] = result['value'] # ".".join(result['value'].split(".")[1:]) result['text'] = result['value'] #.split('.')[-1] result["dtype"] = "function" else: result['text'] = result['value'] result["dtype"] = "" # type(globals().get(code)).__name__ results.append(result) output['output'] = results output['status'] = "complete" sys.stdout.write(json.dumps(output) + '\n') sys.stdout.flush() return def execute(self, execution_id, code, complete=False): if complete==True: return self._complete(execution_id, code) else: result = self._run_code(execution_id, code) if re.match("%?reset", code): # load our monkeypatches... k.client.execute("%matplotlib inline") k.client.execute(vars_patch) return result def get_packages(self): return self.execute("__get_packages()")
def main(): client = BlockingKernelClient() client.load_connection_file(find_connection_file()) code = "; ".join(sys.argv[1:]) if len(sys.argv) > 1 else sys.stdin.read() client.execute(code, allow_stdin=False)
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 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 Kernel(object): def __init__(self, active_dir): # kernel config is stored in a temp file config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, it wasn't up # and running. 1.5 seconds was arbitrarily chosen (but seems to work) time.sleep(1.5) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch) def _run_code(self, code, timeout=0.1): # this function executes some code and waits for it to completely finish before # returning. i don't think that this is neccessarily the best way to do this, but # the IPython documentation isn't very helpful for this particular topic. # # 1) execute code and grab the ID for that execution thread # 2) look for messages coming from the "iopub" channel (this is just a stream of output) # 3) when we get a message that is one of the following, save relevant data to `data`: # - execute_result - content from repr # - stream - content from stdout # - error - ansii encoded stacktrace # the final piece is that we check for when the message indicates that the kernel is idle # and the message's parent is the original execution ID (msg_id) that's associated with # our executing code. if this is the case, we'll return the data and the msg_id and exit msg_id = self.client.execute(code) data = None image = None while True: try: reply = self.client.get_iopub_msg(timeout=timeout) except Empty: continue if "execution_state" in reply['content']: if reply['content']['execution_state']=="idle" and reply['parent_header']['msg_id']==msg_id: if reply['parent_header']['msg_type']=="execute_request": return { "msg_id": msg_id, "output": data, "image": image } elif reply['header']['msg_type']=="execute_result": data = reply['content']['data'].get('text/plain', '') elif reply['header']['msg_type']=="display_data": image = reply['content']['data'].get('image/png', '') elif reply['header']['msg_type']=="stream": data = reply['content'].get('text', '') elif reply['header']['msg_type']=="error": data = "\n".join(reply['content']['traceback']) def execute(self, code): return self._run_code(code) def complete(self, code): # i couldn't figure out how to get the autocomplete working with the ipython # kernel (i couldn't get a completion_reply from the iopub), so we're using # jedi to do the autocompletion. the __autocomplete is defined in `autocomplete_patch` # above. return self.execute("__autocomplete('%s')" % code)
def create_kernel_client(self, ci): kernel_client = BlockingKernelClient() kernel_client.load_connection_info(ci) kernel_client.start_channels() return kernel_client
class Kernel(object): def __init__(self, active_dir): # kernel config is stored in a dot file with the active directory config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) # right now we're spawning a child process for IPython. we can # probably work directly with the IPython kernel API, but the docs # don't really explain how to do it. args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): os.remove(config) atexit.register(remove_config) # i found that if i tried to connect to the kernel immediately, it wasn't # up and running. 1.5 seconds was arbitrarily chosen (but seems to work) time.sleep(1.5) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch) def _run_code(self, code, timeout=0.1): # this function executes some code and waits for it to completely finish # before returning. i don't think that this is neccessarily the best # way to do this, but the IPython documentation isn't very helpful for # this particular topic. # # 1) execute code and grab the ID for that execution thread # 2) look for messages coming from the "iopub" channel (this is just a # stream of output) # 3) when we get a message that is one of the following, save relevant # data to `data`: # - execute_result - content from repr # - stream - content from stdout # - error - ansii encoded stacktrace # the final piece is that we check for when the message indicates that # the kernel is idle and the message's parent is the original execution # ID (msg_id) that's associated with our executing code. if this is the # case, we'll return the data and the msg_id and exit msg_id = self.client.execute(code) data = None image = None while True: try: reply = self.client.get_iopub_msg(timeout=timeout) except Empty: continue if "execution_state" in reply['content']: if reply['content']['execution_state'] == "idle" and reply[ 'parent_header']['msg_id'] == msg_id: if reply['parent_header']['msg_type'] == "execute_request": return { "msg_id": msg_id, "output": data, "image": image } elif reply['header']['msg_type'] == "execute_result": data = reply['content']['data'].get('text/plain', '') elif reply['header']['msg_type'] == "display_data": image = reply['content']['data'].get('image/png', '') elif reply['header']['msg_type'] == "stream": data = reply['content'].get('text', '') elif reply['header']['msg_type'] == "error": data = "\n".join(reply['content']['traceback']) def execute(self, code): return self._run_code(code) def complete(self, code): # i couldn't figure out how to get the autocomplete working with the # ipython kernel (i couldn't get a completion_reply from the iopub), so # we're using jedi to do the autocompletion. the __autocomplete is # defined in `autocomplete_patch` above. return self.execute("__autocomplete('%s')" % code) def get_dataframes(self): return self.execute("__get_variables()")
def setup_kernel(cmd): """start an embedded kernel in a subprocess, and wait for it to be ready Returns ------- kernel_manager: connected KernelManager instance """ def connection_file_ready(connection_file): """Check if connection_file is a readable json file.""" if not os.path.exists(connection_file): return False try: with open(connection_file) as f: json.load(f) return True except ValueError: return False kernel = Popen([sys.executable, "-c", cmd], stdout=PIPE, stderr=PIPE, encoding="utf-8") try: connection_file = os.path.join( paths.jupyter_runtime_dir(), "kernel-%i.json" % kernel.pid, ) # wait for connection file to exist, timeout after 5s tic = time.time() while ( not connection_file_ready(connection_file) and kernel.poll() is None and time.time() < tic + SETUP_TIMEOUT ): time.sleep(0.1) # Wait 100ms for the writing to finish time.sleep(0.1) if kernel.poll() is not None: o, e = kernel.communicate() raise OSError("Kernel failed to start:\n%s" % e) if not os.path.exists(connection_file): if kernel.poll() is None: kernel.terminate() raise OSError("Connection file %r never arrived" % connection_file) client = BlockingKernelClient(connection_file=connection_file) client.load_connection_file() client.start_channels() client.wait_for_ready() try: yield client finally: client.stop_channels() finally: kernel.terminate() kernel.wait() # Make sure all the fds get closed. for attr in ["stdout", "stderr", "stdin"]: fid = getattr(kernel, attr) if fid: fid.close()
def kernel_client(self): self.kc = BlockingKernelClient() self.kc.load_connection_info(self.connection_info) self.kc.start_channels()
'OKCYAN': '\033[96m', 'OKGREEN': '\033[92m', 'WARNING': '\033[93m', 'FAIL': '\033[91m', 'ENDC': '\033[0m', 'BOLD': '\033[1m', 'UNDERLINE': '\033[4m', } print(lookup[color], end='') print(*args, end='') print(lookup['ENDC']) # setup by automatically finding a running kernel cf = find_connection_file() client = BlockingKernelClient(connection_file=cf) client.load_connection_file() client.start_channels() # simplest usage - execute statments and check if OK msgid = client.execute('a = 2') ret = client.get_shell_msg() status = ret['content']['status'] if status == 'ok': print('statement executed ok') elif status == 'error': ename = ret['content']['ename'] print('there was a %s exception, which will also appear on the ' 'iopub channel' % ename) # listen to what's going on in the kernel with blocking calls,
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()
def setup_kernel(cmd): """start an embedded kernel in a subprocess, and wait for it to be ready Returns ------- kernel_manager: connected KernelManager instance """ def connection_file_ready(connection_file): """Check if connection_file is a readable json file.""" if not os.path.exists(connection_file): return False try: with open(connection_file) as f: json.load(f) return True except ValueError: return False kernel = Popen([sys.executable, '-c', cmd], stdout=PIPE, stderr=PIPE) try: # connection_file = os.path.join( # paths.jupyter_runtime_dir(), # 'kernel-%i.json' % kernel.pid, # ) connection_file = "./remotekernel2.json" # # wait for connection file to exist, timeout after 5s # tic = time.time() # while not connection_file_ready(connection_file) \ # and kernel.poll() is None \ # and time.time() < tic + SETUP_TIMEOUT: # time.sleep(0.1) # # # Wait 100ms for the writing to finish # time.sleep(0.1) # # if kernel.poll() is not None: # o,e = kernel.communicate() # e = py3compat.cast_unicode(e) # raise IOError("Kernel failed to start:\n%s" % e) # # if not os.path.exists(connection_file): # if kernel.poll() is None: # kernel.terminate() # raise IOError("Connection file %r never arrived" % connection_file) client = BlockingKernelClient(connection_file=connection_file) client.load_connection_file() # client.start_channels() # client.wait_for_ready() # client. comm_manager = CommManager(parent=client, kernel_client=client) comm = comm_manager.new_comm("neos_comm") # comm.open() nep = Nep(comm, kernel_id=client.kernel_info()) nep.start() # self.nep_start(comm) # print(client.) try: yield client finally: # client.stop_channels() pass finally: # kernel.terminate() pass
class JupyterRAMUsage(Stat): """ tag: ``jupyter.ram_usage`` settings: .. code-block:: javascript { "connection info": "", "query interval [s]": 1 } Tracks the RAM usage of all variables in a user-specified jupyter notebook. If no connection info is given in the settings, take the kernel with the latest start date. ``connection info`` must be a string containing the info displayed when running ``%connect_info`` in a jupyter notebook (you can directly copy-paste it). ``query interval [s]`` specifies how often the thread running in the jupyter notebook should read the variables. The lower this is, the higher the resolution of the stat but it might start affecting the speed of your notebook when too low. Note that RAM tracked in this way is not equal to the actual RAM the OS needs because some further optimization is done by e. g. numpy to reduce the OS memory usage. """ name = 'RAM Usage of objects in a Python Jupyter Notebook [MB]' base_tag = 'ram_usage' default_settings = { 'connection info': '', # how often the memory usage is read in the jupyter notebook 'query interval [s]': 1. } @classmethod def _read_latest_connection_file(cls): """ Reads the latest jupyter kernel connection file. https://jupyter.readthedocs.io/en/latest/projects/jupyter-directories.html. """ runtime_dir = jupyter_runtime_dir() files = glob.glob(os.path.join(runtime_dir, 'kernel-*.json')) if len(files) == 0: return None # use the latest connection file connection_file = max(files, key=os.path.getctime) with open(connection_file, 'r') as f: return json.load(f) @classmethod def get_connection_info(cls): """ Get the target kernel connection info. Returns a dictionary of the connection info supplied in the settings, or the latest started kernel if none is given. Retuns `None` if no kernel has been found. """ if len(cls.settings['connection info']) == 0: return cls._read_latest_connection_file() return json.loads(cls.settings['connection info']) @classmethod def check_availability(cls): # the stat is not available if no suitable connection info # can be found if cls.get_connection_info() is None: raise exceptions.StatNotAvailableError( 'Could not find any running kernel.') def __init__(self, fps): self.config = self.get_connection_info() data_dir = appdirs.user_data_dir('permon', 'bminixhofer') os.makedirs(data_dir, exist_ok=True) self.usage_file = os.path.join(data_dir, 'jupyter_ram_usage.csv') open(self.usage_file, 'w').close() # self.setup_code is the code that is run in the notebook when the # stat is instantiated. It starts a thread which reads the memory # usage of all public variables in a set interval and saves it to a # csv file in the user data directory self.setup_code = f""" if '_permon_running' not in globals() or not _permon_running: import threading import csv import sys import time from pympler import asizeof from types import ModuleType def _permon_get_ram_usage_per_object(): while _permon_running: ram_usage = [] global_vars = [key for key in globals() if not key.startswith('_')] for name in global_vars: value = globals()[name] if name in globals() else None if isinstance(value, ModuleType): continue try: ram_usage.append((name, asizeof.asizeof(value))) except TypeError: continue with open('{self.usage_file}', 'w') as f: writer = csv.writer(f, delimiter=',') for name, ram in ram_usage: writer.writerow([name, ram]) time.sleep({self.settings['query interval [s]']}) _permon_thread = threading.Thread(target=_permon_get_ram_usage_per_object) _permon_running = True _permon_thread.start() """ self.teardown_code = """ _permon_running = False """ self.client = BlockingKernelClient() self.client.load_connection_info(self.config) self.client.start_channels() self.client.execute(self.setup_code) super(JupyterRAMUsage, self).__init__(fps=fps) def __del__(self): # stop the thread running in the jupyter notebook # and stop the connection to the kernel upon deletion self.client.execute(self.teardown_code) self.client.stop_channels() def get_stat(self): # reads the csv file the setup code has written to ram_usage = [] with open(self.usage_file, 'r') as f: reader = csv.reader(f) for row in reader: ram_usage.append((row[0], float(row[1]) / 1000**2)) # sort the ram_usage list so that the largest variables come first ram_usage = sorted(ram_usage, key=lambda x: x[1], reverse=True) # return the sum of RAM usage and the variables taking up the most RAM return sum(x[1] for x in ram_usage), ram_usage[:5] @property def minimum(self): return 0 @property def maximum(self): return None
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
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
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)
import logging import time from queue import Empty # Py 3 from jupyter_client import connect from jupyter_client import BlockingKernelClient logging.basicConfig(level=logging.DEBUG) f = connect.find_connection_file() print("loaded") 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
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 Kernel(object): def __init__(self, active_dir, pyspark): # kernel config is stored in a dot file with the active directory config = os.path.join(active_dir, ".kernel-%s.json" % str(uuid.uuid4())) # right now we're spawning a child process for IPython. we can # probably work directly with the IPython kernel API, but the docs # don't really explain how to do it. log_file = None if pyspark: os.environ["IPYTHON_OPTS"] = "kernel -f %s" % config pyspark = os.path.join(os.environ.get("SPARK_HOME"), "bin/pyspark") spark_log = os.environ.get("SPARK_LOG", None) if spark_log: log_file = open(spark_log, "w") spark_opts = os.environ.get("SPARK_OPTS", "") args = [pyspark] + spark_opts.split() # $SPARK_HOME/bin/pyspark <SPARK_OPTS> p = subprocess.Popen(args, stdout=log_file, stderr=log_file) else: args = [sys.executable, '-m', 'IPython', 'kernel', '-f', config] p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # when __this__ process exits, we're going to remove the ipython config # file and kill the ipython subprocess atexit.register(p.terminate) def remove_config(): os.remove(config) atexit.register(remove_config) def close_file(): if log_file: log_file.close() atexit.register(close_file) # i found that if i tried to connect to the kernel immediately, it wasn't # up and running. 1.5 seconds was arbitrarily chosen (but seems to work) time.sleep(1.5) # fire up the kernel with the appropriate config self.client = BlockingKernelClient(connection_file=config) self.client.load_connection_file() self.client.start_channels() # load our monkeypatches... self.client.execute("%matplotlib inline") self.client.execute(autocomplete_patch) self.client.execute(vars_patch) def _run_code(self, code, timeout=0.1): # this function executes some code and waits for it to completely finish # before returning. i don't think that this is neccessarily the best # way to do this, but the IPython documentation isn't very helpful for # this particular topic. # # 1) execute code and grab the ID for that execution thread # 2) look for messages coming from the "iopub" channel (this is just a # stream of output) # 3) when we get a message that is one of the following, save relevant # data to `data`: # - execute_result - content from repr # - stream - content from stdout # - error - ansii encoded stacktrace # the final piece is that we check for when the message indicates that # the kernel is idle and the message's parent is the original execution # ID (msg_id) that's associated with our executing code. if this is the # case, we'll return the data and the msg_id and exit msg_id = self.client.execute(code) output = { "msg_id": msg_id, "output": None, "image": None, "error": None } while True: try: reply = self.client.get_iopub_msg(timeout=timeout) except Empty: continue if "execution_state" in reply['content']: if reply['content']['execution_state']=="idle" and reply['parent_header']['msg_id']==msg_id: if reply['parent_header']['msg_type']=="execute_request": return output elif reply['header']['msg_type']=="execute_result": output['output'] = reply['content']['data'].get('text/plain', '') elif reply['header']['msg_type']=="display_data": output['image'] = reply['content']['data'].get('image/png', '') elif reply['header']['msg_type']=="stream": output['output'] = reply['content'].get('text', '') elif reply['header']['msg_type']=="error": output['error'] = "\n".join(reply['content']['traceback']) def execute(self, code): return self._run_code(code) def complete(self, code): # i couldn't figure out how to get the autocomplete working with the # ipython kernel (i couldn't get a completion_reply from the iopub), so # we're using jedi to do the autocompletion. the __autocomplete is # defined in `autocomplete_patch` above. return self.execute("__autocomplete('%s')" % code) def get_dataframes(self): return self.execute("__get_variables()")
import logging import time from queue import Empty # Py 3 from jupyter_client import connect from jupyter_client import BlockingKernelClient logging.basicConfig(level=logging.DEBUG) f = connect.find_connection_file() print("loaded") 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):