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
