#!/usr/bin/python
#
#  prologue.user script for Machine/Job Features on Torque/PBS
#
#  Andrew McNab, University of Manchester.
#  Copyright (c) 2016. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or
#  without modification, are permitted provided that the following
#  conditions are met:
#
#    o Redistributions of source code must retain the above
#      copyright notice, this list of conditions and the following
#      disclaimer. 
#    o Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials
#      provided with the distribution. 
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
#  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
#  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
#  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
#  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
#  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#

import os
import re
import sys
import time

#
# PROGRAM MAIN !!!
# 


homdeDir = os.path.expanduser('~')

job_id = os.environ['PBS_JOBID']

jobfeaturesDir = homeDir + '/jobfeatures-' + jobID

os.mkdir(jobfeaturesDir)

open(jobfeaturesDir + '/job_id', 'w').write(job_id)

try:
  allocated_cpu = int(os.environ['PBS_NUM_PPN'])
except:
  allocated_cpu = 1
  
open(jobfeaturesDir + '/job_id', 'w').write(str(allocated_cpu))

try:
  hs06 = float(open('/etc/machinefeatures/hs06','r').readline())
except:
  hs06 = None

try:
  total_cpu = int(open('/etc/machinefeatures/total_cpu','r').readline())
except:
  total_cpu = None
  
if hs06 and total_cpu:
  # Simple pro-rata allocation of total hs06 depending on processors for this job
  hs06_job = (allocated_cpu * hs06) / total_cpu
  open(jobfeaturesDir + '/hs06_job', 'w').write(str(hs06_job))

try: 
  qstatResult = os.popen('qstat -f ' + jobID,'r').read()
except:
  pass

jobstart_secs = int(time.time())
open(jobfeaturesDir + '/jobstart_secs', 'w').write(str(jobstart_secs))