starcluster.plugins.tmux

Package starcluster :: Package plugins :: Module tmux |
[hide private]
[frames] | no frames]

# Source Code for Module starcluster.plugins.tmux

      1  #!/usr/bin/env python
      2
      3  from starcluster import utils
      4  from starcluster import exception
      5  from starcluster import clustersetup
      6  from starcluster.logger import log
      7
      8
      9 -class TmuxControlCenter(clustersetup.DefaultClusterSetup):
     10      """
     11      Starts a TMUX session on StarCluster configured with split panes for all
     12      nodes. This allows you to interactively run commands on all nodes and see
     13      all the output at once.
     14      """
     15      _layouts = ['even-horizontal', 'even-vertical', 'main-horizontal',
     16                  'main-vertical', 'tiled']
     17
     18 -    def __init__(self, envname="starcluster"):
     19          self._envname = envname
     20          self._nodes = None
     21          self._master = None
     22          self._user = None
     23          self._user_shell = None
     24          self._volumes = None
     25
     26 -    def _supports_layout(self, node, envname, layout, window=''):
     27          if layout not in self._layouts:
     28              raise exception.PluginError("unknown layout (options: %s)" %
     29                                          ", ".join(self._layouts))
     30          return self._select_layout(node, envname, layout, window) == 0
     31
     32 -    def _select_layout(self, node, envname, layout="main-vertical", window=''):
     33          if layout not in self._layouts:
     34              raise exception.PluginError("unknown layout (options: %s)" %
     35                                            ", ".join(self._layouts))
     36          cmd = 'tmux select-layout -t %s:%s %s'
     37          return node.ssh.get_status(cmd % (envname, window, layout))
     38
     39 -    def _resize_pane(self, node, envname, pane, units, up=False):
     40          upordown = '-D %s' % units
     41          if up:
     42              upordown = '-D %s' % units
     43          cmd = 'tmux resize-pane -t %s:%s %s' % (envname, pane, upordown)
     44          return node.ssh.execute(cmd)
     45
     46 -    def _split_window(self, node, envname, window='', vertical=False):
     47          cmd = 'tmux split-window'
     48          if vertical:
     49              cmd += ' -h'
     50          return node.ssh.execute('%s -t %s:%s' % (cmd, envname, window))
     51
     52 -    def _rename_window(self, node, envname, window, name):
     53          cmd = 'tmux rename-window -t %s:%s %s' % (envname, window, name)
     54          return node.ssh.execute(cmd)
     55
     56 -    def _has_session(self, node, envname):
     57          status = node.ssh.get_status('tmux has-session -t %s' % envname)
     58          return status == 0
     59
     60 -    def _send_keys(self, node, envname, cmd, window=''):
     61          node.ssh.execute('tmux send-keys -t %s:%s "%s"' % (envname, window,
     62                                                             cmd))
     63          node.ssh.execute('tmux send-keys -t %s:%s "Enter"' % (envname, window))
     64
     65 -    def _new_session(self, node, envname):
     66          node.ssh.execute('tmux new-session -d -s %s' % envname, detach=True)
     67
     68 -    def _kill_session(self, node, envname):
     69          node.ssh.execute('tmux kill-session -t %s' % envname)
     70
     71 -    def _new_window(self, node, envname, title):
     72          node.ssh.execute('tmux new-window -n %s -t %s:' % (title, envname))
     73
     74 -    def _select_window(self, node, envname, window=''):
     75          node.ssh.execute('tmux select-window -t %s:%s' % (envname, window))
     76
     77 -    def _select_pane(self, node, envname, window, pane):
     78          node.ssh.execute('tmux select-pane -t %s:%s.%s' %
     79                           (envname, window, pane))
     80
     81 -    def create_session(self, node, envname, num_windows=5):
     82          if not self._has_session(node, envname):
     83              self._new_session(node, envname)
     84          for i in range(1, num_windows):
     85              self._new_window(node, envname, i)
     86
     87 -    def setup_tmuxcc(self, client=None, nodes=None, user='root',
     88                       layout='tiled'):
     89          log.info("Creating TMUX Control Center for user '%s'" % user)
     90          client = client or self._master
     91          nodes = nodes or self._nodes
     92          envname = self._envname
     93          orig_user = client.ssh._username
     94          if orig_user != user:
     95              client.ssh.connect(username=user)
     96          chunks = [chunk for chunk in utils.chunk_list(nodes, items=8)]
     97          num_windows = len(chunks) + len(nodes)
     98          if len(nodes) == 0:
     99              log.error("Cluster has no nodes, exiting...")
    100              return
    101          self.create_session(client, envname, num_windows=num_windows)
    102          if len(nodes) == 1 and client == nodes[0]:
    103              return
    104          if not self._supports_layout(client, envname, layout, window=0):
    105              log.warn("failed to select layout '%s', defaulting to "
    106                       "'main-vertical'" % layout)
    107              layout = "main-vertical"
    108              status = self._select_layout(client, envname, layout, window=0)
    109              if status != 0:
    110                  raise exception.PluginError("failed to set a layout")
    111          for i, chunk in enumerate(chunks):
    112              self._rename_window(client, envname, i, 'all%s' % i)
    113              for j, node in enumerate(chunk):
    114                  if j != 0:
    115                      self._split_window(client, envname, i)
    116                  self._select_layout(client, envname, window=i, layout=layout)
    117                  if node.alias != client.alias:
    118                      self._send_keys(client, envname, cmd='ssh %s' % node.alias,
    119                                      window="%d.%d" % (i, j))
    120          for i, node in enumerate(nodes):
    121              window = i + len(chunks)
    122              self._rename_window(client, envname, window, node.alias)
    123              if node.alias != client.alias:
    124                  self._send_keys(client, envname, cmd='ssh %s' % node.alias,
    125                                  window=window)
    126          self._select_window(client, envname, window=0)
    127          self._select_pane(client, envname, window=0, pane=0)
    128          if orig_user != user:
    129              client.ssh.connect(username=orig_user)
    130
    131 -    def add_to_utmp_group(self, client, user):
    132          """
    133          Adds user (if exists) to 'utmp' group (if exists)
    134          """
    135          try:
    136              client.add_user_to_group(user, 'utmp')
    137          except exception.BaseException:
    138              pass
    139
    140 -    def run(self, nodes, master, user, user_shell, volumes):
    141          log.info("Starting TMUX Control Center...")
    142          self._nodes = nodes
    143          self._master = master
    144          self._user = user
    145          self._user_shell = user_shell
    146          self._volumes = volumes
    147          self.add_to_utmp_group(master, user)
    148          self.setup_tmuxcc(user=user)
    149          self.setup_tmuxcc(user='root')
    150
    151 -    def on_add_node(self, node, nodes, master, user, user_shell, volumes):
    152          log.info("Adding %s to TMUX Control Center" % node.alias)
    153          #user_home = node.getpwnam(user).pw_dir
    154
    155 -    def on_remove_node(self, node, nodes, master, user, user_shell, volumes):
    156          log.info("Removing %s from TMUX Control Center" % node.alias)
    157