python/cpulimit.py

   1 #!/usr/bin/python2 -u
   2
   3 import sys, os, re, string
   4
   5
   6 TCBASE="/rcfs/taskclass/"
   7 RULES="/rcfs/ce/rules/"
   8
   9 SYSCLASS=TCBASE + "system"
  10 SYSRULE=RULES + "system"
  11 SYSCPUSHARE=100
  12 DEFAULTMAXCPUSHARE=8192
  13
  14 def checkckrm():
  15     checks = [ "/rcfs", TCBASE, RULES ]
  16
  17     for check in checks:
  18         try:
  19             answer = os.stat(check)
  20         except:
  21             print "%s does not exist" % check
  22             return False
  23
  24     return True
  25
  26 def checkclass(tc):
  27     try:
  28         answer = os.stat(TCBASE + tc)
  29         return True
  30
  31     except:
  32         print "%s class does not exist" % tc
  33         return False
  34
  35 def getxid(name):
  36     xid = -1
  37     fp = open('/etc/passwd')
  38     for line in fp.readlines():
  39         rec = string.splitfields(line,':')
  40         if rec[0] == name:
  41             xid = int(rec[2])
  42             break
  43
  44     fp.close()
  45
  46     if xid == -1:
  47         # raise an exception
  48         pass
  49
  50     return xid
  51
  52 def write(filename,s):
  53     fp = os.open(filename,os.O_WRONLY|os.O_CREAT)
  54     os.write(fp,s)
  55     os.close(fp)
  56
  57 def vs2ckrm_on(tc):
  58     xid = getxid(tc)
  59
  60     try:
  61         os.mkdir(TCBASE + tc)
  62     except OSError:
  63         pass # ignore oserror for file exists
  64
  65     s = "xid=%d,class=%s" % (xid,TCBASE+tc)
  66     fname = RULES + tc
  67     write(fname, s)
  68
  69 def vs2ckrm_off(tc):
  70     fname = TCBASE + tc + "/members"
  71     for i in range(1,15):
  72         fp = open(fname)
  73         lines = fp.readlines()
  74         try:
  75             lines.remove("No data to display\n")
  76         except ValueError:
  77             pass
  78         if len(lines) == 0:
  79             try:
  80                 answer = os.stat(RULES + tc)
  81                 os.unlink(RULES + tc)
  82                 answer = os.stat(TCBASE + tc)
  83                 os.rmdir(TCBASE + tc)
  84             except:
  85                 pass
  86             break
  87
  88         else:
  89             print "enter context 1 and kill processes", lines
  90
  91
  92 def cpulimit(tc,limit):
  93     global TCBASE
  94
  95     fname = TCBASE + tc + "/shares"
  96     s = "res=cpu,guarantee=%d\n" % limit
  97     write(fname,s)
  98
  99 def cpuinit():
 100     global TCBASE
 101
 102     fname = TCBASE + "shares"
 103     s = "res=cpu,total_guarantee=%d\n" % DEFAULTMAXCPUSHARE
 104     write(fname,s)
 105
 106 if __name__ == "__main__":
 107     try:
 108         name = sys.argv[1]
 109         limit = int(sys.argv[2])
 110     except:
 111         print "caught exception"
 112
 113     if checkckrm() is True:
 114         cpuinit()
 115         vs2ckrm_on(name)
 116         cpulimit(name,limit)
 117         vs2ckrm_off(name)