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 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 docker_command():
task = TaskInfo(name='test-docker-task',
id=TaskID(value='test-docker-task-id'),
resources=[Cpus(0.1), Mem(64)],
command=CommandInfo(value='echo 100'),
container=ContainerInfo(
type='DOCKER', docker=DockerInfo(image='lensa/satyr')))
return task
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_non_strict_encode_task_info():
t = TaskInfo(name='test-task',
id=TaskID(value='test-task-id'),
resources=[Cpus(0.1), Mem(16)],
command=CommandInfo(value='echo 100'))
t.result = 'some binary data'
t.status = TaskStatus()
p = encode(t)
assert isinstance(p, mesos_pb2.TaskInfo)
assert p.command.value == 'echo 100'
with pytest.raises(AttributeError):
p.status
def test_encode_task_info():
t = TaskInfo(name='test-task',
id=TaskID(value='test-task-id'),
resources=[Cpus(0.1), Mem(16)],
command=CommandInfo(value='echo 100'))
p = encode(t)
assert isinstance(p, mesos_pb2.TaskInfo)
assert p.command.value == 'echo 100'
assert p.name == 'test-task'
assert p.resources[0].name == 'cpus'
assert p.resources[0].scalar.value == 0.1
assert p.task_id.value == 'test-task-id'
def test_status_in_task_info():
t = TaskInfo(name='test-task',
id=TaskID(value='test-task-id'),
resources=[Cpus(0.1), Mem(16)],
command=CommandInfo(value='echo 100'))
assert isinstance(t.status, TaskStatus)
assert t.status.state == 'TASK_STAGING'
p = encode(t)
assert isinstance(p, mesos_pb2.TaskInfo)
with pytest.raises(AttributeError):
p.status
def test_encode_task_info_resources():
task = TaskInfo(name='test-task',
id=TaskID(value='test-task-id'),
resources=[Cpus(0.1), Mem(16)],
command=CommandInfo(value='testcmd'))
pb = encode(task)
assert pb.name == 'test-task'
assert pb.task_id.value == 'test-task-id'
assert pb.resources[0].name == 'cpus'
assert pb.resources[0].scalar.value == 0.1
assert pb.resources[1].name == 'mem'
assert pb.resources[1].scalar.value == 16
assert pb.command.value == 'testcmd'
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_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_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_python_task_status_exception():
status = PythonTaskStatus(task_id=TaskID(value='e'), state='TASK_FAILED')
status.data = (TypeError('test'), 'traceback')
assert isinstance(status.exception, RemoteException)
assert isinstance(status.exception, TypeError)
def command():
task = TaskInfo(name='test-task',
id=TaskID(value='test-task-id'),
resources=[Cpus(0.1), Mem(64)],
command=CommandInfo(value='echo 100'))
return task