def create_resources(cpus, mem, disk, ports, role='*'):
"""Return a list of 'Resource' protobuf for the provided resources."""
cpus_resources = mesos_pb2.Resource()
cpus_resources.name = 'cpus'
cpus_resources.type = mesos_pb2.Value.SCALAR
cpus_resources.role = role
cpus_resources.scalar.value = cpus
mem_resources = mesos_pb2.Resource()
mem_resources.name = 'mem'
mem_resources.type = mesos_pb2.Value.SCALAR
mem_resources.role = role
mem_resources.scalar.value = mem.as_(Data.MB)
disk_resources = mesos_pb2.Resource()
disk_resources.name = 'disk'
disk_resources.type = mesos_pb2.Value.SCALAR
disk_resources.role = role
disk_resources.scalar.value = disk.as_(Data.MB)
ports_resources = mesos_pb2.Resource()
ports_resources.name = 'ports'
ports_resources.type = mesos_pb2.Value.RANGES
ports_resources.role = role
for port in ports:
port_range = ports_resources.ranges.range.add()
port_range.begin = port
port_range.end = port
return [cpus_resources, mem_resources, disk_resources, ports_resources]
class ScalarResource(Resource):
# supports comparison and basic arithmetics with scalars
proto = mesos_pb2.Resource(type=mesos_pb2.Value.SCALAR)
def __init__(self, value=None, **kwargs):
super(Resource, self).__init__(**kwargs)
if value is not None:
self.scalar = Scalar(value=value)
def __cmp__(self, other):
first, second = float(self), float(other)
if first < second:
return -1
elif first > second:
return 1
else:
return 0
def __repr__(self):
return "<{}: {}>".format(self.__class__.__name__, self.scalar.value)
def __float__(self):
return float(self.scalar.value)
@classmethod
def _op(cls, op, first, second):
value = op(float(first), float(second))
return cls(value=value)
def __add__(self, other):
return self._op(operator.add, self, other)
def __radd__(self, other):
return self._op(operator.add, other, self)
def __sub__(self, other):
return self._op(operator.sub, self, other)
def __rsub__(self, other):
return self._op(operator.sub, other, self)
def __mul__(self, other):
return self._op(operator.mul, self, other)
def __rmul__(self, other):
return self._op(operator.mul, other, self)
def __truediv__(self, other):
return self._op(operator.truediv, self, other)
def __rtruediv__(self, other):
return self._op(operator.truediv, other, self)
def __iadd__(self, other):
self.scalar.value = float(self._op(operator.add, self, other))
return self
def __isub__(self, other):
self.scalar.value = float(self._op(operator.sub, self, other))
return self
def test_mesos_executor(self, driver):
# create task queue, empty result queue, task_cpu and task_memory
tasks_queue = Queue()
fake_af_task1 = {"key1", "airflow run tutorial"}
fake_af_task2 = {"key2", "airflow run tutorial2"}
tasks_queue.put(fake_af_task1)
tasks_queue.put(fake_af_task2)
results_queue = Queue()
task_cpu = 2
task_memory = 4
scheduler = AirflowMesosScheduler(tasks_queue,
results_queue,
task_cpu,
task_memory)
# Create Offers
resources = []
fake_cpu_resource = mesos_pb2.Resource(
name='cpus',
type=mesos_pb2.Value.SCALAR,
)
fake_cpu_resource.scalar.value = task_cpu
fake_mem_resource = mesos_pb2.Resource(
name='mem',
type=mesos_pb2.Value.SCALAR,
)
fake_mem_resource.scalar.value = task_memory
resources.append(fake_cpu_resource)
resources.append(fake_mem_resource)
fake_offer = mesos_pb2.Offer(
id=self.offer_id,
framework_id=self.framework_id,
slave_id=self.slave_id,
hostname='fake-host',
resources=resources
)
scheduler.resourceOffers(driver, [fake_offer])
# assertions
self.assertTrue(driver.launchTasks.called)
class Disk(ScalarResource):
proto = mesos_pb2.Resource(name='disk', type=mesos_pb2.Value.SCALAR)
class Mem(ScalarResource):
proto = mesos_pb2.Resource(name='mem', type=mesos_pb2.Value.SCALAR)
class Cpus(ScalarResource):
proto = mesos_pb2.Resource(name='cpus', type=mesos_pb2.Value.SCALAR)