def create_scheduled_task(instance, start_time):
task = ScheduledTask(assignedTask=AssignedTask(
taskId="task_%s" % instance,
slaveId="random_machine_id",
slaveHost="junk.nothing",
task=TaskConfig(job=JobKey(role="nobody",
environment="prod",
name='flibber'),
isService=False,
resources=frozenset([
Resource(numCpus=2),
Resource(ramMb=2048),
Resource(diskMb=4096)
]),
priority=7,
maxTaskFailures=3,
production=False),
assignedPorts={"http": 1001},
instanceId=instance),
status=2,
failureCount=instance + 4,
taskEvents=create_task_events(start_time),
ancestorId="random_task_ancestor%s" %
instance)
return task
def test_simple_config():
job = convert_pystachio_to_thrift(HELLO_WORLD, ports=frozenset(['health']))
expected_key = JobKey(
role=HELLO_WORLD.role().get(),
environment=HELLO_WORLD.environment().get(),
name=HELLO_WORLD.name().get())
assert job.instanceCount == 1
tti = job.taskConfig
assert job.key == expected_key
assert job.owner == Identity(user=getpass.getuser())
assert job.cronSchedule is None
assert tti.job == expected_key
assert tti.isService is False
assert tti.numCpus == 0.1
assert tti.ramMb == 64
assert tti.diskMb == 64
assert tti.requestedPorts == frozenset(['health'])
assert tti.production is False
assert tti.priority == 0
assert tti.maxTaskFailures == 1
assert tti.constraints == set()
assert tti.metadata == set()
assert tti.tier is None
assert Resource(numCpus=0.1) in list(tti.resources)
assert Resource(ramMb=64) in list(tti.resources)
assert Resource(diskMb=64) in list(tti.resources)
assert Resource(namedPort='health') in list(tti.resources)
def test_set_quota(self):
"""Test setting quota."""
api, mock_proxy = self.mock_api()
api.set_quota("role", 1.0, 32, 64)
actual = list(mock_proxy.setQuota.mock_calls[0][1][1].resources)
assert Resource(numCpus=1.0) in actual
assert Resource(ramMb=32) in actual
assert Resource(diskMb=64) in actual
def test_resource_details(self):
details = ResourceManager.resource_details(
[Resource(ramMb=2),
Resource(numCpus=1.0),
Resource(numGpus=1.0)])
assert len(details) == 3
assert details[1] == ResourceDetails(ResourceType.RAM_MB, 2)
assert details[0] == ResourceDetails(ResourceType.CPUS, 1.0)
assert details[2] == ResourceDetails(ResourceType.GPUS, 1)
def test_from_resource(self):
assert ResourceType.from_resource(
Resource(numCpus=1.0)) is ResourceType.CPUS
assert ResourceType.from_resource(
Resource(ramMb=1)) is ResourceType.RAM_MB
assert ResourceType.from_resource(
Resource(diskMb=0)) is ResourceType.DISK_MB
assert ResourceType.from_resource(
Resource(namedPort='http')) is ResourceType.PORTS
def set_quota(self, role, cpu, ram, disk):
log.info("Setting quota for user:%s cpu:%f ram:%d disk: %d"
% (role, cpu, ram, disk))
return self._scheduler_proxy.setQuota(
role,
ResourceAggregate(cpu, ram, disk, frozenset([
Resource(numCpus=cpu),
Resource(ramMb=ram),
Resource(diskMb=disk)])))
def test_config_with_ports():
hwc = HELLO_WORLD(task=HELLO_WORLD.task()(processes=[
Process(
name='hello_world',
cmdline='echo {{thermos.ports[http]}} {{thermos.ports[admin]}}')
]))
config = AuroraConfig(hwc)
job = config.job()
assert Resource(namedPort='http') in list(job.taskConfig.resources)
assert Resource(namedPort='admin') in list(job.taskConfig.resources)
def create_task_config(cls, name):
return TaskConfig(
maxTaskFailures=1,
executorConfig=ExecutorConfig(data='{"fake": "data"}'),
metadata=[],
job=JobKey(role=cls.TEST_ROLE, environment=cls.TEST_ENV,
name=name),
resources=frozenset(
[Resource(numCpus=2),
Resource(ramMb=2),
Resource(diskMb=2)]))
def _backfill_resources(cls, r_object):
resources = list(r_object.resources) if r_object.resources else None
if resources is None:
resources = [
Resource(numCpus=r_object.numCpus),
Resource(ramMb=r_object.ramMb),
Resource(diskMb=r_object.diskMb)
]
if hasattr(r_object, 'requestedPorts'):
resources += [Resource(namedPort=p) for p in r_object.requestedPorts or []]
return resources
def test_quota(self):
quota = ResourceAggregate(resources={
Resource(numCpus=1.0),
Resource(ramMb=2),
Resource(diskMb=3)
})
assert ResourceManager.resource_details_from_quota(quota) == [
ResourceDetails(ResourceType.CPUS, 1.0),
ResourceDetails(ResourceType.RAM_MB, 2),
ResourceDetails(ResourceType.DISK_MB, 3)
]
def setup_mock_quota_call_no_consumption(cls, mock_context):
api = mock_context.get_api('west')
response = cls.create_simple_success_response()
response.result = Result(
getQuotaResult=GetQuotaResult(quota=ResourceAggregate(
resources=frozenset([
Resource(numCpus=5),
Resource(ramMb=20480),
Resource(diskMb=40960)
])),
prodSharedConsumption=None,
prodDedicatedConsumption=None,
nonProdSharedConsumption=None,
nonProdDedicatedConsumption=None))
api.get_quota.return_value = response
def test_increase_quota(self):
"""Tests successful execution of the increase_quota command."""
mock_options = self.setup_mock_options()
with contextlib.nested(
patch('twitter.common.app.get_options',
return_value=mock_options),
patch('apache.aurora.admin.admin.make_admin_client',
return_value=create_autospec(spec=AuroraClientAPI)),
patch('apache.aurora.admin.admin.CLUSTERS',
new=self.TEST_CLUSTERS)) as (_, mock_make_admin_client,
_):
api = mock_make_admin_client.return_value
role = 'test_role'
api.get_quota.return_value = self.create_response(
ResourceAggregate(resources=frozenset([
Resource(numCpus=20.0),
Resource(ramMb=4000),
Resource(diskMb=6000)
])),
ResourceAggregate(resources=frozenset([
Resource(numCpus=15.0),
Resource(ramMb=2000),
Resource(diskMb=3000)
])),
ResourceAggregate(resources=frozenset([
Resource(numCpus=6.0),
Resource(ramMb=200),
Resource(diskMb=600)
])),
)
api.set_quota.return_value = self.create_simple_success_response()
increase_quota([self.TEST_CLUSTER, role, '4.0', '1MB', '1MB'])
api.set_quota.assert_called_with(role, 24.0, 4001, 6001)
assert isinstance(api.set_quota.call_args[0][1], float)
assert isinstance(api.set_quota.call_args[0][2], int)
assert isinstance(api.set_quota.call_args[0][3], int)
def test_static_port_aliasing():
announce = Announcer(primary_port='thrift',
portmap={
'thrift': 8081,
'health': 8300,
'aurora': 'health'
})
config = make_config(announce)
assert config.ports() == set()
for resource in list(config.job().taskConfig.resources):
assert resource.namedPort is None
config = make_config(announce, 'thrift')
assert config.ports() == set()
for resource in list(config.job().taskConfig.resources):
assert resource.namedPort is None
config = make_config(announce, 'thrift', 'health')
assert config.ports() == set()
for resource in list(config.job().taskConfig.resources):
assert resource.namedPort is None
config = make_config(announce, 'derp')
assert config.ports() == set(['derp'])
assert Resource(namedPort='derp') in list(
config.job().taskConfig.resources)
def convert(job, metadata=frozenset(), ports=frozenset()):
"""Convert a Pystachio MesosJob to an Aurora Thrift JobConfiguration."""
owner = Identity(user=getpass.getuser())
key = JobKey(
role=assert_valid_field('role', fully_interpolated(job.role())),
environment=assert_valid_field('environment',
fully_interpolated(job.environment())),
name=assert_valid_field('name', fully_interpolated(job.name())))
task_raw = job.task()
MB = 1024 * 1024
task = TaskConfig()
def not_empty_or(item, default):
return default if item is Empty else fully_interpolated(item)
# job components
task.production = fully_interpolated(job.production(), bool)
task.isService = select_service_bit(job)
task.maxTaskFailures = fully_interpolated(job.max_task_failures())
task.priority = fully_interpolated(job.priority())
task.contactEmail = not_empty_or(job.contact(), None)
task.tier = not_empty_or(job.tier(), None)
if job.has_partition_policy():
task.partitionPolicy = PartitionPolicy(
fully_interpolated(job.partition_policy().reschedule()),
fully_interpolated(job.partition_policy().delay_secs()))
# Add metadata to a task, to display in the scheduler UI.
metadata_set = frozenset()
if job.has_metadata():
customized_metadata = job.metadata()
metadata_set |= frozenset(
(str(fully_interpolated(key_value_metadata.key())),
str(fully_interpolated(key_value_metadata.value())))
for key_value_metadata in customized_metadata)
metadata_set |= frozenset(
(str(key), str(value)) for key, value in metadata)
task.metadata = frozenset(
Metadata(key=key, value=value) for key, value in metadata_set)
# task components
if not task_raw.has_resources():
raise InvalidConfig('Task must specify resources!')
if (fully_interpolated(task_raw.resources().ram()) == 0
or fully_interpolated(task_raw.resources().disk()) == 0):
raise InvalidConfig(
'Must specify ram and disk resources, got ram:%r disk:%r' %
(fully_interpolated(task_raw.resources().ram()),
fully_interpolated(task_raw.resources().disk())))
numCpus = fully_interpolated(task_raw.resources().cpu())
ramMb = fully_interpolated(task_raw.resources().ram()) / MB
diskMb = fully_interpolated(task_raw.resources().disk()) / MB
if numCpus <= 0 or ramMb <= 0 or diskMb <= 0:
raise InvalidConfig(
'Task has invalid resources. cpu/ramMb/diskMb must all be positive: '
'cpu:%r ramMb:%r diskMb:%r' % (numCpus, ramMb, diskMb))
numGpus = fully_interpolated(task_raw.resources().gpu())
task.resources = frozenset([
Resource(numCpus=numCpus),
Resource(ramMb=ramMb),
Resource(diskMb=diskMb)
] + [Resource(namedPort=p)
for p in ports] + ([Resource(numGpus=numGpus)] if numGpus else []))
task.job = key
task.owner = owner
task.taskLinks = {} # See AURORA-739
task.constraints = constraints_to_thrift(
not_empty_or(job.constraints(), {}))
task.container = create_container_config(job.container())
underlying, refs = job.interpolate()
# need to fake an instance id for the sake of schema checking
underlying_checked = underlying.bind(mesos={
'instance': 31337,
'hostname': ''
})
try:
ThermosTaskValidator.assert_valid_task(underlying_checked.task())
except ThermosTaskValidator.InvalidTaskError as e:
raise InvalidConfig('Task is invalid: %s' % e)
if not underlying_checked.check().ok():
raise InvalidConfig('Job not fully specified: %s' %
underlying.check().message())
unbound = []
for ref in refs:
if ref in (THERMOS_TASK_ID_REF, MESOS_INSTANCE_REF,
MESOS_HOSTNAME_REF) or (Ref.subscope(
THERMOS_PORT_SCOPE_REF, ref)):
continue
unbound.append(ref)
if unbound:
raise InvalidConfig('Config contains unbound variables: %s' %
' '.join(map(str, unbound)))
# set the executor that will be used by the Mesos task. Thermos is the default
executor = job.executor_config()
if fully_interpolated(executor.name()) == AURORA_EXECUTOR_NAME:
task.executorConfig = ExecutorConfig(
name=AURORA_EXECUTOR_NAME,
data=filter_aliased_fields(underlying).json_dumps())
else:
task.executorConfig = ExecutorConfig(
name=fully_interpolated(executor.name()),
data=fully_interpolated(executor.data()))
return JobConfiguration(
key=key,
owner=owner,
cronSchedule=not_empty_or(job.cron_schedule(), None),
cronCollisionPolicy=select_cron_policy(job.cron_collision_policy()),
taskConfig=task,
instanceCount=fully_interpolated(job.instances()))
def test_quantity_of(self):
quantity = ResourceManager.quantity_of(
ResourceManager.resource_details(
[Resource(ramMb=2), Resource(numCpus=1.0)]), ResourceType.CPUS)
assert quantity == 1.0
def test_resource_value(self):
assert ResourceType.CPUS.resource_value(Resource(numCpus=1.0)) == 1.0