nordugrid-arc-arex 1.1.1-1 / usr / share / arc / submit-condor-job

#!/usr/bin/perl -w

use File::Spec;
use File::Temp 'tempfile';
use File::Basename;
use Sys::Hostname;
use Cwd;

BEGIN {
    my $oldcwd = getcwd();

    my $basedir = dirname($0);
    chdir $basedir or die "$basedir: $!";
    $basedir = getcwd();

    my $pkgdatadir = $basedir;

    chdir $oldcwd or die "$oldcwd: $!";

    unshift @INC, $pkgdatadir;
}

use ConfigCentral;
use condor_env;

use strict;
use warnings;

# True if the program is run in debugging mode.
my $debug;

# The name of the Condor job description file.  This file is generated from the
# GRAMI file.
my $cmd_filename;

# Pathname of the Condor log.
my $condor_log;

# Pathname of the real executable.
my $real_exe;

# Pathname of the wrapper script.
my $exewrapper;

my %grami;

$0 =~ s#.*/##;
warn "----- starting $0 -----\n";

if (@ARGV > 1 && $ARGV[0] eq '--config') {
    $ENV{ARC_CONFIG} = $ARGV[1];
    shift;
    shift;
}
if (@ARGV > 0 && $ARGV[0] eq '-d') {
    $debug = 1;
    shift;
}
die "Usage: $0 [--config CONFIG_FILE] [-d] GRAMI_FILE\n" unless @ARGV;

parse_grami(my $gramifile = $ARGV[0]);

$grami{joboption_controldir} ||= dirname($gramifile) || getwd();
$grami{joboption_gridid} ||= $1 if basename($gramifile) =~ /job\.(.*)\.grami$/;

my $configfile = $ENV{ARC_CONFIG} ? $ENV{ARC_CONFIG} : '/etc/arc.conf';
my $fullconfig = ConfigCentral::parseConfig($configfile);
my $share = $grami{joboption_queue};
die "$0: ERROR: Job requested invalid share: $share\n"
    unless exists $fullconfig->{shares}{$share};
my %config = ( %$fullconfig, %{$fullconfig->{shares}{$share}} );

# this finds location of condor executables, and sets up environment vars.
configure_condor_env(%config) or die "Condor executables not found\n";

my $condor_bin_path = $ENV{CONDOR_BIN_PATH};
$ENV{RUNTIME_CONFIG_DIR} = $config{runtimedir} || '.'; # avoid undefined warning

if ($debug) {
    # Use a bogus name for the logfile if debugging -- it doesn't matter.
    $condor_log = 'job.log';
} else {
    $condor_log = File::Temp::tempnam($grami{joboption_directory}, 'log.');
}

run_rte0();
create_shell_wrapper();
create_condor_job_description();
submit_condor_job();
warn "$0: job submitted successfully\n",
     "----- exiting $0 -----\n";
exit;

##############################################################################
## Function Definitons
##############################################################################

sub parse_grami {
    local @ARGV = $_[0];
    warn "$0: ----- begin grami file ($_[0]) -----\n";
    while (my $line = <>) {
        chomp $line;

        # Dump every line of the grami file into the log.
        warn "$0: $line\n";

        my ($name, $value) = split /=/, $line, 2;
        next if !$name;

        # Remove outer layer of single quotes.  Backslash escaped single quotes
        # are stripped of their backslashes, and strings protected by single
        # quotes are stripped of the single quotes.  This is supposed to work
        # exactly like Bourne shell quote removal:
        #
        #   foo'bar'     --> foobar
        #   foo\''bar'\' --> foo'bar'
        #
        {
            no warnings 'uninitialized';
            $value =~ s/(?:\\('))?'([^']*)'(?:\\('))?/$1$2$3/g;
        }

        # The variable names are case insensitive, so lowercase them and
        # remember to always refer to them by their lowercase names!
        $grami{lc $name} = $value;
    }
    warn "$0: ----- end grami file ($_[0]) -----\n";
}

#
# Run RTE stage 0 scripts. Also update %grami hash to match the variables set by the
# scripts (but only those beginnig with joboption_)
#
sub run_rte0 {
    for (my $i = 0; my $r = $grami{"joboption_runtime_$i"}; $i++) {

        # Shell script which will source RTE scripts
        my $script = '';

        # construct randon string to use as separator
        my @chars = ("0".."9","A".."Z");
        my $sep = "----";
        $sep .= $chars[rand(@chars)] for 1..16;
        $sep .= "----";

        # initialize joboption_* variables
        $script .= "$_=\Q$grami{$_}\E\n" for (grep /^joboption_[A-Za-z0-9_]+$/, keys %grami);

        # source RTE script
        $script .= ". \Q$ENV{RUNTIME_CONFIG_DIR}/$r\E 0 1>&2\n";
        $script .= "ret=\$?\n";

        # print joboption_* variables to stdout. First export them, then use perl for printing.
        $script .= 'export $(set | sed -n "s/^\(joboption_[A-Za-z0-9_]*\)=.*/\1/p")';
        $script .= "\n/usr/bin/perl -we 'print \"\$_=\$ENV{\$_}$sep\" for grep /^joboption_[A-Za-z0-9_]+\$/, keys \%ENV'\n";
        $script .= "echo STATUS=\$ret\n";

        # write script to temporary file and run it.
        my $dir = $ENV{TMPDIR} || File::Spec->tmpdir;
        my ($fh, $fname) = tempfile('condor_XXXXXXXX', DIR => $dir);
        print $fh $script or die "$0: write $fname: $!\n";
        close $fh           or die "$0: close $fname: $!\n";
        my $output = `/bin/sh $fname`;
        unlink $fname;

        # is output complete?
        if (not $output =~ m/^(?:joboption_\w+=.*?$sep)+STATUS=(\d+)$/s) {
            die "$0: ERROR: RTE script $r exited unexpectedly\n";
        }

        delete $grami{$_} for grep /^joboption_/, keys %grami;
        $grami{$1} = $2 while ($output =~ m/(\w+)=(.*?)$sep/sg);
    }
}

# This function extracts the first string from each line of the file given as
# argument. Handles files like job.*.output which contain two strings
# separated by a space on each line. The second string can be emty. Spaces
# and backslashes within the strings are escaped with a backslash. See the
# FileData class in src/services/a-rex/grid-manager/files/info_types.cpp for
# the c++ implementation.

# OBS: The same can be achieved simpler with the shell built-in 'read'
sub get_outputfiles {
    my $file = shift;
    my @names = ();
    open(F, "<$file") or die "$file: $!\n";
    while (<F>) {
        chomp;
        my @chars = split "";
        my $esc = 0;
        my $str = "";
        for (my $curr = shift @chars; defined $curr; $curr = shift @chars) {
            if ($esc) {
                $esc = 0;
                $str .= $curr;
                next;
            }
            last if $curr eq " ";
            $str .= $curr unless $esc = ($curr eq "\\");
        }
        push @names, $str;
    }
    return @names;
}

#
# Creates a shell script that:
#
#  (1) Sources the runtime scripts with argument "0" before evaluating the job
#      executable.  (This is in case the job refers to variables set by the
#      runtime scripts.)  TODO: should variables be expanded in
#      joboption_runtime_0?
#
#  (2) Sources the runtime scripts with argument "1" before running the job.
#
#  (3) Runs the job, redirecting output as requested in the xRSL.
#
#  (4) Sources the runtime scripts with argument "2" after running the job.
#
#  (5) Exits with the value returned by the job executable in step (3).
#
sub create_shell_wrapper {
    # Create the shell commands to run runtime environment files (stages 1-2).
    my ($setrte1, $setrte2) = ('', '');
    if (notnull($grami{joboption_runtime_0})) {
        for (my $i = 0; notnull(my $r = $grami{"joboption_runtime_$i"}); $i++) {
            $setrte1 .= qq{. "\$RUNTIME_CONFIG_DIR/$r" 1\n};
            $setrte2 .= qq{. "\$RUNTIME_CONFIG_DIR/$r" 2\n};
        }
    }

    # Set $real_exe to the path to the job executable (environment variables
    # expanded).  $exewrapper contains the script which will be submitted to condor.
    $real_exe = $grami{joboption_arg_0};
    $exewrapper = File::Temp::tempnam($grami{joboption_directory}, "condorjob.sh.");

    # Get the name of the stdout file.
    my $stdout = notnull($grami{joboption_stdout}) ?
                 $grami{joboption_stdout} : '/dev/null';
    $stdout =~ s{^\Q$grami{joboption_directory}\E/*}{};

    # Get the name of the stderr file.
    my $stderr = notnull($grami{joboption_stderr}) ?
                 $grami{joboption_stderr} : '/dev/null';
    $stderr =~ s{^\Q$grami{joboption_directory}\E/*}{};

    # Start creating the output script.  Note that the script is created
    # in-memory, instead of being written to file, part by part.  This is
    # because we want to test for all I/O errors, and having just a single
    # write means that there is only one place we have to test for write
    # errors.
    my $output = "#!/bin/sh\n";

    # If the custom RSL attribute 'wrapperdebug' is set, enable command
    # tracing (set -x) and list all files in the session directory.  (This
    # output is sent to stderr.)
    if (notnull($grami{joboption_rsl_wrapperdebug})) {
        $output .= "set -x\nexec &>\Q$stderr\E\nls -la\n";
    }

    # set HOME to the working directory of the job
    $output .= "HOME=\`pwd\`\nexport HOME\n";

    # Overide umask of execution node
    $output .= "umask 077\n";

    # Source runtime scripts with argument 1.
    $output .= $setrte1;

    # Enable the executable bit for non-preinstalled executables.
    if ($real_exe !~ m{^/}) {
        $output .= "eval chmod +x \"\Q$real_exe\E\"\n";
    }

    # Incomplete job command; arguments may follow.
    $output .= "eval \"\Q$real_exe\E\"";

    # Add optional arguments to the command line.
    if (defined $grami{joboption_arg_1}) {
        for (my $i = 1; defined(my $arg = $grami{"joboption_arg_$i"}); $i++) {
            $output .= $arg ne '' ? " \"\Q$arg\E\"" : " ''";
        }
    }

    # Redirect stdout/stderr.  These variables are always set to something
    # (/dev/null if unspecified), so it's safe to unconditionally add these
    # redirections.
    $output .= " >\Q$stdout\E";
    # If we're debugging the wrapper script, we don't do stderr redirection.
    if (!notnull($grami{joboption_rsl_wrapperdebug})) {
        if ($stdout eq $stderr) {
            # We're here if stdout and stderr is redirected to the same file.
            # This will happen when (join = yes) in the xRSL.
            $output .= ' 2>&1';
        } else {
            $output .= " 2>\Q$stderr\E";
        }
    }

    # Always a newline to terminate the job command.
    # Preserve the job's exit code.
    # Run runtime environment files with argument 2.
    $output .= "\n_exitcode=\$?\n$setrte2";

    # Now generate code that removes everything but the requested output.
    # Note that, bashims have been avoided so that there are less strict
    # requirements on /bin/sh on the execute nodes.  (Note that the file
    # utilities used (mkdir, dirname, find, etc.) may still be
    # GNU-centric.  TODO: fix this if we're to support non-x86-Linux.)
    $output .= <<EOF;
find ./ -type l -exec rm -f "{}" ";"
find ./ -type f -exec chmod u+w "{}" ";"
EOF
    my $outputlist = "$grami{joboption_controldir}/job.$grami{joboption_gridid}.output";
    if (-e $outputlist) {
        my @fileslst = grep {$_ ne ""} get_outputfiles($outputlist);
        # Add condor_log (with path stripped) to the list of files to keep.
        push @fileslst, (my $basename) = $condor_log =~ m{([^/]+)$};
        # Remove leading backslashes, if any
        s/^\/*// for @fileslst;
        # Make it safe for shell by replacing single quotes with '\''
        s/'/'\\''/g for @fileslst;

        # Protect from deleting output files including those in the dynamic list
        for my $file ( @fileslst ) {
            if ($file =~ s/^@//) {
                $output .= "dynlist='$file'\n";
                $output .= <<'EOF'
chmod -R u-w "./$dynlist" 2>/dev/null
cat "./$dynlist" | while read name rest; do
  chmod -R u-w "./$name" 2>/dev/null
done
EOF
            } else {
                $output .= "chmod u-w './$file' 2>/dev/null\n";
            }
        }
        $output .= <<EOF;
find ./ -type f -perm +200 -exec rm -f "{}" ";"
find ./ -type f -exec chmod u+w "{}" ";"
EOF
    }

    # Exit with the job's exit code.
    $output .= "exit \$_exitcode\n";

    # Create the actual shell script from $output.
    open EXE, ">$exewrapper"              or die "$0: creat $exewrapper: $!\n";
    print EXE $output                     or die "$0: write $exewrapper: $!\n";
    close EXE                             or die "$0: close $exewrapper: $!\n";
    chmod 0755, $exewrapper               or die "$0: chmod $exewrapper: $!\n";

    # Log the Condor job submission script in gmlog/errors.
    unless ($debug) {
        warn "$0: ----- begin wrapper script ($exewrapper) -----\n";
        warn "$0: $_\n" for split /\n/, $output;
        warn "$0: ----- end wrapper script ($exewrapper) -----\n";
    }
}

#
# Create a Condor job description that submits the wrapper script created
# above.  The Condor job description should mirror the xRSL as much as
# possible.
#
sub create_condor_job_description {
    # As above, the job description is created in-memory, so that only one I/O
    # operation has to be done when writing to disk.
    my $output = "Executable = $exewrapper\n" .
                 "Input = $grami{joboption_stdin}\n";

    $output .= "Log = $condor_log\n";
    $output .= "Output = $grami{joboption_directory}.comment\n";
    $output .= "Error = $grami{joboption_directory}.comment\n";

    my @requirements = ();

    if (notnull($grami{joboption_queue})) {
        my $queue = $grami{joboption_queue};
        $output .= "+NordugridQueue = $queue\n";
    }

    if (notnull($config{condor_rank})) {
        $output .= "Rank = $config{condor_rank}\n";
    }

    if (notnull($config{condor_requirements})) {
        $config{condor_requirements} =~ s/\[separator\]//g;
        push @requirements, $config{condor_requirements};
    }

    # This is a custom RSL attribute used for debugging.  If the xRSL contains
    # (machine = foo), the job will only run on machine "foo".
    if (notnull($grami{joboption_rsl_machine})) {
        push @requirements, "Machine == \"$grami{joboption_rsl_machine}\"";
    }

    if (@requirements) {
        $output .= "Requirements = (" . (join ") && (", @requirements) . ")\n";
    }

    # Option to force Condor to transfer of input and output files by itself
    if ($config{shared_filesystem} =~ /^no/i && notnull($grami{joboption_rsl_inputfiles})) {
        chomp($grami{joboption_rsl_inputfiles});
        my @tmp = split / /, $grami{joboption_rsl_inputfiles};
        $output .= "Transfer_input_files = ";
        for (my $i = 0; $i < @tmp; $i += 2) {
            $output .= ',' if $i > 0;
            $output .= $tmp[$i];
        }
        $output .= "\n";
        $output .= "should_transfer_files = YES\n";
        # No need to specify output files explicitly. Condor will transfer all
        # files from the job directory that were created or modified by the job
        $output .= "When_to_transfer_output = ON_EXIT_OR_EVICT\n";

        if (notnull($grami{joboption_rsl_disk})) {
            push @requirements, "Disk >= " . ($grami{joboption_rsl_disk} * 1024);
        }
    }

    if (notnull($grami{joboption_env_0})) {
        $output .= "Environment = ";
        my $has_globalid = '';
        my $i;
        for ($i = 0; notnull($grami{"joboption_env_$i"}); $i++) {
            $output .= ";" if $i > 0;
            $output .= $grami{"joboption_env_$i"};
            $has_globalid = 1 if $grami{"joboption_env_$i"} =~ m/^GRID_GLOBAL_JOBID=/;
        }
        # guess globalid in case not already provided and export it to the job
        if (not $has_globalid) {
            my $hostname = $config{hostname} || hostname();
            my $gm_port = $config{gm_port} || 2811;
            my $gm_mount_point = $config{gm_mount_point} || "/jobs";
            $output .= ";" if $i > 0;
            $output .= "GRID_GLOBAL_JOBID=gsiftp://$hostname:$gm_port$gm_mount_point/$grami{joboption_gridid}";
        }
        $output .= "\n";
    }

    my $remove="FALSE";
    if (notnull($grami{joboption_cputime})) {
        $output .= "+JobCpuLimit = $grami{joboption_cputime}\n";
        $remove .= " || RemoteUserCpu + RemoteSysCpu > JobCpuLimit";
        warn "$0: Setting CPU limit\n";
    }
    if (notnull($grami{joboption_walltime})) {
        $output .= "+JobTimeLimit = $grami{joboption_walltime}\n";
        $remove .= " || RemoteWallClockTime > JobTimeLimit";
        warn "$0: Setting time limit\n";
    }
    # Set a default memory limit
    $grami{joboption_memory} ||= $config{nodememory} || 1000;
    $output .= "+JobMemoryLimit = ".int(1024*$grami{joboption_memory})."\n";
    $remove .= " || ImageSize > JobMemoryLimit";
    warn "$0: Setting memory limit\n";

    $output .= "GetEnv = True\n" .
               "Universe = vanilla\n" .
               "Notification = Always\n" .
#              "When_to_transfer_output = ON_EXIT\n" .
               "Periodic_remove = $remove\n" .
               "Queue\n";

    if ($debug) {
        print $output;
    } else {
        my $cmd_fh;
        ($cmd_fh, $cmd_filename) = tempfile('XXXXXXXX',
                                            DIR => $grami{joboption_directory},
                                            SUFFIX => '.cmd');
        print $cmd_fh $output  or die "$0: write $cmd_filename: $!\n";
        close $cmd_fh          or die "$0: close $cmd_filename: $!\n";

        # Log the Condor job submission script in gmlog/errors.
        warn "$0: ----- begin condor job description ($cmd_filename) -----\n";
        warn "$0: $_\n" for split /\n/, $output;
        warn "$0: ----- end condor job description ($cmd_filename) -----\n";
    }
}

sub submit_condor_job {
    return if $debug;
    chdir $grami{joboption_directory}
      or die "$0: chdir $grami{joboption_directory}: $!\n";

    my $condor_submit_exe = "$condor_bin_path/condor_submit";
    warn "$0: running $condor_submit_exe $cmd_filename\n";
    my $submit_out = `\Q$condor_submit_exe\E \Q$cmd_filename\E 2>&1`;
    #my $submit_out = `\Qcat\E \Q$cmd_filename\E 2>&1; exit 22`;
    my $err = $?;
    warn "$0: $_\n" for split /\n/, $submit_out;
    die "$0: condor_submit failed!\n" if $err;

    warn "$0: appending local job id to grami file $gramifile\n";
    my ($localid) = $submit_out =~ /submitted to cluster (\d+)\./;
    open GRAMI, ">>$gramifile"                      or die "$0: $gramifile: $!";
    print GRAMI "joboption_jobid=$localid.condor\n" or die "$0: $gramifile: $!";
    print GRAMI "condor_log=$condor_log\n"          or die "$0: $gramifile: $!";
    close GRAMI                                     or die "$0: $gramifile: $!";
}

sub notnull {
    return defined $_[0] && $_[0] ne '';
}