#!/usr/bin/perl -w
#
#  altix_node2vnode - Facilitates migrating custom resources from the single 
#  host "node" as known in PBS Professional 7.1 to the multi-vnode systems
#  as presented in PBS Professional 8.0
#
#
# Copyright (C) 1994-2018 Altair Engineering, Inc.
# For more information, contact Altair at www.altair.com.
#
# This file is part of the PBS Professional ("PBS Pro") software.
#
# Open Source License Information:
#
# PBS Pro is free software. You can redistribute it and/or modify it under the
# terms of the GNU Affero General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# PBS Pro is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.
# See the GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Commercial License Information:
#
# For a copy of the commercial license terms and conditions,
# go to: (http://www.pbspro.com/UserArea/agreement.html)
# or contact the Altair Legal Department.
#
# Altair’s dual-license business model allows companies, individuals, and
# organizations to create proprietary derivative works of PBS Pro and
# distribute them - whether embedded or bundled with other software -
# under a commercial license agreement.
#
# Use of Altair’s trademarks, including but not limited to "PBS™",
# "PBS Professional®", and "PBS Pro™" and Altair’s logos is subject to Altair's
# trademark licensing policies.

#  USAGE:
#  altix_node2vnode /path/to/nodelist
#
# The single argument to this script is a list of all nodes with all defined
# resources, this will be saved by the PBS Professional 8.0 INSTALL 
# script before the 7.1 server was shut down.

if ( $ARGV[0] && $ARGV[0] eq "--version" ) {
    if ( ! defined($ARGV[1]) ) {
        print "pbs_version = @PBS_VERSION@\n";
        exit 0;
    }
}

use strict;
use Switch;
use Env;

my %allnodes;

if($#ARGV != 0 || ! -f $ARGV[0]) {
    print STDERR "Usage $0 /path/to/nodelist\n";
    exit 1;
} else {
    my @rawnodes = `cat $ARGV[0]`;
    my $hostname = undef;
    foreach my $line (@rawnodes) {
        chomp $line;
        my @linearray = split(/ +/,$line);
        if(defined($linearray[1])){
            if(($linearray[1] eq "Host") or ($linearray[1] eq  "Mom")){
                $hostname = $linearray[3];
                my $node = new nodeclass($linearray[3]);
                $allnodes{$linearray[3]}=$node;
            } else {
                switch($linearray[1]) {
                    case "ntype"    { $allnodes{$hostname}->{'NTYPE'} = $linearray[3]; }
                    case "pcpus"    { $allnodes{$hostname}->{'PCPUS'} = $linearray[3]; }
                    case "license"  { $allnodes{$hostname}->{'LICENSE'} = $linearray[3]; }
                    case "state"  { $allnodes{$hostname}->{'STATE'} = $linearray[3]; }
                    case "resources_available.arch"  { $allnodes{$hostname}->{'ARCH'} = $linearray[3]; }
                    case "resources_available.host"  { $allnodes{$hostname}->{'HOST'} = $linearray[3]; }
                    case "resources_available.mem"  { $allnodes{$hostname}->{'MEM'} = $linearray[3]; }
                    case "resources_available.ncpus"  { $allnodes{$hostname}->{'NCPUS'} = $linearray[3]; }
                    case "resources_available.vnode"  { $allnodes{$hostname}->{'VNODE'} = $linearray[3]; }
                    case "resources_available.vmem"  { $allnodes{$hostname}->{'VMEM'} = $linearray[3]; }
                    case (/resources_available.*/) { $allnodes{$hostname}->{'customs'}->{$linearray[1]}=$linearray[3]; }
                }
            }
        }
    }
}

redist_resources("test");
print "\nWould you like to proceed with the above changes? (yes/NO):  ";
my $input = <STDIN>;
chomp($input);

if(lc($input) eq "yes") {
    my $result = redist_resources("no-test");
    if($result ne 0) {
        print "\nERROR: One or more operations failed.\n";
        exit 1;
    }
}


exit 0;

# Function that redistributes resources to Altix systems. Takes one argument which can have the value of
# test or no-test. This value will determine if the user should be shown changes that would be made or if 
# the changes should actually be made.
#
# returns 0 on success and -1 on failure.
sub redist_resources {
    my $mode=shift;
    my $anyfail=0;
    foreach my $key (keys %allnodes) {
        dprint("\nProcessing Host: $key\n");
        if($allnodes{$key}->{'ARCH'} eq "linux_cpuset"){
            dprint("\tHost $key appears to be an Altix\n");
            if(defined($allnodes{$key}->{'customs'})){
                print "\nProposed changes to ".$key.":\n\n" if($mode eq "test");
                my $pbspath = get_pbs_path();
                my @vnodes = `$pbspath/bin/pbsnodes -av | grep resources_available.vnode | grep -i $key | awk \'{ print \$3 }\'`;
                foreach my $newkey (keys %{$allnodes{$key}->{'customs'}}){
                    my $command = sprintf("%s/bin/qmgr -c \"set node %s %s=%s\"\n",$pbspath,$key,$newkey,$allnodes{$key}->{'customs'}->{$newkey});
                    if($mode eq "test"){
                        print $command;
                    } else {
                        print "running: $command";
                        my $result = system($command);
                        dprint("result=$result\n");
                        if($result ne 0){
                            $anyfail=$result;
                        }
                    }
                    foreach my $vnode_key (@vnodes){
                        chomp $vnode_key;
                        if($vnode_key ne $key){
                            my $command = sprintf("%s/bin/qmgr -c \"set node %s %s=@%s\"\n",$pbspath,$vnode_key,$newkey,$key);
                            if($mode eq "test") {
                                print $command;
                            } else {
                                print "running: $command";
                                my $result = system($command);
                                dprint("result=$result\n");
                                if($result ne 0) {
                                    $anyfail=$result;
                                }
                            }
                        }
                    }
                }
            }
        } else { dprint("\tHost $key does NOT appear to be an Altix\n"); }

    }
    return $anyfail;
}

# Attempts to determine the path to qmgr based on /etc/pbs.conf
# Program will abort if this command is not found.
# Return value is the full path including command:
# ex: /usr/pbs/bin/qmgr
sub get_pbs_path {
    # Assuming that PBS_CONF_FILE exists in the users' ENV, if not assuming 
    # the default location of the pbs.conf file.
    my $cf = ($ENV{PBS_CONF_FILE} ? $ENV{PBS_CONF_FILE} : "/etc/pbs.conf" );

    my $result = `cat $cf | grep PBS_EXEC`;
    my @pbs_exec_path = split(/=/,$result);
    chomp($pbs_exec_path[1]);
    if( -d $pbs_exec_path[1]) {
        return $pbs_exec_path[1];
    } else {
       die "ERROR: Unable to find PBS Professional.\n";
    }
}

# Function to print debug output, if environment variable DEBUG is set the string passed to dprint will
# be printed to stderr, otherwise nothing will be printed.
sub dprint {
    printf STDERR @_ if(defined($ENV{'DEBUG'}));
    return;
}

# Simple class that represents a node
package nodeclass;

sub new {
    my $proto = shift;
    my $hostname = shift;
    my $class = ref($proto) || $proto;
    my $self = {};
    $self->{hostname} = $hostname;
    bless ($self,$class);
    return $self;
}

my %customs;
my $NTYPE;
my $PCPUS;
my $LICENSE;
my $STATE;
my $ARCH;
my $HOST;
my $MEM;
my $NCPUS;
my $VNODE;
my $VMEM;

sub DESTROY {
    return;
}
__END__
