<commit_before>import os
import pwd
import threading
import time

from twitter.common import log
from twitter.common.concurrent import defer
from twitter.common.process import ProcessProviderFactory
from twitter.common.quantity import Amount, Time, Data
from twitter.thermos.base.ckpt import CheckpointDispatcher
from twitter.thermos.base.path import TaskPath
from twitter.thermos.runner.inspector import CheckpointInspector
from twitter.thermos.runner.helper import TaskKiller
from twitter.thermos.monitoring.detector import TaskDetector
from twitter.thermos.monitoring.garbage import TaskGarbageCollector

from gen.twitter.mesos.comm.ttypes import (
    AdjustRetainedTasks,
    DeletedTasks,
    SchedulerMessage)
from gen.twitter.mesos.ttypes import ScheduleStatus

from .sandbox_manager import (
    DirectorySandbox,
    AppAppSandbox)
from .executor_base import ThermosExecutorBase
from .executor_detector import ExecutorDetector

from thrift.TSerialization import deserialize as thrift_deserialize
from thrift.TSerialization import serialize as thrift_serialize


class ThermosGCExecutor(ThermosExecutorBase):
  """
    Thermos GC Executor, responsible for:
      - garbage collecting old tasks to make sure they don't clutter up the system
      - state reconciliation with the scheduler (in case it thinks we're running
        something we're not or vice versa.)
  """
  MAX_PID_TIME_DRIFT = Amount(10, Time.SECONDS)
  MAX_CHECKPOINT_TIME_DRIFT = Amount(1, Time.HOURS)  # maximum runner disconnection time
  PERSISTENCE_WAIT = Amount(5, Time.SECONDS)

  # how old a task must be before we're willing to kill it, assuming that there could be
  # slight races in the following scenario:
  #    launch gc with retained_tasks={t1, t2, t3}
  #    launch task t4
  MINIMUM_KILL_AGE = Amount(10, Time.MINUTES)

  def __init__(self, checkpoint_root, verbose=True, task_killer=TaskKiller, clock=time):
    ThermosExecutorBase.__init__(self)
    self._slave_id = None
    self._detector = ExecutorDetector()
    self._collector = TaskGarbageCollector(root=checkpoint_root)
    self._clock = clock
    self._start_time = clock.time()
    self._task_killer = task_killer
    self._checkpoint_root = checkpoint_root
    if 'ANGRYBIRD_THERMOS' in os.environ:
      self._checkpoint_root = os.path.join(os.environ['ANGRYBIRD_THERMOS'], 'thermos', 'run')

  def terminate_task(self, task_id, kill=True):
    """Terminate a task given its task_id."""
    killer = self._task_killer(task_id, self._checkpoint_root)
    self.log('Terminating %s...' % task_id)
    runner_terminate = killer.kill if kill else killer.lose
    try:
      runner_terminate(force=True)
      return True
    except Exception as e:
      self.log('Could not terminate: %s' % e)
      return False

  def partition_tasks(self):
    """Return active/finished tasks as discovered from the checkpoint root."""
    detector = TaskDetector(root=self._checkpoint_root)
    active_tasks = set(t_id for _, t_id in detector.get_task_ids(state='active'))
    finished_tasks = set(t_id for _, t_id in detector.get_task_ids(state='finished'))
    return active_tasks, finished_tasks

  def get_states(self, task_id):
    """Returns the (timestamp, status) tuples of the task or [] if could not replay."""
    pathspec = TaskPath(root=self._checkpoint_root, task_id=task_id)
    runner_ckpt = pathspec.getpath('runner_checkpoint')
    statuses = CheckpointDispatcher.iter_statuses(runner_ckpt)
    return [(state.timestamp_ms / 1000.0, state.state) for state in statuses]

  def maybe_terminate_unknown_task(self, task_id):
    """Terminate a task if we believe the scheduler doesn't know about it.

       It's possible for the scheduler to queue a GC and launch a task afterwards, in which
       case we may see actively running tasks that the scheduler did not report in the
       AdjustRetainedTasks message.
    """
    states = self.get_states(task_id)
    if len(states) > 0:
      task_start_time, _ = states[0]
      if self._start_time - task_start_time > self.MINIMUM_KILL_AGE.as_(Time.SECONDS):
        self.log('Terminating task %s' % task_id)
        self.terminate_task(task_id)

  def reconcile_states(self, driver, retained_tasks):
    """Reconcile states that the scheduler thinks tasks are in vs what they really are in.

       Local      vs   Scheduler  => Action

        ACTIVE          ACTIVE   => no-op
        ACTIVE         TERMINAL  => maybe kill task*
        ACTIVE         !EXISTS   => maybe kill task*
       TERMINAL         ACTIVE   => send actual status**
       TERMINAL        TERMINAL  => no-op
       TERMINAL        !EXISTS   => gc locally
       !EXISTS          ACTIVE   => send LOST**
       !EXISTS         TERMINAL  => gc remotely

       * - Only kill if this does not appear to be a race condition.
       ** - These appear to have no effect

       Side effecting operations:
         ACTIVE   | (TERMINAL / !EXISTS) => maybe kill
         TERMINAL | !EXISTS              => delete
         !EXISTS  | TERMINAL             => delete
    """
    def partition(rt):
      active, finished = set(), set()
      for task_id, schedule_status in rt.items():
        if self.twitter_status_is_terminal(schedule_status):
          finished.add(task_id)
        else:
          active.add(task_id)
      return active, finished

    local_active, local_finished = self.partition_tasks()
    sched_active, sched_finished = partition(retained_tasks)
    local_task_ids = local_active | local_finished
    sched_task_ids = sched_active | sched_finished
    all_task_ids = local_task_ids | sched_task_ids

    self.log('Told to retain the following task ids:')
    for task_id, schedule_status in retained_tasks.items():
      self.log('  => %s as %s' % (task_id, ScheduleStatus._VALUES_TO_NAMES[schedule_status]))

    self.log('Local active tasks:')
    for task_id in local_active:
      self.log('  => %s' % task_id)

    self.log('Local finished tasks:')
    for task_id in local_finished:
      self.log('  => %s' % task_id)

    local_gc, remote_gc = set(), set()
    updates = {}

    for task_id in all_task_ids:
      if task_id in local_active and task_id not in sched_active:
        self.log('Inspecting %s for termination.' % task_id)
        self.maybe_terminate_unknown_task(task_id)
      if task_id in local_finished and task_id not in sched_task_ids:
        self.log('Queueing task_id %s for local deletion.' % task_id)
        local_gc.add(task_id)
      if task_id in local_finished and task_id in sched_active:
        self.log('Task %s finished but scheduler thinks active.' % task_id)
        states = self.get_state(task_id)
        if len(states) > 0:
          _, last_state = states[-1]
          updates[task_id] = self.THERMOS_TO_TWITTER_STATES.get(last_state, ScheduleStatus.UNKNOWN)
          self.send_update(driver, task_id, last_state,
              'Telling scheduler of finished task %s' % task_id)
        else:
          self.log('Task state unavailable for finished task!  Possible checkpoint corruption.')
      if task_id in sched_finished and task_id not in local_task_ids:
        self.log('Queueing task_id %s for remote deletion.' % task_id)
        remote_gc.add(task_id)<commit_msg>Fix bad self.states() call in gc executor causing exception<commit_after>import os
import pwd
import threading
import time

from twitter.common import log
from twitter.common.concurrent import defer
from twitter.common.process import ProcessProviderFactory
from twitter.common.quantity import Amount, Time, Data
from twitter.thermos.base.ckpt import CheckpointDispatcher
from twitter.thermos.base.path import TaskPath
from twitter.thermos.runner.inspector import CheckpointInspector
from twitter.thermos.runner.helper import TaskKiller
from twitter.thermos.monitoring.detector import TaskDetector
from twitter.thermos.monitoring.garbage import TaskGarbageCollector

from gen.twitter.mesos.comm.ttypes import (
    AdjustRetainedTasks,
    DeletedTasks,
    SchedulerMessage)
from gen.twitter.mesos.ttypes import ScheduleStatus

from .sandbox_manager import (
    DirectorySandbox,
    AppAppSandbox)
from .executor_base import ThermosExecutorBase
from .executor_detector import ExecutorDetector

from thrift.TSerialization import deserialize as thrift_deserialize
from thrift.TSerialization import serialize as thrift_serialize


class ThermosGCExecutor(ThermosExecutorBase):
  """
    Thermos GC Executor, responsible for:
      - garbage collecting old tasks to make sure they don't clutter up the system
      - state reconciliation with the scheduler (in case it thinks we're running
        something we're not or vice versa.)
  """
  MAX_PID_TIME_DRIFT = Amount(10, Time.SECONDS)
  MAX_CHECKPOINT_TIME_DRIFT = Amount(1, Time.HOURS)  # maximum runner disconnection time
  PERSISTENCE_WAIT = Amount(5, Time.SECONDS)

  # how old a task must be before we're willing to kill it, assuming that there could be
  # slight races in the following scenario:
  #    launch gc with retained_tasks={t1, t2, t3}
  #    launch task t4
  MINIMUM_KILL_AGE = Amount(10, Time.MINUTES)

  def __init__(self, checkpoint_root, verbose=True, task_killer=TaskKiller, clock=time):
    ThermosExecutorBase.__init__(self)
    self._slave_id = None
    self._detector = ExecutorDetector()
    self._collector = TaskGarbageCollector(root=checkpoint_root)
    self._clock = clock
    self._start_time = clock.time()
    self._task_killer = task_killer
    self._checkpoint_root = checkpoint_root
    if 'ANGRYBIRD_THERMOS' in os.environ:
      self._checkpoint_root = os.path.join(os.environ['ANGRYBIRD_THERMOS'], 'thermos', 'run')

  def terminate_task(self, task_id, kill=True):
    """Terminate a task given its task_id."""
    killer = self._task_killer(task_id, self._checkpoint_root)
    self.log('Terminating %s...' % task_id)
    runner_terminate = killer.kill if kill else killer.lose
    try:
      runner_terminate(force=True)
      return True
    except Exception as e:
      self.log('Could not terminate: %s' % e)
      return False

  def partition_tasks(self):
    """Return active/finished tasks as discovered from the checkpoint root."""
    detector = TaskDetector(root=self._checkpoint_root)
    active_tasks = set(t_id for _, t_id in detector.get_task_ids(state='active'))
    finished_tasks = set(t_id for _, t_id in detector.get_task_ids(state='finished'))
    return active_tasks, finished_tasks

  def get_states(self, task_id):
    """Returns the (timestamp, status) tuples of the task or [] if could not replay."""
    pathspec = TaskPath(root=self._checkpoint_root, task_id=task_id)
    runner_ckpt = pathspec.getpath('runner_checkpoint')
    statuses = CheckpointDispatcher.iter_statuses(runner_ckpt)
    return [(state.timestamp_ms / 1000.0, state.state) for state in statuses]

  def maybe_terminate_unknown_task(self, task_id):
    """Terminate a task if we believe the scheduler doesn't know about it.

       It's possible for the scheduler to queue a GC and launch a task afterwards, in which
       case we may see actively running tasks that the scheduler did not report in the
       AdjustRetainedTasks message.
    """
    states = self.get_states(task_id)
    if len(states) > 0:
      task_start_time, _ = states[0]
      if self._start_time - task_start_time > self.MINIMUM_KILL_AGE.as_(Time.SECONDS):
        self.log('Terminating task %s' % task_id)
        self.terminate_task(task_id)

  def reconcile_states(self, driver, retained_tasks):
    """Reconcile states that the scheduler thinks tasks are in vs what they really are in.

       Local      vs   Scheduler  => Action

        ACTIVE          ACTIVE   => no-op
        ACTIVE         TERMINAL  => maybe kill task*
        ACTIVE         !EXISTS   => maybe kill task*
       TERMINAL         ACTIVE   => send actual status**
       TERMINAL        TERMINAL  => no-op
       TERMINAL        !EXISTS   => gc locally
       !EXISTS          ACTIVE   => send LOST**
       !EXISTS         TERMINAL  => gc remotely

       * - Only kill if this does not appear to be a race condition.
       ** - These appear to have no effect

       Side effecting operations:
         ACTIVE   | (TERMINAL / !EXISTS) => maybe kill
         TERMINAL | !EXISTS              => delete
         !EXISTS  | TERMINAL             => delete
    """
    def partition(rt):
      active, finished = set(), set()
      for task_id, schedule_status in rt.items():
        if self.twitter_status_is_terminal(schedule_status):
          finished.add(task_id)
        else:
          active.add(task_id)
      return active, finished

    local_active, local_finished = self.partition_tasks()
    sched_active, sched_finished = partition(retained_tasks)
    local_task_ids = local_active | local_finished
    sched_task_ids = sched_active | sched_finished
    all_task_ids = local_task_ids | sched_task_ids

    self.log('Told to retain the following task ids:')
    for task_id, schedule_status in retained_tasks.items():
      self.log('  => %s as %s' % (task_id, ScheduleStatus._VALUES_TO_NAMES[schedule_status]))

    self.log('Local active tasks:')
    for task_id in local_active:
      self.log('  => %s' % task_id)

    self.log('Local finished tasks:')
    for task_id in local_finished:
      self.log('  => %s' % task_id)

    local_gc, remote_gc = set(), set()
    updates = {}

    for task_id in all_task_ids:
      if task_id in local_active and task_id not in sched_active:
        self.log('Inspecting %s for termination.' % task_id)
        self.maybe_terminate_unknown_task(task_id)
      if task_id in local_finished and task_id not in sched_task_ids:
        self.log('Queueing task_id %s for local deletion.' % task_id)
        local_gc.add(task_id)
      if task_id in local_finished and task_id in sched_active:
        self.log('Task %s finished but scheduler thinks active.' % task_id)
        states = self.get_states(task_id)
        if len(states) > 0:
          _, last_state = states[-1]
          updates[task_id] = self.THERMOS_TO_TWITTER_STATES.get(last_state, ScheduleStatus.UNKNOWN)
          self.send_update(driver, task_id, last_state,
              'Telling scheduler of finished task %s' % task_id)
        else:
          self.log('Task state unavailable for finished task!  Possible checkpoint corruption.')
      if task_id in sched_finished and task_id not in local_task_ids:
        self.log('Queueing task_id %s for remote deletion.' % task_id)
        remote_gc.add(task_id)
{'fraction_non_alphanumeric': 0.08134961705651005, 'fraction_numerical': 0.002621955426757745, 'mean_word_length': 3.308561236623068, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5541391', 'n_tokens_mistral': 4522, 'n_tokens_neox': 4342, 'n_words': 1253}