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_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 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_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 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={'value': 'test-id'}) proto = encode(task) assert isinstance(proto, mesos_pb2.TaskInfo) assert proto.data == dumped assert proto.task_id.value == 'test-id' 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_submit_get(mocker, docker_python): sched = QueueScheduler() mocker.spy(sched, 'on_update') with Running(sched, name='test-scheduler'): 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)]) result = sched.submit(task) assert result.get(timeout=10) == 11 + i
def test_multiple_submissions(mocker, docker_python): sched = QueueScheduler() mocker.spy(sched, 'on_update') with Running(sched, name='test-scheduler'): results = [None] * 3 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)]) results[i] = sched.submit(task) sched.wait() # block until all tasks finishes assert [r.get() for r in results] == [11, 12, 13]
def test_docker_python_exception(): sched = QueueScheduler() def error(): raise Exception('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'): result = sched.submit(task) with pytest.raises(Exception) as e: result.get() assert e.value.message == 'Dummy exception on executor side!'
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_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_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_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