Index: /gs2-extensions/parallel-building/trunk/src/bin/script/update_data_locality.pl
===================================================================
--- /gs2-extensions/parallel-building/trunk/src/bin/script/update_data_locality.pl	(revision 28646)
+++ /gs2-extensions/parallel-building/trunk/src/bin/script/update_data_locality.pl	(revision 28647)
@@ -1,4 +1,5 @@
 #!/usr/bin/perl
 
+use Cwd;
 use Sort::Key::Natural qw(natsort);
 
@@ -12,136 +13,239 @@
 print "============================================================\n\n";
 
-if (!defined $ARGV[0] || !-d $ARGV[0])
+our $debug = 0;
+
+# Loop through all the arguments, descending into directories to process data
+# locality files
+my $offset = 0;
+my $found_dirs = 0;
+while (defined $ARGV[$offset])
+{
+  my $argument = $ARGV[$offset];
+  # handle -option arguments?
+  if ($argument =~ /^-(.+)$/)
+  {
+    my $option = $1;
+    if ($option eq 'debug')
+    {
+      $debug = 1;
+    }
+  }
+  # handle directories
+  else
+  {
+    my $path = getcwd() . '/' . $argument;
+    if (-d $path)
+    {
+      $found_dirs++;
+      &searchForDataLocalityFiles($path);
+    }
+  }
+  $offset++;
+}
+if ($found_dirs == 0)
 {
   &printUsage('Missing results directory');
 }
 
-# 0. Init
-my $data = {};
-my $data_locality_csv_path = $ARGV[0] . '/data_locality.csv';
-my $hadoop_log_path = $ARGV[0] . '/hadoop.log';
-
-# 1. Read in current data locality comma separated file
-print " * Read current file: data_locality.csv\n";
-if (open(DLIN, '<:utf8', $data_locality_csv_path))
-{
-  my $record_count = 0;
-  while (my $line = <DLIN>)
-  {
-    if ($line =~ /^([^,]+),(\d),\d,"([^"]+)","([^"]+)"$/)
-    {
-      my $taskid       = $1;
-      my $attempt_no   = $2;
-      my $compute_node = $3;
-      my $splits       = $4;
-
-      # Use the taskid to look up the manifest file and therein the filepath
-      my $filepath = '';
-      my $manifest_path = $ARGV[0] . '/manifest' . $taskid . '_' . $attempt_no . '.xml';
-      if (open(MIN, '<:utf8', $manifest_path))
-      {
-        while (my $line2 = <MIN>)
-        {
-          if ($line2 =~ /<Filename>(hdfs|HDFSShell|ThriftFS):(.*)<\/Filename>/)
-          {
-            my $protocol = $1;
-            $filepath = 'hdfs:' . $2;
-          }
-        }
-        close(MIN);
-      }
-      else
-      {
-        die('Failed to open manifest for reading: ' . $manifest_path);
-      }
-
-      $data->{$filepath} = {'taskid' => $taskid,
-                            'attempt' => $attempt_no,
-                            'node' => $compute_node,
-                            'splits' => $splits,
-                            'original_splits' => ''};
-      $record_count++;
-    }
-  }
-  close(DLIN);
-  print ' - read ' . $record_count . " records.\n";
-}
-else
-{
- die('Error! Failed to open file for reading: ' . $data_locality_csv_path);
-}
-
-# 2. Read in hadoop log looking for split mapping information
-print " * Read split information: hadoop.log\n";
-if (open(HLIN, '<:utf8', $hadoop_log_path))
-{
-  my $record_count = 0;
-  my $filepath = '';
-  while (my $line = <HLIN>)
-  {
-    if ($line =~ /^Filepath:\s(.+)$/)
-    {
-      $filepath = $1;
-    }
-    elsif ($filepath ne '')
-    {
-      # this is the information that was unintentionally lost
-      if ($line =~ /^\s-\soriginal\ssplit\s\[(.+)\]$/)
-      {
-        my $original_splits = &parseSplits($1);
-        if (!defined $data->{$filepath})
-        {
-          die('Found split information for file not in data locality log: ' . $filepath);
-        }
-        $data->{$filepath}->{'original_splits'} = $original_splits;
-        $filepath = '';
-        $record_count++;
-      }
-    }
-  }
-  close(HLIN);
-  print ' - updated ' . $record_count . " records.\n";
-}
-else
-{
- die('Error! Failed to open file for reading: ' . $hadoop_log_path);
-}
-
-# 3. Write out the updated data locality file including information about
-#    original splits (should we have found that information) and corrected
-#    data locality flag.
-print " * Writing out updated file: data_locality.csv\n";
-if (open(DLOUT, '>:utf8', $data_locality_csv_path))
-{
-  my $record_count = 0;
-  print DLOUT "TaskID,AttemptNo,Data Local,Compute Node,Splits,OriginalSplits\n";
-  foreach my $filepath (keys %{$data})
-  {
-    my $record = $data->{$filepath};
-    my $compute_node = $record->{'node'};
-    my @dl_nodes = split(',', $record->{'original_splits'});
-    my $data_local = 0;
-    foreach my $dl_node (@dl_nodes)
-    {
-      if ($dl_node eq $compute_node)
-      {
-        $data_local = 1;
-      }
-    }
-    print DLOUT $record->{'taskid'} . ',' . $record->{'attempt'} . ',' . $data_local . ',"' . $compute_node . '","' . $record->{'splits'} . '","' . $record->{'original_splits'} . '"' . "\n";
-    $record_count++;
-  }
-  close(DLOUT);
-  print ' - wrote ' . $record_count . " records\n";
-}
-else
-{
-  die('Error! Failed to open file for writing: ' . $data_locality_csv_path);
-}
-
-# 4. Complete
+# Complete!
 print "\n========================= Complete =========================\n\n";
 exit;
 
+################################################################################
+#################################  Functions  ##################################
+################################################################################
+
+
+## @function searchForDataLocalityFiles()
+#
+sub searchForDataLocalityFiles
+{
+  my ($dir) = @_;
+  # All of the following is *only* valid for directories
+  if (-d $dir)
+  {
+    print ' * Searching for Data Locality information: ' . $dir . "\n";
+    # Does this directory contain data_locality.csv?
+    my $dl_csv_path = $dir . '/data_locality.csv';
+    if (-f $dl_csv_path)
+    {
+      &updateDataLocalityForDirectory($dir);
+    }
+    # Also descend into child directories
+    if (opendir(DIR, $dir))
+    {
+      my @files = readdir(DIR);
+      closedir(DIR);
+      foreach my $file (@files)
+      {
+        my $path = $dir . '/' . $file;
+        if ($file =~ /^\./)
+        {
+        }
+        elsif (-d $path)
+        {
+          &searchForDataLocalityFiles($path);
+        }
+      }
+    }
+    else
+    {
+      die('Error! Failed to open directory for reading: ' . $dir);
+    }
+  }
+}
+## searchForDataLocalityFiles() ##
+
+
+sub updateDataLocalityForDirectory
+{
+  my ($dir_prefix) = @_;
+  # 0. Init
+  my $data = {};
+  my $data_locality_csv_path = $dir_prefix . '/data_locality.csv';
+  my $hadoop_log_path = $dir_prefix . '/hadoop.log';
+
+  # 1. Read in current data locality comma separated file
+  print "   * Read current file: data_locality.csv\n";
+  if (open(DLIN, '<:utf8', $data_locality_csv_path))
+  {
+    my $record_count = 0;
+    while (my $line = <DLIN>)
+    {
+      &printDebug($line);
+      # If the line already indicates this csv has been updated (by the
+      # presence of an "OriginalSplits" column) we don't need to do anything
+      # else
+      if ($line =~ /OriginalSplits/)
+      {
+        print "   ! File already updated...\n";
+        close(DLIN);
+        return;
+      }
+      if ($line =~ /^([^,]+),(\d),\d,"([^"]+)","([^"]+)"$/)
+      {
+        my $taskid       = $1;
+        my $attempt_no   = $2;
+        my $compute_node = $3;
+        my $splits       = $4;
+
+        # Use the taskid to look up the manifest file and therein the filepath
+        my $filepath = '';
+        my $manifest_path = $dir_prefix . '/manifest' . $taskid . '_' . $attempt_no . '.xml';
+        &printDebug('searching for manifest file: ' . $manifest_path);
+        if (open(MIN, '<:utf8', $manifest_path))
+        {
+          while (my $line2 = <MIN>)
+          {
+            if ($line2 =~ /<Filename>(hdfs|HDFSShell|ThriftFS):(.*)<\/Filename>/)
+            {
+              my $protocol = $1;
+              $filepath = 'hdfs:' . $2;
+            }
+          }
+          close(MIN);
+        }
+        else
+        {
+          die('Failed to open manifest for reading: ' . $manifest_path);
+        }
+        &printDebug('filepath: ' . $filepath);
+        $data->{$filepath} = {'taskid' => $taskid,
+                              'attempt' => $attempt_no,
+                              'node' => $compute_node,
+                              'splits' => $splits,
+                              'original_splits' => ''};
+        $record_count++;
+      }
+    }
+    close(DLIN);
+    print '   - read ' . $record_count . " records.\n";
+  }
+  else
+  {
+    die('Error! Failed to open file for reading: ' . $data_locality_csv_path);
+  }
+
+  # 2. Read in hadoop log looking for split mapping information
+  print "   * Read split information: hadoop.log\n";
+  if (open(HLIN, '<:utf8', $hadoop_log_path))
+  {
+    my $record_count = 0;
+    my $filepath = '';
+    while (my $line = <HLIN>)
+    {
+      if ($line =~ /^Filepath:\s(.+)$/)
+      {
+        $filepath = $1;
+      }
+      elsif ($filepath ne '')
+      {
+        # this is the information that was unintentionally lost
+        if ($line =~ /^\s-\soriginal\ssplit\s\[(.+)\]$/)
+        {
+          my $original_splits = &parseSplits($1);
+          if (!defined $data->{$filepath})
+          {
+            die('Found split information for file not in data locality log: ' . $filepath);
+          }
+          $data->{$filepath}->{'original_splits'} = $original_splits;
+          $filepath = '';
+          $record_count++;
+        }
+      }
+    }
+    close(HLIN);
+    print '   - updated ' . $record_count . " records.\n";
+  }
+  else
+  {
+    die('Error! Failed to open file for reading: ' . $hadoop_log_path);
+  }
+
+  # 3. Write out the updated data locality file including information about
+  #    original splits (should we have found that information) and corrected
+  #    data locality flag.
+  # backup old
+  my $backup_path = $data_locality_csv_path . '.bak';
+  if (-f $backup_path)
+  {
+    unlink($backup_path);
+  }
+  rename($data_locality_csv_path, $backup_path);
+  # write new
+  print "   * Writing out updated file: data_locality.csv\n";
+  if (open(DLOUT, '>:utf8', $data_locality_csv_path))
+  {
+    my $record_count = 0;
+    print DLOUT "TaskID,AttemptNo,Data Local,Compute Node,Splits,OriginalSplits\n";
+    foreach my $filepath (keys %{$data})
+    {
+      my $record = $data->{$filepath};
+      my $compute_node = $record->{'node'};
+      my @dl_nodes = split(',', $record->{'original_splits'});
+      my $data_local = 0;
+      foreach my $dl_node (@dl_nodes)
+      {
+        if ($dl_node eq $compute_node)
+        {
+          $data_local = 1;
+        }
+      }
+      print DLOUT $record->{'taskid'} . ',' . $record->{'attempt'} . ',' . $data_local . ',"' . $compute_node . '","' . $record->{'splits'} . '","' . $record->{'original_splits'} . '"' . "\n";
+      $record_count++;
+    }
+    close(DLOUT);
+    print '   - wrote ' . $record_count . " records\n";
+  }
+  else
+  {
+    die('Error! Failed to open file for writing: ' . $data_locality_csv_path);
+  }
+}
+## updateDataLocalityForDirectory() ##
+
+
+## @function parseSplits()
+#
 sub parseSplits
 {
@@ -154,5 +258,9 @@
 {
   my ($msg) = @_;
-  print '[DEBUG] ' . $msg . "\n";
+  if ($debug)
+  {
+    chomp($msg);
+    print '[DEBUG] ' . $msg . "\n";
+  }
 }
 
