#!/usr/local/bin/perl
#
# Modify balance/GrpCpuMins in Slurm (NOT PDB) for a specific account
# See also slurm_modify_project_balance for a project

use strict;
use warnings;

use POSIX qw(floor);
use Getopt::Long;

use lib '/cell_root/software/HPCC/lib/site_perl';
use Glue::HPCC::Cluster;
use Glue::HPCC::Project::Log;
use Glue::HPCC::ClusterUtil;
#Only valid for Slurm
use Glue::HPCC::ClusterUtil::Slurm::Project;


my $slrmproj = 'Glue::HPCC::ClusterUtil::Slurm::Project';
my $GHPLog = 'Glue::HPCC::Project::Log';

my $fmt = "%-10s %-10s %-10s %-10s %-10s\n";
my $ksufmt = "%-0.3f";
my $minfmt = "%-d";

sub usage()
{	print STDERR <<EOF;
$0: Modify ksu balance and/or usage for account in Slurm

This modifies the account balance (GrpCPUMins) and/or usage in Slurm.
It DOES NOT touch the Projects DB (except to update logs)

Usage:

$0 [ -cluster CLUSTER ] [ --force ] [ --dryrun ]  \
	-reason REASON -account ACCOUNT 
	[ -unit UNIT -mode MODE -amount AMOUNT ] [ --zero-usage ]

CLUSTER will default based on system running command on.
REASON will get logged to the PDB
ACCOUNT is the name of the __account__ to modify
UNIT should be one of SU, kSU, CPUmin
MODE should be one of set, add, sub
AMOUNT is the amount (in units of UNIT) to set, add, or sub (see MODE)

Unless --force is given, user will be prompted before 
anything is done.

If --zero-usage given, the usage on the account will be zeroed.
Either all of UNIT, MODE, and AMOUNT must be given (in which case
GrpCPUMins is modified), or none, in which GrpCPUMins in not modified
(only makes sense if --zero-usage given).  You can modify GrpCPUMins
and zero usage in the same command if desired.

EOF
}

sub yorn_abort(;$)
{	my $prompt = shift || 'Proceed? (y/N)?';
	print STDERR "\n$prompt ";
	my $ans = <STDIN>;
	chomp $ans;
	unless ( $ans =~ /^\s*y(es?)?\s*$/i )
	{	die "Aborted at user request.\n";
	}
	return;
}

sub print_confirmation_line(;$$$$$)
{	my $fld = shift;
	my $oldk = shift;
	my $newk = shift;
	my $oldm = shift;
	my $newm = shift;

	if ($fld)
	{	$oldk = sprintf $ksufmt, $oldk;
		$newk = sprintf $ksufmt, $newk if $newk ne '-';
		$oldm = sprintf $minfmt, $oldm;
		$newm = sprintf $minfmt, $newm if $newm ne '-';
	} else
	{	$fld = '';
		$oldk = 'Old(kSU)';
		$newk = 'New(kSU)';
		$oldm = 'Old(min)';
		$newm = 'New(min)';
	}

	printf STDERR $fmt, $fld, $oldk, $newk, $oldm, $newm;
}

#------------------------------------------------------------------
#		Main
#------------------------------------------------------------------

my $help;
my $FORCE;
my $DRYRUN;
my ($reason, $account, $unit, $mode, $amount, $zero);
my $cluster;

my $res = GetOptions(
	'help!' => \$help,
	'force!' => \$FORCE,
	'dryrun!' => \$DRYRUN,

	'zero|zero-usage|zero_uage!' => \$zero,
	'reason=s' => \$reason,
	'account=s' => \$account,
	'unit=s' => \$unit,
	'mode=s' => \$mode,
	'amount=f' => \$amount,
);

unless ($res)
{	usage;
	die "\nError parsing options.\n";
}

if ( $help )
{	usage;
	exit 0;
}

unless ( $account )
{	usage;
	die "\nMissing required parameter 'account'.\n";
}

unless ( $reason )
{	usage;
	die "\nMissing required parameter 'reason'.\n";
}

my $tmp;

#If any of amount, unit or mode given, must give all
if ( defined $amount || $unit || $mode )
{   unless ( defined $amount )
    {	usage;
	die "\nMissing required parameter 'amount'.\n";
    }

    unless ( $unit )
    {	usage;
	die "\nMissing required parameter 'unit'.\n";
    }
    $tmp = lc $unit;
    unless ( $tmp eq 'ksu' || $tmp eq 'su' || $tmp eq 'cpumin' )
    {	usage;
	die "\nIllegal value '$unit' for 'unit', expecting kSU, SU, or CPUmin.\n";
    }
    $unit = $tmp;


    unless ( $mode )
    {	usage;
	die "\nMissing required parameter 'mode'.\n";
    }
    $tmp = lc $mode;
    unless ( $tmp eq 'set' || $tmp eq 'add' || $tmp eq 'sub' )
    {	usage;
	die "\nIllegal value '$mode' for 'mode', expecting set, add, or sub.\n";
    }
    $mode = $tmp;
} else
{   unless ( $zero )
    {	usage;
	die "\nYou really need to either give --amount, --mode, and --unit, or give --zero.  Nothing to do.\n";
    }
}

my $hpc = Glue::HPCC::Cluster->new_from_clustername($cluster);
$cluster = $hpc->db_clustername;
$hpc->dryrun($DRYRUN);

#Convert amount to ksu and cpumin
my ($amtksu, $amtmin);
if ( defined $amount )
{   if ( $unit eq 'ksu' )
    {	$amtksu = $amount;
	$amtmin = floor($amount * 1000 * 60 +0.5); #Must be integer
    } elsif ( $unit eq 'su' )
    {	$amtksu = $amount/1000;
	$amtmin = floor($amount * 60 + 0.5); #Must be integer
    } elsif ( $unit eq 'cpumin' )
    {	$amtksu = $amount/60/1000;
	$amtmin = floor($amount + 0.5); #Must be integer
    } else
    {	die "Unrecognized unit '$unit' at ";
    }
}


#Get current and new balances

my $sshare = $slrmproj->sshare;
$tmp = $sshare->usage_for_account_in_cluster(
	account=>$account, cluster=>$cluster);
unless ( $tmp && ref($tmp) eq 'ARRAY' )
{	die "Error looking up balance for $account in $cluster: $tmp at ";
}

my $cur_cpumin_limit = floor($tmp->[1] + 0.5); #Must be integer
my $cur_ksu_limit = $cur_cpumin_limit/60/1000;
my $cur_cpumin_bal = $tmp->[2];
my $cur_ksu_bal = $cur_cpumin_bal/60/1000;
my $cpusec_used = $tmp->[0];
my $cpumin_used = $cpusec_used/60;
my $ksu_used = $cpumin_used/60/1000;
my $new_cpumin_used = $cpumin_used;
my $new_ksu_used = $ksu_used;
if ( $zero )
{	$new_cpumin_used = 0;
	$new_ksu_used = 0;
}

my ($new_cpumin_limit, $new_ksu_limit);
$mode = '' unless $mode; #Avoid warnings re use of undef value in comparison
if ( $mode eq 'set' )
{	$new_cpumin_limit = $amtmin;
	$new_ksu_limit = $amtksu;
} elsif ( $mode eq 'add' )
{	$new_cpumin_limit = $cur_cpumin_limit + $amtmin;
	$new_ksu_limit = $cur_ksu_limit + $amtksu;
} elsif ( $mode eq 'sub' )
{	$new_cpumin_limit = $cur_cpumin_limit - $amtmin;
	$new_ksu_limit = $cur_ksu_limit - $amtksu;
} else
{	die "Illegal mode '$mode' at " if $mode;
	$new_cpumin_limit = $cur_cpumin_limit;
	$new_ksu_limit = $cur_ksu_limit;
}

my $new_cpumin_bal = $cur_cpumin_bal;
my $new_ksu_bal = $cur_ksu_bal;

$new_cpumin_bal = $new_cpumin_limit - $new_cpumin_used;
$new_ksu_bal = $new_ksu_limit - $new_ksu_used;

my ($tmp_used, $tmp_lim);

#Verify limits > 0
if ( $new_cpumin_limit < 0 )
{       $tmp_lim  = sprintf $ksufmt, $new_ksu_limit;
        print STDERR <<EOF;
***** WARNING ******
Requested transaction would bring the limit on
$account ($tmp_lim kSU) below 0.

EOF
        yorn_abort;
}

# Verify that new limit >= usage
if ( $new_cpumin_used > $new_cpumin_limit )
{       $tmp_used = sprintf $ksufmt, $new_ksu_used;
        $tmp_lim  = sprintf $ksufmt, $new_ksu_limit;
        print STDERR <<EOF;
***** WARNING ******
Requested transaction would bring the limit on
$account ($tmp_lim kSU) below current usage ($tmp_used kSU).

EOF
        yorn_abort;
}


unless ($FORCE)
{	print STDERR <<EOF;
--------------------------------------
Confirmation of ACCOUNT balance modification:

Cluster: $cluster
Account: $account

EOF
	print_confirmation_line;
	print_confirmation_line 'Limit',
		$cur_ksu_limit, $new_ksu_limit,
		$cur_cpumin_limit, $new_cpumin_limit;
	print_confirmation_line 'Used',
		$ksu_used, $new_ksu_used,
		$cpumin_used, $new_cpumin_used;
	print_confirmation_line 'Balance',
		$cur_ksu_bal, $new_ksu_bal,
		$cur_cpumin_bal, $new_cpumin_bal;
	print_confirmation_line 'Delta',
		($new_ksu_bal - $cur_ksu_bal), '-',
		($new_cpumin_bal - $cur_cpumin_bal), '-';


	yorn_abort;
}


my $sacctmgr = $slrmproj->sacctmgr;
$sacctmgr->dryrun($DRYRUN);
$res = $slrmproj->sacctmgr_reset_account_balances(
	account => $account, cluster=> $cluster,
	zero_usage => $zero, cpumin => $new_cpumin_limit);
die "Error reseting account balance on $account: $res at "
	if $res;

#Log what we did
$res = $GHPLog->add_log_entry($hpc,
	action_code => 'modify-acct-usage',
	reason => $reason,
	account => $account,
	cluster => $cluster,
	tablename => 'slurm',
	fieldname => 'usage(cpusec)',
	oldval => $cpusec_used,
	newval => 0,
	description => 'Manually zero usage on account'
) if $zero;
$res = $GHPLog->add_log_entry($hpc,
	action_code => 'modify-acct-bal',
	reason => $reason,
	account => $account,
	cluster => $cluster,
	tablename => 'slurm',
	fieldname => 'GrpCPUMin',
	oldval => $cur_cpumin_limit,
	newval => $new_cpumin_limit,
	description => 'Manually reset GrpCPUMin limit on account'
) if defined $amount;
$hpc->db_commit;

	





	



