Class: OnnxRuby::SessionPool

Inherits:
Object
  • Object
show all
Defined in:
lib/onnx_ruby/session_pool.rb

Defined Under Namespace

Classes: TimeoutError

Instance Method Summary collapse

Constructor Details

#initialize(model_path, size: nil, timeout: nil, **session_opts) ⇒ SessionPool

Returns a new instance of SessionPool.



7
8
9
10
11
12
13
14
15
16
 
# File 'lib/onnx_ruby/session_pool.rb', line 7

def initialize(model_path, size: nil, timeout: nil, **session_opts)
  @model_path = model_path
  @session_opts = session_opts
  @size = size || OnnxRuby.configuration.pool_size
  @timeout = timeout || OnnxRuby.configuration.pool_timeout
  @pool = []
  @mutex = Mutex.new
  @condition = ConditionVariable.new
  @created = 0
end

Instance Method Details

#available ⇒ Object 



38
39
40
 
# File 'lib/onnx_ruby/session_pool.rb', line 38

def available
  @mutex.synchronize { @pool.size }
end

#run(inputs, **kwargs) ⇒ Object

Run inference using a pooled session



29
30
31
 
# File 'lib/onnx_ruby/session_pool.rb', line 29

def run(inputs, **kwargs)
  with_session { |s| s.run(inputs, **kwargs) }
end

#size ⇒ Object

Current pool stats



34
35
36
 
# File 'lib/onnx_ruby/session_pool.rb', line 34

def size
  @mutex.synchronize { @created }
end

#with_session(&block) ⇒ Object

Check out a session, yield it, then check it back in



19
20
21
22
23
24
25
26
 
# File 'lib/onnx_ruby/session_pool.rb', line 19

def with_session(&block)
  session = checkout
  begin
    yield session
  ensure
    checkin(session)
  end
end