kraps 0.1.0

data/lib/kraps/drivers.rb

@@ -0,0 +1,37 @@
+require "attachie"
+require "attachie/s3_driver"
+require "attachie/fake_driver"
+
+module Kraps
+  module Drivers
+    module Driver
+      def with_prefix(path)
+        File.join(*[@prefix, path].compact)
+      end
+    end
+
+    class S3Driver
+      include Driver
+
+      attr_reader :driver, :bucket, :prefix
+
+      def initialize(s3_client:, bucket:, prefix: nil)
+        @driver = Attachie::S3Driver.new(s3_client)
+        @bucket = bucket
+        @prefix = prefix
+      end
+    end
+
+    class FakeDriver
+      include Driver
+
+      attr_reader :driver, :bucket, :prefix
+
+      def initialize(bucket:, prefix: nil)
+        @driver = Attachie::FakeDriver.new
+        @bucket = bucket
+        @prefix = prefix
+      end
+    end
+  end
+end

data/lib/kraps/frame.rb

@@ -0,0 +1,3 @@
+module Kraps
+  Frame = Struct.new(:token, :partitions, keyword_init: true)
+end

data/lib/kraps/interval.rb

@@ -0,0 +1,34 @@
+module Kraps
+  class Interval
+    include MonitorMixin
+
+    def initialize(timeout, &block)
+      super()
+
+      @thread_queue = TimeoutQueue.new
+      @main_queue = TimeoutQueue.new
+      @stopped = false
+
+      @thread = Thread.new do
+        until @stopped
+          item = @thread_queue.deq(timeout: timeout)
+
+          block.call unless @stopped
+
+          @main_queue.enq(1) if item
+        end
+      end
+    end
+
+    def fire(timeout:)
+      @thread_queue.enq(1)
+      @main_queue.deq(timeout: timeout)
+    end
+
+    def stop
+      @stopped = true
+      @thread_queue.enq(nil)
+      @thread.join
+    end
+  end
+end

data/lib/kraps/job.rb

@@ -0,0 +1,62 @@
+module Kraps
+  class Job
+    attr_reader :steps
+
+    def initialize(worker:)
+      @worker = worker
+      @steps = []
+      @partitions = 0
+      @partitioner = MapReduce::HashPartitioner.new(@partitions)
+    end
+
+    def parallelize(partitions:, partitioner: MapReduce::HashPartitioner.new(partitions), worker: @worker, &block)
+      fresh.tap do |job|
+        job.instance_eval do
+          @partitions = partitions
+          @partitioner = partitioner
+
+          @steps << Step.new(action: Actions::PARALLELIZE, args: { partitions: @partitions, partitioner: @partitioner, worker: worker }, block: block)
+        end
+      end
+    end
+
+    def map(partitions: nil, partitioner: nil, worker: @worker, &block)
+      fresh.tap do |job|
+        job.instance_eval do
+          @partitions = partitions if partitions
+          @partitioner = partitioner || MapReduce::HashPartitioner.new(partitions) if partitioner || partitions
+
+          @steps << Step.new(action: Actions::MAP, args: { partitions: @partitions, partitioner: @partitioner, worker: worker }, block: block)
+        end
+      end
+    end
+
+    def reduce(worker: @worker, &block)
+      fresh.tap do |job|
+        job.instance_eval do
+          @steps << Step.new(action: Actions::REDUCE, args: { partitions: @partitions, partitioner: @partitioner, worker: worker }, block: block)
+        end
+      end
+    end
+
+    def each_partition(worker: @worker, &block)
+      fresh.tap do |job|
+        job.instance_eval do
+          @steps << Step.new(action: Actions::EACH_PARTITION, args: { partitions: @partitions, partitioner: @partitioner, worker: worker }, block: block)
+        end
+      end
+    end
+
+    def repartition(partitions:, partitioner: nil, worker: @worker)
+      map(partitions: partitions, partitioner: partitioner, worker: worker) do |key, value, collector|
+        collector.call(key, value)
+      end
+    end
+
+    def fresh
+      dup.tap do |job|
+        job.instance_variable_set(:@steps, @steps.dup)
+      end
+    end
+  end
+end

data/lib/kraps/parallelizer.rb

@@ -0,0 +1,29 @@
+module Kraps
+  class Parallelizer
+    def self.each(enum, num_threads)
+      queue = Queue.new
+
+      enum.each { |element| queue.push element }
+
+      stopped = false
+
+      threads = Array.new(num_threads) do
+        Thread.new do
+          yield queue.pop(true) until stopped || queue.empty?
+        rescue ThreadError
+          # Queue empty
+        rescue StandardError => e
+          stopped = true
+
+          e
+        end
+      end
+
+      threads.each(&:join).each do |thread|
+        raise thread.value if thread.value.is_a?(Exception)
+      end
+
+      enum
+    end
+  end
+end

data/lib/kraps/runner.rb

@@ -0,0 +1,160 @@
+module Kraps
+  class Runner
+    def initialize(klass)
+      @klass = klass
+    end
+
+    def call(*args, **kwargs)
+      Array(@klass.new.call(*args, **kwargs)).tap do |jobs|
+        jobs.each_with_index do |job, job_index|
+          job.steps.each_with_index.inject(nil) do |frame, (_, step_index)|
+            StepRunner.new(
+              klass: @klass,
+              args: args,
+              kwargs: kwargs,
+              jobs: jobs,
+              job_index: job_index,
+              step_index: step_index,
+              frame: frame
+            ).call
+          end
+        end
+      end
+    end
+
+    class StepRunner
+      def initialize(klass:, args:, kwargs:, jobs:, job_index:, step_index:, frame:)
+        @klass = klass
+        @args = args
+        @kwargs = kwargs
+        @jobs = jobs
+        @job_index = job_index
+        @job = @jobs[@job_index]
+        @step_index = step_index
+        @step = @job.steps[@step_index]
+        @frame = frame
+      end
+
+      def call
+        raise(InvalidAction, "Invalid action #{@step.action}") unless Actions::ALL.include?(@step.action)
+
+        @step.frame ||= send(:"perform_#{@step.action}")
+      end
+
+      private
+
+      def perform_parallelize
+        enum = Enumerator.new do |yielder|
+          collector = proc { |item| yielder << item }
+
+          @step.block.call(collector)
+        end
+
+        with_distributed_job do |distributed_job|
+          push_and_wait(distributed_job, enum) do |item, part|
+            enqueue(token: distributed_job.token, part: part, item: item)
+          end
+
+          Frame.new(token: distributed_job.token, partitions: @step.args[:partitions])
+        end
+      end
+
+      def perform_map
+        with_distributed_job do |distributed_job|
+          push_and_wait(distributed_job, 0...@frame.partitions) do |partition, part|
+            enqueue(token: distributed_job.token, part: part, partition: partition)
+          end
+
+          Frame.new(token: distributed_job.token, partitions: @step.args[:partitions])
+        end
+      end
+
+      def perform_reduce
+        with_distributed_job do |distributed_job|
+          push_and_wait(distributed_job, 0...@frame.partitions) do |partition, part|
+            enqueue(token: distributed_job.token, part: part, partition: partition)
+          end
+
+          Frame.new(token: distributed_job.token, partitions: @step.args[:partitions])
+        end
+      end
+
+      def perform_each_partition
+        with_distributed_job do |distributed_job|
+          push_and_wait(distributed_job, 0...@frame.partitions) do |partition, part|
+            enqueue(token: distributed_job.token, part: part, partition: partition)
+          end
+
+          @frame
+        end
+      end
+
+      def enqueue(token:, part:, **rest)
+        Kraps.enqueuer.call(
+          @step.args[:worker],
+          JSON.generate(
+            job_index: @job_index,
+            step_index: @step_index,
+            frame: @frame.to_h,
+            token: token,
+            part: part,
+            klass: @klass,
+            args: @args,
+            kwargs: @kwargs,
+            **rest
+          )
+        )
+      end
+
+      def with_distributed_job
+        distributed_job = Kraps.distributed_job_client.build(token: SecureRandom.hex)
+
+        yield(distributed_job)
+      rescue Interrupt
+        distributed_job&.stop
+        raise
+      end
+
+      def push_and_wait(distributed_job, enum)
+        progress_bar = build_progress_bar("#{@klass}: job #{@job_index + 1}/#{@jobs.size}, step #{@step_index + 1}/#{@job.steps.size}, token #{distributed_job.token}, %a, %c/%C (%p%) => #{@step.action}")
+
+        begin
+          total = 0
+
+          interval = Interval.new(1) do
+            progress_bar.total = total
+          end
+
+          distributed_job.push_each(enum) do |item, part|
+            total += 1
+            interval.fire(timeout: 1)
+
+            yield(item, part)
+          end
+        ensure
+          interval&.stop
+        end
+
+        loop do
+          progress_bar.total = distributed_job.total
+          progress_bar.progress = progress_bar.total - distributed_job.count
+
+          break if distributed_job.finished? || distributed_job.stopped?
+
+          sleep(1)
+        end
+
+        raise(JobStopped, "The job was stopped") if distributed_job.stopped?
+      ensure
+        progress_bar&.stop
+      end
+
+      def build_progress_bar(format)
+        options = { format: format, total: 1, autofinish: false }
+        options[:output] = ProgressBar::Outputs::Null unless Kraps.show_progress?
+
+        ProgressBar.create(options)
+      end
+    end
+  end
+end

data/lib/kraps/step.rb

@@ -0,0 +1,3 @@
+module Kraps
+  Step = Struct.new(:action, :args, :block, :frame, keyword_init: true)
+end

data/lib/kraps/temp_path.rb

@@ -0,0 +1,29 @@
+module Kraps
+  class TempPath
+    attr_reader :path
+
+    def initialize(prefix: nil, suffix: nil)
+      @path = File.join(Dir.tmpdir, [prefix, SecureRandom.hex[0, 16], Process.pid, suffix].compact.join("."))
+
+      File.open(@path, File::CREAT | File::EXCL) {}
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@path))
+
+      return unless block_given?
+
+      begin
+        yield
+      ensure
+        unlink
+      end
+    end
+
+    def unlink
+      FileUtils.rm_f(@path)
+    end
+
+    def self.finalize(path)
+      proc { FileUtils.rm_f(path) }
+    end
+  end
+end

data/lib/kraps/temp_paths.rb

@@ -0,0 +1,34 @@
+module Kraps
+  class TempPaths
+    include MonitorMixin
+    include Enumerable
+
+    def initialize
+      super
+
+      @temp_paths = []
+    end
+
+    def add
+      synchronize do
+        temp_path = TempPath.new
+        @temp_paths << temp_path
+        temp_path
+      end
+    end
+
+    def unlink
+      synchronize do
+        @temp_paths.each(&:unlink)
+      end
+    end
+
+    def each(&block)
+      return enum_for(__method__) unless block_given?
+
+      synchronize do
+        @temp_paths.each(&block)
+      end
+    end
+  end
+end

data/lib/kraps/timeout_queue.rb

@@ -0,0 +1,27 @@
+module Kraps
+  class TimeoutQueue
+    include MonitorMixin
+
+    def initialize
+      super
+
+      @cond = new_cond
+      @queue = []
+    end
+
+    def enq(item)
+      synchronize do
+        @queue << item
+        @cond.signal
+      end
+    end
+
+    def deq(timeout:)
+      synchronize do
+        @cond.wait(timeout) if @queue.empty?
+
+        return @queue.empty? ? nil : @queue.shift
+      end
+    end
+  end
+end

data/lib/kraps/version.rb

@@ -0,0 +1,3 @@
+module Kraps
+  VERSION = "0.1.0"
+end

data/lib/kraps/worker.rb

@@ -0,0 +1,183 @@
+module Kraps
+  class Worker
+    def initialize(json, memory_limit:, chunk_limit:, concurrency:)
+      @args = JSON.parse(json)
+      @memory_limit = memory_limit
+      @chunk_limit = chunk_limit
+      @concurrency = concurrency
+    end
+
+    def call(retries: 3)
+      return if distributed_job.stopped?
+
+      raise(InvalidAction, "Invalid action #{step.action}") unless Actions::ALL.include?(step.action)
+
+      with_retries(retries) do # TODO: allow to use queue based retries
+        send(:"perform_#{step.action}")
+
+        distributed_job.done(@args["part"])
+      end
+    end
+
+    private
+
+    def perform_parallelize
+      implementation = Class.new do
+        def map(key)
+          yield(key, nil)
+        end
+      end
+
+      mapper = MapReduce::Mapper.new(implementation.new, partitioner: partitioner, memory_limit: @memory_limit)
+      mapper.map(@args["item"])
+
+      mapper.shuffle do |partition, tempfile| # TODO: upload in parallel
+        Kraps.driver.driver.store(Kraps.driver.with_prefix("#{@args["token"]}/#{partition}/chunk.#{@args["part"]}.json"), tempfile, Kraps.driver.bucket)
+      end
+    end
+
+    def perform_map
+      temp_paths = TempPaths.new
+
+      files = Kraps.driver.driver.list(Kraps.driver.bucket, prefix: Kraps.driver.with_prefix("#{@args["frame"]["token"]}/#{@args["partition"]}/")).sort
+
+      temp_paths_index = files.each_with_object({}) do |file, hash|
+        hash[file] = temp_paths.add
+      end
+
+      Parallelizer.each(files, @concurrency) do |file|
+        Kraps.driver.driver.download(file, Kraps.driver.bucket, temp_paths_index[file].path)
+      end
+
+      current_step = step
+
+      implementation = Object.new
+      implementation.define_singleton_method(:map) do |key, value, &block|
+        current_step.block.call(key, value, block)
+      end
+
+      mapper = MapReduce::Mapper.new(implementation, partitioner: partitioner, memory_limit: @memory_limit)
+
+      temp_paths.each do |temp_path|
+        File.open(temp_path.path) do |stream|
+          stream.each_line do |line|
+            key, value = JSON.parse(line)
+
+            mapper.map(key, value)
+          end
+        end
+      end
+
+      mapper.shuffle do |partition, tempfile| # TODO: upload in parallel
+        Kraps.driver.driver.store(
+          Kraps.driver.with_prefix("#{@args["token"]}/#{partition}/chunk.#{@args["part"]}.json"), tempfile, Kraps.driver.bucket
+        )
+      end
+    ensure
+      temp_paths&.unlink
+    end
+
+    def perform_reduce
+      current_step = step
+
+      implementation = Object.new
+      implementation.define_singleton_method(:reduce) do |key, value1, value2|
+        current_step.block.call(key, value1, value2)
+      end
+
+      reducer = MapReduce::Reducer.new(implementation)
+
+      Parallelizer.each(Kraps.driver.driver.list(Kraps.driver.bucket, prefix: Kraps.driver.with_prefix("#{@args["frame"]["token"]}/#{@args["partition"]}/")), @concurrency) do |file|
+        Kraps.driver.driver.download(file, Kraps.driver.bucket, reducer.add_chunk)
+      end
+
+      tempfile = Tempfile.new
+
+      reducer.reduce(chunk_limit: @chunk_limit) do |key, value|
+        tempfile.puts(JSON.generate([key, value]))
+      end
+
+      Kraps.driver.driver.store(Kraps.driver.with_prefix("#{@args["token"]}/#{@args["partition"]}/chunk.#{@args["part"]}.json"), tempfile.tap(&:rewind), Kraps.driver.bucket)
+    ensure
+      tempfile&.close(true)
+    end
+
+    def perform_each_partition
+      temp_paths = TempPaths.new
+
+      files = Kraps.driver.driver.list(Kraps.driver.bucket, prefix: Kraps.driver.with_prefix("#{@args["frame"]["token"]}/#{@args["partition"]}/")).sort
+
+      temp_paths_index = files.each_with_object({}) do |file, hash|
+        hash[file] = temp_paths.add
+      end
+
+      Parallelizer.each(files, @concurrency) do |file|
+        Kraps.driver.driver.download(file, Kraps.driver.bucket, temp_paths_index[file].path)
+      end
+
+      enum = Enumerator::Lazy.new(temp_paths) do |yielder, temp_path|
+        File.open(temp_path.path) do |stream|
+          stream.each_line do |line|
+            yielder << JSON.parse(line)
+          end
+        end
+      end
+
+      step.block.call(@args["partition"], enum)
+    ensure
+      temp_paths&.unlink
+    end
+
+    def with_retries(num_retries)
+      retries = 0
+
+      begin
+        yield
+      rescue Kraps::Error
+        distributed_job.stop
+      rescue StandardError
+        sleep(5)
+        retries += 1
+
+        if retries >= num_retries
+          distributed_job.stop
+          raise
+        end
+
+        retry
+      end
+    end
+
+    def jobs
+      @jobs ||= Array(@args["klass"].constantize.new.call(*@args["args"], **@args["kwargs"].transform_keys(&:to_sym)))
+    end
+
+    def job
+      @job ||= begin
+        job_index = @args["job_index"]
+
+        jobs[job_index] || raise(InvalidJob, "Can't find job #{job_index}")
+      end
+    end
+
+    def steps
+      @steps ||= job.steps
+    end
+
+    def step
+      @step ||= begin
+        step_index = @args["step_index"]
+
+        steps[step_index] || raise(InvalidStep, "Can't find step #{step_index}")
+      end
+    end
+
+    def partitioner
+      @partitioner ||= step.args[:partitioner]
+    end
+
+    def distributed_job
+      @distributed_job ||= Kraps.distributed_job_client.build(token: @args["token"])
+    end
+  end
+end

data/lib/kraps.rb

@@ -0,0 +1,48 @@
+require_relative "kraps/version"
+require_relative "kraps/drivers"
+require_relative "kraps/actions"
+require_relative "kraps/parallelizer"
+require_relative "kraps/temp_path"
+require_relative "kraps/temp_paths"
+require_relative "kraps/timeout_queue"
+require_relative "kraps/interval"
+require_relative "kraps/job"
+require_relative "kraps/runner"
+require_relative "kraps/step"
+require_relative "kraps/frame"
+require_relative "kraps/worker"
+require "distributed_job"
+require "ruby-progressbar"
+require "ruby-progressbar/outputs/null"
+require "map_reduce"
+require "redis"
+
+module Kraps
+  class Error < StandardError; end
+  class InvalidAction < Error; end
+  class InvalidStep < Error; end
+  class JobStopped < Error; end
+
+  def self.configure(driver:, redis: Redis.new, namespace: nil, job_ttl: 24 * 60 * 60, show_progress: true, enqueuer: ->(worker, json) { worker.perform_async(json) })
+    @driver = driver
+    @distributed_job_client = DistributedJob::Client.new(redis: redis, namespace: namespace, default_ttl: job_ttl)
+    @show_progress = show_progress
+    @enqueuer = enqueuer
+  end
+
+  def self.driver
+    @driver
+  end
+
+  def self.distributed_job_client
+    @distributed_job_client
+  end
+
+  def self.show_progress?
+    @show_progress
+  end
+
+  def self.enqueuer
+    @enqueuer
+  end
+end