#!/usr/bin/perl

use strict;
use POSIX;

# Torque submit filter to process multicore jobs or jobs requesting large amounts of memory
# This script calculates new memory limits, CPU time limits and/or the number of cores
# The output is identical to the input, except for the recalculated resource limits

# Site-specific limits
my $phys_mem_per_slot = '4096mb';
my $max_cores_per_node = 8;

my %scale_factor = (
    'k' => 1024,
    'm' => 1024*1024,
    'g' => 1024*1024*1024,
    't' => 1024*1024*1024*1024,
);

my $timeout = 10;

my $queue;                  # name of the queue
my %queue_par;              # resource limits on the queue
my %requested;              # resources requested in the input 


# read qsub script from STDIN
my @lines = &read_qsub_input;


# Get resources parameters (default, max) for the selected queue
&get_queue_params( $queue, \%queue_par );

# Get the sanitized version of the requested resources
my %job_resources = &calc_job_resources( \%queue_par, \%requested );

# Finally, print output
# Only modify some #PBS -l resource directives
&print_output( \@lines, \%job_resources );

exit 0;


###############################################################################


# Read the Torque submit script
sub read_qsub_input {
    my @lines;

    @lines = (<>);

    # read comments and PBS directives
    foreach ( @lines ) {
        last unless ( /^#/ );

        # PBS queue specification
        if ( /^#PBS\s+[-]q\s+(\S+)/ ) {
            $queue = $1;
            next;
        }

        # PBS resource specification
        if ( /^#PBS\s+[-]l/ ) {
            if ( /nodes=/ ) {
                # match optionally repeated node/core specs, e.g.:
                # nodes=2:ppn=2+1:ppn=1
                my @node_core = /[=+](?:(\d+)(?::ppn=(\d+))?)/gm;
                for( my $i=0; $i<@node_core; $i+=2 ) {
                    $requested{nodes} += $node_core[$i];
                    $requested{cores} += $node_core[$i+1];
                    $requested{max_cores} = max( $node_core[$i+1], $requested{max_cores} );
                }
            }
            elsif ( /(pmem|vmem|pvmem|mem)=(\d+([kmgt]?[bw]?))/ ) {
                $requested{$1} = $2;
            }
            elsif ( /(cput|walltime)=(\d+)/ ) {
                $requested{$1} = $2;
            }
        }
    }

    return @lines;
}



# Obtain resource definitions for the queue 
sub get_queue_params {
    my ( $q, $qpr ) = @_;

    # Use qmgr to fetch queue parameters
    # Define a timeout for this operation - not to hang submission
    #$SIG{'ALRM'} = sub {
        #print ("ERROR: No response from $hostname (alarm timeout)\n");
        #exit $ERRORS{"UNKNOWN"};
    #};

    #alarm( $timeout );
    my $qmgr_cmd = "qmgr -c 'print queue $q'";
    open QMGR, "$qmgr_cmd|" or die "Failed to execute $qmgr_cmd\n";
    while ( <QMGR> ) {
        if ( /set queue $q resources_(min|default|max)\.(\S+)\s*=\s*(\S+)/ ) {
            $qpr->{$1}{$2} = $3;
        }
    }
    close QMGR;
    #alarm( 0 );
}



# Recalculate requested resource
sub calc_job_resources {
    my ( $qpr, $reqr ) = @_;

    # Job resources
    # Nodes and cores are initialized as requested resources
    my %resources;
    $resources{nodes} = $reqr->{nodes} || 1;
    $resources{cores} = $reqr->{cores} || 1;

    # Maximum cores requested on any of the nodes
    my $maxppn = $reqr->{max_cores} || 1;

    # Under linux, PBS ignores memory requirements ('mem') when the number of nodes is not 1
    # We do the same here
    if ( defined( $reqr->{mem} ) && $resources{nodes} == 1 ) {
        # requested physical memory
        my $phys_mem = &convert_size_to_bytes( $reqr->{mem} );

        # memory per job slot is the global default,
        # unless specified for the queue (pmem)
        my $mem_per_slot = &convert_size_to_bytes( $phys_mem_per_slot );
        if ( defined( $qpr->{default}{pmem} ) ) {
            $mem_per_slot = &convert_size_to_bytes( $qpr->{default}{pmem} ); 
        }

        # Recalculate the number of required cores and max cores on THE node
        if ( $phys_mem > $mem_per_slot ) {
            my $coreholder = max( $resources{cores}, min( $max_cores_per_node, ceil( $phys_mem / $mem_per_slot ) ) );
            $resources{cores} = $coreholder;
            $maxppn = $resources{cores};
        }
    }

    
    # Scale the CPU time with the total number of cores and apply the queue limit
    if ( defined( $qpr->{default}{cput} ) ) {
        if ( defined( $reqr->{cput} ) ) {
            $resources{cput} = $reqr->{cput};
        }
        else {   
            $resources{cput} = &convert_time_to_sec( $qpr->{default}{cput} ) * $resources{cores};
        }
    }
    else {
        if ( defined( $reqr->{cput} ) ) {
            $resources{cput} = $reqr->{cput};
        }
    }
    # Always cap the cput with the max system cput (if defined) or 
    if ( defined( $qpr->{max}{cput} ) ) {
        $resources{cput} ||= &convert_time_to_sec( $qpr->{max}{cput} );
        $resources{cput} = min( $resources{cput}, &convert_time_to_sec( $qpr->{max}{cput} ) );
    }
    if ( defined( $qpr->{max}{walltime} ) ) {
        my $max_scaled_walltime = &convert_time_to_sec( $qpr->{max}{walltime} ) * $resources{cores};
        $resources{cput} = min( $resources{cput}, $max_scaled_walltime );
    }

    # Scale the various memory limits with the total number of cores 
    # OR the maximum number cores on 1 of the nodes
    # Apply the queue limit
    foreach my $f ( 'pvmem', 'pmem', 'vmem', 'mem' ) {
        if ( defined( $qpr->{default}{$f} ) ) {
            if ( defined( $reqr->{$f} ) ) {
                $resources{$f} = $reqr->{$f};
            }
            else {
                # scale per process memory limits with max. cores on a node
                # and per job limits with total number of allocated cores
                my $factor = $resources{cores};
                $factor = 1 if ( $f =~ /^p/ );
                $resources{$f} = $factor * &convert_size_to_bytes( $qpr->{default}{$f} );
            }
        }
        # Always set a max, if it is defined
        if ( defined( $qpr->{max}{$f} ) ) {
            $resources{$f} ||= &convert_size_to_bytes( $qpr->{max}{$f} );
            $resources{$f} = min( $resources{$f}, &convert_size_to_bytes( $qpr->{max}{$f} ) );
        }
    }

    return %resources;

}


# Print the output script
sub print_output {
    my ( $lines, $jobr ) = @_;

    my $printed_resources = 0;      # make sure resource directives are printed just once
    foreach ( @$lines ) {
        if ( /^#PBS\s+[-]l/ ) {
            if ( /nodes=(\d+)/ ) {
                # Ensure that the computed number of cores is printed
                if ( $jobr->{nodes} == 1 && $jobr->{cores} > 1 ) {
                    s!nodes=.*!nodes=$jobr->{nodes}:ppn=$jobr->{cores}!;
                }
                delete $jobr->{nodes};
                delete $jobr->{cores};
            }
            else {
                # In-place substitution of resources that were requested
                # Remove these from the hash holding the job's resources
                foreach my $res ( keys %$jobr ) {
                    if ( /$res=/ ) {
                        s!$res=(\w+)!$res=$jobr->{$res}!;
                        delete $jobr->{$res};
                    }
                }
            }
        }
        elsif ( !/^#/  && $printed_resources == 0 ) {
            # At the end of the PBS directives (the first non-comment line).
            # add the job's resource directives that weren't printed before
            foreach my $res ( keys %$jobr ) {
                print "#PBS -l $res=$jobr->{$res}\n";
            }
            $printed_resources = 1;
        }

        # Write all input lines to stdout
        print;
    }
}


# Return minimum of two numbers
sub min {
    my ( $a, $b ) = @_;
    return ( $a < $b ? $a : $b );
}


# Return maximum of two numbers
sub max {
    my ( $a, $b ) = @_;
    return ( $a > $b ? $a : $b );
}


# Convert a size to bytes using multiples of 1024
sub convert_size_to_bytes {
    my ( $x, $unitr ) = @_;

    my ( $value, $prefix, $factor );

    # match a Torque size, e.g. 2048mb or 2097152
    if ( $x =~ /(\d+)(([kmgt])?[bw])?/ ) {
        $value = $1;
        $prefix = $3;
        $$unitr = $2 if ( defined $unitr );

        # lookup the size prefix (k, m, g, t) or use 1 when not found
        $factor = $scale_factor{$prefix} || 1;
    }
    else {
        die "FATAL: No match in &scalefactor for $x\n";
    }
    return ( $value * $factor );
}


# Convert HH:MM:SS to seconds
sub convert_time_to_sec {
    my ( $tim ) = @_;

    my ( $hours, $minutes, $seconds ) = ( 0, 0, 0 );
    if ( $tim =~ /(\d+):(\d\d):(\d\d)/ ) {
        $hours = $1;
        $minutes = $2;
        $seconds = $3;
    }
    elsif ( $tim =~ /(\d\d):(\d\d)/ ) {
        $minutes = $1;
        $seconds = $2;
    }
    elsif ( $tim =~ /(\d+)/ ) {
        $seconds = $1;
    }
    return ( 60 * (60 * $hours + $minutes) + $seconds );
}
