#!/usr/bin/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;
}
my $getenv_flag = 0;

while ( $_ = shift( @ARGV ) ) {
	SWITCH: {
		if ( /^-h.*/ ) {
			HELP:
			print "usage: $0 [-getenv] [-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;
		}
		if ( /^-getenv/ ) {
			$getenv_flag = 1;
			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";
if ($getenv_flag == 1) {
	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 -terse .condor_submit.$$ |") ||
&abort("Failed to run condor_submit.  Please check your path.\n");
while(<SUBMIT>) {
	if (/^(\d+)\.\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/ ) {
			print STDERR "The job was aborted";
			$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
