def test_python_task_execution():
fn, args, kwargs = sum, [range(5)], {}
task = PythonTask(fn=fn, args=args, kwargs=kwargs,
id='test-id')
task = decode(encode(task))
assert task() == 10
def fn(lst1, lst2):
return sum(lst1) - sum(lst2)
args = [range(5), range(3)]
task = PythonTask(fn=fn, args=args, id='test-id')
task = decode(encode(task))
assert task() == 7
def test_failed_status_updates(mocker):
mocker.patch('threading.Thread', FakeThread)
def failing_function(*args):
raise Exception("Booom!")
driver = mocker.Mock()
task = PythonTask(fn=failing_function, args=['arbitrary', 'args'])
executor = OneOffExecutor()
executor.on_launch(driver, task)
calls = driver.update.call_args_list
args, kwargs = calls[0]
status = args[0]
assert isinstance(status, PythonTaskStatus)
assert status.state == 'TASK_RUNNING'
assert status.data is None
args, kwargs = calls[1]
status = args[0]
assert isinstance(status, PythonTaskStatus)
assert status.state == 'TASK_FAILED'
assert isinstance(status.data, tuple)
assert isinstance(status.exception, RemoteException)
assert status.message == 'Booom!'
def test_python_task_contains_status():
fn, args, kwargs = sum, [range(5)], {}
task = PythonTask(fn=fn, args=args, kwargs=kwargs,
id='test-id',
envs={'TEST': 'value'},
uris=['test_dependency'])
assert isinstance(task.status, PythonTaskStatus)
assert task.status.state == 'TASK_STAGING'
new_status = PythonTaskStatus(task_id=task.id, state='TASK_RUNNING')
task.update(new_status)
assert isinstance(task.status, PythonTaskStatus)
assert task.status.state == 'TASK_RUNNING'
def python_task():
task = PythonTask(id=TaskID(value='test-task-id'),
fn=sum,
args=[range(5)],
resources=[Cpus(0.1), Mem(128),
Disk(0)])
return task
def docker_python():
task = PythonTask(id=TaskID(value='test-python-task-id'),
fn=sum,
args=[range(5)],
name='test-python-task-name',
resources=[Cpus(0.1), Mem(64),
Disk(0)])
return task
def test_python_task_encode():
fn, args, kwargs = sum, [range(5)], {}
data = (fn, args, kwargs)
dumped = cloudpickle.dumps(data)
task = PythonTask(fn=fn, args=args, kwargs=kwargs,
id='test-id',
envs={'TEST': 'value'},
uris=['test_dependency'])
proto = encode(task)
assert isinstance(proto, mesos_pb2.TaskInfo)
assert proto.data == dumped
assert proto.task_id.value == 'test-id'
assert proto.executor.command.uris[0].value == 'test_dependency'
assert proto.executor.command.environment.variables[0].name == 'TEST'
assert proto.executor.command.environment.variables[0].value == 'value'
task = PythonTask(id=TaskID(value='test-id'))
task.data = data
proto = encode(task)
assert isinstance(proto, mesos_pb2.TaskInfo)
assert proto.data == dumped
assert proto.task_id.value == 'test-id'
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_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_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_finished_status_updates(mocker):
mocker.patch('threading.Thread', side_effect=FakeThread)
driver = mocker.Mock()
task = PythonTask(fn=sum, args=[range(5)])
executor = OneOffExecutor()
executor.on_launch(driver, task)
calls = driver.update.call_args_list
args, kwargs = calls[0]
status = args[0]
assert isinstance(status, PythonTaskStatus)
assert status.state == 'TASK_RUNNING'
assert status.data is None
args, kwargs = calls[1]
status = args[0]
assert isinstance(status, PythonTaskStatus)
assert status.state == 'TASK_FINISHED'
assert status.data == 10