add %daskworker

leylines-support/02-dask.ipy

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