<commit_before>require "fileutils"
require "pwrake/logger"
require "pwrake/nbio"
require "pwrake/option/option"
require "pwrake/task/task_wrapper"
require "pwrake/queue/task_queue"
require "pwrake/master/fiber_pool"

module Pwrake

  class Master

    def initialize
      @selector = NBIO::Selector.new
      @hostinfo_by_taskname = {}
      @option = Option.new
      @hdl_set = []
      @channel_by_hostid = {}
      @channels = []
      @hostinfo_by_id = {}
      init_logger
    end

    attr_reader :task_queue
    attr_reader :option
    attr_reader :logger

    def init_logger
      if logdir = @option['LOG_DIR']
        ::FileUtils.mkdir_p(logdir)
        logfile = File.join(logdir,@option['LOG_FILE'])
        @logger = Logger.new(logfile)
      else
        if @option['DEBUG']
          @logger = Logger.new($stderr)
        else
          @logger = Logger.new(File::NULL)
        end
      end

      if @option['DEBUG']
        @logger.level = Logger::DEBUG
      else
        @logger.level = Logger::INFO
      end
    end

    def init(hosts=nil)
      @option.init
      TaskWrapper.init_task_logger(@option)
    end

    def setup_branch_handler(sub_host)
      ior,w0 = IO.pipe
      r2,iow = IO.pipe
      if sub_host == "localhost" && /^(n|f)/i !~ ENV['T']
        @thread = Thread.new(r2,w0,@option) do |r,w,o|
          Rake.application.run_branch_in_thread(r,w,o)
        end
      else
        dir = File.absolute_path(File.dirname($PROGRAM_NAME))
        cmd = "ssh -x -T -q #{sub_host} '" +
          "cd \"#{Dir.pwd}\";"+
          "PATH=#{dir}:${PATH} exec pwrake_branch'"
        Log.debug("BranchCommunicator cmd=#{cmd}")
        spawn(cmd,:pgroup=>true,:out=>w0,:in=>r2)
        w0.close
        r2.close
        Marshal.dump(@option,iow)
        iow.flush
        s = ior.gets
        if !s or s.chomp != "pwrake_branch start"
          raise RuntimeError,"pwrake_branch start failed: receive #{s.inspect}"
        end
      end
      rd = NBIO::Reader.new(@selector,ior)
      wt = NBIO::Writer.new(@selector,iow)
      return NBIO::Handler.new(rd,wt,sub_host)
    end

    def signal_trap(sig)
      case @killed
      when 0
        # log writing failed. can't be called from trap context
        if Rake.application.options.debug
          $stderr.puts "\nSignal trapped. (sig=#{sig} pid=#{Process.pid}"+
            " thread=#{Thread.current} ##{@killed})"
          $stderr.puts caller
        else
          $stderr.puts "\nSignal trapped. (sig=#{sig} pid=#{Process.pid}"+
            " ##{@killed})"
        end
        $stderr.puts "Exiting..."
        @no_more_run = true
        @failed = true
        NBIO::Handler.kill(@hdl_set,sig)
        # @selector.run : not required here
      when 1
        $stderr.puts "\nOnce more Ctrl-C (SIGINT) for exit."
      when 2
        @thread.kill if @thread
      else
        Kernel.exit(false) # must wait for nomral exit
      end
      @killed += 1
    end

    def setup_branches
      sum_ncore = 0
      @option.host_map.each do |sub_host, wk_hosts|
        @hdl_set << hdl = setup_branch_handler(sub_host)
        Fiber.new do
        hdl.put_line "host_list_begin"
        wk_hosts.each do |host_info|
          name = host_info.name
          ncore = host_info.ncore
          host_id = host_info.id
          Log.debug "connecting #{name} ncore=#{ncore} id=#{host_id}"
          hdl.put_line "host:#{host_id} #{name} #{ncore}"
          @channel_by_hostid[host_id] = hdl
          @hostinfo_by_id[host_id] = host_info
        end
        hdl.put_line "host_list_end"
        while s = hdl.get_line
          case s
          when /^ncore:done$/
            break
          when /^ncore:(\d+):(\d+)$/
            id, ncore = $1.to_i, $2.to_i
            Log.debug "worker_id=#{id} ncore=#{ncore}"
            @hostinfo_by_id[id].set_ncore(ncore)
            sum_ncore += ncore
          when /^exited$/
            raise RuntimeError,"Unexpected branch exit"
          else
            msg = "#{hdl.host}:#{s.inspect}"
            raise RuntimeError,"invalid return: #{msg}"
          end
        end
        end.resume
      end
      @selector.run

      Log.info "num_cores=#{sum_ncore}"
      @hostinfo_by_id.each do |id,host|
        if ncore = @hostinfo_by_id[id].idle_cores
          Log.info "#{host} id=#{id} ncore=#{ncore}"
        else
          @hostinfo_by_id.delete(id)
        end
      end
      if @hostinfo_by_id.empty?
        raise RuntimeError,"no worker host"
      end
      queue_class = Pwrake.const_get(@option.queue_class)
      @task_queue = queue_class.new(@hostinfo_by_id)

      @branch_setup_thread = Thread.new do
        #@channels.each do |chan|
        create_fiber(@hdl_set) do |hdl|
          while s = hdl.get_line
            case s
            when /^retire:(\d+)$/
              retire($1.to_i)
            when /^branch_setup:done$/
              break
            else
              raise RuntimeError,"branch_setup failed: s=#{s.inspect}"
            end
          end
        end
        @selector.run
        @killed = 0
        [:TERM,:INT].each do |sig|
          Signal.trap(sig) do
            signal_trap(sig)
          end
        end
      end

    end

    def retire(hid)
      host_info = @hostinfo_by_id[hid.to_i]
      if host_info.decrease(1)
        # all retired
        m = "retired: host #{host_info.name}"
        Log.warn(m)
        $stderr.puts(m)
        drop_host(host_info) # delete from hostinfo_by_id
        if @hostinfo_by_id.empty?
          raise RuntimeError,"no worker host"
        end
      end
    end

    def create_fiber(channels,&blk)
      channels.each do |chan|<commit_msg>check @hostinfo_by_id[hid.to_i] is non-nil<commit_after>require "fileutils"
require "pwrake/logger"
require "pwrake/nbio"
require "pwrake/option/option"
require "pwrake/task/task_wrapper"
require "pwrake/queue/task_queue"
require "pwrake/master/fiber_pool"

module Pwrake

  class Master

    def initialize
      @selector = NBIO::Selector.new
      @hostinfo_by_taskname = {}
      @option = Option.new
      @hdl_set = []
      @channel_by_hostid = {}
      @channels = []
      @hostinfo_by_id = {}
      init_logger
    end

    attr_reader :task_queue
    attr_reader :option
    attr_reader :logger

    def init_logger
      if logdir = @option['LOG_DIR']
        ::FileUtils.mkdir_p(logdir)
        logfile = File.join(logdir,@option['LOG_FILE'])
        @logger = Logger.new(logfile)
      else
        if @option['DEBUG']
          @logger = Logger.new($stderr)
        else
          @logger = Logger.new(File::NULL)
        end
      end

      if @option['DEBUG']
        @logger.level = Logger::DEBUG
      else
        @logger.level = Logger::INFO
      end
    end

    def init(hosts=nil)
      @option.init
      TaskWrapper.init_task_logger(@option)
    end

    def setup_branch_handler(sub_host)
      ior,w0 = IO.pipe
      r2,iow = IO.pipe
      if sub_host == "localhost" && /^(n|f)/i !~ ENV['T']
        @thread = Thread.new(r2,w0,@option) do |r,w,o|
          Rake.application.run_branch_in_thread(r,w,o)
        end
      else
        dir = File.absolute_path(File.dirname($PROGRAM_NAME))
        cmd = "ssh -x -T -q #{sub_host} '" +
          "cd \"#{Dir.pwd}\";"+
          "PATH=#{dir}:${PATH} exec pwrake_branch'"
        Log.debug("BranchCommunicator cmd=#{cmd}")
        spawn(cmd,:pgroup=>true,:out=>w0,:in=>r2)
        w0.close
        r2.close
        Marshal.dump(@option,iow)
        iow.flush
        s = ior.gets
        if !s or s.chomp != "pwrake_branch start"
          raise RuntimeError,"pwrake_branch start failed: receive #{s.inspect}"
        end
      end
      rd = NBIO::Reader.new(@selector,ior)
      wt = NBIO::Writer.new(@selector,iow)
      return NBIO::Handler.new(rd,wt,sub_host)
    end

    def signal_trap(sig)
      case @killed
      when 0
        # log writing failed. can't be called from trap context
        if Rake.application.options.debug
          $stderr.puts "\nSignal trapped. (sig=#{sig} pid=#{Process.pid}"+
            " thread=#{Thread.current} ##{@killed})"
          $stderr.puts caller
        else
          $stderr.puts "\nSignal trapped. (sig=#{sig} pid=#{Process.pid}"+
            " ##{@killed})"
        end
        $stderr.puts "Exiting..."
        @no_more_run = true
        @failed = true
        NBIO::Handler.kill(@hdl_set,sig)
        # @selector.run : not required here
      when 1
        $stderr.puts "\nOnce more Ctrl-C (SIGINT) for exit."
      when 2
        @thread.kill if @thread
      else
        Kernel.exit(false) # must wait for nomral exit
      end
      @killed += 1
    end

    def setup_branches
      sum_ncore = 0
      @option.host_map.each do |sub_host, wk_hosts|
        @hdl_set << hdl = setup_branch_handler(sub_host)
        Fiber.new do
        hdl.put_line "host_list_begin"
        wk_hosts.each do |host_info|
          name = host_info.name
          ncore = host_info.ncore
          host_id = host_info.id
          Log.debug "connecting #{name} ncore=#{ncore} id=#{host_id}"
          hdl.put_line "host:#{host_id} #{name} #{ncore}"
          @channel_by_hostid[host_id] = hdl
          @hostinfo_by_id[host_id] = host_info
        end
        hdl.put_line "host_list_end"
        while s = hdl.get_line
          case s
          when /^ncore:done$/
            break
          when /^ncore:(\d+):(\d+)$/
            id, ncore = $1.to_i, $2.to_i
            Log.debug "worker_id=#{id} ncore=#{ncore}"
            @hostinfo_by_id[id].set_ncore(ncore)
            sum_ncore += ncore
          when /^exited$/
            raise RuntimeError,"Unexpected branch exit"
          else
            msg = "#{hdl.host}:#{s.inspect}"
            raise RuntimeError,"invalid return: #{msg}"
          end
        end
        end.resume
      end
      @selector.run

      Log.info "num_cores=#{sum_ncore}"
      @hostinfo_by_id.each do |id,host|
        if ncore = @hostinfo_by_id[id].idle_cores
          Log.info "#{host} id=#{id} ncore=#{ncore}"
        else
          @hostinfo_by_id.delete(id)
        end
      end
      if @hostinfo_by_id.empty?
        raise RuntimeError,"no worker host"
      end
      queue_class = Pwrake.const_get(@option.queue_class)
      @task_queue = queue_class.new(@hostinfo_by_id)

      @branch_setup_thread = Thread.new do
        #@channels.each do |chan|
        create_fiber(@hdl_set) do |hdl|
          while s = hdl.get_line
            case s
            when /^retire:(\d+)$/
              retire($1.to_i)
            when /^branch_setup:done$/
              break
            else
              raise RuntimeError,"branch_setup failed: s=#{s.inspect}"
            end
          end
        end
        @selector.run
        @killed = 0
        [:TERM,:INT].each do |sig|
          Signal.trap(sig) do
            signal_trap(sig)
          end
        end
      end

    end

    def retire(hid)
      host_info = @hostinfo_by_id[hid.to_i]
      if host_info && host_info.decrease(1)
        # all retired
        m = "retired: host #{host_info.name}"
        Log.warn(m)
        $stderr.puts(m)
        drop_host(host_info) # delete from hostinfo_by_id
        if @hostinfo_by_id.empty?
          raise RuntimeError,"no worker host"
        end
      end
    end

    def create_fiber(channels,&blk)
      channels.each do |chan|
{'fraction_non_alphanumeric': 0.12088395844338792, 'fraction_numerical': 0.003169572107765452, 'mean_word_length': 1.9006639427987742, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5080291', 'n_tokens_mistral': 3879, 'n_tokens_neox': 3800, 'n_words': 763}