#!/usr/bin/env perl

##**************************************************************
##
## Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************




###########################################################################
#
#  This is condor_run, a simple front-end to condor_submit for submitting
#  vanilla jobs to Condor.
#
#  usage: condor_run "shell-cmd"
#
#  examples:
#
#    condor_run "zcat data.gz | awk -f myscript.awk > results.txt"
#    condor_run "cc -o myprog myprog.c"
#    condor_run "make"
#    condor_run "condor_compile cc -o myprog myprog.c"
#
#  condor_run by Jim Basney <jbasney@cs.wisc.edu> 4/3/2000
#
###########################################################################

my @command;
my @appends;
my $universe = 'vanilla';
if (!defined($ARGV[0])) {
    goto HELP;
}
while ( $_ = shift( @ARGV ) ) {
  SWITCH: {
      if ( /^-h.*/ ) {
	HELP:
	  print "usage: $0 [-u <universe>] [-a command]* \"shell-cmd\"\n";
	  print "\twhere shell-cmd is any Unix shell statement.\n";
	  print "\t-u lets you set the universe to which the job is submitted.\n";
	  print "\t   The default is vanilla.\n";
	  print "\t-a lets you add additional submit commands\n";
	  print "\t   e.g. -a request_memory=1024\n";
	  print "\tEnvironment variables CONDOR_ARCH, CONDOR_OPSYS, and\n";
	  print "\tCONDOR_REQUIREMENTS may specify remote machine's Arch,\n";
	  print "\tOpSys, and any additional requirements.\n";
	  exit 1;
      }
      if ( /^-u(.*)/ ) {
          $universe = $1 ? $1 : shift(@ARGV);
          next SWITCH;
      }
      if ( /^-a.*/ ) {
          push(@appends, shift(@ARGV));
          next SWITCH;
      }

      # Anything passed in that isn't -u or -a is part of the command to run
      push(@command, $_);
    }
}
if (!@command) {
    goto HELP;
}

# grab current working directory for initial dir in system using automounter
$pwd = `pwd`;
chomp $pwd;

# set up environment for running something in the current directory in case
# they want to run something in the current working directory and they
# don't specify a "./" infront of it.
$ENV{'PATH'} .= ":.";

# setup cleanup subroutine and error handlers
sub cleanfiles {
    unlink ".condor_run.$$", ".condor_submit.$$", ".condor_log.$$";
    unlink ".condor_out.$$", ".condor_error.$$";
}

sub abort {
    `condor_rm $cluster 2>&1 > /dev/null` if defined($cluster);
    &cleanfiles;
    die @_;
}

sub handler {
    local($sig) = @_;
    &abort("Killed by SIG$sig.\n");
}

$SIG{'HUP'} = 'handler';
$SIG{'INT'} = 'handler';
$SIG{'QUIT'} = 'handler';
$SIG{'TERM'} = 'handler';

# use the user's shell in the script we submit to Condor
$shell = $ENV{'SHELL'};
$shell = "/bin/sh" if (!defined($shell));

# get any requirements from the user's environment
$requirements = $ENV{'CONDOR_REQUIREMENTS'};
$arch = $ENV{'CONDOR_ARCH'};
if (defined($arch)) {
    $requirements .= " && " if (defined($requirements));
    $requirements .= "Arch == \"" . $arch . "\"";
}
$opsys = $ENV{'CONDOR_OPSYS'};
if (defined($opsys)) {
    $requirements .= " && " if (defined($requirements));
    $requirements .= "OpSys == \"" . $opsys . "\"";
}

# create a shell script containing the user's command
open(CMD, ">.condor_run.$$") ||
    &abort("Can't create temporary (CMD) file in current directory.\n");
print CMD "#!", $shell, "\n";
foreach $arg (@command) {
    print CMD $arg, " ";
}
print CMD "\n";
close(CMD) ||
    &abort("Failed to write temporary (CMD) file in current directory.\n");
# fix the permissions to something nice.
chmod 0700, ".condor_run.$$";

# create a job description file to submit the shell script to Condor
open(JDF, ">.condor_submit.$$") ||
    &abort("Can't create temporary (JDF) file in current directory.\n");
print JDF "universe = $universe\n";
print JDF "executable = .condor_run.$$\n";
print JDF "initialdir = $pwd\n";
print JDF "notification = NEVER\n";
print JDF "log = .condor_log.$$\n";
print JDF "output = .condor_out.$$\n";
print JDF "error = .condor_error.$$\n";
print JDF "getenv = True\n";
print JDF "requirements = ", $requirements, "\n" if (defined($requirements));
foreach my $append (@appends) {
    print JDF $append . "\n";
}
print JDF "queue\n";
close(JDF) ||
    &abort("Failed to write temporary (JDF) file in current directory.\n");

# submit the job; $cluster contains cluster number if successful
open(SUBMIT, "condor_submit .condor_submit.$$ 2>&1 |") ||
    &abort("Failed to run condor_submit.  Please check your path.\n");
while(<SUBMIT>) {
    if (/^1 job\(s\) submitted to cluster (\d+)./) {
	($cluster) = $1;
    } elsif (/WARNING.*Invalid log file/ ||
	     /WARNING.*is not writable by condor/) {
	print STDERR $_;
	&abort("Failed to submit Condor job.\n");
    } else {
	$submit_errors .= $_;
    }
}

if (!close(SUBMIT) || !defined($cluster)) {
    print STDERR $submit_errors;
    &abort("Failed to submit Condor job.\n");
}

# watch the Condor log to see when the job completes
$done = 0;
$status = 0;
while (!$done) {
    sleep 5;
    open(LOG, "<.condor_log.$$") ||
    &abort("Failed to open Condor log file.\n");
    while (<LOG>) {
	if (/termination \(return value (\d+)\)/) {
	    $status = $1;
	    $done = 1;
	} elsif (/termination \(signal (\d+)\)/) {
	    print STDERR "Condor job killed by signal ", $1 , ".\n";
	    $done = 1;
        } elsif ( /^009/ && /Job was aborted by the user/ ) {
	    print STDERR "The job was aborted by the user";
	    $done = 1;
	}
    }
    close(LOG) || &abort("Failed to close Condor log file.\n");
}
undef $cluster;			# indicate that job is no longer in queue

# display stdout to the terminal
open(OUTPUT, "<.condor_out.$$") || &abort("Failed to open job output file.\n");
while (<OUTPUT>) {
    print STDOUT;
}
close(OUTPUT) || &abort("Failed to open job output file.\n");

# display stderr to the terminal
open(ERROR, "<.condor_error.$$") || &abort("Failed to open job error file.\n");
while (<ERROR>) {
    print STDERR;
}
close(ERROR) || &abort("Failed to open job output file.\n");

&cleanfiles;			# cleanup temporary files

exit $status;			# exit with job's exit status