#!/usr/bin/env python
#
# This takes pipeline and cluster information and calculates
# the number of instances needed based on the computed number of
# cpu hours needed to complete the task.
#
# Required keys:
# cluster.CLUSTER_NAME
# pipeline.COMPUTED_CPU_HOURS
#
# Optional keys:
# pipeline.DESIRED_RUN_TIME -- Suggestions for the number of hours the pipeline should run in
#                              The default is to run as quickly as possible
# pipeline.EXEC_MAX_INSTANCES -- A cluster has a maximum value associated with it, this overrides
#                                that value
import sys

from vappio.webservice import cluster

CORES_PER_CPU = 8

# Ensure we are dealing with kv pairs
if sys.stdin.readline().strip() != 'kv':
    raise Exception('Header needs to be kv')

# Write our header out
sys.stdout.write('kv\n')

requiredKeys = ['cluster.CLUSTER_NAME',
                'pipeline.COMPUTED_CPU_HOURS'
                ]

optionalKeys = ['pipeline.DESIRED_RUN_TIME',
                'pipeline.EXEC_MAX_INSTANCES'
                ]

allKeys = set(requiredKeys + optionalKeys)
kv = {}

for line in sys.stdin:
    sys.stdout.write(line.strip() + '\n')
    k, v = line.strip().split('=', 1)
    if k in allKeys:
        kv[k] = v

if set(requiredKeys) - set(kv.keys()):
    raise Exception('Not all required keys specified')

#
# We are adding 1 here because almost every calculationg is going to be
# hours + some fraction and we always get charged a whole hour for a fraction
# well add 1 to the cpu hours and then cut it off at the int value.  This is only
# going over if we have calculated the number of CPU hours exactly on an integer
# value, but the chances of that are really slim
computedCPUHours = int(float(kv['pipeline.COMPUTED_CPU_HOURS']) + 1)

cl = cluster.listClusters('localhost', {'cluster_name': kv['cluster.CLUSTER_NAME']})[0]
#
# Artificial number here, this will go away at some point
maxExecNodes = int(kv.get('pipeline.EXEC_MAX_INSTANCES', 0)) or int(cl['config'].get('cluster.exec_max_instances', 10000))
maxExecCores = maxExecNodes * CORES_PER_CPU


if computedCPUHours < maxExecCores:
    # Low ball this one.  If the number of hours is so low that the wanted instances is 0 than make it at least 1
    sys.stdout.write('%s=%d\n' % ('pipeline.EXEC_WANT_INSTANCES', (computedCPUHours / (CORES_PER_CPU * 2)) or 1))
elif computedCPUHours < maxExecCores * 2:
    sys.stdout.write('%s=%d\n' % ('pipeline.EXEC_WANT_INSTANCES', int(maxExecNodes * 0.75)))
else:
    sys.stdout.write('%s=%d\n' % ('pipeline.EXEC_WANT_INSTANCES', maxExecNodes))
