def test_task_callbacks(mocker, python_task, offers):
driver = mocker.Mock()
sched = QueueScheduler(name='test-scheduler')
mocker.spy(python_task, 'on_update')
mocker.spy(python_task, 'on_success')
sched.submit(python_task)
sched.on_offers(driver, offers)
status = PythonTaskStatus(task_id=python_task.id, state='TASK_RUNNING')
sched.on_update(driver, status)
status = PythonTaskStatus(task_id=python_task.id,
state='TASK_FINISHED',
data=python_task())
sched.on_update(driver, status)
update_calls = python_task.on_update.call_args_list
success_calls = python_task.on_success.call_args_list
args, kwargs = update_calls[0]
assert isinstance(args[0], PythonTaskStatus)
assert args[0].state == 'TASK_RUNNING'
args, kwargs = update_calls[1]
assert isinstance(args[0], PythonTaskStatus)
assert args[0].state == 'TASK_FINISHED'
args, kwargs = success_calls[0]
assert isinstance(args[0], PythonTaskStatus)
assert args[0].state == 'TASK_FINISHED'
assert args[0].data == 10
def test_task_result(mocker, python_task, offers):
driver = mocker.Mock()
sched = QueueScheduler(name='test-scheduler')
sched.submit(python_task)
sched.on_offers(driver, offers)
status = PythonTaskStatus(task_id=python_task.id, state='TASK_RUNNING')
sched.on_update(driver, status)
status = PythonTaskStatus(task_id=python_task.id,
state='TASK_FINISHED',
data=python_task())
sched.on_update(driver, status)
assert python_task.status.state == 'TASK_FINISHED'
assert python_task.status.data == 10
def test_sequential_execution(mocker, docker_python):
sched = QueueScheduler()
mocker.spy(sched, 'on_update')
with Running(sched, name='test-scheduler'):
tasks = []
for i in range(3):
task = PythonTask(id=TaskID(value='test-python-task-{}'.format(i)),
fn=sum,
args=[[1, 10, i]],
name='test-python-task-name',
resources=[Cpus(0.1),
Mem(64), Disk(0)])
sched.submit(task)
tasks.append(task)
sched.wait()
assert task.status.data == 11 + i
def test_docker_python(mocker, docker_python):
sched = QueueScheduler()
mocker.spy(sched, 'on_update')
with Running(sched, name='test-scheduler'):
sched.submit(docker_python)
sched.wait() # block until all tasks finishes
calls = sched.on_update.call_args_list
assert len(calls) == 2
args, kwargs = calls[0]
assert args[1].task_id.value == 'test-python-task-id'
assert args[1].state == 'TASK_RUNNING'
args, kwargs = calls[1]
assert args[1].task_id.value == 'test-python-task-id'
assert args[1].state == 'TASK_FINISHED'
def test_parallel_execution(mocker, docker_python):
sched = QueueScheduler()
mocker.spy(sched, 'on_update')
with Running(sched, name='test-scheduler'):
tasks = []
for i in range(3):
task = PythonTask(id=TaskID(value='test-python-task-{}'.format(i)),
fn=sum,
args=[[1, 10, i]],
name='test-python-task-name',
resources=[Cpus(0.1),
Mem(64), Disk(0)])
sched.submit(task)
tasks.append(task)
sched.wait() # block until all tasks finishes
assert [t.status.data for t in tasks] == [11, 12, 13]
def test_docker_python_exception():
sched = QueueScheduler()
def error():
raise TypeError('Dummy exception on executor side!')
task = PythonTask(id=TaskID(value='test-python-task-id'),
fn=error,
name='test-python-task-name',
resources=[Cpus(0.1), Mem(64),
Disk(0)])
with Running(sched, name='test-scheduler'):
sched.submit(task)
sched.wait()
assert task.status.has_failed()
assert isinstance(task.status.exception, RemoteException)
assert isinstance(task.status.exception, TypeError)
def test_launch_decline(mocker, python_task, offers):
driver = mocker.Mock()
sched = QueueScheduler(name='test-scheduler')
sched.submit(python_task)
sched.on_offers(driver, offers)
calls = driver.launch.call_args_list
args, kwargs = calls[0]
assert isinstance(args[0], OfferID)
assert args[0].value == 'first-offer'
assert isinstance(args[1][0], PythonTask)
assert args[1][0].task_id.value == 'test-task-id'
args, kwargs = calls[1]
assert isinstance(args[0], OfferID)
assert args[0].value == 'second-offer'
assert args[1] == [] # declines via launch empty task list
def test_scheduler_retries(mocker):
task = PythonTask(id=TaskID(value='non-existing-docker-image'),
name='test',
fn=lambda: range(int(10e10)),
docker='pina/sen',
resources=[Cpus(0.1), Mem(128),
Disk(0)])
sched = QueueScheduler(name='test-executor-lost', retries=3)
mocker.spy(sched, 'on_update')
with Running(sched, name='test-scheduler'):
sched.submit(task)
sched.wait()
assert sched.on_update.call_count == 3
states = ['TASK_STARTING', 'TASK_STARTING', 'TASK_FAILED']
for ((args, kwargs), state) in zip(sched.on_update.call_args_list, states):
assert args[1].state == state
def test_docker_python_result(mocker, docker_python):
sched = QueueScheduler()
with Running(sched, name='test-scheduler'):
sched.submit(docker_python)
sched.wait() # block until all tasks finishes
assert docker_python.status.data == 10
def test_runner_context_manager():
sched = QueueScheduler(name='test-scheduler')
with Running(sched, name='test-scheduler'):
pass
assert sched