forked from M-Labs/artiq
add rudimentary master/controller/client
This commit is contained in:
parent
2f58cf6aff
commit
59194176a9
|
@ -0,0 +1,26 @@
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from artiq.management.worker import Worker
|
||||||
|
|
||||||
|
|
||||||
|
class Scheduler:
|
||||||
|
def __init__(self, loop):
|
||||||
|
self.loop = loop
|
||||||
|
self.queue = asyncio.Queue()
|
||||||
|
|
||||||
|
def __enter__(self):
|
||||||
|
self.worker = Worker(self.loop)
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, type, value, traceback):
|
||||||
|
self.loop.run_until_complete(self.worker.end_process())
|
||||||
|
del self.worker
|
||||||
|
|
||||||
|
def add_run_once(self, item, timeout):
|
||||||
|
yield from self.queue.put((item, timeout))
|
||||||
|
|
||||||
|
def task(self):
|
||||||
|
while True:
|
||||||
|
item, timeout = yield from self.queue.get()
|
||||||
|
result = yield from self.worker.run(item, timeout)
|
||||||
|
print(result)
|
|
@ -0,0 +1,74 @@
|
||||||
|
import sys
|
||||||
|
import asyncio
|
||||||
|
import subprocess
|
||||||
|
import signal
|
||||||
|
import json
|
||||||
|
|
||||||
|
|
||||||
|
class WorkerFailed(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class Worker:
|
||||||
|
def __init__(self, loop, send_timeout=0.5, start_reply_timeout=1.0,
|
||||||
|
term_timeout=1.0):
|
||||||
|
self.send_timeout = send_timeout
|
||||||
|
self.start_reply_timeout = start_reply_timeout
|
||||||
|
self.term_timeout = term_timeout
|
||||||
|
loop.run_until_complete(self.create_process())
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def create_process(self):
|
||||||
|
self.process = yield from asyncio.create_subprocess_exec(
|
||||||
|
sys.executable, "-m", "artiq.management.worker_impl",
|
||||||
|
stdout=subprocess.PIPE, stdin=subprocess.PIPE)
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def _send(self, obj, timeout):
|
||||||
|
line = json.dumps(obj)
|
||||||
|
self.process.stdin.write(line.encode())
|
||||||
|
self.process.stdin.write("\n".encode())
|
||||||
|
try:
|
||||||
|
fut = self.process.stdin.drain()
|
||||||
|
if fut is not (): # FIXME: why does Python return this?
|
||||||
|
yield from asyncio.wait_for(fut, timeout=timeout)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
raise WorkerFailed("Timeout sending data from worker")
|
||||||
|
except:
|
||||||
|
raise WorkerFailed("Failed to send data to worker")
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def _recv(self, timeout):
|
||||||
|
try:
|
||||||
|
line = yield from asyncio.wait_for(
|
||||||
|
self.process.stdout.readline(), timeout=timeout)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
raise WorkerFailed("Timeout receiving data from worker")
|
||||||
|
if not line:
|
||||||
|
raise WorkerFailed(
|
||||||
|
"Worker ended unexpectedly while trying to receive data")
|
||||||
|
try:
|
||||||
|
obj = json.loads(line.decode())
|
||||||
|
except:
|
||||||
|
raise WorkerFailed("Worker sent invalid JSON data")
|
||||||
|
return obj
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def run(self, run_params, result_timeout):
|
||||||
|
yield from self._send(run_params, self.send_timeout)
|
||||||
|
obj = yield from self._recv(self.start_reply_timeout)
|
||||||
|
if obj != "ack":
|
||||||
|
raise WorkerFailed("Incorrect acknowledgement")
|
||||||
|
result = yield from self._recv(result_timeout)
|
||||||
|
return result
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def end_process(self):
|
||||||
|
if self.process.returncode is not None:
|
||||||
|
return
|
||||||
|
self.process.send_signal(signal.SIGTERM)
|
||||||
|
try:
|
||||||
|
yield from asyncio.wait_for(
|
||||||
|
self.process.wait(), timeout=self.term_timeout)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
self.process.send_signal(signal.SIGKILL)
|
|
@ -0,0 +1,45 @@
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import importlib
|
||||||
|
|
||||||
|
|
||||||
|
def import_in_folder(path, name):
|
||||||
|
try:
|
||||||
|
del sys.modules[name] # force path search
|
||||||
|
except KeyError:
|
||||||
|
pass
|
||||||
|
loader = importlib.find_loader(name, [path])
|
||||||
|
if loader is None:
|
||||||
|
raise ImportError("Could not find loader")
|
||||||
|
return loader.load_module()
|
||||||
|
|
||||||
|
|
||||||
|
def run(path, name):
|
||||||
|
module = import_in_folder(path, name)
|
||||||
|
module.main()
|
||||||
|
|
||||||
|
|
||||||
|
def put_object(obj):
|
||||||
|
ds = json.dumps(obj)
|
||||||
|
sys.__stdout__.write(ds)
|
||||||
|
sys.__stdout__.write("\n")
|
||||||
|
sys.__stdout__.flush()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
sys.stdout = sys.stderr
|
||||||
|
|
||||||
|
while True:
|
||||||
|
line = sys.__stdin__.readline()
|
||||||
|
obj = json.loads(line)
|
||||||
|
put_object("ack")
|
||||||
|
|
||||||
|
try:
|
||||||
|
run(**obj)
|
||||||
|
except Exception as e:
|
||||||
|
put_object({"status": "failed", "message": str(e)})
|
||||||
|
else:
|
||||||
|
put_object({"status": "ok"})
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
|
@ -0,0 +1,45 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import socket
|
||||||
|
import json
|
||||||
|
|
||||||
|
|
||||||
|
def _get_args():
|
||||||
|
parser = argparse.ArgumentParser(description="ARTIQ client")
|
||||||
|
parser.add_argument(
|
||||||
|
"-o", "--run-once", default=[], nargs=3,
|
||||||
|
action="append",
|
||||||
|
help="run experiment once. arguments: <path> <name> <timeout>")
|
||||||
|
parser.add_argument(
|
||||||
|
"-q", "--quit-master", default=False,
|
||||||
|
action="store_true",
|
||||||
|
help="causes the master to quit")
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def _send_obj(sock, obj):
|
||||||
|
line = json.dumps(obj) + "\n"
|
||||||
|
sock.sendall(line.encode())
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
args = _get_args()
|
||||||
|
|
||||||
|
with socket.create_connection(("::1", 8888)) as sock:
|
||||||
|
for path, name, timeout in args.run_once:
|
||||||
|
obj = {
|
||||||
|
"action": "run_once",
|
||||||
|
"run_params": {
|
||||||
|
"path": path,
|
||||||
|
"name": name
|
||||||
|
},
|
||||||
|
"timeout": timeout
|
||||||
|
}
|
||||||
|
_send_obj(sock, obj)
|
||||||
|
if args.quit_master:
|
||||||
|
obj = {"action": "quit"}
|
||||||
|
_send_obj(sock, obj)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
|
@ -0,0 +1,54 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
|
||||||
|
from artiq.management.scheduler import Scheduler
|
||||||
|
|
||||||
|
|
||||||
|
class Server:
|
||||||
|
def __init__(self, loop, scheduler):
|
||||||
|
self.scheduler = scheduler
|
||||||
|
self.terminate_notify = asyncio.Semaphore(0)
|
||||||
|
loop.run_until_complete(
|
||||||
|
asyncio.start_server(self.handle_connection, "::1", 8888))
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def handle_connection(self, reader, writer):
|
||||||
|
while True:
|
||||||
|
line = yield from reader.readline()
|
||||||
|
if not line:
|
||||||
|
writer.close()
|
||||||
|
return
|
||||||
|
obj = json.loads(line.decode())
|
||||||
|
action = obj["action"]
|
||||||
|
if action == "run_once":
|
||||||
|
yield from self.scheduler.add_run_once(obj["run_params"],
|
||||||
|
float(obj["timeout"]))
|
||||||
|
elif action == "quit":
|
||||||
|
self.terminate_notify.release()
|
||||||
|
else:
|
||||||
|
print("warning: unknown action " + action)
|
||||||
|
|
||||||
|
@asyncio.coroutine
|
||||||
|
def task(self, scheduler_task):
|
||||||
|
yield from self.terminate_notify.acquire()
|
||||||
|
scheduler_task.cancel()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
try:
|
||||||
|
with Scheduler(loop) as scheduler:
|
||||||
|
server = Server(loop, scheduler)
|
||||||
|
scheduler_task = asyncio.Task(scheduler.task())
|
||||||
|
server_task = asyncio.Task(server.task(scheduler_task))
|
||||||
|
loop.run_until_complete(asyncio.wait([
|
||||||
|
scheduler_task,
|
||||||
|
server_task
|
||||||
|
]))
|
||||||
|
finally:
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
Loading…
Reference in New Issue