add %daskworker

This commit is contained in:
xenia 2021-06-16 07:55:21 -04:00
parent 7068a88490
commit 61185515f6
1 changed files with 40 additions and 2 deletions

View File

@ -46,7 +46,7 @@ def dask(line, local_ns):
workers = [node for node in db.get_nodes()
if node.id != server_node.id and node.ssh_key is not None]
dest = f"{server_node.ip}:{DEFAULT_PORT}"
print("Connecting to APEX at", dest)
print("connecting to APEX at", dest)
client = Client(dest)
workers_by_ip = {str(node.ip):node for node in workers}
workers_status = {str(node.ip):False for node in workers}
@ -59,4 +59,42 @@ def dask(line, local_ns):
print(f"{node.name} ({node.ip}): down")
local_ns['client'] = client
del dask, ipy
@ipy.core.magic.register_line_magic
@ipy.core.magic.needs_local_scope
def daskworker(line, local_ns):
"picks a worker to launch ipython on"
from distributed.client import default_client
import subprocess
import json
import tempfile
import time
import os
import shutil
splitter = None
if os.environ.get("TMUX", "") != "":
splitter = ["tmux", "split"]
elif shutil.which("gnome-terminal") is not None:
splitter = ["/usr/bin/env", "gnome-terminal", "--"]
else:
raise Exception("don't know how to split terminal!")
client = default_client()
workers = client.scheduler_info()["workers"].items()
sorted_workers = sorted(workers,
key=lambda w: w[1]["metrics"]["memory"] - w[1]["memory_limit"]
+ (100000000000 * w[1]["metrics"]["executing"]))
worker = sorted_workers[0][0]
print("starting ipython kernel on", sorted_workers[0][1]["id"])
info = list(client.start_ipython_workers([worker]).values())[0]
if isinstance(info["key"], bytes):
info["key"] = info["key"].decode()
with tempfile.NamedTemporaryFile(mode="w", prefix="apex-") as f:
json.dump(info, f)
f.flush()
subprocess.check_call(splitter + ["/usr/bin/env", "jupyter", "console", "--existing",
f.name])
time.sleep(1)
del dask, daskworker, ipy