oargridsub

#!/usr/bin/perl

use oargrid_lib;
use oargrid_conflib;
use Data::Dumper;
use Getopt::Long;
use Time::Local;
use POSIX qw(strftime);
use strict;
use warnings;

# Initialize database connection
oargrid_conflib::init_conf(oargrid_lib::get_config_file_name());
my $DB_SERVER = oargrid_conflib::get_conf("DB_HOSTNAME");
my $DB_BASE_NAME= oargrid_conflib::get_conf("DB_BASE_NAME");
my $DB_BASE_LOGIN = oargrid_conflib::get_conf("DB_BASE_LOGIN");
my $DB_BASE_PASSWD = oargrid_conflib::get_conf("DB_BASE_PASSWD");

my $dbh = oargrid_lib::connect($DB_SERVER,$DB_BASE_NAME,$DB_BASE_LOGIN,$DB_BASE_PASSWD);

my $commandTimeout = oargrid_lib::get_command_timeout();

# Print help messages
sub usage($){
    my $dbase = shift;
    
    print <<EOU;
Usage oargridsub -s date [-q queue_name][-p program_to_run][-w walltime][-i identity_file][-T timeout][-d directory][-v][-V][-f file] DESC
    -s give the date of the begining of the reservation, ex: "2005-01-07 11:00:00" (default is NOW)
    -q give the queue
    -t give the type (several -t are allowed)
    -p give the program to run
    -d give the directory where the program will be launched (default is "~")
    -w walltime of the reservation (default is "1:00:00")
    -i import a job key from the specified file
    -T timeout in seconds for remote cluster commands (default: $commandTimeout)
    -F continue even if there is a rejected reservation
    -v turn on verbose mode
    -V print oargrid version and exit
    -f specifie a file where to read reservation description (instead of command line)
Where DESC is in the form :
    clusterAlias1:rdef="/nodes=5/core=4":prop="switch = 'sw4'",clusterAlias2:rdef="/cpu=16/core=1":type="deploy"...

EOU
    print("Available cluster aliases are:\n");
    oargrid_lib::print_aliases_hierarchy($dbase);
}

sub sqldate_to_timestamp($) {
  my $date=shift;
  $date =~ tr/-:/  /;
  my ($year,$mon,$mday,$hour,$min,$sec) = split / /,$date;
  return Time::Local::timelocal_nocheck($sec,$min,$hour,$mday,$mon-1,$year-1900);
}

sub timestamp_to_sqldate($) {
  my $timestamp=shift;
  return(strftime("%F %T",localtime($timestamp)));  
}

$SIG{INT} = 'IGNORE';

# Get the user name
if (!defined($ENV{SUDO_UID})){
    die("[OAR_GRISUBD] I can not get user id\n");
}
my $lusr= getpwuid($ENV{SUDO_UID});

if ($lusr eq 'oargrid' || $lusr eq 'oar' || $lusr eq 'root') {
  die("[OAR_GRISUBD] $lusr can not submit grid jobs!\n");
}

#Store command line for the database
my @initialArguments = @ARGV;

# parse arguments
my $walltime ;
my $timeout;
my $program ;
my $directory ;
my $startDate;
my $sos;
my $verbose;
my $queue;
my @type;
my $file;
my $continue;
my $version;
my $oar_keys=1; # now mandatory
my $jobkey;
Getopt::Long::Configure ("gnu_getopt");
GetOptions ("walltime|w=s" => \$walltime,
            "import-job-key|i=s" => \$jobkey,
            "timeout|T=s" => \$timeout,
            "program|p=s" => \$program,
            "directory|d=s" => \$directory,
            "start_date|s=s" => \$startDate,
            "help|h" => \$sos,
            "verbose|v" => \$verbose,
            "queue|q=s" => \$queue,
            "type|t=s" => \@type,
            "file|f=s" => \$file,
            "FORCE|F" => \$continue,
            "VERSION|V" => \$version
           );

if (defined($version)){
    print("OARGRID version : ".oargrid_lib::get_version()."\n");
    exit(0);
}

if (defined($sos) || (!defined($file) && !defined($ARGV[0]))){
    usage($dbh);
    oargrid_lib::disconnect($dbh);
    exit(1);
}

if (defined($timeout)){
 $commandTimeout=$timeout;
}

# Parse cluster specific parameters and file the hash table %clusterHashArguments
my $argumentDescription = "@ARGV";
print("$argumentDescription\n");
# test if there is a file specified
if (defined($file)) {
    if (-r $file){
        open(FILE, $file);
        my $content;
        while (<FILE>){
            $content .= $_;
        }
        close(FILE);
        #remove comments
        $content =~ s/\s*#.*\n/\n/g;
        #make only one line
        $argumentDescription =~ s/\n/ /g;
        
        if (($content =~ m/WALLTIME\s*=\s*\"(.+)\"\s*/m) && !defined($walltime)){
            $walltime = "$1";
        }
        if (($content =~ m/PROG\s*=\s*\"(.+)\"\s*/m) && !defined($program)){
            $program = "$1";
        }
        if (($content =~ m/DIR\s*=\s*\"(.+)\"\s*/m) && !defined($directory)){
            $directory = "$1";
        }
        if (($content =~ m/STARTDATE\s*=\s*\"(.+)\"\s*/m) && !defined($startDate)){
            $startDate = "$1";
        }
        if (($content =~ m/QUEUE\s*=\s*\"(.+)\"\s*/m) && !defined($queue)){
            $queue = "$1";
        }
        if (($content =~ m/TYPE\s*=\s*\"(.+)\"\s*/m) && !@type){
            push(@type,$1);
        }
        if ($content =~ m/DESC\s*=\s*\"(.+)\"\s*/m){
            my $tmpValue = $1;
            if (!($argumentDescription =~ m/^\s*$/m)){
                $argumentDescription .= ","."$tmpValue";
            }else{
                $argumentDescription = "$tmpValue";
            }
        }
    }else{
        print("/!\\ File $file not readable by user oargrid.\n");
        oargrid_lib::disconnect($dbh);
        exit(4);
    }
}

#Default values
if (!defined($walltime)){
    $walltime = "1:00:00";
}else{
    my ($h,$m,$s) =split (":",$walltime);
    if ((!defined($h)) or (!defined($m)) or (!defined($s))){
        warn("[ERROR] Bad walltime syntax (use something like 13:00:00).\n");
        exit(5);
    }
}
if (!defined($directory)){
    $directory = "~$lusr";
}
if (!defined($queue)){
    $queue = "default";
}
if (!defined($program)){
    $program = "";
}


print("DESC string : $argumentDescription\n") if (defined($verbose));
my @argParsed = split(',', $argumentDescription);
my @clusterArguments;
foreach my $i (@argParsed){
    my %tmpHash = (
                    nodes => 1,
                    weight => 0,
                    properties => "",
                    queue => "",
                    name => "",
                    env => "",
                    part => "",
		    types => [],
                    rdef => ""
                  );
    my @tmpArray = split(":", $i);
    $tmpHash{clusterName} = shift(@tmpArray);
    $tmpHash{clusterName} =~ m/\s*(.+)\s*/m;
    $tmpHash{clusterName} = $1;
    my @types;
    foreach my $j (@tmpArray){
        print "Scanning $j...\n" if (defined($verbose));
        if ($j =~ m/\s*rdef\s*=\s*(.+)/m){
            $tmpHash{rdef} = $1;
        }elsif ($j =~ m/\s*nodes\s*=\s*(\d+)/m){
                $tmpHash{nodes} = $1;
            if (not $j =~ m/\s*rdef\s*=\s*/m){
                print "WARNING: the 'nodes' keyword is deprecated. Please use a resource description with 'rdef'.\n";
            }
        }elsif ($j =~ m/\s*cpu\s*=\s*(\d+)/m){
                $tmpHash{weight} = $1;
            if (not $j =~ m/\s*rdef\s*=\s*/m){
                print "WARNING: the 'cpu' keyword is deprecated. Please use a resource description with 'rdef'.\n";
            }
        }elsif ($j =~ m/\s*prop\s*=\s*(.+)/m){
            $tmpHash{properties} = $1;
        }elsif ($j =~ m/\s*env\s*=\s*(.+)/m){
            $tmpHash{env} = $1;
        }elsif ($j =~ m/\s*part\s*=\s*(.+)/m){
            $tmpHash{part} = $1;
        }elsif ($j =~ m/\s*queue\s*=\s*(.+)/m){
            $tmpHash{queue} = $1;
        }elsif ($j =~ m/\s*name\s*=\s*(.+)/m){
            $tmpHash{name} = $1;
        }elsif ($j =~ m/\s*type\s*=\s*(.+)/m){
            push(@types,$1);
        }
	else {
	  print "ERROR: $j: unknown keyword\n";
	  oargrid_lib::disconnect($dbh);
	  exit(2);
	}
    }
    $tmpHash{types}=[ @types ];
    push(@clusterArguments, \%tmpHash);
}

#test cluster names
my %clusters = oargrid_lib::get_cluster_names($dbh);
foreach my $i (@clusterArguments){
    if (!defined($clusters{$i->{clusterName}})){
        print("[OAR_GRIDSUB] ERROR, the alias $i->{clusterName} is not registered\n");
        print("[OAR_GRIDSUB] right aliases are : \n");
        oargrid_lib::print_aliases_hierarchy($dbh);
        oargrid_lib::disconnect($dbh);
        exit(2);
    }
}

if (defined($startDate) && !($startDate =~ m/^(\d{4}\-\d{1,2}\-\d{1,2})\s(\d{1,2}:\d{1,2}:\d{1,2})$/m)){
    print("\tBAD start date value\n\n");
    usage($dbh);
    oargrid_lib::disconnect($dbh);
    exit(1);
}elsif (!defined($startDate)){
    $startDate = oargrid_lib::get_date($dbh);
}

my $tz_startDate=$startDate;

my $cpuset_name = "oargrid_".oargrid_lib::get_unique_cpuset_id($dbh);
my $tmp_path = "/tmp/oargrid/";
my $ssh_key_name;
my $private_ssh_key;
my $public_ssh_key;
if (defined($oar_keys)){
    system("mkdir -p $tmp_path") and do{warn("[ERROR] Cannot create temporary directory $tmp_path\n");exit(6)};
    $ssh_key_name = $tmp_path."tmp_oargrid_ssh_key_".$cpuset_name;
    if (!defined($jobkey)){
    # Generate ssh keys
        system("ssh-keygen -q -t rsa -N '' -f $ssh_key_name") and do{warn("[ERROR] Cannot create SSH key $ssh_key_name\n");exit(7)};
    }
    else { $ssh_key_name = $jobkey }
    open(FILE,$ssh_key_name) or do{warn("[ERROR] Cannot read SSH key $ssh_key_name\n");exit(8)};
    while (<FILE>) {
        $private_ssh_key .= $_;
    }
    close(FILE);
    open(FILE,$ssh_key_name.".pub") or do{warn("[ERROR] Cannot read SSH key $ssh_key_name.pub\n");exit(9)};
    while (<FILE>) {
        $public_ssh_key .= $_;
    }
    close(FILE);
}
# Launch command on each clusters and parse the job id
my %clusterBatchId;
my $endLoop = 0;
while (($#clusterArguments >= 0) && ($endLoop == 0)){
    my $property_string = '';
    my $i = pop(@clusterArguments);
    my $tmpQueue = $queue;
    if ($i->{queue} ne ""){
        $tmpQueue = $i->{queue};
    }
    my $tmpType = "";
    # Globals types
    if (@type) {
      foreach my $t (@type) {
        $tmpType .= " -t ".$t;
      }
    }
    # Local types (only for that resources definition)
    foreach my $t (@{$i->{types}}){
       $tmpType .= " -t ".$t;
    }
    my $resources_def = "";
    if ($i->{rdef} ne "") {
        $resources_def = $i->{rdef};
    }else{
        $resources_def = "/nodes=$i->{nodes}";
        $resources_def .= "/cpu=$i->{weight}" if ($i->{weight} > 0);
    }
#    my $cmd = "ssh $clusters{$i->{clusterName}}->{hostname} \"/bin/bash -c \\\"cd $directory && sudo -u $lusr oarsub -q $tmpQueue -r \\\\\\\"$startDate\\\\\\\" -l nodes=$i->{nodes},walltime=$walltime";
    #if ($i->{weight} > 0){
    #    $cmd .= "/cpu=$i->{weight}";
    #}
    
    # Date and timezone adjustment
    my $date_cmd="ssh $clusters{$i->{clusterName}}->{hostname} \"/bin/bash -c \\\"date +%s%z\\\"\";date +%s%z";
    my %date_cmdRes=oargrid_lib::launch_command_with_timeout($date_cmd,$commandTimeout);
    print(Dumper(%date_cmdRes)) if (defined($verbose));
    if (!defined($date_cmdRes{status})){
      print("[OAR_GRIDSUB] [$i->{clusterName}] WARNING: Date command timed out. No TZ adjustment.\n");
    }
    elsif ($date_cmdRes{stdout} =~ m/(\d+)([+-]\d+)\d\d\s(\d+)([+-]\d+)\d\d/s) {
      my $deltat=$3+$4*3600-$1-$2*3600;
      print "[OAR_GRIDSUB] [$i->{clusterName}] Date/TZ adjustment: $deltat seconds\n";
      $tz_startDate=timestamp_to_sqldate(sqldate_to_timestamp($startDate)-$deltat);
    }
    else {
      print("[OAR_GRIDSUB] [$i->{clusterName}] WARNING: Could not parse date. No TZ adjustment.\n");
    }
    
    # Submit command
    my $cmd = "ssh $clusters{$i->{clusterName}}->{hostname} \"/bin/bash -c \\\"cd $directory && sudo -u $lusr oarsub -d $directory -q $tmpQueue -r \\\\\\\"$tz_startDate\\\\\\\" -l $resources_def,walltime=$walltime";
    if (defined($oar_keys)){
        #$cmd .= " --job-priv-key \'$private_ssh_key\' --job-pub-key \'$public_ssh_key\'";
        $cmd .= " -k --import-job-key-inline=\'$private_ssh_key\'";
    }else{
        print "[OAR_GRIDSUB] Submit without keys is deprecated. Exiting...\n";
	oargrid_lib::disconnect($dbh);
	exit 10;
        #$cmd .= " --force-cpuset-name $cpuset_name";
    }
    if ((defined $tmpType) && ($tmpType ne "")) {
      $cmd .= " ".$tmpType;
    }
    if ($i->{properties} ne ""){
        $property_string .= "(". $i->{properties} .")";
    }
    if (defined($clusters{$i->{clusterName}}->{properties})) {
        if ($clusters{$i->{clusterName}}->{properties} ne "") {
            if ($property_string ne "") { $property_string .= " and "; }
            $property_string .= "(". $clusters{$i->{clusterName}}->{properties} .")";
        }
    }
    if ($property_string ne "") { $cmd .= " -p \\\\\\\"". $property_string ."\\\\\\\""; }
    if ($program ne ""){
        $cmd .= " '$program'\\\"\"";
    }else{
        $cmd .= " \\\"\"";
    }
    print("[OAR_GRIDSUB] Launch command : $cmd\n") if (defined($verbose));
    my %cmdRes = oargrid_lib::launch_command_with_timeout($cmd,$commandTimeout);
    print(Dumper(%cmdRes)) if (defined($verbose));
    
    if (!defined($cmdRes{status})){
        oargrid_lib::mail_notify("[OAR_GRIDSUB] Command timed out","The command was :\n$cmd\n");
        print("[OAR_GRIDSUB] [$i->{clusterName}] Command timed out\n");
        $endLoop = 1 if (!defined($continue));
    }elsif (($cmdRes{stdout} =~ m/.*JOB_ID\s*=\s*(\d+).*/m) && !($cmdRes{stdout} =~ m/.*--> KO.*/m) && ($cmdRes{status} == 0)){
        my %tmpHash = (
                        "batchId" => $1,
                        "nbNodes" => $i->{nodes},
                        "weight"  => $i->{weight},
                        "properties" => $i->{properties},
                        "queue"   => $tmpQueue,
                        "name"    => $i->{name},
                        "env"     => $i->{env},
                        "part"    => $i->{part},
                        "rdef"    => $resources_def
                      );
        push(@{$clusterBatchId{$i->{clusterName}}}, \%tmpHash);
        #print("[OAR_GRIDSUB] [$i->{clusterName}] Reservation success on $i->{clusterName} : batchId = $1, nbNodes = $i->{nodes}, cpu = $i->{weight}, properties = \"$i->{properties}\", queue = $tmpQueue\n");
        print("[OAR_GRIDSUB] [$i->{clusterName}] Reservation success on $i->{clusterName} : batchId = $1\n");
    }elsif ($cmdRes{stdout} =~ m/.*--> KO.*/m){
        print("[OAR_GRIDSUB] [$i->{clusterName}]Your reservation was rejected when contacting $i->{clusterName}\n");
        $endLoop = 2 if (!defined($continue));
    }else{
        oargrid_lib::mail_notify("[OAR_GRIDSUB] [$i->{clusterName}] I am not able to parse correctly the output","The command was :\n$cmd\n");
        print("[OAR_GRIDSUB] [$i->{clusterName}] I am not able to parse correctly the output, contact administrator (Check if the launching directory is readable and browsable by oargrid user. This is a known limitation but we cannot do in an other way.).\n");
	print ("[OAR_GRIDSUB] [$i->{clusterName}] Use -v option for a more verbose output.\n") if (!defined($verbose));
    print("[OAR_GRIDSUB] [DEBUG] [$cmdRes{status}] $cmdRes{stdout} -- $cmdRes{stderr}\n") if (defined($verbose));
        $endLoop = 3 if (!defined($continue));
    }
}

#print(Dumper(%clusterBatchId));

my $exitCode = 0;
if ($endLoop != 0){
    print("[OAR_GRIDSUB] I delete jobs already submitted\n");
    foreach my $i (keys(%clusterBatchId)){
        foreach my $j (@{$clusterBatchId{$i}}){
            print("[OAR_GRIDSUB] Delete job $j->{batchId} on the cluster $i\n");
            my $deleteCmd = oargrid_lib::get_oardel_command($clusters{$i}->{hostname},$lusr,$j->{batchId});
            print("[OAR_GRIDSUB] Delete command : $deleteCmd\n") if (defined($verbose));
            my %cmdRes = oargrid_lib::launch_command_with_timeout($deleteCmd, $commandTimeout);
        }
    }
    $exitCode = 3;
}else{
    if (%clusterBatchId){
        my $resId = oargrid_lib::add_new_grid_reservation($dbh,$lusr,$directory,$startDate,$walltime,$program,"@initialArguments",\%clusterBatchId);
        print("[OAR_GRIDSUB] Grid reservation id = $resId\n");
        if (defined($oar_keys)){
            my $new_ssh_key = $tmp_path."/oargrid_ssh_key_".$lusr."_$resId";
            system("mv $ssh_key_name $new_ssh_key");
            system("mv $ssh_key_name.pub $new_ssh_key.pub");
            system("sudo chown $lusr $new_ssh_key");
            print("[OAR_GRIDSUB] SSH KEY : $new_ssh_key\n");
            print("\tYou can use this key to connect directly to your OAR nodes with the oar user.\n");
        }
    }
}

oargrid_lib::disconnect($dbh);

exit($exitCode);