#!/usr/bin/perl
# -*- mode: perl; perl-indent-level: 2; indent-tabs-mode: nil; -*-

=head1 NAME

crunch-job: Execute job steps, save snapshots as requested, collate output.

=head1 SYNOPSIS

Obtain job details from Arvados, run tasks on compute nodes (typically
invoked by scheduler on controller):

 crunch-job --job x-y-z

Obtain job details from command line, run tasks on local machine
(typically invoked by application or developer on VM):

 crunch-job --job '{"script_version":"/path/to/tree","script":"scriptname",...}'

=head1 OPTIONS

=over

=item --force-unlock

If the job is already locked, steal the lock and run it anyway.

=item --git-dir

Path to .git directory where the specified commit is found.

=item --job-api-token

Arvados API authorization token to use during the course of the job.

=item --no-clear-tmp

Do not clear per-job/task temporary directories during initial job
setup. This can speed up development and debugging when running jobs
locally.

=back

=head1 RUNNING JOBS LOCALLY

crunch-job's log messages appear on stderr along with the job tasks'
stderr streams. The log is saved in Keep at each checkpoint and when
the job finishes.

If the job succeeds, the job's output locator is printed on stdout.

While the job is running, the following signals are accepted:

=over

=item control-C, SIGINT, SIGQUIT

Save a checkpoint, terminate any job tasks that are running, and stop.

=item SIGALRM

Save a checkpoint and continue.

=item SIGHUP

Refresh node allocation (i.e., check whether any nodes have been added
or unallocated) and attributes of the Job record that should affect
behavior (e.g., cancel job if cancelled_at becomes non-nil).

=back

=cut


use strict;
use POSIX ':sys_wait_h';
use Fcntl qw(F_GETFL F_SETFL O_NONBLOCK);
use Arvados;
use Digest::MD5 qw(md5_hex);
use Getopt::Long;
use IPC::Open2;
use IO::Select;
use File::Temp;
use Fcntl ':flock';

$ENV{"TMPDIR"} ||= "/tmp";
unless (defined $ENV{"CRUNCH_TMP"}) {
  $ENV{"CRUNCH_TMP"} = $ENV{"TMPDIR"} . "/crunch-job";
  if ($ENV{"USER"} ne "crunch" && $< != 0) {
    # use a tmp dir unique for my uid
    $ENV{"CRUNCH_TMP"} .= "-$<";
  }
}
$ENV{"JOB_WORK"} = $ENV{"CRUNCH_TMP"} . "/work";
$ENV{"CRUNCH_INSTALL"} = "$ENV{CRUNCH_TMP}/opt";
$ENV{"CRUNCH_WORK"} = $ENV{"JOB_WORK"}; # deprecated
mkdir ($ENV{"JOB_WORK"});

my $arv_cli;

if (defined $ENV{"ARV_CLI"}) {
  $arv_cli = $ENV{"ARV_CLI"};
}
else {
  $arv_cli = 'arv';
}

my $force_unlock;
my $git_dir;
my $jobspec;
my $job_api_token;
my $no_clear_tmp;
my $resume_stash;
GetOptions('force-unlock' => \$force_unlock,
           'git-dir=s' => \$git_dir,
           'job=s' => \$jobspec,
           'job-api-token=s' => \$job_api_token,
           'no-clear-tmp' => \$no_clear_tmp,
           'resume-stash=s' => \$resume_stash,
    );

if (defined $job_api_token) {
  $ENV{ARVADOS_API_TOKEN} = $job_api_token;
}

my $have_slurm = exists $ENV{SLURM_JOBID} && exists $ENV{SLURM_NODELIST};
my $job_has_uuid = $jobspec =~ /^[-a-z\d]+$/;
my $local_job = !$job_has_uuid;


$SIG{'USR1'} = sub
{
  $main::ENV{CRUNCH_DEBUG} = 1;
};
$SIG{'USR2'} = sub
{
  $main::ENV{CRUNCH_DEBUG} = 0;
};



my $arv = Arvados->new('apiVersion' => 'v1');
my $local_logfile;

my $User = $arv->{'users'}->{'current'}->execute;

my $Job = {};
my $job_id;
my $dbh;
my $sth;
if ($job_has_uuid)
{
  $Job = $arv->{'jobs'}->{'get'}->execute('uuid' => $jobspec);
  if (!$force_unlock) {
    if ($Job->{'is_locked_by_uuid'}) {
      croak("Job is locked: " . $Job->{'is_locked_by_uuid'});
    }
    if ($Job->{'success'} ne undef) {
      croak("Job 'success' flag (" . $Job->{'success'} . ") is not null");
    }
    if ($Job->{'running'}) {
      croak("Job 'running' flag is already set");
    }
    if ($Job->{'started_at'}) {
      croak("Job 'started_at' time is already set (" . $Job->{'started_at'} . ")");
    }
  }
}
else
{
  $Job = JSON::decode_json($jobspec);

  if (!$resume_stash)
  {
    map { croak ("No $_ specified") unless $Job->{$_} }
    qw(script script_version script_parameters);
  }

  $Job->{'is_locked_by_uuid'} = $User->{'uuid'};
  $Job->{'started_at'} = gmtime;

  $Job = $arv->{'jobs'}->{'create'}->execute('job' => $Job);

  $job_has_uuid = 1;
}
$job_id = $Job->{'uuid'};

my $keep_logfile = $job_id . '.log.txt';
$local_logfile = File::Temp->new();

$Job->{'runtime_constraints'} ||= {};
$Job->{'runtime_constraints'}->{'max_tasks_per_node'} ||= 0;
my $max_ncpus = $Job->{'runtime_constraints'}->{'max_tasks_per_node'};


Log (undef, "check slurm allocation");
my @slot;
my @node;
# Should use $ENV{SLURM_TASKS_PER_NODE} instead of sinfo? (eg. "4(x3),2,4(x2)")
my @sinfo;
if (!$have_slurm)
{
  my $localcpus = 0 + `grep -cw ^processor /proc/cpuinfo` || 1;
  push @sinfo, "$localcpus localhost";
}
if (exists $ENV{SLURM_NODELIST})
{
  push @sinfo, `sinfo -h --format='%c %N' --nodes=\Q$ENV{SLURM_NODELIST}\E`;
}
foreach (@sinfo)
{
  my ($ncpus, $slurm_nodelist) = split;
  $ncpus = $max_ncpus if $max_ncpus && $ncpus > $max_ncpus;

  my @nodelist;
  while ($slurm_nodelist =~ s/^([^\[,]+?(\[.*?\])?)(,|$)//)
  {
    my $nodelist = $1;
    if ($nodelist =~ /\[((\d+)(-(\d+))?(,(\d+)(-(\d+))?)*)\]/)
    {
      my $ranges = $1;
      foreach (split (",", $ranges))
      {
	my ($a, $b);
	if (/(\d+)-(\d+)/)
	{
	  $a = $1;
	  $b = $2;
	}
	else
	{
	  $a = $_;
	  $b = $_;
	}
	push @nodelist, map {
	  my $n = $nodelist;
	  $n =~ s/\[[-,\d]+\]/$_/;
	  $n;
	} ($a..$b);
      }
    }
    else
    {
      push @nodelist, $nodelist;
    }
  }
  foreach my $nodename (@nodelist)
  {
    Log (undef, "node $nodename - $ncpus slots");
    my $node = { name => $nodename,
		 ncpus => $ncpus,
		 losing_streak => 0,
		 hold_until => 0 };
    foreach my $cpu (1..$ncpus)
    {
      push @slot, { node => $node,
		    cpu => $cpu };
    }
  }
  push @node, @nodelist;
}



# Ensure that we get one jobstep running on each allocated node before
# we start overloading nodes with concurrent steps

@slot = sort { $a->{cpu} <=> $b->{cpu} } @slot;



my $jobmanager_id;
if ($job_has_uuid)
{
  # Claim this job, and make sure nobody else does
  unless ($Job->update_attributes('is_locked_by_uuid' => $User->{'uuid'}) &&
          $Job->{'is_locked_by_uuid'} == $User->{'uuid'}) {
    croak("Error while updating / locking job");
  }
  $Job->update_attributes('started_at' => scalar gmtime,
                          'running' => 1,
                          'success' => undef,
                          'tasks_summary' => { 'failed' => 0,
                                               'todo' => 1,
                                               'running' => 0,
                                               'done' => 0 });
}


Log (undef, "start");
$SIG{'INT'} = sub { $main::please_freeze = 1; };
$SIG{'QUIT'} = sub { $main::please_freeze = 1; };
$SIG{'TERM'} = \&croak;
$SIG{'TSTP'} = sub { $main::please_freeze = 1; };
$SIG{'ALRM'} = sub { $main::please_info = 1; };
$SIG{'CONT'} = sub { $main::please_continue = 1; };
$SIG{'HUP'} = sub { $main::please_refresh = 1; };

$main::please_freeze = 0;
$main::please_info = 0;
$main::please_continue = 0;
$main::please_refresh = 0;
my $jobsteps_must_output_keys = 0;	# becomes 1 when any task outputs a key

grep { $ENV{$1} = $2 if /^(NOCACHE.*?)=(.*)/ } split ("\n", $$Job{knobs});
$ENV{"CRUNCH_JOB_UUID"} = $job_id;
$ENV{"JOB_UUID"} = $job_id;


my @jobstep;
my @jobstep_todo = ();
my @jobstep_done = ();
my @jobstep_tomerge = ();
my $jobstep_tomerge_level = 0;
my $squeue_checked;
my $squeue_kill_checked;
my $output_in_keep = 0;
my $latest_refresh = scalar time;



if (defined $Job->{thawedfromkey})
{
  thaw ($Job->{thawedfromkey});
}
else
{
  my $first_task = $arv->{'job_tasks'}->{'create'}->execute('job_task' => {
    'job_uuid' => $Job->{'uuid'},
    'sequence' => 0,
    'qsequence' => 0,
    'parameters' => {},
                                                          });
  push @jobstep, { 'level' => 0,
		   'failures' => 0,
                   'arvados_task' => $first_task,
		 };
  push @jobstep_todo, 0;
}


if (!$have_slurm)
{
  must_lock_now("$ENV{CRUNCH_TMP}/.lock", "a job is already running here.");
}


my $build_script;


$ENV{"CRUNCH_SRC_COMMIT"} = $Job->{script_version};

my $skip_install = ($local_job && $Job->{script_version} =~ m{^/});
if ($skip_install)
{
  if (!defined $no_clear_tmp) {
    my $clear_tmp_cmd = 'rm -rf $JOB_WORK $CRUNCH_TMP/opt $CRUNCH_TMP/src*';
    system($clear_tmp_cmd) == 0
	or croak ("`$clear_tmp_cmd` failed: ".($?>>8));
  }
  $ENV{"CRUNCH_SRC"} = $Job->{script_version};
  for my $src_path ("$ENV{CRUNCH_SRC}/arvados/sdk/python") {
    if (-d $src_path) {
      system("virtualenv", "$ENV{CRUNCH_TMP}/opt") == 0
          or croak ("virtualenv $ENV{CRUNCH_TMP}/opt failed: exit ".($?>>8));
      system ("cd $src_path && ./build.sh && \$CRUNCH_TMP/opt/bin/python setup.py install")
          == 0
          or croak ("setup.py in $src_path failed: exit ".($?>>8));
    }
  }
}
else
{
  do {
    local $/ = undef;
    $build_script = <DATA>;
  };
  Log (undef, "Install revision ".$Job->{script_version});
  my $nodelist = join(",", @node);

  if (!defined $no_clear_tmp) {
    # Clean out crunch_tmp/work, crunch_tmp/opt, crunch_tmp/src*

    my $cleanpid = fork();
    if ($cleanpid == 0)
    {
      srun (["srun", "--nodelist=$nodelist", "-D", $ENV{'TMPDIR'}],
	    ['bash', '-c', 'if mount | grep -q $JOB_WORK/; then sudo /bin/umount $JOB_WORK/* 2>/dev/null; fi; sleep 1; rm -rf $JOB_WORK $CRUNCH_TMP/opt $CRUNCH_TMP/src*']);
      exit (1);
    }
    while (1)
    {
      last if $cleanpid == waitpid (-1, WNOHANG);
      freeze_if_want_freeze ($cleanpid);
      select (undef, undef, undef, 0.1);
    }
    Log (undef, "Clean-work-dir exited $?");
  }

  # Install requested code version

  my @execargs;
  my @srunargs = ("srun",
		  "--nodelist=$nodelist",
		  "-D", $ENV{'TMPDIR'}, "--job-name=$job_id");

  $ENV{"CRUNCH_SRC_COMMIT"} = $Job->{script_version};
  $ENV{"CRUNCH_SRC"} = "$ENV{CRUNCH_TMP}/src";

  my $commit;
  my $git_archive;
  my $treeish = $Job->{'script_version'};

  # If we're running under crunch-dispatch, it will have pulled the
  # appropriate source tree into its own repository, and given us that
  # repo's path as $git_dir. If we're running a "local" job, and a
  # script_version was specified, it's up to the user to provide the
  # full path to a local repository in Job->{repository}.
  #
  # TODO: Accept URLs too, not just local paths. Use git-ls-remote and
  # git-archive --remote where appropriate.
  #
  # TODO: Accept a locally-hosted Arvados repository by name or
  # UUID. Use arvados.v1.repositories.list or .get to figure out the
  # appropriate fetch-url.
  my $repo = $git_dir || $ENV{'CRUNCH_DEFAULT_GIT_DIR'} || $Job->{'repository'};

  $ENV{"CRUNCH_SRC_URL"} = $repo;

  if (-d "$repo/.git") {
    # We were given a working directory, but we are only interested in
    # the index.
    $repo = "$repo/.git";
  }

  # If this looks like a subversion r#, look for it in git-svn commit messages

  if ($treeish =~ m{^\d{1,4}$}) {
    my $gitlog = `git --git-dir=\Q$repo\E log --pretty="format:%H" --grep="git-svn-id:.*\@"\Q$treeish\E" " master`;
    chomp $gitlog;
    if ($gitlog =~ /^[a-f0-9]{40}$/) {
      $commit = $gitlog;
      Log (undef, "Using commit $commit for script_version $treeish");
    }
  }

  # If that didn't work, try asking git to look it up as a tree-ish.

  if (!defined $commit) {
    my $found = `git --git-dir=\Q$repo\E rev-list -1 ''\Q$treeish\E`;
    chomp $found;
    if ($found =~ /^[0-9a-f]{40}$/s) {
      $commit = $found;
      if ($commit ne $treeish) {
	# Make sure we record the real commit id in the database,
	# frozentokey, logs, etc. -- instead of an abbreviation or a
	# branch name which can become ambiguous or point to a
	# different commit in the future.
	$ENV{"CRUNCH_SRC_COMMIT"} = $commit;
	Log (undef, "Using commit $commit for tree-ish $treeish");
        if ($commit ne $treeish) {
          $Job->{'script_version'} = $commit;
          !$job_has_uuid or
              $Job->update_attributes('script_version' => $commit) or
              croak("Error while updating job");
        }
      }
    }
  }

  if (defined $commit) {
    $ENV{"CRUNCH_SRC_COMMIT"} = $commit;
    @execargs = ("sh", "-c",
		 "mkdir -p $ENV{CRUNCH_INSTALL} && cd $ENV{CRUNCH_TMP} && perl -");
    $git_archive = `git --git-dir=\Q$repo\E archive ''\Q$commit\E`;
  }
  else {
    croak ("could not figure out commit id for $treeish");
  }

  my $installpid = fork();
  if ($installpid == 0)
  {
    srun (\@srunargs, \@execargs, {}, $build_script . $git_archive);
    exit (1);
  }
  while (1)
  {
    last if $installpid == waitpid (-1, WNOHANG);
    freeze_if_want_freeze ($installpid);
    select (undef, undef, undef, 0.1);
  }
  Log (undef, "Install exited $?");
}

if (!$have_slurm)
{
  # Grab our lock again (we might have deleted and re-created CRUNCH_TMP above)
  must_lock_now("$ENV{CRUNCH_TMP}/.lock", "a job is already running here.");
}

# If this job requires a Docker image, install that.
my $docker_bin = "/usr/bin/docker.io";
my ($docker_locator, $docker_hash);
if ($docker_locator = $Job->{docker_image_locator}) {
  $docker_hash = find_docker_hash($docker_locator);
  if (!$docker_hash)
  {
    croak("No Docker image hash found from locator $docker_locator");
  }
  my $docker_install_script = qq{
if ! $docker_bin images -q --no-trunc | grep -qxF \Q$docker_hash\E; then
    arv-get \Q$docker_locator/$docker_hash.tar\E | $docker_bin load
fi
};
  my $docker_pid = fork();
  if ($docker_pid == 0)
  {
    srun (["srun", "--nodelist=" . join(',', @node)],
          ["/bin/sh", "-ec", $docker_install_script]);
    exit ($?);
  }
  while (1)
  {
    last if $docker_pid == waitpid (-1, WNOHANG);
    freeze_if_want_freeze ($docker_pid);
    select (undef, undef, undef, 0.1);
  }
  if ($? != 0)
  {
    croak("Installing Docker image from $docker_locator returned exit code $?");
  }
}

foreach (qw (script script_version script_parameters runtime_constraints))
{
  Log (undef,
       "$_ " .
       (ref($Job->{$_}) ? JSON::encode_json($Job->{$_}) : $Job->{$_}));
}
foreach (split (/\n/, $Job->{knobs}))
{
  Log (undef, "knob " . $_);
}



$main::success = undef;



ONELEVEL:

my $thisround_succeeded = 0;
my $thisround_failed = 0;
my $thisround_failed_multiple = 0;

@jobstep_todo = sort { $jobstep[$a]->{level} <=> $jobstep[$b]->{level}
		       or $a <=> $b } @jobstep_todo;
my $level = $jobstep[$jobstep_todo[0]]->{level};
Log (undef, "start level $level");



my %proc;
my @freeslot = (0..$#slot);
my @holdslot;
my %reader;
my $progress_is_dirty = 1;
my $progress_stats_updated = 0;

update_progress_stats();



THISROUND:
for (my $todo_ptr = 0; $todo_ptr <= $#jobstep_todo; $todo_ptr ++)
{
  my $id = $jobstep_todo[$todo_ptr];
  my $Jobstep = $jobstep[$id];
  if ($Jobstep->{level} != $level)
  {
    next;
  }

  pipe $reader{$id}, "writer" or croak ($!);
  my $flags = fcntl ($reader{$id}, F_GETFL, 0) or croak ($!);
  fcntl ($reader{$id}, F_SETFL, $flags | O_NONBLOCK) or croak ($!);

  my $childslot = $freeslot[0];
  my $childnode = $slot[$childslot]->{node};
  my $childslotname = join (".",
			    $slot[$childslot]->{node}->{name},
			    $slot[$childslot]->{cpu});
  my $childpid = fork();
  if ($childpid == 0)
  {
    $SIG{'INT'} = 'DEFAULT';
    $SIG{'QUIT'} = 'DEFAULT';
    $SIG{'TERM'} = 'DEFAULT';

    foreach (values (%reader))
    {
      close($_);
    }
    fcntl ("writer", F_SETFL, 0) or croak ($!); # no close-on-exec
    open(STDOUT,">&writer");
    open(STDERR,">&writer");

    undef $dbh;
    undef $sth;

    delete $ENV{"GNUPGHOME"};
    $ENV{"TASK_UUID"} = $Jobstep->{'arvados_task'}->{'uuid'};
    $ENV{"TASK_QSEQUENCE"} = $id;
    $ENV{"TASK_SEQUENCE"} = $level;
    $ENV{"JOB_SCRIPT"} = $Job->{script};
    while (my ($param, $value) = each %{$Job->{script_parameters}}) {
      $param =~ tr/a-z/A-Z/;
      $ENV{"JOB_PARAMETER_$param"} = $value;
    }
    $ENV{"TASK_SLOT_NODE"} = $slot[$childslot]->{node}->{name};
    $ENV{"TASK_SLOT_NUMBER"} = $slot[$childslot]->{cpu};
    $ENV{"TASK_WORK"} = $ENV{"JOB_WORK"}."/$id.$$";
    $ENV{"TASK_KEEPMOUNT"} = $ENV{"TASK_WORK"}.".keep";
    $ENV{"TASK_TMPDIR"} = $ENV{"TASK_WORK"}; # deprecated
    $ENV{"CRUNCH_NODE_SLOTS"} = $slot[$childslot]->{node}->{ncpus};
    $ENV{"PATH"} = $ENV{"CRUNCH_INSTALL"} . "/bin:" . $ENV{"PATH"};

    $ENV{"GZIP"} = "-n";

    my @srunargs = (
      "srun",
      "--nodelist=".$childnode->{name},
      qw(-n1 -c1 -N1 -D), $ENV{'TMPDIR'},
      "--job-name=$job_id.$id.$$",
	);
    my $build_script_to_send = "";
    my $command =
	"if [ -e $ENV{TASK_WORK} ]; then rm -rf $ENV{TASK_WORK}; fi; "
        ."mkdir -p $ENV{JOB_WORK} $ENV{CRUNCH_TMP} $ENV{TASK_WORK} $ENV{TASK_KEEPMOUNT} "
        ."&& chmod og+wrx $ENV{TASK_WORK}"
	."&& cd $ENV{CRUNCH_TMP} ";
    if ($build_script)
    {
      $build_script_to_send = $build_script;
      $command .=
	  "&& perl -";
    }
    $command .= "&& exec arv-mount --allow-other $ENV{TASK_KEEPMOUNT} --exec ";
    if ($docker_hash)
    {
      $command .= "crunchstat -cgroup-root=/sys/fs/cgroup -cgroup-parent=docker -cgroup-cid=$ENV{TASK_WORK}/docker.cid -poll=10000 ";
      $command .= "$docker_bin run -i -a stdin -a stdout -a stderr --cidfile=$ENV{TASK_WORK}/docker.cid ";
      # Dynamically configure the container to use the host system as its
      # DNS server.  Get the host's global addresses from the ip command,
      # and turn them into docker --dns options using gawk.
      $command .=
          q{$(ip -o address show scope global |
              gawk 'match($4, /^([0-9\.:]+)\//, x){print "--dns", x[1]}') };
      $command .= "-v \Q$ENV{TASK_WORK}:/tmp/crunch-job:rw\E ";
      $command .= "-v \Q$ENV{CRUNCH_SRC}:/tmp/crunch-src:ro\E ";
      $command .= "-v \Q$ENV{TASK_KEEPMOUNT}:/mnt:ro\E ";
      while (my ($env_key, $env_val) = each %ENV)
      {
        if ($env_key =~ /^(ARVADOS|JOB|TASK)_/) {
          if ($env_key eq "TASK_WORK") {
            $command .= "-e \QTASK_WORK=/tmp/crunch-job\E ";
          }
          elsif ($env_key eq "TASK_KEEPMOUNT") {
            $command .= "-e \QTASK_KEEPMOUNT=/mnt\E ";
          }
          elsif ($env_key eq "CRUNCH_SRC") {
            $command .= "-e \QCRUNCH_SRC=/tmp/crunch-src\E ";
          }
          else {
            $command .= "-e \Q$env_key=$env_val\E ";
          }
        }
      }
      $command .= "\Q$docker_hash\E ";
      $command .= "stdbuf -o0 -e0 ";
      $command .= "/tmp/crunch-src/crunch_scripts/" . $Job->{"script"};
    } else {
      # Non-docker run
      $command .= "crunchstat -cgroup-root=/sys/fs/cgroup -poll=10000 ";
      $command .= "stdbuf -o0 -e0 ";
      $command .= "$ENV{CRUNCH_SRC}/crunch_scripts/" . $Job->{"script"};
    }

    my @execargs = ('bash', '-c', $command);
    srun (\@srunargs, \@execargs, undef, $build_script_to_send);
    exit (111);
  }
  close("writer");
  if (!defined $childpid)
  {
    close $reader{$id};
    delete $reader{$id};
    next;
  }
  shift @freeslot;
  $proc{$childpid} = { jobstep => $id,
		       time => time,
		       slot => $childslot,
		       jobstepname => "$job_id.$id.$childpid",
		     };
  croak ("assert failed: \$slot[$childslot]->{'pid'} exists") if exists $slot[$childslot]->{pid};
  $slot[$childslot]->{pid} = $childpid;

  Log ($id, "job_task ".$Jobstep->{'arvados_task'}->{'uuid'});
  Log ($id, "child $childpid started on $childslotname");
  $Jobstep->{starttime} = time;
  $Jobstep->{node} = $childnode->{name};
  $Jobstep->{slotindex} = $childslot;
  delete $Jobstep->{stderr};
  delete $Jobstep->{finishtime};

  splice @jobstep_todo, $todo_ptr, 1;
  --$todo_ptr;

  $progress_is_dirty = 1;

  while (!@freeslot
	 ||
	 (@slot > @freeslot && $todo_ptr+1 > $#jobstep_todo))
  {
    last THISROUND if $main::please_freeze;
    if ($main::please_info)
    {
      $main::please_info = 0;
      freeze();
      collate_output();
      save_meta(1);
      update_progress_stats();
    }
    my $gotsome
	= readfrompipes ()
	+ reapchildren ();
    if (!$gotsome)
    {
      check_refresh_wanted();
      check_squeue();
      update_progress_stats();
      select (undef, undef, undef, 0.1);
    }
    elsif (time - $progress_stats_updated >= 30)
    {
      update_progress_stats();
    }
    if (($thisround_failed_multiple >= 8 && $thisround_succeeded == 0) ||
	($thisround_failed_multiple >= 16 && $thisround_failed_multiple > $thisround_succeeded))
    {
      my $message = "Repeated failure rate too high ($thisround_failed_multiple/"
	  .($thisround_failed+$thisround_succeeded)
	  .") -- giving up on this round";
      Log (undef, $message);
      last THISROUND;
    }

    # move slots from freeslot to holdslot (or back to freeslot) if necessary
    for (my $i=$#freeslot; $i>=0; $i--) {
      if ($slot[$freeslot[$i]]->{node}->{hold_until} > scalar time) {
	push @holdslot, (splice @freeslot, $i, 1);
      }
    }
    for (my $i=$#holdslot; $i>=0; $i--) {
      if ($slot[$holdslot[$i]]->{node}->{hold_until} <= scalar time) {
	push @freeslot, (splice @holdslot, $i, 1);
      }
    }

    # give up if no nodes are succeeding
    if (!grep { $_->{node}->{losing_streak} == 0 &&
                    $_->{node}->{hold_count} < 4 } @slot) {
      my $message = "Every node has failed -- giving up on this round";
      Log (undef, $message);
      last THISROUND;
    }
  }
}


push @freeslot, splice @holdslot;
map { $slot[$freeslot[$_]]->{node}->{losing_streak} = 0 } (0..$#freeslot);


Log (undef, "wait for last ".(scalar keys %proc)." children to finish");
while (%proc)
{
  if ($main::please_continue) {
    $main::please_continue = 0;
    goto THISROUND;
  }
  $main::please_info = 0, freeze(), collate_output(), save_meta(1) if $main::please_info;
  readfrompipes ();
  if (!reapchildren())
  {
    check_refresh_wanted();
    check_squeue();
    update_progress_stats();
    select (undef, undef, undef, 0.1);
    killem (keys %proc) if $main::please_freeze;
  }
}

update_progress_stats();
freeze_if_want_freeze();


if (!defined $main::success)
{
  if (@jobstep_todo &&
      $thisround_succeeded == 0 &&
      ($thisround_failed == 0 || $thisround_failed > 4))
  {
    my $message = "stop because $thisround_failed tasks failed and none succeeded";
    Log (undef, $message);
    $main::success = 0;
  }
  if (!@jobstep_todo)
  {
    $main::success = 1;
  }
}

goto ONELEVEL if !defined $main::success;


release_allocation();
freeze();
my $collated_output = &collate_output();

if ($job_has_uuid) {
  $Job->update_attributes('running' => 0,
                          'success' => $collated_output && $main::success,
                          'finished_at' => scalar gmtime)
}

if ($collated_output)
{
  eval {
    open(my $orig_manifest, '-|', 'arv', 'keep', 'get', $collated_output)
        or die "failed to get collated manifest: $!";
    # Read the original manifest, and strip permission hints from it,
    # so we can put the result in a Collection.
    my @stripped_manifest_lines = ();
    my $orig_manifest_text = '';
    while (my $manifest_line = <$orig_manifest>) {
      $orig_manifest_text .= $manifest_line;
      my @words = split(/ /, $manifest_line, -1);
      foreach my $ii (0..$#words) {
        if ($words[$ii] =~ /^[0-9a-f]{32}\+/) {
          $words[$ii] =~ s/\+A[0-9a-f]{40}@[0-9a-f]{8}\b//;
        }
      }
      push(@stripped_manifest_lines, join(" ", @words));
    }
    my $stripped_manifest_text = join("", @stripped_manifest_lines);
    my $output = $arv->{'collections'}->{'create'}->execute('collection' => {
      'uuid' => md5_hex($stripped_manifest_text),
      'manifest_text' => $orig_manifest_text,
    });
    $Job->update_attributes('output' => $output->{uuid});
    if ($Job->{'output_is_persistent'}) {
      $arv->{'links'}->{'create'}->execute('link' => {
        'tail_kind' => 'arvados#user',
        'tail_uuid' => $User->{'uuid'},
        'head_kind' => 'arvados#collection',
        'head_uuid' => $Job->{'output'},
        'link_class' => 'resources',
        'name' => 'wants',
      });
    }
  };
  if ($@) {
    Log (undef, "Failed to register output manifest: $@");
  }
}

Log (undef, "finish");

save_meta();
exit 0;



sub update_progress_stats
{
  $progress_stats_updated = time;
  return if !$progress_is_dirty;
  my ($todo, $done, $running) = (scalar @jobstep_todo,
				 scalar @jobstep_done,
				 scalar @slot - scalar @freeslot - scalar @holdslot);
  $Job->{'tasks_summary'} ||= {};
  $Job->{'tasks_summary'}->{'todo'} = $todo;
  $Job->{'tasks_summary'}->{'done'} = $done;
  $Job->{'tasks_summary'}->{'running'} = $running;
  if ($job_has_uuid) {
    $Job->update_attributes('tasks_summary' => $Job->{'tasks_summary'});
  }
  Log (undef, "status: $done done, $running running, $todo todo");
  $progress_is_dirty = 0;
}



sub reapchildren
{
  my $pid = waitpid (-1, WNOHANG);
  return 0 if $pid <= 0;

  my $whatslot = ($slot[$proc{$pid}->{slot}]->{node}->{name}
		  . "."
		  . $slot[$proc{$pid}->{slot}]->{cpu});
  my $jobstepid = $proc{$pid}->{jobstep};
  my $elapsed = time - $proc{$pid}->{time};
  my $Jobstep = $jobstep[$jobstepid];

  my $childstatus = $?;
  my $exitvalue = $childstatus >> 8;
  my $exitinfo = sprintf("exit %d signal %d%s",
                         $exitvalue,
                         $childstatus & 127,
                         ($childstatus & 128 ? ' core dump' : ''));
  $Jobstep->{'arvados_task'}->reload;
  my $task_success = $Jobstep->{'arvados_task'}->{success};

  Log ($jobstepid, "child $pid on $whatslot $exitinfo success=$task_success");

  if (!defined $task_success) {
    # task did not indicate one way or the other --> fail
    $Jobstep->{'arvados_task'}->{success} = 0;
    $Jobstep->{'arvados_task'}->save;
    $task_success = 0;
  }

  if (!$task_success)
  {
    my $temporary_fail;
    $temporary_fail ||= $Jobstep->{node_fail};
    $temporary_fail ||= ($exitvalue == 111);

    ++$thisround_failed;
    ++$thisround_failed_multiple if $Jobstep->{'failures'} >= 1;

    # Check for signs of a failed or misconfigured node
    if (++$slot[$proc{$pid}->{slot}]->{node}->{losing_streak} >=
	2+$slot[$proc{$pid}->{slot}]->{node}->{ncpus}) {
      # Don't count this against jobstep failure thresholds if this
      # node is already suspected faulty and srun exited quickly
      if ($slot[$proc{$pid}->{slot}]->{node}->{hold_until} &&
	  $elapsed < 5) {
	Log ($jobstepid, "blaming failure on suspect node " .
             $slot[$proc{$pid}->{slot}]->{node}->{name});
        $temporary_fail ||= 1;
      }
      ban_node_by_slot($proc{$pid}->{slot});
    }

    Log ($jobstepid, sprintf('failure (#%d, %s) after %d seconds',
                             ++$Jobstep->{'failures'},
                             $temporary_fail ? 'temporary ' : 'permanent',
                             $elapsed));

    if (!$temporary_fail || $Jobstep->{'failures'} >= 3) {
      # Give up on this task, and the whole job
      $main::success = 0;
      $main::please_freeze = 1;
    }
    else {
      # Put this task back on the todo queue
      push @jobstep_todo, $jobstepid;
    }
    $Job->{'tasks_summary'}->{'failed'}++;
  }
  else
  {
    ++$thisround_succeeded;
    $slot[$proc{$pid}->{slot}]->{node}->{losing_streak} = 0;
    $slot[$proc{$pid}->{slot}]->{node}->{hold_until} = 0;
    push @jobstep_done, $jobstepid;
    Log ($jobstepid, "success in $elapsed seconds");
  }
  $Jobstep->{exitcode} = $childstatus;
  $Jobstep->{finishtime} = time;
  process_stderr ($jobstepid, $task_success);
  Log ($jobstepid, "output " . $Jobstep->{'arvados_task'}->{output});

  close $reader{$jobstepid};
  delete $reader{$jobstepid};
  delete $slot[$proc{$pid}->{slot}]->{pid};
  push @freeslot, $proc{$pid}->{slot};
  delete $proc{$pid};

  # Load new tasks
  my $newtask_list = [];
  my $newtask_results;
  do {
    $newtask_results = $arv->{'job_tasks'}->{'list'}->execute(
      'where' => {
        'created_by_job_task_uuid' => $Jobstep->{'arvados_task'}->{uuid}
      },
      'order' => 'qsequence',
      'offset' => scalar(@$newtask_list),
    );
    push(@$newtask_list, @{$newtask_results->{items}});
  } while (@{$newtask_results->{items}});
  foreach my $arvados_task (@$newtask_list) {
    my $jobstep = {
      'level' => $arvados_task->{'sequence'},
      'failures' => 0,
      'arvados_task' => $arvados_task
    };
    push @jobstep, $jobstep;
    push @jobstep_todo, $#jobstep;
  }

  $progress_is_dirty = 1;
  1;
}

sub check_refresh_wanted
{
  my @stat = stat $ENV{"CRUNCH_REFRESH_TRIGGER"};
  if (@stat && $stat[9] > $latest_refresh) {
    $latest_refresh = scalar time;
    if ($job_has_uuid) {
      my $Job2 = $arv->{'jobs'}->{'get'}->execute('uuid' => $jobspec);
      for my $attr ('cancelled_at',
                    'cancelled_by_user_uuid',
                    'cancelled_by_client_uuid') {
        $Job->{$attr} = $Job2->{$attr};
      }
      if ($Job->{'cancelled_at'}) {
        Log (undef, "Job cancelled at " . $Job->{cancelled_at} .
             " by user " . $Job->{cancelled_by_user_uuid});
        $main::success = 0;
        $main::please_freeze = 1;
      }
    }
  }
}

sub check_squeue
{
  # return if the kill list was checked <4 seconds ago
  if (defined $squeue_kill_checked && $squeue_kill_checked > time - 4)
  {
    return;
  }
  $squeue_kill_checked = time;

  # use killem() on procs whose killtime is reached
  for (keys %proc)
  {
    if (exists $proc{$_}->{killtime}
	&& $proc{$_}->{killtime} <= time)
    {
      killem ($_);
    }
  }

  # return if the squeue was checked <60 seconds ago
  if (defined $squeue_checked && $squeue_checked > time - 60)
  {
    return;
  }
  $squeue_checked = time;

  if (!$have_slurm)
  {
    # here is an opportunity to check for mysterious problems with local procs
    return;
  }

  # get a list of steps still running
  my @squeue = `squeue -s -h -o '%i %j' && echo ok`;
  chop @squeue;
  if ($squeue[-1] ne "ok")
  {
    return;
  }
  pop @squeue;

  # which of my jobsteps are running, according to squeue?
  my %ok;
  foreach (@squeue)
  {
    if (/^(\d+)\.(\d+) (\S+)/)
    {
      if ($1 eq $ENV{SLURM_JOBID})
      {
	$ok{$3} = 1;
      }
    }
  }

  # which of my active child procs (>60s old) were not mentioned by squeue?
  foreach (keys %proc)
  {
    if ($proc{$_}->{time} < time - 60
	&& !exists $ok{$proc{$_}->{jobstepname}}
	&& !exists $proc{$_}->{killtime})
    {
      # kill this proc if it hasn't exited in 30 seconds
      $proc{$_}->{killtime} = time + 30;
    }
  }
}


sub release_allocation
{
  if ($have_slurm)
  {
    Log (undef, "release job allocation");
    system "scancel $ENV{SLURM_JOBID}";
  }
}


sub readfrompipes
{
  my $gotsome = 0;
  foreach my $job (keys %reader)
  {
    my $buf;
    while (0 < sysread ($reader{$job}, $buf, 8192))
    {
      print STDERR $buf if $ENV{CRUNCH_DEBUG};
      $jobstep[$job]->{stderr} .= $buf;
      preprocess_stderr ($job);
      if (length ($jobstep[$job]->{stderr}) > 16384)
      {
	substr ($jobstep[$job]->{stderr}, 0, 8192) = "";
      }
      $gotsome = 1;
    }
  }
  return $gotsome;
}


sub preprocess_stderr
{
  my $job = shift;

  while ($jobstep[$job]->{stderr} =~ /^(.*?)\n/) {
    my $line = $1;
    substr $jobstep[$job]->{stderr}, 0, 1+length($line), "";
    Log ($job, "stderr $line");
    if ($line =~ /srun: error: (SLURM job $ENV{SLURM_JOB_ID} has expired|Unable to confirm allocation for job $ENV{SLURM_JOB_ID})/) {
      # whoa.
      $main::please_freeze = 1;
    }
    elsif ($line =~ /srun: error: (Node failure on|Unable to create job step) /) {
      $jobstep[$job]->{node_fail} = 1;
      ban_node_by_slot($jobstep[$job]->{slotindex});
    }
  }
}


sub process_stderr
{
  my $job = shift;
  my $task_success = shift;
  preprocess_stderr ($job);

  map {
    Log ($job, "stderr $_");
  } split ("\n", $jobstep[$job]->{stderr});
}

sub fetch_block
{
  my $hash = shift;
  my ($keep, $child_out, $output_block);

  my $cmd = "$arv_cli keep get \Q$hash\E";
  open($keep, '-|', $cmd) or die "fetch_block: $cmd: $!";
  sysread($keep, $output_block, 64 * 1024 * 1024);
  close $keep;
  return $output_block;
}

sub collate_output
{
  Log (undef, "collate");

  my ($child_out, $child_in);
  my $pid = open2($child_out, $child_in, $arv_cli, 'keep', 'put', '--raw');
  my $joboutput;
  for (@jobstep)
  {
    next if (!exists $_->{'arvados_task'}->{output} ||
             !$_->{'arvados_task'}->{'success'} ||
             $_->{'exitcode'} != 0);
    my $output = $_->{'arvados_task'}->{output};
    if ($output !~ /^[0-9a-f]{32}(\+\S+)*$/)
    {
      $output_in_keep ||= $output =~ / [0-9a-f]{32}\S*\+K/;
      print $child_in $output;
    }
    elsif (@jobstep == 1)
    {
      $joboutput = $output;
      last;
    }
    elsif (defined (my $outblock = fetch_block ($output)))
    {
      $output_in_keep ||= $outblock =~ / [0-9a-f]{32}\S*\+K/;
      print $child_in $outblock;
    }
    else
    {
      Log (undef, "XXX fetch_block($output) failed XXX");
      $main::success = 0;
    }
  }
  $child_in->close;

  if (!defined $joboutput) {
    my $s = IO::Select->new($child_out);
    if ($s->can_read(120)) {
      sysread($child_out, $joboutput, 64 * 1024 * 1024);
      chomp($joboutput);
    } else {
      Log (undef, "timed out reading from 'arv keep put'");
    }
  }
  waitpid($pid, 0);

  if ($joboutput)
  {
    Log (undef, "output $joboutput");
    $Job->update_attributes('output' => $joboutput) if $job_has_uuid;
  }
  else
  {
    Log (undef, "output undef");
  }
  return $joboutput;
}


sub killem
{
  foreach (@_)
  {
    my $sig = 2;		# SIGINT first
    if (exists $proc{$_}->{"sent_$sig"} &&
	time - $proc{$_}->{"sent_$sig"} > 4)
    {
      $sig = 15;		# SIGTERM if SIGINT doesn't work
    }
    if (exists $proc{$_}->{"sent_$sig"} &&
	time - $proc{$_}->{"sent_$sig"} > 4)
    {
      $sig = 9;			# SIGKILL if SIGTERM doesn't work
    }
    if (!exists $proc{$_}->{"sent_$sig"})
    {
      Log ($proc{$_}->{jobstep}, "sending 2x signal $sig to pid $_");
      kill $sig, $_;
      select (undef, undef, undef, 0.1);
      if ($sig == 2)
      {
	kill $sig, $_;	   # srun wants two SIGINT to really interrupt
      }
      $proc{$_}->{"sent_$sig"} = time;
      $proc{$_}->{"killedafter"} = time - $proc{$_}->{"time"};
    }
  }
}


sub fhbits
{
  my($bits);
  for (@_) {
    vec($bits,fileno($_),1) = 1;
  }
  $bits;
}


sub Log				# ($jobstep_id, $logmessage)
{
  if ($_[1] =~ /\n/) {
    for my $line (split (/\n/, $_[1])) {
      Log ($_[0], $line);
    }
    return;
  }
  my $fh = select STDERR; $|=1; select $fh;
  my $message = sprintf ("%s %d %s %s", $job_id, $$, @_);
  $message =~ s{([^ -\176])}{"\\" . sprintf ("%03o", ord($1))}ge;
  $message .= "\n";
  my $datetime;
  if ($local_logfile || -t STDERR) {
    my @gmtime = gmtime;
    $datetime = sprintf ("%04d-%02d-%02d_%02d:%02d:%02d",
			 $gmtime[5]+1900, $gmtime[4]+1, @gmtime[3,2,1,0]);
  }
  print STDERR ((-t STDERR) ? ($datetime." ".$message) : $message);

  if ($local_logfile) {
    print $local_logfile $datetime . " " . $message;
  }
}


sub croak
{
  my ($package, $file, $line) = caller;
  my $message = "@_ at $file line $line\n";
  Log (undef, $message);
  freeze() if @jobstep_todo;
  collate_output() if @jobstep_todo;
  cleanup();
  save_meta() if $local_logfile;
  die;
}


sub cleanup
{
  return if !$job_has_uuid;
  $Job->update_attributes('running' => 0,
                          'success' => 0,
                          'finished_at' => scalar gmtime);
}


sub save_meta
{
  my $justcheckpoint = shift; # false if this will be the last meta saved
  return if $justcheckpoint;  # checkpointing is not relevant post-Warehouse.pm

  $local_logfile->flush;
  my $cmd = "$arv_cli keep put --filename ''\Q$keep_logfile\E "
      . quotemeta($local_logfile->filename);
  my $loglocator = `$cmd`;
  die "system $cmd failed: $?" if $?;
  chomp($loglocator);

  $local_logfile = undef;   # the temp file is automatically deleted
  Log (undef, "log manifest is $loglocator");
  $Job->{'log'} = $loglocator;
  $Job->update_attributes('log', $loglocator) if $job_has_uuid;
}


sub freeze_if_want_freeze
{
  if ($main::please_freeze)
  {
    release_allocation();
    if (@_)
    {
      # kill some srun procs before freeze+stop
      map { $proc{$_} = {} } @_;
      while (%proc)
      {
	killem (keys %proc);
	select (undef, undef, undef, 0.1);
	my $died;
	while (($died = waitpid (-1, WNOHANG)) > 0)
	{
	  delete $proc{$died};
	}
      }
    }
    freeze();
    collate_output();
    cleanup();
    save_meta();
    exit 0;
  }
}


sub freeze
{
  Log (undef, "Freeze not implemented");
  return;
}


sub thaw
{
  croak ("Thaw not implemented");
}


sub freezequote
{
  my $s = shift;
  $s =~ s/\\/\\\\/g;
  $s =~ s/\n/\\n/g;
  return $s;
}


sub freezeunquote
{
  my $s = shift;
  $s =~ s{\\(.)}{$1 eq "n" ? "\n" : $1}ge;
  return $s;
}


sub srun
{
  my $srunargs = shift;
  my $execargs = shift;
  my $opts = shift || {};
  my $stdin = shift;
  my $args = $have_slurm ? [@$srunargs, @$execargs] : $execargs;
  print STDERR (join (" ",
		      map { / / ? "'$_'" : $_ }
		      (@$args)),
		"\n")
      if $ENV{CRUNCH_DEBUG};

  if (defined $stdin) {
    my $child = open STDIN, "-|";
    defined $child or die "no fork: $!";
    if ($child == 0) {
      print $stdin or die $!;
      close STDOUT or die $!;
      exit 0;
    }
  }

  return system (@$args) if $opts->{fork};

  exec @$args;
  warn "ENV size is ".length(join(" ",%ENV));
  die "exec failed: $!: @$args";
}


sub ban_node_by_slot {
  # Don't start any new jobsteps on this node for 60 seconds
  my $slotid = shift;
  $slot[$slotid]->{node}->{hold_until} = 60 + scalar time;
  $slot[$slotid]->{node}->{hold_count}++;
  Log (undef, "backing off node " . $slot[$slotid]->{node}->{name} . " for 60 seconds");
}

sub must_lock_now
{
  my ($lockfile, $error_message) = @_;
  open L, ">", $lockfile or croak("$lockfile: $!");
  if (!flock L, LOCK_EX|LOCK_NB) {
    croak("Can't lock $lockfile: $error_message\n");
  }
}

sub find_docker_hash {
  # Given a Keep locator, search for a matching link to find the Docker hash
  # of the stored image.
  my $locator = shift;
  my $links_result = $arv->{links}->{list}->execute(
    filters => [["head_uuid", "=", $locator],
                ["link_class", "=", "docker_image_hash"]],
    limit => 1);
  my $docker_hash;
  foreach my $link (@{$links_result->{items}}) {
    $docker_hash = lc($link->{name});
  }
  return $docker_hash;
}

__DATA__
#!/usr/bin/perl

# checkout-and-build

use Fcntl ':flock';

my $destdir = $ENV{"CRUNCH_SRC"};
my $commit = $ENV{"CRUNCH_SRC_COMMIT"};
my $repo = $ENV{"CRUNCH_SRC_URL"};

open L, ">", "$destdir.lock" or die "$destdir.lock: $!";
flock L, LOCK_EX;
if (readlink ("$destdir.commit") eq $commit && -d $destdir) {
    exit 0;
}

unlink "$destdir.commit";
open STDOUT, ">", "$destdir.log";
open STDERR, ">&STDOUT";

mkdir $destdir;
my @git_archive_data = <DATA>;
if (@git_archive_data) {
  open TARX, "|-", "tar", "-C", $destdir, "-xf", "-";
  print TARX @git_archive_data;
  if(!close(TARX)) {
    die "'tar -C $destdir -xf -' exited $?: $!";
  }
}

my $pwd;
chomp ($pwd = `pwd`);
my $install_dir = $ENV{"CRUNCH_INSTALL"} || "$pwd/opt";
mkdir $install_dir;

for my $src_path ("$destdir/arvados/sdk/python") {
  if (-d $src_path) {
    shell_or_die ("virtualenv", $install_dir);
    shell_or_die ("cd $src_path && ./build.sh && $install_dir/bin/python setup.py install");
  }
}

if (-e "$destdir/crunch_scripts/install") {
    shell_or_die ("$destdir/crunch_scripts/install", $install_dir);
} elsif (!-e "./install.sh" && -e "./tests/autotests.sh") {
    # Old version
    shell_or_die ("./tests/autotests.sh", $install_dir);
} elsif (-e "./install.sh") {
    shell_or_die ("./install.sh", $install_dir);
}

if ($commit) {
    unlink "$destdir.commit.new";
    symlink ($commit, "$destdir.commit.new") or die "$destdir.commit.new: $!";
    rename ("$destdir.commit.new", "$destdir.commit") or die "$destdir.commit: $!";
}

close L;

exit 0;

sub shell_or_die
{
  if ($ENV{"DEBUG"}) {
    print STDERR "@_\n";
  }
  system (@_) == 0
      or die "@_ failed: $! exit 0x".sprintf("%x",$?);
}

__DATA__