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