mirror of https://github.com/m-labs/artiq
429 lines
14 KiB
Python
429 lines
14 KiB
Python
import unittest
|
|
import logging
|
|
import asyncio
|
|
import sys
|
|
from time import time, sleep
|
|
|
|
from artiq.experiment import *
|
|
from artiq.master.scheduler import Scheduler
|
|
|
|
|
|
class EmptyExperiment(EnvExperiment):
|
|
def build(self):
|
|
pass
|
|
|
|
def run(self):
|
|
pass
|
|
|
|
|
|
class BackgroundExperiment(EnvExperiment):
|
|
def build(self):
|
|
self.setattr_device("scheduler")
|
|
|
|
def run(self):
|
|
try:
|
|
while True:
|
|
self.scheduler.pause()
|
|
sleep(0.2)
|
|
except TerminationRequested:
|
|
self.set_dataset("termination_ok", True,
|
|
broadcast=True, archive=False)
|
|
|
|
|
|
class CheckPauseBackgroundExperiment(EnvExperiment):
|
|
def build(self):
|
|
self.setattr_device("scheduler")
|
|
|
|
def run(self):
|
|
while True:
|
|
while not self.scheduler.check_pause():
|
|
sleep(0.2)
|
|
self.scheduler.pause()
|
|
|
|
|
|
def _get_expid(name):
|
|
return {
|
|
"log_level": logging.WARNING,
|
|
"file": sys.modules[__name__].__file__,
|
|
"class_name": name,
|
|
"arguments": dict()
|
|
}
|
|
|
|
|
|
def _get_basic_steps(rid, expid, priority=0, flush=False):
|
|
return [
|
|
{"action": "setitem", "key": rid, "value":
|
|
{"pipeline": "main", "status": "pending", "priority": priority,
|
|
"expid": expid, "due_date": None, "flush": flush,
|
|
"repo_msg": None},
|
|
"path": []},
|
|
{"action": "setitem", "key": "status", "value": "preparing",
|
|
"path": [rid]},
|
|
{"action": "setitem", "key": "status", "value": "prepare_done",
|
|
"path": [rid]},
|
|
{"action": "setitem", "key": "status", "value": "running",
|
|
"path": [rid]},
|
|
{"action": "setitem", "key": "status", "value": "run_done",
|
|
"path": [rid]},
|
|
{"action": "setitem", "key": "status", "value": "analyzing",
|
|
"path": [rid]},
|
|
{"action": "setitem", "key": "status", "value": "deleting",
|
|
"path": [rid]},
|
|
{"action": "delitem", "key": rid, "path": []}
|
|
]
|
|
|
|
|
|
class _RIDCounter:
|
|
def __init__(self, next_rid):
|
|
self._next_rid = next_rid
|
|
|
|
def get(self):
|
|
rid = self._next_rid
|
|
self._next_rid += 1
|
|
return rid
|
|
|
|
|
|
class SchedulerCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(self.loop)
|
|
|
|
def test_steps(self):
|
|
loop = self.loop
|
|
scheduler = Scheduler(_RIDCounter(0), dict(), None, None)
|
|
expid = _get_expid("EmptyExperiment")
|
|
|
|
expect = _get_basic_steps(1, expid)
|
|
done = asyncio.Event()
|
|
expect_idx = 0
|
|
def notify(mod):
|
|
nonlocal expect_idx
|
|
self.assertEqual(mod, expect[expect_idx])
|
|
expect_idx += 1
|
|
if expect_idx >= len(expect):
|
|
done.set()
|
|
scheduler.notifier.publish = notify
|
|
|
|
scheduler.start(loop=loop)
|
|
|
|
# Verify that a timed experiment far in the future does not
|
|
# get run, even if it has high priority.
|
|
late = time() + 100000
|
|
expect.insert(0,
|
|
{"action": "setitem", "key": 0, "value":
|
|
{"pipeline": "main", "status": "pending", "priority": 99,
|
|
"expid": expid, "due_date": late, "flush": False,
|
|
"repo_msg": None},
|
|
"path": []})
|
|
scheduler.submit("main", expid, 99, late, False)
|
|
|
|
# This one (RID 1) gets run instead.
|
|
scheduler.submit("main", expid, 0, None, False)
|
|
|
|
loop.run_until_complete(done.wait())
|
|
scheduler.notifier.publish = None
|
|
loop.run_until_complete(scheduler.stop())
|
|
|
|
def test_pending_priority(self):
|
|
"""Check due dates take precedence over priorities when waiting to
|
|
prepare."""
|
|
loop = self.loop
|
|
handlers = {}
|
|
scheduler = Scheduler(_RIDCounter(0), handlers, None, None)
|
|
handlers["scheduler_check_pause"] = scheduler.check_pause
|
|
|
|
expid_empty = _get_expid("EmptyExperiment")
|
|
|
|
expid_bg = _get_expid("CheckPauseBackgroundExperiment")
|
|
# Suppress the SystemExit backtrace when worker process is killed.
|
|
expid_bg["log_level"] = logging.CRITICAL
|
|
|
|
high_priority = 3
|
|
middle_priority = 2
|
|
low_priority = 1
|
|
late = time() + 100000
|
|
early = time() + 1
|
|
|
|
expect = [
|
|
{
|
|
"path": [],
|
|
"action": "setitem",
|
|
"value": {
|
|
"repo_msg": None,
|
|
"priority": low_priority,
|
|
"pipeline": "main",
|
|
"due_date": None,
|
|
"status": "pending",
|
|
"expid": expid_bg,
|
|
"flush": False
|
|
},
|
|
"key": 0
|
|
},
|
|
{
|
|
"path": [],
|
|
"action": "setitem",
|
|
"value": {
|
|
"repo_msg": None,
|
|
"priority": high_priority,
|
|
"pipeline": "main",
|
|
"due_date": late,
|
|
"status": "pending",
|
|
"expid": expid_empty,
|
|
"flush": False
|
|
},
|
|
"key": 1
|
|
},
|
|
{
|
|
"path": [],
|
|
"action": "setitem",
|
|
"value": {
|
|
"repo_msg": None,
|
|
"priority": middle_priority,
|
|
"pipeline": "main",
|
|
"due_date": early,
|
|
"status": "pending",
|
|
"expid": expid_empty,
|
|
"flush": False
|
|
},
|
|
"key": 2
|
|
},
|
|
{
|
|
"path": [0],
|
|
"action": "setitem",
|
|
"value": "preparing",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [0],
|
|
"action": "setitem",
|
|
"value": "prepare_done",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [0],
|
|
"action": "setitem",
|
|
"value": "running",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "preparing",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "prepare_done",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [0],
|
|
"action": "setitem",
|
|
"value": "paused",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "running",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "run_done",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [0],
|
|
"action": "setitem",
|
|
"value": "running",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "analyzing",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [2],
|
|
"action": "setitem",
|
|
"value": "deleting",
|
|
"key": "status"
|
|
},
|
|
{
|
|
"path": [],
|
|
"action": "delitem",
|
|
"key": 2
|
|
},
|
|
]
|
|
done = asyncio.Event()
|
|
expect_idx = 0
|
|
def notify(mod):
|
|
nonlocal expect_idx
|
|
self.assertEqual(mod, expect[expect_idx])
|
|
expect_idx += 1
|
|
if expect_idx >= len(expect):
|
|
done.set()
|
|
scheduler.notifier.publish = notify
|
|
|
|
scheduler.start(loop=loop)
|
|
|
|
scheduler.submit("main", expid_bg, low_priority)
|
|
scheduler.submit("main", expid_empty, high_priority, late)
|
|
scheduler.submit("main", expid_empty, middle_priority, early)
|
|
|
|
loop.run_until_complete(done.wait())
|
|
scheduler.notifier.publish = None
|
|
loop.run_until_complete(scheduler.stop())
|
|
|
|
def test_pause(self):
|
|
loop = self.loop
|
|
|
|
termination_ok = False
|
|
def check_termination(mod):
|
|
nonlocal termination_ok
|
|
self.assertEqual(
|
|
mod,
|
|
{"action": "setitem", "key": "termination_ok",
|
|
"value": (False, True, {}), "path": []})
|
|
termination_ok = True
|
|
handlers = {
|
|
"update_dataset": check_termination
|
|
}
|
|
scheduler = Scheduler(_RIDCounter(0), handlers, None, None)
|
|
|
|
expid_bg = _get_expid("BackgroundExperiment")
|
|
expid = _get_expid("EmptyExperiment")
|
|
|
|
expect = _get_basic_steps(1, expid)
|
|
background_running = asyncio.Event()
|
|
empty_ready = asyncio.Event()
|
|
empty_completed = asyncio.Event()
|
|
background_completed = asyncio.Event()
|
|
expect_idx = 0
|
|
def notify(mod):
|
|
nonlocal expect_idx
|
|
if mod == {"path": [0],
|
|
"value": "running",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
background_running.set()
|
|
if mod == {"path": [0],
|
|
"value": "deleting",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
background_completed.set()
|
|
if mod == {"path": [1],
|
|
"value": "prepare_done",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
empty_ready.set()
|
|
if mod["path"] == [1] or (mod["path"] == [] and mod["key"] == 1):
|
|
self.assertEqual(mod, expect[expect_idx])
|
|
expect_idx += 1
|
|
if expect_idx >= len(expect):
|
|
empty_completed.set()
|
|
scheduler.notifier.publish = notify
|
|
|
|
scheduler.start(loop=loop)
|
|
scheduler.submit("main", expid_bg, -99, None, False)
|
|
loop.run_until_complete(background_running.wait())
|
|
self.assertFalse(scheduler.check_pause(0))
|
|
scheduler.submit("main", expid, 0, None, False)
|
|
self.assertFalse(scheduler.check_pause(0))
|
|
loop.run_until_complete(empty_ready.wait())
|
|
self.assertTrue(scheduler.check_pause(0))
|
|
loop.run_until_complete(empty_completed.wait())
|
|
self.assertFalse(scheduler.check_pause(0))
|
|
|
|
self.assertFalse(termination_ok)
|
|
scheduler.request_termination(0)
|
|
self.assertTrue(scheduler.check_pause(0))
|
|
loop.run_until_complete(background_completed.wait())
|
|
self.assertTrue(termination_ok)
|
|
|
|
loop.run_until_complete(scheduler.stop())
|
|
|
|
def test_close_with_active_runs(self):
|
|
"""Check scheduler exits with experiments still running"""
|
|
loop = self.loop
|
|
|
|
scheduler = Scheduler(_RIDCounter(0), {}, None, None)
|
|
|
|
expid_bg = _get_expid("BackgroundExperiment")
|
|
# Suppress the SystemExit backtrace when worker process is killed.
|
|
expid_bg["log_level"] = logging.CRITICAL
|
|
expid = _get_expid("EmptyExperiment")
|
|
|
|
background_running = asyncio.Event()
|
|
empty_ready = asyncio.Event()
|
|
background_completed = asyncio.Event()
|
|
def notify(mod):
|
|
if mod == {"path": [0],
|
|
"value": "running",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
background_running.set()
|
|
if mod == {"path": [0],
|
|
"value": "deleting",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
background_completed.set()
|
|
if mod == {"path": [1],
|
|
"value": "prepare_done",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
empty_ready.set()
|
|
scheduler.notifier.publish = notify
|
|
|
|
scheduler.start(loop=loop)
|
|
scheduler.submit("main", expid_bg, -99, None, False)
|
|
loop.run_until_complete(background_running.wait())
|
|
|
|
scheduler.submit("main", expid, 0, None, False)
|
|
loop.run_until_complete(empty_ready.wait())
|
|
|
|
# At this point, (at least) BackgroundExperiment is still running; make
|
|
# sure we can stop the scheduler without hanging.
|
|
loop.run_until_complete(scheduler.stop())
|
|
|
|
def test_flush(self):
|
|
loop = self.loop
|
|
scheduler = Scheduler(_RIDCounter(0), dict(), None, None)
|
|
expid = _get_expid("EmptyExperiment")
|
|
|
|
expect = _get_basic_steps(1, expid, 1, True)
|
|
expect.insert(1, {"key": "status",
|
|
"path": [1],
|
|
"value": "flushing",
|
|
"action": "setitem"})
|
|
first_preparing = asyncio.Event()
|
|
done = asyncio.Event()
|
|
expect_idx = 0
|
|
def notify(mod):
|
|
nonlocal expect_idx
|
|
if mod == {"path": [0],
|
|
"value": "preparing",
|
|
"key": "status",
|
|
"action": "setitem"}:
|
|
first_preparing.set()
|
|
if mod["path"] == [1] or (mod["path"] == [] and mod["key"] == 1):
|
|
self.assertEqual(mod, expect[expect_idx])
|
|
expect_idx += 1
|
|
if expect_idx >= len(expect):
|
|
done.set()
|
|
scheduler.notifier.publish = notify
|
|
|
|
scheduler.start(loop=loop)
|
|
scheduler.submit("main", expid, 0, None, False)
|
|
loop.run_until_complete(first_preparing.wait())
|
|
scheduler.submit("main", expid, 1, None, True)
|
|
loop.run_until_complete(done.wait())
|
|
loop.run_until_complete(scheduler.stop())
|
|
|
|
def tearDown(self):
|
|
self.loop.close()
|