pmux 0.1.0

data/lib/pmux/multi_session.rb

@@ -0,0 +1,309 @@
+require 'net/scp'
+
+class MR::Future
+  attr_accessor :addr
+end
+
+module Pmux
+  class SessionWrapper
+    attr_reader :addr
+    attr_accessor :ssh, :scp_session_count
+
+    def initialize addr
+      @addr = addr
+      @ssh = nil
+      @scp = nil
+      @scp_session_count = 0
+    end
+
+    def scp
+      @ssh ? (@scp ||= @ssh.scp) : nil
+    end
+  end
+
+  class MultiSession
+    attr_reader :loop
+    attr_accessor :timeout
+
+    def initialize addrs, options={}, loop=nil
+      @addrs = addrs
+      @options = options
+      @user = options[:user]
+      @loop = loop || Coolio::Loop.default
+      Net::SSH::Compat.coolio_loop = @loop
+      @timeout = 3600 #FIXME
+
+      @err_addrs = []
+      @sessions = {}
+      @channels = {}
+
+      @scptable = {}
+      @scpid = 0
+      @scps = {}
+      @scp_queue = {}
+      @buffers = {}
+
+      @on_error = nil
+    end
+
+    def connect_to_addr addr, cmd=nil
+      return if @sessions[addr]
+      @sessions[addr] = SessionWrapper.new addr
+      @loop.start_ssh(addr, @user) {|ssh|
+        if ssh.respond_to? :open_channel
+          host = ssh.host
+          @sessions[host].ssh = ssh
+          if cmd
+            channel = ssh.open_channel {|ch|
+              ch.exec(cmd) {|ch, success|
+                raise RuntimeError unless success
+                @channels[host] = ch
+
+                if (ary = @buffers[host]) and !ary.empty?
+                  ch.send_data ary.join('')
+                  ary.clear
+                end
+
+                setup_channel ch
+              }
+            }
+          end
+
+          if (queue = @scp_queue[host]) and !queue.empty?
+            queue2 = queue.dup
+            queue.clear
+            scp = @sessions[host].scp
+            queue2.each {|updown, qf, qaddr, qremote, qlocal, qoptions|
+                scp_download_sub scp, qaddr, qf, qremote, qlocal, qoptions}
+          end
+
+          ssh.floop 0.1
+        elsif ssh.kind_of? Exception
+          e = ssh
+          error_on_addr addr, e.inspect
+        else
+          error_on_addr addr, 'failed'
+        end
+      }
+    end
+
+    def setup_channel ch
+      ch.on_data {|c, data|}
+      ch.on_extended_data {|c, type, data|
+        #STDERR.puts c.connection.host+': '+data
+      }
+      #ch.on_close {|c| error_on_addr ch.connection.host}
+    end
+
+    def error_on_addr addr, err=nil
+      @err_addrs.push addr
+      @addrs.delete addr
+      @on_error.call addr, err if @on_error
+    end
+
+    def close_channel addr=nil
+      if addr
+        if (ch = @channels[addr])
+          ch.close
+          ch.wait
+        else
+        end
+      else
+        @channels.each {|addr, ch| ch.close}
+      end
+    end
+
+    def on_error &block
+      @on_error = block
+    end
+
+    def scp_upload_files addr, files, remote, options={}, &block
+      mf = MR::MultiFuture.new
+      for file in files
+        future = scp_upload addr, file, remote, options
+        mf.add future
+      end
+      mf.on_all &block if block
+      mf
+    end
+
+    def scp_upload addr, local, remote, options={}
+      future = MR::Future.new self, @loop
+      queue = (@scp_queue[addr] ||= [])
+
+      if (scp = @sessions[addr].scp)
+      #if (scp = @scps[addr]) or
+      #     (@sessions[addr] and @sessions[addr].ssh and
+      #      scp = (@scps[addr] = @sessions[addr].scp))
+        scp_upload_sub scp, addr, future, local, remote, options
+      else
+        queue.push [:up, future, addr, remote, local, options]
+      end
+      future
+    end
+
+    def scp_upload_sub scp, addr, future, local, remote, options
+      scpid = @scpid
+      @scpid += 1
+      @scptable[scpid] = future
+      scp.upload(local, remote, options) {|ch, name, sent, total|
+        if sent >= total
+          future.set_result nil, sent
+          @scptable.delete scpid
+        end
+      }
+    end
+
+    def scp_download addr, remote, local, options={}
+      future = MR::Future.new self, @loop
+      queue = (@scp_queue[addr] ||= [])
+
+      if (scp = @sessions[addr].scp)
+      #if (scp = @scps[addr]) or
+      #     (@sessions[addr] and @sessions[addr].ssh and
+      #      scp = (@scps[addr] = @sessions[addr].scp))
+        scp_download_sub scp, addr, future, remote, local, options
+      else
+        queue.push [:down, future, addr, remote, local, options]
+      end
+      future
+    end
+
+    def scp_download_sub scp, addr, future, remote, local, options
+      session = @sessions[addr]
+      if !session or session.scp_session_count > 5
+        queue = (@scp_queue[addr] ||= [])
+        queue.push [:down, future, addr, remote, local, options]
+        return
+      end
+      session.scp_session_count += 1
+
+      scpid = @scpid
+      @scpid += 1
+      @scptable[scpid] = future
+      channel = scp.download(remote, local, options)
+      channel.on_eof {|ch|
+        session.scp_session_count -= 1
+        @loop.set_timer(0) {process_scp_queue_once addr}
+
+        future.set_result(nil, options[:set_result])
+        @scptable.delete scpid
+      }
+      channel.on_open_failed {|ch, code, desc|
+        Log.error "#{addr}: scp error: #{desc}"
+        err = RuntimeError.new "scp error: #{desc}"
+        @on_error.call addr, err
+        session.scp_session_count -= 1
+        @loop.set_timer(0) {process_scp_queue_once addr}
+
+        future.set_result(nil, options[:set_result])
+        @scptable.delete scpid
+      }
+    end
+
+    def process_scp_queue_once addr
+      #scp = @scps[addr]
+      scp = @sessions[addr].scp
+      queue = (@scp_queue[addr] ||= [])
+      if scp and !queue.empty?
+        updown, future, addr, remote, local, options = queue.shift
+        case updown
+        when :down
+          scp_download_sub scp, addr, future, remote, local, options
+        when :up
+          scp_upload_sub scp, addr, future, local, remote, options
+        end
+      end
+    end
+  end
+
+  class MRSession < MultiSession
+    include MR::MessageReceiver
+
+    def initialize addrs, options={}, loop=nil
+      super
+
+      @reqtable = {}
+      @seqid = 0
+
+      program_name = options[:program_name] || 'pmux'
+      @cmd = "#{program_name} --server"
+    end
+
+    def connect
+      for addr in @addrs
+        connect_to_addr addr, @cmd + " --ipaddr=#{addr}"
+      end
+    end
+
+    def setup_channel ch
+      pac = MessagePack::Unpacker.new
+      ch.on_data {|c, data|
+        pac.feed_each(data) {|obj| on_message obj}
+      }
+      ch.on_extended_data {|c, type, data|
+        #STDERR.puts c.connection.host+': '+data
+      }
+      ch.on_close {|c| error_on_addr ch.connection.host}
+    end
+
+    def error_on_addr addr, err=nil
+      super
+      err ||= 'closed'
+      @reqtable.select {|msgid, f| f.addr == addr}.each {|msgid, f|
+        f.set_result err, nil
+      }
+    end
+
+    def call_async addr, method, *args
+      send_request addr, method, args
+    end
+
+    def multicast_call_async method, *args
+      mf = MR::MultiFuture.new
+      for addr in @addrs
+        future = send_request addr, method, args
+        mf.add future
+      end
+      mf
+    end
+
+    def on_response msgid, error, result
+      if (future = @reqtable.delete msgid)
+        future.set_result error, result
+      end
+    end
+
+    private
+    def send_request addr, method, param
+      method = method.to_s
+      msgid = @seqid
+      @seqid += 1; if @seqid >= 1<<31 then @seqid = 0 end
+      data = [MR::REQUEST, msgid, method, param].to_msgpack
+      if (ch = @channels[addr])
+        ch.send_data data
+      else
+        (@buffers[addr] ||= []).push data
+      end
+      future = MR::Future.new self, @loop
+      future.addr = addr
+      @reqtable[msgid] = future
+    end
+  end
+end
+
+class MR::MultiFuture
+  # monkey patch for MR::MultiFuture#join_all
+  def join_all
+    @not_joined.dup.each {|future|
+      future.join
+    }
+    @all
+  end
+end
+
+class MR::UNIXClientTransport::ClientSocket
+  def on_close
+    raise MR::TransportError, 'MR::UNIXClientTransport::ClientSocket on close'
+  end
+end

data/lib/pmux/pipeio.rb

@@ -0,0 +1,19 @@
+module Pmux
+  class PipeIO < Coolio::IO
+    def initialize cmd, mode='r'
+      io = IO.popen cmd, mode
+      super io
+      @on_receive = nil
+    end
+
+    def on_receive &block
+      @on_receive = block
+    end
+
+    def on_read data
+      if @on_receive
+        @on_receive.call data
+      end
+    end
+  end
+end

data/lib/pmux/plugin.rb

@@ -0,0 +1,23 @@
+module Pmux
+  class PluginClass
+    def initialize
+    end
+
+    def load_plugins
+      dir = File.join File.dirname(__FILE__), 'plugin'
+      load_plugin_dir dir
+    end
+
+    def load_plugin_dir dir
+      dir = File.expand_path dir
+      return unless File.directory? dir
+      Dir.entries(dir).sort.each {|fname|
+        if fname =~ /\.rb$/
+          require File.join(dir, fname)
+        end
+      }
+    end
+  end
+
+  Plugin = PluginClass.new
+end

data/lib/pmux/q.rb

@@ -0,0 +1,3 @@
+def q(*args)
+  STDERR.puts args.map {|item| item.inspect}.join(', ')
+end

data/lib/pmux/reducer.rb

@@ -0,0 +1,90 @@
+module Pmux
+  class Reducer
+    attr_reader :exitstatus, :output_path
+    attr_accessor :tmp_dir
+
+    def initialize task, tmp_dir, loop=nil
+      @task = task
+      @tmp_dir = tmp_dir
+      @loop = loop
+      @paths = Dir.glob("#{tmp_dir}/t*-#{task['pindex']}")
+      @output_path = nil
+      @on_receive = nil
+      @on_success = nil
+      @on_error = nil
+      raise RuntimeError, 'no input files' if @paths.empty?
+    end
+
+    def do_reduce_task; end
+  end
+
+  class StreamingReducer < Reducer
+    include FixCmdLine
+
+    def on_receive &block
+      @on_receive = block
+    end
+    def on_success &block
+      @on_success = block
+    end
+    def on_error &block
+      @on_error = block
+    end
+
+    def do_reduce_task
+      reducer_cmd = @task['reducer'] || 'cat'
+      @output_path = "#{@tmp_dir}/r#{@task['pindex']}"
+      err_path = "#{@tmp_dir}/.err.#{$$}"
+      err_msg = nil
+      cmd_line = fix_cmd_line reducer_cmd,
+        @paths.join(' '), @output_path, err_path, tmp_dir
+      Log.debug "system: #{cmd_line}"
+      system cmd_line
+      @exitstatus = $?.exitstatus
+      if File.size? err_path
+        err_msg = File.read(err_path).chomp!
+        raise RuntimeError, err_msg
+      end
+      if @exitstatus > 1
+        raise RuntimeError, "failed to execute reducer: #{cmd_line}"
+      end
+      @output_path
+    end
+
+    def do_streaming_reduce_task
+      reducer_cmd = @task['reducer'] || 'cat'
+      @output_path = "#{@tmp_dir}/r#{@task['pindex']}"
+      err_path = "#{@tmp_dir}/.rerr.#{$$}"
+      err_msg = nil
+      cmd_line = fix_cmd_line reducer_cmd,
+        @paths.join(' '), nil, err_path, tmp_dir
+      Log.debug "popen: #{cmd_line}"
+      pipeio = PipeIO.new cmd_line
+      if @on_receive
+        pipeio.on_receive &@on_receive
+      else
+        out = open(@output_path, 'a')
+        pipeio.on_receive {|data|
+          out.write data
+        }
+      end
+      on_success = @on_success
+      on_error = @on_error
+      pipeio.on_close {
+        if out
+          out.close rescue nil
+        end
+        if File.size? err_path
+          err_msg = File.read(err_path).chomp!
+          #raise RuntimeError, err_msg
+          e = RuntimeError.new err_msg
+          e.set_backtrace ['reducer']
+          on_error.call e if on_error
+        else
+          on_success.call self if on_success
+        end
+      }
+      @loop.attach pipeio
+    end
+  end
+end

data/lib/pmux/storage_adapter.rb

@@ -0,0 +1,105 @@
+require 'socket'
+require 'forwardable'
+
+module Pmux
+  class StorageAdapter
+    Adapters = {}
+    extend Forwardable
+    def_delegators :@h, :each, :size, :[], :[]=, :update, :delete
+
+    def self.create name, addrs, options={}
+      adapter_class = Adapters[name] || GlusterFSAdapter
+      adapter_class.new addrs, options
+    end
+
+    attr_reader :addrs
+
+    def initialize addrs=[], options={}
+      @addrs = addrs
+      @options = options
+      @h = init_node_hash addrs
+    end
+
+    def init_node_hash addrs
+      Hash[*(addrs.map {|addr| [addr, {}]}).flatten]
+    end
+
+    def connect_to_storage locator_host, locator_port
+    end
+  end
+end
+
+module Pmux
+  class LocalAdapter < StorageAdapter
+    Adapters['local'] = self
+
+    def get_files args, glob_flag=false
+      if glob_flag
+        Dir.glob args.join("\0")
+      else
+        args
+      end
+    end
+
+    def _get_file_locations files, glob_flag=false
+      files = Dir.glob files.join("\0") if glob_flag
+      files.map {|file|
+        a = addrs.map {|addr| [addr, file]}
+        addrs.push addrs.shift
+        [file, a]
+      }
+    end
+
+    def lookup_file file
+      res = addrs.map {|addr| [addr, file]}
+      addrs.push addrs.shift
+      res
+    end
+  end
+end
+
+module Pmux
+  class GlusterFSAdapter < StorageAdapter
+    Adapters['glusterfs'] = self
+
+    def getaddr host
+      sa = Socket.pack_sockaddr_in 0, host
+      port, addr = Socket.unpack_sockaddr_in sa
+      addr
+    end
+
+    def connect_to_storage locator_host, locator_port
+      locator_port ||= 7076
+      @client = MR::Client.new locator_host, locator_port
+      @client.timeout = 3600 #FIXME
+      @client
+    end
+
+    def get_files args, glob_flag=false
+      raise RuntimeError, 'not connected' unless @client
+      result = @client.call :get_locations, args, glob_flag
+      @locations = result
+
+      # hostname -> addr
+      hosts = {}
+      for vs in result.values
+        for host, path in vs
+          hosts[host] = true
+        end
+      end
+      @host2addr = {}
+      @addrs = hosts.keys.map {|host| addr = getaddr host
+        @host2addr[host] = addr}
+      @h = init_node_hash @addrs
+      result.keys # files
+    end
+
+    def lookup_file file
+      if (res = @locations[file])
+        res.map {|host, path| [@host2addr[host], path]}
+      else
+        nil
+      end
+    end
+  end
+end

data/lib/pmux/task_dispatcher.rb

@@ -0,0 +1,167 @@
+module Pmux
+  class TaskDispatcher
+    attr_reader :options, :adapter
+    attr_reader :msession, :scheduler, :gatherer
+    attr_reader :mf_shuffle
+    attr_reader :jl
+
+    def initialize options, adapter, msession, gatherer=nil
+      @options = options
+      @adapter = adapter
+      @msession = msession
+      @gatherer = gatherer || Gatherer.new(STDOUTWriter.new)
+      @scheduler = TaskScheduler.new adapter
+
+      @verbose = options[:verbose]
+      @on_error = proc {|r|
+        $stderr.write "%s: %s, %s\n" %
+          [r['error'], r['error_message'], r['backtrace']]
+      }
+    end
+
+    def run job
+      if job.num_r.nonzero?
+        @mf_shuffle = MR::MultiFuture.new
+      end
+      @jl = Joblogger.new options[:log_dir], job
+      scheduler.push_job job
+      scheduler.attach_flush_callback {|node_addr, task|
+        t = task.dup
+        [:node_addrs, :alloc_time].each {|k| t.delete k}
+        Log.info "send task #{t[:task_id]} to #{node_addr}"
+        printf "send task %s to %s\n", t[:task_id], node_addr if @verbose
+        future = msession.call_async node_addr, :exec_streaming_task, t
+        future.attach_callback {|f| receive_result job, f.get}
+        # err callback ?
+      }
+
+      # init job
+      wdir = "#{options[:tmp_dir]}/w#{job.id}"
+      Dir.mkdir wdir
+      jl.dump_header
+      jl.sep
+      Log.init "#{options[:log_dir] or wdir}/dispatcher.log"
+
+      puts 'send "init_job"' if @verbose
+      mf_init = msession.multicast_call_async :init_job, job.id
+      mf_init.on_success {|future|
+        addr = future.addr
+        res = future.get
+        job_dir = adapter[addr][:job_dir] = res['job_dir'] # remote job_dir
+        adapter[addr][:num_workers] = res['num_cpu'] || 2
+        printf "%s: remote job_dir = %s\n", addr, res['job_dir'] if @verbose
+        # scp ship_files to remote job_dir
+        if (ship_files = options[:ship_files])
+          mf_scp = msession.scp_upload_files addr, ship_files, job_dir
+          mf_scp.on_all {
+            scheduler.shipped[addr] = true
+            scheduler.process_queue
+          }
+        else
+          puts 'start scheduler' if @verbose
+          scheduler.shipped[addr] = true
+          scheduler.process_queue
+        end
+      }
+      mf_init.on_error {job.taskhash.clear}
+
+      # wait for all map tasks to finish
+      until job.completed?
+        msession.loop.run_once
+      end
+      if job.num_r.zero?
+        gatherer.join_all
+      else
+        mf_shuffle.join_all
+        # reduce phase
+        job.mk_reduce_tasks
+        scheduler.inject_tasks job.tasks
+        scheduler.process_queue
+        # wait for all reduce tasks to finish
+        until job.completed?
+          msession.loop.run_once
+        end
+        gatherer.join_all
+      end
+
+      Log.info "END"
+      job[:end_time] = Time.now
+      jl.dump_footer
+      jl.close
+
+      mf_quit = msession.multicast_call_async :quit
+      mf_quit.join_all
+
+      cleaner = Cleaner.new "#{options[:tmp_dir]}/w*"
+      cleaner.run
+    end
+
+    def on_error &block
+      @on_error = block
+    end
+
+    def receive_result job, result
+      task_id, node_addr, ifbase =
+        result.values_at 'task_id', 'node_addr', 'ifbase'
+      Log.info "receive result #{task_id} from #{node_addr}"
+      jl.add task_id, :node_addr=>node_addr, :ifbase=>ifbase,
+        :welapse=>result['welapse']
+      puts "receive result #{task_id} from #{node_addr}" if @verbose
+      if result['error']
+        if @on_error
+          @on_error.call result
+        else
+        end
+      end
+
+      if (ifbase = result['ifbase'])
+        if job.num_r.zero?
+          # no reducers; get the intermediate file from the mapper node
+          ifpath = ifbase + '-0'
+          local = "#{@options[:tmp_dir]}/w#{job.id}/#{File.basename(ifpath).sub /^[mr]/, 'w'}"
+          if (body = result['result_body'])
+            if body.size > 0
+              open(local, 'w') {|f| f.write body}
+              gatherer.writer.write local
+            end
+          else
+            puts "gather #{node_addr}:#{ifpath} -> #{local}" if @verbose
+            Log.info "gather #{node_addr}:#{ifpath} -> #{local}"
+            gatherer.gather msession, node_addr, ifpath, local
+          end
+        else
+          # send 'notify_reduce' message to the reducer node
+          job.reducers.each_with_index {|reducer_addr, pindex|
+            puts "send notify_reduce #{task_id} to #{reducer_addr}" if @verbose
+            future = msession.call_async reducer_addr, :notify_reduce,
+              :job_id=>job.id, :task_id=>task_id, :pindex=>pindex,
+              :node_addr=>node_addr, :ifbase=>ifbase
+            mf_shuffle.add future
+          }
+        end
+      elsif (output_path = result['output_path'])
+        # reduced result
+        local = "#{@options[:tmp_dir]}/w#{job.id}/#{File.basename(output_path).sub /^[mr]/, 'w'}"
+        gatherer.gather msession, node_addr, output_path, local
+      end
+
+      task = job.get_task_by_id task_id
+      task[:elapse] = Time.now - task[:alloc_time] if task[:alloc_time]
+      task[:welapse] = result['welapse']
+
+      # delete task
+      scheduler.delete_task_from_job job, task, node_addr
+      if (task_keys = result['task_keys'])
+        for tid in task_keys.keys
+          job.delete_task_by_id tid
+        end
+      end
+
+      scheduler.process_queue
+
+      perc = 100 * (job.tasks.size - job.taskhash.size) / job.tasks.size
+      task_type = result['map'] ? 'map' : 'reduce'
+      Log.info "done #{task_type} task #{result['task_id']} (#{perc}%)"
+    end
+  end
+end