modal-rb 0.0.0

data/lib/modal/invocation.rb

@@ -0,0 +1,134 @@
+require_relative 'pickle'
+
+module Modal
+  module Invocation
+    def await_output(timeout = nil); end
+    def retry(retry_count); end
+  end
+
+  class ControlPlaneInvocation
+    include Invocation
+
+    attr_reader :function_call_id
+
+    def initialize(function_call_id, input = nil, function_call_jwt = nil, input_jwt = nil)
+      @function_call_id = function_call_id
+      @input = input
+      @function_call_jwt = function_call_jwt
+      @input_jwt = input_jwt
+    end
+
+    def self.create(function_id, input, invocation_type)
+      function_put_inputs_item = Modal::Client::FunctionPutInputsItem.new(idx: 0, input: input)
+      request = Modal::Client::FunctionMapRequest.new(
+        function_id: function_id,
+        function_call_type: Modal::Client::FunctionCallType::FUNCTION_CALL_TYPE_UNARY,
+        function_call_invocation_type: invocation_type,
+        pipelined_inputs: [function_put_inputs_item]
+      )
+      function_map_response = Modal.client.call(:function_map, request)
+
+      new(
+        function_map_response.function_call_id,
+        input,
+        function_map_response.function_call_jwt,
+        function_map_response.pipelined_inputs[0].input_jwt
+      )
+    end
+
+    def self.from_function_call_id(function_call_id)
+      new(function_call_id)
+    end
+
+    def await_output(timeout = nil)
+      poll_function_output(@function_call_id, timeout)
+    end
+
+    def retry(retry_count)
+      unless @input
+        raise "Cannot retry function invocation - input missing"
+      end
+
+      retry_item = Modal::Client::FunctionRetryInputsItem.new(
+        input_jwt: @input_jwt,
+        input: @input,
+        retry_count: retry_count
+      )
+      request = Modal::Client::FunctionRetryInputsRequest.new(
+        function_call_jwt: @function_call_jwt,
+        inputs: [retry_item]
+      )
+      function_retry_response = Modal.client.call(:function_retry_inputs, request)
+      @input_jwt = function_retry_response.input_jwts[0]
+    end
+
+    private
+
+    def poll_function_output(function_call_id, timeout_ms = nil)
+      start_time = Time.now
+      last_entry_id = ""
+
+      loop do
+        request = Modal::Client::FunctionGetOutputsRequest.new(
+          function_call_id: function_call_id,
+          timeout: 55.0, # seconds
+          last_entry_id: last_entry_id,
+          max_values: 1,
+          clear_on_success: true,
+          requested_at: Time.now.to_f
+        )
+
+        resp = Modal.client.call(:function_get_outputs, request)
+
+        if resp.last_entry_id && !resp.last_entry_id.empty?
+          last_entry_id = resp.last_entry_id
+        end
+
+        if resp.outputs && resp.outputs.any?
+          output_item = resp.outputs.first
+          if output_item.result
+            return process_result(output_item.result, output_item.data_format)
+          end
+        end
+
+        if resp.respond_to?(:num_unfinished_inputs) && resp.num_unfinished_inputs == 0
+          return nil
+        end
+
+        if timeout_ms && (Time.now - start_time) * 1000 > timeout_ms
+          raise FunctionTimeoutError.new("Function call timed out after #{timeout_ms}ms")
+        end
+
+        sleep(1)
+      end
+    rescue GRPC::BadStatus => e
+      if e.code == GRPC::Core::StatusCodes::DEADLINE_EXCEEDED
+        raise FunctionTimeoutError.new("Function call timed out.")
+      else
+        raise e
+      end
+    end
+
+    def process_result(result, data_format)
+      status = result.status.to_s.to_sym
+
+      case status
+      when :GENERIC_STATUS_SUCCESS
+        if result.data && !result.data.empty?
+          return Pickle.load(result.data)
+        elsif result.data_blob_id && !result.data_blob_id.empty?
+          return nil
+        else
+          return nil
+        end
+      when :GENERIC_STATUS_TIMEOUT
+        raise FunctionTimeoutError.new(result.exception || "Function timed out")
+      when :GENERIC_STATUS_INTERNAL_FAILURE
+        raise InternalFailure.new(result.exception || "Internal failure")
+      else
+        error_msg = result.exception || "Unknown error (status: #{result.status})"
+        raise RemoteError.new("Function execution failed: #{error_msg}")
+      end
+    end
+  end
+end

data/lib/modal/pickle.rb

@@ -0,0 +1,50 @@
+require 'pycall'
+require 'json'
+
+module Pickle
+  class PickleError < StandardError; end
+
+  def self.pickle_module
+    @pickle_module ||= PyCall.import_module('pickle')
+  end
+
+  def self.json_module
+    @json_module ||= PyCall.import_module('json')
+  end
+
+  def self.load(data)
+    begin
+      pickle_data = data.respond_to?(:read) ? data.read : data
+      python_obj = pickle_module.loads(pickle_data)
+      json_str = json_module.dumps(python_obj)
+      JSON.parse(json_str.to_s)
+    rescue => e
+      raise PickleError, "Failed to load pickle data: #{e.message}"
+    end
+  end
+
+  def self.dumps(obj)
+    begin
+      json_str = JSON.generate(obj)
+      python_obj = json_module.loads(json_str)
+      pickle_module.dumps(python_obj).to_s
+    rescue => e
+      raise PickleError, "Failed to dump object to pickle: #{e.message}"
+    end
+  end
+
+  def self.dump(obj, file)
+    begin
+      pickled_data = dumps(obj)
+      if file.respond_to?(:write)
+        file.write(pickled_data)
+      else
+        File.open(file, 'wb') do |f|
+          f.write(pickled_data)
+        end
+      end
+    rescue => e
+      raise PickleError, "Failed to dump object to file: #{e.message}"
+    end
+  end
+end

data/lib/modal/sandbox.rb

@@ -0,0 +1,340 @@
+require_relative 'sandbox_filesystem'
+require_relative 'streams'
+require 'ostruct'
+
+module Modal
+  class Sandbox
+    attr_reader :sandbox_id, :stdin, :stdout, :stderr
+    attr_accessor :task_id
+
+    def initialize(sandbox_id)
+      @sandbox_id = sandbox_id
+      @task_id = nil
+
+      @stdin = ModalWriteStream.new(SandboxInputStream.new(sandbox_id))
+      @stdout = ModalReadStream.new(SandboxOutputStream.new(sandbox_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDOUT))
+      @stderr = ModalReadStream.new(SandboxOutputStream.new(sandbox_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDERR))
+    end
+
+    def exec(command, options = {})
+      ensure_task_id
+
+      workdir = options[:workdir]
+      timeout_secs = options[:timeout] ? options[:timeout] / 1000 : 0
+
+      request = Modal::Client::ContainerExecRequest.new(
+        task_id: @task_id,
+        command: command,
+        workdir: workdir,
+        timeout_secs: timeout_secs
+      )
+
+      resp = Modal.client.call(:container_exec, request)
+      ContainerProcess.new(resp.exec_id, "text")
+    end
+
+    def open(path, mode)
+      ensure_task_id
+
+      request = Modal::Client::ContainerFilesystemExecRequest.new(
+        file_open_request: Modal::Client::ContainerFileOpenRequest.new(
+          path: path,
+          mode: mode
+        ),
+        task_id: @task_id
+      )
+      resp = run_filesystem_exec(request)
+      SandboxFile.new(resp.response.file_open_response.file_descriptor, @task_id)
+    end
+
+    def terminate
+      request = Modal::Client::SandboxTerminateRequest.new(sandbox_id: @sandbox_id)
+      Modal.client.call(:sandbox_terminate, request)
+    end
+
+    def wait
+      loop do
+        request = Modal::Client::SandboxWaitRequest.new(
+          sandbox_id: @sandbox_id,
+          timeout: 55 # seconds
+        )
+        resp = Modal.client.call(:sandbox_wait, request)
+        if resp.completed
+          return resp.exit_code || 0
+        end
+        sleep(1) # Poll every second
+      end
+    end
+
+    private
+
+    # Helper to run filesystem exec requests and handle output.
+    def run_filesystem_exec(request)
+      response = Modal.client.call(:container_filesystem_exec, request)
+      if response.respond_to?(:file_descriptor) && response.file_descriptor
+        return OpenStruct.new(
+          response: OpenStruct.new(
+            file_open_response: OpenStruct.new(
+              file_descriptor: response.file_descriptor
+            )
+          )
+        )
+      end
+
+      exec_id = response.exec_id
+      retries = 10
+
+      while retries > 0
+        begin
+          output_request = Modal::Client::ContainerFilesystemExecGetOutputRequest.new(
+            exec_id: exec_id,
+            timeout: 10
+          )
+
+          stream = Modal.client.call(:container_filesystem_exec_get_output, output_request)
+
+          stream.each do |batch|
+            if batch.respond_to?(:error) && batch.error
+              raise SandboxFilesystemError.new(batch.error.error_message)
+            end
+
+            if batch.respond_to?(:eof) && batch.eof
+              return OpenStruct.new(response: response)
+            end
+          end
+
+          retries -= 1
+          sleep(0.1)
+
+        rescue GRPC::BadStatus => e
+          if e.code == GRPC::Core::StatusCodes::DEADLINE_EXCEEDED
+            retries -= 1
+            next
+          else
+            raise e
+          end
+        end
+      end
+
+      raise SandboxFilesystemError.new("Filesystem operation timed out")
+    end
+
+    def ensure_task_id
+      return if @task_id
+
+      request = Modal::Client::SandboxGetTaskIdRequest.new(
+        sandbox_id: @sandbox_id,
+        wait_until_ready: true
+      )
+      resp = Modal.client.call(:sandbox_get_task_id, request)
+
+      if resp.task_id && !resp.task_id.empty?
+        @task_id = resp.task_id
+      else
+        raise "Sandbox #{@sandbox_id} does not have a task ID, it may not be running"
+      end
+    end
+  end
+
+  class ContainerProcess
+    attr_reader :exec_id, :stdin, :stdout, :stderr
+
+    def initialize(exec_id, mode)
+      @exec_id = exec_id
+      @stdin = ModalWriteStream.new(ContainerProcessInputStream.new(exec_id))
+
+      if mode == "text"
+        @stdout = ModalReadStream.new(ContainerProcessOutputStream.new(exec_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDOUT, true))
+        @stderr = ModalReadStream.new(ContainerProcessOutputStream.new(exec_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDERR, true))
+      else
+        @stdout = ModalReadStream.new(ContainerProcessOutputStream.new(exec_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDOUT, false))
+        @stderr = ModalReadStream.new(ContainerProcessOutputStream.new(exec_id, Modal::Client::FileDescriptor::FILE_DESCRIPTOR_STDERR, false))
+      end
+    end
+
+    def wait
+      loop do
+        request = Modal::Client::ContainerExecWaitRequest.new(
+          exec_id: @exec_id,
+          timeout: 55 # seconds
+        )
+        resp = Modal.client.call(:container_exec_wait, request)
+        if resp.completed
+          return resp.exit_code || 0
+        end
+        sleep(1) # Poll every second
+      end
+    end
+  end
+
+  class SandboxInputStream
+    def initialize(sandbox_id)
+      @sandbox_id = sandbox_id
+      @index = 1
+    end
+
+    def write(chunk)
+      request = Modal::Client::SandboxStdinWriteRequest.new(
+        sandbox_id: @sandbox_id,
+        input: chunk.bytes.pack('C*'), # Convert to bytes
+        index: @index
+      )
+      Modal.client.call(:sandbox_stdin_write, request)
+      @index += 1
+    end
+
+    def close
+      request = Modal::Client::SandboxStdinWriteRequest.new(
+        sandbox_id: @sandbox_id,
+        index: @index,
+        eof: true
+      )
+      Modal.client.call(:sandbox_stdin_write, request)
+    end
+  end
+
+  class SandboxOutputStream
+    def initialize(sandbox_id, file_descriptor)
+      @sandbox_id = sandbox_id
+      @file_descriptor = file_descriptor
+      @last_entry_id = ""
+      @data_collected = []
+      @finished = false
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+
+      return if @finished
+
+
+      # make one call and collect all data until EOF
+      request = Modal::Client::SandboxGetLogsRequest.new(
+        sandbox_id: @sandbox_id,
+        file_descriptor: @file_descriptor,
+        timeout: 10, # Give it more time to get all the data
+        last_entry_id: @last_entry_id
+      )
+
+      begin
+        resp = Modal.client.call(:sandbox_get_logs, request)
+
+        # Process the entire streaming response
+        resp.each do |batch|
+
+          # Update last_entry_id
+          if batch.respond_to?(:entry_id) && batch.entry_id && !batch.entry_id.empty?
+            @last_entry_id = batch.entry_id
+          end
+
+          # Collect data from this batch
+          if batch.respond_to?(:items) && batch.items
+            batch.items.each do |item|
+              if item.respond_to?(:data) && item.data && !item.data.empty?
+                @data_collected << item.data
+              end
+            end
+          end
+
+          # Check for EOF
+          if batch.respond_to?(:eof) && batch.eof
+            @finished = true
+            break
+          end
+        end
+
+      rescue GRPC::BadStatus => e
+        if e.code == GRPC::Core::StatusCodes::DEADLINE_EXCEEDED
+          @finished = true
+        else
+          raise e
+        end
+      end
+
+      # Yield all collected data
+      @data_collected.each { |data| yield data }
+
+    end
+  end
+
+  class ContainerProcessInputStream
+    def initialize(exec_id)
+      @exec_id = exec_id
+      @message_index = 1
+    end
+
+    def write(chunk)
+      request = Modal::Client::ContainerExecPutInputRequest.new(
+        exec_id: @exec_id,
+        input: Modal::Client::ContainerExecInput.new(
+          message: chunk.bytes.pack('C*'), # Convert to bytes
+          message_index: @message_index
+        )
+      )
+      Modal.client.call(:container_exec_put_input, request)
+      @message_index += 1
+    end
+
+    def close
+      request = Modal::Client::ContainerExecPutInputRequest.new(
+        exec_id: @exec_id,
+        input: Modal::Client::ContainerExecInput.new(
+          message_index: @message_index,
+          eof: true
+        )
+      )
+      Modal.client.call(:container_exec_put_input, request)
+    end
+  end
+
+  class ContainerProcessOutputStream
+    def initialize(exec_id, file_descriptor, decode_text)
+      @exec_id = exec_id
+      @file_descriptor = file_descriptor
+      @decode_text = decode_text
+      @last_batch_index = 0
+      @finished = false
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+      return if @finished
+
+
+      begin
+        request = Modal::Client::ContainerExecGetOutputRequest.new(
+          exec_id: @exec_id,
+          file_descriptor: @file_descriptor,
+          timeout: 55,
+          get_raw_bytes: true,
+          last_batch_index: @last_batch_index
+        )
+
+        stream = Modal.client.call(:container_exec_get_output, request)
+
+        stream.each do |batch|
+          @last_batch_index = batch.batch_index if batch.respond_to?(:batch_index)
+          if batch.respond_to?(:items) && batch.items
+            batch.items.each do |item|
+              if item.message_bytes && !item.message_bytes.empty?
+                yield item.message_bytes
+              end
+            end
+          end
+
+          if (batch.respond_to?(:has_exit_code) && batch.has_exit_code) || batch.items.empty?
+            break
+          end
+        end
+
+      rescue GRPC::BadStatus => e
+        if e.code == GRPC::Core::StatusCodes::DEADLINE_EXCEEDED
+        else
+          raise e
+        end
+      end
+
+      @finished = true
+    end
+  end
+end