def mock_api(cls): api = AuroraClientAPI(Cluster(name="foo"), 'test-client') mock_proxy = create_autospec(spec=SchedulerProxyApiSpec, spec_set=True, instance=True) api._scheduler_proxy = mock_proxy return api, mock_proxy
def increase_quota(cluster, role, cpu_str, ram_str, disk_str): """usage: increase_quota cluster role cpu ram[unit] disk[unit] Increases the amount of production quota allocated to a user. """ cpu = float(cpu_str) ram = parse_data(ram_str) disk = parse_data(disk_str) options = app.get_options() client = AuroraClientAPI(CLUSTERS[cluster], options.verbosity == "verbose") resp = client.get_quota(role) quota = resp.result.getQuotaResult.quota log.info( "Current quota for %s:\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB" % (role, quota.numCpus, quota.ramMb, quota.diskMb) ) new_cpu = float(cpu + quota.numCpus) new_ram = int((ram + Amount(quota.ramMb, Data.MB)).as_(Data.MB)) new_disk = int((disk + Amount(quota.diskMb, Data.MB)).as_(Data.MB)) log.info( "Attempting to update quota for %s to\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB" % (role, new_cpu, new_ram, new_disk) ) resp = client.set_quota(role, new_cpu, new_ram, new_disk) check_and_log_response(resp)
def increase_quota(cluster, role, cpu_str, ram_str, disk_str): """usage: increase_quota cluster role cpu ram[unit] disk[unit] Increases the amount of production quota allocated to a user. """ cpu = float(cpu_str) ram = parse_data(ram_str) disk = parse_data(disk_str) options = app.get_options() client = AuroraClientAPI(CLUSTERS[cluster], options.verbosity == 'verbose') resp = client.get_quota(role) quota = resp.result.getQuotaResult.quota log.info('Current quota for %s:\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB' % (role, quota.numCpus, quota.ramMb, quota.diskMb)) new_cpu = float(cpu + quota.numCpus) new_ram = int((ram + Amount(quota.ramMb, Data.MB)).as_(Data.MB)) new_disk = int((disk + Amount(quota.diskMb, Data.MB)).as_(Data.MB)) log.info('Attempting to update quota for %s to\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB' % (role, new_cpu, new_ram, new_disk)) resp = client.set_quota(role, new_cpu, new_ram, new_disk) check_and_log_response(resp)
def __init__(self, cluster, role, env, jobs, ssh_user=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._log = log_fn
def sla_probe_hosts(cluster, percentage, duration): """usage: sla_probe_hosts [--filename=filename] [--hosts=hosts] cluster percentage duration Probes individual hosts with respect to their job SLA. Specifically, given a host, outputs all affected jobs with their projected SLAs if the host goes down. In addition, if a job's projected SLA does not clear the specified limits suggests the approximate time when that job reaches its SLA. Output format: HOST JOB PREDICTED_SLA SAFE? PREDICTED_SAFE_IN where: HOST - host being probed. JOB - job that has tasks running on the host being probed. PREDICTED_SLA - predicted effective percentage of up tasks if the host is shut down. SAFE? - PREDICTED_SLA >= percentage PREDICTED_SAFE_IN - expected wait time in seconds for the job to reach requested SLA threshold. """ options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) hosts = parse_hosts(options.filename, options.hosts) vector = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector(hosts) probed_hosts = vector.probe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), hosts) results = [] for host, job_details in sorted(probed_hosts.items()): results.append( "\n".join( [ "%s\t%s\t%.2f\t%s\t%s" % ( host, d.job.to_path(), d.predicted_percentage, d.safe, "n/a" if d.safe_in_secs is None else d.safe_in_secs, ) for d in sorted(job_details) ] ) ) print_results(results)
def __init__(self, cluster, role, env, jobs, ssh_user=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster, user_agent=AURORA_V2_USER_AGENT_NAME) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._log = log_fn
def scheduler_backup_now(cluster): """usage: scheduler_backup_now cluster Immediately initiates a full storage backup. """ options = app.get_options() check_and_log_response(AuroraClientAPI(CLUSTERS[cluster], options.verbosity).perform_backup())
def __init__(self, cluster, role, env, jobs, ssh_user=None): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role
def scheduler_snapshot(cluster): """usage: scheduler_snapshot cluster Request that the scheduler perform a storage snapshot and block until complete. """ options = app.get_options() check_and_log_response(AuroraClientAPI(CLUSTERS[cluster], options.verbosity).snapshot())
def scheduler_unload_recovery(cluster): """usage: scheduler_unload_recovery cluster Unloads a staged recovery. """ options = app.get_options() check_and_log_response(AuroraClientAPI(CLUSTERS[cluster], options.verbosity) .unload_recovery())
def scheduler_stage_recovery(cluster, backup_id): """usage: scheduler_stage_recovery cluster backup_id Stages a backup for recovery. """ options = app.get_options() check_and_log_response( AuroraClientAPI(CLUSTERS[cluster], options.verbosity).stage_recovery(backup_id))
def test_handles_api_auth_error(): context = AuroraCommandContext() mock_scheduler_proxy = mock.create_autospec(spec=SchedulerProxyApiSpec, instance=True) mock_scheduler_proxy.killTasks.side_effect = SchedulerProxy.AuthError() mock_api = AuroraClientAPI(TEST_CLUSTER, "user-agent") mock_api._scheduler_proxy = mock_scheduler_proxy context.apis = {TEST_CLUSTER.name: mock_api} api = context.get_api(TEST_CLUSTER.name, clusters={TEST_CLUSTER.name: TEST_CLUSTER}) with pytest.raises(Context.CommandError) as e: api.kill_job(AuroraJobKey(TEST_CLUSTER.name, "role", "env", "job")) assert e.value.code == EXIT_AUTH_ERROR assert mock_scheduler_proxy.killTasks.call_count == 1
def scheduler_delete_recovery_tasks(cluster, task_ids): """usage: scheduler_delete_recovery_tasks cluster task_ids Deletes a comma-separated list of task IDs from a staged recovery. """ ids = set(task_ids.split(',')) options = app.get_options() check_and_log_response(AuroraClientAPI(CLUSTERS[cluster], options.verbosity) .delete_recovery_tasks(TaskQuery(taskIds=ids)))
def make_admin_client(cluster): if cluster not in CLUSTERS: die('Unknown cluster: %s. Known clusters: %s' % (cluster, ", ".join(CLUSTERS.keys()))) verbose = getattr(app.get_options(), 'verbosity', 'normal') == 'verbose' return AuroraClientAPI(CLUSTERS[cluster], AURORA_ADMIN_USER_AGENT_NAME, verbose=verbose)
def __init__(self, cluster, role, env, jobs, ssh_user=None, ssh_options=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster, user_agent=AURORA_V2_USER_AGENT_NAME) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._ssh_options = ssh_options if ssh_options else [] self._log = log_fn
def get_scheduler(cluster): """usage: get_scheduler CLUSTER Dumps the leading scheduler endpoint URL. """ options = app.get_options() print("Found leading scheduler at: %s" % AuroraClientAPI( CLUSTERS[cluster], options.verbosity).scheduler_proxy.scheduler_client().raw_url)
def test_handles_api_auth_error(): context = AuroraCommandContext() mock_scheduler_proxy = mock.create_autospec(spec=SchedulerProxyApiSpec, instance=True) mock_scheduler_proxy.killTasks.side_effect = SchedulerProxy.AuthError() mock_api = AuroraClientAPI(TEST_CLUSTER, 'user-agent') mock_api._scheduler_proxy = mock_scheduler_proxy context.apis = {TEST_CLUSTER.name: mock_api} api = context.get_api(TEST_CLUSTER.name, clusters={TEST_CLUSTER.name: TEST_CLUSTER}) with pytest.raises(Context.CommandError) as e: api.kill_job(AuroraJobKey(TEST_CLUSTER.name, 'role', 'env', 'job')) assert e.value.code == EXIT_AUTH_ERROR assert mock_scheduler_proxy.killTasks.call_count == 1
def scheduler_list_backups(cluster): """usage: scheduler_list_backups cluster Lists backups available for recovery. """ options = app.get_options() resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).list_backups() check_and_log_response(resp) backups = resp.result.listBackupsResult.backups print('%s available backups:' % len(backups)) for backup in backups: print(backup)
def get_locks(cluster): """usage: get_locks cluster Prints all context/operation locks in the scheduler. """ options = app.get_options() resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).get_locks() check_and_log_response(resp) pp = pprint.PrettyPrinter(indent=2) def pretty_print_lock(lock): return pp.pformat(vars(lock)) print_results([',\n'.join(pretty_print_lock(t) for t in resp.result.getLocksResult.locks)])
def sla_probe_hosts(cluster, percentage, duration): """usage: sla_probe_hosts [--filename=FILENAME] [--grouping=GROUPING] [--hosts=HOSTS] [--min_job_instance_count=COUNT] cluster percentage duration Probes individual hosts with respect to their job SLA. Specifically, given a host, outputs all affected jobs with their projected SLAs if the host goes down. In addition, if a job's projected SLA does not clear the specified limits suggests the approximate time when that job reaches its SLA. Output format: HOST JOB PREDICTED_SLA SAFE? PREDICTED_SAFE_IN where: HOST - host being probed. JOB - job that has tasks running on the host being probed. PREDICTED_SLA - predicted effective percentage of up tasks if the host is shut down. SAFE? - PREDICTED_SLA >= percentage PREDICTED_SAFE_IN - expected wait time in seconds for the job to reach requested SLA threshold. """ options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) hosts = parse_hostnames(options.filename, options.hosts) get_grouping_or_die(options.grouping) vector = AuroraClientAPI( CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector(options.min_instance_count, hosts) groups = vector.probe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), options.grouping) output, _ = format_sla_results(groups) print_results(output)
def sla_probe_hosts(cluster, percentage, duration): """usage: sla_probe_hosts [--filename=FILENAME] [--grouping=GROUPING] [--hosts=HOSTS] [--min_job_instance_count=COUNT] cluster percentage duration Probes individual hosts with respect to their job SLA. Specifically, given a host, outputs all affected jobs with their projected SLAs if the host goes down. In addition, if a job's projected SLA does not clear the specified limits suggests the approximate time when that job reaches its SLA. Output format: HOST JOB PREDICTED_SLA SAFE? PREDICTED_SAFE_IN where: HOST - host being probed. JOB - job that has tasks running on the host being probed. PREDICTED_SLA - predicted effective percentage of up tasks if the host is shut down. SAFE? - PREDICTED_SLA >= percentage PREDICTED_SAFE_IN - expected wait time in seconds for the job to reach requested SLA threshold. """ options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) hosts = parse_hostnames(options.filename, options.hosts) get_grouping_or_die(options.grouping) vector = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector( options.min_instance_count, hosts ) groups = vector.probe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), options.grouping) output, _ = format_sla_results(groups) print_results(output)
def scheduler_print_recovery_tasks(cluster): """usage: scheduler_print_recovery_tasks cluster Prints all active tasks in a staged recovery. """ options = app.get_options() resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).query_recovery( TaskQuery(statuses=ACTIVE_STATES)) check_and_log_response(resp) log.info('Role\tJob\tShard\tStatus\tTask ID') for task in resp.result.queryRecoveryResult.tasks: assigned = task.assignedTask conf = assigned.task log.info('\t'.join((conf.owner.role, conf.jobName, str(assigned.instanceId), ScheduleStatus._VALUES_TO_NAMES[task.status], assigned.taskId)))
def set_quota(cluster, role, cpu_str, ram, disk): """usage: set_quota cluster role cpu ram[MGT] disk[MGT] Alters the amount of production quota allocated to a user. """ try: ram_size = parse_data(ram).as_(Data.MB) disk_size = parse_data(disk).as_(Data.MB) except ValueError as e: die(str(e)) try: cpu = float(cpu_str) ram_mb = int(ram_size) disk_mb = int(disk_size) except ValueError as e: die(str(e)) options = app.get_options() resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).set_quota(role, cpu, ram_mb, disk_mb) check_and_log_response(resp)
def make_admin_client(cluster, verbose=False, bypass_leader_redirect=False): """Creates an API client with the specified options for use in admin commands. :param cluster: The cluster to connect with. :type cluster: Either a string cluster name or a Cluster object. :param verbose: Should the client emit verbose output. :type verbose: bool :type bypass_leader_redirect: Should the client bypass the scheduler's leader redirect filter. :type bypass_leader_redirect: bool :rtype: an AuroraClientAPI instance. """ is_cluster_object = isinstance(cluster, Cluster) if not is_cluster_object and cluster not in CLUSTERS: die('Unknown cluster: %s. Known clusters: %s' % (cluster, ", ".join(CLUSTERS.keys()))) return AuroraClientAPI(cluster if is_cluster_object else CLUSTERS[cluster], AURORA_ADMIN_USER_AGENT_NAME, verbose=verbose, bypass_leader_redirect=bypass_leader_redirect)
def __init__(self, cluster, verbosity, wait_event=None): self._client = AuroraClientAPI(cluster, verbosity == 'verbose') self._wait_event = wait_event or Event()
class HostMaintenance(object): """Submit requests to the scheduler to put hosts into and out of maintenance mode so they can be operated upon without causing LOST tasks. Aurora provides a two-tiered concept of Maintenance. The first step is to initiate maintenance, which will ask the Aurora scheduler to de-prioritize scheduling on a large set of hosts (the ones that will be operated upon during this maintenance window). Once all hosts have been tagged in this manner, the operator can begin draining individual machines, which will have all user-tasks killed and rescheduled. When the tasks get placed onto a new machine, the scheduler will first look for hosts that do not have the maintenance tag, which will help decrease churn and prevent a task from being constantly killed as its hosts go down from underneath it. """ SLA_MIN_JOB_INSTANCE_COUNT = 20 STATUS_POLL_INTERVAL = Amount(5, Time.SECONDS) MAX_STATUS_WAIT = Amount(5, Time.MINUTES) @classmethod def iter_batches(cls, hostnames, grouping_function=DEFAULT_GROUPING): groups = group_hosts(hostnames, grouping_function) groups = sorted(groups.items(), key=lambda v: v[0]) for group in groups: yield Hosts(group[1]) def __init__(self, cluster, verbosity, wait_event=None): self._client = AuroraClientAPI(cluster, verbosity == 'verbose') self._wait_event = wait_event or Event() def _drain_hosts(self, drainable_hosts): """"Drains tasks from the specified hosts. This will move active tasks on these hosts to the DRAINING state, causing them to be rescheduled elsewhere. :param drainable_hosts: Hosts that are in maintenance mode and ready to be drained :type drainable_hosts: gen.apache.aurora.ttypes.Hosts :rtype: set of host names failed to drain """ check_and_log_response(self._client.drain_hosts(drainable_hosts)) drainable_hostnames = [ hostname for hostname in drainable_hosts.hostNames ] total_wait = self.STATUS_POLL_INTERVAL not_drained_hostnames = set(drainable_hostnames) while not self._wait_event.is_set() and not_drained_hostnames: log.info('Waiting for hosts to be in DRAINED: %s' % not_drained_hostnames) self._wait_event.wait(self.STATUS_POLL_INTERVAL.as_(Time.SECONDS)) statuses = self.check_status(list(not_drained_hostnames)) not_drained_hostnames = set(h[0] for h in statuses if h[1] != 'DRAINED') total_wait += self.STATUS_POLL_INTERVAL if not_drained_hostnames and total_wait > self.MAX_STATUS_WAIT: log.warning( 'Failed to move all hosts into DRAINED within %s:\n%s' % (self.MAX_STATUS_WAIT, '\n'.join("\tHost:%s\tStatus:%s" % h for h in sorted(statuses) if h[1] != 'DRAINED'))) break return not_drained_hostnames def _complete_maintenance(self, drained_hosts): """End the maintenance status for a given set of hosts. :param drained_hosts: Hosts that are drained and finished being operated upon :type drained_hosts: gen.apache.aurora.ttypes.Hosts """ check_and_log_response(self._client.end_maintenance(drained_hosts)) resp = self._client.maintenance_status(drained_hosts) for host_status in resp.result.maintenanceStatusResult.statuses: if host_status.mode != MaintenanceMode.NONE: log.warning('%s is DRAINING or in DRAINED' % host_status.host) def _check_sla(self, hostnames, grouping_function, percentage, duration): """Check if the provided list of hosts passes the job uptime SLA check. This is an all-or-nothing check, meaning that all provided hosts must pass their job SLA check for the maintenance to proceed. :param hostnames: list of host names to check SLA for :type hostnames: list of strings :param grouping_function: grouping function to apply to the given hosts :type grouping_function: function :param percentage: SLA uptime percentage override :type percentage: float :param duration: SLA uptime duration override :type duration: twitter.common.quantity.Amount :rtype: set of unsafe hosts """ vector = self._client.sla_get_safe_domain_vector( self.SLA_MIN_JOB_INSTANCE_COUNT, hostnames) host_groups = vector.probe_hosts(percentage, duration.as_(Time.SECONDS), grouping_function) unsafe_hostnames = set() # Given that maintenance is performed 1 group at a time, any result longer than 1 group # should be considered a batch failure. if host_groups: if len(host_groups) > 1: log.error( 'Illegal multiple groups detected in SLA results. Skipping hosts: %s' % hostnames) return set(hostnames) results, unsafe_hostnames = format_sla_results(host_groups, unsafe_only=True) if results: print_results(results) return unsafe_hostnames return unsafe_hostnames def end_maintenance(self, hostnames): """Pull a list of hostnames out of maintenance mode. :param hostnames: List of hosts to operate upon :type hostnames: list of strings """ self._complete_maintenance(Hosts(set(hostnames))) def start_maintenance(self, hostnames): """Put a list of hostnames into maintenance mode, to de-prioritize scheduling. This is part of two-phase draining- tasks will still be running on these hosts until drain_hosts is called upon them. :param hostnames: List of hosts to set for initial maintenance :type hostnames: list of strings :rtype: list of hostnames with the maintenance mode set """ resp = self._client.start_maintenance(Hosts(set(hostnames))) check_and_log_response(resp) result = [ host_status.host for host_status in resp.result.startMaintenanceResult.statuses ] if len(result) != len(hostnames): log.warning('Skipping maintenance for unknown hosts: %s' % (set(hostnames) - set(result))) return result def _operate_on_hosts(self, drained_hosts, callback): """Perform a given operation on a list of hosts that are ready for maintenance. :param drained_hosts: Hosts that have been drained (via _drain_hosts) :type drained_hosts: list of strings :param callback: Function to call one hostname at a time :type callback: function """ for hostname in drained_hosts: callback(hostname) def perform_maintenance(self, hostnames, grouping_function=DEFAULT_GROUPING, percentage=None, duration=None, output_file=None, callback=None): """Put hosts into maintenance mode and drain them. Walk through the process of putting hosts into maintenance and draining them of tasks. The hosts will remain in maintenance mode upon completion. :param hostnames: A list of hostnames to operate upon :type hostnames: list of strings :param grouping_function: How to split up the hostname into groups :type grouping_function: function :param percentage: SLA percentage to use :type percentage: float :param duration: SLA duration to use :type duration: twitter.common.quantity.Time :param output_file: file to write hosts that were not drained due to failed SLA check :type output_file: string :param callback: Function to call once hosts are drained :type callback: function :rtype: set of host names that were successfully drained """ hostnames = self.start_maintenance(hostnames) not_drained_hostnames = set() for hosts in self.iter_batches(hostnames, grouping_function): log.info('Beginning SLA check for %s' % hosts.hostNames) unsafe_hostnames = self._check_sla(list(hosts.hostNames), grouping_function, percentage, duration) if unsafe_hostnames: log.warning( 'Some hosts did not pass SLA check and will not be drained! ' 'Skipping hosts: %s' % unsafe_hostnames) not_drained_hostnames |= unsafe_hostnames drainable_hostnames = hosts.hostNames - unsafe_hostnames if not drainable_hostnames: continue hosts = Hosts(drainable_hostnames) else: log.info('All hosts passed SLA check.') not_drained_hostnames |= self._drain_hosts(hosts) if callback: self._operate_on_hosts(hosts.hostNames - not_drained_hostnames, callback) if not_drained_hostnames: output = '\n'.join(list(not_drained_hostnames)) log.info( 'The following hosts WERE NOT DRAINED due to failed SLA check or external failures:' ) print(output) if output_file: try: with open(output_file, 'w') as fp: fp.write(output) fp.write('\n') log.info('Written unsafe host names into: %s' % output_file) except IOError as e: log.error('Failed to write into the output file: %s' % e) return set(hostnames) - not_drained_hostnames def check_status(self, hostnames): """Query the scheduler to determine the maintenance status for a list of hostnames :param hostnames: Hosts to query for :type hostnames: list of strings :rtype: list of 2-tuples, hostname and MaintenanceMode """ resp = self._client.maintenance_status(Hosts(set(hostnames))) check_and_log_response(resp) statuses = [] for host_status in resp.result.maintenanceStatusResult.statuses: statuses.append( (host_status.host, MaintenanceMode._VALUES_TO_NAMES[host_status.mode])) return statuses
def __init__(self, cluster, verbosity, wait_event=None): self._client = AuroraClientAPI(cluster, verbosity == "verbose") self._wait_event = wait_event or Event()
class DistributedCommandRunner(object): @classmethod def make_executor_path(cls, cluster, executor_name): parameters = cls.sandbox_args(cluster) parameters.update(executor_name=executor_name) return ( posixpath.join( "%(slave_root)s", "slaves/*/frameworks/*/executors/%(executor_name)s/runs", "%(slave_run_directory)s" ) % parameters ) @classmethod def thermos_sandbox(cls, cluster, executor_sandbox=False): sandbox = cls.make_executor_path(cluster, "thermos-{{thermos.task_id}}") return sandbox if executor_sandbox else posixpath.join(sandbox, "sandbox") @classmethod def sandbox_args(cls, cluster): cluster = cluster.with_trait(CommandRunnerTrait) return {"slave_root": cluster.slave_root, "slave_run_directory": cluster.slave_run_directory} @classmethod def substitute_thermos(cls, command, task, cluster, **kw): prefix_command = "cd %s;" % cls.thermos_sandbox(cluster, **kw) thermos_namespace = ThermosContext(task_id=task.assignedTask.taskId, ports=task.assignedTask.assignedPorts) mesos_namespace = MesosContext(instance=task.assignedTask.instanceId) command = String(prefix_command + command) % Environment(thermos=thermos_namespace, mesos=mesos_namespace) return command.get() @classmethod def aurora_sandbox(cls, cluster, executor_sandbox=False): if executor_sandbox: return cls.make_executor_path(cluster, "twitter") else: return "/var/run/nexus/%task_id%/sandbox" @classmethod def substitute_aurora(cls, command, task, cluster, **kw): command = ("cd %s;" % cls.aurora_sandbox(cluster, **kw)) + command command = command.replace("%shard_id%", str(task.assignedTask.instanceId)) command = command.replace("%task_id%", task.assignedTask.taskId) for name, port in task.assignedTask.assignedPorts.items(): command = command.replace("%port:" + name + "%", str(port)) return command @classmethod def substitute(cls, command, task, cluster, **kw): if task.assignedTask.task.executorConfig: return cls.substitute_thermos(command, task, cluster, **kw) else: return cls.substitute_aurora(command, task, cluster, **kw) @classmethod def query_from(cls, role, env, job): return TaskQuery(statuses=LIVE_STATES, jobKeys=[JobKey(role=role, environment=env, name=job)]) def __init__(self, cluster, role, env, jobs, ssh_user=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._log = log_fn def execute(self, args): hostname, role, command = args ssh_command = ["ssh", "-n", "-q", "%s@%s" % (role, hostname), command] self._log(logging.DEBUG, "Running command: %s" % ssh_command) po = subprocess.Popen(ssh_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = po.communicate() return "\n".join("%s: %s" % (hostname, line) for line in output[0].splitlines()) def resolve(self): for job in self._jobs: resp = self._api.query(self.query_from(self._role, self._env, job)) if resp.responseCode != ResponseCode.OK: self._log(logging.ERROR, "Failed to query job: %s" % job) continue for task in resp.result.scheduleStatusResult.tasks: yield task def process_arguments(self, command, **kw): for task in self.resolve(): host = task.assignedTask.slaveHost yield (host, self._ssh_user, self.substitute(command, task, self._cluster, **kw)) def run(self, command, parallelism=1, **kw): threadpool = ThreadPool(processes=parallelism) for result in threadpool.imap_unordered(self.execute, self.process_arguments(command, **kw)): print(result)
class HostMaintenance(object): """Submit requests to the scheduler to put hosts into and out of maintenance mode so they can be operated upon without causing LOST tasks. """ DEFAULT_GROUPING = 'by_host' GROUPING_FUNCTIONS = { 'by_host': group_by_host, } START_MAINTENANCE_DELAY = Amount(30, Time.SECONDS) @classmethod def group_hosts(cls, hostnames, grouping_function=DEFAULT_GROUPING): try: grouping_function = cls.GROUPING_FUNCTIONS[grouping_function] except KeyError: raise ValueError('Unknown grouping function %s!' % grouping_function) groups = defaultdict(set) for hostname in hostnames: groups[grouping_function(hostname)].add(hostname) return groups @classmethod def iter_batches(cls, hostnames, groups_per_batch, grouping_function=DEFAULT_GROUPING): if groups_per_batch <= 0: raise ValueError('Batch size must be > 0!') groups = cls.group_hosts(hostnames, grouping_function) groups = sorted(groups.items(), key=lambda v: v[0]) for k in range(0, len(groups), groups_per_batch): yield Hosts(set.union(*(hostset for (key, hostset) in groups[k:k + groups_per_batch]))) def __init__(self, cluster, verbosity): self._client = AuroraClientAPI(cluster, verbosity == 'verbose') def _drain_hosts(self, drainable_hosts, clock=time): """This will actively turn down tasks running on hosts.""" check_and_log_response(self._client.drain_hosts(drainable_hosts)) not_ready_hosts = [hostname for hostname in drainable_hosts.hostNames] while not_ready_hosts: log.info("Sleeping for %s." % self.START_MAINTENANCE_DELAY) clock.sleep(self.START_MAINTENANCE_DELAY.as_(Time.SECONDS)) resp = self._client.maintenance_status(Hosts(not_ready_hosts)) if not resp.result.maintenanceStatusResult.statuses: not_ready_hosts = None for host_status in resp.result.maintenanceStatusResult.statuses: if host_status.mode != MaintenanceMode.DRAINED: log.warning('%s is currently in status %s' % (host_status.host, MaintenanceMode._VALUES_TO_NAMES[host_status.mode])) else: not_ready_hosts.remove(host_status.host) def _complete_maintenance(self, drained_hosts): """End the maintenance status for a give set of hosts.""" check_and_log_response(self._client.end_maintenance(drained_hosts)) resp = self._client.maintenance_status(drained_hosts) for host_status in resp.result.maintenanceStatusResult.statuses: if host_status.mode != MaintenanceMode.NONE: log.warning('%s is DRAINING or in DRAINED' % host_status.host) def _operate_on_hosts(self, drained_hosts, callback): """Perform a given operation on a list of hosts that are ready for maintenance.""" for host in drained_hosts.hostNames: callback(host) def end_maintenance(self, hosts): """Pull a list of hosts out of maintenance mode.""" self._complete_maintenance(Hosts(set(hosts))) def start_maintenance(self, hosts): """Put a list of hosts into maintenance mode, to de-prioritize scheduling.""" check_and_log_response(self._client.start_maintenance(Hosts(set(hosts)))) def perform_maintenance(self, hosts, groups_per_batch=1, grouping_function=DEFAULT_GROUPING, callback=None): """The wrap a callback in between sending hosts into maintenance mode and back. Walk through the process of putting hosts into maintenance, draining them of tasks, performing an action on them once drained, then removing them from maintenance mode so tasks can schedule. """ self._complete_maintenance(Hosts(set(hosts))) self.start_maintenance(hosts) for hosts in self.iter_batches(hosts, groups_per_batch, grouping_function): self._drain_hosts(hosts) if callback: self._operate_on_hosts(hosts, callback) self._complete_maintenance(hosts) def check_status(self, hosts): resp = self._client.maintenance_status(Hosts(set(hosts))) check_and_log_response(resp) statuses = [] for host_status in resp.result.maintenanceStatusResult.statuses: statuses.append((host_status.host, MaintenanceMode._VALUES_TO_NAMES[host_status.mode])) return statuses
def query(args, options): """usage: query [--force] [--listformat=FORMAT] [--shards=N[,N,...]] [--states=State[,State,...]] cluster [role [job]] Query Mesos about jobs and tasks. """ def _convert_fmt_string(fmtstr): import re def convert(match): return "%%(%s)s" % match.group(1) return re.sub(r'%(\w+)%', convert, fmtstr) def flatten_task(t, d={}): for key in t.__dict__.keys(): val = getattr(t, key) try: val.__dict__.keys() except AttributeError: d[key] = val else: flatten_task(val, d) return d def map_values(d): default_value = lambda v: v mapping = { 'status': lambda v: ScheduleStatus._VALUES_TO_NAMES[v], } return dict( (k, mapping.get(k, default_value)(v)) for (k, v) in d.items() ) for state in options.states.split(','): if state not in ScheduleStatus._NAMES_TO_VALUES: msg = "Unknown state '%s' specified. Valid states are:\n" % state msg += ','.join(ScheduleStatus._NAMES_TO_VALUES.keys()) die(msg) # Role, Job, Instances, States, and the listformat if len(args) == 0: die('Must specify at least cluster.') cluster = args[0] role = args[1] if len(args) > 1 else None job = args[2] if len(args) > 2 else None instances = set(map(int, options.shards.split(','))) if options.shards else set() if options.states: states = set(map(ScheduleStatus._NAMES_TO_VALUES.get, options.states.split(','))) else: states = ACTIVE_STATES | TERMINAL_STATES listformat = _convert_fmt_string(options.listformat) # Figure out "expensive" queries here and bone if they do not have --force # - Does not specify role if not role and not options.force: die('--force is required for expensive queries (no role specified)') # - Does not specify job if not job and not options.force: die('--force is required for expensive queries (no job specified)') # - Specifies status outside of ACTIVE_STATES if not (states <= ACTIVE_STATES) and not options.force: die('--force is required for expensive queries (states outside ACTIVE states') api = AuroraClientAPI(CLUSTERS[cluster], options.verbosity) query_info = api.query(api.build_query(role, job, instances=instances, statuses=states)) if query_info.responseCode != ResponseCode.OK: die('Failed to query scheduler: %s' % query_info.messageDEPRECATED) tasks = query_info.result.scheduleStatusResult.tasks if tasks is None: return try: for task in tasks: d = flatten_task(task) print(listformat % map_values(d)) except KeyError: msg = "Unknown key in format string. Valid keys are:\n" msg += ','.join(d.keys()) die(msg)
def mock_api(cls): api = AuroraClientAPI(Cluster(name="foo")) mock_proxy = Mock() api._scheduler_proxy = mock_proxy return api, mock_proxy
def sla_list_safe_domain(cluster, percentage, duration): """usage: sla_list_safe_domain [--exclude_file=FILENAME] [--exclude_hosts=HOSTS] [--grouping=GROUPING] [--include_file=FILENAME] [--include_hosts=HOSTS] [--list_jobs] [--min_job_instance_count=COUNT] [--override_jobs=FILENAME] cluster percentage duration Returns a list of relevant hosts where it would be safe to kill tasks without violating their job SLA. The SLA is defined as a pair of percentage and duration, where: percentage - Percentage of tasks required to be up within the duration. Applied to all jobs except those listed in --override_jobs file; duration - Time interval (now - value) for the percentage of up tasks. Applied to all jobs except those listed in --override_jobs file. Format: XdYhZmWs (each field is optional but must be in that order.) Examples: 5m, 1d3h45m. NOTE: if --grouping option is specified and is set to anything other than default (by_host) the results will be processed and filtered based on the grouping function on a all-or-nothing basis. In other words, the group is 'safe' IFF it is safe to kill tasks on all hosts in the group at the same time. """ def parse_jobs_file(filename): result = {} with open(filename, 'r') as overrides: for line in overrides: if not line.strip(): continue tokens = line.split() if len(tokens) != 3: die('Invalid line in %s:%s' % (filename, line)) job_key = AuroraJobKey.from_path(tokens[0]) result[job_key] = JobUpTimeLimit( job=job_key, percentage=parse_sla_percentage(tokens[1]), duration_secs=parse_time(tokens[2]).as_(Time.SECONDS) ) return result options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) exclude_hosts = parse_hostnames_optional(options.exclude_hosts, options.exclude_filename) include_hosts = parse_hostnames_optional(options.include_hosts, options.include_filename) override_jobs = parse_jobs_file(options.override_filename) if options.override_filename else {} get_grouping_or_die(options.grouping) vector = AuroraClientAPI( CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector(options.min_instance_count, include_hosts) groups = vector.get_safe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), override_jobs, options.grouping) results = [] for group in groups: for host in sorted(group.keys()): if exclude_hosts and host in exclude_hosts: continue if options.list_jobs: results.append('\n'.join(['%s\t%s\t%.2f\t%d' % (host, d.job.to_path(), d.percentage, d.duration_secs) for d in sorted(group[host])])) else: results.append('%s' % host) print_results(results)
class DistributedCommandRunner(object): @staticmethod def execute(args): hostname, role, command = args ssh_command = ['ssh', '-n', '-q', '%s@%s' % (role, hostname), command] po = subprocess.Popen(ssh_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = po.communicate() return '\n'.join('%s: %s' % (hostname, line) for line in output[0].splitlines()) @classmethod def make_executor_path(cls, cluster, executor_name): parameters = cls.sandbox_args(cluster) parameters.update(executor_name=executor_name) return posixpath.join( '%(slave_root)s', 'slaves/*/frameworks/*/executors/%(executor_name)s/runs', '%(slave_run_directory)s' ) % parameters @classmethod def thermos_sandbox(cls, cluster, executor_sandbox=False): sandbox = cls.make_executor_path(cluster, 'thermos-{{thermos.task_id}}') return sandbox if executor_sandbox else posixpath.join(sandbox, 'sandbox') @classmethod def sandbox_args(cls, cluster): cluster = cluster.with_trait(CommandRunnerTrait) return {'slave_root': cluster.slave_root, 'slave_run_directory': cluster.slave_run_directory} @classmethod def substitute_thermos(cls, command, task, cluster, **kw): prefix_command = 'cd %s;' % cls.thermos_sandbox(cluster, **kw) thermos_namespace = ThermosContext( task_id=task.assignedTask.taskId, ports=task.assignedTask.assignedPorts) mesos_namespace = MesosContext(instance=task.assignedTask.instanceId) command = String(prefix_command + command) % Environment( thermos=thermos_namespace, mesos=mesos_namespace) return command.get() @classmethod def aurora_sandbox(cls, cluster, executor_sandbox=False): if executor_sandbox: return cls.make_executor_path(cluster, 'twitter') else: return '/var/run/nexus/%task_id%/sandbox' @classmethod def substitute_aurora(cls, command, task, cluster, **kw): command = ('cd %s;' % cls.aurora_sandbox(cluster, **kw)) + command command = command.replace('%shard_id%', str(task.assignedTask.instanceId)) command = command.replace('%task_id%', task.assignedTask.taskId) for name, port in task.assignedTask.assignedPorts.items(): command = command.replace('%port:' + name + '%', str(port)) return command @classmethod def substitute(cls, command, task, cluster, **kw): if task.assignedTask.task.executorConfig: return cls.substitute_thermos(command, task, cluster, **kw) else: return cls.substitute_aurora(command, task, cluster, **kw) @classmethod def query_from(cls, role, env, job): return TaskQuery(statuses=LIVE_STATES, owner=Identity(role), jobName=job, environment=env) def __init__(self, cluster, role, env, jobs, ssh_user=None): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role def resolve(self): for job in self._jobs: resp = self._api.query(self.query_from(self._role, self._env, job)) if resp.responseCode != ResponseCode.OK: log.error('Failed to query job: %s' % job) continue for task in resp.result.scheduleStatusResult.tasks: yield task def process_arguments(self, command, **kw): for task in self.resolve(): host = task.assignedTask.slaveHost role = task.assignedTask.task.owner.role yield (host, self._ssh_user, self.substitute(command, task, self._cluster, **kw)) def run(self, command, parallelism=1, **kw): threadpool = ThreadPool(processes=parallelism) for result in threadpool.imap_unordered(self.execute, self.process_arguments(command, **kw)): print result
def sla_list_safe_domain(cluster, percentage, duration): """usage: sla_list_safe_domain [--exclude_file=FILENAME] [--exclude_hosts=HOSTS] [--grouping=GROUPING] [--include_file=FILENAME] [--include_hosts=HOSTS] [--list_jobs] [--min_job_instance_count=COUNT] [--override_jobs=FILENAME] cluster percentage duration Returns a list of relevant hosts where it would be safe to kill tasks without violating their job SLA. The SLA is defined as a pair of percentage and duration, where: percentage - Percentage of tasks required to be up within the duration. Applied to all jobs except those listed in --override_jobs file; duration - Time interval (now - value) for the percentage of up tasks. Applied to all jobs except those listed in --override_jobs file. Format: XdYhZmWs (each field is optional but must be in that order.) Examples: 5m, 1d3h45m. NOTE: if --grouping option is specified and is set to anything other than default (by_host) the results will be processed and filtered based on the grouping function on a all-or-nothing basis. In other words, the group is 'safe' IFF it is safe to kill tasks on all hosts in the group at the same time. """ def parse_jobs_file(filename): result = {} with open(filename, "r") as overrides: for line in overrides: if not line.strip(): continue tokens = line.split() if len(tokens) != 3: die("Invalid line in %s:%s" % (filename, line)) job_key = AuroraJobKey.from_path(tokens[0]) result[job_key] = JobUpTimeLimit( job=job_key, percentage=parse_sla_percentage(tokens[1]), duration_secs=parse_time(tokens[2]).as_(Time.SECONDS), ) return result options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) exclude_hosts = parse_hostnames_optional(options.exclude_hosts, options.exclude_filename) include_hosts = parse_hostnames_optional(options.include_hosts, options.include_filename) override_jobs = parse_jobs_file(options.override_filename) if options.override_filename else {} get_grouping_or_die(options.grouping) vector = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector( options.min_instance_count, include_hosts ) groups = vector.get_safe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), override_jobs, options.grouping) results = [] for group in groups: for host in sorted(group.keys()): if exclude_hosts and host in exclude_hosts: continue if options.list_jobs: results.append( "\n".join( [ "%s\t%s\t%.2f\t%d" % (host, d.job.to_path(), d.percentage, d.duration_secs) for d in sorted(group[host]) ] ) ) else: results.append("%s" % host) print_results(results)
class HostMaintenance(object): """Submit requests to the scheduler to put hosts into and out of maintenance mode so they can be operated upon without causing LOST tasks. Aurora provides a two-tiered concept of Maintenance. The first step is to initiate maintenance, which will ask the Aurora scheduler to de-prioritize scheduling on a large set of hosts (the ones that will be operated upon during this maintenance window). Once all hosts have been tagged in this manner, the operator can begin draining individual machines, which will have all user-tasks killed and rescheduled. When the tasks get placed onto a new machine, the scheduler will first look for hosts that do not have the maintenance tag, which will help decrease churn and prevent a task from being constantly killed as its hosts go down from underneath it. """ START_MAINTENANCE_DELAY = Amount(30, Time.SECONDS) SLA_MIN_JOB_INSTANCE_COUNT = 20 @classmethod def iter_batches(cls, hostnames, grouping_function=DEFAULT_GROUPING): groups = group_hosts(hostnames, grouping_function) groups = sorted(groups.items(), key=lambda v: v[0]) for group in groups: yield Hosts(group[1]) def __init__(self, cluster, verbosity): self._client = AuroraClientAPI(cluster, verbosity == 'verbose') def _drain_hosts(self, drainable_hosts, clock=time): """"Drains tasks from the specified hosts. This will move active tasks on these hosts to the DRAINING state, causing them to be rescheduled elsewhere. :param drainable_hosts: Hosts that are in maintenance mode and ready to be drained :type drainable_hosts: gen.apache.aurora.ttypes.Hosts :param clock: time module for testing :type clock: time """ check_and_log_response(self._client.drain_hosts(drainable_hosts)) not_ready_hostnames = [hostname for hostname in drainable_hosts.hostNames] while not_ready_hostnames: log.info("Sleeping for %s." % self.START_MAINTENANCE_DELAY) clock.sleep(self.START_MAINTENANCE_DELAY.as_(Time.SECONDS)) resp = self._client.maintenance_status(Hosts(set(not_ready_hostnames))) if not resp.result.maintenanceStatusResult.statuses: not_ready_hostnames = None for host_status in resp.result.maintenanceStatusResult.statuses: if host_status.mode != MaintenanceMode.DRAINED: log.warning('%s is currently in status %s' % (host_status.host, MaintenanceMode._VALUES_TO_NAMES[host_status.mode])) else: not_ready_hostnames.remove(host_status.host) def _complete_maintenance(self, drained_hosts): """End the maintenance status for a given set of hosts. :param drained_hosts: Hosts that are drained and finished being operated upon :type drained_hosts: gen.apache.aurora.ttypes.Hosts """ check_and_log_response(self._client.end_maintenance(drained_hosts)) resp = self._client.maintenance_status(drained_hosts) for host_status in resp.result.maintenanceStatusResult.statuses: if host_status.mode != MaintenanceMode.NONE: log.warning('%s is DRAINING or in DRAINED' % host_status.host) def _check_sla(self, hostnames, grouping_function, percentage, duration): """Check if the provided list of hosts passes the job uptime SLA check. This is an all-or-nothing check, meaning that all provided hosts must pass their job SLA check for the maintenance to proceed. :param hostnames: list of host names to check SLA for :type hostnames: list of strings :param grouping_function: grouping function to apply to the given hosts :type grouping_function: function :param percentage: SLA uptime percentage override :type percentage: float :param duration: SLA uptime duration override :type duration: twitter.common.quantity.Amount :rtype: set of unsafe hosts """ vector = self._client.sla_get_safe_domain_vector(self.SLA_MIN_JOB_INSTANCE_COUNT, hostnames) host_groups = vector.probe_hosts( percentage, duration.as_(Time.SECONDS), grouping_function) unsafe_hostnames = set() # Given that maintenance is performed 1 group at a time, any result longer than 1 group # should be considered a batch failure. if host_groups: if len(host_groups) > 1: log.error('Illegal multiple groups detected in SLA results. Skipping hosts: %s' % hostnames) return set(hostnames) results, unsafe_hostnames = format_sla_results(host_groups, unsafe_only=True) if results: print_results(results) return unsafe_hostnames return unsafe_hostnames def end_maintenance(self, hostnames): """Pull a list of hostnames out of maintenance mode. :param hostnames: List of hosts to operate upon :type hostnames: list of strings """ self._complete_maintenance(Hosts(set(hostnames))) def start_maintenance(self, hostnames): """Put a list of hostnames into maintenance mode, to de-prioritize scheduling. This is part of two-phase draining- tasks will still be running on these hosts until drain_hosts is called upon them. :param hostnames: List of hosts to set for initial maintenance :type hostnames: list of strings :rtype: list of hostnames with the maintenance mode set """ resp = self._client.start_maintenance(Hosts(set(hostnames))) check_and_log_response(resp) result = [host_status.host for host_status in resp.result.startMaintenanceResult.statuses] if len(result) != len(hostnames): log.warning('Skipping maintenance for unknown hosts: %s' % (set(hostnames) - set(result))) return result def perform_maintenance(self, hostnames, grouping_function=DEFAULT_GROUPING, percentage=None, duration=None, output_file=None): """Put hosts into maintenance mode and drain them. Walk through the process of putting hosts into maintenance and draining them of tasks. The hosts will remain in maintenance mode upon completion. :param hostnames: A list of hostnames to operate upon :type hostnames: list of strings :param grouping_function: How to split up the hostname into groups :type grouping_function: function :param percentage: SLA percentage to use :type percentage: float :param duration: SLA duration to use :type duration: twitter.common.quantity.Time :param output_file: file to write hosts that were not drained due to failed SLA check :type output_file: string :rtype: set of host names that were successfully drained """ hostnames = self.start_maintenance(hostnames) not_drained_hostnames = set() for hosts in self.iter_batches(hostnames, grouping_function): log.info('Beginning SLA check for %s' % hosts.hostNames) unsafe_hostnames = self._check_sla( list(hosts.hostNames), grouping_function, percentage, duration) if unsafe_hostnames: log.warning('Some hosts did not pass SLA check and will not be drained! ' 'Skipping hosts: %s' % unsafe_hostnames) not_drained_hostnames |= unsafe_hostnames drainable_hostnames = hosts.hostNames - unsafe_hostnames if not drainable_hostnames: continue hosts = Hosts(drainable_hostnames) else: log.info('All hosts passed SLA check.') self._drain_hosts(hosts) if not_drained_hostnames: output = '\n'.join(list(not_drained_hostnames)) log.info('The following hosts did not pass SLA check and were not drained:') print(output) if output_file: try: with open(output_file, 'w') as fp: fp.write(output) fp.write('\n') log.info('Written unsafe host names into: %s' % output_file) except IOError as e: log.error('Failed to write into the output file: %s' % e) return set(hostnames) - not_drained_hostnames def check_status(self, hostnames): """Query the scheduler to determine the maintenance status for a list of hostnames :param hostnames: Hosts to query for :type hostnames: list of strings :rtype: list of 2-tuples, hostname and MaintenanceMode """ resp = self._client.maintenance_status(Hosts(set(hostnames))) check_and_log_response(resp) statuses = [] for host_status in resp.result.maintenanceStatusResult.statuses: statuses.append((host_status.host, MaintenanceMode._VALUES_TO_NAMES[host_status.mode])) return statuses
def __init__(self, cluster, verbosity): self._client = AuroraClientAPI(cluster, verbosity == 'verbose')
class DistributedCommandRunner(object): @classmethod def make_executor_path(cls, cluster, executor_name): parameters = cls.sandbox_args(cluster) parameters.update(executor_name=executor_name) return posixpath.join( '%(slave_root)s', 'slaves/*/frameworks/*/executors/%(executor_name)s/runs', '%(slave_run_directory)s' ) % parameters @classmethod def thermos_sandbox(cls, cluster, executor_sandbox=False): sandbox = cls.make_executor_path(cluster, 'thermos-{{thermos.task_id}}') return sandbox if executor_sandbox else posixpath.join(sandbox, 'sandbox') @classmethod def sandbox_args(cls, cluster): cluster = cluster.with_trait(CommandRunnerTrait) return {'slave_root': cluster.slave_root, 'slave_run_directory': cluster.slave_run_directory} @classmethod def substitute(cls, command, task, cluster, **kw): prefix_command = 'cd %s;' % cls.thermos_sandbox(cluster, **kw) thermos_namespace = ThermosContext( task_id=task.assignedTask.taskId, ports=task.assignedTask.assignedPorts) mesos_namespace = MesosContext(instance=task.assignedTask.instanceId) command = String(prefix_command + command) % Environment( thermos=thermos_namespace, mesos=mesos_namespace) return command.get() @classmethod def query_from(cls, role, env, job): return TaskQuery(statuses=LIVE_STATES, jobKeys=[JobKey(role=role, environment=env, name=job)]) def __init__(self, cluster, role, env, jobs, ssh_user=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI( cluster=cluster, user_agent=AURORA_V2_USER_AGENT_NAME) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._log = log_fn def execute(self, args): hostname, role, command = args ssh_command = ['ssh', '-n', '-q', '%s@%s' % (role, hostname), command] self._log(logging.DEBUG, "Running command: %s" % ssh_command) po = subprocess.Popen(ssh_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = po.communicate() return '\n'.join('%s: %s' % (hostname, line) for line in output[0].splitlines()) def resolve(self): for job in self._jobs: resp = self._api.query(self.query_from(self._role, self._env, job)) if resp.responseCode != ResponseCode.OK: self._log(logging.ERROR, 'Failed to query job: %s' % job) continue for task in resp.result.scheduleStatusResult.tasks: yield task def process_arguments(self, command, **kw): for task in self.resolve(): host = task.assignedTask.slaveHost yield (host, self._ssh_user, self.substitute(command, task, self._cluster, **kw)) def run(self, command, parallelism=1, **kw): threadpool = ThreadPool(processes=parallelism) for result in threadpool.imap_unordered(self.execute, self.process_arguments(command, **kw)): print(result)
class DistributedCommandRunner(object): @classmethod def make_executor_path(cls, cluster, executor_name): parameters = cls.sandbox_args(cluster) parameters.update(executor_name=executor_name) return posixpath.join( '%(slave_root)s', 'slaves/*/frameworks/*/executors/%(executor_name)s/runs', '%(slave_run_directory)s') % parameters @classmethod def thermos_sandbox(cls, cluster, executor_sandbox=False): sandbox = cls.make_executor_path(cluster, 'thermos-{{thermos.task_id}}') return sandbox if executor_sandbox else posixpath.join( sandbox, 'sandbox') @classmethod def sandbox_args(cls, cluster): cluster = cluster.with_trait(CommandRunnerTrait) return { 'slave_root': cluster.slave_root, 'slave_run_directory': cluster.slave_run_directory } @classmethod def substitute_thermos(cls, command, task, cluster, **kw): prefix_command = 'cd %s;' % cls.thermos_sandbox(cluster, **kw) thermos_namespace = ThermosContext( task_id=task.assignedTask.taskId, ports=task.assignedTask.assignedPorts) mesos_namespace = MesosContext(instance=task.assignedTask.instanceId) command = String(prefix_command + command) % Environment( thermos=thermos_namespace, mesos=mesos_namespace) return command.get() @classmethod def aurora_sandbox(cls, cluster, executor_sandbox=False): if executor_sandbox: return cls.make_executor_path(cluster, 'twitter') else: return '/var/run/nexus/%task_id%/sandbox' @classmethod def substitute_aurora(cls, command, task, cluster, **kw): command = ('cd %s;' % cls.aurora_sandbox(cluster, **kw)) + command command = command.replace('%shard_id%', str(task.assignedTask.instanceId)) command = command.replace('%task_id%', task.assignedTask.taskId) for name, port in task.assignedTask.assignedPorts.items(): command = command.replace('%port:' + name + '%', str(port)) return command @classmethod def substitute(cls, command, task, cluster, **kw): if task.assignedTask.task.executorConfig: return cls.substitute_thermos(command, task, cluster, **kw) else: return cls.substitute_aurora(command, task, cluster, **kw) @classmethod def query_from(cls, role, env, job): return TaskQuery( statuses=LIVE_STATES, jobKeys=[JobKey(role=role, environment=env, name=job)]) def __init__(self, cluster, role, env, jobs, ssh_user=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI(cluster=cluster, user_agent=AURORA_V2_USER_AGENT_NAME) self._role = role self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role self._log = log_fn def execute(self, args): hostname, role, command = args ssh_command = ['ssh', '-n', '-q', '%s@%s' % (role, hostname), command] self._log(logging.DEBUG, "Running command: %s" % ssh_command) po = subprocess.Popen(ssh_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = po.communicate() return '\n'.join('%s: %s' % (hostname, line) for line in output[0].splitlines()) def resolve(self): for job in self._jobs: resp = self._api.query(self.query_from(self._role, self._env, job)) if resp.responseCode != ResponseCode.OK: self._log(logging.ERROR, 'Failed to query job: %s' % job) continue for task in resp.result.scheduleStatusResult.tasks: yield task def process_arguments(self, command, **kw): for task in self.resolve(): host = task.assignedTask.slaveHost yield (host, self._ssh_user, self.substitute(command, task, self._cluster, **kw)) def run(self, command, parallelism=1, **kw): threadpool = ThreadPool(processes=parallelism) for result in threadpool.imap_unordered( self.execute, self.process_arguments(command, **kw)): print(result)
def query(args, options): """usage: query [--force] [--listformat=FORMAT] [--shards=N[,N,...]] [--states=State[,State,...]] cluster [role [job]] Query Mesos about jobs and tasks. """ def _convert_fmt_string(fmtstr): import re def convert(match): return "%%(%s)s" % match.group(1) return re.sub(r"%(\w+)%", convert, fmtstr) def flatten_task(t, d={}): for key in t.__dict__.keys(): val = getattr(t, key) try: val.__dict__.keys() except AttributeError: d[key] = val else: flatten_task(val, d) return d def map_values(d): default_value = lambda v: v mapping = {"status": lambda v: ScheduleStatus._VALUES_TO_NAMES[v]} return dict((k, mapping.get(k, default_value)(v)) for (k, v) in d.items()) for state in options.states.split(","): if state not in ScheduleStatus._NAMES_TO_VALUES: msg = "Unknown state '%s' specified. Valid states are:\n" % state msg += ",".join(ScheduleStatus._NAMES_TO_VALUES.keys()) die(msg) # Role, Job, Instances, States, and the listformat if len(args) == 0: die("Must specify at least cluster.") cluster = args[0] role = args[1] if len(args) > 1 else None job = args[2] if len(args) > 2 else None instances = set(map(int, options.shards.split(","))) if options.shards else set() if options.states: states = set(map(ScheduleStatus._NAMES_TO_VALUES.get, options.states.split(","))) else: states = ACTIVE_STATES | TERMINAL_STATES listformat = _convert_fmt_string(options.listformat) # Figure out "expensive" queries here and bone if they do not have --force # - Does not specify role if not role and not options.force: die("--force is required for expensive queries (no role specified)") # - Does not specify job if not job and not options.force: die("--force is required for expensive queries (no job specified)") # - Specifies status outside of ACTIVE_STATES if not (states <= ACTIVE_STATES) and not options.force: die("--force is required for expensive queries (states outside ACTIVE states") api = AuroraClientAPI(CLUSTERS[cluster], options.verbosity) query_info = api.query(api.build_query(role, job, instances=instances, statuses=states)) if query_info.responseCode != ResponseCode.OK: die("Failed to query scheduler: %s" % query_info.messageDEPRECATED) tasks = query_info.result.scheduleStatusResult.tasks if tasks is None: return try: for task in tasks: d = flatten_task(task) print(listformat % map_values(d)) except KeyError: msg = "Unknown key in format string. Valid keys are:\n" msg += ",".join(d.keys()) die(msg)
def sla_list_safe_domain(cluster, percentage, duration): """usage: sla_list_safe_domain [--exclude_hosts=filename] [--include_hosts=filename] [--list_jobs] [--override_jobs=filename] cluster percentage duration Returns a list of relevant hosts where it would be safe to kill tasks without violating their job SLA. The SLA is defined as a pair of percentage and duration, where: percentage - Percentage of tasks required to be up within the duration. Applied to all jobs except those listed in --override_jobs file; duration - Time interval (now - value) for the percentage of up tasks. Applied to all jobs except those listed in --override_jobs file. Format: XdYhZmWs (each field is optional but must be in that order.) Examples: 5m, 1d3h45m. """ def parse_jobs_file(filename): result = {} with open(filename, "r") as overrides: for line in overrides: if not line.strip(): continue tokens = line.split() if len(tokens) != 3: die("Invalid line in %s:%s" % (filename, line)) job_key = AuroraJobKey.from_path(tokens[0]) result[job_key] = DomainUpTimeSlaVector.JobUpTimeLimit( job=job_key, percentage=parse_sla_percentage(tokens[1]), duration_secs=parse_time(tokens[2]).as_(Time.SECONDS), ) return result options = app.get_options() sla_percentage = parse_sla_percentage(percentage) sla_duration = parse_time(duration) exclude_hosts = parse_hosts_optional(options.exclude_hosts, options.exclude_filename) include_hosts = parse_hosts_optional(options.include_hosts, options.include_filename) override_jobs = parse_jobs_file(options.override_filename) if options.override_filename else {} vector = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).sla_get_safe_domain_vector(include_hosts) hosts = vector.get_safe_hosts(sla_percentage, sla_duration.as_(Time.SECONDS), override_jobs) results = [] for host in sorted(hosts.keys()): if exclude_hosts and host in exclude_hosts: continue if options.list_jobs: results.append( "\n".join( [ "%s\t%s\t%.2f\t%d" % (host, d.job.to_path(), d.percentage, d.duration_secs) for d in sorted(hosts[host]) ] ) ) else: results.append("%s" % host) print_results(results)