from __future__ import absolute_import import base64 import os import os.path from enum import Enum from time import sleep, time from aurora.api.constants import AURORA_EXECUTOR_NAME, ACTIVE_STATES from aurora.api.ttypes import ( JobKey, TaskConfig, JobConfiguration, ExecutorConfig, Container, DockerContainer, DockerParameter, Identity, ResponseCode, ScheduleStatus, TaskQuery, JobUpdateRequest, JobUpdateSettings, JobUpdateQuery, JobUpdateStatus, Range, Constraint, TaskConstraint, LimitConstraint, Resource, Metadata, ) from aurora.schema.thermos.schema_base import ( Task as ThermosTask, Process as ThermosProcess, Resources, Logger as ThermosLogger, LoggerMode, RotatePolicy, ) from aurora.schema.aurora.base import ( Announcer, MesosJob as ThermosJob, HealthCheckConfig, HealthCheckerConfig, HttpHealthChecker, DisableLifecycle, ) KB = 1024 MB = 1024 * KB GB = 1024 * MB AURORA_ROLE = "peloton" AURORA_ENVIRONMENT = "production" AURORA_USER = "peloton" MAX_WAIT_TIME_SECONDS = 600 WAIT_INTERVAL_SECONDS = 10 def combine_messages(response): """ Combines the message found in the details of a response. :param response: response to extract messages from. :return: Messages from the details in the response, or an empty string if there were no messages. """ return ", ".join( [d.message or "Unknown error" for d in (response.details or [])] ) class Role(Enum): """ The role of a Peloton app instance """ UNKNOWN = 1 LEADER = 2 FOLLOWER = 3 ALL = 4 class Instance(object): """ Representation of an instance of a Peloton application """ def __init__(self, task, role): assert task.assignedTask self.instance_id = task.assignedTask.instanceId self.host = task.assignedTask.slaveHost self.state = task.status # TODO: query Peloton app endpoint to determine leader/follower role self.role = role class App(object): """ Representation of a Peloton application with multiple instances """ def __init__(self, **kwargs): """ Initializes a Peloton application """ # Default attributes self.enable_debug_logging = False self.enable_secrets = False self.cpu_limit = 4.0 self.mem_limit = 16 * GB self.disk_limit = 16 * GB self.scarce_resource_types = None self.slack_resource_types = None self.enable_revocable_resources = None self.bin_packing = None self.qos_advisor_service_address = None self.stream_only_mode = True self.task_preemption_period = "60s" self.enable_sla_tracking = False self.enable_preemption = False self.respool_path = None self.gpu_respool_path = None self.auth_type = "NOOP" self.auth_config_file = "" self.enable_inplace_update = False self.use_host_pool = False self.http_port = None self.grpc_port = None for k, v in kwargs.iteritems(): setattr(self, k, v) self.client = self.cluster.client self.prod_yaml_location = self.cluster.cfg_file if self.num_instances < 1: raise Exception("App %s has no instances" % self.name) self.job_key = JobKey( role=AURORA_ROLE, environment=AURORA_ENVIRONMENT, name=self.name ) # Generate the new job config for this app self.desired_job_config = self.get_desired_job_config() # Save current job config so that we can rollback to later self.current_job_config = self.get_current_job_config() def get_docker_params(self): """ Returns the docker params for a given Peloton application """ mesos_zk_path = "zk://%s/%s" % ( ",".join(self.cluster.zookeeper), self.cluster.mesos_zk_path, ) peloton_zk_endpoints = "\n".join(self.cluster.zookeeper) # add common variables env_vars = { "ENVIRONMENT": "production", "CONFIG_DIR": "./config", "APP": self.name, # TODO: fix Peloton code to only take self.cluster.mesos_zk_path "MESOS_ZK_PATH": mesos_zk_path, "ENABLE_DEBUG_LOGGING": self.enable_debug_logging, "ELECTION_ZK_SERVERS": peloton_zk_endpoints, "USE_CASSANDRA": self.cluster.cassandra_contact_points is not None, "CASSANDRA_HOSTS": "\n".join( self.cluster.cassandra_contact_points ), "CASSANDRA_STORE": self.cluster.cassandra_keyspace, "CASSANDRA_PORT": self.cluster.cassandra_port, "CLUSTER": self.cluster.name, "DATACENTER": getattr(self.cluster, "datacenter", ""), "MESOS_AGENT_WORK_DIR": self.cluster.mesos_agent_work_dir, "AUTO_MIGRATE": self.cluster.auto_migrate, "ENABLE_SENTRY_LOGGING": self.cluster.enable_sentry_logging, "SECRET_CONFIG_DIR": getattr( self.cluster, "secret_config_dir", "" ), "MESOS_SECRET_FILE": getattr( self.cluster, "mesos_secret_file", "" ), "PELOTON_SECRET_FILE": getattr( self.cluster, "peloton_secret_file", "" ), "ENABLE_SECRETS": self.enable_secrets, "AUTH_TYPE": self.auth_type, "AUTH_CONFIG_FILE": self.auth_config_file, } self.add_app_specific_vars(env_vars) name = self.name if name.startswith("placement_"): name = "placement" # base64 encode the prod config and add it to the PRODUCTION_CONFIG # environment variable inside the container prod_config_path = self.get_app_path().format(name) with open(prod_config_path, "rb") as config_file: env_vars["PRODUCTION_CONFIG"] = base64.b64encode( config_file.read() ) params = [ DockerParameter(name="env", value="%s=%s" % (key, val)) for key, val in env_vars.iteritems() ] volumes = [("/var/log/peloton", "/var/log/peloton", "rw")] # Mount langley secrets if secret dir is specified. if getattr(self.cluster, "secret_config_dir", ""): volumes.append( ( self.cluster.secret_config_dir, self.cluster.secret_config_dir, "ro", ) ) params.extend( DockerParameter(name="volume", value="%s:%s:%s" % (src, dst, mode)) for src, dst, mode in volumes ) return params def add_app_specific_vars(self, env_vars): """ Adds env variables specific to the application """ if self.name == "placement_stateful": env_vars["TASK_TYPE"] = "STATEFUL" env_vars["APP"] = "placement" env_vars["APP_TYPE"] = "placement_stateful" if self.cluster.use_host_pool: env_vars["USE_HOST_POOL"] = self.cluster.use_host_pool if self.name == "placement_stateless": env_vars["TASK_TYPE"] = "STATELESS" env_vars["APP"] = "placement" env_vars["APP_TYPE"] = "placement_stateless" if self.cluster.use_host_pool: env_vars["USE_HOST_POOL"] = self.cluster.use_host_pool if self.name == "placement": if getattr(self, "task_dequeue_limit", ""): env_vars["TASK_DEQUEUE_LIMIT"] = self.task_dequeue_limit if getattr(self, "task_dequeue_period", ""): env_vars["TASK_DEQUEUE_PERIOD"] = self.task_dequeue_period if self.cluster.use_host_pool: env_vars["USE_HOST_POOL"] = self.cluster.use_host_pool if self.name == "jobmgr": env_vars["JOB_TYPE"] = getattr(self, "job_type", "BATCH") env_vars["JOB_RUNTIME_CALCULATION_VIA_CACHE"] = getattr( self, "job_runtime_calculation_via_cache", False ) env_vars["TASK_KILL_RATE_LIMIT"] = getattr( self, "task_kill_rate_limit", 0.0 ) env_vars["TASK_KILL_BURST_LIMIT"] = getattr( self, "task_kill_burst_limit", 0 ) env_vars["TASK_LAUNCH_TIMEOUT"] = getattr( self, "task_launch_timeout", "0" ) env_vars["TASK_START_TIMEOUT"] = getattr( self, "task_start_timeout", "0" ) env_vars["EXECUTOR_SHUTDOWN_RATE_LIMIT"] = getattr( self, "executor_shutdown_rate_limit", 0.0 ) env_vars["EXECUTOR_SHUTDOWN_BURST_LIMIT"] = getattr( self, "executor_shutdown_burst_limit", 0 ) if self.name == "archiver": env_vars["ENABLE_ARCHIVER"] = self.enable_archiver env_vars["STREAM_ONLY_MODE"] = self.stream_only_mode env_vars["POD_EVENTS_CLEANUP"] = self.pod_events_cleanup env_vars["ARCHIVE_AGE"] = self.archive_age env_vars["ARCHIVE_INTERVAL"] = self.archive_interval env_vars["ARCHIVE_STEP_SIZE"] = self.archive_step_size env_vars["KAFKA_TOPIC"] = self.kafka_topic if self.name == "resmgr": env_vars["TASK_PREEMPTION_PERIOD"] = getattr( self, "task_preemption_period", "60s" ) env_vars["ENABLE_SLA_TRACKING"] = getattr( self, "enable_sla_tracking", False ) env_vars["ENABLE_PREEMPTION"] = getattr( self, "enable_preemption", False ) if self.name == "hostmgr": if self.scarce_resource_types: env_vars["SCARCE_RESOURCE_TYPES"] = ",".join( self.scarce_resource_types ) if self.slack_resource_types: env_vars["SLACK_RESOURCE_TYPES"] = ",".join( self.slack_resource_types ) if self.enable_revocable_resources: env_vars[ "ENABLE_REVOCABLE_RESOURCES" ] = self.enable_revocable_resources if self.bin_packing: env_vars["BIN_PACKING"] = self.bin_packing if self.qos_advisor_service_address: env_vars["QOS_ADVISOR_SERVICE_ADDRESS"] \ = self.qos_advisor_service_address if self.cluster.use_host_pool: env_vars["ENABLE_HOST_POOL"] = self.cluster.use_host_pool if self.name == "aurorabridge": if self.respool_path: env_vars["RESPOOL_PATH"] = self.respool_path if self.gpu_respool_path: env_vars["GPU_RESPOOL_PATH"] = self.gpu_respool_path if self.enable_inplace_update: env_vars["ENABLE_INPLACE_UPDATE"] = "true" def get_app_path(self): """ Returns the formatted path for app config """ dirname = os.path.dirname(self.prod_yaml_location) path = os.path.join( dirname, "../..", "config", "{}", "production.yaml" ) return path def get_docker_image(self): """ Returns the docker image path for a Peloton app """ return "%s/%s:%s" % ( self.cluster.docker_registry, self.cluster.docker_repository, self.cluster.version, ) def get_executor_config(self): """ Returns the Thermos executor config for a Peloton app """ host_logdir = "/var/log/peloton/%s" % self.name sandbox_logdir = "$MESOS_DIRECTORY/sandbox/.logs/%s/0" % self.name cmdline = " && ".join( [ "rm -rf %s" % host_logdir, "ln -s %s %s" % (sandbox_logdir, host_logdir), "/bin/entrypoint.sh", ] ) # Override executor command for api server to inject # ports environment variable if self.name == "apiserver" and \ self.http_port is None and \ self.grpc_port is None: cmdline = " && ".join( [ "rm -rf %s" % host_logdir, "ln -s %s %s" % (sandbox_logdir, host_logdir), "exec env HTTP_PORT={{thermos.ports[http]}} " "GRPC_PORT={{thermos.ports[grpc]}} " "/bin/entrypoint.sh", ] ) entrypoint_process = ThermosProcess( name=self.name, cmdline=cmdline, logger=ThermosLogger( mode=LoggerMode("rotate"), rotate=RotatePolicy(log_size=1 * GB, backups=10), ), ) thermos_task = ThermosTask( name=self.name, processes=[entrypoint_process], resources=Resources( cpu=self.cpu_limit, ram=self.mem_limit, disk=self.disk_limit ), ) health_check_config = HealthCheckConfig( health_checker=HealthCheckerConfig( http=HttpHealthChecker( endpoint="/health", expected_response="OK", expected_response_code=200, ) ), initial_interval_secs=15, interval_secs=10, max_consecutive_failures=4, timeout_secs=1, ) announce = Announcer() if self.http_port is not None: announce = Announcer(portmap={"health": self.http_port}) elif self.name == "apiserver": # Use assigned http port for health check if it is api server # and static http port is not configured announce = Announcer(portmap={"health": "{{thermos.ports[http]}}"}) thermos_job = ThermosJob( name=self.name, role=AURORA_ROLE, cluster=self.cluster.name, environment=AURORA_ENVIRONMENT, task=thermos_task, production=False, service=True, health_check_config=health_check_config, announce=announce, ) executor_config = ExecutorConfig( name=AURORA_EXECUTOR_NAME, data=thermos_job.json_dumps() ) return executor_config def get_desired_job_config(self): """ Return the Aurora job configuration defined in Thrift API so that we can create a job via Aurora API. """ # Add docker container container = Container( mesos=None, docker=DockerContainer( image=self.get_docker_image(), parameters=self.get_docker_params(), ), ) host_limit = Constraint( name=self.cluster.constraint, constraint=TaskConstraint(limit=LimitConstraint(limit=1)), ) # Set task metadata if presents in config if hasattr(self, 'metadata'): taskMetaData = set( [Metadata(key=k, value=v) for k, v in self.metadata.items()]) else: taskMetaData = set() # Set task resources resources = set( [ Resource(numCpus=self.cpu_limit), Resource(ramMb=self.mem_limit / MB), Resource(diskMb=self.disk_limit / MB), ] ) if self.name == "apiserver": if self.http_port is None: resources.add(Resource(namedPort="http")) if self.grpc_port is None: resources.add(Resource(namedPort="grpc")) task_config = TaskConfig( job=self.job_key, owner=Identity(user=AURORA_USER), isService=True, numCpus=self.cpu_limit, ramMb=self.mem_limit / MB, diskMb=self.disk_limit / MB, priority=0, maxTaskFailures=0, production=False, tier="preemptible", resources=resources, contactEmail="peloton-oncall-group@uber.com", executorConfig=self.get_executor_config(), container=container, constraints=set([host_limit]), requestedPorts=set(), mesosFetcherUris=set(), taskLinks={}, metadata=taskMetaData, ) job_config = JobConfiguration( key=self.job_key, owner=Identity(user=AURORA_USER), taskConfig=task_config, instanceCount=self.num_instances, ) return job_config def get_current_job_config(self): """ Return the current job config by querying the Aurora API """ resp = self.client.getJobSummary(AURORA_ROLE) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) job_config = None for s in resp.result.jobSummaryResult.summaries: if s.job.key == self.job_key: job_config = s.job break if job_config: instances = self.get_instances() job_config.instanceCount = len(instances.get(Role.ALL, [])) return job_config def get_instances(self): """ Returns all instances grouped by role if exist by querying Aurora API """ # Setup task query for task status of the Aurora job task_query = TaskQuery( role=AURORA_ROLE, environment=AURORA_ENVIRONMENT, jobName=self.name, statuses=ACTIVE_STATES, ) resp = self.client.getTasksWithoutConfigs(task_query) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) instances = {} leader = True for t in resp.result.scheduleStatusResult.tasks: if t.status not in ACTIVE_STATES: # Ignore tasks that are not in active states continue # Temporarily hack to set leader/follower roles # TODO: query Peloton app endpoint to find role role = Role.LEADER if leader else Role.FOLLOWER if leader: leader = False inst = Instance(t, role) instances.setdefault(inst.role, []).append(inst) instances.setdefault(Role.ALL, []).append(inst) return instances def wait_for_running(self, role): """ Wait for the app instances of a given role running """ num_instances = { Role.LEADER: 1, Role.FOLLOWER: self.num_instances - 1, Role.ALL: self.num_instances, }[role] start_time = time() while time() < start_time + MAX_WAIT_TIME_SECONDS: instances = self.get_instances().get(role, []) all_running = True for i in instances: if i.state != ScheduleStatus.RUNNING: all_running = False break print( "Wait for %s %s instances running: %d / %d" % (self.name, role.name, all_running, len(instances)) ) if all_running and len(instances) == num_instances: return True sleep(WAIT_INTERVAL_SECONDS) return False def update_instances(self, instances, job_config): """ Update the task config of the given app instances """ instance_ids = [i.instance_id for i in instances] req = JobUpdateRequest( taskConfig=job_config.taskConfig, instanceCount=self.num_instances, settings=JobUpdateSettings( updateGroupSize=1, maxPerInstanceFailures=3 ), ) if instance_ids: req.settings.updateOnlyTheseInstances = set( Range(i, i) for i in instance_ids ) resp = self.client.startJobUpdate( req, "Update %s instances for %s" % (len(instances), self.name) ) if resp.responseCode == ResponseCode.INVALID_REQUEST: if resp.result is None: raise Exception(combine_messages(resp)) update_key = resp.result.startJobUpdateResult.key update_summary = resp.result.startJobUpdateResult.updateSummary status = update_summary.state.status if status == JobUpdateStatus.ROLLING_FORWARD: # Abort the current update print( "Aborting the update for %s (id=%s)" % (self.name, update_key.id) ) self.client.abortJobUpdate( update_key, "Abort by a new deploy session" ) self.wait_for_update_done(update_key) # Restart the job update resp = self.client.startJobUpdate( req, "Update %s instances for %s" % (len(instances), self.name), ) else: raise Exception( "Invalid Request for job update (status=%s)" % (status, JobUpdateStatus._VALUES_TO_NAMES[status]) ) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) if resp.result is None: # No change for the job update print(resp.details[0].message) return True update_key = resp.result.startJobUpdateResult.key return self.wait_for_update_done(update_key, instance_ids) def wait_for_update_done(self, update_key, instance_ids=[]): """ Wait for the job update to finish """ query = JobUpdateQuery( role=AURORA_ROLE, key=update_key, jobKey=self.job_key ) start_time = time() while time() < start_time + MAX_WAIT_TIME_SECONDS: resp = self.client.getJobUpdateSummaries(query) if resp.responseCode != ResponseCode.OK: print(combine_messages(resp)) sleep(WAIT_INTERVAL_SECONDS) continue result = resp.result.getJobUpdateSummariesResult if len(result.updateSummaries) != 1: raise Exception( "Got multiple update summaries: %s" % str(result.updateSummaries) ) if result.updateSummaries[0].key != update_key: raise Exception( "Mismatch update key, expect %s, received %s" % (update_key, result.updateSummaries[0].key) ) status = result.updateSummaries[0].state.status print( "Updating %s instances %s (status=%s)" % ( self.name, instance_ids, JobUpdateStatus._VALUES_TO_NAMES[status], ) ) if status == JobUpdateStatus.ROLLED_FORWARD: return True elif status in [ JobUpdateStatus.ROLLED_BACK, JobUpdateStatus.ABORTED, JobUpdateStatus.ERROR, JobUpdateStatus.FAILED, ]: return False else: # Wait for 5 seconds sleep(WAIT_INTERVAL_SECONDS) return False def update_or_create_job(self, callback): """ Update the current job for a Peloton app. Create a new job if the job does not exist yet. """ # TODO: find the leader/follower role of each app instance if self.current_job_config is None: # Create the new Job in Aurora and check the response code print( "Creating new job for %s with %s instances" % (self.name, self.num_instances) ) resp = self.client.createJob(self.desired_job_config) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) # Wait for all instances are running retval = self.wait_for_running(Role.ALL) if retval: callback(self) return retval # Get current leader and follower instances cur_instances = self.get_instances() # First updade all existing instances, followers first then leader for role in [Role.FOLLOWER, Role.LEADER]: instances = cur_instances.get(role, []) if role == Role.LEADER and len(instances) > 1: raise Exception( "Found %d leaders for %s" % (len(instances), self.name) ) if len(instances) == 0: print("No %s %s instances to update" % (self.name, role.name)) continue print( "Start updating %d %s %s instances" % (len(instances), self.name, role.name) ) retval = self.update_instances(instances, self.desired_job_config) print( "Finish updating %d %s %s instances -- %s" % ( len(instances), self.name, role.name, "SUCCEED" if retval else "FAILED", ) ) if not retval or not callback(self): # Rollback the update by the caller return False # Then add any missing instances if needed cur_total = len(cur_instances.get(role.ALL, [])) new_total = self.num_instances > cur_total if new_total > 0: print("Start adding %d new %s instances" % (new_total, self.name)) retval = self.update_instances([], self.desired_job_config) print( "Finish adding %d new %s instances -- %s" % (new_total, self.name, "SUCCEED" if retval else "FAILED") ) if not retval or not callback(self): # Rollback the update by the caller return False return True def rollback_job(self): """ Rollback the job config of a Peloton app in case of failures """ if self.current_job_config is None: # Nothing to do if the job doesn't exist before return self.update_instances([], self.current_job_config) class CronApp(App): """ Representation of a CRON job """ def __init__(self, **kwargs): """ Initializes Cron job """ # Default attributes self.enable_debug_logging = False self.enable_secrets = False self.working_dir = None self.cmdline = None self.cron_schedule = "\*/5 * * * *" self.num_instances = 1 super(CronApp, self).__init__(**kwargs) self.cpu_limit = 1.0 self.mem_limit = 1 * GB self.disk_limit = 16 * GB def get_executor_config(self): """ Returns the Thermos executor config for a Peloton app """ if self.working_dir is not None: cmd = " && ".join(["cd " + self.working_dir, self.cmdline]) else: cmd = self.cmdline entrypoint_process = ThermosProcess(name=self.name, cmdline=cmd) thermos_task = ThermosTask( name=self.name, processes=[entrypoint_process], resources=Resources( cpu=self.cpu_limit, ram=self.mem_limit, disk=self.disk_limit ), ) thermos_job = ThermosJob( name=self.name, role=AURORA_ROLE, cluster=self.cluster.name, environment=AURORA_ENVIRONMENT, task=thermos_task, production=False, lifecycle=DisableLifecycle, cron_schedule=self.cron_schedule[1:], ) executor_config = ExecutorConfig( name=AURORA_EXECUTOR_NAME, data=thermos_job.json_dumps() ) return executor_config def get_desired_job_config(self): """ Return the Aurora job configuration defined in Thrift API so that we can create a job via Aurora API. """ # Add docker container container = Container( mesos=None, docker=DockerContainer( image=self.get_docker_image(), parameters=self.get_docker_params(), ), ) host_limit = Constraint( name=self.cluster.constraint, constraint=TaskConstraint(limit=LimitConstraint(limit=1)), ) task_config = TaskConfig( job=self.job_key, owner=Identity(user=AURORA_USER), isService=False, numCpus=self.cpu_limit, ramMb=self.mem_limit / MB, diskMb=self.disk_limit / MB, priority=0, maxTaskFailures=0, production=False, tier="preemptible", resources=set( [ Resource(numCpus=self.cpu_limit), Resource(ramMb=self.mem_limit / MB), Resource(diskMb=self.disk_limit / MB), ] ), contactEmail="peloton-oncall-group@uber.com", executorConfig=self.get_executor_config(), container=container, constraints=set([host_limit]), requestedPorts=set(), mesosFetcherUris=set(), taskLinks={}, metadata=set(), ) job_config = JobConfiguration( key=self.job_key, owner=Identity(user=AURORA_USER), taskConfig=task_config, instanceCount=self.num_instances, cronSchedule=self.cron_schedule[1:], ) return job_config def get_current_job_config(self): """ Return the current job config by querying the Aurora API """ resp = self.client.getJobSummary(AURORA_ROLE) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) job_config = None for s in resp.result.jobSummaryResult.summaries: if s.job.key == self.job_key: job_config = s.job break return job_config def update_or_create_job(self, callback): """ Update the current cron-job for a Peloton app. Create a new cron-job if the job does not exist yet. """ resp = self.client.scheduleCronJob(self.desired_job_config) if resp.responseCode != ResponseCode.OK: raise Exception(combine_messages(resp)) return True def rollback_job(self): """ Rollback the job config of a cron-job in case of failures """ if self.current_job_config is None: # Nothing to do if the job doesn't exist before return resp = self.client.scheduleCronJob(self.current_job_config) if resp.responseCode != ResponseCode.OK: print("Failed to rollback watchdog configuration (%s)" % (resp))