src/nepi/testbeds/planetlab/rspawn.py

   1 # Utility library for spawning remote asynchronous tasks
   2 from nepi.util import server
   3 import getpass
   4
   5 class STDOUT:
   6     """
   7     Special value that when given to remote_spawn in stderr causes stderr to
   8     redirect to whatever stdout was redirected to.
   9     """
  10
  11 class RUNNING:
  12     """
  13     Process is still running
  14     """
  15
  16 class FINISHED:
  17     """
  18     Process is finished
  19     """
  20
  21 class NOT_STARTED:
  22     """
  23     Process hasn't started running yet (this should be very rare)
  24     """
  25
  26 def remote_spawn(command, pidfile, stdout='/dev/null', stderr=STDOUT, stdin='/dev/null', home=None, create_home=False, sudo=False,
  27         host = None, port = None, user = None, agent = None,
  28         ident_key = None, server_key = None,
  29         tty = False):
  30     """
  31     Spawn a remote command such that it will continue working asynchronously.
  32
  33     Parameters:
  34         command: the command to run - it should be a single line.
  35
  36         pidfile: path of a (ideally unique to this task) pidfile for tracking the process.
  37
  38         stdout: path of a file to redirect standard output to - must be a string.
  39             Defaults to /dev/null
  40         stderr: path of a file to redirect standard error to - string or the special STDOUT value
  41             to redirect to the same file stdout was redirected to. Defaults to STDOUT.
  42         stdin: path of a file with input to be piped into the command's standard input
  43
  44         home: path of a folder to use as working directory - should exist, unless you specify create_home
  45
  46         create_home: if True, the home folder will be created first with mkdir -p
  47
  48         sudo: whether the command needs to be executed as root
  49
  50         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
  51
  52     Returns:
  53         (stdout, stderr), process
  54
  55         Of the spawning process, which only captures errors at spawning time.
  56         Usually only useful for diagnostics.
  57     """
  58     # Start process in a "daemonized" way, using nohup and heavy
  59     # stdin/out redirection to avoid connection issues
  60     if stderr is STDOUT:
  61         stderr = '&1'
  62     else:
  63         stderr = ' ' + stderr
  64
  65     daemon_command = '{ { %(command)s  > %(stdout)s 2>%(stderr)s < %(stdin)s & } ; echo $! 1 > %(pidfile)s ; }' % {
  66         'command' : command,
  67         'pidfile' : server.shell_escape(pidfile),
  68
  69         'stdout' : stdout,
  70         'stderr' : stderr,
  71         'stdin' : stdin,
  72     }
  73
  74     cmd = "%(create)s%(gohome)s rm -f %(pidfile)s ; %(sudo)s nohup bash -c %(command)s " % {
  75             'command' : server.shell_escape(daemon_command),
  76
  77             'sudo' : 'sudo -S' if sudo else '',
  78
  79             'pidfile' : server.shell_escape(pidfile),
  80             'gohome' : 'cd %s ; ' % (server.shell_escape(home),) if home else '',
  81             'create' : 'mkdir -p %s ; ' % (server.shell_escape,) if create_home else '',
  82         }
  83     (out,err),proc = server.popen_ssh_command(
  84         cmd,
  85         host = host,
  86         port = port,
  87         user = user,
  88         agent = agent,
  89         ident_key = ident_key,
  90         server_key = server_key,
  91         tty = tty
  92         )
  93
  94     if proc.wait():
  95         raise RuntimeError, "Failed to set up application: %s %s" % (out,err,)
  96
  97     return (out,err),proc
  98
  99 def remote_check_pid(pidfile,
 100         host = None, port = None, user = None, agent = None,
 101         ident_key = None, server_key = None):
 102     """
 103     Check the pidfile of a process spawned with remote_spawn.
 104
 105     Parameters:
 106         pidfile: the pidfile passed to remote_span
 107
 108         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 109
 110     Returns:
 111
 112         A (pid, ppid) tuple useful for calling remote_status and remote_kill,
 113         or None if the pidfile isn't valid yet (maybe the process is still starting).
 114     """
 115
 116     (out,err),proc = server.popen_ssh_command(
 117         "cat %(pidfile)s" % {
 118             'pidfile' : pidfile,
 119         },
 120         host = host,
 121         port = port,
 122         user = user,
 123         agent = agent,
 124         ident_key = ident_key,
 125         server_key = server_key
 126         )
 127
 128     if proc.wait():
 129         return None
 130
 131     if out:
 132         try:
 133             return map(int,out.strip().split(' ',1))
 134         except:
 135             # Ignore, many ways to fail that don't matter that much
 136             return None
 137
 138
 139 def remote_status(pid, ppid,
 140         host = None, port = None, user = None, agent = None,
 141         ident_key = None, server_key = None):
 142     """
 143     Check the status of a process spawned with remote_spawn.
 144
 145     Parameters:
 146         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 147
 148         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 149
 150     Returns:
 151
 152         One of NOT_STARTED, RUNNING, FINISHED
 153     """
 154
 155     (out,err),proc = server.popen_ssh_command(
 156         "ps --ppid %(ppid)d -o pid | grep -c %(pid)d ; true" % {
 157             'ppid' : ppid,
 158             'pid' : pid,
 159         },
 160         host = host,
 161         port = port,
 162         user = user,
 163         agent = agent,
 164         ident_key = ident_key,
 165         server_key = server_key
 166         )
 167
 168     if proc.wait():
 169         return NOT_STARTED
 170
 171     status = False
 172     if out:
 173         try:
 174             status = bool(int(out.strip()))
 175         except:
 176             # Ignore, many ways to fail that don't matter that much
 177             return NOT_STARTED
 178     return RUNNING if status else FINISHED
 179
 180
 181 def remote_kill(pid, ppid, sudo = False,
 182         host = None, port = None, user = None, agent = None,
 183         ident_key = None, server_key = None,
 184         nowait = False):
 185     """
 186     Kill a process spawned with remote_spawn.
 187
 188     First tries a SIGTERM, and if the process does not end in 10 seconds,
 189     it sends a SIGKILL.
 190
 191     Parameters:
 192         pid/ppid: pid and parent-pid of the spawned process. See remote_check_pid
 193
 194         sudo: whether the command was run with sudo - careful killing like this.
 195
 196         host/port/user/agent/ident_key: see nepi.util.server.popen_ssh_command
 197
 198     Returns:
 199
 200         Nothing, should have killed the process
 201     """
 202
 203     cmd = """
 204 %(sudo)s kill %(pid)d
 205 for x in 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 ; do
 206     sleep 0.2
 207     if [ `ps --ppid %(ppid)d -o pid | grep -c %(pid)d` == '0' ]; then
 208         break
 209     fi
 210     sleep 1.8
 211 done
 212 if [ `ps --ppid %(ppid)d -o pid | grep -c %(pid)d` != '0' ]; then
 213     %(sudo)s kill -9 %(pid)d
 214 fi
 215 """
 216     if nowait:
 217         cmd = "{ %s } >/dev/null 2>/dev/null </dev/null &" % (cmd,)
 218
 219     (out,err),proc = server.popen_ssh_command(
 220         cmd % {
 221             'ppid' : ppid,
 222             'pid' : pid,
 223             'sudo' : 'sudo -S' if sudo else ''
 224         },
 225         host = host,
 226         port = port,
 227         user = user,
 228         agent = agent,
 229         ident_key = ident_key,
 230         server_key = server_key
 231         )
 232
 233     # wait, don't leave zombies around
 234     proc.wait()
 235
 236
 237