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