#!/usr/bin/perl
#Parallel dispatch
#Run data parallel jobs, with multiple CPUs per slice

use strict;
use IPC::Open2;
use IO::Select;
use IO::Handle;

#I don't know how to make a proper struct in Perl, so settling for an array
use constant {
  PROCIN    => 0,
  PROCOUT   => 1,
  PROCPID   => 2,
  PROCINDEX => 3
};

#Spawn the job
sub forkjob
{
  my ($machinefileindex,$script,$piece)=@_;
  my $machinefile="machines_".$machinefileindex.".lst";
  my $procin=IO::Handle->new();
  my $procout=IO::Handle->new();
  my $pid;

  $pid=open2($procout,$procin,$script,$machinefile,$piece);

  return [$procin,$procout,$pid,$piece];
}

#Check if have the correct number commandline parameters
if (@ARGV!=4)
{
  print STDERR "pdisp <machinefile> <cpus> <rank> <script>\n";
  exit 1;
}

#Get commandline parameters
my $machinefile=$ARGV[0];
my $cpus=int($ARGV[1]);
my $rank=int($ARGV[2]);
my $script=$ARGV[3];

#Check for reasonable values
if ($cpus<1 || $cpus>99)
{
  print STDERR "cpus should be between 1 and 99.\n";
  exit 1;
}

if ($rank<1 || $rank>9999)
{
  print STDERR "rank should be between 1 and 9999.\n";
  exit 1;
}


#Split machines file into pieces of size $cpus
open(MACHINES,"<",$machinefile);
my @cpulist=<MACHINES>;
close(MACHINES);

my $partitions=int(@cpulist/$cpus);
if ($partitions<1)
{
  print STDERR "More cpus required than are in the machinefile\n";
  exit 1;
}

#Find how many jobs to run in parallel
my $max=$partitions;
if ($rank<$partitions)
{
  $max=$rank;
}

#Write out machinefiles
#Since same numbers are adjacent, this will prefer putting jobs on fewer nodes
for(my $i=0;$i<$max;$i++)
{
  open(OUTMACHINES,">","machines_".$i.".lst");
  print OUTMACHINES @cpulist[($i*$cpus)..((($i+1)*$cpus)-1)];
  close(OUTMACHINES);
}

#Make array of job info
my $atjob=0;
my $sel = IO::Select->new();
my %handle2index;
my @jobs=();
$jobs[$max-1]=(); #Set size of array

#Initialize jobs
for(my $i=0;$i<$max;$i++)
{
  $jobs[$i]=&forkjob($i,$script,$atjob);
  my $procout=$jobs[$i][PROCOUT];
  $sel->add($procout);
  $handle2index{fileno($procout)}=$i;
  $atjob++;
}

#Main wait for I/O loop
my $left=$rank;
while(($left>0) && (my @ready = $sel->can_read))
{
  foreach my $handle (@ready)
  {
    my $i=$handle2index{fileno($handle)};
    while(my $line=<$handle>) {
      chomp($line);
      print $jobs[$i][PROCINDEX].": ".$line."\n";
    }
    if (eof($handle))
    {
      $sel->remove($handle);
      close($jobs[$i][PROCIN]);
      close($jobs[$i][PROCOUT]);
      $left--;
      waitpid $jobs[$i][PROCPID],0;

      if ($atjob<$rank)
      {
        $jobs[$i]=&forkjob($i,$script,$atjob);
        my $procout=$jobs[$i][PROCOUT];
        $sel->add($procout);
        $handle2index{fileno($procout)}=$i;
        $atjob++;
} } } }