|
Brett /
TorqueSubmitFilter(:source lang=python:)
import datetime import sys import os import pwd import grp import getopt import commands import xmlrpclib url = 'https://help.vlsci.unimelb.edu.au/accounts/xmlrpc/' banned_users = ("nobody") help_email = "[email protected]" def get_project(username, requested_project=''): Check if a user is in requested_project or if they have some default project to use
# we may not find a suitable project for the user
result_project = None
# special case for the root user, XXX should this be an error?
if username == "root":
return "FakeRootProject"
try:
# There are several things that could happen here:
# - karaage server can be contacted:
# - user is known to karaage:
# - requested_project != '':
# - user is in requested_project: result_project is requested_project
# - user is not in requested_project:
# - user is in another (default) project, result_project is whatever the default project was
# - user is not in another (default) project, result_project is 'None' (string).
# - requested_project == '':
# - user is in a (default) project, result_project is whatever the default project was
# - user is not in a (default) project, result_project is 'None' (string)
# - user is unknown, result_project is "User '$username' not found"
# - karaage server cannot be contacted, exception is raised
# (eg xmlrpclib.ProtocolError, but possibly others).
# result_project remains None, and we later compare with unix groups
# to see if we can find a project there.
server = xmlrpclib.Server(url)
result_project = server.get_project(username, requested_project)
except:
# there was an error contacting the karaage server
# check if the user is in a unix group with the same name
# as the input requested_project
groups = os.getgroups()
for group in groups:
group_name = grp.getgrgid(group).gr_name
if group_name == requested_project:
result_project = requested_project
break
# user was known to karaage, but not in any projects at all.
if result_project == "None":
print >> sys.stderr, "ERROR: You are not a member of any projects.\nPlease apply to join a project via the VLSCI website."
sys.exit(-1)
# requested_project was specified but karaage does not think user is in that project.
if requested_project != '' and result_project != requested_project:
print >> sys.stderr, "ERROR: You are not a member of project %s.\nUse 'mybalance' to see the projects you are a member of." % (requested_project)
sys.exit(-1)
# this will be None if we couldn't find a project for the user
return result_project
def check_shell(shell): shells = open('/etc/shells', 'r')
valid_shells = shells.readlines()
valid_shells = [x.replace('\n', '') for x in valid_shells]
shells.close()
if shell in valid_shells:
if shell == "/bin/sh":
return "/bin/bash"
return shell
# allow them to run a non-default shell, if it exists and is executable
if os.path.isfile(shell):
if os.access(shell, os.X_OK):
return shell
where = shell.rfind('/')
last_part = shell[where:].rstrip()
for line in valid_shells:
if last_part in line:
return line
return False
if __name__ == "__main__": username = pwd.getpwuid(os.getuid())[0]
interactive = False
queue = False
requests_procs_or_nodes = False
requests_pvmem = False
have_hit_project_already = False
project_line = 0
if username in banned_users:
print >> sys.stderr, "User %s has been blocked from submitting jobs\nIf you feel this is not correct, please email (username, help_email)
sys.exit(-1)
if len(username) > 30:
print >> sys.stderr, "Apologies, but your username generates a torque bug, so we have to stop your jobs for now\nPlease email (help_email)
sys.exit(-1)
#Parse Opts
opts, args = getopt.getopt(sys.argv[1:], 'a:A:b:c:C:d:D:e:hIj:k:l:m:M:N:o:p:q:r:S:t:T:u:xXW:v:Vz')
opts = dict(opts)
if '-I' in opts:
interactive = True
if '-A' in opts:
project = opts['-A']
new_project = get_project(username, project)
if new_project != project:
print >> sys.stderr, "ERROR: You are not a member of project (project)
sys.exit(-1)
if '-q' in opts:
queue = opts['-q']
if '-l' in opts:
if 'nodes=' in opts['-l'] or 'procs=' in opts['-l']:
requests_procs_or_nodes = True
if 'pvmem=' in opts['-l']:
requests_pvmem = True
pbs_in = sys.stdin.readlines()
if len(pbs_in) == 0:
print >> sys.stderr, "ERROR: pbs script is empty"
sys.exit(-1)
first_line = pbs_in[0]
error_and_die = False
newfirstline = ""
insert_first_line = False
if first_line.startswith('#!'):
space = first_line.find(' ')
if space == 2:
nextspace = first_line.find(' ', space + 1)
shell = first_line[space + 1:nextspace].replace('\r', '')
else:
shell = first_line[2:space].replace('\r', '')
newshell = check_shell(shell)
if newshell == False:
error_and_die = True
error = "The shell you specified in your PBS script (%s) does not exist on this cluster.\nYou may wish to email %s to get it installed, or consider using one of the more commonly available shells such as bash.\nUnfortunately, we can't run your job\n" % (shell, help_email)
elif newshell == "/bin/bash" and shell == "/bin/sh":
newfirstline = "#!" + newshell + "\n"
print >> sys.stderr, "Warning: the shell you specified in your PBS script ( (shell)
print >> sys.stderr, "We recommend you use %s instead. Please update your script with the new shell\n" % (newshell)
elif newshell != shell:
newfirstline = "#!" + newshell + "\n"
print >> sys.stderr, "Warning: the shell you specified in your PBS script ( (shell)
print >> sys.stderr, "It can be found instead at (newshell)
print >> sys.stderr, "Please update your script with the correct location\n"
else:
if not interactive:
newfirstline = "#!/bin/bash\n"
insert_first_line = True
print >> sys.stderr, "Warning: you did not specify a shell in the first line of your PBS script"
print >> sys.stderr, "We have assumed you wish to use bash, however please update your script with a valid shell\n"
i = 0
# iterate over rest of lines
for line in pbs_in[1:]:
i += 1
# check if the previous line ends with a backslash and then an enter
# if so, ignore this line
if ord(pbs_in[i - 1][len(pbs_in[i - 1]) - 1]) == 10 and ord(pbs_in[i - 1][len(pbs_in[i - 1]) - 2]) == 92:
continue
# strip out lines with 'please replace' email addresses that users haven't replaced
if line.startswith('#PBS -M YourEmail'):
pbs_in.remove(line)
i -= 1
continue
if line.startswith('#PBS -M [email protected]'):
pbs_in.remove(line)
i -= 1
continue
if line.startswith('#PBS -A'):
project_id = line.split(' ')[2].strip()
project = get_project(username, project_id)
# If the project is None then we couldn't find a project for the user, so we
# leave the project line blank and assume their job will use their default
# project.
if project != None:
pbs_in[i] = "#PBS -A %s\n" % (project)
have_hit_project_already = True
if line.startswith('#PBS -q'):
queue = line.split(' ')[2].strip()
if line.startswith('#PBS -l'):
if line.count('nodes=') + line.count('procs=') > 0:
requests_procs_or_nodes = True
if line.count('pvmem=') > 0:
requests_pvmem = True
if line.startswith('#PBS -S'):
shell = line.split(' ')[2].rstrip()
newshell = check_shell(shell)
if newshell == False:
error_and_die = True
error = "The shell you specified in your PBS script (%s) does not exist on this cluster.\nYou may wish to email %s to get it installed, or consider using one of the more commonly available shells such as bash.\nUnfortunately, we can't run your job\n" % (shell, help_email)
elif newshell == "/bin/bash" and shell == "/bin/sh":
newfirstline = "#!" + newshell + "\n"
print >> sys.stderr, "Warning: the shell you specified in your PBS script ( (shell)
print >> sys.stderr, "We recommend you use %s instead. Please update your script with the new shell\n" % (newshell)
elif newshell != shell:
pbs_in[i] = "#PBS -S " + newshell + "\n"
newfirstline = "#!" + newshell + "\n"
error_and_die = False
print >> sys.stderr, "Warning: the shell you specified in your PBS script ( (shell)
print >> sys.stderr, "It can be found instead at (newshell)
print >> sys.stderr, "Please update your script with the correct location\n"
else:
newfirstline = "#!" + newshell + "\n"
error_and_die = False
if (not line.startswith('#')) and line != "\n" and line != "\r" and line != "\r\n" and line != "\n\r":
if not have_hit_project_already:
project = get_project(username)
project_line = i
have_hit_project_already = True
# If the project is None then we couldn't find a project for the user, so we
# leave the project line blank and assume their job will use their default
# project.
if project_line > 0 and project != None:
pbs_in.insert(project_line, "#PBS -A %s\n" % (project))
if newfirstline != "":
if insert_first_line:
pbs_in.insert(0, newfirstline)
else:
pbs_in[0] = newfirstline
if error_and_die == True:
print >> sys.stderr, error
sys.exit(-1)
if queue == "smp" and (requests_procs_or_nodes or requests_pvmem):
print >> sys.stderr, "ERROR: If using the smp queue, you cannot request procs, nodes or pvmem\nPlease read documentation at: http://www.vlsci.org.au/node/119"
sys.exit(-1)
sys.stdout.writelines(pbs_in)
(:sourceend:) |