This file is a merged representation of the entire codebase, combined into a single document by Repomix.
The content has been processed where content has been compressed (code blocks are separated by ⋮---- delimiter).
This section contains a summary of this file.
This file contains a packed representation of the entire repository's contents.
It is designed to be easily consumable by AI systems for analysis, code review,
or other automated processes.
The content is organized as follows:
1. This summary section
2. Repository information
3. Directory structure
4. Repository files (if enabled)
5. Multiple file entries, each consisting of:
- File path as an attribute
- Full contents of the file
- This file should be treated as read-only. Any changes should be made to the
original repository files, not this packed version.
- When processing this file, use the file path to distinguish
between different files in the repository.
- Be aware that this file may contain sensitive information. Handle it with
the same level of security as you would the original repository.
- Some files may have been excluded based on .gitignore rules and Repomix's configuration
- Binary files are not included in this packed representation. Please refer to the Repository Structure section for a complete list of file paths, including binary files
- Files matching patterns in .gitignore are excluded
- Files matching default ignore patterns are excluded
- Content has been compressed - code blocks are separated by ⋮---- delimiter
- Files are sorted by Git change count (files with more changes are at the bottom)
mix/
tasks/
snakepit.doctor.ex
snakepit.gen.adapter.ex
snakepit.setup.ex
snakepit.status.ex
snakepit/
adapters/
grpc_python.ex
bridge/
session_store.ex
session.ex
tool_registry.ex
error/
device.ex
parser.ex
python_exceptions.ex
python_translation.ex
shape.ex
grpc/
generated/
snakepit_bridge.pb.ex
bridge_server.ex
client_impl.ex
client.ex
endpoint.ex
hardware/
cpu_detector.ex
cuda_detector.ex
detector.ex
mps_detector.ex
rocm_detector.ex
selector.ex
logger/
redaction.ex
pool/
application_cleanup.ex
pool.ex
process_registry.ex
registry.ex
worker_starter_registry.ex
worker_starter.ex
worker_supervisor.ex
telemetry/
handlers/
logger.ex
metrics.ex
control.ex
correlation.ex
events.ex
gpu_profiler.ex
grpc_stream.ex
naming.ex
open_telemetry.ex
safe_metadata.ex
span.ex
worker/
lifecycle_config.ex
lifecycle_manager.ex
taint_registry.ex
worker_profile/
thread/
capacity_store.ex
process.ex
thread.ex
adapter.ex
application.ex
bootstrap.ex
circuit_breaker.ex
compatibility.ex
config.ex
crash_barrier.ex
defaults.ex
env_doctor.ex
error.ex
executor.ex
grpc_worker.ex
hardware.ex
health_monitor.ex
heartbeat_monitor.ex
logger.ex
package_error.ex
process_killer.ex
python_packages.ex
python_runtime.ex
python_thread_limits.ex
python_version.ex
retry_policy.ex
run_id.ex
runtime_cleanup.ex
supervisor_tree.md
telemetry_metrics.ex
telemetry.ex
worker_profile.ex
zero_copy_ref.ex
zero_copy.ex
snakepit.ex
This section contains the contents of the repository's files.
defmodule Mix.Tasks.Snakepit.Doctor do
@moduledoc """
Diagnose the local Python and gRPC tooling required by Snakepit.
"""
use Mix.Task
@shortdoc "Run environment diagnostics for the Python bridge"
@impl true
def run(_args) do
case Snakepit.EnvDoctor.run() do
{:ok, results} ->
Enum.each(results, &print_result/1)
Mix.shell().info("✅ Snakepit environment is ready")
:ok
{:error, results} ->
Enum.each(results, &print_result/1)
Mix.raise("Snakepit environment checks failed. See messages above for remediation steps.")
end
end
defp print_result(%{status: :ok, name: name, message: message}) do
Mix.shell().info("✅ [#{format_name(name)}] #{message}")
end
defp print_result(%{status: :warning, name: name, message: message}) do
Mix.shell().info("⚠️ [#{format_name(name)}] #{message}")
end
defp print_result(%{status: :error, name: name, message: message}) do
Mix.shell().error("❌ [#{format_name(name)}] #{message}")
end
defp format_name(name) do
name
|> Atom.to_string()
|> String.replace("_", " ")
end
end
defmodule Mix.Tasks.Snakepit.Gen.Adapter do
@moduledoc """
Generate a Python adapter skeleton under priv/python.
"""
use Mix.Task
@shortdoc "Generate a Python adapter skeleton"
@impl true
def run(args) do
name =
case args do
[value | _] -> value
_ -> Mix.raise("Usage: mix snakepit.gen.adapter ")
end
adapter_dir = Path.join(["priv", "python", name])
handler_dir = Path.join([adapter_dir, "handlers"])
adapter_file = Path.join(adapter_dir, "adapter.py")
init_file = Path.join(adapter_dir, "__init__.py")
handler_init_file = Path.join(handler_dir, "__init__.py")
if File.exists?(adapter_dir) do
Mix.raise("Adapter directory already exists: #{adapter_dir}")
end
File.mkdir_p!(handler_dir)
File.write!(init_file, "")
File.write!(handler_init_file, "")
File.write!(adapter_file, adapter_template(name))
Mix.shell().info("✅ Adapter created at #{adapter_dir}")
Mix.shell().info("Configure it with:")
Mix.shell().info(~s( adapter_args: ["--adapter", "#{name}.adapter.#{adapter_class(name)}"]))
end
defp adapter_class(name) do
name
|> Macro.camelize()
end
defp adapter_template(name) do
class_name = adapter_class(name)
"""
from snakepit_bridge.adapters.base import BaseAdapter, tool
class #{class_name}(BaseAdapter):
def __init__(self):
super().__init__()
@tool(description="Example tool - replace with your own")
def example(self, payload: dict) -> dict:
return {"ok": True, "payload": payload}
"""
end
end
defmodule Mix.Tasks.Snakepit.Setup do
@moduledoc """
Bootstrap the Snakepit development environment.
This task mirrors `make bootstrap` and prepares both the Elixir and Python
tooling so tests can run without manual steps.
"""
use Mix.Task
@shortdoc "Provision Mix deps, Python venvs, and gRPC stubs"
@impl true
def run(_args) do
case Snakepit.Bootstrap.run() do
:ok -> :ok
{:error, reason} -> Mix.raise("Snakepit bootstrap failed: #{inspect(reason)}")
end
end
end
defmodule Mix.Tasks.Snakepit.Status do
@moduledoc """
Report the current status of Snakepit pools and worker queues.
"""
use Mix.Task
@shortdoc "Show Snakepit pool status"
@impl true
def run(_args) do
ensure_started!()
pooling_enabled = Application.get_env(:snakepit, :pooling_enabled, false)
if pooling_enabled do
report_pool_status()
else
Mix.shell().info("Snakepit pooling is disabled (set :pooling_enabled to true).")
end
end
defp ensure_started! do
case Application.ensure_all_started(:snakepit) do
{:ok, _apps} -> :ok
{:error, {app, reason}} -> Mix.raise("Failed to start #{app}: #{inspect(reason)}")
end
end
defp report_pool_status do
Mix.shell().info("Snakepit pool status")
Mix.shell().info(String.duplicate("-", 32))
case Snakepit.Config.get_pool_configs() do
{:ok, pools} ->
Enum.each(pools, &print_pool_status/1)
{:error, reason} ->
Mix.raise("Unable to load pool configuration: #{inspect(reason)}")
end
end
defp print_pool_status(%{name: pool_name} = pool_config) do
workers =
case Snakepit.Pool.list_workers(Snakepit.Pool, pool_name) do
{:error, :pool_not_found} -> []
list -> list
end
stats =
case Snakepit.Pool.get_stats(Snakepit.Pool, pool_name) do
{:error, :pool_not_found} -> %{}
map -> map
end
profile = Map.get(pool_config, :worker_profile, :process)
Mix.shell().info("")
Mix.shell().info("Pool: #{pool_name} (#{profile})")
Mix.shell().info(" Workers: #{length(workers)}")
Mix.shell().info(" Queued: #{Map.get(stats, :queued, 0)}")
Mix.shell().info(" Requests: #{Map.get(stats, :requests, 0)}")
Mix.shell().info(" Errors: #{Map.get(stats, :errors, 0)}")
Mix.shell().info(" Queue timeouts: #{Map.get(stats, :queue_timeouts, 0)}")
end
end
defmodule Snakepit.Adapters.GRPCPython do
@moduledoc """
gRPC-based Python adapter for Snakepit.
This adapter replaces the stdin/stdout protocol with gRPC for better performance,
streaming capabilities, and more robust communication.
## Configuration
Application.put_env(:snakepit, :adapter_module, Snakepit.Adapters.GRPCPython)
Application.put_env(:snakepit, :grpc_port, 50051)
Application.put_env(:snakepit, :grpc_host, "localhost")
Worker ports are OS-assigned (ephemeral) and reported back during startup.
## Features
- Native streaming support for progressive results
- HTTP/2 multiplexing for concurrent requests
- Built-in health checks and monitoring
- Better error handling with gRPC status codes
- Binary data support without base64 encoding
## Streaming Examples
# Stream ML inference results
Snakepit.execute_stream("batch_inference", %{
batch_items: ["image1.jpg", "image2.jpg", "image3.jpg"]
}, fn chunk ->
handle_chunk(chunk)
end)
# Stream large dataset processing with progress
Snakepit.execute_stream("process_large_dataset", %{
total_rows: 10000,
chunk_size: 500
}, fn chunk ->
handle_progress(chunk)
end)
"""
@behaviour Snakepit.Adapter
alias Snakepit.Bridge.ToolChunk
alias Snakepit.Defaults
alias Snakepit.GRPC.Client
alias Snakepit.Logger, as: SLog
alias Snakepit.PythonRuntime
@log_category :grpc
@impl true
def executable_path do
PythonRuntime.executable_path()
end
@impl true
def script_path do
# Get the application directory
app_dir = Application.app_dir(:snakepit)
# Check if we should use threaded server based on adapter args
# Check both old pool_config format and new pools format
pool_config = Application.get_env(:snakepit, :pool_config, %{})
pools_config = Application.get_env(:snakepit, :pools, [])
# Get adapter args from either source
adapter_args = Map.get(pool_config, :adapter_args, [])
# Also check if any pool is configured with --max-workers in pools config
has_max_workers =
Enum.any?(adapter_args, fn arg ->
is_binary(arg) and String.contains?(arg, "--max-workers")
end) or
Enum.any?(pools_config, fn pool_cfg ->
pool_adapter_args = Map.get(pool_cfg, :adapter_args, [])
Enum.any?(pool_adapter_args, fn arg ->
is_binary(arg) and String.contains?(arg, "--max-workers")
end)
end)
# Use threaded server if --max-workers is specified (indicates threaded mode)
script_name =
if has_max_workers do
"grpc_server_threaded.py"
else
"grpc_server.py"
end
Path.join([app_dir, "priv", "python", script_name])
end
@impl true
def script_args do
# Check if custom adapter args are provided in pool config
pool_config = Application.get_env(:snakepit, :pool_config, %{})
adapter_args = Map.get(pool_config, :adapter_args, nil)
if adapter_args do
# Use custom adapter args if provided
adapter_args
else
# Default to ShowcaseAdapter - fully functional reference implementation
# For custom adapters, set pool_config.adapter_args or use TemplateAdapter as starting point
["--adapter", "snakepit_bridge.adapters.showcase.ShowcaseAdapter"]
end
end
# gRPC-specific functionality
@doc """
Get the gRPC port for this adapter instance.
ROBUST FIX: Use port 0 to let the OS dynamically assign an available port.
This completely eliminates:
- Port collision races
- TIME_WAIT conflicts
- Manual port range management
- Port leak tracking
Python will bind to an OS-assigned port and report it back via the readiness file
(`SNAKEPIT_READY_FILE`).
"""
def get_port do
# Port 0 = "OS, please assign me any available port"
0
end
@doc """
Check if gRPC dependencies are available at runtime.
"""
def grpc_available? do
Code.ensure_loaded?(GRPC.Channel) and Code.ensure_loaded?(Protobuf)
end
@doc """
Initialize gRPC connection for the worker.
Called by GRPCWorker during initialization.
CRITICAL FIX: This includes retry logic to handle the race condition where
the Python process signals readiness before the OS socket is fully bound
and accepting connections. This is common in polyglot systems where the
external process startup timing is non-deterministic.
"""
def init_grpc_connection(port) do
if grpc_available?() do
# Retry up to 5 times with exponential backoff + jitter
# This handles the startup race condition gracefully
retry_connect(port, 5, 50, 1)
else
{:error, :grpc_not_available}
end
end
# Exponential backoff with jitter to prevent thundering herd during concurrent worker startup.
# base_delay: initial retry delay (50ms)
# backoff: multiplier for exponential growth (doubles each retry)
defp retry_connect(_port, 0, _base_delay, _backoff) do
# All retries exhausted
SLog.error(@log_category, "gRPC connection failed after all retries")
{:error, :connection_failed_after_retries}
end
defp retry_connect(port, retries_left, base_delay, backoff) do
case Client.connect(port) do
{:ok, channel} ->
# Connection successful!
SLog.debug(@log_category, "gRPC connection established to port #{port}")
{:ok, %{channel: channel, port: port}}
{:error, reason} when reason in [:connection_refused, :unavailable, :internal] ->
# Socket not ready yet - retry with exponential backoff + jitter
delay = min(base_delay * backoff, 500)
# Add ±25% jitter to prevent synchronized retries (thundering herd)
jitter = :rand.uniform(div(max(delay, 4), 4))
actual_delay = delay + jitter
SLog.debug(
@log_category,
"gRPC connection to port #{port} #{reason}. " <>
"Retrying in #{actual_delay}ms... (#{retries_left - 1} retries left)"
)
# OTP-idiomatic non-blocking wait
receive do
after
actual_delay -> :ok
end
retry_connect(port, retries_left - 1, base_delay, backoff * 2)
{:error, reason} ->
# For any other error, fail immediately (no retry)
SLog.error(
@log_category,
"gRPC connection to port #{port} failed with unexpected reason: #{inspect(reason)}"
)
{:error, reason}
end
end
@doc """
Execute a command via gRPC.
"""
def grpc_execute(connection, session_id, command, args, timeout \\ nil)
def grpc_execute(connection, session_id, command, args, nil) do
grpc_execute(connection, session_id, command, args, Defaults.grpc_command_timeout())
end
def grpc_execute(connection, session_id, command, args, timeout) do
if grpc_available?() do
Client.execute_tool(
connection.channel,
session_id,
command,
args,
timeout: timeout
)
else
{:error, :grpc_not_available}
end
end
@doc """
Execute a streaming command via gRPC with callback.
"""
def grpc_execute_stream(connection, session_id, command, args, callback_fn, timeout \\ nil)
def grpc_execute_stream(connection, session_id, command, args, callback_fn, nil)
when is_function(callback_fn, 1) do
grpc_execute_stream(
connection,
session_id,
command,
args,
callback_fn,
Defaults.grpc_worker_stream_timeout()
)
end
def grpc_execute_stream(connection, session_id, command, args, callback_fn, timeout)
when is_function(callback_fn, 1) do
if grpc_available?() do
connection.channel
|> Client.execute_streaming_tool(
session_id,
command,
args,
timeout: timeout
)
|> consume_stream(callback_fn)
else
{:error, :grpc_not_available}
end
end
@doc """
Check if this adapter uses gRPC.
Returns true only if gRPC dependencies are actually available.
"""
def uses_grpc?, do: grpc_available?()
@impl true
# 5 minutes for ML inference
def command_timeout("batch_inference", _args), do: Defaults.grpc_batch_inference_timeout()
# 10 minutes for large datasets
def command_timeout("process_large_dataset", _args), do: Defaults.grpc_large_dataset_timeout()
# Default
def command_timeout(_command, _args), do: Defaults.grpc_command_timeout()
defp consume_stream({:ok, stream}, callback_fn) do
Enum.reduce_while(stream, :ok, fn message, _acc ->
process_stream_message(message, callback_fn)
end)
end
defp consume_stream({:error, reason}, _callback_fn), do: {:error, reason}
defp process_stream_message({:ok, %ToolChunk{} = chunk}, callback_fn) do
deliver_chunk(chunk, callback_fn)
end
defp process_stream_message(%ToolChunk{} = chunk, callback_fn) do
deliver_chunk(chunk, callback_fn)
end
defp process_stream_message({:error, reason}, _callback_fn), do: {:halt, {:error, reason}}
defp process_stream_message(other, _callback_fn),
do: {:halt, {:error, {:unexpected_stream_item, other}}}
defp deliver_chunk(%ToolChunk{} = chunk, callback_fn) do
payload = build_payload(chunk)
try do
case callback_fn.(payload) do
:halt -> {:halt, :ok}
{:halt, reason} -> {:halt, {:error, reason}}
_ -> {:cont, :ok}
end
rescue
exception ->
{:halt, {:error, {:callback_exception, exception, __STACKTRACE__}}}
end
end
defp build_payload(%ToolChunk{} = chunk) do
base =
case decode_chunk_data(chunk.data) do
:empty -> %{}
{:ok, %{} = map} -> map
{:ok, value} -> %{"data" => value}
{:error, raw} -> %{"raw_data_base64" => raw}
end
base
|> Map.put("is_final", chunk.is_final)
|> maybe_with_metadata(chunk.metadata)
end
defp decode_chunk_data(data) when is_binary(data) do
if data == "" do
:empty
else
case Jason.decode(data) do
{:ok, decoded} -> {:ok, decoded}
{:error, _} -> {:error, Base.encode64(data)}
end
end
end
defp decode_chunk_data(_), do: :empty
defp maybe_with_metadata(payload, metadata) when metadata in [%{}, nil], do: payload
defp maybe_with_metadata(payload, metadata) do
metadata
|> Map.new()
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
|> Map.new()
|> case do
%{} = cleaned when map_size(cleaned) > 0 -> Map.put(payload, "_metadata", cleaned)
_ -> payload
end
end
end
defmodule Snakepit.Bridge.SessionStore do
@moduledoc """
Centralized session store using ETS for high-performance session management.
This GenServer manages a centralized ETS table for storing session data,
providing CRUD operations, TTL-based expiration, and automatic cleanup.
The store is designed for high concurrency with optimized ETS settings.
"""
use GenServer
alias Snakepit.Bridge.Session
alias Snakepit.Defaults
alias Snakepit.Logger, as: SLog
@log_category :bridge
@default_table_name :snakepit_sessions
## Client API
@doc """
Starts the SessionStore GenServer.
## Options
- `:name` - The name to register the GenServer (default: __MODULE__)
- `:table_name` - The ETS table name (default: :snakepit_sessions)
- `:cleanup_interval` - Cleanup interval in milliseconds (default: 60_000)
- `:default_ttl` - Default TTL for sessions in seconds (default: 3600)
"""
@spec start_link(keyword()) :: GenServer.on_start()
def start_link(opts \\ []) do
name = Keyword.get(opts, :name, __MODULE__)
GenServer.start_link(__MODULE__, opts, name: name)
end
@doc """
Creates a new session with the given ID and options.
## Parameters
- `session_id` - Unique session identifier
- `opts` - Keyword list of options passed to Session.new/2
## Returns
`{:ok, session}` if successful, `{:error, reason}` if failed.
## Examples
{:ok, session} = SessionStore.create_session("session_123")
{:ok, session} = SessionStore.create_session("session_456", ttl: 7200)
"""
@spec create_session(String.t(), keyword()) :: {:ok, Session.t()} | {:error, term()}
def create_session(session_id, opts \\ []) when is_binary(session_id) do
GenServer.call(__MODULE__, {:create_session, session_id, opts})
end
@spec create_session(GenServer.server(), String.t(), keyword()) ::
{:ok, Session.t()} | {:error, term()}
def create_session(server, session_id, opts) when is_binary(session_id) do
GenServer.call(server, {:create_session, session_id, opts})
end
@doc """
Gets a session by ID, automatically updating the last_accessed timestamp.
## Parameters
- `session_id` - The session identifier
## Returns
`{:ok, session}` if found, `{:error, :not_found}` if not found.
"""
@spec get_session(String.t()) :: {:ok, Session.t()} | {:error, :not_found}
def get_session(session_id) when is_binary(session_id) do
get_session(__MODULE__, session_id)
end
@spec get_session(GenServer.server(), String.t()) :: {:ok, Session.t()} | {:error, :not_found}
def get_session(server, session_id) when is_binary(session_id) do
GenServer.call(server, {:get_session, session_id})
end
@doc """
Updates a session using the provided update function.
The update function receives the current session and should return
the updated session. The operation is atomic.
## Parameters
- `session_id` - The session identifier
- `update_fn` - Function that takes a session and returns an updated session
## Returns
`{:ok, updated_session}` if successful, `{:error, reason}` if failed.
## Examples
{:ok, session} = SessionStore.update_session("session_123", fn session ->
Map.put(session, :data, %{key: "value"})
end)
"""
@spec update_session(String.t(), (Session.t() -> Session.t())) ::
{:ok, Session.t()} | {:error, term()}
def update_session(session_id, update_fn)
when is_binary(session_id) and is_function(update_fn, 1) do
update_session(__MODULE__, session_id, update_fn)
end
@spec update_session(GenServer.server(), String.t(), (Session.t() -> Session.t())) ::
{:ok, Session.t()} | {:error, term()}
def update_session(server, session_id, update_fn)
when is_binary(session_id) and is_function(update_fn, 1) do
GenServer.call(server, {:update_session, session_id, update_fn})
end
@doc """
Deletes a session by ID.
## Parameters
- `session_id` - The session identifier
## Returns
`:ok` always (idempotent operation).
"""
@spec delete_session(String.t()) :: :ok
def delete_session(session_id) when is_binary(session_id) do
delete_session(__MODULE__, session_id)
end
@spec delete_session(GenServer.server(), String.t()) :: :ok
def delete_session(server, session_id) when is_binary(session_id) do
GenServer.call(server, {:delete_session, session_id})
end
@doc """
Manually triggers cleanup of expired sessions.
## Returns
The number of sessions that were cleaned up.
"""
@spec cleanup_expired_sessions() :: non_neg_integer()
def cleanup_expired_sessions do
cleanup_expired_sessions(__MODULE__)
end
@spec cleanup_expired_sessions(GenServer.server()) :: non_neg_integer()
def cleanup_expired_sessions(server) do
GenServer.call(server, :cleanup_expired_sessions)
end
@doc """
Gets statistics about the session store.
## Returns
A map containing various statistics about the session store.
"""
@spec get_stats() :: map()
def get_stats do
get_stats(__MODULE__)
end
@spec get_stats(GenServer.server()) :: map()
def get_stats(server) do
GenServer.call(server, :get_stats)
end
@doc """
Lists all active session IDs.
## Returns
A list of all active session IDs.
"""
@spec list_sessions() :: [String.t()]
def list_sessions do
list_sessions(__MODULE__)
end
@spec list_sessions(GenServer.server()) :: [String.t()]
def list_sessions(server) do
GenServer.call(server, :list_sessions)
end
@doc """
Checks if a session exists.
## Parameters
- `session_id` - The session identifier
## Returns
`true` if the session exists, `false` otherwise.
"""
@spec session_exists?(String.t()) :: boolean()
def session_exists?(session_id) when is_binary(session_id) do
session_exists?(__MODULE__, session_id)
end
@spec session_exists?(GenServer.server(), String.t()) :: boolean()
def session_exists?(server, session_id) when is_binary(session_id) do
GenServer.call(server, {:session_exists, session_id})
end
@doc """
Stores worker-session affinity mapping.
"""
@spec store_worker_session(String.t(), String.t()) :: :ok
def store_worker_session(session_id, worker_id) do
GenServer.call(__MODULE__, {:upsert_worker_session, session_id, worker_id})
end
## GenServer Callbacks
@impl true
def init(opts) do
table_name = Keyword.get(opts, :table_name, @default_table_name)
table =
:ets.new(table_name, [
:set,
:protected,
:named_table,
{:read_concurrency, true},
{:write_concurrency, true},
{:decentralized_counters, true}
])
cleanup_interval = Keyword.get(opts, :cleanup_interval, Defaults.session_cleanup_interval())
default_ttl = Keyword.get(opts, :default_ttl, Defaults.session_default_ttl())
quota_config = Application.get_env(:snakepit, :session_store, %{})
max_sessions =
resolve_quota(opts, quota_config, :max_sessions, Defaults.session_max_sessions())
# Strict mode for dev/test - warns loudly on session accumulation
strict_mode = Keyword.get(opts, :strict_mode, Map.get(quota_config, :strict_mode, false))
Process.send_after(self(), :cleanup_expired_sessions, cleanup_interval)
state = %{
table: table,
table_name: table_name,
cleanup_interval: cleanup_interval,
default_ttl: default_ttl,
max_sessions: max_sessions,
strict_mode: strict_mode,
stats: %{
sessions_created: 0,
sessions_deleted: 0,
sessions_expired: 0,
cleanup_runs: 0
}
}
SLog.info(@log_category, "SessionStore started with table #{table}")
{:ok, state}
end
@impl true
def handle_call({:create_session, session_id, opts}, _from, state) do
opts = Keyword.put_new(opts, :ttl, state.default_ttl)
session = Session.new(session_id, opts)
with :ok <- Session.validate(session),
:ok <- check_session_quota(state) do
insert_new_session(session_id, session, state)
else
{:error, reason} ->
{:reply, {:error, reason}, state}
end
end
@impl true
def handle_call({:update_session, session_id, update_fn}, _from, state) do
case :ets.lookup(state.table, session_id) do
[{^session_id, {_last_accessed, _ttl, session}}] ->
try do
updated_session = update_fn.(session)
case Session.validate(updated_session) do
:ok ->
touched_session = Session.touch(updated_session)
ets_record =
{session_id,
{touched_session.last_accessed, touched_session.ttl, touched_session}}
:ets.insert(state.table, ets_record)
{:reply, {:ok, touched_session}, state}
{:error, reason} ->
{:reply, {:error, reason}, state}
end
rescue
error ->
SLog.error(@log_category, "Error updating session #{session_id}: #{inspect(error)}")
{:reply, {:error, {:update_failed, error}}, state}
end
[] ->
{:reply, {:error, :not_found}, state}
end
end
@impl true
def handle_call(:cleanup_expired_sessions, _from, state) do
{expired_count, new_stats} = do_cleanup_expired_sessions(state.table, state.stats)
{:reply, expired_count, %{state | stats: new_stats}}
end
@impl true
def handle_call(:get_stats, _from, state) do
current_sessions = :ets.info(state.table, :size)
memory_usage = :ets.info(state.table, :memory) * :erlang.system_info(:wordsize)
stats =
Map.merge(state.stats, %{
current_sessions: current_sessions,
memory_usage_bytes: memory_usage,
table_info: :ets.info(state.table)
})
{:reply, stats, state}
end
@impl true
def handle_call({:get_session, session_id}, _from, state) do
case :ets.lookup(state.table, session_id) do
[{^session_id, {_last_accessed, _ttl, session}}] ->
touched_session = Session.touch(session)
ets_record =
{session_id, {touched_session.last_accessed, touched_session.ttl, touched_session}}
:ets.insert(state.table, ets_record)
{:reply, {:ok, touched_session}, state}
[] ->
{:reply, {:error, :not_found}, state}
end
end
@impl true
def handle_call({:delete_session, session_id}, _from, state) do
:ets.delete(state.table, session_id)
new_stats = Map.update(state.stats, :sessions_deleted, 1, &(&1 + 1))
{:reply, :ok, %{state | stats: new_stats}}
end
@impl true
def handle_call(:list_sessions, _from, state) do
session_ids = :ets.select(state.table, [{{:"$1", :_}, [], [:"$1"]}])
{:reply, session_ids, state}
end
@impl true
def handle_call({:session_exists, session_id}, _from, state) do
exists =
case :ets.lookup(state.table, session_id) do
[{^session_id, _}] -> true
[] -> false
end
{:reply, exists, state}
end
@impl true
def handle_call({:upsert_worker_session, session_id, worker_id}, _from, state) do
case :ets.lookup(state.table, session_id) do
[{^session_id, {_last_accessed, _ttl, session}}] ->
updated_session =
session
|> Map.put(:last_worker_id, worker_id)
|> Session.touch()
ets_record =
{session_id, {updated_session.last_accessed, updated_session.ttl, updated_session}}
:ets.insert(state.table, ets_record)
{:reply, :ok, state}
[] ->
opts = [ttl: state.default_ttl]
session =
Session.new(session_id, opts)
|> Map.put(:last_worker_id, worker_id)
case Session.validate(session) do
:ok ->
ets_record = {session_id, {session.last_accessed, session.ttl, session}}
:ets.insert(state.table, ets_record)
new_stats = Map.update(state.stats, :sessions_created, 1, &(&1 + 1))
{:reply, :ok, %{state | stats: new_stats}}
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to validate session for worker affinity: #{inspect(reason)}"
)
{:reply, :ok, state}
end
end
end
@impl true
def handle_info(:cleanup_expired_sessions, state) do
{expired_count, new_stats} = do_cleanup_expired_sessions(state.table, state.stats)
# Emit telemetry for session pruning
if expired_count > 0 do
emit_cleanup_telemetry(state, expired_count, :ttl)
end
# Check session accumulation thresholds and warn if needed
maybe_warn_session_accumulation(state)
Process.send_after(self(), :cleanup_expired_sessions, state.cleanup_interval)
{:noreply, %{state | stats: new_stats}}
end
@impl true
def handle_info(msg, state) do
SLog.warning(@log_category, "SessionStore received unexpected message: #{inspect(msg)}")
{:noreply, state}
end
## Private Functions
defp do_cleanup_expired_sessions(table, stats) do
current_time = System.monotonic_time(:second)
match_spec = [
{{:_, {:"$1", :"$2", :_}},
[
{:<, {:+, :"$1", :"$2"}, current_time}
], [true]}
]
expired_count = :ets.select_delete(table, match_spec)
if expired_count > 0 do
SLog.debug(
@log_category,
"Cleaned up #{expired_count} expired sessions using high-performance select_delete"
)
end
new_stats =
stats
|> Map.update(:sessions_expired, expired_count, &(&1 + expired_count))
|> Map.update(:cleanup_runs, 1, &(&1 + 1))
{expired_count, new_stats}
end
defp resolve_quota(opts, config, key, default) do
value = Keyword.get(opts, key, Map.get(config, key, default))
normalize_quota(value, default)
end
defp normalize_quota(:infinity, _default), do: :infinity
defp normalize_quota(value, _default) when is_integer(value) and value > 0, do: value
defp normalize_quota(value, default) when is_binary(value) do
case Integer.parse(value) do
{int, _} when int > 0 -> int
_ -> default
end
end
defp normalize_quota(_value, default), do: default
defp session_quota_reached?(%{max_sessions: :infinity}), do: false
defp session_quota_reached?(state) do
:ets.info(state.table, :size) >= state.max_sessions
end
defp check_session_quota(state) do
if session_quota_reached?(state) do
{:error, :session_quota_exceeded}
else
:ok
end
end
defp insert_new_session(session_id, session, state) do
ets_record = {session_id, {session.last_accessed, session.ttl, session}}
case :ets.insert_new(state.table, ets_record) do
true ->
SLog.debug(@log_category, "Created new session: #{session_id}")
new_stats = Map.update(state.stats, :sessions_created, 1, &(&1 + 1))
{:reply, {:ok, session}, %{state | stats: new_stats}}
false ->
SLog.debug(
@log_category,
"Session #{session_id} already exists - reusing (concurrent init)"
)
[{^session_id, {_last_accessed, _ttl, existing_session}}] =
:ets.lookup(state.table, session_id)
{:reply, {:ok, existing_session}, state}
end
end
# Telemetry emission for session pruning events
defp emit_cleanup_telemetry(state, count, reason) do
current_sessions = :ets.info(state.table, :size)
:telemetry.execute(
[:snakepit, :bridge, :session, :pruned],
%{
count: count,
remaining_sessions: current_sessions,
system_time: System.system_time()
},
%{
reason: reason,
table_name: state.table_name,
max_sessions: state.max_sessions
}
)
end
# Warn on session accumulation in strict mode or when approaching limits
defp maybe_warn_session_accumulation(state) do
current_sessions = :ets.info(state.table, :size)
max_sessions = state.max_sessions
cond do
max_sessions == :infinity ->
:ok
current_sessions >= max_sessions ->
emit_accumulation_warning(state, current_sessions, :quota_exceeded)
if state.strict_mode do
SLog.warning(
@log_category,
"""
[STRICT MODE] Session quota exceeded!
Current: #{current_sessions}, Max: #{max_sessions}
Sessions are being rejected. Check for session leaks.
"""
)
end
current_sessions >= trunc(max_sessions * Defaults.session_warning_threshold()) ->
emit_accumulation_warning(state, current_sessions, :threshold_warning)
if state.strict_mode do
utilization = Float.round(current_sessions / max_sessions * 100, 1)
SLog.warning(
@log_category,
"""
[STRICT MODE] High session count warning!
Current: #{current_sessions}, Max: #{max_sessions} (#{utilization}% utilization)
Consider checking for session leaks or increasing max_sessions.
"""
)
end
true ->
:ok
end
end
defp emit_accumulation_warning(state, current_sessions, reason) do
:telemetry.execute(
[:snakepit, :bridge, :session, :accumulation_warning],
%{
current_sessions: current_sessions,
max_sessions: state.max_sessions,
utilization:
if(state.max_sessions != :infinity,
do: current_sessions / state.max_sessions,
else: 0.0
),
system_time: System.system_time()
},
%{
reason: reason,
table_name: state.table_name,
strict_mode: state.strict_mode
}
)
end
end
defmodule Snakepit.Bridge.Session do
@moduledoc """
Session data structure for centralized session management.
Stores program metadata and session state for worker affinity.
"""
@type t :: %__MODULE__{
id: String.t(),
programs: map(),
metadata: map(),
created_at: integer(),
last_accessed: integer(),
last_worker_id: String.t() | nil,
ttl: integer(),
stats: map()
}
@enforce_keys [:id, :created_at, :ttl]
defstruct [
:id,
:created_at,
:last_accessed,
:last_worker_id,
:ttl,
programs: %{},
metadata: %{},
stats: %{
program_count: 0
}
]
@doc """
Creates a new session with the given ID and options.
"""
@spec new(String.t(), keyword()) :: t()
def new(id, opts \\ []) when is_binary(id) do
now = System.monotonic_time(:second)
# 1 hour default
ttl = Keyword.get(opts, :ttl, 3600)
metadata = Keyword.get(opts, :metadata, %{})
%__MODULE__{
id: id,
created_at: now,
last_accessed: now,
ttl: ttl,
metadata: metadata,
programs: Keyword.get(opts, :programs, %{}),
last_worker_id: Keyword.get(opts, :last_worker_id, nil)
}
end
@doc """
Updates the last_accessed timestamp to the current time.
## Parameters
- `session` - The session to touch
## Returns
Updated session with current last_accessed timestamp.
"""
@spec touch(t()) :: t()
def touch(%__MODULE__{} = session) do
%{session | last_accessed: System.monotonic_time(:second)}
end
@doc """
Checks if a session has expired based on its TTL.
## Parameters
- `session` - The session to check
- `current_time` - Optional current time (defaults to current monotonic time)
## Returns
`true` if the session has expired, `false` otherwise.
"""
@spec expired?(t(), integer() | nil) :: boolean()
def expired?(%__MODULE__{} = session, current_time \\ nil) do
current_time = current_time || System.monotonic_time(:second)
session.last_accessed + session.ttl < current_time
end
@doc """
Validates that a session struct has all required fields and valid data.
## Parameters
- `session` - The session to validate
## Returns
`:ok` if valid, `{:error, reason}` if invalid.
"""
@spec validate(t()) :: :ok | {:error, term()}
def validate(%__MODULE__{} = session) do
with :ok <- validate_basic_fields(session),
:ok <- validate_time_fields(session) do
validate_worker_field(session.last_worker_id)
end
end
def validate(_), do: {:error, :not_a_session}
defp validate_basic_fields(session) do
with :ok <- validate_id(session.id),
:ok <- validate_programs(session.programs) do
validate_metadata_field(session.metadata)
end
end
defp validate_time_fields(session) do
with :ok <- validate_created_at(session.created_at),
:ok <- validate_last_accessed(session.last_accessed),
:ok <- validate_ttl(session.ttl) do
validate_timestamps(session.created_at, session.last_accessed)
end
end
defp validate_worker_field(nil), do: :ok
defp validate_worker_field(id) when is_binary(id), do: :ok
defp validate_worker_field(_), do: {:error, :invalid_last_worker_id}
defp validate_id(id) when is_binary(id) and id != "", do: :ok
defp validate_id(_), do: {:error, :invalid_id}
defp validate_programs(programs) when is_map(programs), do: :ok
defp validate_programs(_), do: {:error, :invalid_programs}
defp validate_metadata_field(metadata) when is_map(metadata), do: :ok
defp validate_metadata_field(_), do: {:error, :invalid_metadata}
defp validate_created_at(created_at) when is_integer(created_at), do: :ok
defp validate_created_at(_), do: {:error, :invalid_created_at}
defp validate_last_accessed(last_accessed) when is_integer(last_accessed), do: :ok
defp validate_last_accessed(_), do: {:error, :invalid_last_accessed}
defp validate_ttl(ttl) when is_integer(ttl) and ttl >= 0, do: :ok
defp validate_ttl(_), do: {:error, :invalid_ttl}
defp validate_timestamps(created_at, last_accessed) when last_accessed >= created_at, do: :ok
defp validate_timestamps(_, _), do: {:error, :invalid_timestamps}
@doc """
Adds or updates a program in the session.
## Parameters
- `session` - The session to update
- `program_id` - The program identifier
- `program_data` - The program data to store
## Returns
Updated session with the program added/updated.
"""
@spec put_program(t(), String.t(), term()) :: t()
def put_program(%__MODULE__{} = session, program_id, program_data)
when is_binary(program_id) do
is_update = Map.has_key?(session.programs, program_id)
programs = Map.put(session.programs, program_id, program_data)
stats =
if is_update do
session.stats
else
%{session.stats | program_count: session.stats.program_count + 1}
end
%{session | programs: programs, stats: stats}
end
@doc """
Gets a program from the session.
## Parameters
- `session` - The session to query
- `program_id` - The program identifier
## Returns
`{:ok, program_data}` if found, `{:error, :not_found}` if not found.
"""
@spec get_program(t(), String.t()) :: {:ok, term()} | {:error, :not_found}
def get_program(%__MODULE__{} = session, program_id) when is_binary(program_id) do
case Map.get(session.programs, program_id) do
nil -> {:error, :not_found}
program_data -> {:ok, program_data}
end
end
@doc """
Removes a program from the session.
## Parameters
- `session` - The session to update
- `program_id` - The program identifier to remove
## Returns
Updated session with the program removed.
"""
@spec delete_program(t(), String.t()) :: t()
def delete_program(%__MODULE__{} = session, program_id) when is_binary(program_id) do
programs = Map.delete(session.programs, program_id)
%{session | programs: programs}
end
@doc """
Updates session metadata.
## Parameters
- `session` - The session to update
- `key` - The metadata key
- `value` - The metadata value
## Returns
Updated session with the metadata updated.
"""
@spec put_metadata(t(), term(), term()) :: t()
def put_metadata(%__MODULE__{} = session, key, value) do
metadata = Map.put(session.metadata, key, value)
%{session | metadata: metadata}
end
@doc """
Gets metadata from the session.
## Parameters
- `session` - The session to query
- `key` - The metadata key
- `default` - Default value if key not found
## Returns
The metadata value or the default.
"""
@spec get_metadata(t(), term(), term()) :: term()
def get_metadata(%__MODULE__{} = session, key, default \\ nil) do
Map.get(session.metadata, key, default)
end
@doc """
Gets session statistics.
"""
@spec get_stats(t()) :: map()
def get_stats(%__MODULE__{} = session) do
Map.merge(session.stats, %{
age: System.monotonic_time(:second) - session.created_at,
time_since_access: System.monotonic_time(:second) - session.last_accessed,
total_items: session.stats.program_count
})
end
end
defmodule Snakepit.Bridge.InternalToolSpec do
@moduledoc """
Internal specification for a tool in the registry.
Separate from the protobuf ToolSpec to avoid conflicts.
"""
defstruct name: nil,
# :local or :remote
type: nil,
# Function reference for local tools
handler: nil,
# Worker ID for remote tools
worker_id: nil,
parameters: [],
description: "",
metadata: %{},
exposed_to_python: false
@type t :: %__MODULE__{
name: String.t(),
type: :local | :remote,
handler: (any() -> any()) | nil,
worker_id: String.t() | nil,
parameters: list(map()),
description: String.t(),
metadata: map(),
exposed_to_python: boolean()
}
end
defmodule Snakepit.Bridge.ToolRegistry do
@moduledoc """
Registry for managing tool metadata and execution.
Maintains a registry of both local (Elixir) and remote (Python) tools,
handles tool discovery, registration, and provides execution dispatch.
"""
use GenServer
alias Snakepit.Logger, as: SLog
alias Snakepit.Bridge.InternalToolSpec
@table_name :snakepit_tool_registry
@max_tool_name_length 64
@tool_name_pattern ~r/^[A-Za-z0-9][A-Za-z0-9_\-\.]*$/
@max_metadata_entries 32
@max_metadata_bytes 4_096
@log_category :bridge
# Client API
@doc """
Starts the ToolRegistry GenServer.
"""
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Registers a local Elixir tool.
"""
def register_elixir_tool(session_id, tool_name, handler, metadata \\ %{}) do
GenServer.call(__MODULE__, {:register_elixir_tool, session_id, tool_name, handler, metadata})
end
@doc """
Registers a remote Python tool.
"""
def register_python_tool(session_id, tool_name, worker_id, metadata \\ %{}) do
GenServer.call(
__MODULE__,
{:register_python_tool, session_id, tool_name, worker_id, metadata}
)
end
@doc """
Registers multiple tools at once (used by Python workers on startup).
"""
def register_tools(session_id, tool_specs) do
GenServer.call(__MODULE__, {:register_tools, session_id, tool_specs})
end
@doc """
Gets a specific tool by name.
"""
def get_tool(session_id, tool_name) do
case :ets.lookup(@table_name, {session_id, tool_name}) do
[{_key, tool_spec}] -> {:ok, tool_spec}
[] -> {:error, "Tool #{tool_name} not found for session #{session_id}"}
end
end
@doc """
Lists all tools available for a session.
"""
def list_tools(session_id) do
pattern = {{session_id, :_}, :_}
tools = :ets.match_object(@table_name, pattern)
Enum.map(tools, fn {{_session_id, _tool_name}, tool_spec} -> tool_spec end)
end
@doc """
Lists only Elixir tools exposed to Python for a session.
"""
def list_exposed_elixir_tools(session_id) do
list_tools(session_id)
|> Enum.filter(fn tool -> tool.type == :local && tool.exposed_to_python end)
end
@doc """
Executes a local Elixir tool.
"""
def execute_local_tool(session_id, tool_name, params) do
with {:ok, tool} <- get_tool(session_id, tool_name),
:local <- tool.type do
try do
result = apply(tool.handler, [params])
{:ok, result}
rescue
e -> {:error, "Tool execution failed: #{inspect(e)}"}
end
else
{:error, _} = error -> error
_ -> {:error, "Tool #{tool_name} is not a local tool"}
end
end
@doc """
Removes all tools for a session (cleanup).
"""
def cleanup_session(session_id) do
GenServer.call(__MODULE__, {:cleanup_session, session_id})
end
# Server Callbacks
@impl true
def init(_opts) do
# Create ETS table for fast lookups
:ets.new(@table_name, [:named_table, :set, :protected, read_concurrency: true])
SLog.info(@log_category, "ToolRegistry started with ETS table: #{@table_name}")
{:ok, %{}}
end
@impl true
def handle_call({:register_elixir_tool, session_id, tool_name, handler, metadata}, _from, state) do
with {:ok, normalized_name} <- validate_tool_name(tool_name),
{:ok, normalized_metadata} <- validate_metadata(metadata),
:ok <- ensure_tool_not_registered(session_id, normalized_name) do
tool_spec = %InternalToolSpec{
name: normalized_name,
type: :local,
handler: handler,
parameters: Map.get(normalized_metadata, :parameters, []),
description: Map.get(normalized_metadata, :description, ""),
metadata: normalized_metadata,
exposed_to_python: Map.get(normalized_metadata, :exposed_to_python, false)
}
case :ets.insert_new(@table_name, {{session_id, normalized_name}, tool_spec}) do
true ->
SLog.debug(
@log_category,
"Registered Elixir tool: #{normalized_name} for session: #{session_id}"
)
{:reply, :ok, state}
false ->
{:reply, {:error, {:duplicate_tool, normalized_name}}, state}
end
else
{:error, reason} ->
{:reply, {:error, reason}, state}
end
end
@impl true
def handle_call(
{:register_python_tool, session_id, tool_name, worker_id, metadata},
_from,
state
) do
with {:ok, normalized_name} <- validate_tool_name(tool_name),
{:ok, normalized_metadata} <- validate_metadata(metadata),
:ok <- ensure_tool_not_registered(session_id, normalized_name) do
tool_spec = %InternalToolSpec{
name: normalized_name,
type: :remote,
worker_id: worker_id,
parameters: Map.get(normalized_metadata, :parameters, []),
description: Map.get(normalized_metadata, :description, ""),
metadata: normalized_metadata
}
case :ets.insert_new(@table_name, {{session_id, normalized_name}, tool_spec}) do
true ->
SLog.debug(
@log_category,
"Registered Python tool: #{normalized_name} for session: #{session_id}"
)
{:reply, :ok, state}
false ->
{:reply, {:error, {:duplicate_tool, normalized_name}}, state}
end
else
{:error, reason} ->
{:reply, {:error, reason}, state}
end
end
@impl true
def handle_call({:register_tools, session_id, tool_specs}, _from, state) do
with {:ok, normalized_specs} <- build_remote_specs(tool_specs),
:ok <- ensure_batch_not_registered(session_id, normalized_specs),
{:ok, names} <- insert_tool_batch(session_id, normalized_specs) do
SLog.info(@log_category, "Registered #{length(names)} tools for session: #{session_id}")
{:reply, {:ok, names}, state}
else
{:error, reason} ->
{:reply, {:error, reason}, state}
end
end
@impl true
def handle_call({:cleanup_session, session_id}, _from, state) do
pattern = {{session_id, :_}, :_}
num_deleted = :ets.match_object(@table_name, pattern) |> length()
:ets.match_delete(@table_name, pattern)
SLog.debug(@log_category, "Cleaned up #{num_deleted} tools for session: #{session_id}")
{:reply, :ok, state}
end
defp validate_tool_name(name) when is_binary(name) do
trimmed = String.trim(name)
cond do
trimmed == "" ->
{:error, {:invalid_tool_name, :empty}}
byte_size(trimmed) > @max_tool_name_length ->
{:error, {:invalid_tool_name, :too_long}}
not Regex.match?(@tool_name_pattern, trimmed) ->
{:error, {:invalid_tool_name, :invalid_format}}
true ->
{:ok, trimmed}
end
end
defp validate_tool_name(_), do: {:error, {:invalid_tool_name, :invalid_type}}
defp validate_metadata(nil), do: {:ok, %{}}
defp validate_metadata(%{} = metadata), do: enforce_metadata_constraints(metadata)
defp validate_metadata(metadata) when is_list(metadata) do
metadata
|> Enum.into(%{})
|> enforce_metadata_constraints()
rescue
ArgumentError ->
{:error, {:invalid_metadata, :duplicate_keys}}
end
defp validate_metadata(_), do: {:error, {:invalid_metadata, :unsupported_type}}
defp enforce_metadata_constraints(metadata) do
entry_count = map_size(metadata)
cond do
entry_count > @max_metadata_entries ->
{:error, {:invalid_metadata, :too_many_entries}}
byte_size(:erlang.term_to_binary(metadata)) > @max_metadata_bytes ->
{:error, {:invalid_metadata, :too_large}}
true ->
{:ok, metadata}
end
end
defp ensure_tool_not_registered(session_id, tool_name) do
case :ets.lookup(@table_name, {session_id, tool_name}) do
[] -> :ok
_ -> {:error, {:duplicate_tool, tool_name}}
end
end
defp build_remote_specs(tool_specs) do
tool_specs
|> Enum.reduce_while({:ok, [], MapSet.new()}, &accumulate_remote_spec/2)
|> finalize_remote_specs()
end
defp accumulate_remote_spec(spec, {:ok, acc, names}) do
with {:ok, tool_spec} <- build_remote_tool_spec(spec),
:ok <- check_duplicate_name(tool_spec.name, names) do
{:cont, {:ok, [tool_spec | acc], MapSet.put(names, tool_spec.name)}}
else
{:error, reason} -> {:halt, {:error, reason}}
end
end
defp check_duplicate_name(name, names) do
if MapSet.member?(names, name) do
{:error, {:duplicate_tool, name}}
else
:ok
end
end
defp finalize_remote_specs({:ok, specs, _names}), do: {:ok, Enum.reverse(specs)}
defp finalize_remote_specs({:error, reason}), do: {:error, reason}
defp build_remote_tool_spec(spec) do
metadata = Map.get(spec, :metadata, %{})
with {:ok, normalized_name} <- validate_tool_name(Map.get(spec, :name)),
{:ok, normalized_metadata} <- validate_metadata(metadata) do
{:ok,
%InternalToolSpec{
name: normalized_name,
type: :remote,
worker_id: Map.get(spec, :worker_id),
parameters: Map.get(spec, :parameters, []),
description: Map.get(spec, :description, ""),
metadata: normalized_metadata
}}
else
{:error, _} = error -> error
end
end
defp ensure_batch_not_registered(session_id, specs) do
Enum.reduce_while(specs, :ok, fn spec, :ok ->
case :ets.lookup(@table_name, {session_id, spec.name}) do
[] -> {:cont, :ok}
_ -> {:halt, {:error, {:duplicate_tool, spec.name}}}
end
end)
end
defp insert_tool_batch(session_id, specs) do
specs
|> Enum.reduce_while({:ok, []}, &insert_tool_spec(session_id, &1, &2))
|> finalize_tool_batch()
end
defp insert_tool_spec(session_id, spec, {:ok, inserted_names}) do
case :ets.insert_new(@table_name, {{session_id, spec.name}, spec}) do
true ->
{:cont, {:ok, [spec.name | inserted_names]}}
false ->
rollback_inserted_tools(session_id, inserted_names)
{:halt, {:error, {:duplicate_tool, spec.name}}}
end
end
defp rollback_inserted_tools(session_id, inserted_names) do
Enum.each(inserted_names, fn name ->
:ets.delete(@table_name, {session_id, name})
end)
end
defp finalize_tool_batch({:ok, names}), do: {:ok, Enum.reverse(names)}
defp finalize_tool_batch({:error, reason}), do: {:error, reason}
end
defmodule Snakepit.Error.DeviceMismatch do
@moduledoc """
Device mismatch error for tensor operations.
Raised when tensors on different devices are used in an operation
that requires them to be on the same device.
"""
defexception [
:expected,
:got,
:operation,
:message
]
@type device :: :cpu | :mps | {:cuda, non_neg_integer()} | {:rocm, non_neg_integer()}
@type t :: %__MODULE__{
expected: device() | nil,
got: device() | nil,
operation: String.t() | nil,
message: String.t()
}
@impl true
def message(%__MODULE__{message: msg}), do: msg
end
defmodule Snakepit.Error.OutOfMemory do
@moduledoc """
Out of memory error for GPU operations.
Contains information about the requested allocation, available memory,
and suggestions for recovery.
"""
defexception [
:device,
:requested_bytes,
:available_bytes,
:operation,
:suggestions,
:message
]
@type device :: :cpu | :mps | {:cuda, non_neg_integer()} | {:rocm, non_neg_integer()}
@type t :: %__MODULE__{
device: device(),
requested_bytes: non_neg_integer(),
available_bytes: non_neg_integer(),
operation: String.t() | nil,
suggestions: [String.t()],
message: String.t()
}
@impl true
def message(%__MODULE__{message: msg}), do: msg
end
defmodule Snakepit.Error.Device do
@moduledoc """
Device error creation helpers.
Provides functions for creating device-related errors with
telemetry emission and helpful suggestions.
"""
alias Snakepit.Error.{DeviceMismatch, OutOfMemory}
@doc """
Creates a device mismatch error.
## Examples
error = Device.device_mismatch(:cpu, {:cuda, 0}, "matmul")
"""
@spec device_mismatch(term(), term(), String.t()) :: DeviceMismatch.t()
def device_mismatch(expected, got, operation) do
message =
"Device mismatch in #{operation}: expected #{format_device(expected)}, got #{format_device(got)}"
error = %DeviceMismatch{
expected: expected,
got: got,
operation: operation,
message: message
}
emit_device_telemetry(error)
error
end
@doc """
Creates a device unavailable error.
## Examples
error = Device.device_unavailable({:cuda, 2}, "matrix_multiply")
"""
@spec device_unavailable(term(), String.t()) :: DeviceMismatch.t()
def device_unavailable(device, operation) do
message = "Device #{format_device(device)} unavailable or not found for #{operation}"
error = %DeviceMismatch{
expected: device,
got: nil,
operation: operation,
message: message
}
emit_device_telemetry(error)
error
end
@doc """
Creates an out of memory error with recovery suggestions.
## Examples
error = Device.out_of_memory({:cuda, 0}, 1024 * 1024 * 1024, 512 * 1024 * 1024)
"""
@spec out_of_memory(term(), non_neg_integer(), non_neg_integer(), String.t() | nil) ::
OutOfMemory.t()
def out_of_memory(device, requested_bytes, available_bytes, operation \\ nil) do
suggestions = generate_oom_suggestions(device, requested_bytes, available_bytes)
message =
"OOM on #{format_device(device)}: requested #{format_bytes(requested_bytes)}, " <>
"available #{format_bytes(available_bytes)}"
message =
if operation do
"#{message} during #{operation}"
else
message
end
error = %OutOfMemory{
device: device,
requested_bytes: requested_bytes,
available_bytes: available_bytes,
operation: operation,
suggestions: suggestions,
message: message
}
emit_oom_telemetry(error)
error
end
defp format_device(:cpu), do: "CPU"
defp format_device(:mps), do: "MPS"
defp format_device({:cuda, id}), do: "CUDA:#{id}"
defp format_device({:rocm, id}), do: "ROCm:#{id}"
defp format_device(nil), do: "unknown"
defp format_device(other), do: inspect(other)
defp format_bytes(bytes) when bytes >= 1_073_741_824 do
"#{Float.round(bytes / 1_073_741_824, 2)} GB"
end
defp format_bytes(bytes) when bytes >= 1_048_576 do
"#{Float.round(bytes / 1_048_576, 2)} MB"
end
defp format_bytes(bytes) when bytes >= 1024 do
"#{Float.round(bytes / 1024, 2)} KB"
end
defp format_bytes(bytes), do: "#{bytes} bytes"
defp generate_oom_suggestions(device, requested, available) do
base_suggestions = [
"Reduce batch size",
"Use gradient checkpointing",
"Enable mixed precision training (FP16/BF16)",
"Free unused tensors with del or gc.collect()"
]
device_suggestions =
case device do
{:cuda, _} ->
[
"Use torch.cuda.empty_cache() to clear cached memory",
"Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
]
:mps ->
["Use torch.mps.empty_cache() to clear cached memory"]
_ ->
[]
end
ratio = requested / max(available, 1)
extreme_suggestions =
if ratio > 2.0 do
[
"Consider model parallelism or offloading to CPU",
"Use a smaller model architecture"
]
else
[]
end
base_suggestions ++ device_suggestions ++ extreme_suggestions
end
defp emit_device_telemetry(%DeviceMismatch{} = error) do
:telemetry.execute(
[:snakepit, :error, :device],
%{},
%{
expected_device: error.expected,
actual_device: error.got,
operation: error.operation
}
)
end
defp emit_oom_telemetry(%OutOfMemory{} = error) do
:telemetry.execute(
[:snakepit, :error, :oom],
%{
requested_bytes: error.requested_bytes,
available_bytes: error.available_bytes
},
%{
device: error.device,
operation: error.operation
}
)
end
end
defmodule Snakepit.Error.Parser do
@moduledoc """
Parses Python exception data into structured Elixir errors.
Automatically detects error patterns (OOM, shape mismatch, device errors)
and creates appropriate structured exceptions.
"""
alias Snakepit.Error.{Device, Shape}
@doc """
Parses error data into a structured error.
Accepts a map with "type", "message", and optionally "traceback".
## Examples
{:ok, error} = Parser.parse(%{
"type" => "ValueError",
"message" => "Invalid input"
})
"""
@spec parse(map() | term()) :: {:ok, Exception.t()} | {:error, :invalid_input}
def parse(nil), do: {:error, :invalid_input}
def parse(data) when not is_map(data), do: {:error, :invalid_input}
def parse(%{"type" => type, "message" => message} = data) do
traceback = Map.get(data, "traceback")
error =
cond do
shape_mismatch?(message) ->
parse_shape_mismatch(message)
oom_error?(message) ->
parse_oom_error(message)
device_mismatch?(message) ->
parse_device_mismatch(message)
true ->
create_typed_error(type, message, traceback, data)
end
{:ok, error}
end
def parse(%{}), do: {:error, :invalid_input}
@doc """
Parses a gRPC error response into a structured error.
"""
@spec from_grpc_error(map()) :: {:ok, Exception.t()} | {:error, :invalid_input}
def from_grpc_error(%{status: _status, message: message}) when is_binary(message) do
# Try to parse as JSON first
case Jason.decode(message) do
{:ok, data} when is_map(data) ->
parse(data)
_ ->
# Parse from message string
parse_from_message(message)
end
end
def from_grpc_error(_), do: {:error, :invalid_input}
@doc """
Extracts a shape from a string representation.
## Examples
iex> Parser.extract_shape("[3, 224, 224]")
[3, 224, 224]
iex> Parser.extract_shape("(10, 20)")
[10, 20]
"""
@spec extract_shape(String.t()) :: [integer()] | nil
def extract_shape(str) do
# Try bracket notation [a, b, c]
case Regex.run(~r/\[([0-9,\s]+)\]/, str) do
[_, dims] ->
parse_dims(dims)
nil ->
# Try tuple notation (a, b, c)
case Regex.run(~r/\(([0-9,\s]+)\)/, str) do
[_, dims] -> parse_dims(dims)
nil -> nil
end
end
end
defp parse_dims(dims) do
dims
|> String.split(",")
|> Enum.map(&String.trim/1)
|> Enum.map(&parse_int/1)
|> Enum.filter(&(&1 != nil))
end
defp parse_int(str) do
case Integer.parse(str) do
{n, _} -> n
:error -> nil
end
end
# Pattern detection helpers
defp shape_mismatch?(message) do
String.contains?(String.downcase(message), "shape") and
(String.contains?(message, "mismatch") or
String.contains?(message, "expected") or
String.contains?(message, "got"))
end
defp oom_error?(message) do
lower = String.downcase(message)
String.contains?(lower, "out of memory") or
String.contains?(lower, "oom") or
(String.contains?(lower, "cuda") and String.contains?(lower, "allocate"))
end
defp device_mismatch?(message) do
lower = String.downcase(message)
String.contains?(lower, "device") and
(String.contains?(lower, "mismatch") or
String.contains?(lower, "expected") or
String.contains?(lower, "same device"))
end
# Pattern parsing helpers
defp parse_shape_mismatch(message) do
# Try to extract expected and got shapes
expected = extract_first_shape(message, "expected")
got = extract_first_shape(message, "got")
if expected && got do
Shape.shape_mismatch(expected, got, extract_operation(message))
else
Shape.shape_mismatch([], [], extract_operation(message))
end
end
defp extract_first_shape(message, prefix) do
pattern = ~r/#{prefix}\s*:?\s*(\[[0-9,\s]+\]|\([0-9,\s]+\))/i
case Regex.run(pattern, message) do
[_, shape_str] -> extract_shape(shape_str)
nil -> nil
end
end
defp parse_oom_error(message) do
# Extract device
device = extract_device(message) || {:cuda, 0}
# Try to extract memory values
{requested, available} = extract_memory_values(message)
Device.out_of_memory(device, requested, available)
end
defp extract_memory_values(message) do
# Try to find "allocate X" pattern
requested =
case Regex.run(~r/allocate\s+([\d.]+)\s*(GB|MB|GiB|MiB)/i, message) do
[_, num, unit] -> parse_memory(num, unit)
nil -> 0
end
# Try to find "available X" or "free X" pattern
available =
case Regex.run(~r/(available|free)\s+([\d.]+)\s*(GB|MB|GiB|MiB)/i, message) do
[_, _, num, unit] -> parse_memory(num, unit)
nil -> 0
end
{requested, available}
end
defp parse_memory(num_str, unit) do
{num, _} = Float.parse(num_str)
case String.upcase(unit) do
u when u in ["GB", "GIB"] -> trunc(num * 1_073_741_824)
u when u in ["MB", "MIB"] -> trunc(num * 1_048_576)
_ -> trunc(num)
end
end
defp parse_device_mismatch(message) do
# Try to extract device names
expected = extract_device(message)
got = extract_second_device(message)
Device.device_mismatch(expected || :cpu, got || :cpu, extract_operation(message))
end
defp extract_device(message) do
cond do
String.contains?(message, "cuda:") ->
case Regex.run(~r/cuda:(\d+)/i, message) do
[_, id] -> {:cuda, String.to_integer(id)}
nil -> {:cuda, 0}
end
String.contains?(String.downcase(message), "cuda") ->
{:cuda, 0}
String.contains?(String.downcase(message), "mps") ->
:mps
String.contains?(String.downcase(message), "cpu") ->
:cpu
true ->
nil
end
end
defp extract_second_device(message) do
# Look for second device after "and"
case Regex.run(~r/and\s+(cuda:\d+|cpu|mps)/i, message) do
[_, device_str] -> parse_device_string(device_str)
nil -> nil
end
end
defp parse_device_string(str) do
lower = String.downcase(str)
cond do
String.starts_with?(lower, "cuda:") ->
[_, id] = Regex.run(~r/cuda:(\d+)/i, str)
{:cuda, String.to_integer(id)}
lower == "cuda" ->
{:cuda, 0}
lower == "cpu" ->
:cpu
lower == "mps" ->
:mps
true ->
nil
end
end
defp extract_operation(message) do
# Try to find operation name in common patterns
case Regex.run(~r/(?:in|during|for)\s+(\w+)/i, message) do
[_, op] -> op
nil -> "unknown"
end
end
defp parse_from_message(message) do
# Try to extract type from "TypeName: message" pattern
case Regex.run(~r/^(\w+Error):\s*(.+)$/i, message) do
[_, type, msg] ->
parse(%{"type" => type, "message" => msg})
nil ->
{:ok, %Snakepit.Error.PythonException{python_type: "Unknown", message: message}}
end
end
defp create_typed_error(type, message, traceback, _data) do
case type do
"ValueError" ->
%Snakepit.Error.ValueError{
python_type: type,
message: message,
python_traceback: traceback
}
"TypeError" ->
%Snakepit.Error.TypeError{
python_type: type,
message: message,
python_traceback: traceback
}
"KeyError" ->
%Snakepit.Error.KeyError{
python_type: type,
message: message,
python_traceback: traceback
}
"IndexError" ->
%Snakepit.Error.IndexError{
python_type: type,
message: message,
python_traceback: traceback
}
"AttributeError" ->
%Snakepit.Error.AttributeError{
python_type: type,
message: message,
python_traceback: traceback
}
"RuntimeError" ->
%Snakepit.Error.RuntimeError{
python_type: type,
message: message,
python_traceback: traceback
}
"ImportError" ->
%Snakepit.Error.ImportError{
python_type: type,
message: message,
python_traceback: traceback
}
_ ->
%Snakepit.Error.PythonException{
python_type: type,
message: message,
python_traceback: traceback
}
end
end
end
# credo:disable-for-next-line Credo.Check.Consistency.ExceptionNames
defmodule Snakepit.Error.PythonException do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.ValueError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.KeyError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.IndexError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.TypeError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.AttributeError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.ImportError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.RuntimeError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.NotImplementedError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.FileNotFoundError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.PermissionError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.ZeroDivisionError do
defexception [:message, :context, :stacktrace, :python_type, :python_traceback]
end
defmodule Snakepit.Error.PythonTranslation do
@moduledoc false
alias Snakepit.Error
@mapping %{
"ValueError" => Error.ValueError,
"KeyError" => Error.KeyError,
"IndexError" => Error.IndexError,
"TypeError" => Error.TypeError,
"AttributeError" => Error.AttributeError,
"ImportError" => Error.ImportError,
"RuntimeError" => Error.RuntimeError,
"NotImplementedError" => Error.NotImplementedError,
"FileNotFoundError" => Error.FileNotFoundError,
"PermissionError" => Error.PermissionError,
"ZeroDivisionError" => Error.ZeroDivisionError
}
def from_error_message(error_message) when is_binary(error_message) do
with {:ok, payload} <- Jason.decode(error_message),
true <- is_map(payload),
{:ok, exception} <- from_payload(payload) do
{:ok, exception}
else
_ -> :error
end
end
def from_error_message(_), do: :error
def from_payload(payload) when is_map(payload) do
type = fetch(payload, ["type", :type])
message = fetch(payload, ["message", :message])
if is_binary(type) and is_binary(message) do
context = fetch(payload, ["context", :context]) || %{}
stacktrace = fetch(payload, ["stacktrace", :stacktrace]) || []
python_traceback = fetch(payload, ["traceback", :traceback]) || join_stacktrace(stacktrace)
module = Map.get(@mapping, type, Error.PythonException)
mapped? = Map.has_key?(@mapping, type)
exception =
struct(module, %{
message: message,
context: context,
stacktrace: stacktrace,
python_type: type,
python_traceback: python_traceback
})
emit_exception_event(mapped?, type, context)
{:ok, exception}
else
:error
end
end
def from_payload(_), do: :error
defp fetch(payload, keys) do
Enum.find_value(keys, fn key -> Map.get(payload, key) end)
end
defp join_stacktrace(stacktrace) when is_list(stacktrace) do
Enum.map_join(stacktrace, "", &to_string/1)
end
defp join_stacktrace(_), do: nil
defp emit_exception_event(mapped?, python_type, context) do
event =
if mapped? do
[:snakepit, :python, :exception, :mapped]
else
[:snakepit, :python, :exception, :unmapped]
end
metadata = %{
python_type: python_type,
library: fetch(context, ["library", :library]),
function: fetch(context, ["function", :function])
}
:telemetry.execute(event, %{}, metadata)
end
end
defmodule Snakepit.Error.ShapeMismatch do
@moduledoc """
Shape mismatch error for tensor operations.
Contains detailed information about the shape mismatch including
which dimension differs and what the expected vs actual values were.
"""
defexception [
:expected,
:got,
:dimension,
:expected_dim,
:got_dim,
:operation,
:message
]
@type t :: %__MODULE__{
expected: [integer()] | nil,
got: [integer()] | nil,
dimension: non_neg_integer() | nil,
expected_dim: integer() | nil,
got_dim: integer() | nil,
operation: String.t() | nil,
message: String.t()
}
@impl true
def message(%__MODULE__{message: msg}) when is_binary(msg), do: msg
def message(%__MODULE__{} = error) do
build_message(error)
end
defp build_message(%{expected: expected, got: got, operation: op, dimension: dim}) do
base =
if dim do
"Shape mismatch in #{op || "operation"}: dimension #{dim} expected #{inspect(expected)}, got #{inspect(got)}"
else
"Shape mismatch in #{op || "operation"}: expected #{inspect(expected)}, got #{inspect(got)}"
end
base
end
end
defmodule Snakepit.Error.Shape do
@moduledoc """
Shape error creation helpers.
Provides functions for creating detailed shape mismatch errors
with automatic dimension detection and telemetry emission.
"""
alias Snakepit.Error.ShapeMismatch
@doc """
Creates a shape mismatch error.
Automatically detects which dimension differs and emits telemetry.
## Examples
error = Shape.shape_mismatch([3, 224, 224], [3, 256, 256], "conv2d")
"""
@spec shape_mismatch([integer()], [integer()], String.t()) :: ShapeMismatch.t()
def shape_mismatch(expected, got, operation) when is_list(expected) and is_list(got) do
{dimension, expected_dim, got_dim} = find_mismatch(expected, got)
message =
if length(expected) != length(got) do
"Shape mismatch in #{operation}: rank mismatch - expected #{length(expected)} dimensions #{inspect(expected)}, got #{length(got)} dimensions #{inspect(got)}"
else
"Shape mismatch in #{operation}: expected #{inspect(expected)}, got #{inspect(got)}"
end
error = %ShapeMismatch{
expected: expected,
got: got,
dimension: dimension,
expected_dim: expected_dim,
got_dim: got_dim,
operation: operation,
message: message
}
emit_telemetry(error)
error
end
@doc """
Creates a dimension-specific mismatch error.
Use when you know exactly which dimension has the mismatch.
"""
@spec dimension_mismatch(non_neg_integer(), integer(), integer(), String.t()) ::
ShapeMismatch.t()
def dimension_mismatch(dimension, expected_dim, got_dim, operation) do
message =
"Shape mismatch in #{operation}: dimension #{dimension} expected #{expected_dim}, got #{got_dim}"
error = %ShapeMismatch{
expected: nil,
got: nil,
dimension: dimension,
expected_dim: expected_dim,
got_dim: got_dim,
operation: operation,
message: message
}
emit_telemetry(error)
error
end
@doc """
Creates a broadcast error.
Use when shapes cannot be broadcast together.
"""
@spec broadcast_error([integer()], [integer()], String.t()) :: ShapeMismatch.t()
def broadcast_error(shape1, shape2, operation) do
message =
"Cannot broadcast shapes #{inspect(shape1)} and #{inspect(shape2)} in #{operation}"
error = %ShapeMismatch{
expected: shape1,
got: shape2,
dimension: nil,
expected_dim: nil,
got_dim: nil,
operation: operation,
message: message
}
emit_telemetry(error)
error
end
@spec find_mismatch([integer()], [integer()]) ::
{non_neg_integer() | nil, integer() | nil, integer() | nil}
defp find_mismatch(expected, got) do
if length(expected) != length(got) do
{nil, nil, nil}
else
expected
|> Enum.zip(got)
|> Enum.with_index()
|> Enum.find(fn {{e, g}, _idx} -> e != g end)
|> case do
nil -> {nil, nil, nil}
{{e, g}, idx} -> {idx, e, g}
end
end
end
defp emit_telemetry(%ShapeMismatch{} = error) do
:telemetry.execute(
[:snakepit, :error, :shape_mismatch],
%{},
%{
expected: error.expected,
got: error.got,
dimension: error.dimension,
operation: error.operation
}
)
end
end
defmodule Snakepit.Bridge.PingRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:message, 1, type: :string)
end
defmodule Snakepit.Bridge.PingResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:message, 1, type: :string)
field(:server_time, 2, type: Google.Protobuf.Timestamp, json_name: "serverTime")
end
defmodule Snakepit.Bridge.InitializeSessionRequest.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.InitializeSessionRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:metadata, 2,
repeated: true,
type: Snakepit.Bridge.InitializeSessionRequest.MetadataEntry,
map: true
)
field(:config, 3, type: Snakepit.Bridge.SessionConfig)
end
defmodule Snakepit.Bridge.SessionConfig do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:enable_caching, 1, type: :bool, json_name: "enableCaching")
field(:cache_ttl_seconds, 2, type: :int32, json_name: "cacheTtlSeconds")
field(:enable_telemetry, 3, type: :bool, json_name: "enableTelemetry")
end
defmodule Snakepit.Bridge.InitializeSessionResponse.AvailableToolsEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: Snakepit.Bridge.ToolSpec)
end
defmodule Snakepit.Bridge.InitializeSessionResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:success, 1, type: :bool)
field(:error_message, 2, type: :string, json_name: "errorMessage")
field(:available_tools, 3,
repeated: true,
type: Snakepit.Bridge.InitializeSessionResponse.AvailableToolsEntry,
json_name: "availableTools",
map: true
)
end
defmodule Snakepit.Bridge.CleanupSessionRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:force, 2, type: :bool)
end
defmodule Snakepit.Bridge.CleanupSessionResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:success, 1, type: :bool)
field(:resources_cleaned, 2, type: :int32, json_name: "resourcesCleaned")
end
defmodule Snakepit.Bridge.ToolSpec.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ToolSpec do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:name, 1, type: :string)
field(:description, 2, type: :string)
field(:parameters, 3, repeated: true, type: Snakepit.Bridge.ParameterSpec)
field(:metadata, 4, repeated: true, type: Snakepit.Bridge.ToolSpec.MetadataEntry, map: true)
field(:supports_streaming, 5, type: :bool, json_name: "supportsStreaming")
end
defmodule Snakepit.Bridge.ParameterSpec do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:name, 1, type: :string)
field(:type, 2, type: :string)
field(:description, 3, type: :string)
field(:required, 4, type: :bool)
field(:default_value, 5, type: Google.Protobuf.Any, json_name: "defaultValue")
field(:validation_json, 6, type: :string, json_name: "validationJson")
end
defmodule Snakepit.Bridge.ExecuteToolRequest.ParametersEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: Google.Protobuf.Any)
end
defmodule Snakepit.Bridge.ExecuteToolRequest.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ExecuteToolRequest.BinaryParametersEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :bytes)
end
defmodule Snakepit.Bridge.ExecuteToolRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:tool_name, 2, type: :string, json_name: "toolName")
field(:parameters, 3,
repeated: true,
type: Snakepit.Bridge.ExecuteToolRequest.ParametersEntry,
map: true
)
field(:metadata, 4,
repeated: true,
type: Snakepit.Bridge.ExecuteToolRequest.MetadataEntry,
map: true
)
field(:stream, 5, type: :bool)
field(:binary_parameters, 6,
repeated: true,
type: Snakepit.Bridge.ExecuteToolRequest.BinaryParametersEntry,
json_name: "binaryParameters",
map: true
)
end
defmodule Snakepit.Bridge.ExecuteToolResponse.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ExecuteToolResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:success, 1, type: :bool)
field(:result, 2, type: Google.Protobuf.Any)
field(:error_message, 3, type: :string, json_name: "errorMessage")
field(:metadata, 4,
repeated: true,
type: Snakepit.Bridge.ExecuteToolResponse.MetadataEntry,
map: true
)
field(:execution_time_ms, 5, type: :int64, json_name: "executionTimeMs")
field(:binary_result, 6, type: :bytes, json_name: "binaryResult")
end
defmodule Snakepit.Bridge.ToolChunk.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ToolChunk do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:chunk_id, 1, type: :string, json_name: "chunkId")
field(:data, 2, type: :bytes)
field(:is_final, 3, type: :bool, json_name: "isFinal")
field(:metadata, 4, repeated: true, type: Snakepit.Bridge.ToolChunk.MetadataEntry, map: true)
end
defmodule Snakepit.Bridge.GetSessionRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
end
defmodule Snakepit.Bridge.GetSessionResponse.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.GetSessionResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:metadata, 2,
repeated: true,
type: Snakepit.Bridge.GetSessionResponse.MetadataEntry,
map: true
)
field(:created_at, 3, type: Google.Protobuf.Timestamp, json_name: "createdAt")
field(:tool_count, 4, type: :int32, json_name: "toolCount")
end
defmodule Snakepit.Bridge.HeartbeatRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:client_time, 2, type: Google.Protobuf.Timestamp, json_name: "clientTime")
end
defmodule Snakepit.Bridge.HeartbeatResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:server_time, 1, type: Google.Protobuf.Timestamp, json_name: "serverTime")
field(:session_valid, 2, type: :bool, json_name: "sessionValid")
end
defmodule Snakepit.Bridge.RegisterToolsRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:tools, 2, repeated: true, type: Snakepit.Bridge.ToolRegistration)
field(:worker_id, 3, type: :string, json_name: "workerId")
end
defmodule Snakepit.Bridge.ToolRegistration.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ToolRegistration do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:name, 1, type: :string)
field(:description, 2, type: :string)
field(:parameters, 3, repeated: true, type: Snakepit.Bridge.ParameterSpec)
field(:metadata, 4,
repeated: true,
type: Snakepit.Bridge.ToolRegistration.MetadataEntry,
map: true
)
field(:supports_streaming, 5, type: :bool, json_name: "supportsStreaming")
end
defmodule Snakepit.Bridge.RegisterToolsResponse.ToolIdsEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.RegisterToolsResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:success, 1, type: :bool)
field(:tool_ids, 2,
repeated: true,
type: Snakepit.Bridge.RegisterToolsResponse.ToolIdsEntry,
json_name: "toolIds",
map: true
)
field(:error_message, 3, type: :string, json_name: "errorMessage")
end
defmodule Snakepit.Bridge.GetExposedElixirToolsRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
end
defmodule Snakepit.Bridge.GetExposedElixirToolsResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:tools, 1, repeated: true, type: Snakepit.Bridge.ToolSpec)
end
defmodule Snakepit.Bridge.ExecuteElixirToolRequest.ParametersEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: Google.Protobuf.Any)
end
defmodule Snakepit.Bridge.ExecuteElixirToolRequest.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ExecuteElixirToolRequest do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:session_id, 1, type: :string, json_name: "sessionId")
field(:tool_name, 2, type: :string, json_name: "toolName")
field(:parameters, 3,
repeated: true,
type: Snakepit.Bridge.ExecuteElixirToolRequest.ParametersEntry,
map: true
)
field(:metadata, 4,
repeated: true,
type: Snakepit.Bridge.ExecuteElixirToolRequest.MetadataEntry,
map: true
)
end
defmodule Snakepit.Bridge.ExecuteElixirToolResponse.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.ExecuteElixirToolResponse do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:success, 1, type: :bool)
field(:result, 2, type: Google.Protobuf.Any)
field(:error_message, 3, type: :string, json_name: "errorMessage")
field(:metadata, 4,
repeated: true,
type: Snakepit.Bridge.ExecuteElixirToolResponse.MetadataEntry,
map: true
)
field(:execution_time_ms, 5, type: :int64, json_name: "executionTimeMs")
field(:binary_result, 6, type: :bytes, json_name: "binaryResult")
end
defmodule Snakepit.Bridge.TelemetryEvent.MeasurementsEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: Snakepit.Bridge.TelemetryValue)
end
defmodule Snakepit.Bridge.TelemetryEvent.MetadataEntry do
@moduledoc false
use Protobuf, map: true, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:key, 1, type: :string)
field(:value, 2, type: :string)
end
defmodule Snakepit.Bridge.TelemetryEvent do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:event_parts, 1, repeated: true, type: :string, json_name: "eventParts")
field(:measurements, 2,
repeated: true,
type: Snakepit.Bridge.TelemetryEvent.MeasurementsEntry,
map: true
)
field(:metadata, 3,
repeated: true,
type: Snakepit.Bridge.TelemetryEvent.MetadataEntry,
map: true
)
field(:timestamp_ns, 4, type: :int64, json_name: "timestampNs")
field(:correlation_id, 5, type: :string, json_name: "correlationId")
end
defmodule Snakepit.Bridge.TelemetryValue do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
oneof(:value, 0)
field(:int_value, 1, type: :int64, json_name: "intValue", oneof: 0)
field(:float_value, 2, type: :double, json_name: "floatValue", oneof: 0)
field(:string_value, 3, type: :string, json_name: "stringValue", oneof: 0)
end
defmodule Snakepit.Bridge.TelemetryControl do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
oneof(:control, 0)
field(:toggle, 1, type: Snakepit.Bridge.TelemetryToggle, oneof: 0)
field(:sampling, 2, type: Snakepit.Bridge.TelemetrySamplingUpdate, oneof: 0)
field(:filter, 3, type: Snakepit.Bridge.TelemetryEventFilter, oneof: 0)
end
defmodule Snakepit.Bridge.TelemetryToggle do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:enabled, 1, type: :bool)
end
defmodule Snakepit.Bridge.TelemetrySamplingUpdate do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:sampling_rate, 1, type: :double, json_name: "samplingRate")
field(:event_patterns, 2, repeated: true, type: :string, json_name: "eventPatterns")
end
defmodule Snakepit.Bridge.TelemetryEventFilter do
@moduledoc false
use Protobuf, protoc_gen_elixir_version: "0.15.0", syntax: :proto3
field(:allow, 1, repeated: true, type: :string)
field(:deny, 2, repeated: true, type: :string)
end
defmodule Snakepit.Bridge.BridgeService.Service do
@moduledoc false
use GRPC.Service, name: "snakepit.bridge.BridgeService", protoc_gen_elixir_version: "0.15.0"
rpc(:Ping, Snakepit.Bridge.PingRequest, Snakepit.Bridge.PingResponse)
rpc(
:InitializeSession,
Snakepit.Bridge.InitializeSessionRequest,
Snakepit.Bridge.InitializeSessionResponse
)
rpc(
:CleanupSession,
Snakepit.Bridge.CleanupSessionRequest,
Snakepit.Bridge.CleanupSessionResponse
)
rpc(:GetSession, Snakepit.Bridge.GetSessionRequest, Snakepit.Bridge.GetSessionResponse)
rpc(:Heartbeat, Snakepit.Bridge.HeartbeatRequest, Snakepit.Bridge.HeartbeatResponse)
rpc(:ExecuteTool, Snakepit.Bridge.ExecuteToolRequest, Snakepit.Bridge.ExecuteToolResponse)
rpc(
:ExecuteStreamingTool,
Snakepit.Bridge.ExecuteToolRequest,
stream(Snakepit.Bridge.ToolChunk)
)
rpc(:RegisterTools, Snakepit.Bridge.RegisterToolsRequest, Snakepit.Bridge.RegisterToolsResponse)
rpc(
:GetExposedElixirTools,
Snakepit.Bridge.GetExposedElixirToolsRequest,
Snakepit.Bridge.GetExposedElixirToolsResponse
)
rpc(
:ExecuteElixirTool,
Snakepit.Bridge.ExecuteElixirToolRequest,
Snakepit.Bridge.ExecuteElixirToolResponse
)
rpc(
:StreamTelemetry,
stream(Snakepit.Bridge.TelemetryControl),
stream(Snakepit.Bridge.TelemetryEvent)
)
end
defmodule Snakepit.Bridge.BridgeService.Stub do
@moduledoc false
use GRPC.Stub, service: Snakepit.Bridge.BridgeService.Service
end
defmodule Snakepit.GRPC.BridgeServer do
@moduledoc """
gRPC server implementation for the Snakepit Bridge service.
Handles tool execution and session management through the unified bridge protocol.
"""
use GRPC.Server, service: Snakepit.Bridge.BridgeService.Service
alias Google.Protobuf.{Any, Timestamp}
alias Snakepit.Bridge.{
CleanupSessionRequest,
CleanupSessionResponse,
ExecuteElixirToolRequest,
ExecuteElixirToolResponse,
ExecuteToolRequest,
ExecuteToolResponse,
GetExposedElixirToolsRequest,
GetExposedElixirToolsResponse,
GetSessionRequest,
GetSessionResponse,
HeartbeatRequest,
HeartbeatResponse,
InitializeSessionResponse,
ParameterSpec,
PingRequest,
PingResponse,
RegisterToolsRequest,
RegisterToolsResponse,
ToolChunk,
ToolSpec
}
alias Snakepit.Bridge.SessionStore
alias Snakepit.Bridge.ToolRegistry
alias Snakepit.GRPC.Client, as: GRPCClient
alias Snakepit.GRPCWorker
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Telemetry.Correlation
@log_category :grpc
# Health & Session Management
def ping(%PingRequest{message: message}, _stream) do
SLog.debug(@log_category, "Ping received", message: message)
%PingResponse{
message: "pong: #{message}",
server_time: %Timestamp{seconds: System.system_time(:second), nanos: 0}
}
end
def initialize_session(request, _stream) do
SLog.info(@log_category, "Initializing session", session_id: request.session_id)
case SessionStore.create_session(request.session_id, metadata: request.metadata) do
{:ok, _session} ->
# Success - session was created or already existed (both are fine)
%InitializeSessionResponse{
success: true,
error_message: nil,
available_tools: %{}
}
{:error, reason} ->
# Only raise on actual errors (not :already_exists, which is now handled)
raise GRPC.RPCError,
status: :internal,
message: format_error(reason)
end
end
def cleanup_session(%CleanupSessionRequest{session_id: session_id, force: _force}, _stream) do
SLog.info(@log_category, "Cleaning up session", session_id: session_id)
# NOTE: The force flag is not currently used. SessionStore.delete_session is always idempotent
# and immediately deletes the session regardless of state. Future enhancements could add
# soft-delete or cleanup verification if needed.
SessionStore.delete_session(session_id)
%CleanupSessionResponse{
success: true,
resources_cleaned: 1
}
end
def get_session(%GetSessionRequest{session_id: session_id}, _stream) do
SLog.debug(@log_category, "GetSession", session_id: session_id)
case SessionStore.get_session(session_id) do
{:ok, session} ->
tools = Map.get(session, :tools, %{})
metadata = Map.get(session, :metadata, %{})
tool_count = map_size(tools)
%GetSessionResponse{
session_id: session_id,
metadata: metadata,
created_at: %Timestamp{seconds: session.created_at, nanos: 0},
tool_count: tool_count
}
{:error, :not_found} ->
raise GRPC.RPCError,
status: :not_found,
message: "Session not found: #{session_id}"
end
end
def heartbeat(%HeartbeatRequest{session_id: session_id, client_time: _client_time}, _stream) do
SLog.debug(@log_category, "Heartbeat", session_id: session_id)
# Check if session exists and update last_accessed
session_valid =
case SessionStore.get_session(session_id) do
{:ok, _session} ->
# Getting the session automatically updates last_accessed
true
{:error, :not_found} ->
false
end
%HeartbeatResponse{
server_time: %Timestamp{seconds: System.system_time(:second), nanos: 0},
session_valid: session_valid
}
end
# Tool Execution
def execute_tool(%ExecuteToolRequest{} = request, stream) do
SLog.info(@log_category, "ExecuteTool",
tool_name: request.tool_name,
session_id: request.session_id
)
start_time = System.monotonic_time(:millisecond)
correlation_id = resolve_request_correlation_id(request, stream)
request = ensure_request_correlation(request, correlation_id)
with {:ok, _session} <- SessionStore.get_session(request.session_id),
{:ok, tool} <- ToolRegistry.get_tool(request.session_id, request.tool_name),
{:ok, result} <-
execute_tool_handler(tool, request, request.session_id, correlation_id) do
execution_time = System.monotonic_time(:millisecond) - start_time
{encoded_result, binary_result} = encode_tool_result(result)
%ExecuteToolResponse{
success: true,
result: encoded_result,
binary_result: binary_result || <<>>,
error_message: nil,
metadata: %{
"execution_time" => to_string(execution_time),
"tool_type" => to_string(tool.type)
},
execution_time_ms: execution_time
}
else
{:error, reason} ->
%ExecuteToolResponse{
success: false,
result: nil,
error_message: format_error(reason),
metadata: %{},
execution_time_ms: System.monotonic_time(:millisecond) - start_time
}
end
end
defp execute_tool_handler(%{type: :local} = tool, request, session_id, _correlation_id) do
# Execute local Elixir tool
case decode_tool_parameters(request.parameters, request.binary_parameters) do
{:ok, params} ->
ToolRegistry.execute_local_tool(session_id, tool.name, params)
{:error, reason} ->
{:error, reason}
end
end
defp execute_tool_handler(%{type: :remote} = tool, request, session_id, correlation_id) do
# Forward to Python worker
SLog.debug(@log_category, "Executing remote tool",
tool_name: tool.name,
worker_id: tool.worker_id
)
with {:ok, params} <-
decode_remote_tool_parameters(request.parameters, request.binary_parameters),
{:ok, channel, cleanup} <- ensure_worker_channel(tool.worker_id) do
result =
try do
forward_tool_to_worker(channel, request, session_id, params, correlation_id)
after
cleanup.()
end
case result do
{:ok, response} ->
{:ok, response}
{:error, reason} ->
SLog.error(@log_category, "Failed to execute remote tool",
tool_name: tool.name,
worker_id: tool.worker_id,
reason: reason
)
{:error, {:remote_execution_failed, reason}}
end
else
{:error, {:invalid_parameter, _, _}} = error ->
error
{:error, reason} ->
SLog.error(@log_category, "Failed to execute remote tool",
tool_name: tool.name,
worker_id: tool.worker_id,
reason: reason
)
{:error, {:remote_execution_failed, reason}}
end
end
defp decode_tool_parameters(params, binary_params)
when is_map(params) and (is_map(binary_params) or is_nil(binary_params)) do
decoded =
Enum.reduce_while(params, {:ok, %{}}, fn {key, any_value}, {:ok, acc} ->
case decode_any_value(any_value) do
{:ok, decoded} ->
{:cont, {:ok, Map.put(acc, key, decoded)}}
{:error, reason} ->
{:halt, {:error, {:invalid_parameter, key, reason}}}
end
end)
case decoded do
{:ok, acc} ->
merge_binary_parameters(acc, binary_params || %{})
other ->
other
end
end
defp decode_tool_parameters(_, _), do: {:ok, %{}}
# For remote tools, we decode JSON parameters without merging binary params
# Binary params are forwarded separately via the protobuf message
defp decode_remote_tool_parameters(params, binary_params) do
with {:ok, decoded} <- decode_tool_parameters(params, %{}),
:ok <- validate_binary_parameters(binary_params || %{}) do
{:ok, decoded}
end
end
defp validate_binary_parameters(binary_params) when is_map(binary_params) do
Enum.reduce_while(binary_params, :ok, fn {key, value}, :ok ->
if is_binary(value) do
{:cont, :ok}
else
{:halt, {:error, {:invalid_parameter, normalize_param_key(key), :not_binary}}}
end
end)
end
defp validate_binary_parameters(_), do: :ok
defp decode_any_value(%Any{
type_url: "type.googleapis.com/google.protobuf.StringValue",
value: value
}) do
decode_json(value)
end
defp decode_any_value(%Any{type_url: type_url, value: value}) do
case decode_json(value) do
{:ok, decoded} ->
{:ok, decoded}
{:error, message} ->
{:error, {:invalid_json, message, type_url}}
end
end
defp decode_any_value(value) when is_map(value) or is_list(value) do
{:ok, value}
end
defp merge_binary_parameters(decoded, binary_params) when map_size(binary_params) == 0,
do: {:ok, decoded}
defp merge_binary_parameters(decoded, binary_params) when is_map(binary_params) do
Enum.reduce_while(binary_params, {:ok, decoded}, fn {key, value}, {:ok, acc} ->
if is_binary(value) do
{:cont, {:ok, Map.put(acc, normalize_param_key(key), {:binary, value})}}
else
{:halt, {:error, {:invalid_parameter, normalize_param_key(key), :not_binary}}}
end
end)
end
defp merge_binary_parameters(decoded, _binary_params), do: {:ok, decoded}
defp normalize_param_key(key) when is_atom(key), do: Atom.to_string(key)
defp normalize_param_key(key) when is_binary(key), do: key
defp normalize_param_key(key), do: to_string(key)
defp decode_json(value) when is_binary(value) do
case Jason.decode(value) do
{:ok, decoded} ->
{:ok, decoded}
{:error, %Jason.DecodeError{} = decode_error} ->
{:error, Exception.message(decode_error)}
end
end
defp decode_json(_), do: {:error, "expected JSON encoded string"}
# Helper functions for remote tool execution
defp get_worker_port(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} when is_pid(pid) ->
# Try to get port from worker state - this is a simplified approach
try do
case GenServer.call(pid, :get_port, 1000) do
{:ok, port} -> {:ok, port}
_ -> {:error, "Could not get port from worker"}
end
catch
_exit, _reason -> {:error, "Worker not responding"}
end
{:error, _reason} ->
{:error, "Worker not found: #{worker_id}"}
end
end
defp create_worker_channel(port) do
GRPC.Stub.connect("localhost:#{port}")
rescue
error -> {:error, "Failed to connect to worker: #{inspect(error)}"}
end
defp forward_tool_to_worker(channel, request, session_id, decoded_params, correlation_id) do
worker_metadata = ensure_metadata_correlation(request.metadata, correlation_id)
# Create the request to forward to the worker
worker_request = %ExecuteToolRequest{
session_id: session_id,
tool_name: request.tool_name,
parameters: request.parameters,
metadata: worker_metadata,
binary_parameters: request.binary_parameters
}
binary_params = worker_request.binary_parameters || %{}
opts =
worker_request.metadata
|> tool_call_options()
|> Keyword.put(:binary_parameters, binary_params)
|> Keyword.put(:correlation_id, correlation_id)
case GRPCClient.execute_tool(
channel,
worker_request.session_id,
worker_request.tool_name,
decoded_params,
opts
) do
{:ok, response} ->
{:ok, response}
{:error, reason} ->
{:error, reason}
end
end
defp tool_call_options(metadata) when is_map(metadata) do
value = Map.get(metadata, "timeout_ms") || Map.get(metadata, :timeout_ms)
case parse_timeout_ms(value) do
{:ok, timeout} -> [timeout: timeout]
:error -> []
end
end
defp resolve_request_correlation_id(request, stream) do
request
|> correlation_id_from_metadata()
|> case do
nil -> correlation_id_from_headers(stream)
value -> value
end
|> Correlation.ensure()
end
defp correlation_id_from_metadata(%{metadata: metadata}) when is_map(metadata) do
Map.get(metadata, "correlation_id") || Map.get(metadata, :correlation_id)
end
defp correlation_id_from_metadata(_), do: nil
defp correlation_id_from_headers(nil), do: nil
defp correlation_id_from_headers(stream) do
stream
|> GRPC.Stream.get_headers()
|> Map.get("x-snakepit-correlation-id")
end
defp ensure_request_correlation(%ExecuteToolRequest{} = request, correlation_id) do
metadata = ensure_metadata_correlation(request.metadata, correlation_id)
%{request | metadata: metadata}
end
defp ensure_metadata_correlation(nil, correlation_id) do
%{"correlation_id" => correlation_id}
end
defp ensure_metadata_correlation(metadata, correlation_id) when is_map(metadata) do
Map.put(metadata, "correlation_id", correlation_id)
end
defp parse_timeout_ms(nil), do: :error
defp parse_timeout_ms(value) when is_integer(value) and value > 0, do: {:ok, value}
defp parse_timeout_ms(value) when is_binary(value) do
case Integer.parse(value) do
{int, ""} when int > 0 -> {:ok, int}
_ -> :error
end
end
defp parse_timeout_ms(_), do: :error
defp ensure_worker_channel(worker_id) do
case get_existing_worker_channel(worker_id) do
{:ok, channel} ->
{:ok, channel, fn -> :ok end}
{:error, reason} when is_binary(reason) ->
{:error, reason}
{:error, _reason} ->
with {:ok, port} <- get_worker_port(worker_id),
{:ok, channel} <- create_worker_channel(port) do
{:ok, channel, fn -> disconnect_channel(channel) end}
end
end
end
defp get_existing_worker_channel(worker_id) do
with {:ok, pid, metadata} <- fetch_worker_safely(worker_id),
:ok <- validate_worker_module(metadata) do
GRPCWorker.get_channel(pid)
end
rescue
_ -> {:error, :channel_unavailable}
end
defp fetch_worker_safely(worker_id) do
case PoolRegistry.fetch_worker(worker_id) do
{:ok, pid, metadata} when is_pid(pid) -> {:ok, pid, metadata}
{:error, _} -> {:error, "Worker not found: #{worker_id}"}
end
end
defp validate_worker_module(metadata) do
case Map.get(metadata, :worker_module, GRPCWorker) do
module when module == GRPCWorker -> :ok
_ -> {:error, :unsupported_worker_module}
end
end
defp disconnect_channel(channel) do
_ = GRPC.Stub.disconnect(channel)
:ok
rescue
_ -> :ok
end
def execute_streaming_tool(%ExecuteToolRequest{} = request, stream) do
SLog.info(@log_category, "ExecuteStreamingTool",
tool_name: request.tool_name,
session_id: request.session_id
)
start_time_ms = System.monotonic_time(:millisecond)
correlation_id = resolve_request_correlation_id(request, stream)
request = ensure_request_correlation(request, correlation_id)
with {:ok, _session} <- SessionStore.get_session(request.session_id),
{:ok, tool} <- ToolRegistry.get_tool(request.session_id, request.tool_name),
:ok <- ensure_streaming_supported(tool, request) do
case execute_remote_stream(tool, request, stream, correlation_id, start_time_ms) do
:ok ->
stream
{:error, reason} ->
raise_streaming_rpc_error(request, reason)
end
else
{:error, :not_found} ->
raise GRPC.RPCError,
status: :not_found,
message: "Session not found: #{request.session_id}"
{:error, message} when is_binary(message) ->
raise GRPC.RPCError,
status: :not_found,
message: message
{:error, {:streaming_not_supported, _} = reason} ->
raise_streaming_rpc_error(request, reason)
end
end
# Streaming Support Validation
defp ensure_streaming_supported(%{type: :remote} = tool, _request) do
if tool_supports_streaming?(tool) do
:ok
else
{:error, {:streaming_not_supported, tool.name}}
end
end
defp ensure_streaming_supported(_tool, _request) do
{:error, {:streaming_not_supported, :local_tool}}
end
defp tool_supports_streaming?(%{metadata: metadata}) when is_map(metadata) do
value = Map.get(metadata, "supports_streaming") || Map.get(metadata, :supports_streaming)
case value do
true -> true
"true" -> true
"1" -> true
1 -> true
_ -> false
end
end
defp tool_supports_streaming?(_), do: false
# Remote Stream Execution
defp execute_remote_stream(
%{type: :remote} = tool,
%ExecuteToolRequest{} = request,
stream,
correlation_id,
start_time_ms
) do
with {:ok, decoded_params} <-
decode_remote_tool_parameters(request.parameters, request.binary_parameters),
{:ok, channel, cleanup} <- ensure_worker_channel(tool.worker_id),
{:ok, worker_stream} <-
forward_streaming_tool_to_worker(channel, request, decoded_params, correlation_id) do
try do
forward_worker_stream(worker_stream, stream, start_time_ms, tool)
after
cleanup.()
end
end
end
defp forward_streaming_tool_to_worker(
channel,
%ExecuteToolRequest{} = request,
decoded_params,
correlation_id
) do
worker_metadata = ensure_metadata_correlation(request.metadata, correlation_id)
binary_params = request.binary_parameters || %{}
opts =
worker_metadata
|> tool_call_options()
|> Keyword.put(:binary_parameters, binary_params)
|> Keyword.put(:correlation_id, correlation_id)
channel
|> GRPCClient.execute_streaming_tool(
request.session_id,
request.tool_name,
decoded_params,
opts
)
|> normalize_stream_response()
end
defp normalize_stream_response({:ok, stream, _headers}), do: {:ok, stream}
defp normalize_stream_response({:ok, stream}), do: {:ok, stream}
defp normalize_stream_response({:error, reason}), do: {:error, reason}
defp normalize_stream_response(other), do: {:error, {:unexpected_stream_response, other}}
# Stream Forwarding
defp forward_worker_stream(worker_stream, grpc_stream, start_time_ms, tool) do
acc0 = %{sent: 0, final_seen?: false}
result =
Enum.reduce_while(worker_stream, acc0, fn item, acc ->
case normalize_stream_item(item) do
{:ok, %ToolChunk{} = chunk} ->
final_seen? = acc.final_seen? or chunk.is_final
chunk = maybe_decorate_final_chunk(chunk, start_time_ms, tool)
case safe_send_reply(grpc_stream, chunk) do
:ok ->
{:cont, %{acc | sent: acc.sent + 1, final_seen?: final_seen?}}
{:error, reason} ->
{:halt, {:error, {:stream_send_failed, reason}}}
end
:skip ->
{:cont, acc}
{:error, reason} ->
{:halt, {:error, reason}}
end
end)
case result do
{:error, reason} ->
{:error, reason}
%{final_seen?: true} ->
:ok
%{final_seen?: false} = acc ->
send_synthetic_final_chunk(grpc_stream, start_time_ms, tool, acc.sent)
end
end
# Safe wrapper for GRPC.Server.send_reply that handles exceptions
# Returns :ok on success, {:error, reason} on failure
defp safe_send_reply(nil, _chunk) do
# Test mode: nil stream means no actual sending
:ok
end
defp safe_send_reply(grpc_stream, chunk) do
try do
_stream = GRPC.Server.send_reply(grpc_stream, chunk)
:ok
rescue
e in GRPC.RPCError ->
{:error, {:grpc_error, e.status, e.message}}
e ->
{:error, {:send_exception, Exception.message(e)}}
catch
:exit, reason ->
{:error, {:stream_exit, reason}}
end
end
defp normalize_stream_item({:ok, %ToolChunk{} = chunk}), do: {:ok, chunk}
defp normalize_stream_item(%ToolChunk{} = chunk), do: {:ok, chunk}
defp normalize_stream_item({:trailers, _trailers}), do: :skip
defp normalize_stream_item({:error, reason}), do: {:error, reason}
defp normalize_stream_item(other), do: {:error, {:unexpected_stream_item, other}}
# Chunk Decoration and Synthetic Final
defp maybe_decorate_final_chunk(%ToolChunk{is_final: true} = chunk, start_time_ms, tool) do
exec_ms = System.monotonic_time(:millisecond) - start_time_ms
metadata =
(chunk.metadata || %{})
|> Map.put_new("execution_time_ms", to_string(exec_ms))
|> Map.put_new("tool_type", to_string(tool.type))
|> Map.put_new("worker_id", to_string(tool.worker_id || ""))
%{chunk | metadata: metadata}
end
defp maybe_decorate_final_chunk(chunk, _start_time_ms, _tool), do: chunk
defp send_synthetic_final_chunk(grpc_stream, start_time_ms, tool, sent_count) do
exec_ms = System.monotonic_time(:millisecond) - start_time_ms
chunk = %ToolChunk{
chunk_id: "sp-final-#{:erlang.unique_integer([:positive, :monotonic])}",
data: <<>>,
is_final: true,
metadata: %{
"synthetic_final" => "true",
"execution_time_ms" => to_string(exec_ms),
"tool_type" => to_string(tool.type),
"chunks_sent" => to_string(sent_count)
}
}
case safe_send_reply(grpc_stream, chunk) do
:ok -> :ok
{:error, reason} -> {:error, {:stream_send_failed, reason}}
end
end
# Error Handling for Streaming
defp raise_streaming_rpc_error(%ExecuteToolRequest{} = request, {:streaming_not_supported, _}) do
hint =
"Streaming execution is not enabled for tool #{request.tool_name}. " <>
"Enable streaming support on the adapter (supports_streaming: true) or use ExecuteTool instead."
raise GRPC.RPCError, status: :unimplemented, message: hint
end
defp raise_streaming_rpc_error(_request, {:stream_send_failed, reason}) do
raise GRPC.RPCError,
status: :unavailable,
message: "Client stream closed: #{inspect(reason)}"
end
defp raise_streaming_rpc_error(_request, {:invalid_parameter, _key, _reason} = error) do
raise GRPC.RPCError, status: :invalid_argument, message: format_error(error)
end
defp raise_streaming_rpc_error(_request, reason) do
raise GRPC.RPCError, status: :internal, message: format_error(reason)
end
defp format_error(reason) when is_binary(reason), do: reason
defp format_error(reason) when is_atom(reason), do: to_string(reason)
defp format_error({:error, reason}), do: format_error(reason)
defp format_error({:unknown_type, type}), do: "Unknown type: #{inspect(type)}"
defp format_error({:invalid_constraints, reason}), do: "Invalid constraints: #{reason}"
defp format_error({:invalid_parameter, key, {:invalid_json, message}}) do
"Invalid parameter #{key}: #{message}"
end
defp format_error({:invalid_parameter, key, {:invalid_json, message, type_url}}) do
"Invalid parameter #{key} (#{type_url}): #{message}"
end
defp format_error({:invalid_parameter, key, reason}) do
"Invalid parameter #{key}: #{inspect(reason)}"
end
defp format_error({:remote_execution_failed, reason}) when is_binary(reason) do
"Remote tool execution failed: #{reason}"
end
defp format_error({:remote_execution_failed, reason}) do
"Remote tool execution failed: #{inspect(reason)}"
end
defp format_error({:validation_failed, details}) when is_map(details) do
"Validation failed: #{inspect(details)}"
end
defp format_error(reason), do: inspect(reason)
# Tool Registration & Discovery
def register_tools(%RegisterToolsRequest{} = request, _stream) do
SLog.info(@log_category, "RegisterTools",
session_id: request.session_id,
worker_id: request.worker_id
)
with {:ok, _session} <- SessionStore.get_session(request.session_id),
tool_specs <- convert_proto_tools_to_specs(request.tools, request.worker_id),
{:ok, registered_names} <- ToolRegistry.register_tools(request.session_id, tool_specs) do
tool_ids =
Map.new(registered_names, fn name -> {name, "#{request.session_id}:#{name}"} end)
%RegisterToolsResponse{
success: true,
tool_ids: tool_ids,
error_message: nil
}
else
{:error, reason} ->
%RegisterToolsResponse{
success: false,
tool_ids: %{},
error_message: format_error(reason)
}
end
end
defp convert_proto_tools_to_specs(tools, worker_id) do
Enum.map(tools, fn tool_reg ->
%{
name: tool_reg.name,
description: tool_reg.description,
parameters: tool_reg.parameters,
metadata:
Map.put(
tool_reg.metadata,
"supports_streaming",
to_string(tool_reg.supports_streaming)
),
worker_id: worker_id
}
end)
end
def get_exposed_elixir_tools(%GetExposedElixirToolsRequest{session_id: session_id}, _stream) do
SLog.debug(@log_category, "GetExposedElixirTools", session_id: session_id)
tools = ToolRegistry.list_exposed_elixir_tools(session_id)
tool_specs =
Enum.map(tools, fn tool ->
# Convert metadata, handling different value types
metadata =
Map.new(tool.metadata, fn
{k, v} when is_binary(v) or is_atom(v) or is_number(v) ->
{to_string(k), to_string(v)}
{k, v} when is_list(v) ->
# Don't include complex lists in metadata
{to_string(k), inspect(v)}
{k, v} ->
# For other types, use inspect
{to_string(k), inspect(v)}
end)
# Remove parameters from metadata since they're handled separately
metadata = Map.delete(metadata, "parameters")
%ToolSpec{
name: tool.name,
description: tool.description,
parameters: encode_parameter_specs(tool.parameters),
metadata: metadata,
supports_streaming: Map.get(metadata, "supports_streaming", "false") == "true"
}
end)
%GetExposedElixirToolsResponse{
tools: tool_specs
}
end
def execute_elixir_tool(%ExecuteElixirToolRequest{} = request, _stream) do
SLog.info(@log_category, "ExecuteElixirTool",
tool_name: request.tool_name,
session_id: request.session_id
)
start_time = System.monotonic_time(:millisecond)
with {:ok, _session} <- SessionStore.get_session(request.session_id),
{:ok, tool} <- ToolRegistry.get_tool(request.session_id, request.tool_name),
:local <- tool.type,
{:ok, params} <- decode_tool_parameters(request.parameters, %{}),
{:ok, result} <-
ToolRegistry.execute_local_tool(request.session_id, request.tool_name, params) do
execution_time = System.monotonic_time(:millisecond) - start_time
{encoded_result, binary_result} = encode_tool_result(result)
%ExecuteElixirToolResponse{
success: true,
result: encoded_result,
binary_result: binary_result || <<>>,
error_message: nil,
metadata: %{
"execution_time" => to_string(execution_time)
},
execution_time_ms: execution_time
}
else
:remote ->
%ExecuteElixirToolResponse{
success: false,
result: nil,
error_message: "Tool #{request.tool_name} is not an Elixir tool",
metadata: %{},
execution_time_ms: System.monotonic_time(:millisecond) - start_time
}
{:error, reason} ->
%ExecuteElixirToolResponse{
success: false,
result: nil,
error_message: format_error(reason),
metadata: %{},
execution_time_ms: System.monotonic_time(:millisecond) - start_time
}
end
end
defp encode_parameter_specs(params) when is_list(params) do
Enum.map(params, &encode_single_parameter_spec/1)
end
defp encode_parameter_specs(_), do: []
defp encode_single_parameter_spec(param) do
# Convert atom keys to strings
normalized_param = normalize_param_map(param)
%ParameterSpec{
name: Map.get(normalized_param, "name", ""),
type: to_string(Map.get(normalized_param, "type", "any")),
description: to_string(Map.get(normalized_param, "description", "")),
required: Map.get(normalized_param, "required", false),
default_value: encode_default_value(Map.get(normalized_param, "default")),
validation_json: Jason.encode!(Map.get(normalized_param, "validation", %{}))
}
end
defp normalize_param_map(%{} = param) do
Map.new(param, fn {k, v} -> {to_string(k), v} end)
end
defp normalize_param_map(param), do: param
defp encode_default_value(nil), do: nil
defp encode_default_value(value) do
{any, _binary} = encode_tool_result(value)
any
end
defp encode_tool_result({:binary, data}) when is_binary(data) do
{empty_any(), data}
end
defp encode_tool_result({:binary, data, metadata}) when is_binary(data) do
{encode_any(metadata), data}
end
defp encode_tool_result(value) do
{encode_any(value), nil}
end
defp encode_any(value) do
case Jason.encode(value) do
{:ok, json_string} when is_binary(json_string) ->
%Any{
type_url: "type.googleapis.com/google.protobuf.StringValue",
value: json_string
}
{:error, _} ->
%Any{
type_url: "type.googleapis.com/google.protobuf.StringValue",
value: inspect(value)
}
end
end
# Used when we only have an opaque binary payload and no metadata
defp empty_any do
%Any{
type_url: "type.googleapis.com/google.protobuf.StringValue",
value: "{}"
}
end
end
defmodule Snakepit.GRPC.ClientImpl do
@moduledoc """
Real gRPC client implementation using generated stubs.
"""
alias Snakepit.Bridge
alias Snakepit.Error.PythonTranslation
alias Snakepit.Logger, as: SLog
alias Snakepit.PythonRuntime
alias Snakepit.Telemetry.Correlation
alias Snakepit.ZeroCopyRef
@default_timeout 30_000
@log_category :grpc
def connect(port) when is_integer(port) do
connect("localhost:#{port}")
end
def connect(address) when is_binary(address) do
opts = []
case GRPC.Stub.connect(address, opts) do
{:ok, channel} ->
# Verify connection with ping
case ping(channel, "connection_test") do
{:ok, _} -> {:ok, channel}
error -> error
end
error ->
error
end
end
def ping(channel, message, opts \\ []) do
request = %Bridge.PingRequest{message: message}
timeout = opts[:timeout] || @default_timeout
call_opts = [timeout: timeout]
case Bridge.BridgeService.Stub.ping(channel, request, call_opts) do
{:ok, response, _headers} ->
{:ok,
%{
message: response.message,
server_time: response.server_time
}}
error ->
handle_error(error)
end
end
def initialize_session(channel, session_id, config \\ %{}, opts \\ []) do
metadata = %{
"elixir_node" => to_string(node()),
"initialized_at" => DateTime.to_iso8601(DateTime.utc_now())
}
session_config = %Bridge.SessionConfig{
enable_caching: Map.get(config, :enable_caching, true),
cache_ttl_seconds: Map.get(config, :cache_ttl_seconds, 60),
enable_telemetry: Map.get(config, :enable_telemetry, false)
}
request = %Bridge.InitializeSessionRequest{
session_id: session_id,
metadata: metadata,
config: session_config
}
timeout = opts[:timeout] || @default_timeout
call_opts = [timeout: timeout]
case Bridge.BridgeService.Stub.initialize_session(channel, request, call_opts) do
{:ok, response, _headers} ->
{:ok,
%{
success: response.success,
available_tools: response.available_tools,
error_message: response.error_message
}}
error ->
handle_error(error)
end
end
def cleanup_session(channel, session_id, force \\ false, opts \\ []) do
request = %Bridge.CleanupSessionRequest{
session_id: session_id,
force: force
}
timeout = opts[:timeout] || @default_timeout
call_opts = [timeout: timeout]
case Bridge.BridgeService.Stub.cleanup_session(channel, request, call_opts) do
{:ok, response, _headers} ->
{:ok,
%{
success: response.success,
resources_cleaned: response.resources_cleaned
}}
error ->
handle_error(error)
end
end
def get_session(channel, session_id, opts \\ []) do
request = %Bridge.GetSessionRequest{
session_id: session_id
}
timeout = opts[:timeout] || @default_timeout
call_opts = [timeout: timeout]
case Bridge.BridgeService.Stub.get_session(channel, request, call_opts) do
{:ok, response, _headers} ->
if response.session_id do
{:ok, %{session: decode_session_response(response)}}
else
{:error, :not_found}
end
{:ok, response} ->
# Handle 2-tuple response
if response.session_id do
{:ok, %{session: decode_session_response(response)}}
else
{:error, :not_found}
end
error ->
handle_error(error)
end
end
def heartbeat(channel, session_id, opts \\ []) do
request = %Bridge.HeartbeatRequest{
session_id: session_id
}
timeout = opts[:timeout] || @default_timeout
call_opts = [timeout: timeout]
case Bridge.BridgeService.Stub.heartbeat(channel, request, call_opts) do
{:ok, response, _headers} ->
{:ok, %{success: response.session_valid}}
{:ok, response} ->
# Handle 2-tuple response
{:ok, %{success: response.session_valid}}
error ->
handle_error(error)
end
end
def execute_tool(channel, session_id, tool_name, parameters, opts \\ []) do
binary_params = Keyword.get(opts, :binary_parameters, %{})
case prepare_execute_tool_request(session_id, tool_name, parameters, binary_params, opts) do
{:ok, request, call_opts} ->
case Bridge.BridgeService.Stub.execute_tool(channel, request, call_opts) do
{:ok, response, _headers} -> handle_tool_response(response)
{:ok, response} -> handle_tool_response(response)
{:error, reason} -> handle_error(reason)
end
{:error, reason} ->
{:error, reason}
end
end
def execute_streaming_tool(channel, session_id, tool_name, parameters, opts \\ []) do
binary_params = Keyword.get(opts, :binary_parameters, %{})
case prepare_execute_stream_request(session_id, tool_name, parameters, binary_params, opts) do
{:ok, request, call_opts} ->
Bridge.BridgeService.Stub.execute_streaming_tool(channel, request, call_opts)
{:error, reason} ->
{:error, reason}
end
end
# Helper functions
defp decode_session_response(response) do
%{
id: response.session_id,
# Assume active if we got a response
active: true,
created_at: response.created_at,
# Not provided in response
last_activity: nil,
metadata: Map.new(response.metadata || %{})
}
end
defp handle_error({:error, %GRPC.RPCError{} = error}) do
SLog.error(@log_category, "gRPC error: #{inspect(error)}")
case error.status do
3 -> {:error, :invalid_argument}
5 -> {:error, :not_found}
13 -> {:error, :internal}
14 -> {:error, :unavailable}
_ -> {:error, error}
end
end
defp handle_error(error), do: error
# Simple encoder for tool parameters - just use JSON encoding for now
defp infer_and_encode_any(value) do
case Jason.encode(value) do
{:ok, json_value} ->
{:ok,
%Google.Protobuf.Any{
type_url: "type.googleapis.com/google.protobuf.StringValue",
value: json_value
}}
{:error, %Jason.EncodeError{} = encode_error} ->
{:error, {:invalid_parameter, :json_encode_failed, Exception.message(encode_error)}}
{:error, other} ->
{:error, {:invalid_parameter, :json_encode_failed, inspect(other)}}
end
end
@doc false
def decode_tool_response(response), do: handle_tool_response(response)
defp handle_tool_response(%Bridge.ExecuteToolResponse{
success: true,
result: any_result,
binary_result: binary_result
}) do
if binary_payload?(binary_result) do
metadata = decode_any(any_result)
{:ok, format_binary_result(binary_result, metadata)}
else
{:ok, decode_any(any_result)}
end
end
defp handle_tool_response(%Bridge.ExecuteToolResponse{success: false, error_message: error}) do
case PythonTranslation.from_error_message(error) do
{:ok, translated} -> {:error, translated}
:error -> {:error, error}
end
end
defp binary_payload?(binary) when is_binary(binary), do: byte_size(binary) > 0
defp binary_payload?(_), do: false
defp format_binary_result(binary_result, metadata) do
case metadata do
nil -> {:binary, binary_result}
%{} = map when map_size(map) == 0 -> {:binary, binary_result}
_ -> {:binary, binary_result, metadata}
end
end
defp decode_any(nil), do: nil
defp decode_any(%Google.Protobuf.Any{value: value}) when is_binary(value) do
case Jason.decode(value) do
{:ok, decoded} -> ZeroCopyRef.maybe_from_map(decoded)
{:error, _} -> value
end
end
defp decode_any(%Google.Protobuf.Any{value: value}), do: value
defp build_execute_tool_request(session_id, tool_name, proto_params, binary_params, metadata) do
%Bridge.ExecuteToolRequest{
session_id: session_id,
tool_name: tool_name,
parameters: proto_params,
binary_parameters: binary_params,
metadata: metadata
}
end
defp encode_parameters(parameters) do
Enum.reduce_while(parameters, {:ok, %{}}, fn {k, v}, {:ok, acc} ->
case infer_and_encode_any(v) do
{:ok, proto_any} ->
{:cont, {:ok, Map.put(acc, to_string(k), proto_any)}}
{:error, reason} ->
{:halt, {:error, reason}}
end
end)
end
defp sanitize_parameters(parameters) when is_map(parameters) do
parameters
|> Map.delete(:correlation_id)
|> Map.delete("correlation_id")
end
defp sanitize_parameters(parameters) when is_list(parameters) do
parameters
|> Enum.reject(fn
{:correlation_id, _} -> true
{"correlation_id", _} -> true
_ -> false
end)
end
defp sanitize_parameters(parameters), do: parameters
defp encode_binary_parameters(nil), do: {:ok, %{}}
defp encode_binary_parameters(binary_params) when is_map(binary_params) do
Enum.reduce_while(binary_params, {:ok, %{}}, fn {key, value}, {:ok, acc} ->
if is_binary(value) do
{:cont, {:ok, Map.put(acc, to_string(key), value)}}
else
{:halt, {:error, {:invalid_parameter, key, :not_binary}}}
end
end)
end
defp encode_binary_parameters(_), do: {:ok, %{}}
@doc false
def prepare_execute_tool_request(session_id, tool_name, parameters, binary_params, opts \\ []) do
prepare_execute_request(
session_id,
tool_name,
parameters,
binary_params,
opts,
@default_timeout,
false
)
end
@doc false
def prepare_execute_stream_request(session_id, tool_name, parameters, binary_params, opts \\ []) do
prepare_execute_request(
session_id,
tool_name,
parameters,
binary_params,
opts,
300_000,
true
)
end
defp prepare_execute_request(
session_id,
tool_name,
parameters,
binary_params,
opts,
default_timeout,
stream?
) do
correlation_id = resolve_correlation_id(parameters, opts)
parameters = sanitize_parameters(parameters)
with {:ok, proto_params} <- encode_parameters(parameters),
{:ok, encoded_binary} <- encode_binary_parameters(binary_params) do
metadata = build_request_metadata(correlation_id)
request =
build_execute_tool_request(session_id, tool_name, proto_params, encoded_binary, metadata)
|> maybe_put_stream(stream?)
timeout = opts[:timeout] || default_timeout
call_opts = [timeout: timeout] |> maybe_put_correlation_metadata(correlation_id)
{:ok, request, call_opts}
end
end
defp maybe_put_stream(request, true), do: Map.put(request, :stream, true)
defp maybe_put_stream(request, false), do: request
defp resolve_correlation_id(parameters, opts) do
opts_correlation = Keyword.get(opts, :correlation_id)
(opts_correlation || extract_correlation_id(parameters))
|> Correlation.ensure()
end
defp extract_correlation_id(parameters) when is_map(parameters) do
Map.get(parameters, :correlation_id) || Map.get(parameters, "correlation_id")
end
defp extract_correlation_id(parameters) when is_list(parameters) do
Enum.find_value(parameters, fn
{:correlation_id, value} -> value
{"correlation_id", value} -> value
_ -> nil
end)
end
defp extract_correlation_id(_), do: nil
defp build_request_metadata(correlation_id) when is_binary(correlation_id) do
%{"correlation_id" => correlation_id}
|> Map.merge(PythonRuntime.runtime_metadata())
end
defp maybe_put_correlation_metadata(call_opts, correlation_id) when is_binary(correlation_id) do
existing = Keyword.get(call_opts, :metadata, [])
filtered =
Enum.reject(existing, fn {key, _} ->
String.downcase(to_string(key)) == "x-snakepit-correlation-id"
end)
Keyword.put(call_opts, :metadata, [{"x-snakepit-correlation-id", correlation_id} | filtered])
end
end
defmodule Snakepit.GRPC.Client do
@moduledoc """
gRPC client for the unified bridge protocol.
Delegates to the real implementation when available.
"""
require Logger
alias Snakepit.Defaults
alias Snakepit.GRPC.ClientImpl
# Uncomment when logging is added to this module:
# alias Snakepit.Logger, as: SLog
def connect(port) when is_integer(port) do
connect("localhost:#{port}")
end
def connect(address) when is_binary(address) do
ClientImpl.connect(address)
end
def ping(channel, message, opts \\ []) do
if mock_channel?(channel) do
# Mock implementation
{:ok, %{message: "Pong: #{message}", server_time: DateTime.utc_now()}}
else
ClientImpl.ping(channel, message, opts)
end
end
def initialize_session(channel, session_id, config \\ %{}, opts \\ []) do
if mock_channel?(channel) do
# Mock implementation
{:ok, %{success: true, available_tools: %{}}}
else
ClientImpl.initialize_session(channel, session_id, config, opts)
end
end
def cleanup_session(channel, session_id, force \\ false, opts \\ []) do
if mock_channel?(channel) do
# Mock implementation
{:ok, %{success: true, resources_cleaned: 2}}
else
ClientImpl.cleanup_session(channel, session_id, force, opts)
end
end
def execute_tool(channel, session_id, tool_name, parameters, opts \\ []) do
if mock_channel?(channel) do
if test_pid = Map.get(channel, :test_pid) do
send(test_pid, {:grpc_client_execute_tool, session_id, tool_name, parameters, opts})
end
{:ok, %{success: true, result: %{}, error_message: ""}}
else
ClientImpl.execute_tool(channel, session_id, tool_name, parameters, opts)
end
end
def execute_streaming_tool(channel, session_id, tool_name, parameters, opts \\ []) do
if mock_channel?(channel) do
if test_pid = Map.get(channel, :test_pid) do
send(
test_pid,
{:grpc_client_execute_streaming_tool, session_id, tool_name, parameters, opts}
)
end
# Mock implementation for testing - return a simple stream
stream =
Stream.iterate(1, &(&1 + 1))
|> Stream.take(5)
|> Stream.map(fn i ->
{:ok,
%{
chunk_id: "mock-#{i}",
data: Jason.encode!(%{"step" => i, "total" => 5}),
is_final: i == 5
}}
end)
{:ok, stream}
else
ClientImpl.execute_streaming_tool(
channel,
session_id,
tool_name,
parameters,
opts
)
end
end
# Existing methods for backward compatibility
def execute(channel, command, args, timeout \\ nil) do
timeout = timeout || Defaults.grpc_client_execute_timeout()
# Legacy support - redirect to execute_tool
execute_tool(channel, "default_session", command, args, timeout: timeout)
end
def health(channel, client_id) do
ping(channel, "health_check_#{client_id}")
end
def get_info(_channel) do
# Return mock info for now
{:ok,
%{
version: "1.0.0",
capabilities: ["tools", "streaming"]
}}
end
def get_session(channel, session_id, opts \\ []) do
if mock_channel?(channel) do
# Mock implementation
{:ok, %{session: %{id: session_id, active: true}}}
else
ClientImpl.get_session(channel, session_id, opts)
end
end
def heartbeat(channel, session_id, opts \\ []) do
if mock_channel?(channel) do
# Mock implementation
{:ok, %{success: true}}
else
ClientImpl.heartbeat(channel, session_id, opts)
end
end
defp mock_channel?(channel) when is_map(channel) do
Map.get(channel, :mock, false)
end
defp mock_channel?(_channel), do: true
end
defmodule Snakepit.GRPC.Endpoint do
@moduledoc """
gRPC endpoint for the Snakepit bridge server.
This module defines the gRPC endpoint that handles incoming
requests for the unified bridge protocol.
"""
use GRPC.Endpoint
intercept(GRPC.Server.Interceptors.Logger)
run(Snakepit.GRPC.BridgeServer)
end
defmodule Snakepit.Hardware.CPUDetector do
@moduledoc """
CPU hardware detection.
Detects CPU model, cores, threads, memory, and CPU features (SSE, AVX, etc.).
"""
@type cpu_info :: %{
cores: pos_integer(),
threads: pos_integer(),
model: String.t(),
features: [atom()],
memory_total_mb: non_neg_integer()
}
@doc """
Detects CPU hardware information.
Returns a map with:
- `:cores` - Number of physical CPU cores
- `:threads` - Number of logical threads (cores * hyperthreading)
- `:model` - CPU model name string
- `:features` - List of detected CPU feature atoms (e.g., `:avx`, `:sse4_2`)
- `:memory_total_mb` - Total system memory in MB
"""
@spec detect() :: cpu_info()
def detect do
%{
cores: detect_cores(),
threads: detect_threads(),
model: detect_model(),
features: detect_features(),
memory_total_mb: detect_memory_mb()
}
end
@spec detect_cores() :: pos_integer()
defp detect_cores do
# Try to get physical cores (System.schedulers always returns >= 1)
n = System.schedulers()
max(1, div(n, 2))
end
@spec detect_threads() :: pos_integer()
defp detect_threads do
# System.schedulers_online always returns >= 1
System.schedulers_online()
end
@spec detect_model() :: String.t()
defp detect_model do
case :os.type() do
{:unix, :linux} -> detect_model_linux()
{:unix, :darwin} -> detect_model_darwin()
{:win32, _} -> detect_model_windows()
_ -> "Unknown"
end
end
defp detect_model_linux do
try do
case File.read("/proc/cpuinfo") do
{:ok, content} ->
content
|> String.split("\n")
|> Enum.find(fn line -> String.starts_with?(line, "model name") end)
|> case do
nil -> "Unknown"
line -> line |> String.split(":") |> List.last() |> String.trim()
end
_ ->
"Unknown"
end
rescue
_ -> "Unknown"
end
end
defp detect_model_darwin do
try do
case System.cmd("sysctl", ["-n", "machdep.cpu.brand_string"], stderr_to_stdout: true) do
{output, 0} -> String.trim(output)
_ -> "Apple Silicon"
end
rescue
_ -> "Unknown"
end
end
defp detect_model_windows do
try do
case System.cmd("wmic", ["cpu", "get", "name", "/value"], stderr_to_stdout: true) do
{output, 0} ->
output
|> String.split("=")
|> List.last()
|> String.trim()
_ ->
"Unknown"
end
rescue
_ -> "Unknown"
end
end
@spec detect_features() :: [atom()]
defp detect_features do
case :os.type() do
{:unix, :linux} -> detect_features_linux()
{:unix, :darwin} -> detect_features_darwin()
_ -> []
end
end
defp detect_features_linux do
try do
case File.read("/proc/cpuinfo") do
{:ok, content} ->
content
|> String.split("\n")
|> Enum.find(fn line -> String.starts_with?(line, "flags") end)
|> case do
nil ->
[]
line ->
line
|> String.split(":")
|> List.last()
|> String.split()
|> Enum.map(&normalize_feature/1)
|> Enum.filter(&(&1 != nil))
end
_ ->
[]
end
rescue
_ -> []
end
end
defp detect_features_darwin do
try do
features = []
features =
case System.cmd("sysctl", ["-n", "hw.optional.avx1_0"], stderr_to_stdout: true) do
{"1\n", 0} -> [:avx | features]
_ -> features
end
features =
case System.cmd("sysctl", ["-n", "hw.optional.avx2_0"], stderr_to_stdout: true) do
{"1\n", 0} -> [:avx2 | features]
_ -> features
end
features =
case System.cmd("sysctl", ["-n", "hw.optional.sse4_1"], stderr_to_stdout: true) do
{"1\n", 0} -> [:sse4_1 | features]
_ -> features
end
features =
case System.cmd("sysctl", ["-n", "hw.optional.sse4_2"], stderr_to_stdout: true) do
{"1\n", 0} -> [:sse4_2 | features]
_ -> features
end
# ARM/Apple Silicon features
features =
case System.cmd("sysctl", ["-n", "hw.optional.neon"], stderr_to_stdout: true) do
{"1\n", 0} -> [:neon | features]
_ -> features
end
features
rescue
_ -> []
end
end
@known_features %{
"avx" => :avx,
"avx2" => :avx2,
"avx512f" => :avx512,
"sse4_1" => :sse4_1,
"sse4_2" => :sse4_2,
"sse3" => :sse3,
"ssse3" => :ssse3,
"fma" => :fma,
"f16c" => :f16c,
"aes" => :aes,
"pclmulqdq" => :pclmul,
"neon" => :neon,
"asimd" => :asimd
}
defp normalize_feature(flag) do
Map.get(@known_features, String.downcase(flag))
end
@spec detect_memory_mb() :: non_neg_integer()
defp detect_memory_mb do
# Use :erlang.memory() total as approximation
# This gives the Erlang VM's view but is portable
try do
case :os.type() do
{:unix, :linux} ->
detect_memory_linux()
{:unix, :darwin} ->
detect_memory_darwin()
_ ->
# Fallback to erlang memory
div(:erlang.memory(:total), 1024 * 1024)
end
rescue
_ -> 0
end
end
defp detect_memory_linux do
case File.read("/proc/meminfo") do
{:ok, content} ->
content
|> String.split("\n")
|> Enum.find(fn line -> String.starts_with?(line, "MemTotal") end)
|> case do
nil ->
0
line ->
line
|> String.replace(~r/[^\d]/, "")
|> String.trim()
|> String.to_integer()
|> div(1024)
end
_ ->
0
end
end
defp detect_memory_darwin do
case System.cmd("sysctl", ["-n", "hw.memsize"], stderr_to_stdout: true) do
{output, 0} ->
output
|> String.trim()
|> String.to_integer()
|> div(1024 * 1024)
_ ->
0
end
end
end
defmodule Snakepit.Hardware.CUDADetector do
@moduledoc """
CUDA GPU hardware detection.
Detects NVIDIA CUDA-capable GPUs using nvidia-smi when available.
"""
@type cuda_device :: %{
id: non_neg_integer(),
name: String.t(),
memory_total_mb: non_neg_integer(),
memory_free_mb: non_neg_integer(),
compute_capability: String.t() | nil
}
@type cuda_info :: %{
version: String.t(),
driver_version: String.t(),
devices: [cuda_device()],
cudnn_version: String.t() | nil
}
@doc """
Detects CUDA GPU information.
Returns nil if CUDA is not available, or a map with:
- `:version` - CUDA runtime version (e.g., "12.1")
- `:driver_version` - NVIDIA driver version
- `:devices` - List of CUDA device maps
- `:cudnn_version` - cuDNN version if available, nil otherwise
"""
@spec detect() :: cuda_info() | nil
def detect do
with {:ok, driver_version} <- detect_driver_version(),
{:ok, cuda_version} <- detect_cuda_version(),
{:ok, devices} <- detect_devices() do
%{
version: cuda_version,
driver_version: driver_version,
devices: devices,
cudnn_version: detect_cudnn_version()
}
else
_ -> nil
end
end
@spec detect_driver_version() :: {:ok, String.t()} | :error
defp detect_driver_version do
case run_nvidia_smi(["--query-gpu=driver_version", "--format=csv,noheader,nounits"]) do
{:ok, output} ->
version =
output
|> String.split("\n")
|> List.first()
|> String.trim()
if version != "" do
{:ok, version}
else
:error
end
:error ->
:error
end
end
@spec detect_cuda_version() :: {:ok, String.t()} | :error
defp detect_cuda_version do
# Try nvidia-smi for CUDA version
case run_nvidia_smi(["--query-gpu=cuda_version", "--format=csv,noheader,nounits"]) do
{:ok, output} ->
version =
output
|> String.split("\n")
|> List.first()
|> String.trim()
if version != "" and version != "[N/A]" do
{:ok, version}
else
detect_cuda_version_fallback()
end
:error ->
detect_cuda_version_fallback()
end
end
defp detect_cuda_version_fallback do
# Try nvcc --version as fallback
case System.cmd("nvcc", ["--version"], stderr_to_stdout: true) do
{output, 0} ->
case Regex.run(~r/release (\d+\.\d+)/, output) do
[_, version] -> {:ok, version}
_ -> :error
end
_ ->
:error
end
rescue
_ -> :error
end
@spec detect_devices() :: {:ok, [cuda_device()]} | :error
defp detect_devices do
query = [
"--query-gpu=index,name,memory.total,memory.free,compute_cap",
"--format=csv,noheader,nounits"
]
case run_nvidia_smi(query) do
{:ok, output} ->
devices =
output
|> String.split("\n")
|> Enum.filter(&(&1 != ""))
|> Enum.map(&parse_device_line/1)
|> Enum.filter(&(&1 != nil))
{:ok, devices}
:error ->
:error
end
end
defp parse_device_line(line) do
case String.split(line, ", ") do
[id_str, name, mem_total_str, mem_free_str, compute_cap] ->
%{
id: safe_parse_integer(id_str, 0),
name: String.trim(name),
memory_total_mb: safe_parse_integer(mem_total_str, 0),
memory_free_mb: safe_parse_integer(mem_free_str, 0),
compute_capability: normalize_compute_cap(compute_cap)
}
[id_str, name, mem_total_str, mem_free_str] ->
%{
id: safe_parse_integer(id_str, 0),
name: String.trim(name),
memory_total_mb: safe_parse_integer(mem_total_str, 0),
memory_free_mb: safe_parse_integer(mem_free_str, 0),
compute_capability: nil
}
_ ->
nil
end
end
defp safe_parse_integer(str, default) do
str
|> String.trim()
|> Integer.parse()
|> case do
{n, _} -> n
:error -> default
end
end
defp normalize_compute_cap(cap) do
cap = String.trim(cap)
if cap in ["[N/A]", "N/A", ""] do
nil
else
cap
end
end
@spec detect_cudnn_version() :: String.t() | nil
defp detect_cudnn_version do
# cuDNN version is not easily detectable without Python
# Return nil for now, can be populated via Python probe
nil
end
@spec run_nvidia_smi([String.t()]) :: {:ok, String.t()} | :error
defp run_nvidia_smi(args) do
case System.cmd("nvidia-smi", args, stderr_to_stdout: true) do
{output, 0} -> {:ok, output}
_ -> :error
end
rescue
_ -> :error
end
end
defmodule Snakepit.Hardware.Detector do
@moduledoc """
Unified hardware detection module.
Aggregates CPU, CUDA, MPS, and ROCm detection into a single hardware info structure.
Results are cached in ETS for performance.
"""
alias Snakepit.Hardware.{CPUDetector, CUDADetector, MPSDetector, ROCmDetector}
@hardware_info_key {__MODULE__, :hardware_info}
@capabilities_key {__MODULE__, :capabilities}
@cache_miss :snakepit_cache_miss
@type accelerator :: :cpu | :cuda | :mps | :rocm
@type hardware_info :: %{
accelerator: accelerator(),
cpu: CPUDetector.cpu_info(),
cuda: CUDADetector.cuda_info() | nil,
mps: MPSDetector.mps_info() | nil,
rocm: ROCmDetector.rocm_info() | nil,
platform: String.t()
}
@type capabilities :: %{
cuda: boolean(),
mps: boolean(),
rocm: boolean(),
avx: boolean(),
avx2: boolean(),
avx512: boolean(),
cuda_version: String.t() | nil,
cudnn_version: String.t() | nil,
cudnn: boolean()
}
@doc """
Detects all hardware information.
Returns a map with aggregated hardware info from all detectors.
Results are cached for performance.
"""
@spec detect() :: hardware_info()
def detect do
fetch_cached(@hardware_info_key, &do_detect/0)
end
@doc """
Returns hardware capability flags.
Returns a map of boolean capability flags for quick feature checks.
"""
@spec capabilities() :: capabilities()
def capabilities do
fetch_cached(@capabilities_key, fn ->
detect()
|> build_capabilities()
end)
end
@doc """
Clears the hardware detection cache.
Forces re-detection on next call to detect/0 or capabilities/0.
"""
@spec clear_cache() :: :ok
def clear_cache do
:persistent_term.erase(@hardware_info_key)
:persistent_term.erase(@capabilities_key)
:ok
end
@spec do_detect() :: hardware_info()
defp do_detect do
cpu = CPUDetector.detect()
cuda = CUDADetector.detect()
mps = MPSDetector.detect()
rocm = ROCmDetector.detect()
accelerator = determine_accelerator(cuda, mps, rocm)
platform = build_platform_string()
%{
accelerator: accelerator,
cpu: cpu,
cuda: cuda,
mps: mps,
rocm: rocm,
platform: platform
}
end
@spec determine_accelerator(
CUDADetector.cuda_info() | nil,
MPSDetector.mps_info() | nil,
ROCmDetector.rocm_info() | nil
) :: accelerator()
defp determine_accelerator(cuda, mps, rocm) do
cond do
cuda != nil and cuda.devices != [] -> :cuda
mps != nil and mps.available -> :mps
rocm != nil and rocm.devices != [] -> :rocm
true -> :cpu
end
end
@spec build_platform_string() :: String.t()
defp build_platform_string do
os = os_name()
arch = arch_name()
"#{os}-#{arch}"
end
defp os_name do
case :os.type() do
{:unix, :linux} -> "linux"
{:unix, :darwin} -> "macos"
{:win32, _} -> "windows"
_ -> "unknown"
end
end
defp arch_name do
# :system_architecture always returns a charlist
arch = :erlang.system_info(:system_architecture)
arch_str = List.to_string(arch)
normalize_arch(arch_str)
end
defp normalize_arch(arch_str) do
arch_lower = String.downcase(arch_str)
cond do
String.contains?(arch_lower, "x86_64") or String.contains?(arch_lower, "amd64") ->
"x86_64"
String.contains?(arch_lower, "aarch64") or String.contains?(arch_lower, "arm64") ->
"arm64"
String.contains?(arch_lower, "arm") ->
"arm"
true ->
arch_str
end
end
@spec build_capabilities(hardware_info()) :: capabilities()
defp build_capabilities(info) do
cpu_features = info.cpu.features
%{
cuda: info.cuda != nil and info.cuda.devices != [],
mps: info.mps != nil and info.mps.available,
rocm: info.rocm != nil and info.rocm.devices != [],
avx: :avx in cpu_features,
avx2: :avx2 in cpu_features,
avx512: :avx512 in cpu_features,
cuda_version: get_in(info, [:cuda, :version]),
cudnn_version: get_in(info, [:cuda, :cudnn_version]),
cudnn: get_in(info, [:cuda, :cudnn_version]) != nil
}
end
@spec fetch_cached(term(), (-> term())) :: term()
defp fetch_cached(key, fun) do
case :persistent_term.get(key, @cache_miss) do
@cache_miss ->
value = fun.()
:persistent_term.put(key, value)
value
value ->
value
end
end
end
defmodule Snakepit.Hardware.MPSDetector do
@moduledoc """
Apple Metal Performance Shaders (MPS) hardware detection.
Detects Apple Silicon GPU availability on macOS.
"""
@type mps_info :: %{
available: boolean(),
device_name: String.t(),
memory_total_mb: non_neg_integer()
}
@doc """
Detects MPS (Apple Metal) availability.
Returns nil on non-macOS platforms, or a map with:
- `:available` - true if MPS is available
- `:device_name` - Name of the Metal device
- `:memory_total_mb` - GPU memory (shared memory on Apple Silicon)
"""
@spec detect() :: mps_info() | nil
def detect do
case :os.type() do
{:unix, :darwin} -> detect_darwin()
_ -> nil
end
end
defp detect_darwin do
# Check if we're on Apple Silicon
case detect_apple_silicon() do
{:ok, device_name} ->
memory_mb = detect_unified_memory()
%{
available: true,
device_name: device_name,
memory_total_mb: memory_mb
}
:error ->
# Check for Intel Mac with AMD GPU
case detect_discrete_gpu() do
{:ok, device_name, memory_mb} ->
%{
available: true,
device_name: device_name,
memory_total_mb: memory_mb
}
:error ->
nil
end
end
end
defp detect_apple_silicon do
# Check CPU brand for Apple Silicon
case System.cmd("sysctl", ["-n", "machdep.cpu.brand_string"], stderr_to_stdout: true) do
{output, 0} ->
brand = String.trim(output)
if String.contains?(brand, "Apple") do
# Get specific chip name
case System.cmd("sysctl", ["-n", "hw.model"], stderr_to_stdout: true) do
{model, 0} ->
model = String.trim(model)
chip = extract_chip_name(brand, model)
{:ok, "#{chip} GPU"}
_ ->
{:ok, "Apple GPU"}
end
else
:error
end
_ ->
:error
end
rescue
_ -> :error
end
defp extract_chip_name(brand, _model) do
cond do
String.contains?(brand, "M1") -> "Apple M1"
String.contains?(brand, "M2") -> "Apple M2"
String.contains?(brand, "M3") -> "Apple M3"
String.contains?(brand, "M4") -> "Apple M4"
true -> "Apple Silicon"
end
end
defp detect_unified_memory do
# On Apple Silicon, GPU uses unified memory
case System.cmd("sysctl", ["-n", "hw.memsize"], stderr_to_stdout: true) do
{output, 0} ->
output
|> String.trim()
|> String.to_integer()
|> div(1024 * 1024)
_ ->
0
end
rescue
_ -> 0
end
defp detect_discrete_gpu do
# For Intel Macs with discrete AMD GPU
case System.cmd("system_profiler", ["SPDisplaysDataType"], stderr_to_stdout: true) do
{output, 0} ->
parse_discrete_gpu(output)
_ ->
:error
end
rescue
_ -> :error
end
defp parse_discrete_gpu(output) do
# Look for AMD Radeon or similar
lines = String.split(output, "\n")
gpu_name =
Enum.find_value(lines, fn line ->
cond do
String.contains?(line, "AMD Radeon") ->
line |> String.trim() |> String.replace(":", "")
String.contains?(line, "Metal Family") ->
# Has Metal support
nil
true ->
nil
end
end)
case gpu_name do
nil ->
:error
name ->
# Try to find VRAM
vram =
Enum.find_value(lines, 0, fn line ->
if String.contains?(line, "VRAM") do
case Regex.run(~r/(\d+)\s*MB/, line) do
[_, mb] -> String.to_integer(mb)
_ -> nil
end
end
end)
{:ok, name, vram}
end
end
end
defmodule Snakepit.Hardware.ROCmDetector do
@moduledoc """
AMD ROCm GPU hardware detection.
Detects AMD GPUs with ROCm support using rocm-smi when available.
"""
@type rocm_device :: %{
id: non_neg_integer(),
name: String.t(),
memory_total_mb: non_neg_integer(),
memory_free_mb: non_neg_integer()
}
@type rocm_info :: %{
version: String.t(),
devices: [rocm_device()]
}
@doc """
Detects ROCm GPU information.
Returns nil if ROCm is not available, or a map with:
- `:version` - ROCm version
- `:devices` - List of ROCm device maps
"""
@spec detect() :: rocm_info() | nil
def detect do
with {:ok, version} <- detect_version(),
{:ok, devices} <- detect_devices() do
%{
version: version,
devices: devices
}
else
_ -> nil
end
end
@spec detect_version() :: {:ok, String.t()} | :error
defp detect_version do
# Try rocm-smi for version
case System.cmd("rocm-smi", ["--showversion"], stderr_to_stdout: true) do
{output, 0} ->
case Regex.run(~r/ROCm version:\s*(\d+\.\d+(?:\.\d+)?)/, output) do
[_, version] -> {:ok, version}
_ -> try_rocm_version_file()
end
_ ->
try_rocm_version_file()
end
rescue
_ -> :error
end
defp try_rocm_version_file do
case File.read("/opt/rocm/.info/version") do
{:ok, content} ->
version = String.trim(content)
if version != "", do: {:ok, version}, else: :error
_ ->
:error
end
end
@spec detect_devices() :: {:ok, [rocm_device()]} | :error
defp detect_devices do
case System.cmd("rocm-smi", ["--showid", "--showmeminfo", "vram"], stderr_to_stdout: true) do
{output, 0} ->
devices = parse_rocm_devices(output)
{:ok, devices}
_ ->
:error
end
rescue
_ -> :error
end
defp parse_rocm_devices(output) do
# Parse rocm-smi output format
# This is simplified - actual parsing depends on rocm-smi version
lines = String.split(output, "\n")
lines
|> Enum.chunk_every(4)
|> Enum.with_index()
|> Enum.map(fn {chunk, idx} ->
parse_device_chunk(chunk, idx)
end)
|> Enum.filter(&(&1 != nil))
end
defp parse_device_chunk(chunk, idx) do
name =
Enum.find_value(chunk, "AMD GPU", fn line ->
if String.contains?(line, "GPU") do
line
|> String.replace(~r/GPU\[\d+\]/, "")
|> String.replace(":", "")
|> String.trim()
end
end)
%{
id: idx,
name: if(name == "", do: "AMD GPU #{idx}", else: name),
memory_total_mb: 0,
memory_free_mb: 0
}
end
end
defmodule Snakepit.Hardware.Selector do
@moduledoc """
Device selection logic for hardware abstraction.
Provides intelligent device selection based on availability, preferences,
and fallback strategies.
"""
alias Snakepit.Hardware.Detector
@type device ::
:cpu
| :cuda
| :mps
| :rocm
| {:cuda, non_neg_integer()}
| {:rocm, non_neg_integer()}
@type device_preference :: :auto | :cpu | :cuda | :mps | :rocm | {:cuda, non_neg_integer()}
@doc """
Selects a device based on preference.
## Options
- `:auto` - Automatically select the best available accelerator
- `:cpu` - Select CPU (always available)
- `:cuda` - Select CUDA (fails if not available)
- `:mps` - Select MPS (fails if not available or not on macOS)
- `:rocm` - Select ROCm (fails if not available)
- `{:cuda, device_id}` - Select specific CUDA device
## Returns
- `{:ok, device}` on success
- `{:error, :device_not_available}` if requested device is unavailable
"""
@spec select(device_preference()) :: {:ok, device()} | {:error, :device_not_available}
def select(:auto) do
info = Detector.detect()
device =
cond do
info.cuda != nil and info.cuda.devices != [] ->
{:cuda, 0}
info.mps != nil and info.mps.available ->
:mps
info.rocm != nil and info.rocm.devices != [] ->
{:rocm, 0}
true ->
:cpu
end
{:ok, device}
end
def select(:cpu) do
{:ok, :cpu}
end
def select(:cuda) do
info = Detector.detect()
case info.cuda do
%{devices: [_ | _]} ->
{:ok, {:cuda, 0}}
_ ->
{:error, :device_not_available}
end
end
def select({:cuda, device_id}) when is_integer(device_id) do
info = Detector.detect()
case info.cuda do
%{devices: devices} when is_list(devices) ->
if Enum.any?(devices, &(&1.id == device_id)) do
{:ok, {:cuda, device_id}}
else
{:error, :device_not_available}
end
_ ->
{:error, :device_not_available}
end
end
def select(:mps) do
info = Detector.detect()
case info.mps do
%{available: true} ->
{:ok, :mps}
_ ->
{:error, :device_not_available}
end
end
def select(:rocm) do
info = Detector.detect()
case info.rocm do
%{devices: [_ | _]} ->
{:ok, {:rocm, 0}}
_ ->
{:error, :device_not_available}
end
end
def select({:rocm, device_id}) when is_integer(device_id) do
info = Detector.detect()
case info.rocm do
%{devices: devices} when is_list(devices) ->
if Enum.any?(devices, &(&1.id == device_id)) do
{:ok, {:rocm, device_id}}
else
{:error, :device_not_available}
end
_ ->
{:error, :device_not_available}
end
end
def select(_), do: {:error, :device_not_available}
@doc """
Selects the first available device from a preference list.
Tries each device in order until one is available, returning that device.
If no devices are available, returns `{:error, :no_device}`.
## Examples
iex> Hardware.Selector.select_with_fallback([:cuda, :mps, :cpu])
{:ok, :cpu} # if CUDA and MPS unavailable
"""
@spec select_with_fallback([device_preference()]) ::
{:ok, device()} | {:error, :no_device}
def select_with_fallback([]) do
{:error, :no_device}
end
def select_with_fallback([preference | rest]) do
case select(preference) do
{:ok, device} -> {:ok, device}
{:error, _} -> select_with_fallback(rest)
end
end
@doc """
Returns information about a selected device.
Returns a map with device details useful for logging and telemetry.
"""
@spec device_info(device()) :: map()
def device_info(:cpu) do
info = Detector.detect()
%{
type: :cpu,
name: info.cpu.model,
cores: info.cpu.cores,
threads: info.cpu.threads,
memory_mb: info.cpu.memory_total_mb
}
end
def device_info({:cuda, device_id}) do
info = Detector.detect()
case info.cuda do
%{devices: devices, version: version} ->
device = Enum.find(devices, &(&1.id == device_id))
if device do
%{
type: :cuda,
device_id: device_id,
name: device.name,
memory_total_mb: device.memory_total_mb,
memory_free_mb: device.memory_free_mb,
cuda_version: version,
compute_capability: device.compute_capability
}
else
%{type: :cuda, device_id: device_id, error: :device_not_found}
end
_ ->
%{type: :cuda, device_id: device_id, error: :cuda_not_available}
end
end
def device_info(:cuda) do
device_info({:cuda, 0})
end
def device_info(:mps) do
info = Detector.detect()
case info.mps do
%{available: true} = mps ->
%{
type: :mps,
name: mps.device_name,
memory_mb: mps.memory_total_mb
}
_ ->
%{type: :mps, error: :mps_not_available}
end
end
def device_info({:rocm, device_id}) do
info = Detector.detect()
case info.rocm do
%{devices: devices, version: version} ->
device = Enum.find(devices, &(&1.id == device_id))
if device do
%{
type: :rocm,
device_id: device_id,
name: device.name,
memory_total_mb: device.memory_total_mb,
rocm_version: version
}
else
%{type: :rocm, device_id: device_id, error: :device_not_found}
end
_ ->
%{type: :rocm, device_id: device_id, error: :rocm_not_available}
end
end
def device_info(:rocm) do
device_info({:rocm, 0})
end
end
defmodule Snakepit.Logger.Redaction do
@moduledoc false
@max_map_keys 5
@max_key_length 32
@sample_limit 5
def describe(%struct_module{} = _struct) do
module =
struct_module
|> Atom.to_string()
|> String.trim_leading("Elixir.")
"struct(#{module})"
end
def describe(map) when is_map(map) do
keys =
map
|> Map.delete(:__struct__)
|> Map.keys()
|> Enum.map(&format_key/1)
|> Enum.take(@max_map_keys)
"map(keys: #{inspect(keys)}, count: #{map_size(map)})"
end
def describe(binary) when is_binary(binary) do
"binary(len: #{byte_size(binary)})"
end
def describe(list) when is_list(list) do
length = length(list)
sample_types =
list
|> Enum.take(@sample_limit)
|> Enum.map(&type_label/1)
"list(len: #{length}, sample_types: #{Enum.join(sample_types, ",")})"
end
def describe(tuple) when is_tuple(tuple) do
"tuple(size: #{tuple_size(tuple)})"
end
def describe(number) when is_number(number), do: inspect(number)
def describe(atom) when is_atom(atom), do: Atom.to_string(atom)
def describe(pid) when is_pid(pid), do: inspect(pid)
def describe(ref) when is_reference(ref), do: inspect(ref)
def describe(port) when is_port(port), do: inspect(port)
def describe(fun) when is_function(fun), do: "function/arity=#{:erlang.fun_info(fun)[:arity]}"
def describe(nil), do: "nil"
def describe(term) do
type_label(term)
end
defp format_key(key) when is_atom(key) do
key |> Atom.to_string() |> truncate()
end
defp format_key(key) when is_binary(key) do
key |> truncate()
end
defp format_key(key) do
key |> inspect() |> truncate()
end
defp truncate(value) do
if String.length(value) > @max_key_length do
String.slice(value, 0, @max_key_length) <> "..."
else
value
end
end
defp type_label(%struct_module{}) do
module =
struct_module
|> Atom.to_string()
|> String.trim_leading("Elixir.")
"struct(#{module})"
end
defp type_label(term) when is_map(term), do: "map"
defp type_label(term) when is_binary(term), do: "binary"
defp type_label(term) when is_list(term), do: "list"
defp type_label(term) when is_tuple(term), do: "tuple"
defp type_label(term) when is_integer(term), do: "integer"
defp type_label(term) when is_float(term), do: "float"
defp type_label(term) when is_atom(term), do: Atom.to_string(term)
defp type_label(term) when is_pid(term), do: "pid"
defp type_label(term) when is_reference(term), do: "reference"
defp type_label(term) when is_function(term), do: "function"
defp type_label(nil), do: "nil"
defp type_label(_term), do: "unknown"
end
defmodule Snakepit.Pool.ApplicationCleanup do
@moduledoc """
Provides hard guarantees for worker process cleanup when the application exits.
This module ensures that NO worker processes survive application shutdown,
preventing orphaned processes while still allowing normal pool operations.
"""
use GenServer
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.ProcessRegistry
alias Snakepit.ProcessKiller
@log_category :shutdown
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
def init(_opts) do
# Trap exits so we can cleanup before the VM dies
Process.flag(:trap_exit, true)
# Register for VM shutdown notifications
:erlang.process_flag(:priority, :high)
SLog.info(@log_category, "🛡️ Application cleanup handler started")
{:ok, %{}}
end
# Note: Worker process tracking is handled entirely by ProcessRegistry.
# ApplicationCleanup queries ProcessRegistry during shutdown for process cleanup.
# This is called when the VM is shutting down
#
# IMPORTANT: This is an EMERGENCY handler. It should rarely do actual work.
# The supervision tree (GRPCWorker.terminate + Worker.Starter + Pool) should
# clean up processes during normal shutdown.
#
# If this handler finds orphans, it indicates a bug in the supervision tree.
def terminate(reason, _state) do
SLog.info(@log_category, "🔍 Emergency cleanup check (shutdown reason: #{inspect(reason)})")
SLog.debug(
@log_category,
"ApplicationCleanup.terminate/2 called at: #{System.monotonic_time(:millisecond)}"
)
SLog.debug(@log_category, "ApplicationCleanup process info: #{inspect(Process.info(self()))}")
beam_run_id = ProcessRegistry.get_beam_run_id()
orphaned_pids = find_orphaned_processes(beam_run_id)
if Enum.empty?(orphaned_pids) do
SLog.info(@log_category, "✅ No orphaned processes - supervision tree cleaned up correctly")
emit_telemetry(:cleanup_success, 0)
else
# These are normal during test shutdown - workers that were still starting
SLog.debug(
@log_category,
"Cleanup: Found #{length(orphaned_pids)} processes still starting during shutdown"
)
SLog.debug(@log_category, "Cleanup: Orphaned PIDs: #{inspect(orphaned_pids)}")
emit_telemetry(:orphaned_processes_found, length(orphaned_pids))
# Emergency kill - use SIGKILL directly since supervision already tried SIGTERM
kill_count = emergency_kill_processes(beam_run_id)
if kill_count > 0 do
SLog.debug(@log_category, "Cleanup: Killed #{kill_count} orphaned processes")
emit_telemetry(:emergency_cleanup, kill_count)
end
end
:ok
end
defp find_orphaned_processes(run_id) do
# CRITICAL: Get all registered workers from ProcessRegistry
# A process is only "orphaned" if its Python process is alive BUT
# its Elixir GenServer is dead (supervision tree failed to clean it up)
registered_workers = ProcessRegistry.list_all_workers()
# Find Python processes for this run_id
python_pids = ProcessKiller.find_python_processes()
grpc_pids_for_run = filter_grpc_pids_by_run_id(python_pids, run_id)
# Filter out processes whose Elixir GenServer is still alive
# Those are NOT orphans - the supervision tree will clean them up
Enum.filter(grpc_pids_for_run, fn os_pid ->
orphaned_process?(os_pid, registered_workers)
end)
end
defp filter_grpc_pids_by_run_id(python_pids, run_id) do
Enum.filter(python_pids, fn pid ->
matches_grpc_run_id?(pid, run_id)
end)
end
defp matches_grpc_run_id?(pid, run_id) do
case ProcessKiller.get_process_command(pid) do
{:ok, cmd} ->
has_grpc_server = String.contains?(cmd, "grpc_server.py")
has_old_format = String.contains?(cmd, "--snakepit-run-id #{run_id}")
has_new_format = String.contains?(cmd, "--run-id #{run_id}")
has_grpc_server and (has_old_format or has_new_format)
_ ->
false
end
end
defp orphaned_process?(os_pid, registered_workers) do
worker_entry =
Enum.find(registered_workers, fn {_worker_id, info} ->
Map.get(info, :process_pid) == os_pid
end)
case worker_entry do
{_worker_id, %{elixir_pid: elixir_pid}} ->
check_if_orphan(os_pid, elixir_pid)
nil ->
# Not in registry at all - this IS an orphan
SLog.warning(@log_category, "PID #{os_pid} not in ProcessRegistry - true orphan")
true
end
end
defp check_if_orphan(os_pid, elixir_pid) do
# If the Elixir GenServer is still alive, this is NOT an orphan
# The supervision tree will clean it up - don't interfere!
is_orphan = not Process.alive?(elixir_pid)
if not is_orphan do
SLog.debug(
@log_category,
"Skipping PID #{os_pid} - Elixir GenServer #{inspect(elixir_pid)} still alive, " <>
"supervision tree will handle cleanup"
)
end
is_orphan
end
defp emergency_kill_processes(run_id) do
# Use ProcessKiller with run_id-based cleanup
{:ok, killed_count} = ProcessKiller.kill_by_run_id(run_id)
killed_count
end
defp emit_telemetry(event, count) do
:telemetry.execute(
[:snakepit, :application_cleanup, event],
%{count: count},
%{
beam_run_id: ProcessRegistry.get_beam_run_id(),
timestamp: System.system_time(:second)
}
)
end
end
defmodule Snakepit.Pool do
@moduledoc """
Pool manager for external workers with concurrent initialization.
Features:
- Concurrent worker startup (all workers start in parallel)
- Simple queue-based request distribution
- Non-blocking async execution
- Automatic request queueing when workers are busy
- Adapter-based support for any external process
"""
use GenServer
alias Snakepit.Bridge.SessionStore
alias Snakepit.Config
alias Snakepit.CrashBarrier
alias Snakepit.Defaults
alias Snakepit.Error
alias Snakepit.Logger, as: SLog
alias Snakepit.Logger.Redaction
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Pool.WorkerSupervisor
alias Snakepit.Worker.LifecycleManager
alias Snakepit.WorkerProfile.Thread.CapacityStore
@log_category :pool
# Per-pool state structure
defmodule PoolState do
@moduledoc false
defstruct [
:name,
:size,
:workers,
:available,
:worker_loads,
:worker_capacities,
:capacity_strategy,
:request_queue,
:cancelled_requests,
:stats,
:initialized,
:startup_timeout,
:queue_timeout,
:max_queue_size,
:worker_module,
:adapter_module,
:pool_config,
initialization_waiters: []
]
end
# Top-level state structure
defstruct [
:pools,
# Map of pool_name => PoolState
:affinity_cache,
# Shared across all pools
:default_pool,
# Default pool name for backward compatibility
initializing: false,
# Set to true while async initialization is in progress
init_start_time: nil
# Timestamp for measuring initialization duration
]
# Client API
@doc """
Starts the pool manager.
"""
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: opts[:name] || __MODULE__)
end
@doc """
Executes a command on any available worker.
"""
def execute(command, args, opts \\ []) do
pool = opts[:pool] || __MODULE__
timeout = opts[:timeout] || Defaults.pool_request_timeout()
# Store deadline_ms for queue-aware remaining-time calculations
deadline_ms = System.monotonic_time(:millisecond) + timeout
opts_with_deadline = Keyword.put(opts, :deadline_ms, deadline_ms)
try do
GenServer.call(pool, {:execute, command, args, opts_with_deadline}, timeout)
catch
:exit, {:timeout, _} ->
{:error,
Error.timeout_error("Pool execute timed out", %{
timeout_ms: timeout,
command: command
})}
end
end
# ============================================================================
# Timeout Helpers (Public API for deadline-aware timeout derivation)
# ============================================================================
@doc """
Returns the default timeout for a given call type.
## Call types
- `:execute` - Regular execute operations
- `:execute_stream` - Streaming operations
- `:queue` - Queue wait operations
## Examples
iex> Snakepit.Pool.get_default_timeout_for_call(:execute, %{}, [])
300_000 # from default_timeout()
iex> Snakepit.Pool.get_default_timeout_for_call(:execute, %{}, [timeout: 45_000])
45_000
"""
@spec get_default_timeout_for_call(atom(), map(), Keyword.t()) :: timeout()
def get_default_timeout_for_call(call_type, _args, opts) do
case Keyword.get(opts, :timeout) do
nil -> get_timeout_for_call_type(call_type)
explicit_timeout -> explicit_timeout
end
end
defp get_timeout_for_call_type(:execute), do: Defaults.default_timeout()
defp get_timeout_for_call_type(:execute_stream), do: Defaults.stream_timeout()
defp get_timeout_for_call_type(:queue), do: Defaults.queue_timeout()
defp get_timeout_for_call_type(_), do: Defaults.default_timeout()
@doc """
Derives the RPC timeout from opts, considering deadline if present.
When a request has been queued, time has already elapsed. This function
calculates the remaining time budget for the actual RPC call.
## Examples
# Fresh request with 60s budget
iex> Snakepit.Pool.derive_rpc_timeout_from_opts([], 60_000)
58_800 # 60_000 - 1000 - 200 margins
# Request that waited 500ms in queue
iex> now = System.monotonic_time(:millisecond)
iex> opts = [deadline_ms: now + 59_500]
iex> Snakepit.Pool.derive_rpc_timeout_from_opts(opts, 60_000)
# ~= 58_300 (remaining - margins)
"""
@spec derive_rpc_timeout_from_opts(Keyword.t(), timeout()) :: timeout()
def derive_rpc_timeout_from_opts(_opts, :infinity), do: :infinity
def derive_rpc_timeout_from_opts(opts, default_timeout) when is_integer(default_timeout) do
case Keyword.get(opts, :deadline_ms) do
nil ->
# No deadline set, use full budget
Defaults.rpc_timeout(default_timeout)
deadline_ms ->
# Calculate remaining time
now = System.monotonic_time(:millisecond)
remaining = deadline_ms - now
# Apply margins and floor
Defaults.rpc_timeout(max(remaining, 1))
end
end
@doc """
Computes effective queue timeout considering deadline.
If a deadline is set and less time remains than the configured queue timeout,
returns the remaining time instead.
## Examples
# No deadline - use configured queue timeout
iex> Snakepit.Pool.effective_queue_timeout_ms([], 10_000)
10_000
# Deadline with 5s remaining - use remaining time
iex> now = System.monotonic_time(:millisecond)
iex> opts = [deadline_ms: now + 5_000]
iex> Snakepit.Pool.effective_queue_timeout_ms(opts, 10_000)
# ~= 5_000 (remaining time)
"""
@spec effective_queue_timeout_ms(Keyword.t(), timeout()) :: non_neg_integer()
def effective_queue_timeout_ms(opts, configured_queue_timeout) do
case Keyword.get(opts, :deadline_ms) do
nil ->
configured_queue_timeout
deadline_ms ->
now = System.monotonic_time(:millisecond)
remaining = deadline_ms - now
max(remaining, 0)
end
end
@doc """
Execute a streaming command with callback.
"""
def execute_stream(command, args, callback_fn, opts \\ []) do
pool = opts[:pool] || __MODULE__
timeout = opts[:timeout] || Defaults.pool_streaming_timeout()
pool_identifier = opts[:pool_name] || pool
SLog.debug(@log_category, "[Pool] execute_stream #{command} with #{Redaction.describe(args)}")
case checkout_worker_for_stream(pool, opts) do
{:ok, worker_id} ->
SLog.debug(@log_category, "[Pool] Checked out worker #{worker_id} for streaming")
start_time = System.monotonic_time(:microsecond)
result =
try do
execute_on_worker_stream(worker_id, command, args, callback_fn, timeout)
after
# This block ALWAYS executes, preventing worker leaks on crashes
SLog.debug(
@log_category,
"[Pool] Checking in worker #{worker_id} after stream execution"
)
checkin_worker(pool, worker_id)
end
emit_stream_telemetry(pool_identifier, worker_id, command, result, start_time)
result
{:error, reason} ->
SLog.error(
@log_category,
"[Pool] Failed to checkout worker for streaming: #{inspect(reason)}"
)
{:error, reason}
end
end
defp checkout_worker_for_stream(pool, opts) do
timeout = opts[:checkout_timeout] || Defaults.checkout_timeout()
GenServer.call(pool, {:checkout_worker, opts[:session_id]}, timeout)
end
defp checkin_worker(pool, worker_id) do
GenServer.cast(pool, {:checkin_worker, worker_id})
end
defp execute_on_worker_stream(worker_id, command, args, callback_fn, timeout) do
worker_module = get_worker_module(worker_id)
SLog.debug(@log_category, "[Pool] execute_on_worker_stream using #{inspect(worker_module)}")
if function_exported?(worker_module, :execute_stream, 5) do
SLog.debug(
@log_category,
"[Pool] Invoking #{worker_module}.execute_stream with timeout #{timeout}"
)
result = worker_module.execute_stream(worker_id, command, args, callback_fn, timeout)
SLog.debug(@log_category, "[Pool] execute_stream result: #{Redaction.describe(result)}")
result
else
SLog.error(
@log_category,
"[Pool] Worker module #{worker_module} does not export execute_stream/5"
)
{:error,
Error.worker_error("Streaming not supported by worker", %{
worker_module: worker_module,
worker_id: worker_id
})}
end
end
defp emit_stream_telemetry(pool_identifier, worker_id, command, result, start_time) do
duration_us = System.monotonic_time(:microsecond) - start_time
:telemetry.execute(
[:snakepit, :request, :executed],
%{duration_us: duration_us},
%{
pool: pool_identifier,
worker_id: worker_id,
command: command,
success: result == :ok,
streaming: true
}
)
end
@doc """
Gets pool statistics.
"""
def get_stats(pool \\ __MODULE__) do
GenServer.call(pool, :get_stats)
end
@doc """
Gets statistics for a specific pool name.
"""
def get_stats(pool, pool_name) when is_atom(pool_name) do
GenServer.call(pool, {:get_stats, pool_name})
end
@doc """
Lists all worker IDs in the pool.
Can be called with pool process or pool name:
- `list_workers()` - all workers from all pools
- `list_workers(Snakepit.Pool)` - all workers from all pools
- `list_workers(Snakepit.Pool, :pool_name)` - workers from specific pool
"""
def list_workers(pool \\ __MODULE__)
def list_workers(pool) when is_pid(pool) or is_atom(pool) do
# Call with just pool process (backward compat)
GenServer.call(pool, :list_workers)
end
def list_workers(pool, pool_name) when is_atom(pool_name) do
# Call with pool process and pool name
GenServer.call(pool, {:list_workers, pool_name})
end
@doc """
Waits for the pool to be fully initialized.
Returns `:ok` when all workers are ready, or `{:error, %Snakepit.Error{}}` if
the pool doesn't initialize within the given timeout.
"""
@spec await_ready(atom() | pid(), timeout() | nil) :: :ok | {:error, Error.t()}
def await_ready(pool \\ __MODULE__, timeout \\ nil)
def await_ready(pool, nil), do: await_ready(pool, Defaults.pool_await_ready_timeout())
def await_ready(pool, timeout) do
GenServer.call(pool, :await_ready, timeout)
catch
:exit, {:timeout, _} ->
{:error,
Error.timeout_error("Pool initialization timed out", %{
pool: pool,
timeout_ms: timeout
})}
end
# Server Callbacks
@impl true
def init(opts) do
# CRITICAL: Trap exits to ensure terminate/2 is called
Process.flag(:trap_exit, true)
case resolve_pool_configs() do
{:ok, pool_configs} ->
init_with_configs(opts, pool_configs)
{:error, reason} ->
{:stop, reason}
end
end
defp init_with_configs(opts, pool_configs) do
# PERFORMANCE FIX: Create ETS cache for session affinity to eliminate
# GenServer bottleneck on SessionStore. This provides ~100x faster lookups.
# Shared across ALL pools
affinity_cache =
:ets.new(:worker_affinity_cache, [
:set,
:public,
{:read_concurrency, true}
])
# Check if we're in multi-pool mode (explicit :pools config)
multi_pool_mode? = Application.get_env(:snakepit, :pools) != nil
# Create initial pool states (not yet initialized with workers)
pools =
Enum.reduce(pool_configs, %{}, fn pool_config, acc ->
pool_state = build_pool_state(opts, pool_config, multi_pool_mode?)
Map.put(acc, pool_state.name, pool_state)
end)
# Determine default pool (first pool or :default)
default_pool =
case pool_configs do
[first | _] -> Map.get(first, :name, :default)
[] -> :default
end
state = %__MODULE__{
pools: pools,
affinity_cache: affinity_cache,
default_pool: default_pool
}
# Start concurrent worker initialization for ALL pools
{:ok, state, {:continue, :initialize_workers}}
end
defp build_pool_state(opts, pool_config, multi_pool_mode?) do
pool_name = Map.get(pool_config, :name, :default)
# Extract pool settings with backward-compatible fallbacks
# In multi-pool mode, ALWAYS use per-pool pool_size (ignore opts[:size] from legacy config)
# In legacy mode, use opts[:size] (from application.ex) or fall back to pool_config/default
size = resolve_pool_size(opts, pool_config, multi_pool_mode?)
startup_timeout = Defaults.pool_startup_timeout()
queue_timeout = Defaults.pool_queue_timeout()
max_queue_size = Defaults.pool_max_queue_size()
worker_module = opts[:worker_module] || Snakepit.GRPCWorker
adapter_module =
opts[:adapter_module] ||
Map.get(pool_config, :adapter_module) ||
Application.get_env(:snakepit, :adapter_module)
%PoolState{
name: pool_name,
size: size,
workers: [],
available: MapSet.new(),
worker_loads: %{},
worker_capacities: %{},
capacity_strategy: resolve_capacity_strategy(pool_config),
request_queue: :queue.new(),
cancelled_requests: %{},
stats: %{
requests: 0,
queued: 0,
errors: 0,
queue_timeouts: 0,
pool_saturated: 0
},
initialized: false,
startup_timeout: startup_timeout,
queue_timeout: queue_timeout,
max_queue_size: max_queue_size,
worker_module: worker_module,
adapter_module: adapter_module,
pool_config: pool_config
}
end
defp resolve_pool_size(_opts, pool_config, true = _multi_pool_mode?) do
# Multi-pool mode: use per-pool config only
Map.get(pool_config, :pool_size, Defaults.default_pool_size())
end
defp resolve_pool_size(opts, pool_config, false = _multi_pool_mode?) do
# Legacy mode: opts[:size] takes precedence
opts[:size] || Map.get(pool_config, :pool_size, Defaults.default_pool_size())
end
@impl true
def handle_continue(:initialize_workers, state) do
total_workers = Enum.reduce(state.pools, 0, fn {_name, pool}, acc -> acc + pool.size end)
SLog.info(
@log_category,
"🚀 Starting concurrent initialization of #{total_workers} workers across #{map_size(state.pools)} pool(s)..."
)
start_time = System.monotonic_time(:millisecond)
# DIAGNOSTIC: Capture baseline system resource usage
baseline_resources = capture_resource_metrics()
SLog.info(@log_category, "📊 Baseline resources: #{inspect(baseline_resources)}")
# OTP FIX: Spawn a linked process to do blocking initialization
# This keeps the GenServer responsive to shutdown signals during batch startup
# The linked process will be killed automatically when the GenServer terminates
parent = self()
pools_data = state.pools
# CRITICAL: Capture the Pool GenServer's name BEFORE spawning, so workers
# get the correct pool reference (not the ephemeral init process's pid)
pool_genserver_name = get_pool_genserver_name()
spawn_link(fn ->
updated_pools = do_pool_initialization(pools_data, baseline_resources, pool_genserver_name)
send(parent, {:pool_init_complete, updated_pools})
end)
# Return immediately - GenServer is now responsive to shutdown signals
{:noreply, %{state | initializing: true, init_start_time: start_time}}
end
# Helper to perform blocking pool initialization in a separate process
# pool_genserver_name is the registered name of the Pool GenServer (captured before spawn)
defp do_pool_initialization(pools_data, _baseline_resources, pool_genserver_name) do
Enum.map(pools_data, fn {pool_name, pool_state} ->
SLog.info(
@log_category,
"Initializing pool #{pool_name} with #{pool_state.size} workers..."
)
# Start workers for this pool (may include blocking batch delays)
workers =
start_workers_concurrently(
pool_name,
pool_state.size,
pool_state.startup_timeout,
pool_state.worker_module,
pool_state.adapter_module,
pool_state.pool_config,
pool_genserver_name
)
SLog.info(
@log_category,
"✅ Pool #{pool_name}: Initialized #{length(workers)}/#{pool_state.size} workers"
)
# Update pool state with workers
updated_pool_state =
if Enum.empty?(workers) do
# Pool failed to start any workers
SLog.error(@log_category, "❌ Pool #{pool_name} failed to start any workers!")
pool_state
else
worker_capacities = build_worker_capacities(pool_state, workers)
available = MapSet.new(workers)
%{
pool_state
| workers: workers,
available: available,
worker_capacities: worker_capacities,
worker_loads: %{},
initialized: true
}
end
{pool_name, updated_pool_state}
end)
|> Enum.into(%{})
end
# Handle completion of async pool initialization
@impl true
def handle_info({:pool_init_complete, updated_pools}, state) do
elapsed = System.monotonic_time(:millisecond) - (state.init_start_time || 0)
# DIAGNOSTIC: Capture peak system resource usage after startup
peak_resources = capture_resource_metrics()
baseline_resources = capture_resource_metrics()
resource_delta = calculate_resource_delta(baseline_resources, peak_resources)
SLog.info(@log_category, "✅ All pools initialized in #{elapsed}ms")
SLog.info(@log_category, "📊 Resource usage delta: #{inspect(resource_delta)}")
failed_pools =
Enum.filter(updated_pools, fn {_name, pool} -> Enum.empty?(pool.workers) end)
failed_pool_names = Enum.map(failed_pools, &elem(&1, 0))
if not Enum.empty?(failed_pools) do
SLog.warning(
@log_category,
"⚠️ Pools with zero initialized workers: #{Enum.map_join(failed_pool_names, ", ", &to_string/1)}"
)
end
# Check if any pool successfully started
any_workers_started? =
Enum.any?(updated_pools, fn {_name, pool} -> not Enum.empty?(pool.workers) end)
if any_workers_started? do
# CRITICAL: Get waiters from state.pools (GenServer state), NOT from updated_pools
# (which was returned from the spawned init process and doesn't have waiters)
all_waiters =
Enum.flat_map(state.pools, fn {_name, pool} ->
pool.initialization_waiters
end)
# PERFORMANCE FIX: Stagger replies to waiters from ALL pools
all_waiters
|> Enum.with_index()
|> Enum.each(fn {from, index} ->
# Stagger each reply by 2ms to spread the load
Process.send_after(self(), {:reply_to_waiter, from}, index * 2)
end)
# Merge updated_pools (workers, capacities, initialized status) with cleared waiters
merged_pools =
Enum.map(updated_pools, fn {name, updated_pool} ->
{name, %{updated_pool | initialization_waiters: []}}
end)
|> Enum.into(%{})
new_state = %{state | pools: merged_pools, initializing: false, init_start_time: nil}
{:noreply, new_state}
else
SLog.error(
@log_category,
"❌ All configured pools failed to start workers (#{Enum.map_join(failed_pool_names, ", ", &to_string/1)})."
)
{:stop, :no_workers_started, state}
end
end
# Handle EXIT from the init process during async initialization
# This occurs when the GenServer is shutting down while initialization is in progress
@impl true
def handle_info({:EXIT, _pid, reason}, %{initializing: true} = state) do
case reason do
:normal ->
# Init process completed normally but we haven't received the result yet
# This shouldn't happen in normal operation
{:noreply, state}
:shutdown ->
# Clean shutdown - supervisor is terminating
SLog.info(@log_category, "Pool initialization interrupted by shutdown")
{:stop, :shutdown, state}
{:shutdown, _} ->
# Clean shutdown with reason
SLog.info(@log_category, "Pool initialization interrupted by shutdown")
{:stop, :shutdown, state}
:killed ->
# Init process was killed (e.g., supervisor timeout)
SLog.warning(@log_category, "Pool initialization process killed")
{:stop, :killed, state}
other ->
# Init process crashed
SLog.error(@log_category, "Pool initialization failed: #{inspect(other)}")
{:stop, {:init_failed, other}, state}
end
end
# Handle EXIT from other linked processes (workers, etc.) - passthrough
def handle_info({:EXIT, _pid, _reason}, state) do
# Let the default behavior handle this
{:noreply, state}
end
# queue_timeout - WITH pool_name
def handle_info({:queue_timeout, pool_name, from}, state) do
case Map.get(state.pools, pool_name) do
nil ->
{:noreply, state}
pool_state ->
now = System.monotonic_time(:millisecond)
retention_ms = cancellation_retention_ms(pool_state.queue_timeout)
{pruned_queue, dropped?} =
drop_request_from_queue(pool_state.request_queue, from)
if dropped? do
GenServer.reply(from, {:error, :queue_timeout})
SLog.debug(
@log_category,
"Removed timed out request #{inspect(from)} from queue in pool #{pool_name}"
)
new_cancelled =
record_cancelled_request(pool_state.cancelled_requests, from, now, retention_ms)
updated_stats = Map.update!(pool_state.stats, :queue_timeouts, &(&1 + 1))
updated_pool_state = %{
pool_state
| request_queue: pruned_queue,
cancelled_requests: new_cancelled,
stats: updated_stats
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
else
SLog.debug(
@log_category,
"Queue timeout fired for #{inspect(from)} in pool #{pool_name} after request was already handled"
)
{:noreply, state}
end
end
end
# Legacy queue_timeout WITHOUT pool_name
def handle_info({:queue_timeout, from}, state) do
handle_info({:queue_timeout, state.default_pool, from}, state)
end
# Worker death - find which pool it belongs to
def handle_info({:DOWN, _ref, :process, pid, reason}, state) do
case Snakepit.Pool.Registry.get_worker_id_by_pid(pid) do
{:error, :not_found} ->
{:noreply, state}
{:ok, worker_id} ->
handle_worker_down(worker_id, pid, reason, state)
end
end
def handle_info({:reply_to_waiter, from}, state) do
# PERFORMANCE FIX: Staggered reply to avoid thundering herd
GenServer.reply(from, :ok)
{:noreply, state}
end
@doc false
# Handles completion messages from tasks started via Task.Supervisor.async_nolink.
# These are used for fire-and-forget operations (like replying to callers or
# kicking off worker respawns), so we can safely ignore the completion message.
def handle_info({ref, _result}, state) when is_reference(ref) do
{:noreply, state}
end
def handle_info(msg, state) do
SLog.debug(@log_category, "Pool received unexpected message: #{inspect(msg)}")
{:noreply, state}
end
@impl true
def handle_call({:execute, command, args, opts}, from, state) do
# Backward compatibility: route to default pool
pool_name = opts[:pool_name] || state.default_pool
handle_execute_for_pool(pool_name, command, args, opts, from, state)
end
def handle_call({:execute, pool_name, command, args, opts}, from, state) do
handle_execute_for_pool(pool_name, command, args, opts, from, state)
end
def handle_call({:checkout_worker, session_id}, _from, state) do
# Backward compat: use default pool
pool_name = state.default_pool
case Map.get(state.pools, pool_name) do
nil ->
{:reply, {:error, :pool_not_found}, state}
pool_state ->
# Pass affinity_cache from top-level state
case checkout_worker(pool_state, session_id, state.affinity_cache) do
{:ok, worker_id, new_pool_state} ->
updated_pools = Map.put(state.pools, pool_name, new_pool_state)
{:reply, {:ok, worker_id}, %{state | pools: updated_pools}}
{:error, :no_workers} ->
{:reply, {:error, :no_workers_available}, state}
end
end
end
def handle_call({:checkout_worker, pool_name, session_id}, _from, state)
when is_atom(pool_name) do
case Map.get(state.pools, pool_name) do
nil ->
{:reply, {:error, :pool_not_found}, state}
pool_state ->
case checkout_worker(pool_state, session_id, state.affinity_cache) do
{:ok, worker_id, new_pool_state} ->
updated_pools = Map.put(state.pools, pool_name, new_pool_state)
{:reply, {:ok, worker_id}, %{state | pools: updated_pools}}
{:error, :no_workers} ->
{:reply, {:error, :no_workers_available}, state}
end
end
end
def handle_call(:get_stats, _from, state) do
# Backward compat: aggregate stats from all pools
aggregate_stats =
Enum.reduce(
state.pools,
%{
requests: 0,
queued: 0,
errors: 0,
queue_timeouts: 0,
pool_saturated: 0,
workers: 0,
available: 0,
busy: 0
},
fn {_name, pool}, acc ->
%{
requests: acc.requests + Map.get(pool.stats, :requests, 0),
queued: acc.queued + Map.get(pool.stats, :queued, 0) + :queue.len(pool.request_queue),
errors: acc.errors + Map.get(pool.stats, :errors, 0),
queue_timeouts: acc.queue_timeouts + Map.get(pool.stats, :queue_timeouts, 0),
pool_saturated: acc.pool_saturated + Map.get(pool.stats, :pool_saturated, 0),
workers: acc.workers + length(pool.workers),
available: acc.available + MapSet.size(pool.available),
busy: acc.busy + busy_worker_count(pool)
}
end
)
{:reply, aggregate_stats, state}
end
def handle_call({:get_stats, pool_name}, _from, state) do
case Map.get(state.pools, pool_name) do
nil ->
{:reply, {:error, :pool_not_found}, state}
pool_state ->
stats =
Map.merge(pool_state.stats, %{
workers: length(pool_state.workers),
available: MapSet.size(pool_state.available),
busy: busy_worker_count(pool_state),
queued: :queue.len(pool_state.request_queue)
})
{:reply, stats, state}
end
end
def handle_call(:list_workers, _from, state) do
# Backward compat: return workers from ALL pools
all_workers =
Enum.flat_map(state.pools, fn {_name, pool} ->
pool.workers
end)
{:reply, all_workers, state}
end
def handle_call({:list_workers, pool_name}, _from, state) do
case Map.get(state.pools, pool_name) do
nil ->
{:reply, {:error, :pool_not_found}, state}
pool_state ->
{:reply, pool_state.workers, state}
end
end
@impl true
def handle_call(:await_ready, from, state) do
if all_pools_initialized?(state.pools) do
{:reply, :ok, state}
else
updated_pools = add_waiter_to_uninitialized_pools(state.pools, from)
{:noreply, %{state | pools: updated_pools}}
end
end
def handle_call({:await_ready, pool_name}, from, state) do
case Map.get(state.pools, pool_name) do
nil ->
{:reply, {:error, :pool_not_found}, state}
pool_state ->
if pool_state.initialized do
{:reply, :ok, state}
else
# Add waiter to this pool
updated_pool = %{
pool_state
| initialization_waiters: [from | pool_state.initialization_waiters]
}
updated_pools = Map.put(state.pools, pool_name, updated_pool)
{:noreply, %{state | pools: updated_pools}}
end
end
end
def handle_call({:worker_ready, worker_id}, _from, state) do
SLog.info(@log_category, "Worker #{worker_id} reported ready. Processing queued work.")
# Find which pool this worker belongs to by worker_id prefix
pool_name = extract_pool_name_from_worker_id(worker_id)
case Map.get(state.pools, pool_name) do
nil ->
SLog.error(
@log_category,
"Worker #{worker_id} reported ready but pool #{pool_name} not found!"
)
{:reply, {:error, :pool_not_found}, state}
pool_state ->
# Ensure worker is in workers list
new_workers =
if Enum.member?(pool_state.workers, worker_id) do
pool_state.workers
else
[worker_id | pool_state.workers]
end
updated_pool_state =
pool_state
|> Map.put(:workers, new_workers)
|> ensure_worker_capacity(worker_id)
|> ensure_worker_available(worker_id)
# CRITICAL FIX: Immediately drive the queue by treating this as a checkin
GenServer.cast(self(), {:checkin_worker, pool_name, worker_id, :skip_decrement})
CrashBarrier.maybe_emit_restart(pool_name, worker_id)
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:reply, :ok, %{state | pools: updated_pools}}
end
end
defp all_pools_initialized?(pools) do
Enum.all?(pools, fn {_name, pool} -> pool.initialized end)
end
defp add_waiter_to_uninitialized_pools(pools, from) do
Enum.map(pools, fn {name, pool} ->
if pool.initialized do
{name, pool}
else
{name, %{pool | initialization_waiters: [from | pool.initialization_waiters]}}
end
end)
|> Enum.into(%{})
end
defp handle_execute_for_pool(pool_name, command, args, opts, from, state) do
with {:ok, pool_state} <- get_pool(state.pools, pool_name),
:ok <- check_pool_initialized(pool_state) do
handle_execute_in_pool(pool_name, pool_state, command, args, opts, from, state)
else
{:error, :pool_not_found} ->
{:reply, {:error, {:pool_not_found, pool_name}}, state}
{:error, :pool_not_initialized} ->
{:reply, {:error, :pool_not_initialized}, state}
end
end
defp get_pool(pools, pool_name) do
case Map.get(pools, pool_name) do
nil -> {:error, :pool_not_found}
pool_state -> {:ok, pool_state}
end
end
defp check_pool_initialized(pool_state) do
if pool_state.initialized, do: :ok, else: {:error, :pool_not_initialized}
end
defp handle_execute_in_pool(pool_name, pool_state, command, args, opts, from, state) do
pool_state = compact_pool_queue(pool_state)
session_id = opts[:session_id]
case checkout_worker(pool_state, session_id, state.affinity_cache) do
{:ok, worker_id, new_pool_state} ->
execute_with_worker(
pool_name,
worker_id,
new_pool_state,
command,
args,
opts,
from,
state
)
{:error, :no_workers} ->
handle_no_workers_available(pool_name, pool_state, command, args, opts, from, state)
end
end
defp execute_with_worker(pool_name, worker_id, new_pool_state, command, args, opts, from, state) do
spawn_execution_task(
pool_name,
worker_id,
command,
args,
opts,
from,
new_pool_state.pool_config
)
updated_pool_state = %{
new_pool_state
| stats: Map.update!(new_pool_state.stats, :requests, &(&1 + 1))
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
defp spawn_execution_task(pool_name, worker_id, command, args, opts, from, pool_config) do
async_with_context(fn ->
{client_pid, _tag} = from
ref = Process.monitor(client_pid)
start_time = System.monotonic_time(:microsecond)
case monitor_client_status(ref, client_pid) do
{:down, reason} ->
handle_client_already_down(pool_name, worker_id, command, ref, start_time, reason)
:alive ->
exec_ctx = %{
pool_pid: self(),
pool_name: pool_name,
worker_id: worker_id,
command: command,
args: args,
opts: opts,
from: from,
ref: ref,
client_pid: client_pid,
pool_config: pool_config
}
handle_client_alive(exec_ctx, start_time)
end
end)
end
defp handle_client_already_down(pool_name, worker_id, command, ref, start_time, reason) do
Process.demonitor(ref, [:flush])
duration_us = System.monotonic_time(:microsecond) - start_time
:telemetry.execute(
[:snakepit, :request, :executed],
%{duration_us: duration_us},
%{
pool: pool_name,
worker_id: worker_id,
command: command,
success: false,
aborted: true,
reason: :client_down,
client_down_reason: reason
}
)
SLog.debug(@log_category, "Client was already down; skipping work on worker #{worker_id}")
GenServer.cast(__MODULE__, {:checkin_worker, pool_name, worker_id})
end
defp handle_client_alive(exec_ctx, start_time) do
%{
pool_pid: pool_pid,
pool_name: pool_name,
worker_id: worker_id,
command: command,
args: args,
opts: opts,
from: from,
ref: ref,
client_pid: client_pid,
pool_config: pool_config
} = exec_ctx
{result, final_worker_id} =
execute_with_crash_barrier(
pool_pid,
pool_name,
worker_id,
command,
args,
opts,
pool_config
)
telemetry_worker_id = final_worker_id || worker_id
duration_us = System.monotonic_time(:microsecond) - start_time
:telemetry.execute(
[:snakepit, :request, :executed],
%{duration_us: duration_us},
%{
pool: pool_name,
worker_id: telemetry_worker_id,
command: command,
success: match?({:ok, _}, result)
}
)
handle_client_reply(pool_name, final_worker_id, from, ref, client_pid, result)
end
defp handle_client_reply(pool_name, worker_id, from, ref, client_pid, result) do
receive do
{:DOWN, ^ref, :process, ^client_pid, _reason} ->
SLog.warning(
@log_category,
"Client #{inspect(client_pid)} died before receiving reply. " <>
"Checking in worker #{inspect(worker_id)}."
)
maybe_checkin_worker(pool_name, worker_id)
after
0 ->
Process.demonitor(ref, [:flush])
GenServer.reply(from, result)
maybe_checkin_worker(pool_name, worker_id)
end
end
defp maybe_checkin_worker(_pool_name, nil), do: :ok
defp maybe_checkin_worker(pool_name, worker_id) do
GenServer.cast(__MODULE__, {:checkin_worker, pool_name, worker_id})
end
defp handle_no_workers_available(pool_name, pool_state, command, args, opts, from, state) do
current_queue_size = :queue.len(pool_state.request_queue)
if current_queue_size >= pool_state.max_queue_size do
handle_pool_saturated(pool_name, pool_state, current_queue_size, state)
else
queue_request(pool_name, pool_state, command, args, opts, from, state)
end
end
defp handle_pool_saturated(pool_name, pool_state, current_queue_size, state) do
updated_pool_state = %{
pool_state
| stats: Map.update!(pool_state.stats, :pool_saturated, &(&1 + 1))
}
:telemetry.execute(
[:snakepit, :pool, :saturated],
%{queue_size: current_queue_size, max_queue_size: pool_state.max_queue_size},
%{
pool: pool_name,
available_workers: MapSet.size(pool_state.available),
busy_workers: busy_worker_count(pool_state)
}
)
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:reply, {:error, :pool_saturated}, %{state | pools: updated_pools}}
end
defp queue_request(pool_name, pool_state, command, args, opts, from, state) do
timer_ref =
Process.send_after(self(), {:queue_timeout, pool_name, from}, pool_state.queue_timeout)
request = {from, command, args, opts, System.monotonic_time(), timer_ref}
new_queue = :queue.in(request, pool_state.request_queue)
updated_stats =
pool_state.stats
|> Map.update!(:requests, &(&1 + 1))
|> Map.update!(:queued, &(&1 + 1))
updated_pool_state = %{
pool_state
| request_queue: new_queue,
stats: updated_stats
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
# checkin_worker - WITH pool_name parameter
@impl true
def handle_cast({:checkin_worker, pool_name, worker_id, :skip_decrement}, state)
when is_atom(pool_name) do
do_handle_checkin(pool_name, worker_id, state, false)
end
def handle_cast({:checkin_worker, pool_name, worker_id}, state) when is_atom(pool_name) do
do_handle_checkin(pool_name, worker_id, state, true)
end
# Legacy checkin_worker WITHOUT pool_name (infer from worker_id)
def handle_cast({:checkin_worker, worker_id}, state) do
pool_name = extract_pool_name_from_worker_id(worker_id)
handle_cast({:checkin_worker, pool_name, worker_id}, state)
end
defp do_handle_checkin(pool_name, worker_id, state, decrement?) do
case Map.get(state.pools, pool_name) do
nil ->
SLog.error(@log_category, "checkin_worker: pool #{pool_name} not found!")
{:noreply, state}
pool_state ->
process_checkin(pool_name, worker_id, pool_state, state, decrement?)
end
end
defp process_checkin(pool_name, worker_id, pool_state, state, decrement?) do
now = System.monotonic_time(:millisecond)
retention_ms = cancellation_retention_ms(pool_state.queue_timeout)
pruned_cancelled = prune_cancelled_requests(pool_state.cancelled_requests, now, retention_ms)
pool_state =
if decrement? do
decrement_worker_load(pool_state, worker_id)
else
pool_state
end
process_next_queued_request(pool_name, worker_id, pool_state, pruned_cancelled, state)
end
defp process_next_queued_request(pool_name, worker_id, pool_state, pruned_cancelled, state) do
case select_queue_worker(pool_state, worker_id) do
{:ok, queue_worker} ->
case :queue.out(pool_state.request_queue) do
{{:value, request}, new_queue} ->
handle_queued_request(
pool_name,
queue_worker,
pool_state,
request,
new_queue,
pruned_cancelled,
state
)
{:empty, _} ->
updated_pool_state = %{pool_state | cancelled_requests: pruned_cancelled}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
:no_worker ->
updated_pool_state = %{pool_state | cancelled_requests: pruned_cancelled}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
end
defp handle_queued_request(
pool_name,
worker_id,
pool_state,
{queued_from, command, args, opts, _queued_at, timer_ref},
new_queue,
pruned_cancelled,
state
) do
cancel_queue_timer(timer_ref)
ctx = %{
pool_name: pool_name,
worker_id: worker_id,
pool_state: pool_state,
queued_from: queued_from,
new_queue: new_queue,
pruned_cancelled: pruned_cancelled,
state: state
}
if Map.has_key?(pruned_cancelled, queued_from) do
handle_cancelled_request(ctx)
else
handle_valid_request(ctx, command, args, opts)
end
end
defp handle_cancelled_request(ctx) do
%{
pool_name: pool_name,
worker_id: worker_id,
pool_state: pool_state,
queued_from: queued_from,
new_queue: new_queue,
pruned_cancelled: pruned_cancelled,
state: state
} = ctx
SLog.debug(@log_category, "Skipping cancelled request from #{inspect(queued_from)}")
new_cancelled = drop_cancelled_request(pruned_cancelled, queued_from)
updated_pool_state = %{
pool_state
| request_queue: new_queue,
cancelled_requests: new_cancelled
}
GenServer.cast(self(), {:checkin_worker, pool_name, worker_id, :skip_decrement})
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
defp handle_valid_request(ctx, command, args, opts) do
%{queued_from: queued_from} = ctx
{client_pid, _tag} = queued_from
if Process.alive?(client_pid) do
execute_queued_request(ctx, client_pid, command, args, opts)
else
handle_dead_client(ctx, client_pid)
end
end
defp execute_queued_request(ctx, client_pid, command, args, opts) do
%{
pool_name: pool_name,
worker_id: worker_id,
queued_from: queued_from,
pool_state: pool_state,
new_queue: new_queue,
pruned_cancelled: pruned_cancelled,
state: state
} = ctx
pool_state = increment_worker_load(pool_state, worker_id, nil)
pool_pid = self()
async_with_context(fn ->
ref = Process.monitor(client_pid)
{result, final_worker_id} =
execute_with_crash_barrier(
pool_pid,
pool_name,
worker_id,
command,
args,
opts,
pool_state.pool_config
)
checkin_worker_id = final_worker_id
receive do
{:DOWN, ^ref, :process, ^client_pid, _reason} ->
SLog.warning(
@log_category,
"Queued client #{inspect(client_pid)} died during execution."
)
maybe_checkin_worker(pool_name, checkin_worker_id)
after
0 ->
Process.demonitor(ref, [:flush])
GenServer.reply(queued_from, result)
maybe_checkin_worker(pool_name, checkin_worker_id)
end
end)
updated_pool_state = %{
pool_state
| request_queue: new_queue,
cancelled_requests: pruned_cancelled
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
defp handle_dead_client(ctx, client_pid) do
%{
pool_name: pool_name,
worker_id: worker_id,
pool_state: pool_state,
new_queue: new_queue,
pruned_cancelled: pruned_cancelled,
state: state
} = ctx
SLog.debug(@log_category, "Discarding request from dead client #{inspect(client_pid)}")
GenServer.cast(self(), {:checkin_worker, pool_name, worker_id, :skip_decrement})
updated_pool_state = %{
pool_state
| request_queue: new_queue,
cancelled_requests: pruned_cancelled
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
{:noreply, %{state | pools: updated_pools}}
end
defp handle_worker_down(worker_id, pid, reason, state) do
SLog.error(
@log_category,
"Worker #{worker_id} (pid: #{inspect(pid)}) died: #{inspect(reason)}"
)
:ets.match_delete(state.affinity_cache, {:_, worker_id, :_})
pool_name = extract_pool_name_from_worker_id(worker_id)
case Map.get(state.pools, pool_name) do
nil ->
SLog.warning(
@log_category,
"Dead worker #{worker_id} belongs to unknown pool #{pool_name}"
)
{:noreply, state}
pool_state ->
maybe_taint_on_crash(pool_name, worker_id, reason, pool_state)
updated_state =
state
|> remove_worker_from_pool(pool_name, pool_state, worker_id)
|> tap(fn _ ->
SLog.debug(@log_category, "Removed dead worker #{worker_id} from pool #{pool_name}")
end)
{:noreply, updated_state}
end
end
defp maybe_taint_on_crash(pool_name, worker_id, reason, pool_state) do
crash_config = CrashBarrier.config(pool_state.pool_config)
with true <- CrashBarrier.enabled?(crash_config),
{:ok, info} <- CrashBarrier.crash_info({:error, {:worker_exit, reason}}, crash_config) do
maybe_taint_worker(pool_name, worker_id, info, crash_config)
else
_ -> :ok
end
end
defp remove_worker_from_pool(state, pool_name, pool_state, worker_id) do
new_workers = List.delete(pool_state.workers, worker_id)
new_available = MapSet.delete(pool_state.available, worker_id)
new_loads = Map.delete(pool_state.worker_loads, worker_id)
new_capacities = Map.delete(pool_state.worker_capacities, worker_id)
updated_pool_state = %{
pool_state
| workers: new_workers,
available: new_available,
worker_loads: new_loads,
worker_capacities: new_capacities
}
updated_pools = Map.put(state.pools, pool_name, updated_pool_state)
%{state | pools: updated_pools}
end
@impl true
def terminate(reason, state) do
SLog.info(@log_category, "🛑 Pool manager terminating with reason: #{inspect(reason)}.")
# Log state of pools during shutdown (debug level)
Enum.each(state.pools, fn {pool_name, pool_state} ->
SLog.debug(@log_category, """
Pool #{pool_name} shutdown state:
Initialized: #{pool_state.initialized}
Workers: #{length(pool_state.workers)}
Available: #{MapSet.size(pool_state.available)}
Busy: #{busy_worker_count(pool_state)}
Queued: #{:queue.len(pool_state.request_queue)}
Waiters: #{length(pool_state.initialization_waiters)}
""")
if not Enum.empty?(pool_state.initialization_waiters) do
SLog.warning(
@log_category,
"Pool #{pool_name}: #{length(pool_state.initialization_waiters)} processes still waiting for pool init!"
)
end
end)
# Supervision tree will handle worker shutdown via WorkerSupervisor
:ok
end
# REMOVE the wait_for_ports_to_exit/2 helper functions.
# Private Functions
defp compact_pool_queue(pool_state) do
now_ms = System.monotonic_time(:millisecond)
retention_ms = cancellation_retention_ms(pool_state.queue_timeout)
{new_queue, new_cancelled} =
compact_request_queue(
pool_state.request_queue,
pool_state.cancelled_requests,
now_ms,
retention_ms
)
%{pool_state | request_queue: new_queue, cancelled_requests: new_cancelled}
end
defp compact_request_queue(queue, cancelled_requests, now_ms, retention_ms) do
pruned_cancelled = prune_cancelled_requests(cancelled_requests, now_ms, retention_ms)
{filtered, updated_cancelled} =
queue
|> :queue.to_list()
|> Enum.reduce({[], pruned_cancelled}, fn
{from, _command, _args, _opts, _queued_at, timer_ref} = request,
{acc, current_cancelled} ->
cond do
Map.has_key?(current_cancelled, from) ->
cancel_queue_timer(timer_ref)
{acc, drop_cancelled_request(current_cancelled, from)}
not alive_from?(from) ->
cancel_queue_timer(timer_ref)
{acc, drop_cancelled_request(current_cancelled, from)}
true ->
{[request | acc], current_cancelled}
end
end)
new_queue =
filtered
|> Enum.reverse()
|> :queue.from_list()
{new_queue, updated_cancelled}
end
defp drop_request_from_queue(queue, from) do
{remaining, dropped?} =
queue
|> :queue.to_list()
|> Enum.reduce({[], false}, fn
{queued_from, _command, _args, _opts, _queued_at, timer_ref}, {acc, _}
when queued_from == from ->
cancel_queue_timer(timer_ref)
{acc, true}
request, {acc, dropped?} ->
{[request | acc], dropped?}
end)
new_queue =
remaining
|> Enum.reverse()
|> :queue.from_list()
{new_queue, dropped?}
end
defp alive_from?({pid, _ref}) when is_pid(pid), do: Process.alive?(pid)
defp alive_from?(_), do: false
defp cancellation_retention_ms(queue_timeout)
when is_integer(queue_timeout) and queue_timeout > 0 do
retention = queue_timeout * Defaults.pool_cancelled_retention_multiplier()
max(retention, queue_timeout)
end
defp cancellation_retention_ms(_queue_timeout) do
Defaults.pool_queue_timeout() * Defaults.pool_cancelled_retention_multiplier()
end
defp prune_cancelled_requests(cancelled_requests, _now_ms, _retention_ms)
when cancelled_requests == %{} do
cancelled_requests
end
defp prune_cancelled_requests(cancelled_requests, now_ms, retention_ms) do
cutoff = now_ms - retention_ms
cancelled_requests
|> Enum.reject(fn {_from, recorded_at} -> recorded_at < cutoff end)
|> Map.new()
end
defp record_cancelled_request(cancelled_requests, from, now_ms, retention_ms) do
cancelled_requests
|> prune_cancelled_requests(now_ms, retention_ms)
|> Map.put(from, now_ms)
|> trim_cancelled_requests()
end
defp drop_cancelled_request(cancelled_requests, from) do
Map.delete(cancelled_requests, from)
end
defp cancel_queue_timer(nil), do: :ok
defp cancel_queue_timer(timer_ref) do
Process.cancel_timer(timer_ref, async: true, info: false)
:ok
end
defp trim_cancelled_requests(cancelled_requests) do
max_entries = Defaults.pool_max_cancelled_entries()
if map_size(cancelled_requests) <= max_entries do
cancelled_requests
else
entries_to_keep = max_entries
drop_count = map_size(cancelled_requests) - entries_to_keep
cancelled_requests
|> Enum.sort_by(fn {_from, recorded_at} -> recorded_at end)
|> Enum.drop(drop_count)
|> Map.new()
end
end
defp resolve_pool_configs do
case Config.get_pool_configs() do
{:ok, configs} when is_list(configs) and configs != [] ->
SLog.info(@log_category, "Initializing #{length(configs)} pool(s)")
{:ok, configs}
{:ok, []} ->
SLog.warning(@log_category, "No pool configs found, using legacy defaults")
{:ok, [%{name: :default}]}
{:error, reason} ->
SLog.error(@log_category, "Pool configuration error: #{inspect(reason)}")
{:error, {:invalid_pool_config, reason}}
end
end
defp start_workers_concurrently(
pool_name,
count,
startup_timeout,
worker_module,
adapter_module,
pool_config,
pool_genserver_name
) do
actual_count = enforce_max_workers(count, pool_config)
log_worker_startup_info(actual_count, worker_module)
# NOTE: pool_genserver_name is now passed in from the caller
# (captured before spawning the init process)
batch_config = get_batch_config(pool_config)
start_worker_batches(
pool_name,
actual_count,
startup_timeout,
worker_module,
adapter_module,
pool_config,
pool_genserver_name,
batch_config
)
end
defp enforce_max_workers(count, pool_config) do
legacy_pool_config = Application.get_env(:snakepit, :pool_config, %{})
max_workers =
Map.get(pool_config, :max_workers) ||
Map.get(legacy_pool_config, :max_workers, Defaults.pool_max_workers())
actual_count = min(count, max_workers)
if actual_count < count do
SLog.warning(
@log_category,
"⚠️ Requested #{count} workers but limiting to #{actual_count} (max_workers=#{max_workers})"
)
SLog.warning(
@log_category,
"⚠️ To increase this limit, set :pool_config.max_workers in config/config.exs"
)
end
actual_count
end
defp log_worker_startup_info(actual_count, worker_module) do
SLog.info(@log_category, "🚀 Starting concurrent initialization of #{actual_count} workers...")
SLog.info(@log_category, "📦 Using worker type: #{inspect(worker_module)}")
end
defp get_pool_genserver_name do
case Process.info(self(), :registered_name) do
{:registered_name, name} -> name
nil -> self()
end
end
defp get_batch_config(pool_config) do
legacy_pool_config = Application.get_env(:snakepit, :pool_config, %{})
batch_size =
Map.get(pool_config, :startup_batch_size) ||
Map.get(legacy_pool_config, :startup_batch_size, Defaults.pool_startup_batch_size())
batch_delay =
Map.get(pool_config, :startup_batch_delay_ms) ||
Map.get(
legacy_pool_config,
:startup_batch_delay_ms,
Defaults.pool_startup_batch_delay_ms()
)
%{size: batch_size, delay: batch_delay}
end
defp start_worker_batches(
pool_name,
actual_count,
startup_timeout,
worker_module,
adapter_module,
pool_config,
pool_genserver_name,
batch_config
) do
worker_ctx = %{
pool_name: pool_name,
actual_count: actual_count,
startup_timeout: startup_timeout,
worker_module: worker_module,
adapter_module: adapter_module,
pool_config: pool_config,
pool_genserver_name: pool_genserver_name,
batch_config: batch_config
}
1..actual_count
|> Enum.chunk_every(batch_config.size)
|> Enum.with_index()
|> Enum.flat_map(fn {batch, batch_num} ->
start_single_batch(worker_ctx, batch, batch_num)
end)
end
defp start_single_batch(worker_ctx, batch, batch_num) do
# CRITICAL: Check if WorkerSupervisor is still alive before starting a batch
# Short-circuit if app is shutting down to prevent cascading errors
if supervisor_alive?() do
%{actual_count: actual_count, startup_timeout: startup_timeout, batch_config: batch_config} =
worker_ctx
batch_start = batch_num * batch_config.size + 1
batch_end = min(batch_start + length(batch) - 1, actual_count)
SLog.info(
@log_category,
"Starting batch #{batch_num + 1}: workers #{batch_start}-#{batch_end}"
)
workers =
batch
|> Task.async_stream(
fn i -> start_worker_in_batch(worker_ctx, i) end,
timeout: startup_timeout,
max_concurrency: batch_config.size,
on_timeout: :kill_task
)
|> Enum.map(&handle_worker_start_result/1)
|> Enum.filter(&(&1 != nil))
maybe_delay_between_batches(batch_num, actual_count, batch_config)
workers
else
SLog.warning(
@log_category,
"Skipping batch #{batch_num + 1}: WorkerSupervisor terminated during startup"
)
[]
end
end
defp start_worker_in_batch(worker_ctx, i) do
# CRITICAL: Check if WorkerSupervisor is still alive before attempting to start
# This prevents crashes when app is shutting down during batch initialization
if supervisor_alive?() do
%{
pool_name: pool_name,
actual_count: actual_count,
worker_module: worker_module,
adapter_module: adapter_module,
pool_config: pool_config,
pool_genserver_name: pool_genserver_name
} = worker_ctx
worker_id = "#{pool_name}_worker_#{i}_#{:erlang.unique_integer([:positive])}"
result =
if Map.has_key?(pool_config, :worker_profile) do
start_worker_with_profile(
worker_id,
pool_name,
worker_module,
adapter_module,
pool_config,
pool_genserver_name
)
else
start_worker_legacy(
worker_id,
pool_name,
worker_module,
adapter_module,
pool_genserver_name
)
end
handle_worker_start_result_with_log(result, worker_id, i, actual_count)
else
{:error, :supervisor_terminated}
end
end
# Check if the WorkerSupervisor is alive - prevents crashes during shutdown
defp supervisor_alive? do
case Process.whereis(Snakepit.Pool.WorkerSupervisor) do
nil -> false
pid -> Process.alive?(pid)
end
end
defp start_worker_with_profile(
worker_id,
pool_name,
worker_module,
adapter_module,
pool_config,
pool_genserver_name
) do
profile_module = Config.get_profile_module(pool_config)
worker_config =
pool_config
|> Map.put(:worker_id, worker_id)
|> Map.put(:worker_module, worker_module)
|> Map.put(:adapter_module, adapter_module)
|> Map.put(:pool_name, pool_genserver_name)
|> Map.put(:pool_identifier, pool_name)
profile_module.start_worker(worker_config)
end
defp start_worker_legacy(
worker_id,
pool_name,
worker_module,
adapter_module,
pool_genserver_name
) do
WorkerSupervisor.start_worker(
worker_id,
worker_module,
adapter_module,
pool_genserver_name,
%{pool_identifier: pool_name}
)
end
defp handle_worker_start_result_with_log(result, worker_id, i, actual_count) do
case result do
{:ok, _pid} ->
SLog.info(@log_category, "✅ Worker #{i}/#{actual_count} ready: #{worker_id}")
worker_id
{:error, reason} ->
SLog.error(@log_category, "❌ Worker #{i}/#{actual_count} failed: #{inspect(reason)}")
nil
end
end
defp handle_worker_start_result({:ok, worker_id}), do: worker_id
defp handle_worker_start_result({:exit, reason}) do
SLog.error(@log_category, "Worker startup task failed: #{inspect(reason)}")
nil
end
defp maybe_delay_between_batches(batch_num, actual_count, batch_config) do
unless batch_num == div(actual_count - 1, batch_config.size) do
wait_for_batch_delay(batch_config.delay)
end
end
defp wait_for_batch_delay(delay_ms) when delay_ms <= 0, do: :ok
defp wait_for_batch_delay(delay_ms) do
ref = make_ref()
Process.send_after(self(), {:startup_batch_delay, ref}, delay_ms)
receive do
{:startup_batch_delay, ^ref} -> :ok
end
end
defp checkout_worker(pool_state, session_id, affinity_cache) do
case try_checkout_preferred_worker(pool_state, session_id, affinity_cache) do
{:ok, worker_id, new_state} ->
{:ok, worker_id, new_state}
:no_preferred_worker ->
case next_available_worker(pool_state) do
{:ok, worker_id} ->
new_pool_state = increment_worker_load(pool_state, worker_id, session_id)
{:ok, worker_id, new_pool_state}
:no_workers ->
{:error, :no_workers}
end
end
end
defp try_checkout_preferred_worker(_pool_state, nil, _affinity_cache), do: :no_preferred_worker
defp try_checkout_preferred_worker(pool_state, session_id, affinity_cache) do
# PERFORMANCE FIX: Use shared affinity_cache from top-level state
case get_preferred_worker(session_id, affinity_cache) do
{:ok, preferred_worker_id} ->
# Check if preferred worker is available and not tainted
if MapSet.member?(pool_state.available, preferred_worker_id) and
not CrashBarrier.worker_tainted?(preferred_worker_id) do
new_pool_state = increment_worker_load(pool_state, preferred_worker_id, session_id)
SLog.debug(
@log_category,
"Using preferred worker #{preferred_worker_id} for session #{session_id}"
)
{:ok, preferred_worker_id, new_pool_state}
else
:no_preferred_worker
end
{:error, :not_found} ->
:no_preferred_worker
end
end
defp increment_worker_load(pool_state, worker_id, session_id) do
pool_state = ensure_worker_capacity(pool_state, worker_id)
new_load = worker_load(pool_state, worker_id) + 1
capacity = effective_capacity(pool_state, worker_id)
new_loads = Map.put(pool_state.worker_loads, worker_id, new_load)
new_available =
if new_load < capacity do
MapSet.put(pool_state.available, worker_id)
else
MapSet.delete(pool_state.available, worker_id)
end
pool_state = %{
pool_state
| worker_loads: new_loads,
available: new_available
}
maybe_track_capacity(pool_state, worker_id, :increment)
if session_id do
store_session_affinity(session_id, worker_id)
end
pool_state
end
defp next_available_worker(pool_state) do
pool_state.available
|> Enum.find(fn worker_id -> not CrashBarrier.worker_tainted?(worker_id) end)
|> case do
nil -> :no_workers
worker_id -> {:ok, worker_id}
end
end
defp select_queue_worker(pool_state, worker_id) do
if CrashBarrier.worker_tainted?(worker_id) do
case next_available_worker(pool_state) do
{:ok, available} -> {:ok, available}
:no_workers -> :no_worker
end
else
{:ok, worker_id}
end
end
defp decrement_worker_load(pool_state, worker_id) do
pool_state = ensure_worker_capacity(pool_state, worker_id)
current_load = worker_load(pool_state, worker_id)
new_load = max(current_load - 1, 0)
capacity = effective_capacity(pool_state, worker_id)
new_loads =
if new_load > 0 do
Map.put(pool_state.worker_loads, worker_id, new_load)
else
Map.delete(pool_state.worker_loads, worker_id)
end
new_available =
if new_load < capacity and Enum.member?(pool_state.workers, worker_id) do
MapSet.put(pool_state.available, worker_id)
else
pool_state.available
end
pool_state = %{
pool_state
| worker_loads: new_loads,
available: new_available
}
maybe_track_capacity(pool_state, worker_id, :decrement)
pool_state
end
defp worker_load(pool_state, worker_id) do
Map.get(pool_state.worker_loads, worker_id, 0)
end
defp worker_capacity(pool_state, worker_id) do
Map.get(
pool_state.worker_capacities,
worker_id,
resolve_worker_capacity(pool_state, worker_id)
)
end
defp effective_capacity(pool_state, worker_id) do
capacity = worker_capacity(pool_state, worker_id)
case pool_state.capacity_strategy do
:profile -> 1
_ -> capacity
end
end
defp build_worker_capacities(pool_state, workers) do
Enum.reduce(workers, pool_state.worker_capacities, fn worker_id, acc ->
Map.put_new(acc, worker_id, resolve_worker_capacity(pool_state, worker_id))
end)
end
defp ensure_worker_capacity(pool_state, worker_id) do
if Map.has_key?(pool_state.worker_capacities, worker_id) do
pool_state
else
capacity = resolve_worker_capacity(pool_state, worker_id)
%{
pool_state
| worker_capacities: Map.put(pool_state.worker_capacities, worker_id, capacity)
}
end
end
defp ensure_worker_available(pool_state, worker_id) do
load = worker_load(pool_state, worker_id)
capacity = effective_capacity(pool_state, worker_id)
if load < capacity do
%{pool_state | available: MapSet.put(pool_state.available, worker_id)}
else
pool_state
end
end
defp resolve_worker_capacity(pool_state, _worker_id) do
pool_config = pool_state.pool_config || %{}
capacity =
if Config.thread_profile?(pool_config) do
Map.get(pool_config, :threads_per_worker, 1)
else
1
end
max(capacity, 1)
end
defp busy_worker_count(pool_state) do
map_size(pool_state.worker_loads)
end
defp resolve_capacity_strategy(pool_config) do
Map.get(pool_config, :capacity_strategy) ||
Application.get_env(:snakepit, :capacity_strategy, :pool)
end
defp maybe_track_capacity(pool_state, worker_id, :increment) do
if should_track_capacity?(pool_state) do
track_capacity_increment(worker_id)
end
end
defp maybe_track_capacity(pool_state, worker_id, :decrement) do
if pool_state.capacity_strategy == :hybrid and
Config.thread_profile?(pool_state.pool_config) do
case Snakepit.Pool.Registry.get_worker_pid(worker_id) do
{:ok, pid} ->
_ = CapacityStore.decrement_load(pid)
{:error, _} ->
:ok
end
end
end
defp should_track_capacity?(pool_state) do
pool_state.capacity_strategy == :hybrid and Config.thread_profile?(pool_state.pool_config)
end
defp track_capacity_increment(worker_id) do
_ = CapacityStore.ensure_started()
with {:ok, pid} <- Snakepit.Pool.Registry.get_worker_pid(worker_id),
result <- CapacityStore.check_and_increment_load(pid) do
handle_capacity_increment_result(result, pid)
else
{:error, _} -> :ok
end
end
defp handle_capacity_increment_result({:ok, capacity, new_load}, pid) do
if new_load == capacity do
emit_capacity_reached_telemetry(capacity, new_load, pid, false)
end
end
defp handle_capacity_increment_result({:at_capacity, capacity, load}, pid) do
emit_capacity_reached_telemetry(capacity, load, pid, true)
end
defp handle_capacity_increment_result({:error, :unknown_worker}, pid) do
SLog.warning(@log_category, "Worker #{inspect(pid)} not found in capacity store")
end
defp emit_capacity_reached_telemetry(capacity, load, pid, rejected?) do
metadata = %{worker_pid: pid, profile: :thread}
metadata = if rejected?, do: Map.put(metadata, :rejected, true), else: metadata
:telemetry.execute(
[:snakepit, :pool, :capacity_reached],
%{capacity: capacity, load: load},
metadata
)
end
# PERFORMANCE FIX: ETS-cached session affinity lookup
# Eliminates GenServer bottleneck by caching session->worker mappings with TTL
defp get_preferred_worker(session_id, cache_table) do
current_time = System.monotonic_time(:second)
case lookup_cached_worker(cache_table, session_id, current_time) do
{:ok, worker_id} ->
{:ok, worker_id}
:cache_miss ->
fetch_and_cache_worker(session_id, cache_table, current_time)
end
end
defp lookup_cached_worker(cache_table, session_id, current_time) do
case :ets.lookup(cache_table, session_id) do
[{^session_id, worker_id, expires_at}] when expires_at > current_time ->
{:ok, worker_id}
_ ->
:cache_miss
end
end
defp fetch_and_cache_worker(session_id, cache_table, current_time) do
case SessionStore.get_session(session_id) do
{:ok, session} ->
extract_and_cache_worker_id(session, session_id, cache_table, current_time)
{:error, :not_found} ->
{:error, :not_found}
end
end
defp extract_and_cache_worker_id(session, session_id, cache_table, current_time) do
case Map.get(session, :last_worker_id) do
nil ->
{:error, :not_found}
worker_id ->
expires_at = current_time + 60
:ets.insert(cache_table, {session_id, worker_id, expires_at})
{:ok, worker_id}
end
end
defp store_session_affinity(session_id, worker_id) do
# Store the worker affinity in a supervised task for better error logging
async_with_context(fn ->
:ok = SessionStore.store_worker_session(session_id, worker_id)
SLog.debug(@log_category, "Stored session affinity: #{session_id} -> #{worker_id}")
:ok
end)
end
defp execute_on_worker(worker_id, command, args, opts) do
timeout = get_command_timeout(command, args, opts)
worker_module = get_worker_module(worker_id)
try do
result = worker_module.execute(worker_id, command, args, timeout)
# Increment request count for lifecycle management (on success only)
case result do
{:ok, _} ->
LifecycleManager.increment_request_count(worker_id)
_ ->
:ok
end
result
catch
:exit, {:timeout, _} ->
{:error, :worker_timeout}
:exit, reason ->
{:error, {:worker_exit, reason}}
end
end
defp execute_with_crash_barrier(
pool_pid,
pool_name,
worker_id,
command,
args,
opts,
pool_config
) do
crash_config = CrashBarrier.config(pool_config)
if CrashBarrier.enabled?(crash_config) do
attempt_with_retry(pool_pid, pool_name, worker_id, command, args, opts, crash_config, 0)
else
{execute_on_worker(worker_id, command, args, opts), worker_id}
end
end
defp attempt_with_retry(
pool_pid,
pool_name,
worker_id,
command,
args,
opts,
crash_config,
attempt
) do
result = execute_on_worker(worker_id, command, args, opts)
retry_context = %{
pool_pid: pool_pid,
pool_name: pool_name,
command: command,
args: args,
opts: opts,
crash_config: crash_config
}
case CrashBarrier.crash_info(result, crash_config) do
{:ok, info} ->
maybe_taint_worker(pool_name, worker_id, info, crash_config)
handle_crash_retry(result, info, retry_context, attempt)
:error ->
{result, worker_id}
end
end
defp maybe_taint_worker(pool_name, worker_id, info, crash_config) do
if CrashBarrier.worker_tainted?(worker_id) do
:ok
else
CrashBarrier.taint_worker(pool_name, worker_id, info, crash_config)
end
end
defp handle_crash_retry(result, info, retry_context, attempt) do
if CrashBarrier.retry_allowed?(
retry_context.crash_config,
CrashBarrier.idempotent?(retry_context.args),
attempt
) do
maybe_wait_backoff(CrashBarrier.retry_backoff(retry_context.crash_config, attempt + 1))
retry_with_worker(result, info, retry_context, attempt)
else
{CrashBarrier.normalize_crash_error(result, info), nil}
end
end
defp retry_with_worker(result, info, retry_context, attempt) do
case checkout_worker_for_retry(
retry_context.pool_pid,
retry_context.pool_name,
retry_context.args,
retry_context.opts
) do
{:ok, next_worker} ->
attempt_with_retry(
retry_context.pool_pid,
retry_context.pool_name,
next_worker,
retry_context.command,
retry_context.args,
retry_context.opts,
retry_context.crash_config,
attempt + 1
)
{:error, _reason} ->
{CrashBarrier.normalize_crash_error(result, info), nil}
end
end
defp checkout_worker_for_retry(pool_pid, pool_name, args, opts) do
session_id =
opts[:session_id] ||
Map.get(args, :session_id) ||
Map.get(args, "session_id")
GenServer.call(
pool_pid,
{:checkout_worker, pool_name, session_id},
Defaults.crash_barrier_checkout_timeout()
)
end
defp maybe_wait_backoff(delay_ms) when is_integer(delay_ms) and delay_ms > 0 do
receive do
after
delay_ms -> :ok
end
end
defp maybe_wait_backoff(_), do: :ok
defp get_worker_module(worker_id) do
# Try to determine the worker module from registry or configuration
case PoolRegistry.fetch_worker(worker_id) do
{:ok, _pid, %{worker_module: module}} when is_atom(module) ->
module
_ ->
# Fallback: use GRPCWorker
Snakepit.GRPCWorker
end
end
defp get_command_timeout(command, args, opts) do
# Prefer explicit client timeout, then adapter timeout, then global default
case opts[:timeout] do
nil ->
case get_adapter_timeout(command, args) do
# Global default
nil -> Defaults.default_command_timeout()
adapter_timeout -> adapter_timeout
end
client_timeout ->
client_timeout
end
end
defp get_adapter_timeout(command, args) do
case Application.get_env(:snakepit, :adapter_module) do
nil ->
nil
adapter_module ->
if function_exported?(adapter_module, :command_timeout, 2) do
try do
adapter_module.command_timeout(command, args)
rescue
# Fall back to default if adapter timeout fails
_ -> nil
end
else
nil
end
end
end
defp monitor_client_status(ref, client_pid) do
if Process.alive?(client_pid) do
await_client_down(ref, client_pid)
else
case await_client_down(ref, client_pid) do
:alive -> {:down, :unknown}
other -> other
end
end
end
defp await_client_down(ref, client_pid) do
receive do
{:DOWN, ^ref, :process, ^client_pid, reason} -> {:down, reason}
after
0 -> :alive
end
end
defp async_with_context(fun) when is_function(fun, 0) do
ctx = :otel_ctx.get_current()
Task.Supervisor.async_nolink(Snakepit.TaskSupervisor, fn ->
token = :otel_ctx.attach(ctx)
try do
fun.()
after
:otel_ctx.detach(token)
end
end)
end
# DIAGNOSTIC: Resource monitoring helpers
defp capture_resource_metrics do
%{
beam_processes: length(:erlang.processes()),
beam_ports: length(:erlang.ports()),
memory_total_mb: div(:erlang.memory(:total), 1_024 * 1_024),
memory_processes_mb: div(:erlang.memory(:processes), 1_024 * 1_024),
ets_tables: length(:ets.all()),
timestamp: System.monotonic_time(:millisecond)
}
end
defp calculate_resource_delta(baseline, peak) do
%{
processes_delta: peak.beam_processes - baseline.beam_processes,
ports_delta: peak.beam_ports - baseline.beam_ports,
memory_delta_mb: peak.memory_total_mb - baseline.memory_total_mb,
memory_processes_delta_mb: peak.memory_processes_mb - baseline.memory_processes_mb,
ets_tables_delta: peak.ets_tables - baseline.ets_tables,
time_elapsed_ms: peak.timestamp - baseline.timestamp
}
end
@doc false
def extract_pool_name_from_worker_id(worker_id) do
case lookup_pool_from_registry(worker_id) do
{:ok, pool_name} ->
pool_name
{:error, reason} ->
inferred = infer_pool_from_id(worker_id)
SLog.warning(
@log_category,
"Falling back to worker_id parsing for #{worker_id}: #{inspect(reason)}. Using #{inspect(inferred)}"
)
inferred
end
end
defp lookup_pool_from_registry(worker_id) do
case PoolRegistry.fetch_worker(worker_id) do
{:ok, pid, metadata} ->
extract_pool_from_metadata(metadata, pid)
{:error, reason} ->
{:error, reason}
end
rescue
_ -> {:error, :registry_lookup_failed}
end
defp extract_pool_from_metadata(metadata, pid) do
case Map.get(metadata, :pool_identifier) do
pool_identifier when is_atom(pool_identifier) and not is_nil(pool_identifier) ->
{:ok, pool_identifier}
_ ->
extract_pool_from_name(metadata, pid)
end
end
defp extract_pool_from_name(metadata, pid) do
pool_name = Map.get(metadata, :pool_name)
extract_pool_name_by_type(pool_name, metadata, pid)
end
defp extract_pool_name_by_type(pool_name, _metadata, _pid) when is_atom(pool_name) do
validate_atom_pool_name(pool_name)
end
defp extract_pool_name_by_type(pool_pid, _metadata, _pid) when is_pid(pool_pid) do
extract_pool_name_from_pid(pool_pid)
end
defp extract_pool_name_by_type(pool_name, _metadata, _pid) when is_binary(pool_name) do
convert_string_to_pool_name(pool_name)
end
defp extract_pool_name_by_type(nil, metadata, pid) do
{:error, {:pool_metadata_missing, %{metadata_keys: Map.keys(metadata), worker_pid: pid}}}
end
defp validate_atom_pool_name(pool_name) do
if module_atom?(pool_name) do
{:error, {:pool_metadata_module_atom, pool_name}}
else
{:ok, pool_name}
end
end
defp extract_pool_name_from_pid(pool_pid) do
case Process.info(pool_pid, :registered_name) do
{:registered_name, name} when is_atom(name) ->
{:ok, name}
_ ->
{:error, {:pool_metadata_not_atom, pool_pid}}
end
end
defp convert_string_to_pool_name(pool_name) do
atom_name = String.to_existing_atom(pool_name)
validate_atom_pool_name(atom_name)
rescue
ArgumentError -> {:error, {:pool_metadata_not_atom, pool_name}}
end
defp module_atom?(atom) when is_atom(atom) do
atom |> Atom.to_string() |> String.starts_with?("Elixir.")
end
defp infer_pool_from_id(worker_id) do
case String.split(worker_id, "_worker_", parts: 2) do
[pool_name_str, _rest] ->
safe_to_existing_atom(pool_name_str)
_ ->
:default
end
end
defp safe_to_existing_atom(pool_name_str) do
String.to_existing_atom(pool_name_str)
rescue
ArgumentError -> :default
end
end
defmodule Snakepit.Pool.ProcessRegistry do
@moduledoc """
Registry for tracking external worker processes with OS-level PID management.
This module maintains a mapping between:
- Worker IDs
- Elixir worker PIDs
- External process PIDs
- Process fingerprints
Enables robust orphaned process detection and cleanup.
"""
use GenServer
alias Snakepit.Defaults
alias Snakepit.Logger, as: SLog
@table_name :snakepit_pool_process_registry
@log_category :pool
defstruct [
:table,
:dets_table,
:beam_run_id,
:beam_os_pid
]
# Client API
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Reserves a worker slot before spawning the process.
This ensures we can track the process even if we crash during spawn.
"""
def reserve_worker(worker_id) do
GenServer.call(__MODULE__, {:reserve_worker, worker_id})
end
@doc """
Activates a reserved worker with its actual process information.
This is a synchronous call that blocks until the worker is registered.
This ensures the happens-before relationship: worker registration completes
before the worker is considered ready for work.
"""
def activate_worker(worker_id, elixir_pid, process_pid, fingerprint) do
activate_worker(worker_id, elixir_pid, process_pid, fingerprint, [])
end
def activate_worker(worker_id, elixir_pid, process_pid, fingerprint, opts) when is_list(opts) do
pgid = Keyword.get(opts, :pgid, process_pid)
process_group? = Keyword.get(opts, :process_group?, false)
GenServer.call(
__MODULE__,
{:activate_worker, worker_id, elixir_pid, process_pid, fingerprint, pgid, process_group?},
5000
)
end
@doc """
Unregisters a worker from tracking.
Returns :ok regardless of whether the worker was registered.
"""
def unregister_worker(worker_id) do
GenServer.cast(__MODULE__, {:unregister, worker_id})
end
@doc """
Checks if a worker is currently registered.
"""
def worker_registered?(worker_id) do
case :ets.lookup(@table_name, worker_id) do
[{^worker_id, _}] -> true
[] -> false
end
end
@doc """
Gets all active external process PIDs from registered workers.
"""
def get_active_process_pids do
GenServer.call(__MODULE__, :get_active_process_pids)
end
@doc """
Gets all registered external process PIDs, regardless of worker status.
This is useful during shutdown when workers may have been terminated
but external processes still need cleanup.
"""
def get_all_process_pids do
GenServer.call(__MODULE__, :get_all_process_pids)
end
@doc """
Gets all registered worker information.
"""
def list_all_workers do
:ets.tab2list(@table_name)
end
@doc """
Gets all registered worker entries for the current BEAM run.
"""
def current_run_entries do
GenServer.call(__MODULE__, :current_run_entries)
end
@doc """
Gets information for a specific worker.
"""
def get_worker_info(worker_id) do
GenServer.call(__MODULE__, {:get_worker_info, worker_id})
end
@doc """
Gets workers with specific fingerprints.
"""
def get_workers_by_fingerprint(fingerprint) do
:ets.tab2list(@table_name)
|> Enum.filter(fn {_id, %{fingerprint: fp}} -> fp == fingerprint end)
end
@doc """
Validates that all registered workers are still alive.
Returns a list of dead workers that should be cleaned up.
"""
def validate_workers do
:ets.tab2list(@table_name)
|> Enum.filter(fn {_id, %{elixir_pid: pid}} -> not Process.alive?(pid) end)
|> Enum.map(fn {worker_id, worker_info} -> {worker_id, worker_info} end)
end
@doc """
Cleans up dead worker entries from the registry.
"""
def cleanup_dead_workers do
GenServer.call(__MODULE__, :cleanup_dead_workers)
end
@doc """
Get the current BEAM run ID.
"""
def get_beam_run_id do
GenServer.call(__MODULE__, :get_beam_run_id)
end
@doc """
Gets registry statistics.
"""
def get_stats do
all_workers = :ets.tab2list(@table_name)
alive_workers =
Enum.filter(all_workers, fn {_id, %{elixir_pid: pid}} -> Process.alive?(pid) end)
%{
total_registered: length(all_workers),
alive_workers: length(alive_workers),
dead_workers: length(all_workers) - length(alive_workers),
active_process_pids: length(get_active_process_pids())
}
end
@doc """
Manually trigger orphan cleanup. Useful for testing and debugging.
"""
def manual_orphan_cleanup do
GenServer.call(__MODULE__, :manual_orphan_cleanup)
end
@doc """
Debug function to show all DETS entries with their BEAM run IDs.
"""
def debug_show_all_entries do
GenServer.call(__MODULE__, :debug_show_all_entries)
end
@doc """
Returns the number of entries currently stored in the DETS table.
"""
def dets_table_size do
GenServer.call(__MODULE__, :dets_table_size)
end
# Server Callbacks
@impl true
def init(_opts) do
# Generate short 7-character run ID for this BEAM instance
# This will be embedded in Python CLI commands for reliable tracking
run_id = Snakepit.RunID.generate()
# Keep as beam_run_id for backward compatibility
beam_run_id = run_id
# CRITICAL: Get the BEAM OS PID for robust stale entry detection
# This allows us to verify if the BEAM that created an entry is still running
beam_os_pid = System.pid() |> String.to_integer()
# Create a proper file path for DETS
# Include node name to prevent conflicts between multiple BEAM instances
priv_dir = :code.priv_dir(:snakepit) |> to_string()
# Sanitize node name for filesystem usage
node_name = node() |> to_string() |> String.replace(~r/[^a-zA-Z0-9_-]/, "_")
dets_file = Path.join([priv_dir, "data", "process_registry_#{node_name}.dets"])
# Ensure directory exists for DETS file
dets_dir = Path.dirname(dets_file)
File.mkdir_p!(dets_dir)
# Open DETS for persistence with repair option
# Generate an unguessable table identifier so callers cannot mutate DETS directly.
dets_table_name =
:crypto.strong_rand_bytes(8)
|> Base.encode32(case: :lower)
|> then(&:"snakepit_process_registry_dets_#{&1}")
dets_result =
:dets.open_file(dets_table_name, [
{:file, to_charlist(dets_file)},
{:type, :set},
# Auto-save every 1000ms
{:auto_save, 1000},
# Automatically repair corrupted files
{:repair, true}
])
dets_table =
case dets_result do
{:ok, table} ->
table
{:error, reason} ->
SLog.error(
@log_category,
"Failed to open DETS file: #{inspect(reason)}. Deleting and recreating..."
)
File.rm(dets_file)
{:ok, table} =
:dets.open_file(dets_table_name, [
{:file, to_charlist(dets_file)},
{:type, :set},
{:auto_save, 1000}
])
table
end
# Create ETS table for worker tracking - protected so only GenServer can write
table =
:ets.new(@table_name, [
:set,
:protected,
:named_table,
{:read_concurrency, true}
])
SLog.info(
@log_category,
"Snakepit Pool Process Registry started with BEAM run ID: #{beam_run_id}, BEAM OS PID: #{beam_os_pid}"
)
# Perform startup cleanup of orphaned processes FIRST
cleanup_orphaned_processes(dets_table, beam_run_id, beam_os_pid)
# Load current run's processes into ETS
load_current_run_processes(dets_table, table, beam_run_id)
# Schedule periodic cleanup
schedule_cleanup()
{:ok,
%__MODULE__{
table: table,
dets_table: dets_table,
beam_run_id: beam_run_id,
beam_os_pid: beam_os_pid
}}
end
@impl true
def handle_cast({:unregister, worker_id}, state) do
case :ets.lookup(state.table, worker_id) do
[{^worker_id, %{process_pid: process_pid} = info}] ->
if process_alive?(process_pid) do
updated = mark_terminating(info)
:ets.insert(state.table, {worker_id, updated})
:dets.insert(state.dets_table, {worker_id, updated})
:dets.sync(state.dets_table)
schedule_unregister_cleanup(
worker_id,
process_pid,
Defaults.process_registry_unregister_cleanup_attempts()
)
SLog.debug(
@log_category,
"Deferring unregister for #{worker_id}; external process #{process_pid} still alive"
)
else
:ets.delete(state.table, worker_id)
:dets.delete(state.dets_table, worker_id)
:dets.sync(state.dets_table)
SLog.info(
@log_category,
"🚮 Unregistered worker #{worker_id} with external process PID #{process_pid}"
)
end
[] ->
# Also check DETS in case ETS was cleared
:dets.delete(state.dets_table, worker_id)
:dets.sync(state.dets_table)
# Defensive check: Only log at debug level for unknown workers
# This is expected during certain race conditions and shouldn't be a warning
SLog.debug(
@log_category,
"Attempted to unregister unknown worker #{worker_id} - ignoring (worker may have failed to register)"
)
end
{:noreply, state}
end
@impl true
def handle_call(
{:activate_worker, worker_id, elixir_pid, process_pid, fingerprint, pgid, process_group?},
_from,
state
) do
worker_info = %{
status: :active,
elixir_pid: elixir_pid,
process_pid: process_pid,
fingerprint: fingerprint,
registered_at: System.system_time(:second),
beam_run_id: state.beam_run_id,
beam_os_pid: state.beam_os_pid,
pgid: pgid,
process_group?: process_group?
}
# Update both ETS and DETS atomically
:ets.insert(state.table, {worker_id, worker_info})
:dets.insert(state.dets_table, {worker_id, worker_info})
# Sync immediately for activation too
:dets.sync(state.dets_table)
SLog.debug(
@log_category,
"🆕 Worker activated: #{worker_id} | PID #{process_pid} | " <>
"BEAM run #{state.beam_run_id} | Elixir PID: #{inspect(elixir_pid)}"
)
# Reply :ok to unblock the caller - worker is now fully registered
{:reply, :ok, state}
end
@impl true
def handle_call({:reserve_worker, worker_id}, _from, state) do
reservation_info = %{
status: :reserved,
reserved_at: System.system_time(:second),
beam_run_id: state.beam_run_id,
beam_os_pid: state.beam_os_pid
}
# CRITICAL: Persist to DETS immediately and sync
:dets.insert(state.dets_table, {worker_id, reservation_info})
# Force immediate write to disk
:dets.sync(state.dets_table)
SLog.info(
@log_category,
"Reserved worker slot #{worker_id} for BEAM run #{state.beam_run_id}"
)
{:reply, :ok, state}
end
@impl true
def handle_call({:get_worker_info, worker_id}, _from, state) do
reply =
case :ets.lookup(state.table, worker_id) do
[{^worker_id, worker_info}] -> {:ok, worker_info}
[] -> {:error, :not_found}
end
{:reply, reply, state}
end
@impl true
def handle_call(:get_active_process_pids, _from, state) do
pids =
:ets.tab2list(state.table)
|> Enum.filter(fn
{_id, %{status: :active, elixir_pid: pid}} -> Process.alive?(pid)
# Legacy entries without status
{_id, %{elixir_pid: pid}} -> Process.alive?(pid)
_ -> false
end)
|> Enum.map(fn {_id, info} -> Map.get(info, :process_pid) end)
|> Enum.filter(&(&1 != nil))
{:reply, pids, state}
end
@impl true
def handle_call(:get_all_process_pids, _from, state) do
pids =
:ets.tab2list(state.table)
|> Enum.filter(fn
{_id, %{status: :active}} -> true
# Legacy entries
{_id, %{process_pid: _}} -> true
_ -> false
end)
|> Enum.map(fn {_id, info} -> Map.get(info, :process_pid) end)
|> Enum.filter(&(&1 != nil))
{:reply, pids, state}
end
@impl true
def handle_call(:current_run_entries, _from, state) do
entries =
:ets.tab2list(state.table)
|> Enum.filter(fn {_id, info} -> Map.get(info, :beam_run_id) == state.beam_run_id end)
{:reply, entries, state}
end
@impl true
def handle_call(:get_beam_run_id, _from, state) do
{:reply, state.beam_run_id, state}
end
@impl true
def handle_call(:cleanup_dead_workers, _from, state) do
dead_count = do_cleanup_dead_workers(state)
{:reply, dead_count, state}
end
@impl true
def handle_call(:manual_orphan_cleanup, _from, state) do
SLog.info(@log_category, "Manual orphan cleanup triggered")
cleanup_orphaned_processes(state.dets_table, state.beam_run_id, state.beam_os_pid)
{:reply, :ok, state}
end
@impl true
def handle_call(:debug_show_all_entries, _from, state) do
all_entries = :dets.match_object(state.dets_table, :_)
entries_info =
Enum.map(all_entries, fn {worker_id, info} ->
%{
worker_id: worker_id,
process_pid: info.process_pid,
beam_run_id: info.beam_run_id,
is_current_run: info.beam_run_id == state.beam_run_id,
process_alive: process_alive?(info.process_pid)
}
end)
{:reply, {entries_info, state.beam_run_id}, state}
end
@impl true
def handle_call(:dets_table_size, _from, state) do
reply =
case state.dets_table do
nil ->
{:error, :not_initialized}
table ->
{:ok, :dets.info(table, :size)}
end
{:reply, reply, state}
end
@impl true
def handle_info(:cleanup_dead_workers, state) do
dead_count = do_cleanup_dead_workers(state)
if dead_count > 0 do
SLog.info(@log_category, "Cleaned up #{dead_count} dead worker entries")
end
# Schedule next cleanup
schedule_cleanup()
{:noreply, state}
end
@impl true
def handle_info({:unregister_retry, worker_id, process_pid, attempts_left}, state) do
case :ets.lookup(state.table, worker_id) do
[{^worker_id, %{process_pid: ^process_pid}}] ->
handle_unregister_retry_match(state, worker_id, process_pid, attempts_left)
_ ->
:ok
end
{:noreply, state}
end
@impl true
def handle_info(msg, state) do
SLog.debug(@log_category, "ProcessRegistry received unexpected message: #{inspect(msg)}")
{:noreply, state}
end
defp handle_unregister_retry_match(state, worker_id, process_pid, attempts_left) do
cond do
process_alive?(process_pid) and attempts_left > 1 ->
schedule_unregister_cleanup(worker_id, process_pid, attempts_left - 1)
process_alive?(process_pid) ->
:ok
true ->
:ets.delete(state.table, worker_id)
:dets.delete(state.dets_table, worker_id)
:dets.sync(state.dets_table)
SLog.info(
@log_category,
"🚮 Unregistered worker #{worker_id} with external process PID #{process_pid}"
)
end
end
@impl true
def terminate(reason, state) do
SLog.info(@log_category, "Snakepit Pool Process Registry terminating: #{inspect(reason)}")
# Log current state before closing
all_entries = :dets.match_object(state.dets_table, :_)
SLog.info(
@log_category,
"ProcessRegistry terminating with #{length(all_entries)} entries in DETS"
)
# Ensure DETS is properly synced and closed
if state.dets_table do
:dets.sync(state.dets_table)
:dets.close(state.dets_table)
end
:ok
end
# Private Functions
defp schedule_cleanup do
Process.send_after(
self(),
:cleanup_dead_workers,
Defaults.process_registry_cleanup_interval()
)
end
defp schedule_unregister_cleanup(worker_id, process_pid, attempts_left)
when attempts_left > 0 do
Process.send_after(
self(),
{:unregister_retry, worker_id, process_pid, attempts_left},
Defaults.process_registry_unregister_cleanup_delay()
)
end
defp cleanup_orphaned_processes(dets_table, current_beam_run_id, current_beam_os_pid) do
SLog.debug(
@log_category,
"Starting orphan cleanup for BEAM run #{current_beam_run_id}, BEAM OS PID #{current_beam_os_pid}"
)
all_entries = :dets.match_object(dets_table, :_)
SLog.info(@log_category, "Total entries in DETS: #{length(all_entries)}")
stale_entries = find_stale_entries(all_entries, current_beam_run_id, current_beam_os_pid)
SLog.info(
@log_category,
"Found #{length(stale_entries)} stale entries to remove (from previous runs)"
)
old_run_orphans = find_old_run_orphans(dets_table, current_beam_run_id)
SLog.info(
@log_category,
"Found #{length(old_run_orphans)} active processes from previous BEAM runs"
)
abandoned_reservations = find_abandoned_reservations(all_entries, current_beam_run_id)
SLog.info(@log_category, "Found #{length(abandoned_reservations)} abandoned reservations")
killed_count = kill_orphaned_processes(old_run_orphans, current_beam_run_id)
abandoned_killed = kill_abandoned_reservation_processes(abandoned_reservations)
entries_to_remove =
combine_entries_to_remove(old_run_orphans, abandoned_reservations, stale_entries)
remove_dets_entries(dets_table, entries_to_remove)
rogue_killed = cleanup_rogue_processes(current_beam_run_id)
log_cleanup_summary(
killed_count,
abandoned_killed,
rogue_killed,
stale_entries,
entries_to_remove
)
end
defp find_stale_entries(all_entries, current_beam_run_id, current_beam_os_pid) do
Enum.filter(all_entries, fn {_worker_id, info} ->
stale_entry?(info, current_beam_run_id, current_beam_os_pid)
end)
end
defp stale_entry?(info, current_beam_run_id, current_beam_os_pid) do
cond do
Map.has_key?(info, :beam_os_pid) ->
check_beam_os_pid_stale(info, current_beam_run_id, current_beam_os_pid)
Map.has_key?(info, :beam_run_id) ->
info.beam_run_id != current_beam_run_id
true ->
true
end
end
defp check_beam_os_pid_stale(info, current_beam_run_id, current_beam_os_pid) do
beam_dead = not Snakepit.ProcessKiller.process_alive?(info.beam_os_pid)
if beam_dead do
SLog.info(@log_category, "Entry is stale: BEAM OS PID #{info.beam_os_pid} is dead")
true
else
check_same_beam_stale(info, current_beam_run_id, current_beam_os_pid)
end
end
defp check_same_beam_stale(info, current_beam_run_id, current_beam_os_pid) do
if info.beam_os_pid == current_beam_os_pid do
info.beam_run_id != current_beam_run_id
else
false
end
end
defp find_old_run_orphans(dets_table, current_beam_run_id) do
:dets.select(dets_table, [
{{:"$1", :"$2"},
[
{:andalso, {:"/=", {:map_get, :beam_run_id, :"$2"}, current_beam_run_id},
{
:orelse,
{:==, {:map_get, :status, :"$2"}, :active},
{:==, {:map_size, :"$2"}, 6}
}}
], [{{:"$1", :"$2"}}]}
])
end
defp find_abandoned_reservations(all_entries, current_beam_run_id) do
now = System.system_time(:second)
Enum.filter(all_entries, fn {_id, info} ->
abandoned_reservation?(info, current_beam_run_id, now)
end)
end
defp abandoned_reservation?(info, current_beam_run_id, now) do
Map.get(info, :status) == :reserved and
(info.beam_run_id != current_beam_run_id or
now - Map.get(info, :reserved_at, 0) > 60)
end
defp kill_orphaned_processes(old_run_orphans, current_beam_run_id) do
Enum.reduce(old_run_orphans, 0, fn {worker_id, info}, acc ->
attempt_kill_orphaned_process(worker_id, info, current_beam_run_id, acc)
end)
end
defp attempt_kill_orphaned_process(worker_id, info, current_beam_run_id, acc) do
if Snakepit.ProcessKiller.process_alive?(info.process_pid) do
log_orphan_found(info.process_pid, worker_id, info.beam_run_id)
verify_and_kill_process(info.process_pid, info.beam_run_id, current_beam_run_id)
acc + 1
else
SLog.debug(
@log_category,
"Orphaned entry #{worker_id} with PID #{info.process_pid} already dead"
)
acc
end
end
defp log_orphan_found(process_pid, worker_id, beam_run_id) do
SLog.warning(
@log_category,
"Found orphaned process #{process_pid} (worker: #{worker_id}) from previous " <>
"BEAM run #{beam_run_id}. Terminating..."
)
end
defp verify_and_kill_process(process_pid, expected_run_id, current_beam_run_id) do
case Snakepit.ProcessKiller.get_process_command(process_pid) do
{:ok, cmd} ->
handle_process_command(process_pid, cmd, expected_run_id, current_beam_run_id)
{:error, _} ->
SLog.debug(@log_category, "Process #{process_pid} not found, already dead")
end
end
defp handle_process_command(process_pid, cmd, expected_run_id, current_beam_run_id) do
has_grpc_server = String.contains?(cmd, "grpc_server.py")
has_old_run_id = String.contains?(cmd, "--snakepit-run-id #{expected_run_id}")
has_new_run_id = String.contains?(cmd, "--run-id #{expected_run_id}")
log_pid_reuse_check(process_pid, expected_run_id, current_beam_run_id, cmd)
if has_grpc_server and (has_old_run_id or has_new_run_id) do
kill_confirmed_orphan(process_pid)
else
log_pid_reuse_detected(process_pid, has_grpc_server, cmd)
end
end
defp log_pid_reuse_check(process_pid, expected_run_id, current_beam_run_id, cmd) do
SLog.warning(
@log_category,
"PID REUSE CHECK: PID #{process_pid} | Expected run_id: #{expected_run_id} | " <>
"Current BEAM run_id: #{current_beam_run_id} | Process cmd: #{String.trim(cmd)}"
)
end
defp kill_confirmed_orphan(process_pid) do
SLog.info(
@log_category,
"Confirmed PID #{process_pid} is a grpc_server process with matching run_id"
)
case Snakepit.ProcessKiller.kill_with_escalation(process_pid) do
:ok ->
SLog.info(@log_category, "Process #{process_pid} successfully terminated")
{:error, reason} ->
SLog.error(@log_category, "Failed to kill process #{process_pid}: #{inspect(reason)}")
end
end
defp log_pid_reuse_detected(process_pid, has_grpc_server, cmd) do
if has_grpc_server do
SLog.warning(
@log_category,
"PID #{process_pid} is a grpc_server process but with DIFFERENT beam_run_id. " <>
"OS reused PID for new worker! Skipping kill. Command: #{String.trim(cmd)}"
)
else
SLog.debug(
@log_category,
"PID #{process_pid} is not a grpc_server process, skipping: #{String.trim(cmd)}"
)
end
end
defp kill_abandoned_reservation_processes(abandoned_reservations) do
abandoned_reservations
|> Enum.map(fn {worker_id, info} ->
kill_abandoned_reservation(worker_id, info)
end)
|> Enum.sum()
end
defp kill_abandoned_reservation(worker_id, info) do
SLog.debug(
@log_category,
"Found abandoned reservation #{worker_id} from run #{info.beam_run_id}. " <>
"Attempting cleanup..."
)
{:ok, count} = Snakepit.ProcessKiller.kill_by_run_id(info.beam_run_id)
SLog.info(@log_category, "Killed #{count} processes for run #{info.beam_run_id}")
count
end
defp combine_entries_to_remove(old_run_orphans, abandoned_reservations, stale_entries) do
(old_run_orphans ++ abandoned_reservations ++ stale_entries)
|> Enum.uniq_by(fn {worker_id, _} -> worker_id end)
end
defp remove_dets_entries(dets_table, entries_to_remove) do
Enum.each(entries_to_remove, fn {worker_id, _info} ->
:dets.delete(dets_table, worker_id)
end)
end
defp log_cleanup_summary(
killed_count,
abandoned_killed,
rogue_killed,
stale_entries,
entries_to_remove
) do
SLog.info(
@log_category,
"Orphan cleanup complete. Killed #{killed_count} orphaned processes, " <>
"killed #{abandoned_killed} from abandoned reservations, " <>
"killed #{rogue_killed} rogue processes, " <>
"removed #{length(stale_entries)} stale entries, " <>
"total removed: #{length(entries_to_remove)} entries."
)
end
defp cleanup_rogue_processes(current_beam_run_id) do
cleanup_config =
Application.get_env(:snakepit, :rogue_cleanup, enabled: true)
|> normalize_cleanup_config()
if cleanup_config[:enabled] == false do
SLog.info(
@log_category,
"Skipping rogue process cleanup (disabled via :rogue_cleanup config)"
)
0
else
do_cleanup_rogue_processes(current_beam_run_id, cleanup_config)
end
end
defp do_cleanup_rogue_processes(current_beam_run_id, cleanup_config) do
scripts = Map.get(cleanup_config, :scripts, default_cleanup_scripts())
run_markers = Map.get(cleanup_config, :run_markers, default_run_markers())
python_commands = get_python_commands()
owned_processes = filter_owned_processes(python_commands, scripts, run_markers)
SLog.info(
@log_category,
"Found #{length(owned_processes)} snakepit grpc_server processes with run markers"
)
rogue_processes =
find_rogue_processes(owned_processes, current_beam_run_id, scripts, run_markers)
log_rogue_processes(rogue_processes)
kill_rogue_processes(rogue_processes)
end
defp get_python_commands do
Snakepit.ProcessKiller.find_python_processes()
|> Enum.reduce([], fn pid, acc ->
case Snakepit.ProcessKiller.get_process_command(pid) do
{:ok, cmd} -> [{pid, cmd} | acc]
_ -> acc
end
end)
end
defp filter_owned_processes(python_commands, scripts, run_markers) do
Enum.filter(python_commands, fn {_pid, cmd} ->
snakepit_command?(cmd, scripts) and has_run_marker?(cmd, run_markers)
end)
end
defp find_rogue_processes(owned_processes, current_beam_run_id, scripts, run_markers) do
Enum.filter(owned_processes, fn {_pid, cmd} ->
cleanup_candidate?(cmd, current_beam_run_id, scripts: scripts, run_markers: run_markers)
end)
end
defp log_rogue_processes([]), do: :ok
defp log_rogue_processes(rogue_processes) do
rogue_pids = Enum.map(rogue_processes, fn {pid, _cmd} -> pid end)
SLog.warning(
@log_category,
"Found #{length(rogue_processes)} rogue grpc_server processes not belonging to current run"
)
SLog.warning(@log_category, "Rogue PIDs: #{inspect(rogue_pids)}")
end
defp kill_rogue_processes(rogue_processes) do
Enum.reduce(rogue_processes, 0, fn {pid, cmd}, acc ->
kill_rogue_process(pid, cmd, acc)
end)
end
defp kill_rogue_process(pid, cmd, acc) do
SLog.warning(@log_category, "Killing rogue process #{pid}: #{String.trim(cmd)}")
case Snakepit.ProcessKiller.kill_with_escalation(pid) do
:ok ->
acc + 1
{:error, reason} ->
SLog.error(@log_category, "Failed to kill rogue process #{pid}: #{inspect(reason)}")
acc
end
end
@doc false
def cleanup_candidate?(command, current_run_id, opts \\ []) when is_binary(command) do
scripts = Keyword.get(opts, :scripts, default_cleanup_scripts())
markers = Keyword.get(opts, :run_markers, default_run_markers())
snakepit_command?(command, scripts) and has_run_marker?(command, markers) and
not has_run_id?(command, current_run_id, markers)
end
defp snakepit_command?(command, scripts) do
Enum.any?(scripts, &String.contains?(command, &1))
end
defp has_run_marker?(command, markers) do
Enum.any?(markers, &String.contains?(command, &1))
end
defp has_run_id?(command, run_id, markers) when is_binary(run_id) do
Enum.any?(markers, fn marker ->
String.contains?(command, "#{marker} #{run_id}")
end)
end
defp normalize_cleanup_config(%{} = config) do
config
|> Map.put_new(:enabled, true)
|> Map.put_new(:scripts, default_cleanup_scripts())
|> Map.put_new(:run_markers, default_run_markers())
end
defp normalize_cleanup_config(config) when is_list(config),
do: Enum.into(config, %{}) |> normalize_cleanup_config()
defp normalize_cleanup_config(_),
do: %{
enabled: true,
scripts: default_cleanup_scripts(),
run_markers: default_run_markers()
}
defp default_cleanup_scripts, do: ["grpc_server.py", "grpc_server_threaded.py"]
defp default_run_markers, do: ["--snakepit-run-id", "--run-id"]
defp load_current_run_processes(dets_table, ets_table, beam_run_id) do
# Load only processes from current BEAM run into ETS
current_processes =
:dets.select(dets_table, [
{{:"$1", :"$2"}, [{:==, {:map_get, :beam_run_id, :"$2"}, beam_run_id}],
[{{:"$1", :"$2"}}]}
])
Enum.each(current_processes, fn {worker_id, info} ->
:ets.insert(ets_table, {worker_id, info})
end)
SLog.info(
@log_category,
"Loaded #{length(current_processes)} processes from current BEAM run"
)
end
# Delegate to ProcessKiller for process checking
defp process_alive?(pid), do: Snakepit.ProcessKiller.process_alive?(pid)
# Helper function for cleanup that operates on the table
defp do_cleanup_dead_workers(state) do
dead_workers =
:ets.tab2list(state.table)
|> Enum.filter(fn {_id, %{elixir_pid: pid}} -> not Process.alive?(pid) end)
{count, dirty} =
Enum.reduce(dead_workers, {0, false}, fn {worker_id, %{process_pid: process_pid} = info},
{acc, dirty?} ->
if process_alive?(process_pid) do
updated = mark_terminating(info)
:ets.insert(state.table, {worker_id, updated})
:dets.insert(state.dets_table, {worker_id, updated})
:dets.sync(state.dets_table)
{acc, dirty?}
else
:ets.delete(state.table, worker_id)
:dets.delete(state.dets_table, worker_id)
SLog.info(
@log_category,
"Cleaned up dead worker #{worker_id} with external process PID #{process_pid}"
)
{acc + 1, true}
end
end)
if dirty do
:dets.sync(state.dets_table)
end
count
end
defp mark_terminating(info) do
info
|> Map.put(:terminating?, true)
|> Map.put(:terminated_at, System.system_time(:second))
end
end
defmodule Snakepit.Pool.Registry do
@moduledoc """
Registry for pool worker processes.
This is a thin wrapper around Elixir's Registry that provides:
- Consistent naming for worker processes
- Easy migration path to distributed registry (Horde)
- Helper functions for worker lookup
## Canonical Metadata
All workers store a metadata map containing the following canonical keys:
* `:worker_module` – module that owns the worker implementation (usually `Snakepit.GRPCWorker`)
* `:pool_name` – atom name of the logical pool (e.g. `:default`)
* `:pool_identifier` – optional human-friendly identifier used in docs/metrics
* `:adapter_module` – adapter used to launch the Python worker
Higher-level helpers (pool, diagnostics, worker profiles) should prefer
`Snakepit.Pool.Registry.fetch_worker/1` so these keys stay authoritative.
"""
alias Snakepit.Logger, as: SLog
@registry_name __MODULE__
@metadata_keys [:worker_module, :pool_name, :pool_identifier, :adapter_module]
@log_category :pool
@doc """
Returns the child spec for the registry.
"""
def child_spec(_opts) do
Registry.child_spec(
keys: :unique,
name: @registry_name
)
end
@doc """
Returns a via tuple for registering/looking up a worker process.
## Examples
iex> Snakepit.Pool.Registry.via_tuple("worker_123")
{:via, Registry, {Snakepit.Pool.Registry, "worker_123"}}
"""
def via_tuple(worker_id) when is_binary(worker_id) do
{:via, Registry, {@registry_name, worker_id}}
end
@doc """
Lists all registered worker IDs.
"""
def list_workers do
Registry.select(@registry_name, [{{:"$1", :_, :_}, [], [:"$1"]}])
end
@doc """
Checks if a worker is registered.
"""
def worker_exists?(worker_id) do
match?({:ok, _pid, _meta}, fetch_worker(worker_id))
end
@doc """
Gets the PID for a worker ID.
"""
def get_worker_pid(worker_id) do
with {:ok, pid, _metadata} <- fetch_worker(worker_id) do
{:ok, pid}
end
end
@doc """
Counts the number of registered workers.
"""
def worker_count do
Registry.count(@registry_name)
end
@doc """
Returns the list of canonical metadata keys maintained for each worker.
"""
def metadata_keys, do: @metadata_keys
@doc """
Register a worker with metadata for O(1) reverse lookups.
This is only used for manual registration - workers started with via_tuple are already registered.
"""
def register_worker(_worker_id, _pid) do
# Workers started with via_tuple are already registered automatically
# This is a no-op for compatibility
:ok
end
@doc """
Adds or updates metadata for a registered worker.
Accepts maps to keep metadata consistent across callers. When `Registry`
has `nil` metadata (the default when using `:via` tuples), this function
replaces it with the provided map. Future updates merge with the existing map.
Returns `:ok` on success or `{:error, :not_registered}` if the worker has
not been registered yet (best-effort semantics).
"""
def put_metadata(worker_id, metadata) when is_binary(worker_id) and is_map(metadata) do
sanitized = normalize_metadata(metadata)
try do
case Registry.update_value(@registry_name, worker_id, fn
current when is_map(current) -> Map.merge(current, sanitized)
_ -> sanitized
end) do
{_, _} ->
:ok
:error ->
SLog.debug(
@log_category,
"Pool.Registry.put_metadata/2 attempted to update #{inspect(worker_id)} before registration"
)
{:error, :not_registered}
end
rescue
ArgumentError ->
SLog.debug(
@log_category,
"Pool.Registry.put_metadata/2 attempted to update #{inspect(worker_id)} before registration"
)
{:error, :not_registered}
end
end
def put_metadata(_worker_id, _metadata), do: :ok
@doc """
Returns `{pid, metadata}` for a registered worker.
"""
def fetch_worker(worker_id) when is_binary(worker_id) do
case Registry.lookup(@registry_name, worker_id) do
[{pid, metadata}] ->
{:ok, pid, normalize_metadata(metadata)}
[] ->
{:error, :not_found}
end
end
def fetch_worker(_worker_id), do: {:error, :invalid_worker_id}
@doc """
Returns only the metadata for a worker.
"""
def get_worker_metadata(worker_id) do
with {:ok, _pid, metadata} <- fetch_worker(worker_id) do
{:ok, metadata}
end
end
@doc """
Get worker_id from PID for O(1) lookups in :DOWN messages.
"""
def get_worker_id_by_pid(pid) do
# Use Registry's keys/2 function for O(1) reverse lookup
case Registry.keys(@registry_name, pid) do
[worker_id] -> {:ok, worker_id}
[] -> {:error, :not_found}
end
end
defp normalize_metadata(metadata) when is_map(metadata), do: metadata
defp normalize_metadata(_metadata), do: %{}
end
defmodule Snakepit.Pool.Worker.StarterRegistry do
@moduledoc """
Registry for worker starter supervisors.
This registry provides a clean separation between worker processes and
their starter supervisors, making debugging and process tracking easier.
Worker starters are registered with their worker_id as the key, allowing
for easy lookup and management of individual starter supervisors.
"""
@registry_name __MODULE__
@doc """
Returns the child spec for the starter registry.
"""
def child_spec(_opts) do
Registry.child_spec(
keys: :unique,
name: @registry_name
)
end
@doc """
Returns a via tuple for registering/looking up a worker starter supervisor.
## Examples
iex> Snakepit.Pool.Worker.StarterRegistry.via_tuple("worker_123")
{:via, Registry, {Snakepit.Pool.Worker.StarterRegistry, "worker_123"}}
"""
def via_tuple(worker_id) when is_binary(worker_id) do
{:via, Registry, {@registry_name, worker_id}}
end
@doc """
Lists all registered worker starter IDs.
"""
def list_starters do
Registry.select(@registry_name, [{{:"$1", :_, :_}, [], [:"$1"]}])
end
@doc """
Checks if a worker starter is registered.
"""
def starter_exists?(worker_id) do
case Registry.lookup(@registry_name, worker_id) do
[] -> false
_ -> true
end
end
@doc """
Gets the PID for a worker starter supervisor.
"""
def get_starter_pid(worker_id) do
case Registry.lookup(@registry_name, worker_id) do
[{pid, _}] -> {:ok, pid}
[] -> {:error, :not_found}
end
end
@doc """
Counts the number of registered worker starters.
"""
def starter_count do
Registry.count(@registry_name)
end
end
defmodule Snakepit.Pool.Worker.Starter do
@moduledoc """
Supervisor wrapper for individual workers that provides automatic restart capability.
This module implements the "Permanent Wrapper" pattern for managing workers that
control external OS processes (Python gRPC servers).
## Architecture Decision
**See**: `docs/architecture/adr-001-worker-starter-supervision-pattern.md` for
detailed rationale, alternatives considered, and trade-offs.
## Why This Pattern?
**TL;DR**: Workers manage external Python processes, not just Elixir state.
This pattern provides:
- Automatic restart without Pool intervention
- Atomic resource cleanup (worker + Python process)
- Future extensibility for per-worker resources
**Trade-off**: Extra process (~1KB) per worker for better encapsulation.
## Architecture
```
DynamicSupervisor (WorkerSupervisor)
└── Worker.Starter (Supervisor, :permanent)
└── GRPCWorker (GenServer, :transient)
└── Port → Python grpc_server.py
```
## Lifecycle
**When GRPCWorker crashes**:
1. Worker.Starter detects crash via :one_for_one strategy
2. Worker.Starter automatically restarts GRPCWorker
3. Pool notified via :DOWN but doesn't manage restart
4. New GRPCWorker spawns new Python process and re-registers
**When Worker.Starter terminates**:
1. GRPCWorker receives shutdown signal
2. GRPCWorker.terminate sends SIGTERM to Python
3. Python process exits gracefully
4. Worker.Starter confirms all children stopped
5. Clean atomic shutdown
This decouples Pool (availability management) from Worker lifecycle (crash/restart).
## Related
- **Issue #2**: Community feedback questioning this complexity
- **ADR-001**: Full architecture decision record with alternatives
- **External Process Design**: `docs/20251007_external_process_supervision_design.md`
"""
use Supervisor
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.Worker.StarterRegistry
@log_category :pool
@doc """
Starts a worker starter supervisor.
## Parameters
* `worker_id` - Unique identifier for the worker
"""
def start_link(worker_id) when is_binary(worker_id) do
Supervisor.start_link(__MODULE__, {worker_id, Snakepit.Pool.Worker},
name: via_name(worker_id)
)
end
def start_link({worker_id, worker_module}) when is_binary(worker_id) do
Supervisor.start_link(__MODULE__, {worker_id, worker_module}, name: via_name(worker_id))
end
def start_link({worker_id, worker_module, adapter_module}) when is_binary(worker_id) do
Supervisor.start_link(__MODULE__, {worker_id, worker_module, adapter_module, nil},
name: via_name(worker_id)
)
end
def start_link({worker_id, worker_module, adapter_module, pool_name})
when is_binary(worker_id) do
Supervisor.start_link(__MODULE__, {worker_id, worker_module, adapter_module, pool_name, %{}},
name: via_name(worker_id)
)
end
def start_link({worker_id, worker_module, adapter_module, pool_name, worker_config})
when is_binary(worker_id) do
Supervisor.start_link(
__MODULE__,
{worker_id, worker_module, adapter_module, pool_name, worker_config},
name: via_name(worker_id)
)
end
@doc """
Returns a via tuple for this starter supervisor.
"""
def via_name(worker_id) do
StarterRegistry.via_tuple(worker_id)
end
@impl true
def init({worker_id, worker_module}) do
init({worker_id, worker_module, nil, nil, %{}})
end
def init({worker_id, worker_module, adapter_module}) do
init({worker_id, worker_module, adapter_module, nil, %{}})
end
def init({worker_id, worker_module, adapter_module, pool_name}) do
init({worker_id, worker_module, adapter_module, pool_name, %{}})
end
def init({worker_id, worker_module, adapter_module, pool_name, worker_config}) do
# Check if the Pool is already terminating
# For dynamic pools, we can't check a specific name, so skip this check if pool_name is a PID
should_check_global_pool = pool_name == nil || pool_name == Snakepit.Pool
if should_check_global_pool do
case Process.whereis(Snakepit.Pool) do
nil ->
# Global pool is dead, don't start workers
SLog.debug(
@log_category,
"Aborting worker starter for #{worker_id} - Global pool is dead"
)
:ignore
_pid ->
do_init_worker(worker_id, worker_module, adapter_module, pool_name, worker_config)
end
else
# Using a custom pool (like in tests), always proceed
do_init_worker(worker_id, worker_module, adapter_module, pool_name, worker_config)
end
end
defp do_init_worker(worker_id, worker_module, adapter_module, pool_name, worker_config) do
SLog.debug(
@log_category,
"Starting worker starter for #{worker_id} with module #{inspect(worker_module)}"
)
adapter = adapter_module || Application.get_env(:snakepit, :adapter_module)
# CRITICAL FIX: Pass pool_name to worker so it knows which pool to notify when ready
# Default to Snakepit.Pool for backward compatibility (production use)
# v0.6.0: Pass worker_config for lifecycle management
worker_opts =
[
id: worker_id,
adapter: adapter,
pool_name: pool_name || Snakepit.Pool,
worker_config: worker_config
]
|> maybe_put_pool_identifier(pool_name, worker_config)
children = [
%{
id: worker_id,
start: {worker_module, :start_link, [worker_opts]},
# Within this supervisor, the worker restarts on crashes but not during shutdown
restart: :transient,
# CRITICAL: Give worker time to gracefully shutdown (send SIGTERM, wait for Python).
# Derived from :graceful_shutdown_timeout_ms + margin to stay in sync with GRPCWorker.
# Default: 6000ms (graceful) + 2000ms (margin) = 8000ms
shutdown: supervisor_shutdown_timeout(),
type: :worker
}
]
Supervisor.init(children, strategy: :one_for_one)
end
defp maybe_put_pool_identifier(opts, pool_name, worker_config) do
identifier =
if is_map(worker_config) do
worker_config
|> Map.get(:pool_identifier)
|> normalize_identifier()
else
normalize_identifier(pool_name)
end
if identifier do
Keyword.put(opts, :pool_identifier, identifier)
else
opts
end
end
defp normalize_identifier(value) when is_atom(value) do
if Atom.to_string(value) |> String.starts_with?("Elixir.") do
nil
else
value
end
end
defp normalize_identifier(value) when is_binary(value) do
normalize_identifier(String.to_existing_atom(value))
rescue
ArgumentError -> nil
end
defp normalize_identifier(_), do: nil
# Derive supervisor shutdown timeout from the same config as GRPCWorker.
# This ensures consistency: if a user sets :graceful_shutdown_timeout_ms,
# both the worker's terminate/2 and the supervisor's shutdown are aligned.
@default_graceful_shutdown_timeout 6000
@shutdown_margin 2000
defp supervisor_shutdown_timeout do
graceful =
Application.get_env(
:snakepit,
:graceful_shutdown_timeout_ms,
@default_graceful_shutdown_timeout
)
graceful + @shutdown_margin
end
end
defmodule Snakepit.Pool.WorkerSupervisor do
@moduledoc """
DynamicSupervisor for pool worker processes.
This supervisor manages the lifecycle of workers:
- Starts workers on demand
- Handles crashes with automatic restarts
- Provides clean shutdown of workers
"""
use DynamicSupervisor
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Pool.Worker.StarterRegistry
@log_category :pool
@doc """
Starts the worker supervisor.
"""
def start_link(init_arg) do
DynamicSupervisor.start_link(__MODULE__, init_arg, name: __MODULE__)
end
@impl true
def init(_init_arg) do
DynamicSupervisor.init(
strategy: :one_for_one,
extra_arguments: []
)
end
@doc """
Starts a new pool worker with the given ID.
## Examples
iex> Snakepit.Pool.WorkerSupervisor.start_worker("worker_123")
{:ok, #PID<0.123.0>}
"""
def start_worker(
worker_id,
worker_module \\ Snakepit.GRPCWorker,
adapter_module \\ nil,
pool_name \\ nil,
worker_config \\ %{}
)
when is_binary(worker_id) do
# Start the permanent starter supervisor, not the transient worker directly
# This gives us automatic worker restarts without Pool intervention
# CRITICAL FIX: Pass pool_name to Worker.Starter so workers know which pool to notify
# v0.6.0: Pass worker_config for lifecycle management
child_spec =
{Snakepit.Pool.Worker.Starter,
{worker_id, worker_module, adapter_module, pool_name, worker_config}}
case DynamicSupervisor.start_child(__MODULE__, child_spec) do
{:ok, starter_pid} ->
SLog.info(
@log_category,
"Started worker starter for #{worker_id} with PID #{inspect(starter_pid)}"
)
{:ok, starter_pid}
{:error, {:already_started, starter_pid}} ->
SLog.debug(
@log_category,
"Worker starter for #{worker_id} already running with PID #{inspect(starter_pid)}"
)
{:ok, starter_pid}
{:error, reason} = error ->
SLog.error(
@log_category,
"Failed to start worker starter for #{worker_id}: #{inspect(reason)}"
)
error
end
end
@doc """
Stops a worker gracefully.
"""
def stop_worker(worker_pid) when is_pid(worker_pid) do
case PoolRegistry.get_worker_id_by_pid(worker_pid) do
{:ok, worker_id} -> stop_worker(worker_id)
{:error, :not_found} -> {:error, :worker_not_found}
end
end
def stop_worker(worker_id) when is_binary(worker_id) do
case StarterRegistry.get_starter_pid(worker_id) do
{:ok, starter_pid} ->
DynamicSupervisor.terminate_child(__MODULE__, starter_pid)
{:error, :not_found} ->
{:error, :worker_not_found}
end
end
@doc """
Lists all supervised workers.
"""
def list_workers do
DynamicSupervisor.which_children(__MODULE__)
|> Enum.map(fn {_, pid, _, _} -> pid end)
end
@doc """
Returns the count of active workers.
"""
def worker_count do
DynamicSupervisor.count_children(__MODULE__).active
end
@doc """
Restarts a worker by ID.
"""
def restart_worker(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, old_pid} ->
# Get port metadata before terminating so we can check if it's released
%{current_port: current_port, requested_port: requested_port} =
get_worker_port_info(old_pid)
# Worker exists, terminate it and wait for resource cleanup
with :ok <- stop_worker(worker_id),
:ok <- wait_for_resource_cleanup(worker_id, current_port, requested_port) do
start_worker(worker_id)
else
# Propagate termination/cleanup errors
{:error, :worker_not_found} -> start_worker(worker_id)
error -> error
end
{:error, :not_found} ->
# Worker doesn't exist, so we just need to start it
start_worker(worker_id)
end
end
defp cleanup_retry_interval_ms do
Application.get_env(:snakepit, :cleanup_retry_interval_ms, 50)
end
defp cleanup_max_retries do
Application.get_env(:snakepit, :cleanup_max_retries, 20)
end
# Wait for external resources to be released after worker termination.
#
# This is necessary because:
# 1. DynamicSupervisor.terminate_child waits for Elixir process termination
# 2. But external OS process + ports may still be shutting down
# 3. Starting a new worker immediately can cause port binding conflicts
#
# We check:
# - Port availability (can we bind to it?)
# - Registry cleanup (entry removed?)
#
# This prevents race conditions on worker restart.
# Uses exponential backoff for efficient polling: starts fast, backs off gradually.
defp wait_for_resource_cleanup(
worker_id,
current_port,
requested_port,
retries \\ cleanup_max_retries(),
backoff \\ cleanup_retry_interval_ms()
) do
if retries > 0 do
check_and_wait_for_cleanup(worker_id, current_port, requested_port, retries, backoff)
else
handle_cleanup_timeout(worker_id)
end
end
defp check_and_wait_for_cleanup(worker_id, current_port, requested_port, retries, backoff) do
port_to_probe = port_probe_target(current_port, requested_port)
probe_port? = should_probe_port?(requested_port) and port_to_probe not in [nil, 0]
maybe_delay_initial_probe(probe_port?, retries, backoff)
port_released? = check_port_released(worker_id, port_to_probe, probe_port?, retries)
if port_released? and registry_cleaned?(worker_id) do
SLog.debug(@log_category, "Resources released for #{worker_id}, safe to restart")
:ok
else
retry_cleanup_check(worker_id, current_port, requested_port, retries, backoff)
end
end
defp maybe_delay_initial_probe(probe_port?, retries, backoff) do
if probe_port? and retries == cleanup_max_retries() do
initial_delay = min(backoff, 50)
receive do
after
initial_delay -> :ok
end
end
end
defp check_port_released(worker_id, port_to_probe, probe_port?, retries) do
if probe_port? do
SLog.debug(@log_category, "Probing port #{port_to_probe} before restarting #{worker_id}")
port_available?(port_to_probe)
else
log_ephemeral_port_skip(worker_id, retries)
true
end
end
defp log_ephemeral_port_skip(worker_id, retries) do
if retries == cleanup_max_retries() do
SLog.info(
@log_category,
"Skipping port availability probe for #{worker_id}; worker requested an ephemeral port"
)
end
end
defp retry_cleanup_check(worker_id, current_port, requested_port, retries, backoff) do
delay = min(backoff, 200)
receive do
after
delay -> :ok
end
wait_for_resource_cleanup(
worker_id,
current_port,
requested_port,
retries - 1,
backoff * 2
)
end
defp handle_cleanup_timeout(worker_id) do
SLog.warning(
@log_category,
"Resource cleanup timeout for #{worker_id} after #{cleanup_max_retries()} retries, " <>
"proceeding with restart anyway"
)
{:error, :cleanup_timeout}
end
defp get_worker_port_info(worker_pid) do
case GenServer.call(worker_pid, :get_port_metadata, 1000) do
{:ok, %{current_port: port} = info} ->
%{
current_port: port,
requested_port: Map.get(info, :requested_port)
}
_ ->
legacy_port_info(worker_pid)
end
catch
:exit, _ -> %{current_port: nil, requested_port: nil}
end
defp legacy_port_info(worker_pid) do
case GenServer.call(worker_pid, :get_port, 1000) do
{:ok, port} -> %{current_port: port, requested_port: port}
_ -> %{current_port: nil, requested_port: nil}
end
catch
:exit, _ -> %{current_port: nil, requested_port: nil}
end
defp port_available?(port) when is_integer(port) do
# Try to bind to the port to verify it's available
case :gen_tcp.listen(port, [:binary, active: false, reuseaddr: true]) do
{:ok, socket} ->
:gen_tcp.close(socket)
true
{:error, :eaddrinuse} ->
false
{:error, _other} ->
# Other errors (permission, etc) - assume unavailable
false
end
end
# No port to check
defp port_available?(nil), do: true
defp registry_cleaned?(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:error, :not_found} -> true
{:ok, _pid} -> false
end
end
@doc false
def port_probe_target(current_port, requested_port) do
cond do
current_port not in [nil, 0] -> current_port
requested_port not in [nil, 0] -> requested_port
true -> nil
end
end
defp should_probe_port?(requested_port) do
requested_port not in [nil, 0]
end
end
defmodule Snakepit.Telemetry.Handlers.Logger do
@moduledoc """
Telemetry handler that logs ML-related events.
Provides structured logging for hardware detection, circuit breaker
state changes, GPU profiling, and error events.
"""
alias Snakepit.Logger, as: SLog
alias Snakepit.Telemetry.Events
@handler_id "snakepit-ml-logger"
@log_category :telemetry
@doc """
Attaches the logger handler to all ML events.
"""
@spec attach() :: :ok
def attach do
# Detach first to allow re-attachment
detach()
events = Events.all_ml_events()
:telemetry.attach_many(
@handler_id,
events,
&handle_event/4,
%{level: log_level()}
)
:ok
end
@doc """
Detaches the logger handler.
"""
@spec detach() :: :ok
def detach do
:telemetry.detach(@handler_id)
:ok
rescue
_ -> :ok
end
# Hardware events
defp handle_event(
[:snakepit, :hardware, :detect, :stop],
%{duration: duration},
%{accelerator: acc, platform: platform},
config
) do
duration_ms = System.convert_time_unit(duration, :native, :millisecond)
log(
config.level,
"Hardware detection completed in #{duration_ms}ms: accelerator=#{acc} platform=#{platform}"
)
end
defp handle_event(
[:snakepit, :hardware, :select, :stop],
%{duration: duration},
%{device: device, success: success},
config
) do
duration_us = System.convert_time_unit(duration, :native, :microsecond)
log(
config.level,
"Device selection: #{inspect(device)} success=#{success} (#{duration_us}μs)"
)
end
# Circuit breaker events
defp handle_event(
[:snakepit, :circuit_breaker, :opened],
%{failure_count: count},
%{pool: pool, reason: reason},
config
) do
log(
:warning,
"Circuit breaker OPENED for pool=#{pool}: #{count} failures, reason=#{reason}"
)
_ = config
end
defp handle_event(
[:snakepit, :circuit_breaker, :closed],
_measurements,
%{pool: pool},
config
) do
log(config.level, "Circuit breaker CLOSED for pool=#{pool}")
end
defp handle_event(
[:snakepit, :circuit_breaker, :half_open],
_measurements,
%{pool: pool},
config
) do
log(config.level, "Circuit breaker HALF-OPEN for pool=#{pool}")
end
defp handle_event(
[:snakepit, :circuit_breaker, :call, :rejected],
_measurements,
%{pool: pool, state: state},
config
) do
log(:warning, "Circuit breaker rejected call for pool=#{pool} state=#{state}")
_ = config
end
# GPU profiler events
defp handle_event(
[:snakepit, :gpu, :memory, :sampled],
%{used_mb: used, total_mb: total},
%{device: device},
config
) do
percent = Float.round(used / max(total, 1) * 100, 1)
log(
config.level,
"GPU memory: #{used}/#{total}MB (#{percent}%) device=#{inspect(device)}"
)
end
# Error events
defp handle_event(
[:snakepit, :error, :shape_mismatch],
_measurements,
%{expected: expected, got: got, operation: op},
_config
) do
log(:warning, "Shape mismatch in #{op}: expected #{inspect(expected)}, got #{inspect(got)}")
end
defp handle_event(
[:snakepit, :error, :oom],
%{requested_bytes: requested, available_bytes: available},
%{device: device},
_config
) do
req_mb = Float.round(requested / 1_048_576, 1)
avail_mb = Float.round(available / 1_048_576, 1)
log(
:error,
"OOM on device=#{inspect(device)}: requested #{req_mb}MB, available #{avail_mb}MB"
)
end
defp handle_event(
[:snakepit, :error, :device],
_measurements,
%{expected_device: expected, actual_device: actual},
_config
) do
log(
:warning,
"Device mismatch: expected #{inspect(expected)}, got #{inspect(actual)}"
)
end
# Retry events
defp handle_event(
[:snakepit, :retry, :attempt],
%{attempt: attempt, delay_ms: delay},
%{pool: pool},
config
) do
log(config.level, "Retry attempt #{attempt} for pool=#{pool} after #{delay}ms delay")
end
defp handle_event(
[:snakepit, :retry, :exhausted],
%{attempts: attempts},
%{pool: pool, last_error: error},
_config
) do
log(
:error,
"Retries exhausted for pool=#{pool} after #{attempts} attempts: #{inspect(error)}"
)
end
# Catch-all for unhandled events
defp handle_event(event, measurements, metadata, config) do
log(
config.level,
"Telemetry: #{inspect(event)} measurements=#{inspect(measurements)} metadata=#{inspect(metadata)}"
)
end
defp log(level, message) do
case level do
:debug -> SLog.debug(@log_category, message)
:info -> SLog.info(@log_category, message)
:warning -> SLog.warning(@log_category, message)
:error -> SLog.error(@log_category, message)
_ -> SLog.debug(@log_category, message)
end
end
defp log_level do
Application.get_env(:snakepit, :telemetry_log_level, :debug)
end
end
defmodule Snakepit.Telemetry.Handlers.Metrics do
@moduledoc """
Telemetry metrics definitions for ML-related events.
Provides `telemetry_metrics` compatible metric definitions for
hardware detection, circuit breaker, GPU profiling, and error events.
"""
import Telemetry.Metrics
@doc """
Returns all ML-related telemetry metrics definitions.
These can be used with `TelemetryMetricsPrometheus` or other
telemetry metrics reporters.
"""
@spec definitions() :: [Telemetry.Metrics.t()]
def definitions do
hardware_metrics() ++
circuit_breaker_metrics() ++
gpu_profiler_metrics() ++
error_metrics() ++
retry_metrics()
end
@doc """
Returns Prometheus-compatible metric definitions.
Same as `definitions/0` but ensures all metrics have names
compatible with Prometheus naming conventions.
"""
@spec prometheus_definitions() :: [Telemetry.Metrics.t()]
def prometheus_definitions do
definitions()
end
@spec hardware_metrics() :: [Telemetry.Metrics.t()]
defp hardware_metrics do
[
# Hardware detection timing
summary(
"snakepit.hardware.detect.duration",
event_name: [:snakepit, :hardware, :detect, :stop],
measurement: :duration,
unit: {:native, :millisecond},
tags: [:accelerator, :platform],
description: "Duration of hardware detection"
),
# Device selection timing
summary(
"snakepit.hardware.select.duration",
event_name: [:snakepit, :hardware, :select, :stop],
measurement: :duration,
unit: {:native, :microsecond},
tags: [:device],
description: "Duration of device selection"
),
# Cache hit/miss counts
counter(
"snakepit.hardware.cache.hits.total",
event_name: [:snakepit, :hardware, :cache, :hit],
description: "Hardware cache hits"
),
counter(
"snakepit.hardware.cache.misses.total",
event_name: [:snakepit, :hardware, :cache, :miss],
description: "Hardware cache misses"
)
]
end
@spec circuit_breaker_metrics() :: [Telemetry.Metrics.t()]
defp circuit_breaker_metrics do
[
# State transition counts
counter(
"snakepit.circuit_breaker.opened.total",
event_name: [:snakepit, :circuit_breaker, :opened],
tags: [:pool, :reason],
description: "Circuit breaker open events"
),
counter(
"snakepit.circuit_breaker.closed.total",
event_name: [:snakepit, :circuit_breaker, :closed],
tags: [:pool],
description: "Circuit breaker close events"
),
counter(
"snakepit.circuit_breaker.half_open.total",
event_name: [:snakepit, :circuit_breaker, :half_open],
tags: [:pool],
description: "Circuit breaker half-open events"
),
# Call metrics
counter(
"snakepit.circuit_breaker.calls.allowed.total",
event_name: [:snakepit, :circuit_breaker, :call, :allowed],
tags: [:pool],
description: "Calls allowed through circuit breaker"
),
counter(
"snakepit.circuit_breaker.calls.rejected.total",
event_name: [:snakepit, :circuit_breaker, :call, :rejected],
tags: [:pool],
description: "Calls rejected by circuit breaker"
),
summary(
"snakepit.circuit_breaker.call.duration",
event_name: [:snakepit, :circuit_breaker, :call, :success],
measurement: :duration,
unit: {:native, :millisecond},
tags: [:pool],
description: "Duration of successful circuit breaker calls"
),
# Failure count at open time
last_value(
"snakepit.circuit_breaker.failure_count",
event_name: [:snakepit, :circuit_breaker, :opened],
measurement: :failure_count,
tags: [:pool],
description: "Failure count when circuit breaker opened"
)
]
end
@spec gpu_profiler_metrics() :: [Telemetry.Metrics.t()]
defp gpu_profiler_metrics do
[
# GPU memory
last_value(
"snakepit.gpu.memory.used_mb",
event_name: [:snakepit, :gpu, :memory, :sampled],
measurement: :used_mb,
tags: [:device],
description: "GPU memory used in MB"
),
last_value(
"snakepit.gpu.memory.total_mb",
event_name: [:snakepit, :gpu, :memory, :sampled],
measurement: :total_mb,
tags: [:device],
description: "GPU total memory in MB"
),
last_value(
"snakepit.gpu.memory.free_mb",
event_name: [:snakepit, :gpu, :memory, :sampled],
measurement: :free_mb,
tags: [:device],
description: "GPU free memory in MB"
),
# GPU utilization
last_value(
"snakepit.gpu.utilization.percent",
event_name: [:snakepit, :gpu, :utilization, :sampled],
measurement: :gpu_percent,
tags: [:device],
description: "GPU utilization percentage"
),
# GPU temperature
last_value(
"snakepit.gpu.temperature.celsius",
event_name: [:snakepit, :gpu, :temperature, :sampled],
measurement: :celsius,
tags: [:device],
description: "GPU temperature in Celsius"
),
# GPU power
last_value(
"snakepit.gpu.power.watts",
event_name: [:snakepit, :gpu, :power, :sampled],
measurement: :watts,
tags: [:device],
description: "GPU power usage in watts"
)
]
end
@spec error_metrics() :: [Telemetry.Metrics.t()]
defp error_metrics do
[
# Error type counts
counter(
"snakepit.errors.shape_mismatch.total",
event_name: [:snakepit, :error, :shape_mismatch],
tags: [:operation],
description: "Shape mismatch errors"
),
counter(
"snakepit.errors.device.total",
event_name: [:snakepit, :error, :device],
description: "Device mismatch errors"
),
counter(
"snakepit.errors.oom.total",
event_name: [:snakepit, :error, :oom],
tags: [:device],
description: "Out of memory errors"
),
counter(
"snakepit.errors.dtype_mismatch.total",
event_name: [:snakepit, :error, :dtype_mismatch],
description: "Data type mismatch errors"
),
counter(
"snakepit.errors.python_exception.total",
event_name: [:snakepit, :error, :python_exception],
tags: [:type],
description: "Python exceptions"
)
]
end
@spec retry_metrics() :: [Telemetry.Metrics.t()]
defp retry_metrics do
[
# Retry attempts
summary(
"snakepit.retry.attempts",
event_name: [:snakepit, :retry, :success],
measurement: :attempts,
tags: [:pool],
description: "Number of retry attempts before success"
),
# Retry exhaustion
counter(
"snakepit.retry.exhausted.total",
event_name: [:snakepit, :retry, :exhausted],
tags: [:pool],
description: "Retry exhaustion events"
),
# Backoff delays
summary(
"snakepit.retry.backoff.delay_ms",
event_name: [:snakepit, :retry, :backoff],
measurement: :delay_ms,
tags: [:pool],
description: "Retry backoff delay in milliseconds"
)
]
end
end
defmodule Snakepit.Telemetry.Control do
@moduledoc """
Helper functions for creating telemetry control messages.
Control messages flow from Elixir to Python workers over the gRPC
telemetry stream to adjust telemetry behavior at runtime.
"""
alias Snakepit.Bridge.{
TelemetryControl,
TelemetryEventFilter,
TelemetrySamplingUpdate,
TelemetryToggle
}
@doc """
Creates a control message to enable or disable telemetry.
## Examples
iex> Snakepit.Telemetry.Control.toggle(true)
%Snakepit.Bridge.TelemetryControl{
control: {:toggle, %Snakepit.Bridge.TelemetryToggle{enabled: true}}
}
"""
def toggle(enabled) when is_boolean(enabled) do
%TelemetryControl{
control: {:toggle, %TelemetryToggle{enabled: enabled}}
}
end
@doc """
Creates a control message to adjust sampling rate.
The sampling rate must be between 0.0 and 1.0, where:
- 0.0 = no events emitted
- 1.0 = all events emitted
- 0.1 = 10% of events emitted
Event patterns use glob-style matching (e.g., "python.*").
## Examples
iex> Snakepit.Telemetry.Control.sampling(0.5)
%Snakepit.Bridge.TelemetryControl{
control: {:sampling, %Snakepit.Bridge.TelemetrySamplingUpdate{
sampling_rate: 0.5,
event_patterns: []
}}
}
iex> Snakepit.Telemetry.Control.sampling(0.1, ["python.call.*"])
%Snakepit.Bridge.TelemetryControl{
control: {:sampling, %Snakepit.Bridge.TelemetrySamplingUpdate{
sampling_rate: 0.1,
event_patterns: ["python.call.*"]
}}
}
"""
def sampling(rate, patterns \\ [])
when is_number(rate) and rate >= 0.0 and rate <= 1.0 and is_list(patterns) do
%TelemetryControl{
control:
{:sampling,
%TelemetrySamplingUpdate{
sampling_rate: rate,
event_patterns: Enum.map(patterns, &to_string/1)
}}
}
end
@doc """
Creates a control message to filter events.
Allows explicit whitelisting or blacklisting of events.
## Examples
iex> Snakepit.Telemetry.Control.filter(allow: ["python.call.start"])
%Snakepit.Bridge.TelemetryControl{
control: {:filter, %Snakepit.Bridge.TelemetryEventFilter{
allow: ["python.call.start"],
deny: []
}}
}
iex> Snakepit.Telemetry.Control.filter(deny: ["python.memory.sampled"])
%Snakepit.Bridge.TelemetryControl{
control: {:filter, %Snakepit.Bridge.TelemetryEventFilter{
allow: [],
deny: ["python.memory.sampled"]
}}
}
"""
def filter(opts \\ []) do
allow = Keyword.get(opts, :allow, []) |> Enum.map(&to_string/1)
deny = Keyword.get(opts, :deny, []) |> Enum.map(&to_string/1)
%TelemetryControl{
control:
{:filter,
%TelemetryEventFilter{
allow: allow,
deny: deny
}}
}
end
end
defmodule Snakepit.Telemetry.Correlation do
@moduledoc """
Utilities for generating and propagating correlation identifiers.
"""
@prefix "sp"
@doc """
Generates a new correlation identifier.
"""
@spec new_id() :: String.t()
def new_id do
random = :crypto.strong_rand_bytes(12)
encoded = Base.encode16(random, case: :lower)
"#{@prefix}-#{encoded}"
end
@doc """
Ensures a non-empty correlation identifier is present.
"""
@spec ensure(String.t() | nil) :: String.t()
def ensure(nil), do: new_id()
def ensure(""), do: new_id()
def ensure(id) when is_binary(id), do: id
def ensure(_other), do: new_id()
end
defmodule Snakepit.Telemetry.Events do
@moduledoc """
ML-specific telemetry event definitions.
Defines telemetry events for hardware detection, GPU profiling,
circuit breaker operations, and structured exceptions.
"""
@type event :: [atom()]
@type measurement_type :: :integer | :float | :monotonic_time | :system_time
@type metadata_type :: :string | :atom | :integer | :map | :list | :any
@type schema :: %{
measurements: %{atom() => measurement_type()},
metadata: %{atom() => metadata_type()}
}
@doc """
Returns all hardware-related telemetry events.
"""
@spec hardware_events() :: [event()]
def hardware_events do
[
[:snakepit, :hardware, :detect, :start],
[:snakepit, :hardware, :detect, :stop],
[:snakepit, :hardware, :select, :start],
[:snakepit, :hardware, :select, :stop],
[:snakepit, :hardware, :cache, :hit],
[:snakepit, :hardware, :cache, :miss]
]
end
@doc """
Returns all circuit breaker telemetry events.
"""
@spec circuit_breaker_events() :: [event()]
def circuit_breaker_events do
[
[:snakepit, :circuit_breaker, :opened],
[:snakepit, :circuit_breaker, :closed],
[:snakepit, :circuit_breaker, :half_open],
[:snakepit, :circuit_breaker, :call, :allowed],
[:snakepit, :circuit_breaker, :call, :rejected],
[:snakepit, :circuit_breaker, :call, :success],
[:snakepit, :circuit_breaker, :call, :failure]
]
end
@doc """
Returns all exception/error telemetry events.
"""
@spec exception_events() :: [event()]
def exception_events do
[
[:snakepit, :error, :shape_mismatch],
[:snakepit, :error, :device],
[:snakepit, :error, :oom],
[:snakepit, :error, :dtype_mismatch],
[:snakepit, :error, :dimension_error],
[:snakepit, :error, :python_exception]
]
end
@doc """
Returns all GPU profiler telemetry events.
"""
@spec gpu_profiler_events() :: [event()]
def gpu_profiler_events do
[
[:snakepit, :gpu, :memory, :sampled],
[:snakepit, :gpu, :utilization, :sampled],
[:snakepit, :gpu, :temperature, :sampled],
[:snakepit, :gpu, :power, :sampled]
]
end
@doc """
Returns all retry/backoff telemetry events.
"""
@spec retry_events() :: [event()]
def retry_events do
[
[:snakepit, :retry, :attempt],
[:snakepit, :retry, :success],
[:snakepit, :retry, :exhausted],
[:snakepit, :retry, :backoff]
]
end
@doc """
Returns all ML-related telemetry events.
This combines hardware, circuit breaker, exception, GPU profiler,
and retry events.
"""
@spec all_ml_events() :: [event()]
def all_ml_events do
hardware_events() ++
circuit_breaker_events() ++
exception_events() ++
gpu_profiler_events() ++
retry_events()
end
@doc """
Returns the schema for a given event.
Returns nil for unknown events.
"""
@spec event_schema(event()) :: schema() | nil
def event_schema(event) do
Map.get(schemas(), event)
end
@spec schemas() :: %{event() => schema()}
defp schemas do
%{
# Hardware events
[:snakepit, :hardware, :detect, :start] => %{
measurements: %{system_time: :system_time},
metadata: %{}
},
[:snakepit, :hardware, :detect, :stop] => %{
measurements: %{duration: :monotonic_time},
metadata: %{accelerator: :atom, platform: :string}
},
[:snakepit, :hardware, :select, :start] => %{
measurements: %{system_time: :system_time},
metadata: %{preference: :any}
},
[:snakepit, :hardware, :select, :stop] => %{
measurements: %{duration: :monotonic_time},
metadata: %{device: :any, success: :atom}
},
[:snakepit, :hardware, :cache, :hit] => %{
measurements: %{},
metadata: %{key: :atom}
},
[:snakepit, :hardware, :cache, :miss] => %{
measurements: %{},
metadata: %{key: :atom}
},
# Circuit breaker events
[:snakepit, :circuit_breaker, :opened] => %{
measurements: %{failure_count: :integer},
metadata: %{pool: :atom, reason: :atom}
},
[:snakepit, :circuit_breaker, :closed] => %{
measurements: %{},
metadata: %{pool: :atom}
},
[:snakepit, :circuit_breaker, :half_open] => %{
measurements: %{},
metadata: %{pool: :atom}
},
[:snakepit, :circuit_breaker, :call, :allowed] => %{
measurements: %{},
metadata: %{pool: :atom, state: :atom}
},
[:snakepit, :circuit_breaker, :call, :rejected] => %{
measurements: %{},
metadata: %{pool: :atom, state: :atom}
},
[:snakepit, :circuit_breaker, :call, :success] => %{
measurements: %{duration: :monotonic_time},
metadata: %{pool: :atom}
},
[:snakepit, :circuit_breaker, :call, :failure] => %{
measurements: %{duration: :monotonic_time},
metadata: %{pool: :atom, error: :any}
},
# Exception events
[:snakepit, :error, :shape_mismatch] => %{
measurements: %{},
metadata: %{
expected: :list,
got: :list,
dimension: :integer,
operation: :string
}
},
[:snakepit, :error, :device] => %{
measurements: %{},
metadata: %{
expected_device: :any,
actual_device: :any,
operation: :string
}
},
[:snakepit, :error, :oom] => %{
measurements: %{
requested_bytes: :integer,
available_bytes: :integer
},
metadata: %{device: :any, operation: :string}
},
[:snakepit, :error, :dtype_mismatch] => %{
measurements: %{},
metadata: %{expected: :atom, got: :atom}
},
[:snakepit, :error, :dimension_error] => %{
measurements: %{},
metadata: %{expected_dims: :integer, got_dims: :integer}
},
[:snakepit, :error, :python_exception] => %{
measurements: %{},
metadata: %{type: :string, message: :string, traceback: :string}
},
# GPU profiler events
[:snakepit, :gpu, :memory, :sampled] => %{
measurements: %{
used_mb: :integer,
total_mb: :integer,
free_mb: :integer
},
metadata: %{device: :any, utilization: :float}
},
[:snakepit, :gpu, :utilization, :sampled] => %{
measurements: %{gpu_percent: :float, memory_percent: :float},
metadata: %{device: :any}
},
[:snakepit, :gpu, :temperature, :sampled] => %{
measurements: %{celsius: :float},
metadata: %{device: :any}
},
[:snakepit, :gpu, :power, :sampled] => %{
measurements: %{watts: :float, limit_watts: :float},
metadata: %{device: :any}
},
# Retry events
[:snakepit, :retry, :attempt] => %{
measurements: %{attempt: :integer, delay_ms: :integer},
metadata: %{pool: :atom, operation: :any}
},
[:snakepit, :retry, :success] => %{
measurements: %{attempts: :integer, total_duration: :monotonic_time},
metadata: %{pool: :atom}
},
[:snakepit, :retry, :exhausted] => %{
measurements: %{attempts: :integer, total_duration: :monotonic_time},
metadata: %{pool: :atom, last_error: :any}
},
[:snakepit, :retry, :backoff] => %{
measurements: %{delay_ms: :integer},
metadata: %{pool: :atom, attempt: :integer}
}
}
end
end
defmodule Snakepit.Telemetry.GPUProfiler do
@moduledoc """
GPU memory and utilization profiler.
Periodically samples GPU metrics and emits telemetry events.
Supports NVIDIA CUDA GPUs via nvidia-smi.
"""
use GenServer
require Logger
alias Snakepit.Hardware
@default_interval_ms 5_000
@min_interval_ms 100
@type state :: %{
interval_ms: pos_integer(),
enabled: boolean(),
sample_count: non_neg_integer(),
last_sample_time: integer() | nil,
timer_ref: reference() | nil,
devices: [Hardware.Selector.device()]
}
# Client API
@doc """
Starts the GPU profiler.
## Options
- `:interval_ms` - Sampling interval in milliseconds (default: 5000)
- `:enabled` - Whether to start sampling immediately (default: true)
- `:name` - GenServer name (default: __MODULE__)
"""
@spec start_link(keyword()) :: GenServer.on_start()
def start_link(opts \\ []) do
name = Keyword.get(opts, :name, __MODULE__)
GenServer.start_link(__MODULE__, opts, name: name)
end
@doc """
Triggers an immediate GPU sample.
"""
@spec sample_now(GenServer.server()) :: :ok | {:error, :no_gpu}
def sample_now(server \\ __MODULE__) do
GenServer.call(server, :sample_now)
end
@doc """
Returns profiler statistics.
"""
@spec get_stats(GenServer.server()) :: map()
def get_stats(server \\ __MODULE__) do
GenServer.call(server, :get_stats)
end
@doc """
Enables GPU sampling.
"""
@spec enable(GenServer.server()) :: :ok
def enable(server \\ __MODULE__) do
GenServer.call(server, :enable)
end
@doc """
Disables GPU sampling.
"""
@spec disable(GenServer.server()) :: :ok
def disable(server \\ __MODULE__) do
GenServer.call(server, :disable)
end
@doc """
Updates the sampling interval.
"""
@spec set_interval(GenServer.server(), pos_integer()) :: :ok | {:error, :invalid_interval}
def set_interval(server \\ __MODULE__, interval_ms) do
GenServer.call(server, {:set_interval, interval_ms})
end
# GenServer callbacks
@impl true
def init(opts) do
interval_ms = Keyword.get(opts, :interval_ms, @default_interval_ms)
enabled = Keyword.get(opts, :enabled, true)
state = %{
interval_ms: interval_ms,
enabled: enabled,
sample_count: 0,
last_sample_time: nil,
timer_ref: nil,
devices: detect_gpu_devices()
}
state =
if enabled and state.devices != [] do
schedule_sample(state)
else
state
end
{:ok, state}
end
@impl true
def handle_call(:sample_now, _from, state) do
case do_sample(state) do
{:ok, new_state} ->
{:reply, :ok, new_state}
{:error, reason} ->
{:reply, {:error, reason}, state}
end
end
def handle_call(:get_stats, _from, state) do
stats = %{
interval_ms: state.interval_ms,
enabled: state.enabled,
sample_count: state.sample_count,
last_sample_time: state.last_sample_time,
device_count: length(state.devices)
}
{:reply, stats, state}
end
def handle_call(:enable, _from, state) do
state = %{state | enabled: true}
state = if state.devices != [], do: schedule_sample(state), else: state
{:reply, :ok, state}
end
def handle_call(:disable, _from, state) do
state = cancel_timer(state)
state = %{state | enabled: false}
{:reply, :ok, state}
end
def handle_call({:set_interval, interval_ms}, _from, state)
when interval_ms >= @min_interval_ms do
state = cancel_timer(state)
state = %{state | interval_ms: interval_ms}
state = if state.enabled and state.devices != [], do: schedule_sample(state), else: state
{:reply, :ok, state}
end
def handle_call({:set_interval, _}, _from, state) do
{:reply, {:error, :invalid_interval}, state}
end
@impl true
def handle_info(:sample, state) do
state =
case do_sample(state) do
{:ok, new_state} -> new_state
{:error, _} -> state
end
state = if state.enabled, do: schedule_sample(state), else: state
{:noreply, state}
end
def handle_info(_msg, state) do
{:noreply, state}
end
# Private functions
defp detect_gpu_devices do
info = Hardware.detect()
case info.cuda do
%{devices: devices} when is_list(devices) and devices != [] ->
Enum.map(devices, fn d -> {:cuda, d.id} end)
_ ->
[]
end
end
defp schedule_sample(state) do
state = cancel_timer(state)
ref = Process.send_after(self(), :sample, state.interval_ms)
%{state | timer_ref: ref}
end
defp cancel_timer(%{timer_ref: nil} = state), do: state
defp cancel_timer(%{timer_ref: ref} = state) do
Process.cancel_timer(ref)
%{state | timer_ref: nil}
end
defp do_sample(%{devices: []} = _state) do
{:error, :no_gpu}
end
defp do_sample(state) do
now = System.monotonic_time(:millisecond)
Enum.each(state.devices, fn device ->
case sample_device(device) do
{:ok, metrics} ->
emit_metrics(device, metrics)
{:error, _reason} ->
:ok
end
end)
new_state = %{
state
| sample_count: state.sample_count + 1,
last_sample_time: now
}
{:ok, new_state}
end
defp sample_device({:cuda, device_id}) do
query = [
"--id=#{device_id}",
"--query-gpu=memory.used,memory.total,memory.free,utilization.gpu,temperature.gpu,power.draw",
"--format=csv,noheader,nounits"
]
case System.cmd("nvidia-smi", query, stderr_to_stdout: true) do
{output, 0} ->
parse_nvidia_smi_output(output)
_ ->
{:error, :nvidia_smi_failed}
end
rescue
_ -> {:error, :nvidia_smi_not_found}
end
defp sample_device(_), do: {:error, :unsupported_device}
defp parse_nvidia_smi_output(output) do
case String.split(String.trim(output), ", ") do
[used, total, free, gpu_util, temp, power] ->
{:ok,
%{
memory_used_mb: parse_int(used),
memory_total_mb: parse_int(total),
memory_free_mb: parse_int(free),
gpu_utilization: parse_float(gpu_util),
temperature: parse_float(temp),
power_watts: parse_float(power)
}}
_ ->
{:error, :parse_error}
end
end
defp parse_int(str) do
case Integer.parse(String.trim(str)) do
{n, _} -> n
:error -> 0
end
end
defp parse_float(str) do
case Float.parse(String.trim(str)) do
{f, _} -> f
:error -> 0.0
end
end
defp emit_metrics(device, metrics) do
# Memory event
:telemetry.execute(
[:snakepit, :gpu, :memory, :sampled],
%{
used_mb: metrics.memory_used_mb,
total_mb: metrics.memory_total_mb,
free_mb: metrics.memory_free_mb
},
%{
device: device,
utilization: metrics.memory_used_mb / max(metrics.memory_total_mb, 1)
}
)
# Utilization event
:telemetry.execute(
[:snakepit, :gpu, :utilization, :sampled],
%{
gpu_percent: metrics.gpu_utilization,
memory_percent: metrics.memory_used_mb / max(metrics.memory_total_mb, 1) * 100
},
%{device: device}
)
# Temperature event
if metrics.temperature > 0 do
:telemetry.execute(
[:snakepit, :gpu, :temperature, :sampled],
%{celsius: metrics.temperature},
%{device: device}
)
end
# Power event
if metrics.power_watts > 0 do
:telemetry.execute(
[:snakepit, :gpu, :power, :sampled],
%{watts: metrics.power_watts, limit_watts: 0.0},
%{device: device}
)
end
end
end
defmodule Snakepit.Telemetry.GrpcStream do
@moduledoc """
Manages gRPC telemetry streams from Python workers.
This GenServer maintains bidirectional telemetry streams with Python workers,
translating Python telemetry events into Elixir `:telemetry` events.
Features:
- Automatic stream registration when workers connect
- Dynamic sampling rate adjustments
- Event filtering
- Graceful handling of worker disconnections
"""
use GenServer
alias Snakepit.Logger, as: SLog
alias GRPC.Channel
alias Snakepit.Bridge.{BridgeService, TelemetryEvent}
alias Snakepit.Telemetry.{Control, Naming, SafeMetadata}
@log_category :telemetry
@type worker_ctx :: %{
worker_id: String.t(),
pool_name: atom(),
python_pid: integer() | nil
}
## Client API
@doc """
Starts the telemetry stream manager.
"""
def start_link(opts) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Registers a worker for telemetry streaming.
Automatically initiates a telemetry stream with the worker and starts
consuming events.
## Examples
iex> channel = connect_to_worker()
iex> Snakepit.Telemetry.GrpcStream.register_worker(channel, %{
...> worker_id: "worker_1",
...> pool_name: :default,
...> python_pid: 12345
...> })
:ok
"""
def register_worker(channel, worker_ctx) do
if stream_capable_channel?(channel) do
GenServer.cast(__MODULE__, {:register_worker, channel, worker_ctx})
else
SLog.debug(
@log_category,
"Skipping telemetry stream registration; channel unsupported",
worker_id: worker_ctx.worker_id,
pool_name: worker_ctx.pool_name,
channel_type: describe_channel(channel)
)
:ok
end
end
@doc """
Removes a worker from telemetry streaming.
Called when a worker disconnects or terminates.
"""
def unregister_worker(worker_id) do
GenServer.cast(__MODULE__, {:unregister_worker, worker_id})
end
@doc """
Updates the sampling rate for a specific worker.
## Examples
iex> Snakepit.Telemetry.GrpcStream.update_sampling("worker_1", 0.1)
:ok
iex> Snakepit.Telemetry.GrpcStream.update_sampling("worker_1", 0.5, ["python.call.*"])
:ok
"""
def update_sampling(worker_id, rate, patterns \\ []) do
GenServer.cast(__MODULE__, {:update_sampling, worker_id, rate, patterns})
end
@doc """
Enables or disables telemetry for a specific worker.
"""
def toggle(worker_id, enabled) do
GenServer.cast(__MODULE__, {:toggle, worker_id, enabled})
end
@doc """
Updates event filters for a specific worker.
"""
def update_filter(worker_id, opts) do
GenServer.cast(__MODULE__, {:update_filter, worker_id, opts})
end
@doc """
Gets the current state of all registered streams.
"""
def list_streams do
case Process.whereis(__MODULE__) do
nil -> []
pid -> GenServer.call(pid, :list_streams)
end
end
## Server Callbacks
@impl true
def init(_opts) do
{:ok, %{streams: %{}}}
end
@impl true
def handle_cast({:register_worker, channel, worker_ctx}, state) do
case initiate_stream(channel, worker_ctx) do
{:ok, stream_info} ->
new_state = put_in(state, [:streams, worker_ctx.worker_id], stream_info)
SLog.info(
@log_category,
"Telemetry stream registered for worker #{worker_ctx.worker_id}",
worker_id: worker_ctx.worker_id,
pool_name: worker_ctx.pool_name
)
{:noreply, new_state}
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to register telemetry stream for worker #{worker_ctx.worker_id}: #{inspect(reason)}",
worker_id: worker_ctx.worker_id,
reason: reason
)
{:noreply, state}
end
end
@impl true
def handle_cast({:unregister_worker, worker_id}, state) do
case Map.get(state.streams, worker_id) do
nil ->
{:noreply, state}
stream_info ->
# Cancel the consumer task
if stream_info.task && Process.alive?(stream_info.task.pid) do
Task.shutdown(stream_info.task, :brutal_kill)
end
new_state = update_in(state, [:streams], &Map.delete(&1, worker_id))
SLog.debug(@log_category, "Telemetry stream unregistered for worker #{worker_id}",
worker_id: worker_id
)
{:noreply, new_state}
end
end
@impl true
def handle_cast({:update_sampling, worker_id, rate, patterns}, state) do
case Map.get(state.streams, worker_id) do
nil ->
SLog.debug(@log_category, "Cannot update sampling for unknown worker #{worker_id}")
{:noreply, state}
%{stream: stream} ->
control_msg = Control.sampling(rate, patterns)
new_state =
case send_control_request(stream, control_msg) do
{:ok, updated_stream} ->
SLog.debug(@log_category, "Updated sampling for worker #{worker_id} to #{rate}",
worker_id: worker_id,
rate: rate,
patterns: patterns
)
put_in(state, [:streams, worker_id, :stream], updated_stream)
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to update sampling for worker #{worker_id}: #{inspect(reason)}",
worker_id: worker_id,
reason: reason
)
state
end
{:noreply, new_state}
end
end
@impl true
def handle_cast({:toggle, worker_id, enabled}, state) do
case Map.get(state.streams, worker_id) do
nil ->
{:noreply, state}
%{stream: stream} ->
control_msg = Control.toggle(enabled)
new_state =
case send_control_request(stream, control_msg) do
{:ok, updated_stream} ->
SLog.debug(@log_category, "Toggled telemetry for worker #{worker_id} to #{enabled}",
worker_id: worker_id,
enabled: enabled
)
put_in(state, [:streams, worker_id, :stream], updated_stream)
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to toggle telemetry for worker #{worker_id}: #{inspect(reason)}",
worker_id: worker_id,
reason: reason
)
state
end
{:noreply, new_state}
end
end
@impl true
def handle_cast({:update_filter, worker_id, opts}, state) do
case Map.get(state.streams, worker_id) do
nil ->
{:noreply, state}
%{stream: stream} ->
control_msg = Control.filter(opts)
new_state =
case send_control_request(stream, control_msg) do
{:ok, updated_stream} ->
SLog.debug(@log_category, "Updated filters for worker #{worker_id}",
worker_id: worker_id
)
put_in(state, [:streams, worker_id, :stream], updated_stream)
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to update filters for worker #{worker_id}: #{inspect(reason)}",
worker_id: worker_id,
reason: reason
)
state
end
{:noreply, new_state}
end
end
@impl true
def handle_call(:list_streams, _from, state) do
stream_info =
Enum.map(state.streams, fn {worker_id, info} ->
%{
worker_id: worker_id,
pool_name: info.worker_ctx.pool_name,
task_alive: info.task && Process.alive?(info.task.pid)
}
end)
{:reply, stream_info, state}
end
@impl true
def handle_info({ref, :stream_completed}, state) when is_reference(ref) do
# Task completed successfully
Process.demonitor(ref, [:flush])
{:noreply, state}
end
@impl true
def handle_info({:DOWN, _ref, :process, _pid, reason}, state) do
# Task crashed or was killed
SLog.debug(@log_category, "Telemetry stream consumer task terminated: #{inspect(reason)}")
{:noreply, state}
end
@impl true
def handle_info({:gun_response, _pid, _stream_ref, _fin, status, headers}, state) do
SLog.debug(@log_category, "Telemetry stream HTTP response received",
status: status,
headers: headers
)
{:noreply, state}
end
@impl true
def handle_info({:gun_data, _pid, _stream_ref, _is_fin, _data}, state) do
# gRPC data frames are consumed by GRPC.Stub.recv/2; ignore low-level messages.
{:noreply, state}
end
@impl true
def handle_info({:gun_down, _pid, _proto, _reason, _killed_streams, _}, state) do
SLog.debug(@log_category, "Telemetry stream HTTP connection closed by gun")
{:noreply, state}
end
@impl true
def handle_info({:gun_error, _pid, _stream_ref, reason}, state) do
SLog.debug(@log_category, "Telemetry stream HTTP error from gun", reason: reason)
{:noreply, state}
end
@impl true
def handle_info(_message, state) do
{:noreply, state}
end
## Private Helpers
defp send_control_request(stream, control_msg) do
{:ok, GRPC.Stub.send_request(stream, control_msg)}
rescue
error ->
{:error, error}
catch
:exit, reason ->
{:error, reason}
end
defp initiate_stream(channel, worker_ctx) do
# Use longer timeout for stream operations
case channel |> open_telemetry_stream() |> normalize_stream_response() do
{:ok, stream} ->
# Send initial toggle message to enable telemetry
stream = GRPC.Stub.send_request(stream, Control.toggle(true))
# Start async task to consume events
task =
Task.Supervisor.async_nolink(
Snakepit.TaskSupervisor,
fn -> consume_stream(stream, worker_ctx) end
)
{:ok,
%{
stream: stream,
task: task,
worker_ctx: worker_ctx,
started_at: System.monotonic_time()
}}
{:error, reason} ->
{:error, reason}
end
end
defp open_telemetry_stream(channel) do
BridgeService.Stub.stream_telemetry(channel, timeout: :infinity)
rescue
exception ->
{:error, {:invalid_channel, exception}}
end
defp normalize_stream_response({:ok, %GRPC.Client.Stream{} = stream}), do: {:ok, stream}
defp normalize_stream_response(%GRPC.Client.Stream{} = stream), do: {:ok, stream}
defp normalize_stream_response({:error, _reason} = error), do: error
defp normalize_stream_response(other), do: {:error, {:unexpected_stream_response, other}}
defp consume_stream(stream, worker_ctx) do
case GRPC.Stub.recv(stream, timeout: :infinity) do
{:ok, enum} ->
Enum.each(enum, fn
{:ok, %TelemetryEvent{} = event} ->
translate_and_emit(event, worker_ctx)
{:error, reason} ->
SLog.warning(
@log_category,
"Telemetry stream error for worker #{worker_ctx.worker_id}: #{inspect(reason)}",
worker_id: worker_ctx.worker_id,
reason: reason
)
{:trailers, trailers} ->
SLog.debug(@log_category, "Telemetry stream trailers: #{inspect(trailers)}",
worker_id: worker_ctx.worker_id
)
end)
SLog.debug(@log_category, "Telemetry stream completed for worker #{worker_ctx.worker_id}",
worker_id: worker_ctx.worker_id
)
{:error, reason} ->
log_stream_closed(worker_ctx, reason)
end
end
defp stream_capable_channel?(%Channel{}), do: true
defp stream_capable_channel?(%{__struct__: module}) when is_atom(module) do
String.starts_with?(Atom.to_string(module), "GRPC.")
end
defp stream_capable_channel?(_), do: false
defp describe_channel(%Channel{}), do: "GRPC.Channel"
defp describe_channel(%{__struct__: module}) when is_atom(module), do: Atom.to_string(module)
defp describe_channel(channel) when is_reference(channel), do: "reference"
defp describe_channel(channel) when is_pid(channel), do: "pid"
defp describe_channel(channel) when is_map(channel), do: "map"
defp describe_channel(channel) when is_binary(channel), do: "binary"
defp describe_channel(channel) when is_list(channel), do: "list"
defp describe_channel(channel), do: inspect(channel)
defp log_stream_closed(worker_ctx, reason) do
log_fun =
if shutdown_reason?(reason), do: &SLog.debug/3, else: &SLog.warning/3
log_fun.(
@log_category,
"Telemetry stream closed for worker #{worker_ctx.worker_id}: #{inspect(reason)}",
worker_id: worker_ctx.worker_id,
reason: reason
)
end
defp shutdown_reason?(%GRPC.RPCError{message: message}) when is_binary(message) do
message
|> String.downcase()
|> String.contains?("shutdown")
end
defp shutdown_reason?(%GRPC.RPCError{}), do: false
defp shutdown_reason?(:shutdown), do: true
defp shutdown_reason?({:shutdown, _}), do: true
defp shutdown_reason?({:down, :shutdown}), do: true
defp shutdown_reason?({:error, reason}), do: shutdown_reason?(reason)
defp shutdown_reason?(_), do: false
defp translate_and_emit(event, worker_ctx) do
with {:ok, event_name} <- Naming.from_parts(event.event_parts),
{:ok, measurements} <- cast_measurements(event.measurements),
{:ok, metadata} <-
SafeMetadata.enrich(event.metadata,
node: node(),
worker_id: worker_ctx.worker_id,
pool_name: worker_ctx.pool_name,
python_pid: worker_ctx.python_pid,
correlation_id: blank_to_nil(event.correlation_id),
timestamp_ns: event.timestamp_ns
) do
:telemetry.execute(event_name, measurements, metadata)
else
{:error, reason} ->
SLog.debug(
@log_category,
"Skipping telemetry event #{inspect(event.event_parts)}: #{inspect(reason)}",
worker_id: worker_ctx.worker_id,
event_parts: event.event_parts,
reason: reason
)
end
end
defp cast_measurements(measurements) do
Enum.reduce_while(measurements, {:ok, %{}}, fn {key, value}, {:ok, acc} ->
case Naming.measurement_key(key) do
{:ok, atom_key} ->
val = extract_measurement_value(value.value)
{:cont, {:ok, Map.put(acc, atom_key, val)}}
{:error, reason} ->
{:halt, {:error, {:invalid_measurement_key, key, reason}}}
end
end)
end
defp extract_measurement_value({:int_value, v}), do: v
defp extract_measurement_value({:float_value, v}), do: v
defp extract_measurement_value({:string_value, v}), do: v
defp extract_measurement_value(nil), do: nil
defp blank_to_nil(""), do: nil
defp blank_to_nil(value), do: value
end
defmodule Snakepit.Telemetry.Naming do
@moduledoc """
Event catalog and naming validation for Snakepit telemetry.
This module ensures atom safety by maintaining a curated catalog of all
valid telemetry events and measurement keys. Python-originated events
must pass through this module to prevent arbitrary atom creation.
## Python Event Catalog
`python_event_catalog/0` lists the event strings emitted by `snakepit_bridge`
and the measurement keys they are expected to use. When adding a new Python
telemetry event, update that catalog and the allowlist in
`snakepit_bridge.telemetry.stream` together so both languages agree on the
schema.
"""
# Layer 1: Infrastructure Events (Elixir-originated)
@pool_events [
:initialized,
:status,
:queue_enqueued,
:queue_dequeued,
:queue_timeout,
:worker_spawn_started,
:worker_spawned,
:worker_spawn_failed,
:worker_terminated,
:worker_restarted
]
@session_events [
:created,
:destroyed,
:affinity_assigned,
:affinity_broken
]
# Layer 2: Python Execution Events (Python-originated, folded by Elixir)
@python_events [
:call_start,
:call_stop,
:call_exception,
:memory_sampled,
:cpu_sampled,
:gc_completed,
:error_occurred
]
# Layer 3: gRPC Bridge Events (Elixir-originated)
@grpc_events [
:call_start,
:call_stop,
:call_exception,
:stream_opened,
:stream_message,
:stream_closed,
:connection_established,
:connection_lost,
:connection_reconnected
]
@python_event_catalog [
%{
name: "python.call.start",
event: [:snakepit, :python, :call, :start],
measurements: [:count]
},
%{
name: "python.call.stop",
event: [:snakepit, :python, :call, :stop],
measurements: [:duration, :count]
},
%{
name: "python.call.exception",
event: [:snakepit, :python, :call, :exception],
measurements: [:count]
},
%{
name: "python.memory.sampled",
event: [:snakepit, :python, :memory, :sampled],
measurements: [:rss_bytes, :vms_bytes]
},
%{
name: "python.cpu.sampled",
event: [:snakepit, :python, :cpu, :sampled],
measurements: [:cpu_percent]
},
%{
name: "python.gc.completed",
event: [:snakepit, :python, :gc, :completed],
measurements: [:count, :generation]
},
%{
name: "python.error.occurred",
event: [:snakepit, :python, :error, :occurred],
measurements: [:count]
},
%{
name: "tool.execution.start",
event: [:snakepit, :python, :tool, :execution, :start],
measurements: [:count]
},
%{
name: "tool.execution.stop",
event: [:snakepit, :python, :tool, :execution, :stop],
measurements: [:duration, :count]
},
%{
name: "tool.execution.exception",
event: [:snakepit, :python, :tool, :execution, :exception],
measurements: [:count]
},
%{
name: "tool.result_size",
event: [:snakepit, :python, :tool, :result_size],
measurements: [:message_size]
}
]
# Valid measurement keys (atom-safe)
@measurement_keys [
:duration,
:duration_ms,
:bytes,
:system_time,
:queue_depth,
:queue_time,
:available_workers,
:busy_workers,
:total_workers,
:worker_count,
:lifetime,
:total_commands,
:command_count,
:downtime,
:restart_count,
:retry_count,
:request_size,
:response_size,
:message_size,
:sequence_number,
:message_count,
:network_time,
:uptime,
:call_count,
:affinity_duration,
:commands_with_affinity,
:rss_bytes,
:vms_bytes,
:cpu_percent,
:collected,
:generation,
:latency_ms,
:count,
:python_pid
]
@doc """
Convert Python event parts to a valid Elixir telemetry event name.
Returns `{:ok, event_name}` if the parts map to a known event,
`{:error, reason}` otherwise.
## Examples
iex> Snakepit.Telemetry.Naming.from_parts(["python", "call", "start"])
{:ok, [:snakepit, :python, :call, :start]}
iex> Snakepit.Telemetry.Naming.from_parts(["unknown", "event"])
{:error, :unknown_event}
"""
def from_parts(["python", "call", "start"]) do
{:ok, [:snakepit, :python, :call, :start]}
end
def from_parts(["python", "call", "stop"]) do
{:ok, [:snakepit, :python, :call, :stop]}
end
def from_parts(["python", "call", "exception"]) do
{:ok, [:snakepit, :python, :call, :exception]}
end
def from_parts(["python", "memory", "sampled"]) do
{:ok, [:snakepit, :python, :memory, :sampled]}
end
def from_parts(["python", "cpu", "sampled"]) do
{:ok, [:snakepit, :python, :cpu, :sampled]}
end
def from_parts(["python", "gc", "completed"]) do
{:ok, [:snakepit, :python, :gc, :completed]}
end
def from_parts(["python", "error", "occurred"]) do
{:ok, [:snakepit, :python, :error, :occurred]}
end
def from_parts(["tool", "execution", action]) when action in ["start", "stop", "exception"] do
{:ok, [:snakepit, :python, :tool, :execution, String.to_existing_atom(action)]}
rescue
ArgumentError -> {:error, :invalid_atom}
end
def from_parts(["tool", "result_size"]) do
{:ok, [:snakepit, :python, :tool, :result_size]}
end
def from_parts(parts) when is_list(parts) do
{:error, :unknown_event}
end
@doc """
Validate a measurement key and convert to atom if it's in the allowlist.
## Examples
iex> Snakepit.Telemetry.Naming.measurement_key("duration")
{:ok, :duration}
iex> Snakepit.Telemetry.Naming.measurement_key("unknown_key")
{:error, :unknown_measurement_key}
"""
def measurement_key(key) when is_binary(key) do
atom_key = String.to_existing_atom(key)
if atom_key in @measurement_keys do
{:ok, atom_key}
else
{:error, :unknown_measurement_key}
end
rescue
ArgumentError -> {:error, :invalid_atom}
end
def measurement_key(key) when is_atom(key) do
if key in @measurement_keys do
{:ok, key}
else
{:error, :unknown_measurement_key}
end
end
@doc """
Get all valid pool events.
"""
def pool_events, do: @pool_events
@doc """
Get all valid session events.
"""
def session_events, do: @session_events
@doc """
Get all valid Python events.
"""
def python_events, do: @python_events
@doc """
Return the catalog describing Python event names, their telemetry atoms, and expected measurements.
"""
def python_event_catalog, do: @python_event_catalog
@doc """
Get all valid gRPC events.
"""
def grpc_events, do: @grpc_events
@doc """
Get all valid measurement keys.
"""
def measurement_keys, do: @measurement_keys
@doc """
Build an event name from components.
## Examples
iex> Snakepit.Telemetry.Naming.event(:pool, :worker, :spawned)
[:snakepit, :pool, :worker, :spawned]
"""
def event(component, resource, action) do
[:snakepit, component, resource, action]
end
@doc """
Build a pool event name.
"""
def pool_event(:initialized), do: [:snakepit, :pool, :initialized]
def pool_event(:status), do: [:snakepit, :pool, :status]
def pool_event(:queue_enqueued), do: [:snakepit, :pool, :queue, :enqueued]
def pool_event(:queue_dequeued), do: [:snakepit, :pool, :queue, :dequeued]
def pool_event(:queue_timeout), do: [:snakepit, :pool, :queue, :timeout]
def pool_event(:worker_spawn_started), do: [:snakepit, :pool, :worker, :spawn_started]
def pool_event(:worker_spawned), do: [:snakepit, :pool, :worker, :spawned]
def pool_event(:worker_spawn_failed), do: [:snakepit, :pool, :worker, :spawn_failed]
def pool_event(:worker_terminated), do: [:snakepit, :pool, :worker, :terminated]
def pool_event(:worker_restarted), do: [:snakepit, :pool, :worker, :restarted]
@doc """
Build a session event name.
"""
def session_event(action) when action in @session_events do
case action do
:created -> [:snakepit, :session, :created]
:destroyed -> [:snakepit, :session, :destroyed]
:affinity_assigned -> [:snakepit, :session, :affinity, :assigned]
:affinity_broken -> [:snakepit, :session, :affinity, :broken]
end
end
@doc """
Build a Python event name.
"""
def python_event(action) when action in @python_events do
case action do
:call_start -> [:snakepit, :python, :call, :start]
:call_stop -> [:snakepit, :python, :call, :stop]
:call_exception -> [:snakepit, :python, :call, :exception]
:memory_sampled -> [:snakepit, :python, :memory, :sampled]
:cpu_sampled -> [:snakepit, :python, :cpu, :sampled]
:gc_completed -> [:snakepit, :python, :gc, :completed]
:error_occurred -> [:snakepit, :python, :error, :occurred]
end
end
@doc """
Build a gRPC event name.
"""
def grpc_event(resource, action) when action in @grpc_events do
[:snakepit, :grpc, resource, action]
end
end
defmodule Snakepit.Telemetry.OpenTelemetry do
@moduledoc """
Bootstraps OpenTelemetry tracing and telemetry bridges for Snakepit.
When enabled via `:snakepit, :opentelemetry` configuration this module ensures
the OpenTelemetry runtime is started, exporters are configured, and telemetry
events are mapped to spans and span events. Exporters remain opt-in; by default
spans are created but not shipped anywhere.
"""
require OpenTelemetry.Tracer, as: Tracer
alias OpenTelemetry.Span
alias OpentelemetryTelemetry, as: OTelBridge
alias Snakepit.Logger, as: SLog
@grpc_handler_id "snakepit-otel-grpc-worker"
@heartbeat_handler_id "snakepit-otel-heartbeat"
@log_category :telemetry
@doc """
Configures OpenTelemetry and attaches telemetry handlers when enabled.
"""
@spec setup() :: :ok
def setup do
config = load_config()
if config.enabled do
configure_resource(config)
case maybe_ensure_runtime(config) do
:ok ->
attach_grpc_handlers(config)
attach_heartbeat_handlers(config)
:ok
{:error, reason} ->
SLog.warning(
@log_category,
"OpenTelemetry runtime unavailable (#{inspect(reason)}); continuing without spans"
)
:ok
end
else
detach(@grpc_handler_id)
detach(@heartbeat_handler_id)
:ok
end
end
defp ensure_runtime(config) do
with :ok <- start_app(:opentelemetry),
:ok <- maybe_start_exporter_app(config) do
configure_exporter(config.exporters)
end
rescue
error -> {:error, error}
catch
kind, reason -> {:error, {kind, reason}}
end
defp start_app(app) do
case Application.ensure_all_started(app) do
{:ok, _apps} -> :ok
{:error, {:already_started, _}} -> :ok
{:error, reason} -> {:error, reason}
end
end
defp maybe_start_exporter_app(%{exporters: %{otlp: %{enabled: true}}}) do
start_app(:opentelemetry_exporter)
end
defp maybe_start_exporter_app(_config), do: :ok
defp configure_exporter(%{otlp: %{enabled: true} = otlp}) do
opts = otlp_exporter_opts(otlp)
:otel_batch_processor.set_exporter(:opentelemetry_exporter, opts)
:ok
end
defp configure_exporter(%{console: %{enabled: true}}) do
:otel_batch_processor.set_exporter(:otel_exporter_stdout, %{})
:ok
end
defp configure_exporter(_), do: :ok
defp otlp_exporter_opts(otlp) do
endpoint =
otlp
|> Map.get(:endpoint, "http://localhost:4318")
|> to_string()
headers =
otlp
|> Map.get(:headers, [])
|> Enum.map(fn
{k, v} -> {to_string(k), to_string(v)}
other -> other
end)
protocol =
otlp
|> Map.get(:protocol, :http_protobuf)
|> normalize_protocol()
opts = %{endpoints: [endpoint], protocol: protocol}
if headers != [] do
Map.put(opts, :headers, headers)
else
opts
end
end
defp normalize_protocol(value) when value in [:http_protobuf, "http_protobuf"],
do: :http_protobuf
defp normalize_protocol(value) when value in [:grpc, "grpc"], do: :grpc
defp normalize_protocol(_), do: :http_protobuf
defp configure_resource(%{resource: resource}) when resource != %{} do
Application.put_env(:opentelemetry, :resource, resource)
end
defp configure_resource(_), do: :ok
defp attach_grpc_handlers(config) do
detach(@grpc_handler_id)
events = [
[:snakepit, :grpc_worker, :execute, :start],
[:snakepit, :grpc_worker, :execute, :stop],
[:snakepit, :grpc_worker, :execute, :exception]
]
handler_config = %{tracer_id: config.tracer_id, debug_pid: Map.get(config, :debug_pid)}
case :telemetry.attach_many(
@grpc_handler_id,
events,
&handle_grpc_worker_event/4,
handler_config
) do
:ok -> :ok
{:error, :already_exists} -> :ok
end
rescue
error ->
SLog.error(
@log_category,
"Failed to attach OpenTelemetry handler for GRPC worker: #{Exception.message(error)}"
)
:ok
end
defp attach_heartbeat_handlers(config) do
detach(@heartbeat_handler_id)
heartbeat_events = [
[:snakepit, :heartbeat, :ping_sent],
[:snakepit, :heartbeat, :pong_received],
[:snakepit, :heartbeat, :monitor_failure],
[:snakepit, :heartbeat, :heartbeat_timeout]
]
case :telemetry.attach_many(
@heartbeat_handler_id,
heartbeat_events,
&handle_heartbeat_event/4,
config
) do
:ok -> :ok
{:error, :already_exists} -> :ok
end
rescue
error ->
SLog.warning(
@log_category,
"Failed to attach heartbeat OpenTelemetry handler: #{Exception.message(error)}"
)
:ok
end
defp handle_grpc_worker_event(
[:snakepit, :grpc_worker, :execute, :start],
measurements,
metadata,
%{tracer_id: tracer_id} = handler_config
) do
attributes = span_attributes(metadata)
start_opts = %{start_time: measurements.monotonic_time, attributes: attributes, kind: :client}
OTelBridge.start_telemetry_span(tracer_id, span_name(metadata), metadata, start_opts)
debug(handler_config, {:grpc_execute_start, metadata})
:ok
rescue
error ->
SLog.debug(@log_category, "OpenTelemetry start handler failed: #{Exception.message(error)}")
:ok
end
defp handle_grpc_worker_event(
[:snakepit, :grpc_worker, :execute, :stop],
measurements,
metadata,
%{tracer_id: tracer_id} = handler_config
) do
ctx = OTelBridge.set_current_telemetry_span(tracer_id, metadata)
duration_ms = Map.get(measurements, :duration_ms)
Tracer.set_attributes(stop_attributes(measurements, metadata))
maybe_set_status(ctx, metadata)
if duration_ms do
Tracer.add_event("snakepit.grpc.duration", duration: duration_ms)
end
OTelBridge.end_telemetry_span(tracer_id, metadata)
debug(handler_config, {:grpc_execute_stop, metadata, measurements})
:ok
rescue
error ->
SLog.debug(@log_category, "OpenTelemetry stop handler failed: #{Exception.message(error)}")
:ok
end
defp handle_grpc_worker_event(
[:snakepit, :grpc_worker, :execute, :exception],
measurements,
metadata,
%{tracer_id: tracer_id} = handler_config
) do
ctx = OTelBridge.set_current_telemetry_span(tracer_id, metadata)
kind = Map.get(metadata, :kind, :error)
reason = Map.get(metadata, :reason, :unknown)
stacktrace = Map.get(metadata, :stacktrace, [])
_ = Map.get(measurements, :duration)
Span.record_exception(ctx, kind, reason, stacktrace)
Tracer.set_status(:error, format_reason(reason))
OTelBridge.end_telemetry_span(tracer_id, metadata)
debug(handler_config, {:grpc_execute_exception, metadata, measurements})
:ok
rescue
error ->
SLog.debug(
@log_category,
"OpenTelemetry exception handler failed: #{Exception.message(error)}"
)
:ok
end
defp handle_grpc_worker_event(_event, _measurements, _metadata, _config), do: :ok
defp handle_heartbeat_event([:snakepit, :heartbeat, event], measurements, metadata, _config) do
attributes = heartbeat_attributes(event, measurements, metadata)
name = "snakepit.heartbeat.#{event}"
case Tracer.current_span_ctx() do
:undefined -> :ok
_ctx -> Tracer.add_event(name, attributes)
end
rescue
error ->
SLog.debug(
@log_category,
"Heartbeat OpenTelemetry event failed: #{Exception.message(error)}"
)
:ok
end
defp span_name(metadata) do
command = metadata[:command] || "execute"
"snakepit.grpc.#{command}"
end
defp span_attributes(metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]},
{"snakepit.worker.pid", pid_attribute(metadata[:worker_pid])},
{"snakepit.pool", pool_attribute(metadata[:pool])},
{"snakepit.command", metadata[:command]},
{"snakepit.session_id", metadata[:session_id]},
{"snakepit.correlation_id", metadata[:correlation_id]},
{"snakepit.telemetry.operation", metadata[:operation]}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp stop_attributes(measurements, metadata) do
[
{"snakepit.grpc.duration_ms", measurements[:duration_ms]},
{"snakepit.grpc.status", metadata[:status]},
{"snakepit.grpc.error", format_reason(metadata[:error])},
{"snakepit.grpc.error_kind", metadata[:error_kind]},
{"snakepit.grpc.executions", measurements[:executions]},
{"snakepit.grpc.errors", measurements[:errors]}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp heartbeat_attributes(:pong_received, measurements, metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]},
{"snakepit.heartbeat.latency_ms", measurements[:latency_ms]},
{"snakepit.heartbeat.missed", metadata[:missed_heartbeats]}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp heartbeat_attributes(:heartbeat_timeout, measurements, metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]},
{"snakepit.heartbeat.missed", measurements[:missed]},
{"snakepit.heartbeat.timeouts", measurements[:timeouts]}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp heartbeat_attributes(:monitor_failure, _measurements, metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]},
{"snakepit.heartbeat.failure_reason", format_reason(metadata[:failure_reason])}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp heartbeat_attributes(:ping_sent, _measurements, metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]}
]
end
defp heartbeat_attributes(_event, _measurements, metadata) do
[
{"snakepit.worker.id", metadata[:worker_id]}
]
end
defp maybe_set_status(ctx, %{status: :error, error: error}) do
Span.set_status(ctx, OpenTelemetry.status(:error, format_reason(error)))
end
defp maybe_set_status(_ctx, _metadata), do: :ok
defp pid_attribute(nil), do: nil
defp pid_attribute(pid) when is_pid(pid), do: inspect(pid)
defp pid_attribute(other), do: other
defp pool_attribute(nil), do: nil
defp pool_attribute(atom) when is_atom(atom), do: Atom.to_string(atom)
defp pool_attribute(other), do: inspect(other)
defp format_reason(nil), do: nil
defp format_reason({kind, reason}), do: "#{inspect(kind)}: #{inspect(reason)}"
defp format_reason(reason), do: inspect(reason)
defp detach(handler_id) do
:telemetry.detach(handler_id)
rescue
_ -> :ok
end
defp maybe_ensure_runtime(%{skip_runtime?: true}), do: :ok
defp maybe_ensure_runtime(config), do: ensure_runtime(config)
defp load_config do
defaults = %{
enabled: false,
tracer_id: :snakepit_grpc_worker,
skip_runtime?: false,
exporters: %{
otlp: %{
enabled: false,
endpoint: "http://localhost:4318",
protocol: :http_protobuf,
headers: []
},
console: %{
enabled: false
}
},
resource: %{}
}
config =
:snakepit
|> Application.get_env(:opentelemetry, %{})
|> to_map()
|> deep_merge(defaults)
|> Map.update!(:enabled, &truthy?/1)
cond do
Map.get(config, :force?, false) ->
config
Application.get_env(:snakepit, :enable_otlp?, false) ->
config
true ->
Map.put(config, :enabled, false)
end
end
defp to_map(value) when is_map(value) do
value
|> Enum.map(fn {k, v} -> {normalize_key(k), to_map(v)} end)
|> Enum.into(%{})
end
defp to_map(list) when is_list(list) do
if Keyword.keyword?(list) do
list
|> Enum.map(fn {k, v} -> {normalize_key(k), to_map(v)} end)
|> Enum.into(%{})
else
Enum.map(list, &to_map/1)
end
end
defp to_map(other), do: other
defp normalize_key(key) when is_atom(key), do: key
defp normalize_key(key) when is_binary(key), do: String.to_atom(key)
defp normalize_key(other), do: other
defp deep_merge(map, defaults) when is_map(map) and is_map(defaults) do
Map.merge(defaults, map, fn _key, default_val, user_val ->
deep_merge(user_val, default_val)
end)
end
defp deep_merge(value, _default), do: value
defp truthy?(value) when is_boolean(value), do: value
defp truthy?(value) when is_binary(value) do
normalized = String.downcase(String.trim(value))
normalized in ["true", "1", "yes", "on"]
end
defp truthy?(value) when is_integer(value), do: value != 0
defp truthy?(true), do: true
defp truthy?(_), do: false
defp debug(%{debug_pid: pid}, message) when is_pid(pid) do
send(pid, {:snakepit_otel, message})
:ok
end
defp debug(_config, _message), do: :ok
end
defmodule Snakepit.Telemetry.SafeMetadata do
@moduledoc """
Safe metadata handling for telemetry events.
This module ensures that metadata from Python workers doesn't create
new atoms at runtime, which could exhaust the BEAM atom table.
Only keys from the allowlist are converted to atoms; everything else
remains as strings.
"""
alias Snakepit.Telemetry.Naming
# Metadata keys that are safe to convert to atoms
@allowed_atom_keys [
:node,
:pool_name,
:worker_id,
:session_id,
:command,
:correlation_id,
:worker_pid,
:python_port,
:python_pid,
:mode,
:reason,
:planned,
:previous_pid,
:new_pid,
:retry_count,
:timeout_ms,
:result,
:error_type,
:error_category,
:error_message,
:traceback,
:rpc_method,
:grpc_status,
:stream_id,
:direction,
:is_final,
:tool,
:library,
:function,
:python_type,
:model,
:operation,
:timestamp_ns,
:worker_module,
:adapter_module,
:size,
:failure_reason,
:kind,
:device,
:dtype,
:shape,
:exit_code,
:payload_version,
:call_type,
:idempotent,
:retry
]
@doc """
Enriches metadata from Python with Elixir context.
Only allowed keys are converted to atoms; unknown keys remain as strings.
## Examples
iex> Snakepit.Telemetry.SafeMetadata.enrich(
...> %{"tool" => "predict"},
...> [node: :nonode@nohost, worker_id: "worker_1"]
...> )
{:ok, %{tool: "predict", node: :nonode@nohost, worker_id: "worker_1"}}
"""
def enrich(python_metadata, elixir_context)
when is_map(python_metadata) and is_list(elixir_context) do
with {:ok, safe_python} <- sanitize(python_metadata),
{:ok, safe_elixir} <- sanitize(Map.new(elixir_context)) do
{:ok, Map.merge(safe_python, safe_elixir)}
end
end
@doc """
Sanitizes a metadata map, converting only allowed keys to atoms.
## Examples
iex> Snakepit.Telemetry.SafeMetadata.sanitize(%{"node" => "test@host", "unknown" => "value"})
{:ok, %{node: "test@host", "unknown" => "value"}}
"""
def sanitize(metadata) when is_map(metadata) do
safe_metadata =
Enum.reduce(metadata, %{}, fn {key, value}, acc ->
safe_key = safe_key(key)
Map.put(acc, safe_key, value)
end)
{:ok, safe_metadata}
end
@doc """
Validates and converts measurements map.
All measurement keys must be from the allowlist (enforced by Naming module).
## Examples
iex> Snakepit.Telemetry.SafeMetadata.measurements(%{"duration" => 1000})
{:ok, %{duration: 1000}}
"""
def measurements(measurements) when is_map(measurements) do
result =
Enum.reduce_while(measurements, {:ok, %{}}, fn {key, value}, {:ok, acc} ->
case Naming.measurement_key(key) do
{:ok, atom_key} ->
{:cont, {:ok, Map.put(acc, atom_key, value)}}
{:error, reason} ->
{:halt, {:error, {:invalid_measurement_key, key, reason}}}
end
end)
result
end
# Private Helpers
defp safe_key(key) when is_atom(key) do
if key in @allowed_atom_keys do
key
else
Atom.to_string(key)
end
end
defp safe_key(key) when is_binary(key) do
atom_key = String.to_existing_atom(key)
if atom_key in @allowed_atom_keys do
atom_key
else
key
end
rescue
ArgumentError -> key
end
defp safe_key(key), do: to_string(key)
@doc """
Merges two metadata maps safely.
## Examples
iex> Snakepit.Telemetry.SafeMetadata.merge(%{"tool" => "predict"}, %{node: :nonode@nohost})
{:ok, %{"tool" => "predict", node: :nonode@nohost}}
"""
def merge(metadata1, metadata2) when is_map(metadata1) and is_map(metadata2) do
with {:ok, safe1} <- sanitize(metadata1),
{:ok, safe2} <- sanitize(metadata2) do
{:ok, Map.merge(safe1, safe2)}
end
end
@doc """
Returns the list of allowed atom keys.
"""
def allowed_atom_keys, do: @allowed_atom_keys
end
defmodule Snakepit.Telemetry.Span do
@moduledoc """
Telemetry span helpers for wrapping operations.
Provides convenient helpers for emitting start/stop/exception
telemetry events around function calls.
## Usage
# Automatic span with function
result = Snakepit.Telemetry.Span.span(
[:snakepit, :my_operation],
%{pool: :default},
fn -> do_operation() end
)
# Manual span management
span_ref = Snakepit.Telemetry.Span.start_span([:snakepit, :operation], %{})
# ... do work ...
Snakepit.Telemetry.Span.end_span(span_ref)
"""
@type event :: [atom()]
@type metadata :: map()
@type span_ref :: %{
event: event(),
start_time: integer(),
metadata: metadata()
}
@doc """
Executes a function wrapped in telemetry span events.
Emits `event ++ [:start]` before the function runs,
and `event ++ [:stop]` after it completes successfully.
If the function raises, throws, or exits, emits `event ++ [:exception]`.
## Examples
Span.span([:myapp, :operation], %{user_id: 123}, fn ->
perform_operation()
end)
"""
@spec span(event(), metadata(), (-> result)) :: result when result: any()
def span(event, metadata, fun) when is_function(fun, 0) do
start_time = System.monotonic_time()
:telemetry.execute(
event ++ [:start],
%{system_time: System.system_time()},
metadata
)
try do
result = fun.()
:telemetry.execute(
event ++ [:stop],
%{duration: System.monotonic_time() - start_time},
metadata
)
result
rescue
e ->
:telemetry.execute(
event ++ [:exception],
%{duration: System.monotonic_time() - start_time},
Map.merge(metadata, %{kind: :error, reason: e, stacktrace: __STACKTRACE__})
)
reraise e, __STACKTRACE__
catch
kind, reason ->
:telemetry.execute(
event ++ [:exception],
%{duration: System.monotonic_time() - start_time},
Map.merge(metadata, %{kind: kind, reason: reason, stacktrace: __STACKTRACE__})
)
:erlang.raise(kind, reason, __STACKTRACE__)
end
end
@doc """
Starts a telemetry span.
Returns a span reference that should be passed to `end_span/1` or `end_span/2`.
Emits `event ++ [:start]` immediately.
## Examples
span_ref = Span.start_span([:myapp, :operation], %{user_id: 123})
# ... do work ...
Span.end_span(span_ref)
"""
@spec start_span(event(), metadata()) :: span_ref()
def start_span(event, metadata \\ %{}) do
start_time = System.monotonic_time()
:telemetry.execute(
event ++ [:start],
%{system_time: System.system_time()},
metadata
)
%{
event: event,
start_time: start_time,
metadata: metadata
}
end
@doc """
Ends a telemetry span.
Emits `event ++ [:stop]` with the duration measurement.
## Examples
span_ref = Span.start_span([:myapp, :operation], %{})
# ... do work ...
Span.end_span(span_ref)
"""
@spec end_span(span_ref()) :: :ok
def end_span(span_ref) do
end_span(span_ref, %{})
end
@doc """
Ends a telemetry span with additional metadata.
Merges the additional metadata with the original span metadata
before emitting the stop event.
## Examples
span_ref = Span.start_span([:myapp, :operation], %{})
result = do_work()
Span.end_span(span_ref, %{result: result, items_processed: 100})
"""
@spec end_span(span_ref(), metadata()) :: :ok
def end_span(span_ref, additional_metadata) do
duration = System.monotonic_time() - span_ref.start_time
metadata = Map.merge(span_ref.metadata, additional_metadata)
:telemetry.execute(
span_ref.event ++ [:stop],
%{duration: duration},
metadata
)
:ok
end
@doc """
Ends a span with an exception.
Use this when you catch an exception but want to emit the
exception telemetry event before re-raising or handling it.
## Examples
span_ref = Span.start_span([:myapp, :operation], %{})
try do
do_risky_work()
rescue
e ->
Span.end_span_exception(span_ref, :error, e, __STACKTRACE__)
handle_error(e)
end
"""
@spec end_span_exception(span_ref(), :error | :exit | :throw, term(), Exception.stacktrace()) ::
:ok
def end_span_exception(span_ref, kind, reason, stacktrace) do
duration = System.monotonic_time() - span_ref.start_time
metadata =
Map.merge(span_ref.metadata, %{
kind: kind,
reason: reason,
stacktrace: stacktrace
})
:telemetry.execute(
span_ref.event ++ [:exception],
%{duration: duration},
metadata
)
:ok
end
end
defmodule Snakepit.Worker.LifecycleConfig do
@moduledoc """
Canonical configuration for lifecycle-managed workers.
Pools assemble rich `worker_config` maps that flow through the worker
pipeline. The lifecycle manager only needs a stable subset of those values
to make recycling decisions and to start replacement workers. This module
normalizes that subset into a struct so the contract is explicit and tested.
"""
@enforce_keys [:pool_name, :worker_module, :adapter_module, :profile_module]
defstruct [
:pool_name,
:pool_identifier,
:worker_module,
:adapter_module,
:worker_profile,
:profile_module,
:raw_worker_ttl,
:worker_ttl_seconds,
:worker_max_requests,
:memory_threshold_mb,
adapter_args: [],
adapter_env: [],
base_worker_config: %{}
]
@type t :: %__MODULE__{
pool_name: term(),
pool_identifier: atom() | nil,
worker_module: module(),
adapter_module: module(),
worker_profile: atom() | module(),
profile_module: module(),
raw_worker_ttl: term(),
worker_ttl_seconds: :infinity | non_neg_integer(),
worker_max_requests: :infinity | pos_integer(),
memory_threshold_mb: nil | pos_integer(),
adapter_args: list(),
adapter_env: list(),
base_worker_config: map()
}
@doc """
Ensures lifecycle config is represented as a `%LifecycleConfig{}` struct.
Accepts either an existing struct or a worker_config map. The optional
`pool_name` argument acts as a fallback when the map does not include one.
"""
@spec ensure(term(), map() | t(), keyword()) :: t()
def ensure(pool_name, config, opts \\ [])
def ensure(_pool_name, %__MODULE__{} = config, _opts), do: config
def ensure(pool_name, config, opts) when is_map(config), do: build(pool_name, config, opts)
@doc """
Builds a worker_config map for a replacement worker using the canonical data.
"""
@spec to_worker_config(t(), String.t()) :: map()
def to_worker_config(%__MODULE__{} = config, worker_id) when is_binary(worker_id) do
config.base_worker_config
|> Map.put(:worker_id, worker_id)
end
defp build(pool_name_arg, config, opts) do
pool_name = pool_name_arg || Map.get(config, :pool_name) || Snakepit.Pool
pool_identifier = Map.get(config, :pool_identifier) || opts[:pool_identifier]
worker_module = Map.get(config, :worker_module) || opts[:worker_module] || Snakepit.GRPCWorker
adapter_module =
Map.get(config, :adapter_module) ||
opts[:adapter_module] ||
Application.get_env(:snakepit, :adapter_module) ||
Snakepit.Adapters.GRPCPython
worker_profile_value = Map.get(config, :worker_profile, :process)
{worker_profile, profile_module} = resolve_profile(worker_profile_value)
adapter_args = Map.get(config, :adapter_args, [])
adapter_env = Map.get(config, :adapter_env, [])
raw_worker_ttl = Map.get(config, :worker_ttl, :infinity)
worker_ttl_seconds = normalize_ttl(raw_worker_ttl)
worker_max_requests = Map.get(config, :worker_max_requests, :infinity)
memory_threshold_mb = Map.get(config, :memory_threshold_mb)
base_worker_config =
config
|> Map.drop([:worker_id, :lifecycle_config])
|> Map.put(:pool_name, pool_name)
|> maybe_put(:pool_identifier, pool_identifier)
|> Map.put(:worker_module, worker_module)
|> Map.put(:adapter_module, adapter_module)
|> Map.put(:worker_profile, worker_profile)
|> Map.put(:adapter_args, adapter_args)
|> Map.put(:adapter_env, adapter_env)
|> Map.put(:worker_ttl, raw_worker_ttl)
|> Map.put(:worker_max_requests, worker_max_requests)
|> Map.put(:memory_threshold_mb, memory_threshold_mb)
%__MODULE__{
pool_name: pool_name,
pool_identifier: pool_identifier,
worker_module: worker_module,
adapter_module: adapter_module,
worker_profile: worker_profile,
profile_module: profile_module,
raw_worker_ttl: raw_worker_ttl,
worker_ttl_seconds: worker_ttl_seconds,
worker_max_requests: worker_max_requests,
memory_threshold_mb: memory_threshold_mb,
adapter_args: adapter_args,
adapter_env: adapter_env,
base_worker_config: base_worker_config
}
end
defp maybe_put(map, _key, nil), do: map
defp maybe_put(map, key, value), do: Map.put(map, key, value)
defp resolve_profile(:thread), do: {:thread, Snakepit.WorkerProfile.Thread}
defp resolve_profile(:process), do: {:process, Snakepit.WorkerProfile.Process}
defp resolve_profile(module) when is_atom(module), do: {module, module}
defp normalize_ttl(:infinity), do: :infinity
defp normalize_ttl({value, :seconds}) when is_integer(value), do: value
defp normalize_ttl({value, :minutes}) when is_integer(value), do: value * 60
defp normalize_ttl({value, :hours}) when is_integer(value), do: value * 3600
defp normalize_ttl({value, :days}) when is_integer(value), do: value * 86_400
defp normalize_ttl(value) when is_integer(value) and value >= 0, do: value
defp normalize_ttl(_), do: :infinity
end
defmodule Snakepit.Worker.LifecycleManager do
@moduledoc """
Worker lifecycle manager for automatic recycling and health monitoring.
Manages worker lifecycle events:
- **TTL-based recycling**: Recycle workers after configured time
- **Request-count recycling**: Recycle after N requests
- **Memory monitoring**: Recycle when the BEAM worker process exceeds a configurable threshold (optional)
- **Health checks**: Monitor worker health and restart if needed
## Why Worker Recycling?
Long-running Python processes can accumulate memory due to:
- Memory fragmentation
- Cache growth
- Subtle memory leaks in C libraries
- ML model weight accumulation
Automatic recycling prevents these issues from impacting production. The current
implementation samples the BEAM `Snakepit.GRPCWorker` process memory via
`:get_memory_usage`; Python child process memory is not yet measured directly.
## Configuration
config :snakepit,
pools: [
%{
name: :hpc_pool,
worker_profile: :thread,
worker_ttl: {3600, :seconds}, # Recycle hourly
worker_max_requests: 1000, # Or after 1000 requests
memory_threshold_mb: 2048 # Or at 2GB (optional)
}
]
## Usage
The LifecycleManager runs automatically when started in the supervision tree.
It monitors all workers across all pools.
# Manual worker recycling
Snakepit.Worker.LifecycleManager.recycle_worker(pool_name, worker_id)
# Get lifecycle statistics
Snakepit.Worker.LifecycleManager.get_stats()
## Implementation
- Runs periodic health checks (every 60 seconds)
- Tracks worker metadata (start time, request count)
- Gracefully replaces workers when recycling
- Emits telemetry events for monitoring
"""
use GenServer
alias Snakepit.Defaults
alias Snakepit.Logger, as: SLog
alias Snakepit.Worker.LifecycleConfig
@log_category :worker
defstruct [
:workers,
:check_ref,
:health_ref,
:memory_recycle_counts
]
# Client API
@doc """
Start the lifecycle manager.
"""
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Track a worker for lifecycle management.
Called automatically when workers start.
"""
def track_worker(pool_name, worker_id, worker_pid, config) do
GenServer.cast(__MODULE__, {:track, pool_name, worker_id, worker_pid, config})
end
@doc """
Untrack a worker (called when worker stops).
"""
def untrack_worker(worker_id) do
GenServer.cast(__MODULE__, {:untrack, worker_id})
end
@doc """
Manually recycle a worker.
"""
def recycle_worker(pool_name, worker_id) do
GenServer.call(__MODULE__, {:recycle, pool_name, worker_id})
end
@doc """
Increment request count for a worker.
Called after each successful request.
"""
def increment_request_count(worker_id) do
GenServer.cast(__MODULE__, {:increment_requests, worker_id})
end
@doc """
Get lifecycle statistics.
"""
def get_stats do
GenServer.call(__MODULE__, :get_stats)
end
@doc """
Returns a map of pools to the number of memory-threshold-based recycles observed
since the lifecycle manager started.
"""
def memory_recycle_counts do
GenServer.call(__MODULE__, :memory_recycle_counts)
end
# Server Callbacks
@impl true
def init(_opts) do
# Schedule periodic checks
check_ref = schedule_lifecycle_check()
health_ref = schedule_health_check()
state = %__MODULE__{
workers: %{},
check_ref: check_ref,
health_ref: health_ref,
memory_recycle_counts: %{}
}
SLog.info(@log_category, "Worker LifecycleManager started")
{:ok, state}
end
@impl true
def handle_cast({:track, pool_name, worker_id, worker_pid, config}, state) do
lifecycle_config =
LifecycleConfig.ensure(pool_name, config, worker_module: Snakepit.GRPCWorker)
ttl = lifecycle_config.worker_ttl_seconds
max_requests = lifecycle_config.worker_max_requests
memory_threshold = lifecycle_config.memory_threshold_mb
worker_state = %{
pool: pool_name,
worker_id: worker_id,
pid: worker_pid,
started_at: System.monotonic_time(:second),
request_count: 0,
ttl: ttl,
max_requests: max_requests,
memory_threshold: memory_threshold,
config: lifecycle_config
}
# Monitor the worker process
Process.monitor(worker_pid)
new_workers = Map.put(state.workers, worker_id, worker_state)
SLog.debug(
@log_category,
"Tracking worker #{worker_id} (TTL: #{inspect(ttl)}, max_requests: #{inspect(max_requests)})"
)
{:noreply, %{state | workers: new_workers}}
end
@impl true
def handle_cast({:untrack, worker_id}, state) do
new_workers = Map.delete(state.workers, worker_id)
SLog.debug(@log_category, "Untracked worker #{worker_id}")
{:noreply, %{state | workers: new_workers}}
end
@impl true
def handle_cast({:increment_requests, worker_id}, state) do
case Map.get(state.workers, worker_id) do
nil ->
# Worker not tracked (may have been recycled)
{:noreply, state}
worker_state ->
updated_state = %{worker_state | request_count: worker_state.request_count + 1}
new_workers = Map.put(state.workers, worker_id, updated_state)
# Check if we hit max requests
if should_recycle_requests?(updated_state) do
SLog.info(
@log_category,
"Worker #{worker_id} reached max requests (#{updated_state.request_count}), scheduling recycle"
)
# Schedule recycle asynchronously
GenServer.cast(self(), {:recycle_worker, worker_id, :max_requests})
end
{:noreply, %{state | workers: new_workers}}
end
end
@impl true
def handle_cast({:recycle_worker, worker_id, reason}, state) do
case Map.get(state.workers, worker_id) do
nil ->
SLog.debug(@log_category, "Worker #{worker_id} already recycled")
{:noreply, state}
worker_state ->
SLog.info(@log_category, "Recycling worker #{worker_id} (reason: #{reason})")
# Emit telemetry
emit_recycle_telemetry(worker_state, reason)
# Perform recycling
do_recycle_worker(worker_state)
# Remove from tracking
new_workers = Map.delete(state.workers, worker_id)
{:noreply, %{state | workers: new_workers}}
end
end
@impl true
def handle_call({:recycle, pool_name, worker_id}, _from, state) do
case Map.get(state.workers, worker_id) do
nil ->
{:reply, {:error, :worker_not_found}, state}
worker_state ->
if worker_state.pool == pool_name do
SLog.info(@log_category, "Manual recycle requested for worker #{worker_id}")
# Emit telemetry
emit_recycle_telemetry(worker_state, :manual)
# Perform recycling
do_recycle_worker(worker_state)
# Remove from tracking
new_workers = Map.delete(state.workers, worker_id)
{:reply, :ok, %{state | workers: new_workers}}
else
{:reply, {:error, :pool_mismatch}, state}
end
end
end
@impl true
def handle_call(:get_stats, _from, state) do
stats = %{
total_workers: map_size(state.workers),
workers_by_pool:
state.workers
|> Enum.group_by(fn {_id, worker} -> worker.pool end)
|> Map.new(fn {pool, workers} -> {pool, length(workers)} end),
total_requests:
state.workers
|> Enum.map(fn {_id, worker} -> worker.request_count end)
|> Enum.sum(),
workers_near_ttl: count_workers_near_ttl(state.workers),
workers_near_max_requests: count_workers_near_max_requests(state.workers),
memory_recycles_by_pool: state.memory_recycle_counts
}
{:reply, stats, state}
end
@impl true
def handle_call(:memory_recycle_counts, _from, state) do
{:reply, state.memory_recycle_counts, state}
end
@impl true
def handle_info(:lifecycle_check, state) do
now = System.monotonic_time(:second)
# Check all workers for recycling conditions
{recycled_workers, memory_recycle_counts} =
Enum.reduce(state.workers, {[], state.memory_recycle_counts}, fn
{worker_id, worker_state}, {acc, memory_counts} ->
case recycle_decision(worker_state, now) do
{:recycle, reason, extra_metadata} ->
log_recycle_reason(worker_id, reason, extra_metadata)
emit_recycle_telemetry(worker_state, reason, extra_metadata)
do_recycle_worker(worker_state)
new_counts = maybe_track_memory_recycle(memory_counts, worker_state.pool, reason)
{[worker_id | acc], new_counts}
:keep ->
{acc, memory_counts}
end
end)
# Remove recycled workers from tracking
new_workers =
Enum.reduce(recycled_workers, state.workers, fn worker_id, workers ->
Map.delete(workers, worker_id)
end)
# Schedule next check
check_ref = schedule_lifecycle_check()
{:noreply,
%{
state
| workers: new_workers,
check_ref: check_ref,
memory_recycle_counts: memory_recycle_counts
}}
end
@impl true
def handle_info(:health_check, state) do
# Perform health checks on all workers
Enum.each(state.workers, fn {worker_id, worker_state} ->
check_worker_health(worker_id, worker_state)
end)
# Schedule next health check
health_ref = schedule_health_check()
{:noreply, %{state | health_ref: health_ref}}
end
@impl true
def handle_info({:DOWN, _ref, :process, pid, reason}, state) do
# Find worker by PID
case Enum.find(state.workers, fn {_id, worker} -> worker.pid == pid end) do
nil ->
# Not a tracked worker
{:noreply, state}
{worker_id, worker_state} ->
if Application.get_env(:snakepit, :test_mode, false) do
SLog.debug(
@log_category,
"Worker #{worker_id} (#{inspect(pid)}) died: #{inspect(reason)}"
)
else
SLog.warning(
@log_category,
"Worker #{worker_id} (#{inspect(pid)}) died: #{inspect(reason)}"
)
end
# Emit telemetry
emit_recycle_telemetry(worker_state, :worker_died)
# Remove from tracking (supervisor will restart automatically)
new_workers = Map.delete(state.workers, worker_id)
{:noreply, %{state | workers: new_workers}}
end
end
# Private Functions
defp schedule_lifecycle_check do
Process.send_after(self(), :lifecycle_check, Defaults.lifecycle_check_interval())
end
defp schedule_health_check do
Process.send_after(self(), :health_check, Defaults.lifecycle_health_check_interval())
end
defp recycle_decision(worker_state, now) do
cond do
should_recycle_ttl?(worker_state, now) ->
{:recycle, :ttl_expired, %{}}
should_recycle_requests?(worker_state) ->
{:recycle, :max_requests, %{}}
true ->
case memory_recycle_decision(worker_state) do
nil -> :keep
extra -> {:recycle, :memory_threshold, extra}
end
end
end
defp should_recycle_ttl?(worker_state, now) do
case worker_state.ttl do
:infinity -> false
ttl_seconds -> now - worker_state.started_at >= ttl_seconds
end
end
defp should_recycle_requests?(worker_state) do
case worker_state.max_requests do
:infinity -> false
max -> worker_state.request_count >= max
end
end
defp memory_recycle_decision(worker_state) do
case worker_state.memory_threshold do
nil ->
nil
threshold_mb ->
case get_worker_memory_mb(worker_state.pid) do
{:ok, memory_mb} when memory_mb >= threshold_mb ->
%{memory_mb: memory_mb, memory_threshold_mb: threshold_mb}
{:ok, _memory_mb} ->
nil
{:error, reason} ->
SLog.warning(
@log_category,
"Memory probe for #{worker_state.worker_id} failed: #{inspect(reason)} (threshold #{threshold_mb} MB)"
)
nil
end
end
end
defp log_recycle_reason(worker_id, :ttl_expired, _extra) do
SLog.info(@log_category, "Worker #{worker_id} TTL expired, recycling...")
end
defp log_recycle_reason(worker_id, :max_requests, _extra) do
SLog.info(@log_category, "Worker #{worker_id} reached max requests, recycling...")
end
defp log_recycle_reason(worker_id, :memory_threshold, %{
memory_mb: memory_mb,
memory_threshold_mb: threshold_mb
}) do
SLog.info(
@log_category,
"Worker #{worker_id} exceeded memory threshold (#{memory_mb} MB >= #{threshold_mb} MB), recycling..."
)
end
defp log_recycle_reason(worker_id, other_reason, _extra) do
SLog.info(@log_category, "Worker #{worker_id} recycling due to #{inspect(other_reason)}")
end
defp maybe_track_memory_recycle(counts, _pool, reason) when reason != :memory_threshold,
do: counts
defp maybe_track_memory_recycle(counts, pool, :memory_threshold) do
Map.update(counts, pool, 1, &(&1 + 1))
end
defp do_recycle_worker(worker_state) do
pool_name = worker_state.pool
worker_id = worker_state.worker_id
# Stop the old worker
SLog.debug(@log_category, "Stopping worker #{worker_id} for recycling...")
# Get the profile module for this worker
profile_module = lifecycle_profile_module(worker_state.config)
# Stop via profile
case profile_module.stop_worker(worker_state.pid) do
:ok ->
SLog.debug(@log_category, "Worker #{worker_id} stopped successfully")
# Start a replacement
case start_replacement_worker(pool_name, worker_state.config) do
{:ok, new_pid} ->
SLog.info(
@log_category,
"Worker #{worker_id} recycled successfully (new PID: #{inspect(new_pid)})"
)
:ok
{:error, reason} ->
SLog.error(
@log_category,
"Failed to start replacement for #{worker_id}: #{inspect(reason)}"
)
{:error, reason}
end
error ->
SLog.error(@log_category, "Failed to stop worker #{worker_id}: #{inspect(error)}")
error
end
end
defp start_replacement_worker(pool_name, %LifecycleConfig{} = config) do
# Generate new worker ID
worker_id = "pool_worker_#{:erlang.unique_integer([:positive])}"
profile_module = config.profile_module
# Build config for new worker
worker_config = LifecycleConfig.to_worker_config(config, worker_id)
# Start via profile
case profile_module.start_worker(worker_config) do
{:ok, pid} ->
# Track the new worker
track_worker(pool_name, worker_id, pid, config)
{:ok, pid}
error ->
error
end
end
defp start_replacement_worker(pool_name, config) when is_map(config) do
lifecycle_config = LifecycleConfig.ensure(pool_name, config)
start_replacement_worker(pool_name, lifecycle_config)
end
defp check_worker_health(worker_id, worker_state) do
profile_module = lifecycle_profile_module(worker_state.config)
case profile_module.health_check(worker_state.pid) do
:ok ->
SLog.debug(@log_category, "Worker #{worker_id} health check passed")
{:error, reason} ->
SLog.warning(@log_category, "Worker #{worker_id} health check failed: #{inspect(reason)}")
# Emit telemetry
:telemetry.execute(
[:snakepit, :worker, :health_check_failed],
%{count: 1},
%{
worker_id: worker_id,
pool: worker_state.pool,
reason: reason
}
)
end
end
defp get_worker_memory_mb(worker_pid) do
# Try to get memory usage from worker
# This requires worker to expose memory info
case GenServer.call(worker_pid, :get_memory_usage, 1000) do
{:ok, memory_bytes} ->
{:ok, div(memory_bytes, 1024 * 1024)}
_ ->
{:error, :not_available}
end
catch
:exit, _ -> {:error, :worker_not_responding}
end
defp count_workers_near_ttl(workers) do
now = System.monotonic_time(:second)
Enum.count(workers, fn {_id, worker} ->
case worker.ttl do
:infinity ->
false
ttl_seconds ->
age = now - worker.started_at
# Within 10% of TTL
age >= ttl_seconds * 0.9
end
end)
end
defp count_workers_near_max_requests(workers) do
Enum.count(workers, fn {_id, worker} ->
case worker.max_requests do
:infinity ->
false
max ->
# Within 10% of max
worker.request_count >= max * 0.9
end
end)
end
defp emit_recycle_telemetry(worker_state, reason, extra_metadata \\ %{}) do
measurements =
case Map.get(extra_metadata, :memory_mb) do
nil -> %{count: 1}
memory_mb -> %{count: 1, memory_mb: memory_mb}
end
metadata =
%{
worker_id: worker_state.worker_id,
pool: worker_state.pool,
reason: reason,
uptime_seconds: System.monotonic_time(:second) - worker_state.started_at,
request_count: worker_state.request_count
}
|> maybe_put_metadata(:memory_threshold_mb, Map.get(extra_metadata, :memory_threshold_mb))
|> maybe_put_metadata(:memory_mb, Map.get(extra_metadata, :memory_mb))
:telemetry.execute([:snakepit, :worker, :recycled], measurements, metadata)
end
defp lifecycle_profile_module(%LifecycleConfig{profile_module: module}), do: module
defp lifecycle_profile_module(config) when is_map(config) do
case Map.get(config, :worker_profile, :process) do
:process -> Snakepit.WorkerProfile.Process
:thread -> Snakepit.WorkerProfile.Thread
module when is_atom(module) -> module
end
end
defp maybe_put_metadata(map, _key, nil), do: map
defp maybe_put_metadata(map, key, value), do: Map.put(map, key, value)
end
defmodule Snakepit.Worker.TaintRegistry do
@moduledoc """
Tracks tainted workers and devices after crash classification.
"""
@table :snakepit_worker_taints
def taint_worker(worker_id, opts) when is_binary(worker_id) do
ensure_table()
now_ms = System.monotonic_time(:millisecond)
duration = Keyword.get(opts, :duration_ms, 60_000)
record = %{
tainted_until: now_ms + duration,
reason: Keyword.get(opts, :reason),
exit_code: Keyword.get(opts, :exit_code),
device: Keyword.get(opts, :device),
crashed_at: now_ms,
restart_notified: false
}
:ets.insert(@table, {worker_id, record})
:ok
end
def worker_tainted?(worker_id) when is_binary(worker_id) do
ensure_table()
case :ets.lookup(@table, worker_id) do
[{^worker_id, record}] ->
if expired?(record) do
:ets.delete(@table, worker_id)
false
else
true
end
_ ->
false
end
end
def worker_info(worker_id) when is_binary(worker_id) do
ensure_table()
case :ets.lookup(@table, worker_id) do
[{^worker_id, record}] ->
if expired?(record) do
:ets.delete(@table, worker_id)
:error
else
{:ok, record}
end
_ ->
:error
end
end
def consume_restart(worker_id) when is_binary(worker_id) do
ensure_table()
case :ets.lookup(@table, worker_id) do
[{^worker_id, record}] ->
cond do
expired?(record) ->
:ets.delete(@table, worker_id)
:error
record.restart_notified ->
:error
true ->
updated = Map.put(record, :restart_notified, true)
:ets.insert(@table, {worker_id, updated})
{:ok, record}
end
_ ->
:error
end
end
def clear_worker(worker_id) when is_binary(worker_id) do
ensure_table()
:ets.delete(@table, worker_id)
:ok
end
defp expired?(record) do
now_ms = System.monotonic_time(:millisecond)
now_ms > Map.get(record, :tainted_until, 0)
end
defp ensure_table do
case :ets.whereis(@table) do
:undefined ->
try do
:ets.new(@table, [:named_table, :set, :public, {:read_concurrency, true}])
rescue
ArgumentError ->
:ok
end
@table
_ ->
@table
end
end
end
defmodule Snakepit.WorkerProfile.Thread.CapacityStore do
@moduledoc false
use GenServer
alias Snakepit.Logger, as: SLog
@table_name :snakepit_worker_capacity
@log_category :worker
## Client API
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
def ensure_started do
case Process.whereis(__MODULE__) do
nil ->
case start_link([]) do
{:ok, pid} -> {:ok, pid}
{:error, {:already_started, pid}} -> {:ok, pid}
other -> other
end
pid ->
{:ok, pid}
end
end
def track_worker(worker_pid, capacity) when is_pid(worker_pid) and capacity > 0 do
GenServer.call(__MODULE__, {:track_worker, worker_pid, capacity})
end
def untrack_worker(worker_pid) when is_pid(worker_pid) do
GenServer.call(__MODULE__, {:untrack_worker, worker_pid})
end
def check_and_increment_load(worker_pid) when is_pid(worker_pid) do
GenServer.call(__MODULE__, {:check_and_increment_load, worker_pid})
end
def decrement_load(worker_pid) when is_pid(worker_pid) do
GenServer.call(__MODULE__, {:decrement_load, worker_pid})
end
def get_capacity(worker_pid) when is_pid(worker_pid) do
GenServer.call(__MODULE__, {:get_capacity, worker_pid})
end
def get_load(worker_pid) when is_pid(worker_pid) do
GenServer.call(__MODULE__, {:get_load, worker_pid})
end
def table_name, do: @table_name
## Server callbacks
@impl true
def init(_opts) do
table =
:ets.new(@table_name, [
:set,
:protected,
:named_table,
{:read_concurrency, true}
])
SLog.debug(@log_category, "Thread capacity store started with ETS table #{inspect(table)}")
{:ok, %{table: table}}
end
@impl true
def handle_call({:track_worker, worker_pid, capacity}, _from, state) do
:ets.insert(state.table, {worker_pid, capacity, 0})
{:reply, :ok, state}
end
@impl true
def handle_call({:untrack_worker, worker_pid}, _from, state) do
:ets.delete(state.table, worker_pid)
{:reply, :ok, state}
end
@impl true
def handle_call({:check_and_increment_load, worker_pid}, _from, state) do
reply =
case :ets.lookup(state.table, worker_pid) do
[{^worker_pid, capacity, load}] when load < capacity ->
:ets.insert(state.table, {worker_pid, capacity, load + 1})
{:ok, capacity, load + 1}
[{^worker_pid, capacity, load}] ->
{:at_capacity, capacity, load}
[] ->
{:error, :unknown_worker}
end
{:reply, reply, state}
end
@impl true
def handle_call({:get_capacity, worker_pid}, _from, state) do
capacity =
case :ets.lookup(state.table, worker_pid) do
[{^worker_pid, capacity, _load}] -> capacity
[] -> 1
end
{:reply, capacity, state}
end
@impl true
def handle_call({:get_load, worker_pid}, _from, state) do
load =
case :ets.lookup(state.table, worker_pid) do
[{^worker_pid, _capacity, load}] -> load
[] -> 0
end
{:reply, load, state}
end
@impl true
def handle_call({:decrement_load, worker_pid}, _from, state) do
new_load =
case :ets.lookup(state.table, worker_pid) do
[{^worker_pid, capacity, load}] when load > 0 ->
:ets.insert(state.table, {worker_pid, capacity, load - 1})
load - 1
[{^worker_pid, capacity, _load}] ->
:ets.insert(state.table, {worker_pid, capacity, 0})
0
[] ->
0
end
{:reply, new_load, state}
end
end
defmodule Snakepit.WorkerProfile.Process do
@moduledoc """
Multi-process worker profile (default).
Each worker is a separate OS process, providing:
- **Process isolation**: Crashes don't affect other workers
- **GIL compatibility**: Works with all Python versions
- **High concurrency**: Optimal for 100+ workers with I/O-bound tasks
This is the default profile and maintains 100% backward compatibility
with Snakepit v0.5.x configurations.
## Configuration
config :snakepit,
pools: [
%{
name: :default,
worker_profile: :process, # Explicit (or omit for default)
pool_size: 100,
adapter_module: Snakepit.Adapters.GRPCPython,
adapter_env: [
{"OPENBLAS_NUM_THREADS", "1"},
{"OMP_NUM_THREADS", "1"}
]
}
]
## Implementation Details
- Each worker runs a single-threaded Python process
- Workers are single-capacity (one request at a time)
- Environment variables enforce single-threading in scientific libraries
- Startup is batched to prevent resource exhaustion
"""
@behaviour Snakepit.WorkerProfile
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Pool.WorkerSupervisor
@log_category :worker
@impl true
def start_worker(config) do
worker_id = Map.fetch!(config, :worker_id)
worker_module = Map.get(config, :worker_module, Snakepit.GRPCWorker)
adapter_module = Map.fetch!(config, :adapter_module)
pool_name = Map.get(config, :pool_name, Snakepit.Pool)
# Build adapter environment with single-threading enforcement
config_with_env = apply_adapter_env(config)
# Start the worker via the WorkerSupervisor, passing worker_config for lifecycle management
case WorkerSupervisor.start_worker(
worker_id,
worker_module,
adapter_module,
pool_name,
config_with_env
) do
{:ok, pid} ->
SLog.debug(@log_category, "Process profile started worker #{worker_id}: #{inspect(pid)}")
{:ok, pid}
error ->
error
end
end
@impl true
def stop_worker(worker_pid) when is_pid(worker_pid) do
case PoolRegistry.get_worker_id_by_pid(worker_pid) do
{:ok, worker_id} ->
case WorkerSupervisor.stop_worker(worker_id) do
{:error, :worker_not_found} -> :ok
other -> other
end
{:error, :not_found} ->
# Worker not found, may already be stopped
:ok
end
end
def stop_worker(worker_id) when is_binary(worker_id) do
WorkerSupervisor.stop_worker(worker_id)
end
@impl true
def execute_request(worker_pid, request, timeout) when is_pid(worker_pid) do
command = Map.fetch!(request, :command)
args = Map.get(request, :args, %{})
# Use the worker module's execute function
worker_module = get_worker_module(worker_pid)
worker_module.execute(worker_pid, command, args, timeout)
end
def execute_request(worker_id, request, timeout) when is_binary(worker_id) do
# Lookup PID from worker_id
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
execute_request(pid, request, timeout)
{:error, _} ->
{:error, :worker_not_found}
end
end
@impl true
def get_capacity(_worker_handle) do
# Process profile: single-threaded, capacity = 1
1
end
@impl true
def get_load(_worker_handle) do
# For process profile, load is binary: 0 (idle) or 1 (busy)
# This information is tracked by the pool, not the worker itself
# Return 0 as workers don't maintain their own load state
# The pool's busy/available sets track actual load
0
end
@impl true
def health_check(worker_handle) when is_pid(worker_handle) do
if Process.alive?(worker_handle) do
# Could optionally send a ping command
:ok
else
{:error, :worker_dead}
end
end
def health_check(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
health_check(pid)
{:error, _} ->
{:error, :worker_not_found}
end
end
@impl true
def get_metadata(worker_handle) when is_pid(worker_handle) do
{:ok,
%{
profile: :process,
capacity: 1,
worker_type: "single-process",
threading: "single-threaded"
}}
end
def get_metadata(worker_id) when is_atom(worker_id) do
# Handle atom input (for tests with :fake_worker, etc.)
{:ok,
%{
profile: :process,
capacity: 1,
worker_type: "single-process",
threading: "single-threaded"
}}
end
def get_metadata(worker_id) when is_binary(worker_id) do
{:ok,
%{
profile: :process,
capacity: 1,
worker_type: "single-process",
threading: "single-threaded",
worker_id: worker_id
}}
end
# Private helpers
@all_thread_control_vars [
"OPENBLAS_NUM_THREADS",
"MKL_NUM_THREADS",
"OMP_NUM_THREADS",
"NUMEXPR_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
# macOS Accelerate
"GRPC_POLL_STRATEGY"
]
@doc false
def apply_adapter_env(config) when is_map(config) do
Map.put(config, :adapter_env, build_process_env(config))
end
@doc false
def build_process_env(config) do
defaults = build_process_env_defaults()
user_env = Map.get(config, :adapter_env, [])
merge_env(defaults, user_env)
end
defp build_process_env_defaults do
Enum.map(@all_thread_control_vars, fn var ->
{var, System.get_env(var) || default_thread_value(var)}
end)
end
defp default_thread_value("GRPC_POLL_STRATEGY"), do: "poll"
defp default_thread_value(_var), do: "1"
defp merge_env(defaults, user_env) do
merged =
Enum.reduce(normalize_env_entries(user_env), Map.new(defaults), fn {key, val}, acc ->
Map.put(acc, key, val)
end)
Map.to_list(merged)
end
defp normalize_env_entries(nil), do: []
defp normalize_env_entries(env) when is_map(env) do
env
|> Map.to_list()
|> normalize_env_entries()
end
defp normalize_env_entries(env) when is_list(env) do
Enum.flat_map(env, fn
{key, value} -> [{to_string(key), to_string(value)}]
key when is_binary(key) -> [{key, ""}]
key when is_atom(key) -> [{Atom.to_string(key), ""}]
_ -> []
end)
end
defp get_worker_module(worker_pid) do
with {:ok, worker_id} <-
Snakepit.Pool.Registry.get_worker_id_by_pid(worker_pid),
{:ok, _pid, %{worker_module: module}} <-
PoolRegistry.fetch_worker(worker_id) do
module
else
_ ->
# Default to GRPCWorker
Snakepit.GRPCWorker
end
end
end
defmodule Snakepit.WorkerProfile.Thread do
@moduledoc """
Multi-threaded worker profile (Python 3.13+ optimized).
Each worker is a Python process with a thread pool, providing:
- **Shared memory**: Zero-copy data sharing within worker
- **CPU parallelism**: True multi-threading without GIL (Python 3.13+)
- **Lower memory**: One interpreter vs many
- **High throughput**: Optimal for CPU-bound tasks
## Configuration
config :snakepit,
pools: [
%{
name: :hpc_pool,
worker_profile: :thread,
pool_size: 4, # 4 processes
threads_per_worker: 16, # 16 threads each = 64 total capacity
adapter_module: Snakepit.Adapters.GRPCPython,
adapter_args: ["--mode", "threaded", "--max-workers", "16"],
adapter_env: [
# Allow multi-threading in libraries
{"OPENBLAS_NUM_THREADS", "16"},
{"OMP_NUM_THREADS", "16"}
],
worker_ttl: {3600, :seconds}, # Recycle hourly
worker_max_requests: 1000, # Or after 1000 requests
thread_safety_checks: true # Enable runtime validation
}
]
## Requirements
- Python 3.13+ for optimal performance (free-threading)
- Thread-safe Python adapters
- Thread-safe ML libraries (NumPy, PyTorch, etc.)
## Status
Thread profile is fully supported when paired with Python 3.13+ and thread-safe adapters.
## Implementation Notes
The thread profile:
1. Starts fewer Python processes (4-16 instead of 100+)
2. Runs a ThreadPoolExecutor per worker process
3. Tracks per-worker capacity via `threads_per_worker` for pool scheduling
4. Supports optional CapacityStore telemetry with `capacity_strategy: :hybrid`
5. Allows concurrent requests to the same worker via HTTP/2 multiplexing
"""
@behaviour Snakepit.WorkerProfile
alias Snakepit.Logger, as: SLog
alias Snakepit.Logger.Redaction
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Pool.WorkerSupervisor
alias Snakepit.WorkerProfile.Thread.CapacityStore
@log_category :worker
@impl true
def start_worker(config) do
worker_id = Map.fetch!(config, :worker_id)
worker_module = Map.get(config, :worker_module, Snakepit.GRPCWorker)
adapter_module = Map.fetch!(config, :adapter_module)
pool_name = Map.get(config, :pool_name, Snakepit.Pool)
threads_per_worker = Map.get(config, :threads_per_worker, 10)
# Ensure capacity tracking table exists
ensure_capacity_table()
# Build adapter args and env for threaded mode
adapter_args = build_adapter_args(config)
adapter_env = build_adapter_env(config)
SLog.info(
@log_category,
"Starting threaded worker #{worker_id} with #{threads_per_worker} threads"
)
SLog.debug(@log_category, "Thread worker adapter_args: #{Redaction.describe(adapter_args)}")
# Create enhanced worker config with thread profile settings
worker_config =
config
|> Map.put(:adapter_args, adapter_args)
|> Map.put(:adapter_env, adapter_env)
# Start the worker via WorkerSupervisor with full config
case WorkerSupervisor.start_worker(
worker_id,
worker_module,
adapter_module,
pool_name,
worker_config
) do
{:ok, pid} ->
:ok = CapacityStore.track_worker(pid, threads_per_worker)
SLog.info(
@log_category,
"Thread profile started worker #{worker_id}: #{inspect(pid)} with capacity #{threads_per_worker}"
)
{:ok, pid}
error ->
SLog.error(
@log_category,
"Failed to start threaded worker #{worker_id}: #{inspect(error)}"
)
error
end
end
@impl true
def stop_worker(worker_pid) when is_pid(worker_pid) do
# Remove from capacity table
ensure_capacity_table()
:ok = CapacityStore.untrack_worker(worker_pid)
case PoolRegistry.get_worker_id_by_pid(worker_pid) do
{:ok, worker_id} ->
case WorkerSupervisor.stop_worker(worker_id) do
{:error, :worker_not_found} -> :ok
other -> other
end
{:error, :not_found} ->
# Worker not found, may already be stopped
:ok
end
end
def stop_worker(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} -> stop_worker(pid)
{:error, :not_found} -> :ok
end
end
@impl true
def execute_request(worker_pid, request, timeout) when is_pid(worker_pid) do
# Check capacity before executing
case check_and_increment_load(worker_pid) do
:ok ->
command = Map.fetch!(request, :command)
args = Map.get(request, :args, %{})
# Execute via worker module
worker_module = get_worker_module(worker_pid)
try do
result = worker_module.execute(worker_pid, command, args, timeout)
result
after
# Always decrement load, even on error
decrement_load(worker_pid)
end
{:error, :at_capacity} ->
{:error, :worker_at_capacity}
end
end
def execute_request(worker_id, request, timeout) when is_binary(worker_id) do
# Lookup PID from worker_id
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
execute_request(pid, request, timeout)
{:error, _} ->
{:error, :worker_not_found}
end
end
@impl true
def get_capacity(worker_pid) when is_pid(worker_pid) do
ensure_capacity_table()
CapacityStore.get_capacity(worker_pid)
end
def get_capacity(worker_id) when is_atom(worker_id) do
# Handle atom input (for tests)
1
end
def get_capacity(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} -> get_capacity(pid)
{:error, _} -> 1
end
end
@impl true
def get_load(worker_pid) when is_pid(worker_pid) do
ensure_capacity_table()
CapacityStore.get_load(worker_pid)
end
def get_load(worker_id) when is_atom(worker_id) do
# Handle atom input (for tests with :fake_worker, etc.)
0
end
def get_load(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} -> get_load(pid)
{:error, _} -> 0
end
end
@impl true
def health_check(worker_handle) when is_pid(worker_handle) do
if Process.alive?(worker_handle) do
# Optionally could send a ping command
:ok
else
{:error, :worker_dead}
end
end
def health_check(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
health_check(pid)
{:error, _} ->
{:error, :worker_not_found}
end
end
@impl true
def get_metadata(worker_pid) when is_pid(worker_pid) do
capacity = get_capacity(worker_pid)
load = get_load(worker_pid)
{:ok,
%{
profile: :thread,
capacity: capacity,
load: load,
available_capacity: capacity - load,
worker_type: "multi-threaded",
threading: "thread-pool"
}}
end
def get_metadata(worker_id) when is_atom(worker_id) do
# Handle atom input (for tests with :fake_worker, etc.)
{:ok,
%{
profile: :thread,
capacity: 1,
load: 0,
available_capacity: 1,
worker_type: "multi-threaded",
threading: "thread-pool"
}}
end
def get_metadata(worker_id) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
capacity = get_capacity(pid)
load = get_load(pid)
{:ok,
%{
profile: :thread,
capacity: capacity,
load: load,
available_capacity: capacity - load,
worker_type: "multi-threaded",
threading: "thread-pool",
worker_id: worker_id
}}
{:error, _} ->
{:error, :worker_not_found}
end
end
# Private helpers
defp ensure_capacity_table do
case CapacityStore.ensure_started() do
{:ok, _pid} ->
:ok
{:error, reason} ->
SLog.warning(@log_category, "Capacity store failed to start: #{inspect(reason)}")
:ok
end
end
defp build_adapter_args(config) do
threads = Map.get(config, :threads_per_worker, 10)
adapter_spec = get_adapter_spec(config)
thread_safety_checks = Map.get(config, :thread_safety_checks, false)
# Base args for threaded mode
base_args = [
"--adapter",
adapter_spec,
"--max-workers",
"#{threads}"
]
# Add thread safety checking if enabled
base_args =
if thread_safety_checks do
base_args ++ ["--thread-safety-check"]
else
base_args
end
# Merge with user-provided args
user_args = Map.get(config, :adapter_args, [])
# User args can override base args
merge_args(base_args, user_args)
end
defp get_adapter_spec(config) do
# Try multiple sources for adapter spec
Map.get(config, :adapter_spec) ||
extract_adapter_from_args(Map.get(config, :adapter_args, [])) ||
"snakepit_bridge.adapters.threaded_showcase.ThreadedShowcaseAdapter"
end
defp extract_adapter_from_args(args) do
# Look for --adapter flag in user args
case Enum.find_index(args, &(&1 == "--adapter")) do
nil -> nil
index -> Enum.at(args, index + 1)
end
end
defp merge_args(base_args, user_args) do
# Simple merge: user args override base args
# For more sophisticated merging, could parse flags
(user_args ++ base_args)
|> Enum.chunk_every(2)
|> Enum.uniq_by(fn
[flag, _] -> flag
[flag] -> flag
end)
|> List.flatten()
end
defp build_adapter_env(config) do
threads = Map.get(config, :threads_per_worker, 10)
# Default env for threaded mode (allow multi-threading)
default_env = [
{"OPENBLAS_NUM_THREADS", "#{threads}"},
{"OMP_NUM_THREADS", "#{threads}"},
{"MKL_NUM_THREADS", "#{threads}"},
{"NUMEXPR_NUM_THREADS", "#{threads}"}
]
# Get user-specified environment (overrides defaults)
user_env = Map.get(config, :adapter_env, [])
# User env takes precedence
merged =
Enum.reduce(user_env, Map.new(default_env), fn {key, val}, acc ->
Map.put(acc, key, val)
end)
Map.to_list(merged)
end
defp check_and_increment_load(worker_pid) do
ensure_capacity_table()
case CapacityStore.check_and_increment_load(worker_pid) do
{:ok, capacity, new_load} ->
if new_load == capacity do
:telemetry.execute(
[:snakepit, :pool, :capacity_reached],
%{capacity: capacity, load: new_load},
%{worker_pid: worker_pid, profile: :thread}
)
end
:ok
{:at_capacity, capacity, load} ->
:telemetry.execute(
[:snakepit, :pool, :capacity_reached],
%{capacity: capacity, load: load},
%{worker_pid: worker_pid, profile: :thread, rejected: true}
)
{:error, :at_capacity}
{:error, :unknown_worker} ->
SLog.warning(@log_category, "Worker #{inspect(worker_pid)} not found in capacity store")
:ok
end
end
defp decrement_load(worker_pid) do
ensure_capacity_table()
CapacityStore.decrement_load(worker_pid)
end
defp get_worker_module(worker_pid) do
with {:ok, worker_id} <-
Snakepit.Pool.Registry.get_worker_id_by_pid(worker_pid),
{:ok, _pid, %{worker_module: module}} <-
PoolRegistry.fetch_worker(worker_id) do
module
else
_ ->
# Default to GRPCWorker
Snakepit.GRPCWorker
end
end
end
defmodule Snakepit.Adapter do
@moduledoc """
Behaviour for implementing adapters in Snakepit.
Adapters define how to communicate with external processes (Python, Node.js, etc.).
This allows Snakepit to be truly generalized and support multiple ML frameworks
or external systems.
## Required Callbacks
- `executable_path/0` - Returns the path to the runtime executable (python3, node, etc.)
- `script_path/0` - Returns the path to the external script to execute
- `script_args/0` - Returns additional arguments for the script
## Example Implementation
defmodule MyApp.PythonMLAdapter do
@behaviour Snakepit.Adapter
def executable_path, do: System.find_executable("python3") || System.find_executable("python")
def script_path, do: Path.join(:code.priv_dir(:my_app), "python/ml_bridge.py")
def script_args, do: ["--mode", "pool-worker"]
end
"""
@doc """
Returns the path to the runtime executable.
This is the interpreter or runtime that will execute the script.
Examples: "python3", "node", "ruby", "R", etc.
"""
@callback executable_path() :: String.t()
@doc """
Returns the path to the external script that will be executed.
This should be an absolute path to a script that implements the
bridge protocol for communication with Snakepit.
"""
@callback script_path() :: String.t()
@doc """
Returns additional command-line arguments for the script.
These arguments will be passed to the script when it's started.
Common examples: ["--mode", "pool-worker"], ["--config", "prod"]
"""
@callback script_args() :: [String.t()]
@doc """
Optional callback to get a command-specific timeout in milliseconds.
This allows adapters to specify appropriate timeouts for different
commands based on their expected execution time.
"""
@callback command_timeout(command :: String.t(), args :: map()) :: pos_integer()
@optional_callbacks [command_timeout: 2]
end
defmodule Snakepit.Application do
@moduledoc """
Application supervisor for Snakepit pooler.
Starts the core infrastructure:
- Registry for worker process registration
- StarterRegistry for worker starter supervisors
- ProcessRegistry for external PID tracking
- SessionStore for session management
- WorkerSupervisor for managing worker processes
- Pool manager for request distribution
"""
use Application
require Logger
alias Snakepit.Defaults
alias Snakepit.Logger, as: SLog
alias Snakepit.PythonThreadLimits
alias Snakepit.Telemetry.OpenTelemetry
@runtime_env Application.compile_env(:snakepit, :environment, :prod)
@impl true
def start(_type, _args) do
configure_logging()
# Configure threading limits for Python scientific libraries and gRPC
# This prevents fork bombs when spawning many workers concurrently
# Each Python worker tries to spawn threads from multiple sources:
# - OpenBLAS: 24 threads (numpy/scipy)
# - gRPC: CPU cores threads (grpcio polling)
# - Other libraries (absl, protobuf, etc.)
# With 250 workers, this can create 6,000+ threads causing "Cannot fork" errors
thread_limits =
:snakepit
|> Application.get_env(:python_thread_limits)
|> PythonThreadLimits.resolve()
# Scientific computing libraries
System.put_env("OPENBLAS_NUM_THREADS", thread_limits[:openblas] |> to_string())
System.put_env("OMP_NUM_THREADS", thread_limits[:omp] |> to_string())
System.put_env("MKL_NUM_THREADS", thread_limits[:mkl] |> to_string())
System.put_env("NUMEXPR_NUM_THREADS", thread_limits[:numexpr] |> to_string())
# gRPC library threading
# Use single-threaded polling
System.put_env("GRPC_POLL_STRATEGY", "poll")
# Reduce logging overhead
System.put_env("GRPC_VERBOSITY", "ERROR")
# Python threading behavior
# Unbuffered output for better logging
System.put_env("PYTHONUNBUFFERED", "1")
SLog.info(
:startup,
"Set Python thread limits",
openblas: thread_limits[:openblas],
omp: thread_limits[:omp],
mkl: thread_limits[:mkl],
numexpr: thread_limits[:numexpr],
grpc_poll_strategy: "poll"
)
# Check if pooling is enabled (default: false to prevent auto-start issues)
pooling_enabled = Application.get_env(:snakepit, :pooling_enabled, false)
if pooling_enabled do
ensure_python_ready()
end
if Application.get_env(:snakepit, :enable_otlp?, false) do
SLog.info(:startup, "OTLP telemetry enabled", enabled: true)
OpenTelemetry.setup()
else
SLog.debug(:startup, "OTLP telemetry disabled", enabled: false)
end
SLog.debug(
:startup,
"Snakepit.Application.start/2",
pooling_enabled: pooling_enabled,
environment: @runtime_env
)
# Get gRPC config for the Elixir server
grpc_port = Defaults.grpc_port()
# Always start SessionStore as it's needed for tests and bridge functionality
telemetry_children = Snakepit.TelemetryMetrics.reporter_children()
base_children = [
Snakepit.Bridge.SessionStore,
Snakepit.Bridge.ToolRegistry,
# Process registry for PID tracking (always available for cleanup)
Snakepit.Pool.ProcessRegistry,
# Application cleanup for hard process termination guarantees
# Runs after pool children terminate to catch any stragglers
Snakepit.Pool.ApplicationCleanup
]
pool_children =
if pooling_enabled do
pool_config = Application.get_env(:snakepit, :pool_config, %{})
pool_size = Map.get(pool_config, :pool_size, System.schedulers_online() * 2)
SLog.info(:startup, "Starting Snakepit with pooling enabled", pool_size: pool_size)
[
# GRPC client supervisor - required for connecting to Python workers
# Must be started before any gRPC client connections are attempted
{GRPC.Client.Supervisor, []},
# Start the central gRPC server that manages state
# DIAGNOSTIC: Increase backlog to handle high concurrent connection load (200+ workers)
# Default Cowboy backlog is ~128, which causes connection refusals during startup
{GRPC.Server.Supervisor,
endpoint: Snakepit.GRPC.Endpoint,
port: grpc_port,
start_server: true,
adapter_opts: [
num_acceptors: Defaults.grpc_num_acceptors(),
max_connections: Defaults.grpc_max_connections(),
socket_opts: [backlog: Defaults.grpc_socket_backlog()]
]},
# Task supervisor for async pool operations
{Task.Supervisor, name: Snakepit.TaskSupervisor},
# Telemetry gRPC stream manager (for Python worker telemetry)
Snakepit.Telemetry.GrpcStream,
# Registry for worker process registration
Snakepit.Pool.Registry,
# Registry for worker starter supervisors
Snakepit.Pool.Worker.StarterRegistry,
# Thread profile capacity tracking
Snakepit.WorkerProfile.Thread.CapacityStore,
# Worker supervisor for managing worker processes
Snakepit.Pool.WorkerSupervisor,
# Worker lifecycle manager for automatic recycling
Snakepit.Worker.LifecycleManager,
# Main pool manager
{Snakepit.Pool, [size: pool_size]}
]
else
SLog.info(:startup, "Starting Snakepit with pooling disabled", pooling_enabled: false)
[]
end
children = telemetry_children ++ base_children ++ pool_children
opts = [strategy: :one_for_one, name: Snakepit.Supervisor]
result = Supervisor.start_link(children, opts)
SLog.debug(:startup, "Snakepit.Application started",
started_at_ms: System.monotonic_time(:millisecond)
)
result
end
@impl true
def stop(_state) do
SLog.debug(:shutdown, "Snakepit.Application.stop/1",
stopped_at_ms: System.monotonic_time(:millisecond)
)
maybe_cleanup_on_stop()
:ok
end
defp maybe_cleanup_on_stop do
if Application.get_env(:snakepit, :cleanup_on_stop, true) do
if Process.whereis(Snakepit.Pool.ProcessRegistry) do
timeout_ms = Defaults.cleanup_on_stop_timeout_ms()
poll_interval_ms = Defaults.cleanup_poll_interval_ms()
try do
Snakepit.RuntimeCleanup.cleanup_current_run(
timeout_ms: timeout_ms,
poll_interval_ms: poll_interval_ms
)
rescue
error ->
SLog.warning(:shutdown, "Shutdown cleanup failed", error: error)
catch
:exit, reason ->
SLog.warning(:shutdown, "Shutdown cleanup exited", reason: reason)
end
end
end
end
defp ensure_python_ready do
doctor = Application.get_env(:snakepit, :env_doctor_module, Snakepit.EnvDoctor)
doctor.ensure_python!()
rescue
error ->
reraise error, __STACKTRACE__
end
defp configure_logging do
grpc_level =
Application.get_env(
:snakepit,
:grpc_log_level,
default_grpc_log_level()
)
if grpc_level do
Logger.put_application_level(:grpc, grpc_level)
end
end
defp default_grpc_log_level do
if Application.get_env(:snakepit, :library_mode, true) do
:error
else
nil
end
end
end
defmodule Snakepit.Bootstrap do
@moduledoc """
Provisioning workflow for development and CI environments.
It installs Mix dependencies, prepares the default Python virtual
environments, regenerates gRPC stubs, and surfaces the interpreter path the
application will use at runtime.
"""
alias Snakepit.Adapters.GRPCPython
alias Snakepit.Bootstrap.Runner
alias Snakepit.PythonRuntime
@requirements_path ["priv", "python", "requirements.txt"]
@setup_script ["scripts", "setup_test_pythons.sh"]
@grpc_script ["priv", "python", "generate_grpc.sh"]
@doc """
Execute the bootstrap workflow.
Options:
* `:project_root` - overrides the working directory (defaults to `File.cwd!/0`)
* `:runner` - injects a custom runner, useful for tests
* `:skip_mix_deps` - skips `mix deps.get` (useful for test bootstrapping)
"""
@spec run(Keyword.t()) :: :ok | {:error, term()}
def run(opts \\ []) do
# Prevent concurrent bootstrap runs - use fixed lock key
lock = {__MODULE__, :bootstrap_lock}
case :global.set_lock(lock, [node()], 0) do
true ->
run_with_lock(lock, fn -> do_run(opts) end)
false ->
Mix.shell().info("⚠️ Bootstrap already running, waiting...")
:global.trans(lock, fn -> do_run(opts) end, [node()])
end
end
defp run_with_lock(lock, fun) do
fun.()
after
:global.del_lock(lock, [node()])
end
defp do_run(opts) do
state = build_state(opts)
skip_mix_deps = Keyword.get(opts, :skip_mix_deps, false)
with :ok <- maybe_fetch_mix_deps(state, skip_mix_deps),
:ok <- ensure_managed_python(state),
:ok <- ensure_primary_python(state),
:ok <- run_script(state, @setup_script, :setup_pythons),
:ok <- run_script(state, @grpc_script, :generate_grpc) do
print_python_summary()
Mix.shell().info("✅ Snakepit bootstrap complete")
:ok
else
{:error, reason} = error ->
Mix.shell().error("❌ Snakepit bootstrap failed: #{format_reason(reason)}")
error
end
end
defp build_state(opts) do
project_root =
opts[:project_root] ||
Application.get_env(:snakepit, :bootstrap_project_root) ||
File.cwd!()
runner =
opts[:runner] ||
Application.get_env(:snakepit, :bootstrap_runner, Snakepit.Bootstrap.Runner.System)
%{
project_root: project_root,
runner: runner
}
end
defp fetch_mix_deps(%{runner: runner}) do
Mix.shell().info("📦 mix deps.get")
ensure_hex_started()
case runner.mix("deps.get", []) do
:ok -> :ok
{:error, reason} -> {:error, {:mix_failure, "deps.get", reason}}
end
end
defp maybe_fetch_mix_deps(_state, true), do: :ok
defp maybe_fetch_mix_deps(state, false), do: fetch_mix_deps(state)
defp ensure_primary_python(%{project_root: root, runner: runner}) do
requirements = Path.join([root | @requirements_path])
if File.exists?(requirements) do
with :ok <- create_primary_venv(root, runner) do
install_requirements(root, requirements, runner)
end
else
return_missing(:requirements, requirements)
end
end
defp ensure_managed_python(%{project_root: root, runner: runner}) do
case PythonRuntime.install_managed(runner, project_root: root) do
:ok -> :ok
{:error, reason} -> {:error, reason}
end
end
defp create_primary_venv(root, runner) do
venv_dir = Path.join(root, ".venv")
if File.dir?(venv_dir) do
:ok
else
create_new_venv(root, runner)
end
end
defp create_new_venv(root, runner) do
python = discover_python()
if python do
create_venv_with_python(root, runner, python)
else
{:error, :python_not_found}
end
end
defp create_venv_with_python(root, runner, python) do
Mix.shell().info("🐍 Creating .venv with #{python}")
case runner.cmd(python, ["-m", "venv", ".venv"], cd: root) do
:ok ->
Mix.shell().info("✅ .venv created successfully")
:ok
error ->
error
end
end
defp install_requirements(root, requirements, runner) do
Mix.shell().info("📦 Installing Python requirements")
pip = Path.join([root, ".venv", "bin", "pip"])
# Use -q to suppress "Requirement already satisfied" noise
runner.cmd(pip, ["install", "-q", "-r", requirements], cd: root)
end
defp run_script(%{project_root: root, runner: runner}, parts, step) do
script = Path.join([root | parts])
if File.exists?(script) do
Mix.shell().info("▶️ #{describe_step(step)}")
runner.cmd(script, [], cd: root)
else
return_missing(step, script)
end
end
defp describe_step(:setup_pythons), do: "scripts/setup_test_pythons.sh"
defp describe_step(:generate_grpc), do: "priv/python/generate_grpc.sh"
defp print_python_summary do
python = GRPCPython.executable_path()
if python do
Mix.shell().info("🐍 Detected Python interpreter: #{python}")
else
Mix.shell().info("🐍 No interpreter detected (set SNAKEPIT_PYTHON=/path/to/python)")
end
Mix.shell().info("ℹ️ Override via SNAKEPIT_PYTHON=/path/to/python")
end
defp discover_python do
System.find_executable("python3") || System.find_executable("python")
end
defp ensure_hex_started do
Mix.ensure_application!(:hex)
rescue
_ -> :ok
end
defp return_missing(kind, path), do: {:error, {:missing, kind, path}}
defp format_reason({:mix_failure, task, reason}),
do: "mix #{task} failed: #{inspect(reason)}"
defp format_reason({:missing, kind, path}),
do: "missing #{kind} at #{path}"
defp format_reason(:python_not_found),
do: "python3/python not found in PATH"
defp format_reason({:command_failed, command, status}),
do: "#{command} exited with status #{status}"
defp format_reason(reason), do: inspect(reason)
defmodule Runner do
@moduledoc """
Behaviour for executing bootstrap steps. Allows tests to inject fakes.
"""
@callback mix(task :: String.t(), args :: [String.t()]) :: :ok | {:error, term()}
@callback cmd(command :: String.t(), args :: [String.t()], keyword()) ::
:ok | {:error, term()}
defmodule System do
@moduledoc false
@behaviour Snakepit.Bootstrap.Runner
# Alias the Elixir System module to avoid name collision
alias Elixir.System, as: ErlangSystem
@impl true
def mix(task, args) do
Mix.Task.run(task, args)
:ok
catch
kind, reason ->
{:error, {kind, reason}}
end
@impl true
def cmd(command, args, opts) do
opts = Keyword.put_new(opts, :stderr_to_stdout, true)
# For shell scripts, call through bash to handle shebang properly
{actual_command, actual_args} =
if String.ends_with?(command, ".sh") do
{"bash", [command | args]}
else
{command, args}
end
# Use ErlangSystem to avoid recursive call!
case ErlangSystem.cmd(actual_command, actual_args, opts) do
{output, 0} ->
if String.trim(output) != "" do
IO.write(output)
end
:ok
{output, status} ->
if String.trim(output) != "" do
IO.write(output)
end
{:error, {:command_failed, command, status}}
end
end
end
end
end
defmodule Snakepit.CircuitBreaker do
@moduledoc """
Circuit breaker for Python worker fault tolerance.
Implements the circuit breaker pattern to prevent cascading failures
when workers are experiencing issues.
## States
- `:closed` - Normal operation, all calls allowed
- `:open` - Failure threshold exceeded, calls rejected
- `:half_open` - Testing if service recovered, limited calls allowed
## Usage
{:ok, cb} = CircuitBreaker.start_link(name: :my_cb, failure_threshold: 5)
case CircuitBreaker.call(cb, fn -> risky_operation() end) do
{:ok, result} -> handle_success(result)
{:error, :circuit_open} -> handle_circuit_open()
{:error, reason} -> handle_error(reason)
end
"""
use GenServer
alias Snakepit.Defaults
require Logger
@type state :: :closed | :open | :half_open
@type t :: %{
state: state(),
failure_count: non_neg_integer(),
success_count: non_neg_integer(),
failure_threshold: pos_integer(),
reset_timeout_ms: pos_integer(),
half_open_max_calls: pos_integer(),
half_open_calls: non_neg_integer(),
last_failure_time: integer() | nil,
name: atom() | nil
}
# Client API
@doc """
Starts a circuit breaker.
## Options
- `:name` - GenServer name (optional)
- `:failure_threshold` - Number of failures before opening (default: 5)
- `:reset_timeout_ms` - Time before transitioning to half-open (default: 30000)
- `:half_open_max_calls` - Max calls allowed in half-open state (default: 1)
"""
@spec start_link(keyword()) :: GenServer.on_start()
def start_link(opts \\ []) do
name = Keyword.get(opts, :name)
gen_opts = if name, do: [name: name], else: []
GenServer.start_link(__MODULE__, opts, gen_opts)
end
@doc """
Returns the current circuit state.
"""
@spec state(GenServer.server()) :: state()
def state(server) do
GenServer.call(server, :get_state)
end
@doc """
Checks if a call is allowed through the circuit.
"""
@spec allow_call?(GenServer.server()) :: boolean()
def allow_call?(server) do
GenServer.call(server, :allow_call?)
end
@doc """
Records a successful call.
"""
@spec record_success(GenServer.server()) :: :ok
def record_success(server) do
GenServer.cast(server, :record_success)
end
@doc """
Records a failed call.
"""
@spec record_failure(GenServer.server()) :: :ok
def record_failure(server) do
GenServer.cast(server, :record_failure)
end
@doc """
Executes a function through the circuit breaker.
Returns `{:error, :circuit_open}` if the circuit is open.
"""
@spec call(GenServer.server(), (-> any())) :: any()
def call(server, fun) do
case GenServer.call(server, :try_call) do
:allowed ->
try do
result = fun.()
case result do
{:ok, _} ->
record_success(server)
result
{:error, _} ->
record_failure(server)
result
other ->
record_success(server)
other
end
rescue
e ->
record_failure(server)
reraise e, __STACKTRACE__
catch
kind, reason ->
record_failure(server)
:erlang.raise(kind, reason, __STACKTRACE__)
end
:rejected ->
{:error, :circuit_open}
end
end
@doc """
Returns circuit breaker statistics.
"""
@spec stats(GenServer.server()) :: map()
def stats(server) do
GenServer.call(server, :get_stats)
end
@doc """
Resets the circuit breaker to closed state.
"""
@spec reset(GenServer.server()) :: :ok
def reset(server) do
GenServer.cast(server, :reset)
end
# GenServer callbacks
@impl true
def init(opts) do
state = %{
state: :closed,
failure_count: 0,
success_count: 0,
failure_threshold:
Keyword.get(opts, :failure_threshold, Defaults.circuit_breaker_failure_threshold()),
reset_timeout_ms:
Keyword.get(opts, :reset_timeout_ms, Defaults.circuit_breaker_reset_timeout_ms()),
half_open_max_calls:
Keyword.get(opts, :half_open_max_calls, Defaults.circuit_breaker_half_open_max_calls()),
half_open_calls: 0,
last_failure_time: nil,
name: Keyword.get(opts, :name)
}
{:ok, state}
end
@impl true
def handle_call(:get_state, _from, state) do
state = maybe_transition_to_half_open(state)
{:reply, state.state, state}
end
def handle_call(:allow_call?, _from, state) do
state = maybe_transition_to_half_open(state)
allowed =
case state.state do
:closed -> true
:half_open -> state.half_open_calls < state.half_open_max_calls
:open -> false
end
{:reply, allowed, state}
end
def handle_call(:try_call, _from, state) do
state = maybe_transition_to_half_open(state)
case state.state do
:closed ->
{:reply, :allowed, state}
:half_open ->
if state.half_open_calls < state.half_open_max_calls do
{:reply, :allowed, %{state | half_open_calls: state.half_open_calls + 1}}
else
{:reply, :rejected, state}
end
:open ->
{:reply, :rejected, state}
end
end
def handle_call(:get_stats, _from, state) do
stats = %{
state: state.state,
failure_count: state.failure_count,
success_count: state.success_count,
failure_threshold: state.failure_threshold,
half_open_calls: state.half_open_calls
}
{:reply, stats, state}
end
@impl true
def handle_cast(:record_success, state) do
state = %{state | success_count: state.success_count + 1}
state =
case state.state do
:half_open ->
# Success in half-open transitions to closed
emit_closed_event(state)
%{state | state: :closed, failure_count: 0, half_open_calls: 0}
_ ->
state
end
{:noreply, state}
end
def handle_cast(:record_failure, state) do
now = System.monotonic_time(:millisecond)
state = %{state | failure_count: state.failure_count + 1, last_failure_time: now}
state =
case state.state do
:closed ->
if state.failure_count >= state.failure_threshold do
emit_opened_event(state)
%{state | state: :open}
else
state
end
:half_open ->
# Failure in half-open transitions back to open
emit_opened_event(state)
%{state | state: :open, half_open_calls: 0}
:open ->
state
end
{:noreply, state}
end
def handle_cast(:reset, state) do
emit_closed_event(state)
{:noreply,
%{
state
| state: :closed,
failure_count: 0,
half_open_calls: 0,
last_failure_time: nil
}}
end
# Private functions
defp maybe_transition_to_half_open(%{state: :open} = state) do
now = System.monotonic_time(:millisecond)
if state.last_failure_time &&
now - state.last_failure_time >= state.reset_timeout_ms do
emit_half_open_event(state)
%{state | state: :half_open, half_open_calls: 0}
else
state
end
end
defp maybe_transition_to_half_open(state), do: state
defp emit_opened_event(state) do
:telemetry.execute(
[:snakepit, :circuit_breaker, :opened],
%{failure_count: state.failure_count},
%{pool: state.name, reason: :failures}
)
end
defp emit_closed_event(state) do
:telemetry.execute(
[:snakepit, :circuit_breaker, :closed],
%{},
%{pool: state.name}
)
end
defp emit_half_open_event(state) do
:telemetry.execute(
[:snakepit, :circuit_breaker, :half_open],
%{},
%{pool: state.name}
)
end
end
defmodule Snakepit.Compatibility do
@moduledoc """
Thread-safety compatibility matrix for common Python libraries.
"""
@type thread_safety :: true | false | :conditional
@type library_info :: %{
thread_safe: thread_safety(),
notes: String.t()
}
@libraries %{
"numpy" => %{thread_safe: true, notes: "Releases GIL during computation"},
"scipy" => %{thread_safe: true, notes: "Releases GIL for numerical ops"},
"torch" => %{thread_safe: true, notes: "Configure with torch.set_num_threads/1"},
"tensorflow" => %{thread_safe: true, notes: "Use tf.config.threading APIs"},
"scikit-learn" => %{thread_safe: :conditional, notes: "Set n_jobs=1 per estimator"},
"polars" => %{thread_safe: true, notes: "Thread-safe DataFrame library"},
"requests" => %{thread_safe: true, notes: "Use separate Session per thread"},
"httpx" => %{thread_safe: true, notes: "Async-first; thread-safe clients"},
"aiohttp" => %{thread_safe: :conditional, notes: "One ClientSession per thread"},
"grpcio" => %{thread_safe: true, notes: "Thread-safe client with shared channels"},
"numexpr" => %{thread_safe: true, notes: "Releases GIL for expression evaluation"},
"onnxruntime" => %{thread_safe: true, notes: "Thread-safe inference sessions"},
"jax" => %{thread_safe: :conditional, notes: "Avoid shared mutable state"},
"opencv" => %{thread_safe: :conditional, notes: "Avoid shared mutable state"},
"pillow" => %{thread_safe: :conditional, notes: "Use separate Image objects"},
"pandas" => %{thread_safe: false, notes: "Not thread-safe; lock DataFrame ops"},
"matplotlib" => %{thread_safe: false, notes: "Global state; prefer process mode"},
"sqlite3" => %{thread_safe: false, notes: "Use separate connections per thread"},
"sqlalchemy" => %{thread_safe: :conditional, notes: "Use per-thread sessions"},
"spacy" => %{thread_safe: false, notes: "Models share global state"},
"fasttext" => %{thread_safe: false, notes: "Model objects are not thread-safe"},
"xgboost" => %{thread_safe: :conditional, notes: "Set num_threads and avoid shared state"},
"lightgbm" => %{thread_safe: :conditional, notes: "Set num_threads and avoid shared state"},
"cupy" => %{thread_safe: :conditional, notes: "Manage CUDA context per thread"},
"faiss" => %{thread_safe: :conditional, notes: "Avoid shared index mutation"},
"ray" => %{thread_safe: false, notes: "Prefer process-based workers"},
"celery" => %{thread_safe: false, notes: "Use process-based workers"}
}
@spec check(String.t() | atom(), :thread | :process) ::
{:ok, String.t()} | {:warning, String.t()} | {:error, String.t()}
def check(library, profile) when profile in [:thread, :process] do
name = normalize_name(library)
case Map.get(@libraries, name) do
nil ->
{:warning, "Unknown library: #{name}"}
_info when profile == :process ->
{:ok, "Process profile isolates workers"}
%{thread_safe: true} ->
{:ok, "Thread-safe"}
%{thread_safe: false, notes: notes} ->
{:error, "Not thread-safe: #{notes}"}
%{thread_safe: :conditional, notes: notes} ->
{:warning, "Conditionally thread-safe: #{notes}"}
end
end
def check(_library, _profile), do: {:error, "Unknown profile"}
@spec get_library_info(String.t() | atom()) :: library_info() | nil
def get_library_info(library) do
name = normalize_name(library)
Map.get(@libraries, name)
end
@spec list_all(:thread_safe | :thread_unsafe | :conditional | :all) :: [String.t()]
def list_all(:thread_safe) do
select_names(fn {_name, info} -> info.thread_safe == true end)
end
def list_all(:thread_unsafe) do
select_names(fn {_name, info} -> info.thread_safe == false end)
end
def list_all(:conditional) do
select_names(fn {_name, info} -> info.thread_safe == :conditional end)
end
def list_all(:all), do: Map.keys(@libraries)
def list_all(_), do: []
@spec generate_report([String.t() | atom()], :thread | :process) ::
{:ok, map()} | {:error, term()}
def generate_report(libraries, profile) when is_list(libraries) do
report =
Enum.reduce(libraries, %{safe: [], unsafe: [], conditional: [], unknown: []}, fn library,
acc ->
name = normalize_name(library)
case Map.get(@libraries, name) do
nil ->
Map.update!(acc, :unknown, &[name | &1])
%{thread_safe: true} ->
Map.update!(acc, :safe, &[name | &1])
%{thread_safe: false} ->
Map.update!(acc, :unsafe, &[name | &1])
%{thread_safe: :conditional} ->
Map.update!(acc, :conditional, &[name | &1])
end
end)
|> Map.put(:profile, profile)
|> normalize_report()
{:ok, report}
end
def generate_report(_libraries, _profile), do: {:error, :invalid_libraries}
defp normalize_name(library) when is_atom(library),
do: library |> Atom.to_string() |> normalize_name()
defp normalize_name(library) when is_binary(library) do
library
|> String.trim()
|> String.downcase()
|> String.replace("_", "-")
end
defp select_names(fun) do
@libraries
|> Enum.filter(fun)
|> Enum.map(fn {name, _info} -> name end)
end
defp normalize_report(report) do
report
|> Map.update!(:safe, &Enum.reverse/1)
|> Map.update!(:unsafe, &Enum.reverse/1)
|> Map.update!(:conditional, &Enum.reverse/1)
|> Map.update!(:unknown, &Enum.reverse/1)
end
end
defmodule Snakepit.Config do
@moduledoc """
Configuration management for Snakepit pools.
Handles validation and normalization of pool configurations,
supporting both legacy single-pool and new multi-pool configurations.
## Backward Compatibility
Existing v0.5.x configurations continue to work:
# Legacy config (v0.5.x) - still works!
config :snakepit,
pooling_enabled: true,
adapter_module: Snakepit.Adapters.GRPCPython,
pool_size: 100
## New Multi-Pool Configuration (v0.6.0+)
config :snakepit,
pools: [
%{
name: :default,
worker_profile: :process,
pool_size: 100,
adapter_module: Snakepit.Adapters.GRPCPython
},
%{
name: :hpc,
worker_profile: :thread,
pool_size: 4,
threads_per_worker: 16
}
]
## Configuration Schema
Per-pool configuration options:
### Required
- `name` - Pool identifier (atom)
- `adapter_module` - Adapter module
### Profile Selection
- `worker_profile` - `:process` or `:thread` (default: `:process`)
### Common Options
- `pool_size` - Number of workers
- `adapter_args` - CLI arguments for adapter
- `adapter_env` - Environment variables
- `capacity_strategy` - `:pool`, `:profile`, or `:hybrid` (default: `:pool`)
### Process Profile Specific
- `startup_batch_size` - Workers per batch (default: 8)
- `startup_batch_delay_ms` - Delay between batches (default: 750)
### Thread Profile Specific
- `threads_per_worker` - Thread pool size per worker
- `thread_safety_checks` - Enable runtime checks
### Lifecycle Management
- `worker_ttl` - Time-to-live (`:infinity` or `{value, :seconds/:minutes/:hours}`)
- `worker_max_requests` - Max requests before recycling (`:infinity` or integer)
Heartbeat options are mirrored in `snakepit_bridge.heartbeat.HeartbeatConfig`,
so any new keys added here must be added to the Python struct and documented
in the heartbeat guides to keep both sides in sync.
## Normalized Shape
`Snakepit.Config.normalize_pool_config/1` converts user input into a canonical
map that downstream components rely on. The resulting structure (documented
under `t:normalized_pool_config/0`) always includes heartbeat defaults,
adapter metadata, and profile-specific knobs so pool, worker supervisor, and
diagnostics modules never have to pattern-match on partial user input.
"""
require Logger
alias Snakepit.Defaults
alias Snakepit.Logger, as: SLog
@typedoc """
Normalized pool configuration returned by `normalize_pool_config/1`.
```
%{
name: atom(),
worker_profile: :process | :thread,
pool_size: pos_integer(),
adapter_module: module(),
adapter_args: list(),
adapter_env: list(),
capacity_strategy: :pool | :profile | :hybrid,
pool_identifier: atom() | nil,
worker_ttl: :infinity | {integer(), :seconds | :minutes | :hours},
worker_max_requests: :infinity | pos_integer(),
heartbeat: map(),
# Profile-specific fields:
startup_batch_size: pos_integer(),
startup_batch_delay_ms: non_neg_integer(),
threads_per_worker: pos_integer(),
thread_safety_checks: boolean()
}
```
"""
@type pool_config :: map()
@type normalized_pool_config :: map()
@type validation_result :: {:ok, [pool_config()]} | {:error, term()}
# Base heartbeat config - actual runtime defaults are in Snakepit.Defaults
@base_heartbeat_config_template %{
enabled: true,
dependent: true
}
@heartbeat_known_keys [
:enabled,
:ping_interval_ms,
:timeout_ms,
:max_missed_heartbeats,
:initial_delay_ms,
:dependent
]
@heartbeat_string_keys Enum.map(@heartbeat_known_keys, &Atom.to_string/1)
@doc """
Get and validate pool configurations from application environment.
Supports both legacy single-pool and new multi-pool configurations.
Returns `{:ok, [pool_configs]}` or `{:error, reason}`.
## Examples
# With legacy config
{:ok, [%{name: :default, worker_profile: :process, ...}]}
# With multi-pool config
{:ok, [%{name: :default, ...}, %{name: :hpc, ...}]}
"""
@spec get_pool_configs() :: validation_result()
def get_pool_configs do
case Application.get_env(:snakepit, :pools) do
nil ->
# Legacy configuration - convert to new format
convert_legacy_config()
pools when is_list(pools) ->
# New multi-pool configuration
validate_pool_configs(pools)
invalid ->
{:error, {:invalid_pools_config, invalid}}
end
end
@doc """
Validate a single pool configuration.
Returns `{:ok, normalized_config}` or `{:error, reason}`.
"""
@spec validate_pool_config(map()) :: {:ok, pool_config()} | {:error, term()}
def validate_pool_config(config) when is_map(config) do
with :ok <- validate_required_fields(config),
:ok <- validate_profile(config),
:ok <- validate_capacity_strategy(config),
:ok <- validate_pool_size(config),
:ok <- validate_lifecycle_options(config) do
{:ok, normalize_pool_config(config)}
end
end
def validate_pool_config(invalid) do
{:error, {:invalid_pool_config, invalid}}
end
@doc """
Normalize a pool configuration by filling in defaults.
## Examples
iex> Snakepit.Config.normalize_pool_config(%{name: :test})
%{
name: :test,
worker_profile: :process,
pool_size: 16,
# ... other defaults
}
"""
@spec normalize_pool_config(map()) :: pool_config()
def normalize_pool_config(config) do
profile = Map.get(config, :worker_profile, Defaults.default_worker_profile())
capacity_strategy =
Map.get(config, :capacity_strategy) ||
Application.get_env(:snakepit, :capacity_strategy, Defaults.default_capacity_strategy())
base_config =
config
|> Map.put_new(:worker_profile, Defaults.default_worker_profile())
|> Map.put_new(:pool_size, Defaults.default_pool_size())
|> Map.put_new(:adapter_args, [])
|> Map.put_new(:adapter_env, [])
|> Map.put_new(:worker_ttl, :infinity)
|> Map.put_new(:worker_max_requests, :infinity)
|> Map.put(:capacity_strategy, capacity_strategy)
heartbeat =
config
|> Map.get(:heartbeat, %{})
|> normalize_heartbeat_overrides()
|> merge_with_heartbeat_defaults(heartbeat_defaults())
base_config = Map.put(base_config, :heartbeat, heartbeat)
# Add profile-specific defaults
case profile do
:process ->
base_config
|> Map.put_new(:startup_batch_size, Defaults.config_default_batch_size())
|> Map.put_new(:startup_batch_delay_ms, Defaults.config_default_batch_delay())
:thread ->
base_config
|> Map.put_new(:threads_per_worker, Defaults.config_default_threads_per_worker())
|> Map.put_new(:thread_safety_checks, false)
_ ->
base_config
end
end
@doc """
Get configuration for a specific named pool.
Returns `{:ok, config}` or `{:error, reason}`.
The error can be `:pool_not_found` if the pool doesn't exist, or any error
from `get_pool_configs/0` if there's a configuration issue.
## Examples
iex> Snakepit.Config.get_pool_config(:default)
{:ok, %{name: :default, worker_profile: :process, ...}}
"""
@spec get_pool_config(atom()) :: {:ok, pool_config()} | {:error, term()}
def get_pool_config(pool_name) when is_atom(pool_name) do
with {:ok, pools} <- get_pool_configs() do
find_pool_by_name(pools, pool_name)
end
end
defp find_pool_by_name(pools, pool_name) do
case Enum.find(pools, fn pool -> Map.get(pool, :name) == pool_name end) do
nil -> {:error, :pool_not_found}
config -> {:ok, config}
end
end
@doc """
Check if a pool configuration is using the thread profile.
## Examples
iex> Snakepit.Config.thread_profile?(%{worker_profile: :thread})
true
iex> Snakepit.Config.thread_profile?(%{worker_profile: :process})
false
"""
@spec thread_profile?(pool_config()) :: boolean()
def thread_profile?(config) when is_map(config) do
Map.get(config, :worker_profile) == :thread
end
@doc """
Get the profile module for a pool configuration.
Returns the module that implements the WorkerProfile behaviour.
## Examples
iex> Snakepit.Config.get_profile_module(%{worker_profile: :process})
Snakepit.WorkerProfile.Process
iex> Snakepit.Config.get_profile_module(%{worker_profile: :thread})
Snakepit.WorkerProfile.Thread
"""
@spec get_profile_module(pool_config()) :: module()
def get_profile_module(config) when is_map(config) do
case Map.get(config, :worker_profile, :process) do
:process -> Snakepit.WorkerProfile.Process
:thread -> Snakepit.WorkerProfile.Thread
end
end
@doc """
Returns the normalized default heartbeat configuration, merged with application env overrides.
This shape is shared with `snakepit_bridge.heartbeat.HeartbeatConfig`. When adding new keys,
update both modules to keep the cross-language schema aligned.
"""
@spec heartbeat_defaults() :: map()
def heartbeat_defaults do
Application.get_env(:snakepit, :heartbeat, %{})
|> normalize_heartbeat_overrides()
|> merge_with_heartbeat_defaults()
end
# Private functions
defp convert_legacy_config do
# Build a default pool from legacy config options
base_pool = %{
name: :default,
worker_profile: :process,
pool_size: Application.get_env(:snakepit, :pool_size, Defaults.default_pool_size()),
adapter_module: Application.get_env(:snakepit, :adapter_module),
adapter_args: [],
adapter_env: [],
capacity_strategy:
Application.get_env(:snakepit, :capacity_strategy, Defaults.default_capacity_strategy())
}
# Add pool_config if present
legacy_pool =
if pool_config = Application.get_env(:snakepit, :pool_config) do
Map.merge(base_pool, %{
startup_batch_size:
Map.get(pool_config, :startup_batch_size, Defaults.config_default_batch_size()),
startup_batch_delay_ms:
Map.get(pool_config, :startup_batch_delay_ms, Defaults.config_default_batch_delay()),
max_workers: Map.get(pool_config, :max_workers, 1000)
})
else
base_pool
end
case validate_pool_config(legacy_pool) do
{:ok, config} ->
SLog.info(:startup, "Converted legacy configuration to pool config", pool_name: :default)
{:ok, [config]}
error ->
error
end
end
defp validate_pool_configs(pools) do
results =
Enum.map(pools, fn pool ->
validate_pool_config(pool)
end)
errors =
Enum.filter(results, fn
{:error, _} -> true
_ -> false
end)
if Enum.empty?(errors) do
configs = Enum.map(results, fn {:ok, config} -> config end)
check_duplicate_names(configs)
else
{:error, {:validation_failed, errors}}
end
end
defp check_duplicate_names(configs) do
names = Enum.map(configs, & &1.name)
duplicates = names -- Enum.uniq(names)
if Enum.empty?(duplicates) do
{:ok, configs}
else
{:error, {:duplicate_pool_names, duplicates}}
end
end
defp validate_required_fields(config) do
required = [:name]
missing =
Enum.filter(required, fn field ->
not Map.has_key?(config, field) or is_nil(Map.get(config, field))
end)
if Enum.empty?(missing) do
:ok
else
{:error, {:missing_required_fields, missing}}
end
end
defp validate_profile(config) do
case Map.get(config, :worker_profile, Defaults.default_worker_profile()) do
profile when profile in [:process, :thread] ->
:ok
invalid ->
{:error, {:invalid_profile, invalid, "must be :process or :thread"}}
end
end
defp validate_capacity_strategy(config) do
case Map.get(config, :capacity_strategy) do
nil ->
:ok
strategy when strategy in [:pool, :profile, :hybrid] ->
:ok
invalid ->
{:error, {:invalid_capacity_strategy, invalid}}
end
end
defp validate_pool_size(config) do
case Map.get(config, :pool_size, Defaults.default_pool_size()) do
size when is_integer(size) and size > 0 ->
:ok
invalid ->
{:error, {:invalid_pool_size, invalid, "must be positive integer"}}
end
end
defp validate_lifecycle_options(config) do
with :ok <- validate_ttl(config) do
validate_max_requests(config)
end
end
defp validate_ttl(config) do
case Map.get(config, :worker_ttl, :infinity) do
:infinity ->
:ok
{value, unit}
when is_integer(value) and value > 0 and unit in [:seconds, :minutes, :hours] ->
:ok
invalid ->
{:error,
{:invalid_worker_ttl, invalid,
"must be :infinity or {value, :seconds | :minutes | :hours}"}}
end
end
defp validate_max_requests(config) do
case Map.get(config, :worker_max_requests, :infinity) do
:infinity ->
:ok
count when is_integer(count) and count > 0 ->
:ok
invalid ->
{:error, {:invalid_worker_max_requests, invalid, "must be :infinity or positive integer"}}
end
end
defp normalize_heartbeat_overrides(overrides) when is_map(overrides) do
Enum.reduce(overrides, %{}, fn {key, value}, acc ->
Map.put(acc, normalize_heartbeat_key(key), value)
end)
end
defp normalize_heartbeat_overrides(_), do: %{}
defp default_heartbeat_config do
Map.merge(@base_heartbeat_config_template, %{
ping_interval_ms: Defaults.heartbeat_ping_interval_ms(),
timeout_ms: Defaults.heartbeat_timeout_ms(),
max_missed_heartbeats: Defaults.heartbeat_max_missed(),
initial_delay_ms: Defaults.heartbeat_initial_delay_ms()
})
end
defp merge_with_heartbeat_defaults(overrides, base \\ nil) do
base = base || default_heartbeat_config()
defaults =
Enum.reduce(base, %{}, fn {key, value}, acc ->
Map.put(acc, key, value)
end)
Map.merge(defaults, overrides, fn _key, _default, override -> override end)
end
defp normalize_heartbeat_key(key) when is_atom(key) do
if key in @heartbeat_known_keys do
key
else
key
end
end
defp normalize_heartbeat_key(key) when is_binary(key) do
if key in @heartbeat_string_keys do
String.to_existing_atom(key)
else
key
end
rescue
ArgumentError ->
key
end
end
defmodule Snakepit.CrashBarrier do
@moduledoc """
Crash barrier policy for worker failures.
Classifies worker crashes, taints unstable workers, and determines retry eligibility.
"""
alias Snakepit.Defaults
alias Snakepit.Error
alias Snakepit.Worker.TaintRegistry
defp default_config do
%{
enabled: false,
retry: :idempotent,
max_restarts: Defaults.crash_barrier_max_restarts(),
taint_duration_ms: Defaults.crash_barrier_taint_duration_ms(),
backoff_ms: Defaults.crash_barrier_backoff_ms(),
mark_on: [:segfault, :oom, :gpu],
taint_on_exit_codes: [],
taint_on_error_types: [],
taint_device_on_cuda_fatal: true
}
end
@segfault_codes [139]
@abort_codes [134]
@oom_codes [137]
def config(pool_config \\ %{}) do
global =
:snakepit
|> Application.get_env(:crash_barrier, [])
|> Map.new()
pool_override =
pool_config
|> Map.get(:crash_barrier, %{})
|> Map.new()
default_config()
|> Map.merge(global)
|> Map.merge(pool_override)
|> normalize_config()
end
def enabled?(config), do: config.enabled == true
def idempotent?(args) when is_map(args) do
Map.get(args, :idempotent) || Map.get(args, "idempotent") || false
end
def idempotent?(_), do: false
def crash_info({:error, {:worker_exit, reason}}, config) do
classify_reason(reason, config)
end
def crash_info({:error, {:worker_crash, info}}, _config) do
{:ok, info}
end
def crash_info(_result, _config), do: :error
def retry_allowed?(config, idempotent, attempt) do
max_retries = retry_limit(config)
cond do
attempt >= max_retries -> false
config.retry == :always -> true
config.retry == :idempotent -> idempotent
true -> false
end
end
def retry_backoff(config, attempt) do
backoffs = List.wrap(config.backoff_ms)
case Enum.at(backoffs, attempt - 1) do
nil -> List.last(backoffs) || 0
value -> value
end
end
def taint_worker(pool_name, worker_id, info, config) do
TaintRegistry.taint_worker(worker_id,
duration_ms: config.taint_duration_ms,
reason: info.reason,
exit_code: info.exit_code,
device: info.device
)
emit_worker_tainted(pool_name, worker_id, info, config)
:ok
end
def worker_tainted?(worker_id), do: TaintRegistry.worker_tainted?(worker_id)
def maybe_emit_restart(pool_name, worker_id) do
case TaintRegistry.consume_restart(worker_id) do
{:ok, info} ->
:telemetry.execute(
[:snakepit, :worker, :restarted],
%{},
%{
worker_id: worker_id,
pool: pool_name,
reason: info.reason,
exit_code: info.exit_code,
device: info.device
}
)
:error ->
:ok
end
end
def normalize_crash_error({:error, {:worker_exit, reason}}, info) do
{:error,
Error.worker_error("Python worker crashed", %{
type: :worker_crash,
reason: reason,
exit_code: info.exit_code,
classification: info.classification
})}
end
def normalize_crash_error(result, _info), do: result
defp classify_reason(reason, config) do
exit_code = exit_code_from_reason(reason)
reason_string = inspect(reason)
error_type = error_type_from_reason(reason_string, config)
classification =
cond do
exit_code in @segfault_codes -> :segfault
exit_code in @abort_codes -> :abort
exit_code in @oom_codes -> :oom
error_type == :gpu -> :gpu
true -> :unknown
end
if should_taint?(config, exit_code, error_type, classification) do
{:ok,
%{
reason: reason,
exit_code: exit_code,
classification: classification,
device: device_from_reason(reason)
}}
else
:error
end
end
defp should_taint?(config, exit_code, error_type, classification) do
mark_on = normalize_mark_on(config.mark_on)
cond do
exit_code in taint_exit_codes(config, mark_on) ->
true
error_type == :gpu and :gpu in mark_on ->
true
classification in mark_on ->
true
true ->
false
end
end
defp taint_exit_codes(config, mark_on) do
config.taint_on_exit_codes
|> List.wrap()
|> Kernel.++(exit_codes_for(mark_on))
|> Enum.uniq()
end
defp exit_codes_for(mark_on) do
[]
|> maybe_append(:segfault in mark_on, @segfault_codes)
|> maybe_append(:abort in mark_on, @abort_codes)
|> maybe_append(:oom in mark_on, @oom_codes)
end
defp maybe_append(list, true, values), do: list ++ values
defp maybe_append(list, false, _values), do: list
defp normalize_mark_on(mark_on) do
mark_on
|> List.wrap()
|> Enum.map(fn
value when is_atom(value) -> value
value when is_binary(value) -> normalize_mark_atom(value)
_ -> :unknown
end)
end
defp normalize_mark_atom(value) do
case String.downcase(value) do
"segfault" -> :segfault
"oom" -> :oom
"gpu" -> :gpu
"abort" -> :abort
_ -> :unknown
end
end
defp error_type_from_reason(reason_string, config) do
reason_string = String.downcase(reason_string)
configured = config.taint_on_error_types |> List.wrap() |> Enum.map(&String.downcase/1)
cond do
Enum.any?(configured, &String.contains?(reason_string, &1)) -> :gpu
String.contains?(reason_string, "cuda") -> :gpu
String.contains?(reason_string, "gpu") -> :gpu
true -> :unknown
end
end
defp device_from_reason(reason) do
if String.contains?(inspect(reason), "CUDA") do
:cuda
else
nil
end
end
defp exit_code_from_reason({:grpc_server_exited, status}) when is_integer(status), do: status
defp exit_code_from_reason({:exit_status, status}) when is_integer(status), do: status
defp exit_code_from_reason({:external_process_died, {:exit_status, status}}), do: status
defp exit_code_from_reason({:shutdown, {:grpc_server_exited, status}}), do: status
defp exit_code_from_reason(_), do: nil
defp retry_limit(config) do
Map.get(config, :max_restarts) || Map.get(config, :max_retries) || 1
end
defp normalize_config(config) do
config
|> Map.update(:retry, :idempotent, fn
true -> :always
false -> :never
value -> value
end)
|> Map.merge(retry_from_legacy(config))
end
defp retry_from_legacy(config) do
case Map.fetch(config, :retry_idempotent) do
{:ok, true} -> %{retry: :idempotent}
{:ok, false} -> %{retry: :never}
:error -> %{}
end
end
defp emit_worker_tainted(pool_name, worker_id, info, config) do
:telemetry.execute(
[:snakepit, :worker, :crash],
%{},
%{
worker_id: worker_id,
pool: pool_name,
reason: info.reason,
exit_code: info.exit_code,
device: info.device
}
)
:telemetry.execute(
[:snakepit, :worker, :tainted],
%{duration_ms: config.taint_duration_ms},
%{
worker_id: worker_id,
pool: pool_name,
reason: info.reason,
exit_code: info.exit_code,
device: info.device
}
)
end
end
defmodule Snakepit.Defaults do
@moduledoc """
Centralized defaults for all configurable values in Snakepit.
This module provides functions that read from `Application.get_env(:snakepit, key, default)`
for every configurable value. This allows operators to override defaults via application
configuration while maintaining backward compatibility.
All defaults are the EXACT values that were previously hardcoded throughout the codebase.
Snakepit behaves identically before and after this change unless configuration is explicitly
provided.
## Configuration Example
# config/runtime.exs
config :snakepit,
# Timeouts
default_command_timeout: 30_000,
pool_request_timeout: 60_000,
pool_streaming_timeout: 300_000,
pool_startup_timeout: 10_000,
pool_queue_timeout: 5_000,
checkout_timeout: 5_000,
grpc_worker_execute_timeout: 30_000,
grpc_worker_stream_timeout: 300_000,
grpc_command_timeout: 30_000,
executor_batch_timeout: 30_000,
health_check_interval: 30_000,
circuit_breaker_reset_timeout: 30_000,
graceful_shutdown_timeout_ms: 6_000,
# Pool settings
pool_max_queue_size: 1000,
pool_max_workers: 150,
pool_max_cancelled_entries: 1024,
pool_cancelled_retention_multiplier: 4,
pool_startup_batch_size: 10,
pool_startup_batch_delay_ms: 500,
# Retry settings
retry_max_attempts: 3,
retry_max_backoff_ms: 30_000,
retry_jitter_factor: 0.25,
retry_backoff_sequence: [100, 200, 400, 800, 1600],
# Circuit breaker settings
circuit_breaker_failure_threshold: 5,
circuit_breaker_half_open_max_calls: 1,
# Crash barrier settings
crash_barrier_taint_duration_ms: 60_000,
crash_barrier_max_restarts: 1,
crash_barrier_backoff_ms: [50, 100, 200],
# Health monitor settings
health_monitor_crash_window_ms: 60_000,
health_monitor_max_crashes: 10,
# Lifecycle manager settings
lifecycle_check_interval: 60_000,
lifecycle_health_check_interval: 300_000,
# Session store settings
session_cleanup_interval: 60_000,
session_default_ttl: 3600,
session_max_sessions: 10_000,
session_warning_threshold: 0.8,
# Process registry settings
process_registry_cleanup_interval: 30_000,
process_registry_unregister_cleanup_delay: 500,
process_registry_unregister_cleanup_attempts: 10,
# gRPC settings
grpc_num_acceptors: 20,
grpc_max_connections: 1000,
grpc_socket_backlog: 512,
# Heartbeat settings
heartbeat_ping_interval_ms: 2_000,
heartbeat_timeout_ms: 10_000,
heartbeat_max_missed: 3,
heartbeat_initial_delay_ms: 0
## Usage
Instead of hardcoding values like `30_000`, modules now call:
Snakepit.Defaults.default_command_timeout()
This returns the configured value or the original default if not configured.
## Timeout Profiles (v0.8.8+)
Snakepit supports profile-based timeout configuration for different deployment scenarios:
| Profile | default_timeout | stream_timeout | queue_timeout |
|---------|-----------------|----------------|---------------|
| :balanced | 300_000 (5m) | 900_000 (15m) | 10_000 (10s) |
| :production | 300_000 (5m) | 900_000 (15m) | 10_000 (10s) |
| :production_strict | 60_000 (60s) | 300_000 (5m) | 5_000 (5s) |
| :development | 900_000 (15m) | 3_600_000 (60m) | 60_000 (60s) |
| :ml_inference | 900_000 (15m) | 3_600_000 (60m) | 60_000 (60s) |
| :batch | 3_600_000 (60m) | :infinity | 300_000 (5m) |
Configure via:
config :snakepit, timeout_profile: :production
Legacy per-key configuration is still supported and takes precedence when set.
"""
# ============================================================================
# Timeout Profiles (NEW API)
# ============================================================================
@timeout_profiles %{
balanced: %{
default_timeout: 300_000,
stream_timeout: 900_000,
queue_timeout: 10_000
},
production: %{
default_timeout: 300_000,
stream_timeout: 900_000,
queue_timeout: 10_000
},
production_strict: %{
default_timeout: 60_000,
stream_timeout: 300_000,
queue_timeout: 5_000
},
development: %{
default_timeout: 900_000,
stream_timeout: 3_600_000,
queue_timeout: 60_000
},
ml_inference: %{
default_timeout: 900_000,
stream_timeout: 3_600_000,
queue_timeout: 60_000
},
batch: %{
default_timeout: 3_600_000,
stream_timeout: :infinity,
queue_timeout: 300_000
}
}
@doc """
Returns all available timeout profiles.
Each profile contains `default_timeout`, `stream_timeout`, and `queue_timeout` values.
"""
@spec timeout_profiles() :: %{atom() => %{atom() => timeout()}}
def timeout_profiles, do: @timeout_profiles
@doc """
Returns the currently configured timeout profile.
Defaults to `:balanced` if not configured.
"""
@spec timeout_profile() :: atom()
def timeout_profile do
Application.get_env(:snakepit, :timeout_profile, :balanced)
end
@doc """
Returns the default timeout for regular execute operations based on the current profile.
This is the primary user-facing timeout API. Legacy getters derive from this value
when not explicitly configured.
"""
@spec default_timeout() :: timeout()
def default_timeout do
profile = timeout_profile()
get_in(@timeout_profiles, [profile, :default_timeout]) || 300_000
end
@doc """
Returns the default timeout for streaming operations based on the current profile.
"""
@spec stream_timeout() :: timeout()
def stream_timeout do
profile = timeout_profile()
get_in(@timeout_profiles, [profile, :stream_timeout]) || 900_000
end
@doc """
Returns the default queue timeout based on the current profile.
"""
@spec queue_timeout() :: timeout()
def queue_timeout do
profile = timeout_profile()
get_in(@timeout_profiles, [profile, :queue_timeout]) || 10_000
end
# ============================================================================
# Margin Configuration
# ============================================================================
@doc """
Margin reserved for GenServer.call overhead when routing to workers.
This is subtracted from the total timeout budget to derive the RPC timeout.
Default: 1000 ms
"""
@spec worker_call_margin_ms() :: pos_integer()
def worker_call_margin_ms do
Application.get_env(:snakepit, :worker_call_margin_ms, 1000)
end
@doc """
Margin reserved for pool reply overhead.
This is subtracted from the total timeout budget to derive the RPC timeout.
Default: 200 ms
"""
@spec pool_reply_margin_ms() :: pos_integer()
def pool_reply_margin_ms do
Application.get_env(:snakepit, :pool_reply_margin_ms, 200)
end
@doc """
Derives the RPC (inner) timeout from the total timeout budget.
Formula: `rpc_timeout = total_timeout - worker_call_margin_ms - pool_reply_margin_ms`
This ensures inner timeouts expire before outer GenServer.call timeouts,
producing structured error returns instead of unhandled exits.
## Examples
iex> Snakepit.Defaults.rpc_timeout(60_000)
58_800 # 60_000 - 1000 - 200
iex> Snakepit.Defaults.rpc_timeout(:infinity)
:infinity
"""
@spec rpc_timeout(timeout()) :: timeout()
def rpc_timeout(:infinity), do: :infinity
def rpc_timeout(total_timeout) when is_integer(total_timeout) do
margins = worker_call_margin_ms() + pool_reply_margin_ms()
result = total_timeout - margins
# Floor at minimum usable timeout
max(result, 1)
end
# ============================================================================
# Pool Timeouts (Legacy API - now derives from new budgets)
# ============================================================================
@doc """
Default timeout for pool execute calls.
Used in `Snakepit.Pool.execute/3`.
When not explicitly configured, derives from `default_timeout/0` based on the
current timeout profile.
Default: derived from profile (300_000 ms for :balanced)
"""
@spec pool_request_timeout() :: timeout()
def pool_request_timeout do
Application.get_env(:snakepit, :pool_request_timeout) || default_timeout()
end
@doc """
Default timeout for pool streaming calls.
Used in `Snakepit.Pool.execute_stream/4`.
When not explicitly configured, derives from `stream_timeout/0` based on the
current timeout profile.
Default: derived from profile (900_000 ms for :balanced)
"""
@spec pool_streaming_timeout() :: timeout()
def pool_streaming_timeout do
Application.get_env(:snakepit, :pool_streaming_timeout) || stream_timeout()
end
@doc """
Default timeout for worker startup.
Used in pool initialization.
Default: 10_000 ms (10 seconds)
"""
@spec pool_startup_timeout() :: timeout()
def pool_startup_timeout do
Application.get_env(:snakepit, :pool_startup_timeout, 10_000)
end
@doc """
Default timeout for queued requests.
Used in `Snakepit.Pool` for queue management.
When not explicitly configured, derives from `queue_timeout/0` based on the
current timeout profile.
Default: derived from profile (10_000 ms for :balanced)
"""
@spec pool_queue_timeout() :: timeout()
def pool_queue_timeout do
Application.get_env(:snakepit, :pool_queue_timeout) || queue_timeout()
end
@doc """
Default timeout for checking out a worker for streaming.
Used in `Snakepit.Pool` for worker checkout during streaming operations.
When not explicitly configured, derives from `queue_timeout/0` based on the
current timeout profile.
Default: derived from profile (10_000 ms for :balanced)
"""
@spec checkout_timeout() :: timeout()
def checkout_timeout do
Application.get_env(:snakepit, :checkout_timeout) || queue_timeout()
end
@doc """
Default command timeout for worker execute operations.
Used in `Snakepit.Pool` for command timeout calculation.
When not explicitly configured, derives from `rpc_timeout(default_timeout())` based on the
current timeout profile.
Default: derived from profile (rpc_timeout of default_timeout)
"""
@spec default_command_timeout() :: timeout()
def default_command_timeout do
Application.get_env(:snakepit, :default_command_timeout) || rpc_timeout(default_timeout())
end
# ============================================================================
# Pool Sizing
# ============================================================================
@doc """
Default pool size based on system schedulers.
Used when no explicit pool_size is configured.
Default: System.schedulers_online() * 2
"""
@spec default_pool_size() :: pos_integer()
def default_pool_size do
Application.get_env(:snakepit, :default_pool_size, System.schedulers_online() * 2)
end
@doc """
Maximum queue size for pending requests.
Used in `Snakepit.Pool` for queue management.
Default: 1000
"""
@spec pool_max_queue_size() :: pos_integer()
def pool_max_queue_size do
Application.get_env(:snakepit, :pool_max_queue_size, 1000)
end
@doc """
Maximum number of workers allowed per pool.
Used in `Snakepit.Pool` for worker limit enforcement.
Default: 150
"""
@spec pool_max_workers() :: pos_integer()
def pool_max_workers do
Application.get_env(:snakepit, :pool_max_workers, 150)
end
@doc """
Maximum number of cancelled request entries to track.
Used in `Snakepit.Pool` for cancelled request management.
Default: 1024
"""
@spec pool_max_cancelled_entries() :: pos_integer()
def pool_max_cancelled_entries do
Application.get_env(:snakepit, :pool_max_cancelled_entries, 1024)
end
@doc """
Multiplier for cancelled request retention time.
Retention time = queue_timeout * this multiplier.
Default: 4
"""
@spec pool_cancelled_retention_multiplier() :: pos_integer()
def pool_cancelled_retention_multiplier do
Application.get_env(:snakepit, :pool_cancelled_retention_multiplier, 4)
end
@doc """
Number of workers to start per batch during pool initialization.
Used in `Snakepit.Pool` for batched startup.
Default: 10
"""
@spec pool_startup_batch_size() :: pos_integer()
def pool_startup_batch_size do
Application.get_env(:snakepit, :pool_startup_batch_size, 10)
end
@doc """
Delay between worker startup batches in milliseconds.
Used in `Snakepit.Pool` for batched startup.
Default: 500 ms
"""
@spec pool_startup_batch_delay_ms() :: non_neg_integer()
def pool_startup_batch_delay_ms do
Application.get_env(:snakepit, :pool_startup_batch_delay_ms, 500)
end
# ============================================================================
# gRPC Worker
# ============================================================================
@doc """
Default timeout for GRPCWorker execute calls.
Used in `Snakepit.GRPCWorker.execute/4`.
When not explicitly configured, derives from `rpc_timeout(default_timeout())` based on the
current timeout profile.
Default: derived from profile (rpc_timeout of default_timeout)
"""
@spec grpc_worker_execute_timeout() :: timeout()
def grpc_worker_execute_timeout do
Application.get_env(:snakepit, :grpc_worker_execute_timeout) || rpc_timeout(default_timeout())
end
@doc """
Default timeout for GRPCWorker streaming calls.
Used in `Snakepit.GRPCWorker.execute_stream/5`.
Default: 300_000 ms (5 minutes)
"""
@spec grpc_worker_stream_timeout() :: timeout()
def grpc_worker_stream_timeout do
Application.get_env(:snakepit, :grpc_worker_stream_timeout, 300_000)
end
@doc """
Graceful shutdown timeout for Python process termination.
Must be >= Python's shutdown envelope: server.stop(2s) + wait_for_termination(3s) = 5s.
Default: 6_000 ms (6 seconds)
"""
@spec graceful_shutdown_timeout_ms() :: pos_integer()
def graceful_shutdown_timeout_ms do
Application.get_env(:snakepit, :graceful_shutdown_timeout_ms, 6_000)
end
@doc """
Margin added to graceful_shutdown_timeout for supervisor shutdown.
This gives the worker time to complete its terminate/2 callback.
Default: 2_000 ms (2 seconds)
"""
@spec shutdown_margin_ms() :: pos_integer()
def shutdown_margin_ms do
Application.get_env(:snakepit, :shutdown_margin_ms, 2_000)
end
@doc """
Interval for health checks in GRPCWorker.
Used in `Snakepit.GRPCWorker` for periodic health check scheduling.
Default: 30_000 ms (30 seconds)
"""
@spec grpc_worker_health_check_interval() :: pos_integer()
def grpc_worker_health_check_interval do
Application.get_env(:snakepit, :grpc_worker_health_check_interval, 30_000)
end
# ============================================================================
# Heartbeat
# ============================================================================
@doc """
Default heartbeat ping interval.
Used in `Snakepit.GRPCWorker` heartbeat configuration.
Default: 2_000 ms (2 seconds)
"""
@spec heartbeat_ping_interval_ms() :: pos_integer()
def heartbeat_ping_interval_ms do
Application.get_env(:snakepit, :heartbeat_ping_interval_ms, 2_000)
end
@doc """
Default heartbeat timeout.
Used in `Snakepit.GRPCWorker` heartbeat configuration.
Default: 10_000 ms (10 seconds)
"""
@spec heartbeat_timeout_ms() :: pos_integer()
def heartbeat_timeout_ms do
Application.get_env(:snakepit, :heartbeat_timeout_ms, 10_000)
end
@doc """
Maximum missed heartbeats before worker is considered unhealthy.
Used in `Snakepit.GRPCWorker` heartbeat configuration.
Default: 3
"""
@spec heartbeat_max_missed() :: pos_integer()
def heartbeat_max_missed do
Application.get_env(:snakepit, :heartbeat_max_missed, 3)
end
@doc """
Initial delay before starting heartbeat monitoring.
Used in `Snakepit.GRPCWorker` heartbeat configuration.
Default: 0 ms
"""
@spec heartbeat_initial_delay_ms() :: non_neg_integer()
def heartbeat_initial_delay_ms do
Application.get_env(:snakepit, :heartbeat_initial_delay_ms, 0)
end
# ============================================================================
# gRPC Python Adapter
# ============================================================================
@doc """
Default command timeout for gRPC adapter.
Used in `Snakepit.Adapters.GRPCPython` for default command timeouts.
When not explicitly configured, derives from `rpc_timeout(default_timeout())` based on the
current timeout profile.
Default: derived from profile (rpc_timeout of default_timeout)
"""
@spec grpc_command_timeout() :: timeout()
def grpc_command_timeout do
Application.get_env(:snakepit, :grpc_command_timeout) || rpc_timeout(default_timeout())
end
@doc """
Timeout for batch inference commands.
Used in `Snakepit.Adapters.GRPCPython` for batch inference operations.
Default: 300_000 ms (5 minutes)
"""
@spec grpc_batch_inference_timeout() :: timeout()
def grpc_batch_inference_timeout do
Application.get_env(:snakepit, :grpc_batch_inference_timeout, 300_000)
end
@doc """
Timeout for large dataset processing commands.
Used in `Snakepit.Adapters.GRPCPython` for large dataset processing operations.
Default: 600_000 ms (10 minutes)
"""
@spec grpc_large_dataset_timeout() :: timeout()
def grpc_large_dataset_timeout do
Application.get_env(:snakepit, :grpc_large_dataset_timeout, 600_000)
end
# ============================================================================
# Executor
# ============================================================================
@doc """
Default timeout for batch operations in Executor.
Used in `Snakepit.Executor.execute_batch/2`.
Default: 30_000 ms (30 seconds)
"""
@spec executor_batch_timeout() :: timeout()
def executor_batch_timeout do
Application.get_env(:snakepit, :executor_batch_timeout, 30_000)
end
# ============================================================================
# Health Monitor
# ============================================================================
@doc """
Default interval for health monitor cleanup.
Used in `Snakepit.HealthMonitor`.
Default: 30_000 ms (30 seconds)
"""
@spec health_monitor_check_interval() :: pos_integer()
def health_monitor_check_interval do
Application.get_env(:snakepit, :health_monitor_check_interval, 30_000)
end
@doc """
Default crash window for health monitor.
Rolling window for crash counting.
Default: 60_000 ms (1 minute)
"""
@spec health_monitor_crash_window_ms() :: pos_integer()
def health_monitor_crash_window_ms do
Application.get_env(:snakepit, :health_monitor_crash_window_ms, 60_000)
end
@doc """
Default max crashes before pool is considered unhealthy.
Used in `Snakepit.HealthMonitor`.
Default: 10
"""
@spec health_monitor_max_crashes() :: pos_integer()
def health_monitor_max_crashes do
Application.get_env(:snakepit, :health_monitor_max_crashes, 10)
end
# ============================================================================
# Retry Policy
# ============================================================================
@doc """
Default maximum retry attempts.
Used in `Snakepit.RetryPolicy`.
Default: 3
"""
@spec retry_max_attempts() :: pos_integer()
def retry_max_attempts do
Application.get_env(:snakepit, :retry_max_attempts, 3)
end
@doc """
Default backoff sequence for retries.
Used in `Snakepit.RetryPolicy`.
Default: [100, 200, 400, 800, 1600]
"""
@spec retry_backoff_sequence() :: [pos_integer()]
def retry_backoff_sequence do
Application.get_env(:snakepit, :retry_backoff_sequence, [100, 200, 400, 800, 1600])
end
@doc """
Default base backoff for retry calculations.
Used in `Snakepit.RetryPolicy`.
Default: 100 ms
"""
@spec retry_base_backoff_ms() :: pos_integer()
def retry_base_backoff_ms do
Application.get_env(:snakepit, :retry_base_backoff_ms, 100)
end
@doc """
Default maximum backoff delay.
Used in `Snakepit.RetryPolicy`.
Default: 30_000 ms (30 seconds)
"""
@spec retry_max_backoff_ms() :: pos_integer()
def retry_max_backoff_ms do
Application.get_env(:snakepit, :retry_max_backoff_ms, 30_000)
end
@doc """
Default backoff multiplier for exponential backoff.
Used in `Snakepit.RetryPolicy`.
Default: 2.0
"""
@spec retry_backoff_multiplier() :: float()
def retry_backoff_multiplier do
Application.get_env(:snakepit, :retry_backoff_multiplier, 2.0)
end
@doc """
Default jitter factor for retry delays.
Used in `Snakepit.RetryPolicy`.
Default: 0.25 (25%)
"""
@spec retry_jitter_factor() :: float()
def retry_jitter_factor do
Application.get_env(:snakepit, :retry_jitter_factor, 0.25)
end
# ============================================================================
# Circuit Breaker
# ============================================================================
@doc """
Default failure threshold before circuit opens.
Used in `Snakepit.CircuitBreaker`.
Default: 5
"""
@spec circuit_breaker_failure_threshold() :: pos_integer()
def circuit_breaker_failure_threshold do
Application.get_env(:snakepit, :circuit_breaker_failure_threshold, 5)
end
@doc """
Default reset timeout before transitioning to half-open.
Used in `Snakepit.CircuitBreaker`.
Default: 30_000 ms (30 seconds)
"""
@spec circuit_breaker_reset_timeout_ms() :: pos_integer()
def circuit_breaker_reset_timeout_ms do
Application.get_env(:snakepit, :circuit_breaker_reset_timeout_ms, 30_000)
end
@doc """
Default max calls allowed in half-open state.
Used in `Snakepit.CircuitBreaker`.
Default: 1
"""
@spec circuit_breaker_half_open_max_calls() :: pos_integer()
def circuit_breaker_half_open_max_calls do
Application.get_env(:snakepit, :circuit_breaker_half_open_max_calls, 1)
end
# ============================================================================
# Crash Barrier
# ============================================================================
@doc """
Default taint duration for crashed workers.
Used in `Snakepit.CrashBarrier`.
Default: 60_000 ms (1 minute)
"""
@spec crash_barrier_taint_duration_ms() :: pos_integer()
def crash_barrier_taint_duration_ms do
Application.get_env(:snakepit, :crash_barrier_taint_duration_ms, 60_000)
end
@doc """
Default max restarts for crash barrier retry.
Used in `Snakepit.CrashBarrier`.
Default: 1
"""
@spec crash_barrier_max_restarts() :: pos_integer()
def crash_barrier_max_restarts do
Application.get_env(:snakepit, :crash_barrier_max_restarts, 1)
end
@doc """
Default backoff sequence for crash barrier retries.
Used in `Snakepit.CrashBarrier`.
Default: [50, 100, 200]
"""
@spec crash_barrier_backoff_ms() :: [pos_integer()]
def crash_barrier_backoff_ms do
Application.get_env(:snakepit, :crash_barrier_backoff_ms, [50, 100, 200])
end
# ============================================================================
# Lifecycle Manager
# ============================================================================
@doc """
Default interval for lifecycle checks.
Used in `Snakepit.Worker.LifecycleManager`.
Default: 60_000 ms (1 minute)
"""
@spec lifecycle_check_interval() :: pos_integer()
def lifecycle_check_interval do
Application.get_env(:snakepit, :lifecycle_check_interval, 60_000)
end
@doc """
Default interval for health checks in lifecycle manager.
Used in `Snakepit.Worker.LifecycleManager`.
Default: 300_000 ms (5 minutes)
"""
@spec lifecycle_health_check_interval() :: pos_integer()
def lifecycle_health_check_interval do
Application.get_env(:snakepit, :lifecycle_health_check_interval, 300_000)
end
# ============================================================================
# Session Store
# ============================================================================
@doc """
Default cleanup interval for expired sessions.
Used in `Snakepit.Bridge.SessionStore`.
Default: 60_000 ms (1 minute)
"""
@spec session_cleanup_interval() :: pos_integer()
def session_cleanup_interval do
Application.get_env(:snakepit, :session_cleanup_interval, 60_000)
end
@doc """
Default TTL for sessions in seconds.
Used in `Snakepit.Bridge.SessionStore`.
Default: 3600 seconds (1 hour)
"""
@spec session_default_ttl() :: pos_integer()
def session_default_ttl do
Application.get_env(:snakepit, :session_default_ttl, 3600)
end
@doc """
Default maximum number of sessions.
Used in `Snakepit.Bridge.SessionStore`.
Default: 10_000
"""
@spec session_max_sessions() :: pos_integer() | :infinity
def session_max_sessions do
Application.get_env(:snakepit, :session_max_sessions, 10_000)
end
@doc """
Session warning threshold as a fraction of max_sessions.
When session count exceeds this percentage, warnings are emitted.
Default: 0.8 (80%)
"""
@spec session_warning_threshold() :: float()
def session_warning_threshold do
Application.get_env(:snakepit, :session_warning_threshold, 0.8)
end
# ============================================================================
# Process Registry
# ============================================================================
@doc """
Default cleanup interval for process registry.
Used in `Snakepit.Pool.ProcessRegistry`.
Default: 30_000 ms (30 seconds)
"""
@spec process_registry_cleanup_interval() :: pos_integer()
def process_registry_cleanup_interval do
Application.get_env(:snakepit, :process_registry_cleanup_interval, 30_000)
end
@doc """
Delay before retrying unregister when external process is still alive.
Used in `Snakepit.Pool.ProcessRegistry`.
Default: 500 ms
"""
@spec process_registry_unregister_cleanup_delay() :: pos_integer()
def process_registry_unregister_cleanup_delay do
Application.get_env(:snakepit, :process_registry_unregister_cleanup_delay, 500)
end
@doc """
Maximum attempts to retry unregister cleanup.
Used in `Snakepit.Pool.ProcessRegistry`.
Default: 10
"""
@spec process_registry_unregister_cleanup_attempts() :: pos_integer()
def process_registry_unregister_cleanup_attempts do
Application.get_env(:snakepit, :process_registry_unregister_cleanup_attempts, 10)
end
# ============================================================================
# gRPC Server Configuration
# ============================================================================
@doc """
Default number of acceptors for gRPC server.
Used in `Snakepit.Application`.
Default: 20
"""
@spec grpc_num_acceptors() :: pos_integer()
def grpc_num_acceptors do
Application.get_env(:snakepit, :grpc_num_acceptors, 20)
end
@doc """
Default maximum connections for gRPC server.
Used in `Snakepit.Application`.
Default: 1000
"""
@spec grpc_max_connections() :: pos_integer()
def grpc_max_connections do
Application.get_env(:snakepit, :grpc_max_connections, 1000)
end
@doc """
Default socket backlog for gRPC server.
Used in `Snakepit.Application`.
Default: 512
"""
@spec grpc_socket_backlog() :: pos_integer()
def grpc_socket_backlog do
Application.get_env(:snakepit, :grpc_socket_backlog, 512)
end
# ============================================================================
# Config Module Defaults
# ============================================================================
@doc """
Default worker profile.
Used in `Snakepit.Config`.
Default: :process
"""
@spec default_worker_profile() :: :process | :thread
def default_worker_profile do
Application.get_env(:snakepit, :default_worker_profile, :process)
end
@doc """
Default batch size for process profile.
Used in `Snakepit.Config`.
Default: 8
"""
@spec config_default_batch_size() :: pos_integer()
def config_default_batch_size do
Application.get_env(:snakepit, :config_default_batch_size, 8)
end
@doc """
Default batch delay for process profile.
Used in `Snakepit.Config`.
Default: 750 ms
"""
@spec config_default_batch_delay() :: pos_integer()
def config_default_batch_delay do
Application.get_env(:snakepit, :config_default_batch_delay, 750)
end
@doc """
Default threads per worker for thread profile.
Used in `Snakepit.Config`.
Default: 10
"""
@spec config_default_threads_per_worker() :: pos_integer()
def config_default_threads_per_worker do
Application.get_env(:snakepit, :config_default_threads_per_worker, 10)
end
@doc """
Default capacity strategy.
Used in `Snakepit.Config`.
Default: :pool
"""
@spec default_capacity_strategy() :: :pool | :profile | :hybrid
def default_capacity_strategy do
Application.get_env(:snakepit, :default_capacity_strategy, :pool)
end
# ============================================================================
# gRPC Client
# ============================================================================
@doc """
Default timeout for gRPC client execute calls.
Used in `Snakepit.GRPC.Client`.
Default: 30_000 ms (30 seconds)
"""
@spec grpc_client_execute_timeout() :: timeout()
def grpc_client_execute_timeout do
Application.get_env(:snakepit, :grpc_client_execute_timeout, 30_000)
end
# ============================================================================
# Application
# ============================================================================
@doc """
Default gRPC port for Elixir server.
Used in `Snakepit.Application`.
Default: 50_051
"""
@spec grpc_port() :: pos_integer()
def grpc_port do
Application.get_env(:snakepit, :grpc_port, 50_051)
end
@doc """
Timeout for cleanup on stop.
Used in `Snakepit.Application`.
Default: 3_000 ms (3 seconds)
"""
@spec cleanup_on_stop_timeout_ms() :: pos_integer()
def cleanup_on_stop_timeout_ms do
Application.get_env(:snakepit, :cleanup_on_stop_timeout_ms, 3_000)
end
@doc """
Poll interval for cleanup operations.
Used in `Snakepit.Application`.
Default: 50 ms
"""
@spec cleanup_poll_interval_ms() :: pos_integer()
def cleanup_poll_interval_ms do
Application.get_env(:snakepit, :cleanup_poll_interval_ms, 50)
end
# ============================================================================
# Crash Barrier Worker Checkout
# ============================================================================
@doc """
Timeout for checking out worker during crash barrier retry.
Used in `Snakepit.Pool` crash barrier retry logic.
Default: 5_000 ms (5 seconds)
"""
@spec crash_barrier_checkout_timeout() :: timeout()
def crash_barrier_checkout_timeout do
Application.get_env(:snakepit, :crash_barrier_checkout_timeout, 5_000)
end
# ============================================================================
# Pool Ready Timeout
# ============================================================================
@doc """
Default timeout for awaiting pool readiness.
Used in `Snakepit.Pool.await_ready/2`.
Default: 15_000 ms (15 seconds)
"""
@spec pool_await_ready_timeout() :: timeout()
def pool_await_ready_timeout do
Application.get_env(:snakepit, :pool_await_ready_timeout, 15_000)
end
# ============================================================================
# Worker Ready Notification
# ============================================================================
@doc """
Timeout for worker ready notification to pool.
Used in `Snakepit.GRPCWorker` when notifying pool of readiness.
Default: 30_000 ms (30 seconds)
"""
@spec worker_ready_timeout() :: timeout()
def worker_ready_timeout do
Application.get_env(:snakepit, :worker_ready_timeout, 30_000)
end
# ============================================================================
# gRPC Server Ready Wait
# ============================================================================
@doc """
Timeout for waiting for gRPC server to become ready.
Used in `Snakepit.GRPCWorker` during initialization.
Default: 30_000 ms (30 seconds)
"""
@spec grpc_server_ready_timeout() :: timeout()
def grpc_server_ready_timeout do
Application.get_env(:snakepit, :grpc_server_ready_timeout, 30_000)
end
# ============================================================================
# Session Affinity Cache
# ============================================================================
@doc """
TTL for session affinity cache entries in seconds.
Used in `Snakepit.Pool` for ETS affinity caching.
Default: 60 seconds (1 minute)
"""
@spec affinity_cache_ttl_seconds() :: pos_integer()
def affinity_cache_ttl_seconds do
Application.get_env(:snakepit, :affinity_cache_ttl_seconds, 60)
end
end
defmodule Snakepit.EnvDoctor do
@moduledoc """
Environment diagnostics for the Python bridge.
Provides both a Mix task integration (`mix snakepit.doctor`) and runtime
guardrails via `ensure_python!/1`.
"""
import Bitwise
@type check_result :: %{name: atom(), status: :ok | :warning | :error, message: String.t()}
@default_checks [
:python_exec,
:grpc_import,
:venv,
:venv_py313,
:grpc_server,
:adapter_imports,
:grpc_port
]
@runtime_checks [:python_exec, :grpc_import, :grpc_server, :adapter_imports]
alias Snakepit.Adapters.GRPCPython
alias Snakepit.PythonRuntime
@doc """
Run the full doctor suite. Returns `{:ok, results}` or `{:error, results}`.
"""
@spec run(Keyword.t()) :: {:ok, [check_result()]} | {:error, [check_result()]}
def run(opts \\ []) do
run_checks(@default_checks, opts)
end
@doc """
Ensure the Python runtime is ready. Raises if any critical check fails.
"""
@spec ensure_python!(Keyword.t()) :: :ok | no_return()
def ensure_python!(opts \\ []) do
case run_checks(@runtime_checks, opts) do
{:ok, _results} ->
:ok
{:error, results} ->
message =
results
|> Enum.filter(&(&1.status == :error))
|> Enum.map_join("\n", &"* #{&1.message}")
raise RuntimeError,
"Python environment is not ready:\n" <> message
end
end
defp run_checks(names, opts) do
state = build_state(opts)
{results, status} =
Enum.reduce(names, {[], :ok}, fn name, {acc, acc_status} ->
result = run_check(name, state)
new_status = if result.status == :error, do: :error, else: acc_status
{[result | acc], new_status}
end)
results = Enum.reverse(results)
case status do
:ok -> {:ok, results}
:error -> {:error, results}
end
end
defp build_state(opts) do
project_root =
opts[:project_root] ||
Application.get_env(:snakepit, :bootstrap_project_root) ||
File.cwd!()
python_runtime = PythonRuntime.config()
python_path = opts[:python_path] || GRPCPython.executable_path()
runner =
opts[:runner] ||
Application.get_env(:snakepit, :env_doctor_runner, Snakepit.Bootstrap.Runner.System)
require_python_313? =
Keyword.get(
opts,
:require_python_313?,
Application.get_env(:snakepit, :require_python_313?, false)
)
grpc_port = opts[:grpc_port] || Application.get_env(:snakepit, :grpc_port, 50_051)
%{
project_root: project_root,
python_path: python_path,
python_runtime: python_runtime,
runner: runner,
require_python_313?: require_python_313?,
grpc_port: grpc_port
}
end
defp run_check(:python_exec, state) do
case python_path_for_check(state) do
{:ok, path} -> ok(:python_exec, "Python executable found at #{path}")
{:error, message} -> error(:python_exec, message)
end
end
defp run_check(:grpc_import, state) do
case python_path_for_check(state) do
{:ok, _} ->
run_python(
state,
["-c", "import grpc"],
:grpc_import,
"Importing grpc failed. Run mix snakepit.setup (or make bootstrap)."
)
{:error, message} ->
error(:grpc_import, message)
end
end
defp run_check(:venv, %{project_root: root, python_runtime: runtime}) do
if PythonRuntime.managed?(runtime) do
warning(:venv, "Managed Python enabled; .venv check skipped.")
else
case File.dir?(Path.join(root, ".venv")) do
true ->
ok(:venv, ".venv present (Python 3.12)")
false ->
error(
:venv,
".venv missing. Run mix snakepit.setup (or make bootstrap) to create the default Python environment."
)
end
end
end
defp run_check(:venv_py313, %{
project_root: root,
require_python_313?: required?,
python_runtime: runtime
}) do
if PythonRuntime.managed?(runtime) do
warning(:venv_py313, "Managed Python enabled; .venv-py313 check skipped.")
else
path = Path.join(root, ".venv-py313")
cond do
File.dir?(path) ->
ok(:venv_py313, ".venv-py313 ready (Python 3.13)")
required? ->
error(
:venv_py313,
".venv-py313 missing. Run mix snakepit.setup (or make bootstrap) to enable free-threaded tests."
)
true ->
warning(
:venv_py313,
".venv-py313 missing. Thread-profile tests will be skipped until you run mix snakepit.setup (or make bootstrap)."
)
end
end
end
defp run_check(:grpc_server, state) do
case python_path_for_check(state) do
{:ok, _} ->
case grpc_server_root(state) do
nil ->
error(
:grpc_server,
"priv/python/grpc_server.py missing. Run mix snakepit.setup (or make bootstrap) to regenerate assets."
)
root ->
script = Path.join(root, "priv/python/grpc_server.py")
args = [script, "--health-check"] ++ default_adapter_args()
run_python(
state,
args,
:grpc_server,
"gRPC server health check failed. Regenerate stubs or reinstall deps.",
root
)
end
{:error, message} ->
error(:grpc_server, message)
end
end
defp run_check(:adapter_imports, state) do
case python_path_for_check(state) do
{:ok, python_path} ->
adapters = configured_adapter_paths()
case adapters do
[] ->
warning(
:adapter_imports,
"No adapter configured; default ShowcaseAdapter will be used."
)
_ ->
check_adapter_imports(state, python_path, adapters)
end
{:error, message} ->
error(:adapter_imports, message)
end
end
defp run_check(:grpc_port, %{grpc_port: port}) do
case :gen_tcp.listen(port, [:binary, active: false, reuseaddr: true]) do
{:ok, socket} ->
:gen_tcp.close(socket)
ok(:grpc_port, "Port #{port} available for the Elixir gRPC server")
{:error, :eaddrinuse} ->
error(
:grpc_port,
"Port #{port} is already in use. Stop the conflicting service or adjust :grpc_port."
)
{:error, reason} ->
warning(:grpc_port, "Unable to verify port #{port}: #{inspect(reason)}")
end
end
defp run_python(state, args, name, failure_message, root_override \\ nil) do
%{python_path: path, runner: runner, project_root: root} = state
root = root_override || root
case runner.cmd(path, args, cd: root, env: python_env(root)) do
:ok -> ok(name, "#{humanize(name)} check passed")
{:error, reason} -> error(name, "#{failure_message} (#{format_reason(reason)})")
end
end
defp grpc_server_root(state) do
candidates =
[
state.project_root,
snakepit_app_root()
]
|> Enum.reject(&is_nil/1)
Enum.find(candidates, fn root ->
File.exists?(Path.join(root, "priv/python/grpc_server.py"))
end)
end
defp snakepit_app_root do
case :code.priv_dir(:snakepit) do
{:error, _} -> nil
priv_dir -> priv_dir |> to_string() |> Path.dirname()
end
end
defp python_env(root) do
path_sep = path_separator()
path =
[
Path.join(root, "priv/python"),
snakepit_priv_python(),
snakebridge_priv_python(),
System.get_env("PYTHONPATH")
]
|> Enum.reject(&blank?/1)
|> Enum.uniq()
|> Enum.join(path_sep)
[{"PYTHONPATH", path} | python_log_level_env()]
end
defp snakepit_priv_python do
case :code.priv_dir(:snakepit) do
{:error, _} -> nil
priv_dir -> Path.join([to_string(priv_dir), "python"])
end
end
defp snakebridge_priv_python do
case :code.priv_dir(:snakebridge) do
{:error, _} -> nil
priv_dir -> Path.join([to_string(priv_dir), "python"])
end
end
defp blank?(value), do: value in [nil, ""]
defp python_log_level_env do
level = Application.get_env(:snakepit, :log_level, :error)
[{"SNAKEPIT_LOG_LEVEL", elixir_to_python_level(level)}]
end
defp elixir_to_python_level(:debug), do: "debug"
defp elixir_to_python_level(:info), do: "info"
defp elixir_to_python_level(:warning), do: "warning"
defp elixir_to_python_level(:error), do: "error"
defp elixir_to_python_level(:none), do: "none"
defp elixir_to_python_level(_), do: "error"
defp path_separator do
case :os.type() do
{:win32, _} -> ";"
_ -> ":"
end
end
defp configured_adapter_paths do
case Snakepit.Config.get_pool_configs() do
{:ok, pools} ->
pools
|> Enum.map(&extract_adapter_path(&1))
|> Enum.reject(&is_nil/1)
|> ensure_default_adapter()
{:error, _} ->
[]
|> ensure_default_adapter()
end
end
defp extract_adapter_path(%{adapter_args: adapter_args}) do
parse_adapter_from_args(adapter_args)
end
defp extract_adapter_path(_), do: nil
defp parse_adapter_from_args(args) when is_list(args) do
args
|> Enum.reduce({nil, false}, fn arg, {found, expecting} ->
cond do
expecting ->
{arg, false}
is_binary(arg) and String.starts_with?(arg, "--adapter=") ->
{String.replace_prefix(arg, "--adapter=", ""), false}
arg == "--adapter" ->
{found, true}
true ->
{found, false}
end
end)
|> elem(0)
end
defp parse_adapter_from_args(_), do: nil
defp ensure_default_adapter(adapters) do
case adapters do
[] ->
default = parse_adapter_from_args(GRPCPython.script_args() || [])
if default do
[default]
else
[]
end
_ ->
Enum.uniq(adapters)
end
end
defp check_adapter_imports(state, python_path, adapters) do
root = grpc_server_root(state) || state.project_root
script = Path.join(root, "priv/python/grpc_server.py")
env = python_env(root)
{ok_adapters, failed_adapters} =
Enum.reduce(adapters, {[], []}, fn adapter, {oks, fails} ->
args = [script, "--health-check", "--adapter", adapter]
case state.runner.cmd(python_path, args, cd: root, env: env) do
:ok -> {[adapter | oks], fails}
{:error, reason} -> {oks, [{adapter, reason} | fails]}
end
end)
case failed_adapters do
[] ->
ok(
:adapter_imports,
"Adapter import checks passed (#{length(ok_adapters)} adapters)"
)
_ ->
failures =
Enum.map_join(failed_adapters, ", ", fn {adapter, reason} ->
"#{adapter} (#{format_reason(reason)})"
end)
error(
:adapter_imports,
"Adapter import checks failed: #{failures}"
)
end
end
defp default_adapter_args do
GRPCPython.script_args()
end
defp ok(name, message), do: %{name: name, status: :ok, message: message}
defp warning(name, message), do: %{name: name, status: :warning, message: message}
defp error(name, message), do: %{name: name, status: :error, message: message}
defp format_reason({:command_failed, command, status}),
do: "#{command} exited with #{status}"
defp format_reason(reason), do: inspect(reason)
defp humanize(name) do
name
|> Atom.to_string()
|> String.replace("_", " ")
|> String.capitalize()
end
defp executable?(path) do
case File.stat(path) do
{:ok, %File.Stat{mode: mode}} -> band(mode, 0o111) != 0
_ -> false
end
end
defp python_path_for_check(%{python_path: nil, python_runtime: runtime}) do
case PythonRuntime.missing_reason(runtime) do
{:error, message} ->
{:error, message}
_ ->
{:error,
"Python not configured. Run mix snakepit.setup (or make bootstrap) or set SNAKEPIT_PYTHON=/path/to/python."}
end
end
defp python_path_for_check(%{python_path: path}) do
cond do
not File.exists?(path) ->
{:error,
"Configured interpreter not found at #{path}. Run mix snakepit.setup (or make bootstrap)."}
not executable?(path) ->
{:error,
"Interpreter at #{path} is not executable. Fix permissions or recreate the venv."}
true ->
{:ok, path}
end
end
end
defmodule Snakepit.Error do
@moduledoc """
Structured error type for Snakepit operations.
Provides detailed context for debugging cross-language and distributed system issues.
Python exceptions translated from the gRPC bridge are returned as
`Snakepit.Error.*` exception structs (see `Snakepit.Error.PythonException`).
`Snakepit.Error` remains the structured error type for Snakepit runtime failures.
## Error Categories
- `:worker` - Worker process errors (not found, crashed, etc.)
- `:timeout` - Operation timed out
- `:python_error` - Exception from Python code
- `:grpc_error` - gRPC communication error
- `:validation` - Input validation error
- `:pool` - Pool management error
## Examples
# Create a worker error
error = Snakepit.Error.worker_error("Worker not found", %{worker_id: "w1"})
# Create a Python exception error
error = Snakepit.Error.python_error(
"ValueError",
"Invalid input",
traceback_string,
%{function: "process_data"}
)
# Pattern match in your code
case Snakepit.execute("command", %{}) do
{:ok, result} -> result
{:error, %Snakepit.Error{category: :timeout}} -> retry()
{:error, %Snakepit.Error{category: :python_error} = error} ->
Snakepit.Logger.error("Python error: \#{error.message}")
Snakepit.Logger.debug("Traceback: \#{error.python_traceback}")
{:error, error} -> {:error, error}
end
"""
@type category :: :worker | :timeout | :python_error | :grpc_error | :validation | :pool
@type t :: %__MODULE__{
category: category(),
message: String.t(),
details: map(),
python_traceback: String.t() | nil,
grpc_status: atom() | nil
}
defstruct [:category, :message, :details, :python_traceback, :grpc_status]
@doc """
Creates a worker-related error.
## Examples
iex> Snakepit.Error.worker_error("Worker crashed")
%Snakepit.Error{category: :worker, message: "Worker crashed", details: %{}}
iex> Snakepit.Error.worker_error("Worker not found", %{worker_id: "w1"})
%Snakepit.Error{category: :worker, message: "Worker not found", details: %{worker_id: "w1"}}
"""
@spec worker_error(String.t(), map()) :: t()
def worker_error(message, details \\ %{}) do
%__MODULE__{
category: :worker,
message: message,
details: details
}
end
@doc """
Creates a timeout error.
## Examples
iex> Snakepit.Error.timeout_error("Request timed out", %{timeout_ms: 5000})
%Snakepit.Error{category: :timeout, message: "Request timed out", details: %{timeout_ms: 5000}}
"""
@spec timeout_error(String.t(), map()) :: t()
def timeout_error(message, details \\ %{}) do
%__MODULE__{
category: :timeout,
message: message,
details: details
}
end
@doc """
Creates a Python exception error with traceback.
## Examples
iex> Snakepit.Error.python_error("ValueError", "Invalid input", "Traceback...")
%Snakepit.Error{
category: :python_error,
message: "ValueError: Invalid input",
python_traceback: "Traceback...",
details: %{exception_type: "ValueError"}
}
"""
@spec python_error(String.t(), String.t(), String.t(), map()) :: t()
def python_error(exception_type, message, traceback, details \\ %{}) do
%__MODULE__{
category: :python_error,
message: "#{exception_type}: #{message}",
details: Map.put(details, :exception_type, exception_type),
python_traceback: traceback
}
end
@doc """
Creates a gRPC communication error.
## Examples
iex> Snakepit.Error.grpc_error(:unavailable, "Service unavailable")
%Snakepit.Error{
category: :grpc_error,
message: "Service unavailable",
grpc_status: :unavailable
}
"""
@spec grpc_error(atom(), String.t(), map()) :: t()
def grpc_error(status, message, details \\ %{}) do
%__MODULE__{
category: :grpc_error,
message: message,
grpc_status: status,
details: details
}
end
@doc """
Creates a pool management error.
## Examples
iex> Snakepit.Error.pool_error("Pool not found", %{pool_name: :test})
%Snakepit.Error{category: :pool, message: "Pool not found", details: %{pool_name: :test}}
"""
@spec pool_error(String.t(), map()) :: t()
def pool_error(message, details \\ %{}) do
%__MODULE__{
category: :pool,
message: message,
details: details
}
end
@doc """
Creates a validation error.
## Examples
iex> Snakepit.Error.validation_error("Invalid field", %{field: "user_id"})
%Snakepit.Error{category: :validation, message: "Invalid field", details: %{field: "user_id"}}
"""
@spec validation_error(String.t(), map()) :: t()
def validation_error(message, details \\ %{}) do
%__MODULE__{
category: :validation,
message: message,
details: details
}
end
defimpl String.Chars do
def to_string(%Snakepit.Error{} = error) do
base = "[#{error.category}] #{error.message}"
details_str =
if map_size(error.details) > 0 do
"\nDetails: #{inspect(error.details)}"
else
""
end
traceback_str =
if error.python_traceback do
"\n\nPython Traceback:\n#{error.python_traceback}"
else
""
end
grpc_str =
if error.grpc_status do
"\ngRPC Status: #{error.grpc_status}"
else
""
end
base <> details_str <> grpc_str <> traceback_str
end
end
end
defmodule Snakepit.Executor do
@moduledoc """
Execution helpers with retry, circuit breaker, and timeout support.
Provides various execution strategies for running operations
with fault tolerance.
## Usage
# Simple execution with retry
result = Executor.execute_with_retry(
fn -> risky_operation() end,
max_attempts: 3,
backoff_ms: [100, 200, 400]
)
# With circuit breaker
result = Executor.execute_with_circuit_breaker(cb, fn ->
external_call()
end)
# With timeout
result = Executor.execute_with_timeout(
fn -> slow_operation() end,
timeout_ms: 5000
)
"""
alias Snakepit.{CircuitBreaker, Defaults, RetryPolicy}
@doc """
Executes a function directly.
"""
@spec execute((-> any()), keyword()) :: any()
def execute(fun, _opts \\ []) when is_function(fun, 0) do
fun.()
end
@doc """
Executes a function with retry on transient failures.
## Options
- `:max_attempts` - Maximum attempts (default: 3)
- `:backoff_ms` - List of backoff delays (default: [100, 200, 400])
- `:retriable_errors` - Errors to retry (default: [:timeout, :unavailable])
- `:jitter` - Add random jitter (default: false)
"""
@spec execute_with_retry((-> any()), keyword()) :: any()
def execute_with_retry(fun, opts \\ []) when is_function(fun, 0) do
policy = RetryPolicy.new(opts)
do_retry(fun, policy, 1)
end
defp do_retry(fun, policy, attempt) do
result = fun.()
cond do
match?({:ok, _}, result) or match?(:ok, result) ->
emit_retry_success(policy, attempt)
result
RetryPolicy.retry_for_error?(policy, result) and
RetryPolicy.should_retry?(policy, attempt) ->
delay = RetryPolicy.backoff_for_attempt(policy, attempt)
emit_retry_attempt(policy, attempt, delay)
Process.sleep(delay)
do_retry(fun, policy, attempt + 1)
RetryPolicy.should_retry?(policy, attempt) == false ->
emit_retry_exhausted(policy, attempt, result)
result
true ->
# Non-retriable error
result
end
end
@doc """
Executes a function through a circuit breaker.
"""
@spec execute_with_circuit_breaker(GenServer.server(), (-> any()), keyword()) :: any()
def execute_with_circuit_breaker(circuit_breaker, fun, _opts \\ [])
when is_function(fun, 0) do
CircuitBreaker.call(circuit_breaker, fun)
end
@doc """
Executes a function with a timeout.
Returns `{:error, :timeout}` if the function doesn't complete in time.
## Options
- `:timeout_ms` - Timeout in milliseconds (required)
"""
@spec execute_with_timeout((-> any()), keyword()) :: any()
def execute_with_timeout(fun, opts) when is_function(fun, 0) do
timeout_ms = Keyword.fetch!(opts, :timeout_ms)
task = Task.async(fn -> fun.() end)
case Task.yield(task, timeout_ms) do
{:ok, result} ->
result
nil ->
Task.shutdown(task, :brutal_kill)
{:error, :timeout}
end
end
@doc """
Executes a function asynchronously.
Returns a Task that can be awaited.
"""
@spec execute_async((-> any()), keyword()) :: Task.t()
def execute_async(fun, _opts \\ []) when is_function(fun, 0) do
Task.async(fun)
end
@doc """
Executes multiple functions in parallel.
Returns results in the same order as the input functions.
## Options
- `:timeout_ms` - Timeout for all operations (default: 30000)
- `:max_concurrency` - Maximum concurrent operations (default: unlimited)
"""
@spec execute_batch([(-> any())], keyword()) :: [any()]
def execute_batch(functions, opts \\ []) when is_list(functions) do
timeout_ms = Keyword.get(opts, :timeout_ms, Defaults.executor_batch_timeout())
max_concurrency = Keyword.get(opts, :max_concurrency, length(functions))
functions
|> Task.async_stream(
fn fun -> fun.() end,
max_concurrency: max_concurrency,
timeout: timeout_ms
)
|> Enum.map(fn
{:ok, result} -> result
{:exit, reason} -> {:error, {:exit, reason}}
end)
end
@doc """
Executes with retry and circuit breaker.
Combines retry logic with circuit breaker protection.
"""
@spec execute_with_protection(GenServer.server(), (-> any()), keyword()) :: any()
def execute_with_protection(circuit_breaker, fun, opts \\ []) do
execute_with_retry(
fn -> execute_with_circuit_breaker(circuit_breaker, fun) end,
opts
)
end
# Telemetry helpers
defp emit_retry_attempt(_policy, attempt, delay) do
:telemetry.execute(
[:snakepit, :retry, :attempt],
%{attempt: attempt, delay_ms: delay},
%{pool: nil, operation: nil}
)
end
defp emit_retry_success(_policy, attempts) do
:telemetry.execute(
[:snakepit, :retry, :success],
%{attempts: attempts, total_duration: 0},
%{pool: nil}
)
end
defp emit_retry_exhausted(_policy, attempts, error) do
:telemetry.execute(
[:snakepit, :retry, :exhausted],
%{attempts: attempts, total_duration: 0},
%{pool: nil, last_error: error}
)
end
end
defmodule Snakepit.GRPCWorker do
@moduledoc """
A GenServer that manages gRPC connections to external processes.
This worker can handle both traditional request/response and streaming operations
via gRPC instead of stdin/stdout communication.
## Features
- Automatic gRPC connection management
- Health check monitoring
- Streaming support with callback-based API
- Session affinity for stateful operations
- Graceful fallback to traditional workers if gRPC unavailable
## Usage
# Start a gRPC worker
{:ok, worker} = Snakepit.GRPCWorker.start_link(adapter: Snakepit.Adapters.GRPCPython)
# Simple execution
{:ok, result} = Snakepit.GRPCWorker.execute(worker, "ping", %{})
# Streaming execution
Snakepit.GRPCWorker.execute_stream(worker, "batch_inference", %{
batch_items: ["img1.jpg", "img2.jpg"]
}, fn chunk ->
handle_chunk(chunk)
end)
"""
use GenServer
require Logger
alias Snakepit.Adapters.GRPCPython
alias Snakepit.Defaults
alias Snakepit.Error
alias Snakepit.GRPC.Client
alias Snakepit.Logger, as: SLog
alias Snakepit.Logger.Redaction
alias Snakepit.Pool.ProcessRegistry
alias Snakepit.Pool.Registry, as: PoolRegistry
alias Snakepit.Telemetry.Correlation
alias Snakepit.Telemetry.GrpcStream
alias Snakepit.Worker.LifecycleManager
require OpenTelemetry.Tracer, as: Tracer
def child_spec(opts) when is_list(opts) do
%{
id: Keyword.get(opts, :id, __MODULE__),
start: {__MODULE__, :start_link, [opts]},
restart: :transient,
type: :worker,
# Must give worker time for graceful Python shutdown.
# Derived from :graceful_shutdown_timeout_ms + margin.
shutdown: supervisor_shutdown_timeout()
}
end
@type worker_state :: %{
adapter: module(),
connection: map() | nil,
port: integer(),
process_pid: integer() | nil,
pgid: integer() | nil,
process_group?: boolean(),
server_port: port() | nil,
id: String.t(),
pool_name: atom() | pid(),
health_check_ref: reference() | nil,
heartbeat_monitor: pid() | nil,
heartbeat_config: map(),
ready_file: String.t(),
stats: map(),
session_id: String.t(),
worker_config: map(),
shutting_down: boolean()
}
# Base heartbeat defaults - actual values are retrieved via Defaults module
# to allow runtime configuration. These are the compile-time fallbacks.
@base_heartbeat_defaults_template %{
enabled: true,
ping_fun: nil,
test_pid: nil,
dependent: true
}
defp base_heartbeat_defaults do
Map.merge(@base_heartbeat_defaults_template, %{
ping_interval_ms: Defaults.heartbeat_ping_interval_ms(),
timeout_ms: Defaults.heartbeat_timeout_ms(),
max_missed_heartbeats: Defaults.heartbeat_max_missed(),
initial_delay_ms: Defaults.heartbeat_initial_delay_ms()
})
end
@heartbeat_known_keys [
:enabled,
:ping_interval_ms,
:timeout_ms,
:max_missed_heartbeats,
:ping_fun,
:test_pid,
:initial_delay_ms,
:dependent
]
@heartbeat_known_key_strings Enum.map(@heartbeat_known_keys, &Atom.to_string/1)
@log_category :grpc
# Client API
@doc """
Start a gRPC worker with the given adapter.
"""
def start_link(opts) do
worker_id = Keyword.get(opts, :id)
pool_name = Keyword.get(opts, :pool_name, Snakepit.Pool)
pool_identifier = resolve_pool_identifier(opts, pool_name)
metadata =
%{worker_module: __MODULE__, pool_name: pool_name}
|> maybe_put_pool_identifier(pool_identifier)
opts_with_metadata =
opts
|> Keyword.put(:registry_metadata, metadata)
|> maybe_put_pool_identifier_opt(pool_identifier)
name = build_worker_name(worker_id)
GenServer.start_link(__MODULE__, opts_with_metadata, name: name)
end
@doc """
Execute a command and return the result.
"""
# Header for default values
def execute(worker, command, args, timeout \\ nil)
def execute(worker, command, args, nil) do
execute(worker, command, args, Defaults.grpc_worker_execute_timeout())
end
def execute(worker_id, command, args, timeout) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
GenServer.call(pid, {:execute, command, args, timeout}, timeout + 1_000)
{:error, _} ->
{:error,
Error.worker_error("Worker not found", %{worker_id: worker_id, command: command})}
end
end
def execute(worker_pid, command, args, timeout) when is_pid(worker_pid) do
GenServer.call(worker_pid, {:execute, command, args, timeout}, timeout + 1_000)
end
@doc """
Execute a streaming command with callback.
"""
def execute_stream(worker, command, args, callback_fn, timeout \\ nil)
def execute_stream(worker, command, args, callback_fn, nil) do
execute_stream(worker, command, args, callback_fn, Defaults.grpc_worker_stream_timeout())
end
def execute_stream(worker_id, command, args, callback_fn, timeout) when is_binary(worker_id) do
case PoolRegistry.get_worker_pid(worker_id) do
{:ok, pid} ->
GenServer.call(
pid,
{:execute_stream, command, args, callback_fn, timeout},
timeout + 1_000
)
{:error, _} ->
{:error,
Error.worker_error("Worker not found", %{worker_id: worker_id, command: command})}
end
end
def execute_stream(worker_pid, command, args, callback_fn, timeout) when is_pid(worker_pid) do
GenServer.call(
worker_pid,
{:execute_stream, command, args, callback_fn, timeout},
timeout + 1_000
)
end
@doc """
Execute a command in a specific session.
"""
def execute_in_session(worker, session_id, command, args, timeout \\ nil)
def execute_in_session(worker, session_id, command, args, nil) do
execute_in_session(worker, session_id, command, args, Defaults.grpc_worker_execute_timeout())
end
def execute_in_session(worker, session_id, command, args, timeout) do
GenServer.call(
worker,
{:execute_session, session_id, command, args, timeout},
timeout + 1_000
)
end
@doc """
Get worker health and statistics.
"""
def get_health(worker) do
GenServer.call(worker, :get_health)
end
@doc """
Get worker information and capabilities.
"""
def get_info(worker) do
GenServer.call(worker, :get_info)
end
@doc """
Get the gRPC channel for direct client usage.
"""
def get_channel(worker) do
GenServer.call(worker, :get_channel)
end
@doc """
Get the session ID for this worker.
"""
def get_session_id(worker) do
GenServer.call(worker, :get_session_id)
end
defp resolve_pool_identifier(opts, pool_name) do
case Keyword.get(opts, :pool_identifier) do
identifier when is_atom(identifier) ->
identifier
identifier when is_binary(identifier) ->
string_to_existing_atom_safe(identifier)
_ ->
infer_pool_identifier(pool_name)
end
end
defp string_to_existing_atom_safe(identifier) do
String.to_existing_atom(identifier)
rescue
ArgumentError -> nil
end
defp infer_pool_identifier(pool_name) when is_atom(pool_name), do: pool_name
defp infer_pool_identifier(pool_name) when is_pid(pool_name) do
case Process.info(pool_name, :registered_name) do
{:registered_name, name} when is_atom(name) -> name
_ -> nil
end
end
defp infer_pool_identifier(_), do: nil
defp build_worker_name(nil), do: nil
defp build_worker_name(worker_id) do
{:via, Registry, {Snakepit.Pool.Registry, worker_id}}
end
defp maybe_put_pool_identifier(metadata, nil), do: metadata
defp maybe_put_pool_identifier(metadata, identifier),
do: Map.put(metadata, :pool_identifier, identifier)
defp maybe_put_pool_identifier_opt(opts, nil), do: opts
defp maybe_put_pool_identifier_opt(opts, identifier),
do: Keyword.put(opts, :pool_identifier, identifier)
defp ensure_registry_metadata(metadata, pool_name, pool_identifier) do
metadata
|> Map.put(:worker_module, __MODULE__)
|> Map.put(:pool_name, pool_name)
|> maybe_put_pool_identifier(pool_identifier)
end
defp maybe_attach_registry_metadata(worker_id, metadata) when is_binary(worker_id) do
case PoolRegistry.put_metadata(worker_id, metadata) do
:ok ->
:ok
{:error, :not_registered} ->
SLog.debug(
@log_category,
"Pool.Registry missing entry for #{worker_id} while attaching metadata"
)
:ok
end
rescue
_ -> :ok
end
defp maybe_attach_registry_metadata(_worker_id, _metadata), do: :ok
defp normalize_worker_config(config, pool_name, adapter_module, pool_identifier) do
config
|> Map.put(:worker_module, __MODULE__)
|> Map.put_new(:adapter_module, adapter_module)
|> Map.put(:pool_name, pool_name)
|> maybe_put_pool_identifier(pool_identifier)
end
defp current_process_memory_bytes do
case Process.info(self(), :memory) do
{:memory, bytes} when is_integer(bytes) and bytes >= 0 -> bytes
_ -> 0
end
end
# Server callbacks
@impl true
def init(opts) do
# CRITICAL: Trap exits so terminate/2 is called on shutdown
# Without this, the GenServer is brutally killed and Python processes are orphaned!
Process.flag(:trap_exit, true)
adapter = Keyword.fetch!(opts, :adapter)
worker_id = Keyword.fetch!(opts, :id)
pool_name = Keyword.get(opts, :pool_name, Snakepit.Pool)
pool_identifier = Keyword.get(opts, :pool_identifier)
Logger.metadata(worker_id: worker_id, pool_name: pool_name, adapter: adapter)
metadata =
opts
|> Keyword.get(:registry_metadata, %{})
|> ensure_registry_metadata(pool_name, pool_identifier)
maybe_attach_registry_metadata(worker_id, metadata)
case ProcessRegistry.reserve_worker(worker_id) do
:ok ->
init_worker(opts, adapter, worker_id, pool_name, pool_identifier)
{:error, reason} ->
SLog.error(
@log_category,
"Failed to reserve worker slot for #{worker_id}: #{inspect(reason)}"
)
{:stop, {:reservation_failed, reason}}
end
end
defp init_worker(opts, adapter, worker_id, pool_name, pool_identifier) do
SLog.debug(@log_category, "Reserved worker slot for #{worker_id}")
session_id = generate_session_id()
Logger.metadata(session_id: session_id)
elixir_address = build_elixir_address()
port = adapter.get_port()
worker_config =
opts
|> Keyword.get(:worker_config, %{})
|> normalize_worker_config(pool_name, adapter, pool_identifier)
heartbeat_config =
worker_config
|> get_worker_config_section(:heartbeat)
|> normalize_heartbeat_config()
spawn_config =
build_spawn_config(
adapter,
worker_config,
heartbeat_config,
port,
elixir_address,
worker_id
)
server_port = spawn_grpc_server(spawn_config)
process_pid = extract_and_log_pid(server_port, port)
{pgid, process_group?} = resolve_process_group(process_pid, spawn_config)
register_worker_pid(worker_id, process_pid, pgid, process_group?)
state_params = %{
worker_id: worker_id,
pool_name: pool_name,
adapter: adapter,
port: port,
server_port: server_port,
process_pid: process_pid,
pgid: pgid,
process_group?: process_group?,
session_id: session_id,
worker_config: worker_config,
heartbeat_config: heartbeat_config,
ready_file: spawn_config.ready_file
}
state = build_initial_state(state_params)
{:ok, state, {:continue, :connect_and_wait}}
end
defp generate_session_id do
"session_#{:erlang.unique_integer([:positive, :monotonic])}_#{:erlang.system_time(:microsecond)}"
end
defp build_elixir_address do
elixir_grpc_host = Application.get_env(:snakepit, :grpc_host, "localhost")
elixir_grpc_port = Application.get_env(:snakepit, :grpc_port, 50_051)
"#{elixir_grpc_host}:#{elixir_grpc_port}"
end
defp build_spawn_config(
adapter,
worker_config,
heartbeat_config,
port,
elixir_address,
worker_id
) do
python_executable = adapter.executable_path()
adapter_args = resolve_adapter_args(adapter, worker_config)
adapter_env = worker_config |> Map.get(:adapter_env, []) |> merge_with_default_adapter_env()
heartbeat_env_json = encode_heartbeat_env(heartbeat_config)
ready_file = build_ready_file(worker_id)
script_path = determine_script_path(adapter, adapter_args)
args = build_spawn_args(adapter_args, port, elixir_address)
SLog.info(
@log_category,
"Starting gRPC server: #{python_executable} #{script_path || ""} #{Enum.join(args, " ")}"
)
%{
executable: python_executable,
script_path: script_path,
args: args,
process_group?: process_group_spawn?(),
adapter_env: adapter_env,
heartbeat_env_json: heartbeat_env_json,
ready_file: ready_file
}
end
defp resolve_adapter_args(adapter, worker_config) do
worker_adapter_args = Map.get(worker_config, :adapter_args, [])
if worker_adapter_args == [] do
adapter.script_args() || []
else
worker_adapter_args
end
end
defp determine_script_path(adapter, adapter_args) do
if Enum.any?(adapter_args, fn arg ->
is_binary(arg) and String.contains?(arg, "--max-workers")
end) do
app_dir = Application.app_dir(:snakepit)
Path.join([app_dir, "priv", "python", "grpc_server_threaded.py"])
else
adapter.script_path()
end
end
defp build_spawn_args(adapter_args, port, elixir_address) do
adapter_args
|> maybe_add_arg("--port", to_string(port))
|> maybe_add_arg("--elixir-address", elixir_address)
|> add_run_id_arg()
end
defp build_ready_file(worker_id) do
base_dir = System.tmp_dir() || File.cwd!()
safe_worker_id =
worker_id
|> to_string()
|> String.replace(~r/[^a-zA-Z0-9_.-]/, "_")
unique = :erlang.unique_integer([:positive, :monotonic])
Path.join(base_dir, "snakepit_ready_#{safe_worker_id}_#{unique}")
end
defp maybe_add_arg(args, flag, value) do
if Enum.any?(args, &String.contains?(&1, flag)) do
args
else
args ++ [flag, value]
end
end
defp add_run_id_arg(args) do
run_id = ProcessRegistry.get_beam_run_id()
args ++ ["--snakepit-run-id", run_id]
end
defp spawn_grpc_server(%{
executable: executable,
script_path: script_path,
args: args,
adapter_env: adapter_env,
heartbeat_env_json: heartbeat_env_json,
ready_file: ready_file
}) do
{spawn_args, port_opts} = build_port_config(script_path, args)
port_opts = apply_env_to_port_opts(port_opts, adapter_env, heartbeat_env_json, ready_file)
server_port = Port.open({:spawn_executable, executable}, [{:args, spawn_args} | port_opts])
Port.monitor(server_port)
server_port
end
defp build_port_config(script_path, args) do
case script_path do
path when is_binary(path) and byte_size(path) > 0 ->
{
[path | args],
[:binary, :exit_status, :use_stdio, :stderr_to_stdout, {:cd, Path.dirname(path)}]
}
_ ->
{args, [:binary, :exit_status, :use_stdio, :stderr_to_stdout]}
end
end
defp process_group_spawn? do
Application.get_env(:snakepit, :process_group_kill, true) and
Snakepit.ProcessKiller.process_group_supported?()
end
defp resolve_process_group(process_pid, %{process_group?: true})
when is_integer(process_pid) do
with {:ok, pgid} <- Snakepit.ProcessKiller.get_process_group_id(process_pid),
true <- pgid == process_pid do
{pgid, true}
else
_ -> {nil, false}
end
end
defp resolve_process_group(_process_pid, _spawn_config), do: {nil, false}
defp apply_env_to_port_opts(port_opts, adapter_env, heartbeat_env_json, ready_file) do
env_entries =
adapter_env
|> maybe_put_heartbeat_env(heartbeat_env_json)
|> maybe_put_ready_file_env(ready_file)
|> maybe_put_log_level_env()
if env_entries != [] do
env_tuples = Enum.map(env_entries, &to_env_tuple/1)
port_opts ++ [{:env, env_tuples}]
else
port_opts
end
end
defp extract_and_log_pid(server_port, port) do
case Port.info(server_port, :os_pid) do
{:os_pid, pid} ->
SLog.info(@log_category, "Started gRPC server process, will listen on TCP port #{port}")
pid
error ->
SLog.error(@log_category, "Failed to get gRPC server process PID: #{inspect(error)}")
nil
end
end
defp register_worker_pid(_worker_id, nil, _pgid, _process_group?), do: :ok
defp register_worker_pid(worker_id, process_pid, pgid, process_group?) do
case ProcessRegistry.activate_worker(worker_id, self(), process_pid, "grpc_worker",
pgid: pgid,
process_group?: process_group?
) do
:ok ->
SLog.debug(
@log_category,
"Registered Python PID #{process_pid} for worker #{worker_id} in ProcessRegistry"
)
{:error, reason} ->
SLog.error(
@log_category,
"Failed to register Python PID #{process_pid} for worker #{worker_id}: #{inspect(reason)}"
)
end
end
defp build_initial_state(%{
worker_id: worker_id,
pool_name: pool_name,
adapter: adapter,
port: port,
server_port: server_port,
process_pid: process_pid,
pgid: pgid,
process_group?: process_group?,
session_id: session_id,
worker_config: worker_config,
heartbeat_config: heartbeat_config,
ready_file: ready_file
}) do
%{
id: worker_id,
pool_name: pool_name,
adapter: adapter,
port: port,
server_port: server_port,
process_pid: process_pid,
pgid: pgid,
process_group?: process_group?,
session_id: session_id,
requested_port: port,
worker_config: worker_config,
heartbeat_config: heartbeat_config,
ready_file: ready_file,
heartbeat_monitor: nil,
connection: nil,
health_check_ref: nil,
python_output_buffer: "",
# Track whether we initiated shutdown (to distinguish expected vs unexpected exits)
shutting_down: false,
stats: %{
requests: 0,
errors: 0,
start_time: System.monotonic_time()
}
}
end
defp resolve_pool_pid(pool_name) when is_atom(pool_name), do: Process.whereis(pool_name)
defp resolve_pool_pid(pool_name), do: pool_name
defp verify_pool_alive(nil, worker_id), do: {:error, {:pool_dead, worker_id}}
defp verify_pool_alive(pool_pid, worker_id) do
if Process.alive?(pool_pid) do
:ok
else
{:error, {:pool_dead, worker_id}}
end
end
defp complete_worker_initialization(state, connection, actual_port) do
health_ref = schedule_health_check()
LifecycleManager.track_worker(state.pool_name, state.id, self(), state.worker_config)
SLog.info(
@log_category,
"✅ gRPC worker #{state.id} initialization complete and acknowledged."
)
maybe_initialize_session(connection, state.session_id)
register_telemetry_stream(connection, state)
emit_worker_spawned_telemetry(state, actual_port)
new_state =
state
|> Map.put(:connection, connection)
|> Map.put(:port, actual_port)
|> Map.put(:health_check_ref, health_ref)
|> maybe_start_heartbeat_monitor()
{:noreply, new_state}
end
defp emit_worker_spawned_telemetry(state, actual_port) do
start_time = Map.get(state.stats, :start_time, System.monotonic_time())
:telemetry.execute(
[:snakepit, :pool, :worker, :spawned],
%{
duration: System.monotonic_time() - start_time,
system_time: System.system_time()
},
%{
node: node(),
pool_name: state.pool_name,
worker_id: state.id,
worker_pid: self(),
python_port: actual_port,
python_pid: state.process_pid,
mode: :process
}
)
end
@impl true
def handle_continue(:connect_and_wait, state) do
with {:ok, actual_port} <-
wait_for_server_ready(
state.server_port,
state.ready_file,
Defaults.grpc_server_ready_timeout()
),
{:ok, connection} <-
wrap_grpc_connection_result(state.adapter.init_grpc_connection(actual_port)),
pool_pid <- resolve_pool_pid(state.pool_name),
:ok <- verify_pool_alive(pool_pid, state.id),
:ok <- notify_pool_ready(pool_pid, state.id) do
complete_worker_initialization(state, connection, actual_port)
else
{:error, :shutdown} ->
SLog.debug(@log_category, "gRPC server exited during startup (shutdown)")
{:stop, :shutdown, state}
{:error, {:exit_status, status}} when status in [137] ->
SLog.error(@log_category, "gRPC server exited during startup with status #{status}")
{:stop, {:grpc_server_failed, {:exit_status, status}}, state}
{:error, {:pool_dead, _}} ->
SLog.debug(
@log_category,
"Worker #{state.id} finished starting but Pool is shut down. Stopping gracefully."
)
{:stop, :normal, state}
{:error, {:pool_handshake_failed, reason}} ->
SLog.debug(
@log_category,
"Pool handshake failed for worker #{state.id}: #{inspect(reason)}"
)
{:stop, :shutdown, state}
{:error, {:grpc_connection_failed, reason}} ->
SLog.error(@log_category, "Failed to connect to gRPC server: #{reason}")
{:stop, {:grpc_connection_failed, reason}, state}
{:error, reason} ->
SLog.error(@log_category, "Failed to start gRPC server: #{inspect(reason)}")
{:stop, {:grpc_server_failed, reason}, state}
end
end
defp wrap_grpc_connection_result({:ok, connection}), do: {:ok, connection}
defp wrap_grpc_connection_result({:error, reason}),
do: {:error, {:grpc_connection_failed, reason}}
@impl true
def handle_call({:execute, command, args, timeout}, _from, state) do
args_with_corr = ensure_correlation(args)
case instrument_execute(
:execute,
state,
command,
args_with_corr,
timeout,
fn instrumented_args ->
state.adapter.grpc_execute(
state.connection,
state.session_id,
command,
instrumented_args,
timeout
)
end
) do
{:ok, result} ->
new_state = update_stats(state, :success)
{:reply, {:ok, result}, new_state}
{:error, reason} ->
new_state = update_stats(state, :error)
{:reply, {:error, reason}, new_state}
end
end
@impl true
def handle_call({:execute_stream, command, args, callback_fn, timeout}, _from, state) do
SLog.debug(
@log_category,
"[GRPCWorker] execute_stream #{command} with args #{Redaction.describe(args)}"
)
args_with_corr = ensure_correlation(args)
result =
state.adapter.grpc_execute_stream(
state.connection,
state.session_id,
command,
args_with_corr,
callback_fn,
timeout
)
SLog.debug(@log_category, "[GRPCWorker] execute_stream result: #{Redaction.describe(result)}")
new_state =
case result do
:ok -> update_stats(state, :success)
{:error, _reason} -> update_stats(state, :error)
end
{:reply, result, new_state}
end
@impl true
def handle_call(:get_port, _from, state) do
{:reply, {:ok, state.port}, state}
end
@impl true
def handle_call(:get_port_metadata, _from, state) do
info = %{
current_port: state.port,
requested_port: Map.get(state, :requested_port, state.port)
}
{:reply, {:ok, info}, state}
end
@impl true
def handle_call({:execute_session, session_id, command, args, timeout}, _from, state) do
session_args =
args
|> Map.put(:session_id, session_id)
|> ensure_correlation()
case instrument_execute(
:execute_session,
state,
command,
session_args,
timeout,
fn instrumented_args ->
state.adapter.grpc_execute(
state.connection,
state.session_id,
command,
instrumented_args,
timeout
)
end
) do
{:ok, result} ->
new_state = update_stats(state, :success)
{:reply, {:ok, result}, new_state}
{:error, reason} ->
new_state = update_stats(state, :error)
{:reply, {:error, reason}, new_state}
end
end
@impl true
def handle_call(:get_health, _from, state) do
# Make gRPC health check call
health_result = make_health_check(state)
{:reply, health_result, state}
end
@impl true
def handle_call(:get_info, _from, state) do
# Make gRPC info call
info_result = make_info_call(state)
{:reply, info_result, state}
end
@impl true
def handle_call(:get_channel, _from, state) do
if state.connection do
{:reply, {:ok, state.connection.channel}, state}
else
{:reply,
{:error,
Error.grpc_error(:not_connected, "Not connected to gRPC server", %{worker_id: state.id})},
state}
end
end
@impl true
def handle_call(:get_session_id, _from, state) do
{:reply, {:ok, state.session_id}, state}
end
def handle_call(:get_memory_usage, _from, state) do
{:reply, {:ok, current_process_memory_bytes()}, state}
end
@impl true
def handle_info(:health_check, state) do
case make_health_check(state) do
{:ok, _health} ->
# Health check passed, schedule next one
health_ref = schedule_health_check()
{:noreply, %{state | health_check_ref: health_ref}}
{:error, reason} ->
SLog.warning(@log_category, "Health check failed: #{reason}")
# Could implement reconnection logic here
health_ref = schedule_health_check()
{:noreply, %{state | health_check_ref: health_ref}}
end
end
@impl true
def handle_info({:DOWN, _ref, :port, port, reason}, %{server_port: port} = state) do
# Use same shutdown detection as exit_status handler to avoid race conditions.
# :DOWN can arrive before or instead of exit_status on some platforms.
effective_shutting_down? =
state.shutting_down or
shutdown_pending_in_mailbox?() or
not pool_alive?(state.pool_name)
state =
if effective_shutting_down? and not state.shutting_down do
%{state | shutting_down: true}
else
state
end
if effective_shutting_down? do
SLog.debug(@log_category, """
gRPC port DOWN during shutdown
Worker: #{state.id}
Reason: #{inspect(reason)}
""")
{:stop, :shutdown, state}
else
SLog.error(@log_category, """
External gRPC process died unexpectedly
Worker: #{state.id}
Reason: #{inspect(reason)}
""")
{:stop, {:external_process_died, reason}, state}
end
end
@impl true
def handle_info({:EXIT, monitor_pid, exit_reason}, %{heartbeat_monitor: monitor_pid} = state) do
SLog.warning(
@log_category,
"Heartbeat monitor for #{state.id} exited with #{inspect(exit_reason)}; terminating worker"
)
{:stop, {:shutdown, exit_reason}, %{state | heartbeat_monitor: nil}}
end
@impl true
def handle_info({port, {:data, data}}, %{server_port: port} = state) do
output = to_string(data)
buffer = append_startup_output(state.python_output_buffer, output)
if log_python_output?() do
trimmed = String.trim(output)
if trimmed != "" do
SLog.info(@log_category, "gRPC server output: #{trimmed}")
end
end
{:noreply, %{state | python_output_buffer: buffer}}
end
@impl true
def handle_info({port, {:exit_status, status}}, %{server_port: port} = state) do
# DIAGNOSTIC: Drain any remaining error output from the port buffer
remaining_output = drain_port_buffer(port, 200)
last_output =
state.python_output_buffer
|> append_startup_output(remaining_output)
|> String.trim()
last_output =
if last_output == "" do
""
else
last_output
end
# Compute effective shutdown status to handle mailbox race conditions.
# The port exit message may arrive before the {:EXIT, _, :shutdown} message is processed.
# We check multiple signals to determine if we're in a shutdown scenario:
# 1. state.shutting_down was already set
# 2. A shutdown EXIT message is pending in the mailbox
# 3. The pool is no longer alive (system is shutting down)
effective_shutting_down? =
state.shutting_down or
shutdown_pending_in_mailbox?() or
not pool_alive?(state.pool_name)
# Update state if we detected shutdown via mailbox peek or pool check
state =
if effective_shutting_down? and not state.shutting_down do
%{state | shutting_down: true}
else
state
end
# Shutdown exit codes: 0 (clean), 143 (SIGTERM: 128+15), 137 (SIGKILL: 128+9)
# These are expected during shutdown and should not be treated as errors.
case {status, effective_shutting_down?} do
{s, true} when s in [0, 137, 143] ->
# Expected shutdown - Python exited with a normal shutdown code
SLog.debug(@log_category, """
Python gRPC server exited during shutdown (status #{s})
Worker: #{state.id}
Port: #{state.port}
PID: #{state.process_pid}
""")
{:stop, :shutdown, state}
{0, false} ->
# Unexpected but clean exit - Python decided to exit on its own
# This could be idle timeout, internal shutdown, or other reason
SLog.warning(@log_category, """
Python gRPC server exited unexpectedly (status 0)
Worker: #{state.id}
Port: #{state.port}
PID: #{state.process_pid}
Last output: #{last_output}
""")
# Use an abnormal reason so Worker.Starter (with :transient) will restart.
# This maintains pool capacity when Python exits unexpectedly.
{:stop, {:grpc_server_exited_unexpectedly, 0}, state}
{_nonzero, _} ->
# Real crash - non-zero exit status (not a shutdown code)
SLog.error(@log_category, """
🔴 Python gRPC server crashed with status #{status}
Worker: #{state.id}
Port: #{state.port}
PID: #{state.process_pid}
Last output: #{last_output}
""")
{:stop, {:grpc_server_exited, status}, state}
end
end
# Handle shutdown signals from supervisor.
# Matches both :shutdown and {:shutdown, term} which supervisors use.
# Does not match :normal since that can come from other linked processes (like Tasks).
@impl true
def handle_info({:EXIT, _from, reason}, state) when reason == :shutdown do
SLog.debug(@log_category, """
Received shutdown signal for worker #{state.id}
Reason: #{inspect(reason)}
Setting shutting_down flag and stopping gracefully
""")
{:stop, :shutdown, %{state | shutting_down: true}}
end
@impl true
def handle_info({:EXIT, _from, {:shutdown, term} = reason}, state) do
SLog.debug(@log_category, """
Received shutdown signal for worker #{state.id}
Reason: #{inspect(reason)}
Setting shutting_down flag and stopping gracefully
""")
{:stop, {:shutdown, term}, %{state | shutting_down: true}}
end
@impl true
def handle_info(msg, state) do
SLog.debug(@log_category, "Unexpected message: #{inspect(msg)}")
{:noreply, state}
end
# Graceful shutdown timeout for Python process termination.
# Must be >= Python's shutdown envelope: server.stop(2s) + wait_for_termination(3s) = 5s
# We use 6s as default to provide margin. Configurable via :graceful_shutdown_timeout_ms.
@default_graceful_shutdown_timeout 6000
# Margin added to graceful_shutdown_timeout for supervisor shutdown.
# This gives the worker time to complete its terminate/2 callback.
@shutdown_margin 2000
defp graceful_shutdown_timeout do
Application.get_env(
:snakepit,
:graceful_shutdown_timeout_ms,
@default_graceful_shutdown_timeout
)
end
@doc """
Returns the recommended supervisor shutdown timeout.
This is `graceful_shutdown_timeout + margin` to ensure supervisors give workers
enough time to complete their terminate/2 callback (which includes graceful
Python process termination).
Use this value for:
- `shutdown:` in child_spec
- `shutdown:` in Worker.Starter
- Any other supervisor that manages GRPCWorker processes
## Example
children = [
%{
id: MyWorker,
start: {Snakepit.GRPCWorker, :start_link, [opts]},
shutdown: Snakepit.GRPCWorker.supervisor_shutdown_timeout()
}
]
"""
def supervisor_shutdown_timeout do
graceful_shutdown_timeout() + @shutdown_margin
end
@impl true
def terminate(reason, state) do
SLog.debug(
@log_category,
"GRPCWorker.terminate/2 called for #{state.id}, reason: #{inspect(reason)}, PID: #{state.process_pid}"
)
SLog.debug(@log_category, "gRPC worker #{state.id} terminating: #{inspect(reason)}")
emit_worker_terminated_telemetry(state, reason)
cleanup_heartbeat(state, reason)
kill_python_process(state, reason)
cleanup_ready_file(state.ready_file)
cleanup_resources(state)
:ok
end
defp emit_worker_terminated_telemetry(state, reason) do
start_time = Map.get(state.stats, :start_time, 0)
total_commands = Map.get(state.stats, :requests, 0)
:telemetry.execute(
[:snakepit, :pool, :worker, :terminated],
%{
lifetime: System.monotonic_time() - start_time,
total_commands: total_commands
},
%{
node: node(),
pool_name: state.pool_name,
worker_id: state.id,
worker_pid: self(),
reason: reason,
planned: shutdown_reason?(reason) or reason == :normal
}
)
end
defp cleanup_heartbeat(state, reason) do
maybe_stop_heartbeat_monitor(state.heartbeat_monitor)
maybe_notify_test_pid(state.heartbeat_config, {:heartbeat_monitor_stopped, state.id, reason})
end
# Group all kill_python_process clauses together
defp kill_python_process(%{process_pid: nil}, _reason), do: :ok
defp kill_python_process(state, reason) when reason == :normal do
do_graceful_kill(state)
end
defp kill_python_process(state, :shutdown) do
do_graceful_kill(state)
end
defp kill_python_process(state, {:shutdown, _}) do
do_graceful_kill(state)
end
defp kill_python_process(state, reason) do
SLog.warning(
@log_category,
"Non-graceful termination (#{inspect(reason)}), immediately killing PID #{state.process_pid}"
)
result =
if use_process_group_kill?(state) do
Snakepit.ProcessKiller.kill_process_group(state.pgid, :sigkill)
else
Snakepit.ProcessKiller.kill_process(state.process_pid, :sigkill)
end
case result do
:ok ->
SLog.debug(@log_category, "✅ Immediately killed gRPC server PID #{state.process_pid}")
{:error, kill_reason} ->
SLog.warning(
@log_category,
"Failed to kill #{state.process_pid}: #{inspect(kill_reason)}"
)
end
end
defp do_graceful_kill(state) do
SLog.debug(
@log_category,
"Starting graceful shutdown of external gRPC process PID: #{state.process_pid}..."
)
result =
if use_process_group_kill?(state) do
Snakepit.ProcessKiller.kill_process_group_with_escalation(
state.pgid,
graceful_shutdown_timeout()
)
else
Snakepit.ProcessKiller.kill_with_escalation(
state.process_pid,
graceful_shutdown_timeout()
)
end
case result do
:ok ->
SLog.debug(@log_category, "✅ gRPC server PID #{state.process_pid} terminated gracefully")
{:error, kill_reason} ->
SLog.warning(
@log_category,
"Failed to gracefully kill #{state.process_pid}: #{inspect(kill_reason)}"
)
end
end
defp use_process_group_kill?(%{process_group?: true, pgid: pgid})
when is_integer(pgid) do
Application.get_env(:snakepit, :process_group_kill, true)
end
defp use_process_group_kill?(_state), do: false
defp cleanup_resources(state) do
disconnect_connection(state.connection)
cancel_health_check_timer(state.health_check_ref)
close_server_port(state.server_port)
GrpcStream.unregister_worker(state.id)
ProcessRegistry.unregister_worker(state.id)
end
defp cancel_health_check_timer(nil), do: :ok
defp cancel_health_check_timer(health_check_ref) do
Process.cancel_timer(health_check_ref)
end
defp close_server_port(nil), do: :ok
defp close_server_port(server_port) do
safe_close_port(server_port)
end
defp notify_pool_ready(nil, worker_id),
do:
{:error,
{:pool_handshake_failed, Error.pool_error("Pool not found", %{worker_id: worker_id})}}
defp notify_pool_ready(pool_pid, worker_id) when is_pid(pool_pid) do
GenServer.call(pool_pid, {:worker_ready, worker_id}, Defaults.worker_ready_timeout())
catch
:exit, {:noproc, _} ->
{:error,
{:pool_handshake_failed,
Error.pool_error("Pool not found", %{worker_id: worker_id, pool_pid: pool_pid})}}
:exit, {:shutdown, _} = reason ->
{:error, {:pool_handshake_failed, reason}}
:exit, {:killed, _} = reason ->
{:error, {:pool_handshake_failed, reason}}
:exit, {:timeout, _} = reason ->
{:error, {:pool_handshake_failed, reason}}
:exit, reason ->
{:error, {:pool_handshake_failed, reason}}
else
:ok ->
:ok
other ->
{:error, {:pool_handshake_failed, {:unexpected_reply, other}}}
end
defp maybe_start_heartbeat_monitor(state) do
config = normalize_heartbeat_config(state.heartbeat_config)
cond do
not config[:enabled] ->
maybe_stop_heartbeat_monitor(state.heartbeat_monitor)
%{state | heartbeat_config: config, heartbeat_monitor: nil}
heartbeat_monitor_running?(state.heartbeat_monitor) ->
%{state | heartbeat_config: config}
state.connection == nil ->
%{state | heartbeat_config: config}
true ->
monitor_opts = [
{:worker_pid, self()},
{:worker_id, state.id},
{:ping_interval_ms, config[:ping_interval_ms]},
{:timeout_ms, config[:timeout_ms]},
{:max_missed_heartbeats, config[:max_missed_heartbeats]},
{:initial_delay_ms, config[:initial_delay_ms]},
{:dependent, config[:dependent]},
{:ping_fun, config[:ping_fun] || build_default_ping_fun(state, config)}
]
case Snakepit.HeartbeatMonitor.start_link(monitor_opts) do
{:ok, monitor_pid} ->
maybe_notify_test_pid(config, {:heartbeat_monitor_started, state.id, monitor_pid})
%{state | heartbeat_monitor: monitor_pid, heartbeat_config: config}
{:error, {:already_started, monitor_pid}} when is_pid(monitor_pid) ->
maybe_notify_test_pid(config, {:heartbeat_monitor_started, state.id, monitor_pid})
%{state | heartbeat_monitor: monitor_pid, heartbeat_config: config}
{:error, reason} ->
SLog.error(
@log_category,
"Failed to start heartbeat monitor for #{state.id}: #{inspect(reason)}"
)
maybe_notify_test_pid(config, {:heartbeat_monitor_failed, state.id, reason})
%{state | heartbeat_monitor: nil, heartbeat_config: config}
end
end
end
defp heartbeat_monitor_running?(pid) when is_pid(pid) do
Process.alive?(pid)
end
defp heartbeat_monitor_running?(_), do: false
defp build_default_ping_fun(state, config) do
connection = state.connection
adapter = state.adapter
session_id = state.session_id
channel = connection && Map.get(connection, :channel)
fn timestamp ->
result =
cond do
function_exported?(adapter, :grpc_heartbeat, 3) ->
adapter.grpc_heartbeat(connection, session_id, config)
function_exported?(adapter, :grpc_heartbeat, 2) ->
adapter.grpc_heartbeat(connection, session_id)
heartbeat_channel_available?(channel) ->
Client.heartbeat(channel, session_id, timeout: config[:timeout_ms])
true ->
{:error,
Error.grpc_error(:no_heartbeat_transport, "No heartbeat transport available", %{
adapter: adapter,
session_id: session_id
})}
end
handle_heartbeat_response(self(), timestamp, result)
end
end
defp heartbeat_channel_available?(channel) when is_map(channel), do: true
defp heartbeat_channel_available?(channel) when is_struct(channel), do: true
defp heartbeat_channel_available?(channel) when is_reference(channel), do: true
defp heartbeat_channel_available?(channel) when is_pid(channel), do: true
defp heartbeat_channel_available?(channel) when is_binary(channel), do: byte_size(channel) > 0
defp heartbeat_channel_available?(_), do: false
defp maybe_initialize_session(connection, session_id) do
channel = connection && Map.get(connection, :channel)
if heartbeat_channel_available?(channel) do
try do
_ = Client.initialize_session(channel, session_id, %{})
:ok
rescue
exception ->
SLog.debug(
@log_category,
"Heartbeat session initialization failed: #{inspect(exception)}"
)
:error
catch
:exit, reason ->
SLog.debug(@log_category, "Heartbeat session initialization exited: #{inspect(reason)}")
:error
end
else
:error
end
end
defp register_telemetry_stream(connection, state) do
channel = connection && Map.get(connection, :channel)
if channel do
try do
worker_ctx = %{
worker_id: state.id,
pool_name: state.pool_name,
python_pid: state.process_pid
}
GrpcStream.register_worker(channel, worker_ctx)
SLog.debug(@log_category, "Registered telemetry stream for worker #{state.id}")
:ok
rescue
exception ->
SLog.warning(
@log_category,
"Failed to register telemetry stream for worker #{state.id}: #{inspect(exception)}"
)
:error
catch
:exit, reason ->
SLog.warning(
@log_category,
"Telemetry stream registration exited for worker #{state.id}: #{inspect(reason)}"
)
:error
end
else
SLog.debug(@log_category, "No channel available for telemetry stream registration")
:error
end
end
defp handle_heartbeat_response(monitor_pid, timestamp, :ok) do
Snakepit.HeartbeatMonitor.notify_pong(monitor_pid, timestamp)
:ok
end
defp handle_heartbeat_response(monitor_pid, timestamp, {:ok, %{success: success}})
when success in [true, true, 1] do
Snakepit.HeartbeatMonitor.notify_pong(monitor_pid, timestamp)
:ok
end
defp handle_heartbeat_response(_monitor_pid, _timestamp, {:ok, %{success: false} = payload}) do
{:error, {:heartbeat_failed, payload}}
end
defp handle_heartbeat_response(monitor_pid, timestamp, {:ok, _response}) do
Snakepit.HeartbeatMonitor.notify_pong(monitor_pid, timestamp)
:ok
end
defp handle_heartbeat_response(_monitor_pid, _timestamp, {:error, reason}) do
{:error, reason}
end
defp handle_heartbeat_response(_monitor_pid, _timestamp, other) do
{:error, other}
end
defp maybe_notify_test_pid(%{test_pid: pid}, message) when is_pid(pid) do
send(pid, message)
:ok
end
defp maybe_notify_test_pid(%{"test_pid" => pid}, message) when is_pid(pid) do
send(pid, message)
:ok
end
defp maybe_notify_test_pid(_config, _message), do: :ok
defp maybe_stop_heartbeat_monitor(nil), do: :ok
defp maybe_stop_heartbeat_monitor(pid) when is_pid(pid) do
if Process.alive?(pid) do
try do
GenServer.stop(pid, :shutdown)
catch
:exit, _ -> :ok
end
end
:ok
end
defp normalize_heartbeat_config(config) when is_map(config) do
defaults = default_heartbeat_config()
normalized =
Enum.reduce(@heartbeat_known_keys, %{}, fn key, acc ->
Map.put(acc, key, get_config_value(config, key, Map.get(defaults, key)))
end)
extras =
config
|> Enum.reject(fn {key, _value} -> heartbeat_known_key?(key) end)
|> Map.new()
Map.merge(extras, normalized)
end
defp normalize_heartbeat_config(_config), do: default_heartbeat_config()
defp default_heartbeat_config do
Map.merge(
base_heartbeat_defaults(),
Snakepit.Config.heartbeat_defaults(),
fn _key, _base, override -> override end
)
end
defp get_config_value(config, key, default) when is_atom(key) do
cond do
Map.has_key?(config, key) ->
Map.get(config, key)
Map.has_key?(config, Atom.to_string(key)) ->
Map.get(config, Atom.to_string(key))
true ->
default
end
end
defp heartbeat_known_key?(key) when is_atom(key) do
key in @heartbeat_known_keys
end
defp heartbeat_known_key?(key) when is_binary(key) do
key in @heartbeat_known_key_strings
end
defp disconnect_connection(nil), do: :ok
defp disconnect_connection(%{channel: %GRPC.Channel{} = channel}) do
GRPC.Stub.disconnect(channel)
rescue
_ -> :ok
catch
:exit, _ -> :ok
end
defp disconnect_connection(%{channel: channel}) when not is_nil(channel), do: :ok
defp disconnect_connection(_), do: :ok
defp get_worker_config_section(config, key) when is_map(config) and is_atom(key) do
cond do
Map.has_key?(config, key) ->
Map.get(config, key)
Map.has_key?(config, Atom.to_string(key)) ->
Map.get(config, Atom.to_string(key))
true ->
nil
end
end
defp get_worker_config_section(_config, _key), do: nil
defp encode_heartbeat_env(config) when is_map(config) do
config
|> Map.new()
|> Map.take([
:enabled,
:ping_interval_ms,
:timeout_ms,
:max_missed_heartbeats,
:initial_delay_ms,
:dependent
])
|> Enum.reduce(%{}, fn
{:ping_interval_ms, value}, acc -> Map.put(acc, "interval_ms", value)
{key, value}, acc -> Map.put(acc, to_string(key), value)
end)
|> Enum.reject(fn {_key, value} -> is_nil(value) end)
|> Map.new()
|> case do
%{} = map when map == %{} -> nil
map -> Jason.encode!(map)
end
end
defp normalize_adapter_env_entries(nil), do: []
defp normalize_adapter_env_entries(env) when is_map(env) do
env
|> Map.to_list()
|> normalize_adapter_env_entries()
end
defp normalize_adapter_env_entries(env) when is_list(env) do
Enum.flat_map(env, fn
{key, value} -> [{to_string(key), to_string(value)}]
key when is_binary(key) -> [{key, ""}]
key when is_atom(key) -> [{Atom.to_string(key), ""}]
_ -> []
end)
end
defp merge_with_default_adapter_env(env) do
existing = normalize_adapter_env_entries(env)
defaults = default_adapter_env()
existing_keys =
existing
|> Enum.map(fn {key, _} -> String.downcase(key) end)
|> MapSet.new()
defaults
|> Enum.reject(fn {key, _value} -> MapSet.member?(existing_keys, String.downcase(key)) end)
|> Kernel.++(existing)
end
defp default_adapter_env do
priv_python =
:code.priv_dir(:snakepit)
|> to_string()
|> Path.join("python")
repo_priv_python =
Path.join(File.cwd!(), "priv/python")
snakebridge_priv_python =
case :code.priv_dir(:snakebridge) do
{:error, _} -> nil
priv_dir -> Path.join([to_string(priv_dir), "python"])
end
path_sep = path_separator()
pythonpath =
[System.get_env("PYTHONPATH"), priv_python, repo_priv_python, snakebridge_priv_python]
|> Enum.reject(&(&1 in [nil, ""]))
|> Enum.uniq()
|> Enum.join(path_sep)
interpreter =
Application.get_env(:snakepit, :python_executable) ||
System.get_env("SNAKEPIT_PYTHON") ||
GRPCPython.executable_path()
process_group_env =
if Application.get_env(:snakepit, :process_group_kill, true) and
Snakepit.ProcessKiller.process_group_supported?() do
[{"SNAKEPIT_PROCESS_GROUP", "1"}]
else
[]
end
base =
[]
|> maybe_cons("PYTHONPATH", pythonpath)
|> maybe_cons("SNAKEPIT_PYTHON", interpreter)
extra_env =
Snakepit.PythonRuntime.config()
|> Map.get(:extra_env, %{})
|> normalize_adapter_env_entries()
base ++ process_group_env ++ extra_env ++ Snakepit.PythonRuntime.runtime_env()
end
defp maybe_cons(acc, _key, value) when value in [nil, ""], do: acc
defp maybe_cons(acc, key, value), do: [{key, value} | acc]
defp path_separator do
case :os.type() do
{:win32, _} -> ";"
_ -> ":"
end
end
defp elixir_to_python_level(:debug), do: "debug"
defp elixir_to_python_level(:info), do: "info"
defp elixir_to_python_level(:warning), do: "warning"
defp elixir_to_python_level(:error), do: "error"
defp elixir_to_python_level(:none), do: "none"
defp elixir_to_python_level(_), do: "error"
defp maybe_put_heartbeat_env(entries, nil), do: entries
defp maybe_put_heartbeat_env(entries, json),
do: maybe_put_env(entries, "SNAKEPIT_HEARTBEAT_CONFIG", json)
defp maybe_put_ready_file_env(entries, ready_file),
do: maybe_put_env(entries, "SNAKEPIT_READY_FILE", ready_file)
defp maybe_put_log_level_env(entries) do
level = Application.get_env(:snakepit, :log_level, :error)
maybe_put_env(entries, "SNAKEPIT_LOG_LEVEL", elixir_to_python_level(level))
end
defp maybe_put_env(entries, _key, value) when value in [nil, ""], do: entries
defp maybe_put_env(entries, key, value) do
filtered =
Enum.reject(entries, fn {existing_key, _value} ->
String.downcase(existing_key) == String.downcase(key)
end)
[{key, value} | filtered]
end
defp to_env_tuple({key, value}) do
{String.to_charlist(key), String.to_charlist(value)}
end
# CRITICAL FIX: Defensive port cleanup that handles all exit scenarios
defp safe_close_port(port) do
Port.close(port)
rescue
# ArgumentError is raised if the port is already closed
ArgumentError -> :ok
# Catch any other exceptions
_ -> :ok
catch
# Handle exits (e.g., from brutal :kill)
:exit, _ -> :ok
# Handle throws
:throw, _ -> :ok
end
# Private functions
# Drain remaining output from port buffer to capture error messages
defp drain_port_buffer(port, timeout) do
drain_port_buffer(port, timeout, [])
end
defp drain_port_buffer(port, timeout, acc) do
receive do
{^port, {:data, data}} ->
output = to_string(data)
drain_port_buffer(port, timeout, [output | acc])
after
timeout ->
# No more data, return accumulated output
acc
|> Enum.reverse()
|> Enum.join("")
|> String.trim()
end
end
defp wait_for_server_ready(port, ready_file, timeout, output_buffer \\ "") do
deadline = System.monotonic_time(:millisecond) + timeout
wait_for_server_ready_loop(port, ready_file, deadline, timeout, output_buffer)
end
defp wait_for_server_ready_loop(port, ready_file, deadline, timeout, output_buffer) do
case read_ready_file(ready_file) do
{:ok, actual_port} ->
cleanup_ready_file(ready_file)
{:ok, actual_port}
{:error, reason} ->
cleanup_ready_file(ready_file)
log_startup_output(output_buffer)
SLog.error(@log_category, "Failed to read readiness file: #{inspect(reason)}")
{:error, {:ready_file, reason}}
:not_ready ->
remaining = max(deadline - System.monotonic_time(:millisecond), 0)
if remaining == 0 do
cleanup_ready_file(ready_file)
log_startup_output(output_buffer)
SLog.error(
@log_category,
"Timeout waiting for Python gRPC server to start after #{timeout}ms"
)
{:error, :timeout}
else
receive do
{^port, {:data, data}} ->
output = to_string(data)
output_buffer = append_startup_output(output_buffer, output)
if String.trim(output) != "" do
SLog.debug(
@log_category,
"Python server output during startup: #{String.trim(output)}"
)
end
wait_for_server_ready_loop(port, ready_file, deadline, timeout, output_buffer)
{^port, {:exit_status, status}} ->
cleanup_ready_file(ready_file)
case status do
0 ->
SLog.debug(
@log_category,
"Python gRPC server process exited with status 0 during startup (shutdown)"
)
{:error, :shutdown}
143 ->
SLog.debug(
@log_category,
"Python gRPC server process exited with status 143 during startup (shutdown)"
)
{:error, :shutdown}
_ ->
log_startup_output(output_buffer)
SLog.error(
@log_category,
"Python gRPC server process exited with status #{status} during startup"
)
{:error, {:exit_status, status}}
end
{:DOWN, _ref, :port, ^port, reason} ->
cleanup_ready_file(ready_file)
log_startup_output(output_buffer)
SLog.error(
@log_category,
"Python gRPC server port died during startup: #{inspect(reason)}"
)
{:error, {:port_died, reason}}
after
min(remaining, 50) ->
wait_for_server_ready_loop(port, ready_file, deadline, timeout, output_buffer)
end
end
end
end
defp read_ready_file(path) do
case File.read(path) do
{:ok, contents} ->
case Integer.parse(String.trim(contents)) do
{port, _} ->
{:ok, port}
:error ->
# Empty or invalid content - file may still be mid-write (atomic rename race)
# Treat as not ready and keep polling
:not_ready
end
{:error, :enoent} ->
:not_ready
{:error, reason} ->
{:error, reason}
end
end
defp cleanup_ready_file(path) do
File.rm(path)
:ok
rescue
_ -> :ok
end
defp append_startup_output(buffer, output) do
max_bytes = 4096
buffer = buffer <> output
if byte_size(buffer) > max_bytes do
String.slice(buffer, byte_size(buffer) - max_bytes, max_bytes)
else
buffer
end
end
defp log_startup_output(buffer) do
trimmed = String.trim(buffer)
if trimmed != "" do
SLog.error(@log_category, "Python server output during startup:\n#{trimmed}")
end
end
defp log_python_output? do
Application.get_env(:snakepit, :log_python_output, false)
end
defp schedule_health_check do
Process.send_after(self(), :health_check, Defaults.grpc_worker_health_check_interval())
end
defp make_health_check(state) do
case Client.health(state.connection.channel, inspect(self())) do
{:ok, health_response} ->
{:ok, health_response}
{:error, reason} ->
{:error,
Error.grpc_error(:health_check_failed, "Health check failed", %{
worker_id: state.id,
reason: reason
})}
end
end
defp make_info_call(state) do
case Client.get_info(state.connection.channel) do
{:ok, info_response} ->
{:ok, info_response}
end
end
defp update_stats(state, result) do
stats =
case result do
:success ->
%{state.stats | requests: state.stats.requests + 1}
:error ->
%{
state.stats
| requests: state.stats.requests + 1,
errors: state.stats.errors + 1
}
end
%{state | stats: stats}
end
defp instrument_execute(kind, state, command, args, timeout, fun) when is_function(fun, 1) do
correlation_id = correlation_id_from(args)
metadata = base_execute_metadata(kind, state, command, args, correlation_id, timeout)
span_name = otel_span_name(kind, command)
span_attributes = otel_start_attributes(state, command, args, correlation_id, timeout)
:telemetry.span([:snakepit, :grpc_worker, kind], metadata, fn ->
Tracer.with_span span_name, %{attributes: span_attributes, kind: :client} do
start_time = System.monotonic_time()
result = fun.(args)
duration_native = System.monotonic_time() - start_time
duration_ms = System.convert_time_unit(duration_native, :native, :millisecond)
measurements =
%{duration_ms: duration_ms, executions: 1}
|> maybe_track_error_measurement(result)
stop_metadata = build_stop_metadata(metadata, result)
Tracer.set_attributes(otel_stop_attributes(result, duration_ms, stop_metadata))
maybe_set_span_status(result, stop_metadata)
{result, measurements, stop_metadata}
end
end)
end
defp build_stop_metadata(metadata, {:error, {kind, reason}}) do
metadata
|> Map.put(:status, :error)
|> Map.put(:error_kind, kind)
|> Map.put(:error, reason)
end
defp build_stop_metadata(metadata, {:error, reason}) do
metadata
|> Map.put(:status, :error)
|> Map.put(:error, reason)
end
defp build_stop_metadata(metadata, _result) do
Map.put(metadata, :status, :ok)
end
defp maybe_track_error_measurement(measurements, {:error, _reason}) do
Map.put(measurements, :errors, 1)
end
defp maybe_track_error_measurement(measurements, _result), do: measurements
defp ensure_correlation(nil) do
id = Correlation.new_id()
%{"correlation_id" => id, correlation_id: id}
end
defp ensure_correlation(args) when is_map(args) do
existing =
Map.get(args, :correlation_id) ||
Map.get(args, "correlation_id")
id = Correlation.ensure(existing)
args
|> Map.put(:correlation_id, id)
|> Map.put("correlation_id", id)
end
defp ensure_correlation(args) when is_list(args) do
args
|> Map.new()
|> ensure_correlation()
end
defp correlation_id_from(%{} = args) do
args
|> Map.get(:correlation_id)
|> case do
nil -> Map.get(args, "correlation_id")
value -> value
end
|> Correlation.ensure()
end
defp base_execute_metadata(kind, state, command, args, correlation_id, timeout) do
session_id =
Map.get(args, :session_id) ||
Map.get(args, "session_id") ||
state.session_id
%{
operation: kind,
worker_id: state.id,
worker_pid: self(),
command: command,
adapter: adapter_name(state.adapter),
adapter_module: state.adapter,
pool: state.pool_name,
session_id: session_id,
correlation_id: correlation_id,
timeout_ms: timeout,
span_kind: :client,
rpc_system: :grpc,
telemetry_source: :snakepit_grpc_worker
}
end
defp otel_span_name(kind, command) do
operation = kind |> Atom.to_string() |> String.replace("_", "-")
"snakepit.grpc.#{operation}.#{command}"
end
defp otel_start_attributes(state, command, args, correlation_id, timeout) do
session_id = Map.get(args, :session_id) || Map.get(args, "session_id") || state.session_id
[
{"snakepit.worker.id", state.id},
{"snakepit.pool", pool_attribute(state.pool_name)},
{"snakepit.command", command},
{"snakepit.session_id", session_id},
{"snakepit.correlation_id", correlation_id},
{"snakepit.timeout_ms", timeout}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp otel_stop_attributes(result, duration_ms, metadata) do
[
{"snakepit.grpc.duration_ms", duration_ms},
{"snakepit.grpc.status", metadata[:status]},
{"snakepit.grpc.error", format_reason(metadata[:error] || error_from_result(result))},
{"snakepit.grpc.error_kind", metadata[:error_kind]}
]
|> Enum.reject(fn {_k, v} -> is_nil(v) end)
end
defp maybe_set_span_status({:error, _}, metadata) do
reason = format_reason(metadata[:error]) || "snakepit.grpc.error"
Tracer.set_status(:error, reason)
end
defp maybe_set_span_status(_result, _metadata), do: :ok
defp pool_attribute(nil), do: nil
defp pool_attribute(pool) when is_atom(pool), do: Atom.to_string(pool)
defp pool_attribute(pool), do: inspect(pool)
defp format_reason(nil), do: nil
defp format_reason({kind, reason}), do: "#{inspect(kind)}: #{inspect(reason)}"
defp format_reason(reason), do: inspect(reason)
defp error_from_result({:error, {kind, reason}}), do: {kind, reason}
defp error_from_result({:error, reason}), do: reason
defp error_from_result(_), do: nil
defp adapter_name(module) when is_atom(module) do
module
|> Atom.to_string()
|> String.replace_prefix("Elixir.", "")
end
defp adapter_name(other), do: inspect(other)
# Shutdown detection helpers
# These eliminate race conditions between port exit messages and shutdown signals
@doc false
# Matches both :shutdown and {:shutdown, term} which supervisors use
defp shutdown_reason?(:shutdown), do: true
defp shutdown_reason?({:shutdown, _}), do: true
defp shutdown_reason?(_), do: false
@doc false
# Peek into mailbox to detect if a shutdown signal is pending but not yet processed.
# This handles the race where port exit arrives before the EXIT message is processed.
# Only called on rare port-exit path, not on hot request paths.
defp shutdown_pending_in_mailbox? do
case Process.info(self(), :messages) do
{:messages, msgs} ->
Enum.any?(msgs, fn
{:EXIT, _from, reason} -> shutdown_reason?(reason)
_ -> false
end)
_ ->
false
end
end
@doc false
# Check if the pool is still alive - if not, we're in system shutdown
defp pool_alive?(pool_name) do
case resolve_pool_pid(pool_name) do
pid when is_pid(pid) -> Process.alive?(pid)
_ -> false
end
end
end
defmodule Snakepit.Hardware do
@moduledoc """
Hardware abstraction layer for Snakepit.
Provides unified hardware detection and device selection for ML workloads.
Supports CPU, NVIDIA CUDA, Apple MPS, and AMD ROCm accelerators.
## Features
- **Automatic Detection**: Detects available hardware at startup
- **Device Selection**: Intelligent device selection with fallback strategies
- **Caching**: Results are cached for performance
- **Lock File Support**: Identity map for lock file generation
## Usage
# Detect all hardware
info = Snakepit.Hardware.detect()
# => %{accelerator: :cuda, cpu: %{...}, cuda: %{...}, ...}
# Check capabilities
caps = Snakepit.Hardware.capabilities()
# => %{cuda: true, mps: false, avx2: true, ...}
# Select device
{:ok, device} = Snakepit.Hardware.select(:auto)
# => {:ok, {:cuda, 0}}
# Select with fallback
{:ok, device} = Snakepit.Hardware.select_with_fallback([:cuda, :mps, :cpu])
# => {:ok, :cpu}
## Identity Map
The `identity/0` function returns a map suitable for lock file generation:
identity = Snakepit.Hardware.identity()
# => %{"platform" => "linux-x86_64", "accelerator" => "cuda", ...}
This can be serialized to JSON/YAML for lock files that need to track
the hardware environment.
"""
alias Snakepit.Hardware.{Detector, Selector}
@type device :: Selector.device()
@type device_preference :: Selector.device_preference()
@type hardware_info :: Detector.hardware_info()
@type capabilities :: Detector.capabilities()
@doc """
Detects all hardware information.
Returns a map with:
- `:accelerator` - Primary accelerator type (`:cpu`, `:cuda`, `:mps`, `:rocm`)
- `:cpu` - CPU information (cores, threads, model, features, memory)
- `:cuda` - NVIDIA CUDA info or nil
- `:mps` - Apple MPS info or nil
- `:rocm` - AMD ROCm info or nil
- `:platform` - Platform string (e.g., "linux-x86_64")
## Examples
info = Snakepit.Hardware.detect()
info.accelerator
#=> :cuda
info.cpu.cores
#=> 8
"""
@spec detect() :: hardware_info()
defdelegate detect(), to: Detector
@doc """
Alias for `detect/0`.
Returns the same hardware info map as detect/0.
"""
@spec info() :: hardware_info()
def info, do: detect()
@doc """
Returns hardware capability flags.
Returns a map of boolean flags for quick feature checks:
- `:cuda` - CUDA available
- `:mps` - Apple MPS available
- `:rocm` - AMD ROCm available
- `:avx` - AVX instruction set available
- `:avx2` - AVX2 instruction set available
- `:avx512` - AVX-512 instruction set available
- `:cuda_version` - CUDA version string or nil
- `:cudnn_version` - cuDNN version string or nil
- `:cudnn` - cuDNN available
## Examples
caps = Snakepit.Hardware.capabilities()
if caps.cuda do
cuda_version = caps.cuda_version
end
"""
@spec capabilities() :: capabilities()
defdelegate capabilities(), to: Detector
@doc """
Clears the hardware detection cache.
Forces re-detection on next call. Useful after hardware changes
or for testing.
## Examples
Snakepit.Hardware.clear_cache()
:ok
"""
@spec clear_cache() :: :ok
defdelegate clear_cache(), to: Detector
@doc """
Selects a device based on preference.
## Options
- `:auto` - Automatically select best available accelerator
- `:cpu` - Select CPU (always available)
- `:cuda` - Select CUDA (fails if not available)
- `:mps` - Select MPS (fails if not macOS with Apple Silicon)
- `:rocm` - Select ROCm (fails if not available)
- `{:cuda, device_id}` - Select specific CUDA device by ID
## Returns
- `{:ok, device}` on success
- `{:error, :device_not_available}` if requested device is unavailable
## Examples
# Auto-select best device
{:ok, device} = Snakepit.Hardware.select(:auto)
# Request specific device
case Snakepit.Hardware.select(:cuda) do
{:ok, {:cuda, 0}} -> :ok
{:error, :device_not_available} -> :error
end
"""
@spec select(device_preference()) :: {:ok, device()} | {:error, :device_not_available}
defdelegate select(preference), to: Selector
@doc """
Selects the first available device from a preference list.
Tries each device in order until one is available. This is useful
for graceful degradation strategies.
## Examples
# Prefer CUDA, fall back to MPS, then CPU
{:ok, device} = Snakepit.Hardware.select_with_fallback([:cuda, :mps, :cpu])
# Returns :cpu if CUDA and MPS are unavailable
"""
@spec select_with_fallback([device_preference()]) :: {:ok, device()} | {:error, :no_device}
defdelegate select_with_fallback(preferences), to: Selector
@doc """
Returns device information for a selected device.
Returns a map with device-specific details useful for logging,
telemetry, and diagnostics.
## Examples
info = Snakepit.Hardware.device_info({:cuda, 0})
# => %{type: :cuda, device_id: 0, name: "NVIDIA GeForce RTX 3080", ...}
"""
@spec device_info(device()) :: map()
defdelegate device_info(device), to: Selector
@doc """
Returns a hardware identity map for lock files.
The identity map contains string keys and is suitable for
serialization to JSON/YAML lock files that need to track
the hardware environment.
## Keys
- `"platform"` - Platform string (e.g., "linux-x86_64")
- `"accelerator"` - Primary accelerator type as string
- `"cpu_features"` - List of CPU feature strings
- `"gpu_count"` - Number of GPUs detected
## Examples
identity = Snakepit.Hardware.identity()
Jason.encode!(identity)
# => "{\\"platform\\":\\"linux-x86_64\\",\\"accelerator\\":\\"cuda\\",...}"
"""
@spec identity() :: map()
def identity do
info = detect()
gpu_count =
cond do
info.cuda != nil -> length(info.cuda.devices)
info.rocm != nil -> length(info.rocm.devices)
info.mps != nil and info.mps.available -> 1
true -> 0
end
cpu_features =
info.cpu.features
|> Enum.map(&Atom.to_string/1)
%{
"platform" => info.platform,
"accelerator" => Atom.to_string(info.accelerator),
"cpu_features" => cpu_features,
"gpu_count" => gpu_count
}
end
end
defmodule Snakepit.HealthMonitor do
@moduledoc """
Monitors worker health and crash patterns.
Tracks crashes within a rolling window and determines overall pool health.
Can be used to trigger circuit breaker actions or alerting.
## Usage
{:ok, hm} = HealthMonitor.start_link(
name: :my_pool_health,
pool: :default,
max_crashes: 10,
crash_window_ms: 60_000
)
HealthMonitor.record_crash(hm, "worker_1", %{reason: :segfault})
if HealthMonitor.healthy?(hm) do
# Pool is healthy
else
# Too many crashes, consider action
end
"""
use GenServer
alias Snakepit.Defaults
require Logger
@type worker_stats :: %{
crash_count: non_neg_integer(),
last_crash_time: integer() | nil,
crash_reasons: [term()]
}
@type t :: %{
pool: atom(),
workers: %{String.t() => worker_stats()},
crash_window_ms: pos_integer(),
max_crashes: pos_integer(),
total_crashes: non_neg_integer(),
check_interval_ms: pos_integer(),
check_timer: reference() | nil
}
# Client API
@doc """
Starts a health monitor.
## Options
- `:name` - GenServer name (required)
- `:pool` - Pool name to monitor (required)
- `:check_interval_ms` - Health check interval (default: 30000)
- `:crash_window_ms` - Rolling window for crash counting (default: 60000)
- `:max_crashes` - Max crashes in window before unhealthy (default: 10)
"""
@spec start_link(keyword()) :: GenServer.on_start()
def start_link(opts) do
name = Keyword.fetch!(opts, :name)
GenServer.start_link(__MODULE__, opts, name: name)
end
@doc """
Records a worker crash.
"""
@spec record_crash(GenServer.server(), String.t(), map()) :: :ok
def record_crash(server, worker_id, info \\ %{}) do
GenServer.cast(server, {:record_crash, worker_id, info})
end
@doc """
Returns whether the pool is considered healthy.
"""
@spec healthy?(GenServer.server()) :: boolean()
def healthy?(server) do
GenServer.call(server, :healthy?)
end
@doc """
Returns health status for a specific worker.
"""
@spec worker_health(GenServer.server(), String.t()) :: map()
def worker_health(server, worker_id) do
GenServer.call(server, {:worker_health, worker_id})
end
@doc """
Returns comprehensive health statistics.
"""
@spec stats(GenServer.server()) :: map()
def stats(server) do
GenServer.call(server, :get_stats)
end
# GenServer callbacks
@impl true
def init(opts) do
pool = Keyword.fetch!(opts, :pool)
state = %{
pool: pool,
workers: %{},
crash_window_ms:
Keyword.get(opts, :crash_window_ms, Defaults.health_monitor_crash_window_ms()),
max_crashes: Keyword.get(opts, :max_crashes, Defaults.health_monitor_max_crashes()),
total_crashes: 0,
check_interval_ms:
Keyword.get(opts, :check_interval_ms, Defaults.health_monitor_check_interval()),
check_timer: nil
}
# Schedule periodic cleanup
timer = schedule_cleanup(state.check_interval_ms)
{:ok, %{state | check_timer: timer}}
end
@impl true
def handle_call(:healthy?, _from, state) do
crashes_in_window = count_crashes_in_window(state)
healthy = crashes_in_window < state.max_crashes
{:reply, healthy, state}
end
def handle_call({:worker_health, worker_id}, _from, state) do
worker_stats = Map.get(state.workers, worker_id, default_worker_stats())
health = %{
healthy: worker_stats.crash_count < 3,
crash_count: worker_stats.crash_count,
last_crash_time: worker_stats.last_crash_time
}
{:reply, health, state}
end
def handle_call(:get_stats, _from, state) do
crashes_in_window = count_crashes_in_window(state)
stats = %{
pool: state.pool,
total_crashes: state.total_crashes,
crashes_in_window: crashes_in_window,
workers_with_crashes: map_size(state.workers),
max_crashes: state.max_crashes,
crash_window_ms: state.crash_window_ms,
is_healthy: crashes_in_window < state.max_crashes
}
{:reply, stats, state}
end
@impl true
def handle_cast({:record_crash, worker_id, info}, state) do
now = System.monotonic_time(:millisecond)
worker_stats =
state.workers
|> Map.get(worker_id, default_worker_stats())
|> update_worker_crash(info, now)
workers = Map.put(state.workers, worker_id, worker_stats)
state = %{
state
| workers: workers,
total_crashes: state.total_crashes + 1
}
# Emit telemetry
:telemetry.execute(
[:snakepit, :worker, :crash],
%{},
%{
pool: state.pool,
worker_id: worker_id,
reason: Map.get(info, :reason)
}
)
{:noreply, state}
end
@impl true
def handle_info(:cleanup, state) do
# Remove old crash data
state = cleanup_old_crashes(state)
# Reschedule
timer = schedule_cleanup(state.check_interval_ms)
{:noreply, %{state | check_timer: timer}}
end
def handle_info(_msg, state) do
{:noreply, state}
end
# Private functions
defp schedule_cleanup(interval_ms) do
Process.send_after(self(), :cleanup, interval_ms)
end
defp default_worker_stats do
%{
crash_count: 0,
last_crash_time: nil,
crash_reasons: [],
crash_times: []
}
end
defp update_worker_crash(stats, info, now) do
reason = Map.get(info, :reason)
%{
stats
| crash_count: stats.crash_count + 1,
last_crash_time: now,
crash_reasons: [reason | Enum.take(stats.crash_reasons, 9)],
crash_times: [now | Enum.take(Map.get(stats, :crash_times, []), 99)]
}
end
defp count_crashes_in_window(state) do
now = System.monotonic_time(:millisecond)
window_start = now - state.crash_window_ms
state.workers
|> Map.values()
|> Enum.flat_map(fn stats -> Map.get(stats, :crash_times, []) end)
|> Enum.count(fn time -> time >= window_start end)
end
defp cleanup_old_crashes(state) do
now = System.monotonic_time(:millisecond)
window_start = now - state.crash_window_ms
workers =
state.workers
|> Enum.map(fn {id, stats} ->
crash_times =
stats
|> Map.get(:crash_times, [])
|> Enum.filter(fn time -> time >= window_start end)
{id, Map.put(stats, :crash_times, crash_times)}
end)
|> Enum.filter(fn {_id, stats} ->
# Keep workers with recent crashes
Map.get(stats, :crash_times, []) != []
end)
|> Map.new()
%{state | workers: workers}
end
end
defmodule Snakepit.HeartbeatMonitor do
@moduledoc """
Monitors a worker process using a configurable heartbeat protocol.
The monitor periodically invokes a ping function and expects the worker
to send a pong via `notify_pong/2`. Missed heartbeats trigger worker
termination, allowing supervisors to restart the worker.
"""
use GenServer
alias Snakepit.Logger, as: SLog
@default_ping_interval 2_000
@default_timeout 10_000
@default_max_missed 3
@log_category :worker
@type start_option ::
{:worker_pid, pid()}
| {:worker_id, String.t()}
| {:ping_interval_ms, non_neg_integer()}
| {:timeout_ms, non_neg_integer()}
| {:max_missed_heartbeats, non_neg_integer()}
| {:ping_fun, (integer() -> :ok | {:ok, term()} | {:error, term()} | term())}
| {:dependent, boolean()}
defstruct [
:worker_pid,
:worker_id,
:ping_interval,
:timeout,
:max_missed_heartbeats,
:ping_fun,
:ping_timer,
:timeout_timer,
:last_ping_timestamp,
:initial_delay,
dependent: true,
missed_heartbeats: 0,
stats: %{pings_sent: 0, pongs_received: 0, timeouts: 0}
]
@spec start_link([start_option()]) :: GenServer.on_start()
def start_link(opts) when is_list(opts) do
GenServer.start_link(__MODULE__, opts)
end
@doc """
Notify the monitor that a pong response has been received.
"""
@spec notify_pong(pid(), integer()) :: :ok
def notify_pong(monitor_pid, timestamp) when is_pid(monitor_pid) do
GenServer.cast(monitor_pid, {:pong, timestamp})
end
@impl true
def init(opts) do
worker_pid = Keyword.fetch!(opts, :worker_pid)
worker_id = Keyword.fetch!(opts, :worker_id)
ping_interval = Keyword.get(opts, :ping_interval_ms, @default_ping_interval)
timeout = Keyword.get(opts, :timeout_ms, @default_timeout)
max_missed = Keyword.get(opts, :max_missed_heartbeats, @default_max_missed)
initial_delay = Keyword.get(opts, :initial_delay_ms, 0)
dependent = Keyword.get(opts, :dependent, true)
ping_fun =
Keyword.get(opts, :ping_fun, &default_ping_fun/1)
state = %__MODULE__{
worker_pid: worker_pid,
worker_id: worker_id,
ping_interval: ping_interval,
timeout: timeout,
max_missed_heartbeats: max_missed,
ping_fun: ping_fun,
initial_delay: initial_delay,
dependent: dependent
}
Process.monitor(worker_pid)
new_state = schedule_initial_ping(state)
emit_event(:monitor_started, new_state, %{})
{:ok, new_state}
end
@impl true
def handle_cast({:pong, ping_timestamp}, state) do
now = System.monotonic_time(:millisecond)
if state.timeout_timer do
Process.cancel_timer(state.timeout_timer)
end
new_stats =
state.stats
|> Map.update!(:pongs_received, &(&1 + 1))
new_state =
%{
state
| missed_heartbeats: 0,
timeout_timer: nil,
stats: new_stats,
last_ping_timestamp: ping_timestamp
}
|> schedule_next_ping()
emit_event(:pong_received, new_state, %{latency_ms: now - ping_timestamp})
{:noreply, new_state}
end
@impl true
def handle_info(:send_ping, state) do
timestamp = System.monotonic_time(:millisecond)
result =
try do
state.ping_fun.(timestamp)
rescue
exception ->
{:error, {exception, __STACKTRACE__}}
catch
kind, reason ->
{:error, {kind, reason}}
end
case normalize_ping_result(result) do
:ok ->
stats = Map.update!(state.stats, :pings_sent, &(&1 + 1))
new_state =
%{state | stats: stats, last_ping_timestamp: timestamp}
|> schedule_timeout()
emit_event(:ping_sent, new_state, %{})
{:noreply, new_state}
{:error, reason} ->
SLog.warning(
@log_category,
"Heartbeat ping failed for #{state.worker_id}: #{inspect(reason)}"
)
handle_worker_failure(state, :ping_failed)
end
end
@impl true
def handle_info(:heartbeat_timeout, state) do
missed = state.missed_heartbeats + 1
stats = Map.update!(state.stats, :timeouts, &(&1 + 1))
emit_event(:heartbeat_timeout, %{state | missed_heartbeats: missed, stats: stats}, %{
missed_count: missed
})
if missed >= state.max_missed_heartbeats do
log_message =
"Worker #{state.worker_id} missed #{missed} heartbeat(s); initiating termination"
if state.dependent do
SLog.error(@log_category, log_message)
else
SLog.warning(@log_category, "#{log_message} (worker configured as heartbeat-independent)")
end
handle_worker_failure(
%{state | missed_heartbeats: missed, stats: stats},
:heartbeat_timeout
)
else
new_state =
%{state | missed_heartbeats: missed, stats: stats, timeout_timer: nil}
|> schedule_next_ping()
{:noreply, new_state}
end
end
@impl true
def handle_info({:DOWN, _ref, :process, pid, reason}, %{worker_pid: pid} = state) do
SLog.debug(
@log_category,
"Heartbeat monitor observed worker #{state.worker_id} exit: #{inspect(reason)}"
)
{:stop, {:worker_down, reason}, state}
end
@impl true
def handle_info(message, state) do
SLog.debug(@log_category, "Unhandled heartbeat monitor message: #{inspect(message)}")
{:noreply, state}
end
@impl true
def terminate(reason, state) do
cancel_timer(state.ping_timer)
cancel_timer(state.timeout_timer)
emit_event(:monitor_stopped, state, %{reason: reason})
:ok
end
defp default_ping_fun(_timestamp) do
{:error, :not_implemented}
end
defp schedule_initial_ping(state) do
delay = max(state.initial_delay || 0, 0)
timer = Process.send_after(self(), :send_ping, delay)
%{state | ping_timer: timer, initial_delay: 0}
end
defp schedule_next_ping(state) do
cancel_timer(state.ping_timer)
timer = Process.send_after(self(), :send_ping, state.ping_interval)
%{state | ping_timer: timer}
end
defp schedule_timeout(state) do
cancel_timer(state.timeout_timer)
timer = Process.send_after(self(), :heartbeat_timeout, state.timeout)
%{state | timeout_timer: timer}
end
defp cancel_timer(nil), do: :ok
defp cancel_timer(timer_ref) do
Process.cancel_timer(timer_ref, async: true, info: false)
:ok
end
defp handle_worker_failure(state, reason) do
emit_event(:monitor_failure, state, %{failure_reason: reason})
if state.dependent do
Process.exit(state.worker_pid, {:shutdown, reason})
{:stop, {:shutdown, reason}, state}
else
SLog.debug(
@log_category,
"Heartbeat monitor for #{state.worker_id} suppressing termination (independent worker, reason=#{inspect(reason)})"
)
cancel_timer(state.timeout_timer)
new_state =
%{state | timeout_timer: nil}
|> schedule_next_ping()
{:noreply, new_state}
end
end
defp normalize_ping_result(:ok), do: :ok
defp normalize_ping_result({:ok, _data}), do: :ok
defp normalize_ping_result({:error, reason}), do: {:error, reason}
defp normalize_ping_result(:error), do: {:error, :unknown}
defp normalize_ping_result(other), do: {:error, other}
defp emit_event(event, state, metadata) do
event_name = [:snakepit, :heartbeat, event]
measurements = heartbeat_measurements(event, state, metadata)
meta =
Map.merge(
%{
worker_id: state.worker_id,
worker_pid: state.worker_pid,
missed_heartbeats: state.missed_heartbeats,
dependent: state.dependent
},
metadata
)
:telemetry.execute(event_name, measurements, meta)
end
defp heartbeat_measurements(:ping_sent, _state, _metadata) do
%{timestamp: System.monotonic_time(:millisecond), count: 1, pings: 1}
end
defp heartbeat_measurements(:pong_received, _state, metadata) do
base = %{timestamp: System.monotonic_time(:millisecond), count: 1, pongs: 1}
case Map.get(metadata, :latency_ms) do
nil -> base
latency -> Map.put(base, :latency_ms, latency)
end
end
defp heartbeat_measurements(:monitor_failure, _state, _metadata) do
%{timestamp: System.monotonic_time(:millisecond), count: 1, failures: 1}
end
defp heartbeat_measurements(:heartbeat_timeout, state, metadata) do
missed = Map.get(metadata, :missed_count, state.missed_heartbeats)
%{
timestamp: System.monotonic_time(:millisecond),
count: 1,
timeouts: 1,
missed: missed
}
end
defp heartbeat_measurements(_event, _state, _metadata) do
%{timestamp: System.monotonic_time(:millisecond), count: 1}
end
end
defmodule Snakepit.Logger do
@moduledoc """
Centralized, silent-by-default logging for Snakepit.
## Configuration
# Silent (default) - only errors
config :snakepit, log_level: :error
# Warnings and errors
config :snakepit, log_level: :warning
# Verbose - info, warnings, errors
config :snakepit, log_level: :info
# Debug - everything
config :snakepit, log_level: :debug
# Completely silent (not even errors)
config :snakepit, log_level: :none
## Categories
Enable specific categories for targeted debugging:
config :snakepit, log_categories: [:lifecycle, :grpc]
## Process-Level Isolation (for Testing)
Log levels can be set per-process to avoid race conditions in async tests:
# Set log level for current process only
Snakepit.Logger.set_process_level(:debug)
# Execute with temporary log level
Snakepit.Logger.with_level(:warning, fn ->
# ... code that should log at warning level
end)
# Clear process-level override
Snakepit.Logger.clear_process_level()
The resolution order is:
1. Process-level override (via `set_process_level/1`)
2. Elixir Logger process level (via `Logger.put_process_level/2`)
3. Application config (via `config :snakepit, log_level: ...`)
"""
require Logger
@type category ::
:lifecycle
| :pool
| :grpc
| :bridge
| :worker
| :startup
| :shutdown
| :telemetry
| :general
@type level :: :debug | :info | :warning | :error | :none
@default_level :error
@default_category :general
@process_level_key :snakepit_log_level_override
@category_whitelist [
:lifecycle,
:pool,
:grpc,
:bridge,
:worker,
:startup,
:shutdown,
:telemetry,
:general
]
@doc """
Log at debug level if configured log level allows it.
"""
def debug(category, message, metadata) when category in @category_whitelist do
log(:debug, category, message, metadata)
end
def debug(category, message) when category in @category_whitelist do
log(:debug, category, message, [])
end
def debug(message, metadata), do: log(:debug, @default_category, message, metadata)
def debug(message), do: log(:debug, @default_category, message, [])
@doc """
Log at info level if configured log level allows it.
"""
def info(category, message, metadata) when category in @category_whitelist do
log(:info, category, message, metadata)
end
def info(category, message) when category in @category_whitelist do
log(:info, category, message, [])
end
def info(message, metadata), do: log(:info, @default_category, message, metadata)
def info(message), do: log(:info, @default_category, message, [])
@doc """
Log at warning level if configured log level allows it.
"""
def warning(category, message, metadata) when category in @category_whitelist do
log(:warning, category, message, metadata)
end
def warning(category, message) when category in @category_whitelist do
log(:warning, category, message, [])
end
def warning(message, metadata), do: log(:warning, @default_category, message, metadata)
def warning(message), do: log(:warning, @default_category, message, [])
@doc """
Log at error level if configured log level allows it.
"""
def error(category, message, metadata) when category in @category_whitelist do
log(:error, category, message, metadata)
end
def error(category, message) when category in @category_whitelist do
log(:error, category, message, [])
end
def error(message, metadata), do: log(:error, @default_category, message, metadata)
def error(message), do: log(:error, @default_category, message, [])
@doc """
Check if logging at the given level is enabled.
"""
def should_log?(level), do: should_log_level?(level)
@doc """
Check if logging at the given level/category is enabled.
"""
def should_log?(level, category) do
should_log_level?(level) and category_allowed?(level, category)
end
@doc """
Set the log level for the current process only.
This is useful for test isolation - each test process can have its own
log level without affecting other concurrent tests.
## Examples
Snakepit.Logger.set_process_level(:debug)
# All logging in this process now uses :debug level
Snakepit.Logger.set_process_level(:none)
# All logging in this process is now suppressed
"""
@spec set_process_level(level()) :: :ok
def set_process_level(level) when level in [:debug, :info, :warning, :error, :none] do
Process.put(@process_level_key, level)
:ok
end
@doc """
Get the effective log level for the current process.
Returns the log level in priority order:
1. Process-level override (set via `set_process_level/1`)
2. Elixir Logger process level
3. Application config
"""
@spec get_process_level() :: level()
def get_process_level do
resolve_effective_level()
end
@doc """
Clear the process-level log level override.
After calling this, the process will use the global Application config.
"""
@spec clear_process_level() :: :ok
def clear_process_level do
Process.delete(@process_level_key)
:ok
end
@doc """
Execute a function with a temporary log level for the current process.
The previous log level is restored after the function completes,
even if it raises an exception.
## Examples
Snakepit.Logger.with_level(:debug, fn ->
# Debug logs are enabled here
Snakepit.Logger.debug(:pool, "detailed info")
end)
# Previous log level is restored
"""
@spec with_level(level(), (-> result)) :: result when result: term()
def with_level(level, fun) when is_function(fun, 0) do
previous = Process.get(@process_level_key)
try do
set_process_level(level)
fun.()
after
if previous do
Process.put(@process_level_key, previous)
else
Process.delete(@process_level_key)
end
end
end
defp log(level, category, message, metadata) do
if should_log?(level, category) do
Logger.log(level, message, with_category(metadata, category))
end
end
defp should_log_level?(level) do
effective_level = resolve_effective_level()
case effective_level do
:none -> false
:error -> level == :error
:warning -> level in [:error, :warning]
:info -> level in [:error, :warning, :info]
:debug -> true
_ -> level == :error
end
end
# Resolve the effective log level using priority order:
# 1. Process-level override (highest priority)
# 2. Elixir Logger process level
# 3. Application config (lowest priority)
defp resolve_effective_level do
case Process.get(@process_level_key) do
nil ->
case Logger.get_process_level(self()) do
nil ->
Application.get_env(:snakepit, :log_level, @default_level)
# Map Elixir Logger levels to our levels
logger_level when logger_level in [:emergency, :alert, :critical, :error] ->
:error
:warning ->
:warning
:notice ->
:warning
:info ->
:info
:debug ->
:debug
:all ->
:debug
:none ->
:none
_ ->
Application.get_env(:snakepit, :log_level, @default_level)
end
level ->
level
end
end
defp category_allowed?(level, category) when level in [:debug, :info] do
case Application.get_env(:snakepit, :log_categories) do
nil -> true
categories when is_list(categories) -> category in categories
_ -> true
end
end
defp category_allowed?(_level, _category), do: true
defp with_category(metadata, category) do
metadata
|> normalize_metadata()
|> Keyword.put_new(:category, category)
end
defp normalize_metadata(metadata) when is_list(metadata), do: metadata
defp normalize_metadata(metadata) when is_map(metadata), do: Map.to_list(metadata)
defp normalize_metadata(_metadata), do: []
end
defmodule Snakepit.PackageError do
@moduledoc """
Structured error for Python package installation and inspection.
"""
defexception [:type, :packages, :message, :suggestion, :output]
@type type ::
:not_installed
| :install_failed
| :version_mismatch
| :invalid_requirement
@type t :: %__MODULE__{
type: type(),
packages: [String.t()],
message: String.t(),
suggestion: String.t() | nil,
output: String.t() | nil
}
@impl true
def message(%__MODULE__{message: message, suggestion: suggestion}) do
base = message || "Python package error"
if suggestion do
base <> "\nSuggestion: " <> suggestion
else
base
end
end
end
defmodule Snakepit.ProcessKiller do
@moduledoc """
Robust OS process management using Erlang primitives.
No shell commands, pure Erlang/Elixir.
This module provides POSIX-compliant process management that works
across Linux, macOS, and BSD systems without relying on shell-specific
features like pkill.
"""
alias Snakepit.Logger, as: SLog
@kill_command_candidates ["/bin/kill", "/usr/bin/kill"]
@ps_command_candidates ["/bin/ps", "/usr/bin/ps"]
@log_category :shutdown
@doc """
Returns true if the platform supports process group kill semantics.
"""
def process_group_supported? do
case :os.type() do
{:unix, _} ->
true
_ ->
false
end
end
@doc """
Returns the path to the setsid executable, or {:error, :not_found}.
"""
def setsid_executable do
case System.find_executable("setsid") do
nil -> {:error, :not_found}
path -> {:ok, path}
end
end
@doc """
Returns the setsid executable path or raises if not available.
"""
def setsid_executable! do
case setsid_executable() do
{:ok, path} -> path
{:error, _} -> raise "setsid executable not found"
end
end
@doc """
Kills a process by PID using proper Erlang signals.
## Parameters
- `os_pid`: OS process ID (integer)
- `signal`: :sigterm | :sigkill | :sighup
## Returns
- `:ok` if kill succeeded
- `{:error, reason}` if kill failed
"""
def kill_process(os_pid, signal \\ :sigterm) when is_integer(os_pid) do
signal_num = signal_to_number(signal)
# DEBUG: Log all kills to find who's killing workers during startup
caller = Process.info(self(), :registered_name)
SLog.debug(
@log_category,
"ProcessKiller.kill_process: PID=#{os_pid}, signal=#{signal}, caller=#{inspect(caller)}"
)
with {:ok, kill_path} <- require_executable("kill", @kill_command_candidates),
{:ok, output, code} <-
run_command(kill_path, ["-#{signal_num}", Integer.to_string(os_pid)]) do
trimmed = String.trim(output || "")
cond do
code == 0 ->
:ok
String.contains?(trimmed, "No such process") ->
:ok
true ->
{:error, if(trimmed == "", do: {:exit_status, code}, else: trimmed)}
end
else
{:error, reason} ->
SLog.warning(@log_category, "Failed to execute kill command: #{inspect(reason)}")
{:error, reason}
end
end
@doc """
Kills a process group by PGID using proper Erlang signals.
## Parameters
- `pgid`: Process group ID (integer)
- `signal`: :sigterm | :sigkill | :sighup
"""
def kill_process_group(pgid, signal \\ :sigterm) when is_integer(pgid) do
signal_num = signal_to_number(signal)
with {:ok, kill_path} <- require_executable("kill", @kill_command_candidates),
{:ok, output, code} <-
run_command(kill_path, ["-#{signal_num}", "--", "-#{pgid}"]) do
trimmed = String.trim(output || "")
cond do
code == 0 ->
:ok
String.contains?(trimmed, "No such process") ->
:ok
true ->
{:error, if(trimmed == "", do: {:exit_status, code}, else: trimmed)}
end
else
{:error, reason} ->
SLog.warning(@log_category, "Failed to execute kill command: #{inspect(reason)}")
{:error, reason}
end
end
@doc """
Checks if a process is alive.
Uses kill -0 (signal 0) which doesn't kill but checks existence.
"""
def process_alive?(os_pid) when is_integer(os_pid) do
case :os.type() do
{:unix, :linux} ->
process_alive_via_proc(os_pid)
{:unix, _} ->
process_alive_via_ps(os_pid)
_ ->
false
end
end
def process_alive?(_), do: false
defp process_alive_via_ps(os_pid) do
with {:ok, ps_path} <- require_executable("ps", @ps_command_candidates),
{:ok, output, 0} <-
run_command(ps_path, ["-p", Integer.to_string(os_pid), "-o", "pid="]) do
String.trim(output || "") != ""
else
_ -> false
end
end
defp process_alive_via_proc(os_pid) do
stat_path = "/proc/#{os_pid}/stat"
case File.read(stat_path) do
{:ok, content} ->
case parse_proc_state(content) do
{:ok, state} -> state not in ["Z", "X", "x"]
:error -> File.exists?(stat_path)
end
{:error, :enoent} ->
false
{:error, _} ->
File.exists?(stat_path)
end
end
defp parse_proc_state(content) when is_binary(content) do
case String.split(content, ") ", parts: 2) do
[_prefix, rest] ->
case String.split(rest, " ", parts: 2) do
[state | _] -> {:ok, state}
_ -> :error
end
_ ->
:error
end
end
@doc """
Gets the command line of a process.
POSIX-compliant using /proc on Linux, ps on macOS/BSD.
"""
def get_process_command(os_pid) when is_integer(os_pid) do
# Try Linux /proc first (fastest)
proc_file = "/proc/#{os_pid}/cmdline"
if File.exists?(proc_file) do
# Linux: Read /proc/PID/cmdline
case File.read(proc_file) do
{:ok, content} ->
# cmdline uses null bytes as separators
command = content |> String.split(<<0>>) |> Enum.join(" ") |> String.trim()
{:ok, command}
{:error, _} ->
get_process_command_ps(os_pid)
end
else
# macOS/BSD: Use ps command
get_process_command_ps(os_pid)
end
end
@doc """
Gets the process group ID (PGID) for a process.
"""
def get_process_group_id(os_pid) when is_integer(os_pid) do
case :os.type() do
{:unix, :linux} ->
get_process_group_id_proc(os_pid)
{:unix, _} ->
get_process_group_id_ps(os_pid)
_ ->
{:error, :not_supported}
end
end
defp get_process_group_id_proc(os_pid) do
stat_path = "/proc/#{os_pid}/stat"
case File.read(stat_path) do
{:ok, content} -> parse_stat_pgid(content)
{:error, _} -> {:error, :not_found}
end
end
defp parse_stat_pgid(content) do
case String.split(content, ") ", parts: 2) do
[_prefix, rest] -> extract_pgid_from_stat(rest)
_ -> {:error, :parse_error}
end
end
defp extract_pgid_from_stat(rest) do
case String.split(rest, " ", trim: true) do
[_state, _ppid, pgrp | _] -> parse_pgid(pgrp)
_ -> {:error, :parse_error}
end
end
defp parse_pgid(pgrp) do
case Integer.parse(pgrp) do
{pgid, ""} -> {:ok, pgid}
_ -> {:error, :parse_error}
end
end
defp get_process_group_id_ps(os_pid) do
with {:ok, ps_path} <- require_executable("ps", @ps_command_candidates),
{:ok, output, 0} <-
run_command(ps_path, ["-p", Integer.to_string(os_pid), "-o", "pgid="]) do
case Integer.parse(String.trim(output || "")) do
{pgid, ""} -> {:ok, pgid}
_ -> {:error, :not_found}
end
else
{:error, {:executable_not_found, _cmd}} ->
{:error, :not_found}
_ ->
{:error, :not_found}
end
end
defp get_process_command_ps(os_pid) do
with {:ok, ps_path} <- require_executable("ps", @ps_command_candidates),
{:ok, output, 0} <-
run_command(ps_path, ["-p", Integer.to_string(os_pid), "-o", "args="]) do
case String.trim(output || "") do
"" -> {:error, :not_found}
cmd -> {:ok, cmd}
end
else
{:error, {:executable_not_found, _cmd}} ->
{:error, :not_found}
_ ->
{:error, :not_found}
end
end
@doc """
Kills all processes matching a run ID.
Pure Erlang implementation, no pkill.
"""
def kill_by_run_id(run_id) when is_binary(run_id) do
SLog.warning(@log_category, "🔪 Killing all processes with run_id: #{run_id}")
SLog.debug(@log_category, "kill_by_run_id called at: #{System.monotonic_time(:millisecond)}")
caller_info = Process.info(self(), [:registered_name, :current_stacktrace])
SLog.debug(@log_category, "Called from: #{inspect(caller_info)}")
# Get all Python processes
python_pids = find_python_processes()
# Filter by run_id in command line
# Support both --snakepit-run-id (current) and --run-id (future)
matching_pids =
python_pids
|> Enum.filter(fn pid ->
case get_process_command(pid) do
{:ok, cmd} ->
has_grpc_server = String.contains?(cmd, "grpc_server.py")
has_old_format = String.contains?(cmd, "--snakepit-run-id #{run_id}")
has_new_format = String.contains?(cmd, "--run-id #{run_id}")
has_grpc_server and (has_old_format or has_new_format)
_ ->
false
end
end)
SLog.info(@log_category, "Found #{length(matching_pids)} processes to kill")
# Kill with escalation
killed_count =
Enum.reduce(matching_pids, 0, fn pid, acc ->
case kill_with_escalation(pid) do
:ok ->
acc + 1
{:error, reason} ->
SLog.warning(@log_category, "Failed to kill #{pid}: #{inspect(reason)}")
acc
end
end)
{:ok, killed_count}
end
defp find_python_processes_linux do
case File.ls("/proc") do
{:ok, entries} ->
entries
|> Enum.reduce([], &collect_python_pid/2)
|> Enum.uniq()
{:error, _} ->
find_python_processes_posix()
end
end
defp collect_python_pid(entry, acc) do
case Integer.parse(entry) do
{pid, ""} ->
if python_command?(pid), do: [pid | acc], else: acc
_ ->
acc
end
end
defp python_command?(pid) do
comm_path = "/proc/#{pid}/comm"
cmdline_path = "/proc/#{pid}/cmdline"
if File.exists?(comm_path) do
case File.read(comm_path) do
{:ok, comm} ->
comm
|> String.trim()
|> String.downcase()
|> String.contains?("python")
_ ->
python_cmdline?(cmdline_path)
end
else
python_cmdline?(cmdline_path)
end
rescue
_ -> false
end
defp python_cmdline?(path) do
case File.read(path) do
{:ok, content} ->
try do
content
|> String.replace(<<0>>, " ")
|> String.downcase()
|> String.contains?("python")
rescue
ArgumentError -> false
end
_ ->
false
end
end
defp find_python_processes_posix do
with {:ok, ps_path} <- require_executable("ps", @ps_command_candidates),
{:ok, output, 0} <- run_command(ps_path, ["-eo", "pid,comm"]) do
output
|> String.split("\n", trim: true)
|> Enum.reduce([], fn line, acc ->
trimmed = String.trim_leading(line)
case Regex.split(~r/\s+/, trimmed, parts: 2) do
[pid_str, command] ->
with {pid, ""} <- Integer.parse(pid_str),
true <- String.contains?(String.downcase(command), "python") do
[pid | acc]
else
_ -> acc
end
_ ->
acc
end
end)
|> Enum.reverse()
else
{:error, {:executable_not_found, _cmd}} ->
SLog.warning(@log_category, "ps command not available; skipping python process discovery")
[]
_ ->
[]
end
end
@doc """
Finds all Python processes on the system.
Returns a list of OS PIDs.
"""
def find_python_processes do
case :os.type() do
{:unix, :linux} -> find_python_processes_linux()
{:unix, _} -> find_python_processes_posix()
_ -> []
end
end
@doc """
Kills a process with escalation: SIGTERM -> wait -> SIGKILL
"""
def kill_with_escalation(os_pid, timeout_ms \\ 2000) do
# Try SIGTERM first (graceful)
case kill_process(os_pid, :sigterm) do
:ok ->
# Wait for process to die
if wait_for_death(os_pid, timeout_ms) do
SLog.debug(@log_category, "✅ Process #{os_pid} terminated gracefully")
:ok
else
# Escalate to SIGKILL
SLog.warning(@log_category, "⏰ Process #{os_pid} didn't die, escalating to SIGKILL")
kill_process(os_pid, :sigkill)
end
error ->
error
end
end
@doc """
Kills a process group with escalation: SIGTERM -> wait -> SIGKILL.
"""
def kill_process_group_with_escalation(pgid, timeout_ms \\ 2000) when is_integer(pgid) do
case kill_process_group(pgid, :sigterm) do
:ok ->
if wait_for_death(pgid, timeout_ms) do
SLog.debug(@log_category, "✅ Process group #{pgid} terminated gracefully")
:ok
else
SLog.warning(@log_category, "⏰ Process group #{pgid} didn't die, escalating to SIGKILL")
kill_process_group(pgid, :sigkill)
end
error ->
error
end
end
defp wait_for_death(os_pid, timeout_ms) do
deadline = System.monotonic_time(:millisecond) + timeout_ms
wait_for_death_loop(os_pid, deadline, 1)
end
# Non-blocking polling with exponential backoff using receive after.
# Starts at 1ms, doubles to 2ms, 4ms, 8ms, capping at 100ms.
# This is the OTP-correct way to implement timed waits without blocking the scheduler.
defp wait_for_death_loop(os_pid, deadline, backoff) do
if System.monotonic_time(:millisecond) >= deadline do
false
else
if process_alive?(os_pid) do
delay = min(backoff, 100)
# OTP-idiomatic non-blocking wait - integrates with process mailbox and scheduler
receive do
after
delay -> :ok
end
wait_for_death_loop(os_pid, deadline, backoff * 2)
else
true
end
end
end
defp signal_to_number(:sigterm), do: 15
defp signal_to_number(:sigkill), do: 9
defp signal_to_number(:sighup), do: 1
defp signal_to_number(n) when is_integer(n), do: n
defp require_executable(cmd, fallback_paths) when is_list(fallback_paths) do
case System.find_executable(cmd) do
nil ->
fallback_paths
|> Enum.find(&File.exists?/1)
|> case do
nil -> {:error, {:executable_not_found, cmd}}
path -> {:ok, path}
end
path ->
{:ok, path}
end
end
defp run_command(path, args) when is_binary(path) and is_list(args) do
{output, status} = System.cmd(path, args, stderr_to_stdout: true)
{:ok, output, status}
rescue
error -> {:error, error}
end
end
defmodule Snakepit.PythonPackages do
@moduledoc """
Package installation and inspection for Snakepit-managed Python runtimes.
Uses uv when available or configured, with pip as the fallback.
## Examples
Snakepit.PythonPackages.ensure!({:list, ["numpy~=1.26", "scipy~=1.11"]})
{:ok, :all_installed} =
Snakepit.PythonPackages.check_installed(["numpy~=1.26", "scipy~=1.11"])
{:ok, metadata} = Snakepit.PythonPackages.lock_metadata(["numpy~=1.26"])
Snakepit.PythonPackages.ensure!({:file, "requirements.txt"}, upgrade: true)
"""
alias Snakepit.PackageError
alias Snakepit.PythonRuntime
@type requirement :: String.t()
@type requirements_spec ::
{:list, [requirement()]}
| {:file, Path.t()}
@default_env %{
"PYTHONNOUSERSITE" => "1",
"PIP_DISABLE_PIP_VERSION_CHECK" => "1",
"PIP_NO_INPUT" => "1",
"PIP_NO_WARN_SCRIPT_LOCATION" => "1",
"UV_NO_PROGRESS" => "1",
"PYTHONDONTWRITEBYTECODE" => "1"
}
@default_timeout 300_000
@doc """
Ensure all packages in the requirements spec are installed.
Options:
* `:upgrade` - upgrade matching packages
* `:quiet` - suppress installer output
* `:timeout` - install timeout in ms
"""
@spec ensure!(requirements_spec(), keyword()) :: :ok | no_return()
def ensure!(spec, opts \\ []) do
requirements = normalize_spec!(spec)
case check_installed(requirements, opts) do
{:ok, :all_installed} -> :ok
{:ok, {:missing, missing}} -> install!(missing, opts)
end
end
@doc """
Check which packages are installed.
Returns `{:ok, :all_installed}` when every requirement is present, or
`{:ok, {:missing, requirements}}` when any are missing.
"""
@spec check_installed([requirement()], keyword()) ::
{:ok, :all_installed} | {:ok, {:missing, [requirement()]}}
def check_installed(requirements, opts \\ []) do
requirements = normalize_list(requirements)
validate_requirements!(requirements)
do_check_installed(requirements, opts)
end
@doc """
Return package metadata for lockfiles.
The result maps package name to `%{version: version, hash: hash}` entries.
"""
@spec lock_metadata([requirement()], keyword()) :: {:ok, map()} | {:error, term()}
def lock_metadata(requirements, opts \\ []) do
requirements = normalize_list(requirements)
validate_requirements!(requirements)
if requirements == [] do
{:ok, %{}}
else
python = package_python!(opts)
case freeze(installer(), python, opts) do
{:ok, output} -> {:ok, build_metadata(requirements, output)}
{:error, %PackageError{} = error} -> {:error, error}
end
end
end
@doc """
Return the active installer (`:uv` or `:pip`).
"""
@spec installer() :: :uv | :pip
def installer do
config = config()
case config.installer do
:auto -> detect_installer()
:uv -> :uv
:pip -> :pip
_ -> detect_installer()
end
end
@doc """
Install the given package requirements.
Prefer `ensure!/2` unless you already know which requirements are missing.
"""
@spec install!([requirement()], keyword()) :: :ok | no_return()
def install!(requirements, opts \\ []) do
requirements = normalize_list(requirements)
validate_requirements!(requirements)
if requirements == [] do
:ok
else
python = package_python!(opts)
case installer() do
:uv -> install_with_uv(python, requirements, opts)
:pip -> install_with_pip(python, requirements, opts)
end
end
end
defmodule Runner do
@moduledoc false
@callback cmd(String.t(), [String.t()], keyword()) :: {String.t(), non_neg_integer()}
end
defmodule Runner.System do
@moduledoc false
@behaviour Runner
@impl true
def cmd(command, args, opts) do
{timeout, opts} = Keyword.pop(opts, :timeout)
run = fn -> System.cmd(command, args, opts) end
if is_integer(timeout) do
task = Task.async(run)
case Task.yield(task, timeout) || Task.shutdown(task, :brutal_kill) do
{:ok, result} -> result
nil -> {"Command timed out after #{timeout}ms", 124}
end
else
run.()
end
rescue
error in ErlangError ->
{"#{Exception.message(error)}", 127}
end
end
defp config do
raw =
:snakepit
|> Application.get_env(:python_packages, [])
|> normalize_config_input()
runtime_env = normalize_env_input(PythonRuntime.config().extra_env)
env_override = normalize_env_input(Map.get(raw, :env, %{}))
%{
installer: Map.get(raw, :installer, :auto),
timeout: Map.get(raw, :timeout, @default_timeout),
env: @default_env |> Map.merge(runtime_env) |> Map.merge(env_override),
runner: Map.get(raw, :runner, Runner.System),
env_dir: Map.get(raw, :env_dir)
}
end
defp normalize_spec!({:list, requirements}) when is_list(requirements) do
normalize_list(requirements)
end
defp normalize_spec!({:file, path}) do
case File.read(path) do
{:ok, contents} ->
contents
|> String.split("\n")
|> normalize_list()
{:error, reason} ->
raise PackageError,
type: :invalid_requirement,
packages: [path],
message: "Could not read requirements file: #{inspect(reason)}",
suggestion: "Check the file path and permissions."
end
end
defp normalize_spec!(spec) do
raise PackageError,
type: :invalid_requirement,
packages: [inspect(spec)],
message: "Unsupported requirements spec",
suggestion: "Use {:list, requirements} or {:file, path}."
end
defp normalize_list(nil), do: []
defp normalize_list(requirements) when is_list(requirements) do
requirements
|> Enum.map(&normalize_requirement!/1)
|> Enum.reject(&(&1 == ""))
end
defp normalize_list(other) do
raise PackageError,
type: :invalid_requirement,
packages: [inspect(other)],
message: "Requirements must be a list of strings",
suggestion: "Provide a list like [\"numpy~=1.26\"]."
end
defp normalize_requirement!(requirement) when is_binary(requirement) do
requirement
|> String.split("#", parts: 2)
|> hd()
|> String.trim()
end
defp normalize_requirement!(requirement) do
raise PackageError,
type: :invalid_requirement,
packages: [inspect(requirement)],
message: "Requirement must be a string",
suggestion: "Use PEP-440 requirement strings."
end
defp validate_requirements!(requirements) do
invalid =
requirements
|> Enum.reject(&is_binary/1)
if invalid != [] do
raise PackageError,
type: :invalid_requirement,
packages: invalid,
message: "Requirements must be strings",
suggestion: "Use PEP-440 requirement strings."
end
end
defp do_check_installed([], _opts), do: {:ok, :all_installed}
defp do_check_installed(requirements, opts) do
python = package_python!(opts)
missing =
requirements
|> Enum.reduce([], fn requirement, acc ->
name = requirement_name!(requirement)
if package_installed?(installer(), python, name, opts) do
acc
else
[requirement | acc]
end
end)
|> Enum.reverse()
case missing do
[] -> {:ok, :all_installed}
_ -> {:ok, {:missing, missing}}
end
end
defp requirement_name!(requirement) do
name =
requirement
|> String.split(";", parts: 2)
|> hd()
|> String.split(~r/\s*@\s*/, parts: 2)
|> hd()
|> String.trim()
|> String.split("[", parts: 2)
|> hd()
|> String.split(~r/[<>=!~]/, parts: 2)
|> hd()
|> String.trim()
if name == "" do
raise PackageError,
type: :invalid_requirement,
packages: [requirement],
message: "Could not derive package name from requirement",
suggestion: "Use PEP-440 requirement strings."
end
name
end
defp base_python! do
case PythonRuntime.resolve_executable() do
{:ok, python, _meta} ->
python
{:error, reason} ->
raise PackageError,
type: :install_failed,
packages: [],
message: "Python runtime unavailable: #{inspect(reason)}",
suggestion: "Run mix snakepit.setup or set SNAKEPIT_PYTHON."
end
end
defp package_python!(opts) do
base = base_python!()
case env_dir(config()) do
nil ->
base
env_dir ->
ensure_venv!(base, env_dir, opts)
venv_python!(env_dir)
end
end
defp package_installed?(:uv, python, package, opts) do
{_, status} =
run_cmd(
uv_path!(),
["pip", "show", package, "--python", python],
opts
)
status == 0
end
defp package_installed?(:pip, python, package, opts) do
{_, status} = run_cmd(python, ["-m", "pip", "show", package], opts)
status == 0
end
defp install_with_uv(python, requirements, opts) do
args =
["pip", "install", "--python", python] ++
build_install_args(opts) ++
requirements
{output, status} = run_cmd(uv_path!(), args, opts)
if status == 0 do
:ok
else
raise PackageError,
type: :install_failed,
packages: requirements,
message: "UV install failed with exit code #{status}",
output: output,
suggestion: "Check package names and network connectivity."
end
end
defp install_with_pip(python, requirements, opts) do
args =
["-m", "pip", "install"] ++
build_install_args(opts) ++
requirements
{output, status} = run_cmd(python, args, opts)
if status == 0 do
:ok
else
raise PackageError,
type: :install_failed,
packages: requirements,
message: "Pip install failed with exit code #{status}",
output: output,
suggestion: "Check package names and network connectivity."
end
end
defp freeze(:uv, python, opts) do
{output, status} =
run_cmd(
uv_path!(),
["pip", "freeze", "--python", python],
opts
)
if status == 0 do
{:ok, output}
else
{:error,
%PackageError{
type: :install_failed,
packages: [],
message: "UV freeze failed with exit code #{status}",
output: output,
suggestion: "Verify the Python environment is accessible."
}}
end
end
defp freeze(:pip, python, opts) do
{output, status} = run_cmd(python, ["-m", "pip", "freeze"], opts)
if status == 0 do
{:ok, output}
else
{:error,
%PackageError{
type: :install_failed,
packages: [],
message: "Pip freeze failed with exit code #{status}",
output: output,
suggestion: "Verify the Python environment is accessible."
}}
end
end
defp build_metadata(requirements, output) do
versions = parse_freeze_output(output)
Enum.reduce(requirements, %{}, fn requirement, acc ->
name = requirement_name!(requirement)
version = Map.get(versions, String.downcase(name))
Map.put(acc, name, %{version: version, hash: nil})
end)
end
defp parse_freeze_output(output) do
output
|> String.split("\n", trim: true)
|> Enum.reduce(%{}, fn line, acc ->
case String.split(line, "==", parts: 2) do
[name, version] ->
Map.put(acc, String.downcase(String.trim(name)), String.trim(version))
_ ->
acc
end
end)
end
defp build_install_args(opts) do
args = []
args = if Keyword.get(opts, :upgrade, false), do: args ++ ["--upgrade"], else: args
args = if Keyword.get(opts, :quiet, false), do: args ++ ["--quiet"], else: args
args
end
defp run_cmd(command, args, opts) do
runner = opts[:runner] || config().runner
runner.cmd(command, args, cmd_opts(opts))
end
defp cmd_opts(opts) do
config = config()
timeout = Keyword.get(opts, :timeout, config.timeout)
[
env: Map.to_list(config.env),
stderr_to_stdout: true,
timeout: timeout
]
end
defp env_dir(config) do
case Map.get(config, :env_dir) do
false -> nil
nil -> default_env_dir()
value -> Path.expand(value, File.cwd!())
end
end
defp default_env_dir do
runtime_dir = PythonRuntime.config().runtime_dir || "priv/snakepit/python"
Path.expand(Path.join(runtime_dir, "venv"), File.cwd!())
end
defp ensure_venv!(base_python, env_dir, opts) do
if venv_python_paths(env_dir) |> Enum.any?(&File.exists?/1) do
:ok
else
File.mkdir_p!(env_dir)
{output, status} = run_cmd(base_python, ["-m", "venv", env_dir], opts)
if status == 0 do
:ok
else
raise PackageError,
type: :install_failed,
packages: [],
message: "Failed to create virtual environment in #{env_dir}",
output: output,
suggestion: "Ensure python includes venv and the directory is writable."
end
end
end
defp venv_python!(env_dir) do
case Enum.find(venv_python_paths(env_dir), &File.exists?/1) do
nil ->
raise PackageError,
type: :install_failed,
packages: [],
message: "Virtual environment missing Python executable in #{env_dir}",
suggestion: "Remove the directory and retry setup."
python ->
python
end
end
defp venv_python_paths(env_dir) do
[
Path.join([env_dir, "bin", "python3"]),
Path.join([env_dir, "bin", "python"]),
Path.join([env_dir, "Scripts", "python.exe"]),
Path.join([env_dir, "Scripts", "python"])
]
end
defp detect_installer do
if uv_path() do
:uv
else
:pip
end
end
defp uv_path do
PythonRuntime.config().uv_path || System.find_executable("uv")
end
defp uv_path! do
case uv_path() do
nil ->
raise PackageError,
type: :install_failed,
packages: [],
message: "uv not found",
suggestion: "Install uv or set :python, uv_path: \"/path/to/uv\"."
path ->
path
end
end
defp normalize_config_input(nil), do: %{}
defp normalize_config_input(%{} = map), do: map
defp normalize_config_input(list) when is_list(list), do: Map.new(list)
defp normalize_config_input(_), do: %{}
defp normalize_env_input(nil), do: %{}
defp normalize_env_input(%{} = map), do: map
defp normalize_env_input(list) when is_list(list), do: Map.new(list)
defp normalize_env_input(_), do: %{}
end
defmodule Snakepit.PythonRuntime do
@moduledoc """
Resolve and manage the Python runtime used by Snakepit.
"""
@default_config %{
strategy: :system,
managed: false,
python_version: "3.11.8",
runtime_dir: "priv/snakepit/python",
cache_dir: "priv/snakepit/python/cache",
uv_path: nil,
extra_env: %{"PYTHONNOUSERSITE" => "1"}
}
@identity_key {__MODULE__, :runtime_identity}
def config do
config =
:snakepit
|> Application.get_env(:python, [])
|> normalize_config_input()
@default_config
|> Map.merge(config)
|> normalize_strategy()
end
def managed?(config \\ config()) do
config.strategy == :uv and config.managed == true
end
def executable_path do
case resolve_executable() do
{:ok, path, _meta} -> path
{:error, _reason} -> nil
end
end
def resolve_executable do
config = config()
case override_python() do
{:ok, path} ->
{:ok, path, %{source: :override}}
:none ->
case package_env_python(config) do
{:ok, path} ->
{:ok, path, %{source: :package_env}}
:none ->
resolve_managed_or_fallback(config)
end
end
end
defp resolve_managed_or_fallback(config) do
case managed_executable(config) do
{:ok, path} -> {:ok, path, %{source: :managed}}
{:error, :not_managed} -> resolve_fallback(config)
{:error, reason} -> {:error, reason}
end
end
def runtime_identity do
case :persistent_term.get(@identity_key, nil) do
nil -> compute_and_cache_identity()
identity -> identity
end
end
def runtime_metadata do
case runtime_identity() do
{:ok, identity} ->
%{
"python_runtime_hash" => identity.hash,
"python_version" => identity.version,
"python_platform" => identity.platform
}
_ ->
%{}
end
end
def runtime_env do
case runtime_identity() do
{:ok, identity} ->
[
{"SNAKEPIT_PYTHON_RUNTIME_HASH", identity.hash},
{"SNAKEPIT_PYTHON_VERSION", identity.version},
{"SNAKEPIT_PYTHON_PLATFORM", identity.platform}
]
_ ->
[]
end
end
def missing_reason(config \\ config()) do
if managed?(config) do
uv = uv_path(config)
cond do
is_nil(uv) ->
{:error, "uv not found. Install uv or set :python, uv_path: \"/path/to/uv\"."}
not runtime_installed?(config) ->
{:error,
"Managed Python missing in #{runtime_dir(config)}. Run mix snakepit.setup to install it."}
true ->
:ok
end
else
:ok
end
end
def install_managed(runner, opts \\ []) do
config = config()
uv = uv_path(config)
cond do
not managed?(config) ->
:ok
is_nil(uv) ->
{:error, :uv_not_found}
true ->
do_install_managed(runner, config, uv, opts)
end
end
defp normalize_strategy(config) do
strategy =
case Map.get(config, :strategy) do
nil ->
if Map.get(config, :managed) == true do
:uv
else
:system
end
value ->
value
end
Map.put(config, :strategy, strategy)
end
defp managed_executable(config) do
if managed?(config) do
find_managed_executable(config)
else
{:error, :not_managed}
end
end
defp find_managed_executable(config) do
runtime_dir = runtime_dir(config)
candidates = [
Path.join([runtime_dir, "bin", "python3"]),
Path.join([runtime_dir, "bin", "python"])
]
case Enum.find(candidates, &File.exists?/1) do
nil ->
find_nested_executable(runtime_dir)
path ->
{:ok, path}
end
end
defp find_nested_executable(runtime_dir) do
nested =
Path.wildcard(Path.join([runtime_dir, "**", "bin", "python3"])) ++
Path.wildcard(Path.join([runtime_dir, "**", "bin", "python"]))
case Enum.find(nested, &File.exists?/1) do
nil -> {:error, :managed_missing}
path -> {:ok, path}
end
end
defp do_install_managed(runner, config, uv, opts) do
runtime_dir = runtime_dir(config)
cache_dir = cache_dir(config)
project_root = opts[:project_root] || File.cwd!()
File.mkdir_p!(runtime_dir)
File.mkdir_p!(cache_dir)
env =
config.extra_env
|> Map.new()
|> Map.merge(%{
"UV_PYTHON_INSTALL_DIR" => runtime_dir,
"UV_CACHE_DIR" => cache_dir
})
args = ["python", "install", config.python_version]
case runner.cmd(uv, args, cd: project_root, env: env) do
:ok -> :ok
{:error, reason} -> {:error, {:uv_install_failed, reason}}
end
end
defp resolve_fallback(_config) do
cond do
venv = find_venv_python() -> {:ok, venv, %{source: :venv}}
system = system_python() -> {:ok, system, %{source: :system}}
true -> {:error, :not_found}
end
end
defp override_python do
case Application.get_env(:snakepit, :python_executable) || System.get_env("SNAKEPIT_PYTHON") do
nil -> :none
path -> {:ok, path}
end
end
defp package_env_python(config) do
case python_packages_env_dir(config) do
nil ->
:none
env_dir ->
case Enum.find(venv_python_paths(env_dir), &File.exists?/1) do
nil -> :none
path -> {:ok, path}
end
end
end
defp python_packages_env_dir(config) do
case Map.get(python_packages_config(), :env_dir) do
false -> nil
nil -> Path.join(runtime_dir(config), "venv")
value -> Path.expand(value, project_root())
end
end
defp venv_python_paths(env_dir) do
[
Path.join([env_dir, "bin", "python3"]),
Path.join([env_dir, "bin", "python"]),
Path.join([env_dir, "Scripts", "python.exe"]),
Path.join([env_dir, "Scripts", "python"])
]
end
defp python_packages_config do
:snakepit
|> Application.get_env(:python_packages, [])
|> normalize_config_input()
end
defp normalize_config_input(nil), do: %{}
defp normalize_config_input(%{} = map), do: map
defp normalize_config_input(list) when is_list(list), do: Map.new(list)
defp normalize_config_input(_), do: %{}
defp find_venv_python do
candidates = [
".venv/bin/python3",
"../.venv/bin/python3",
System.get_env("VIRTUAL_ENV") &&
Path.join([System.get_env("VIRTUAL_ENV"), "bin", "python3"])
]
Enum.find_value(candidates, fn path ->
path && File.exists?(Path.expand(path)) && Path.expand(path)
end)
end
defp system_python do
System.find_executable("python3") || System.find_executable("python")
end
defp runtime_dir(config) do
Path.expand(config.runtime_dir, project_root())
end
defp cache_dir(config) do
Path.expand(config.cache_dir, project_root())
end
defp runtime_installed?(config) do
case managed_executable(config) do
{:ok, _} -> true
_ -> false
end
end
defp uv_path(config) do
config.uv_path || System.find_executable("uv")
end
defp project_root do
Application.get_env(:snakepit, :bootstrap_project_root) || File.cwd!()
end
defp compute_and_cache_identity do
identity =
case resolve_executable() do
{:ok, path, _meta} -> build_identity(path)
{:error, reason} -> {:error, reason}
end
:persistent_term.put(@identity_key, identity)
identity
end
defp build_identity(path) do
version = python_version(path)
platform = system_platform()
hash = binary_hash(path)
{:ok, %{path: path, version: version, platform: platform, hash: hash}}
end
defp python_version(path) do
case System.cmd(path, ["-c", "import sys; print(sys.version.split()[0])"],
stderr_to_stdout: true
) do
{output, 0} -> String.trim(output)
_ -> "unknown"
end
end
defp system_platform do
:erlang.system_info(:system_architecture)
|> to_string()
end
defp binary_hash(path) do
case File.read(path) do
{:ok, contents} ->
:crypto.hash(:sha256, contents)
|> Base.encode16(case: :lower)
_ ->
"unknown"
end
end
end
defmodule Snakepit.PythonThreadLimits do
@moduledoc """
Normalizes Python threading configuration with safe defaults.
Resolves partial overrides pulled from application environment and
produces a complete map ready for runtime consumption.
"""
@typedoc "Thread limit configuration keyed by known library identifiers."
@type t :: %{
optional(:openblas) => pos_integer(),
optional(:omp) => pos_integer(),
optional(:mkl) => pos_integer(),
optional(:numexpr) => pos_integer(),
optional(:grpc_poll_threads) => pos_integer()
}
@defaults %{
openblas: 1,
omp: 1,
mkl: 1,
numexpr: 1,
grpc_poll_threads: 1
}
@doc "Default thread limit configuration."
@spec defaults() :: t()
def defaults, do: @defaults
@doc """
Merge a user-supplied configuration with defaults.
Accepts `nil`, maps with atom keys, or keyword lists.
Unknown keys are ignored; non-integer values are coerced with `String.to_integer/1`
when possible.
"""
@spec resolve(nil | map() | keyword()) :: t()
def resolve(nil), do: defaults()
def resolve(config) when is_list(config) do
config
|> Enum.into(%{})
|> resolve()
end
def resolve(config) when is_map(config) do
config
|> Enum.reduce(defaults(), fn {key, value}, acc ->
key = normalize_key(key)
case key do
nil ->
acc
normalized_key ->
Map.put(acc, normalized_key, normalize_value(value))
end
end)
end
def resolve(_unknown), do: defaults()
@known_keys %{
"openblas" => :openblas,
"omp" => :omp,
"mkl" => :mkl,
"numexpr" => :numexpr,
"grpc_poll_threads" => :grpc_poll_threads
}
defp normalize_key(key) when is_atom(key) do
if Map.has_key?(defaults(), key), do: key, else: nil
end
defp normalize_key(key) when is_binary(key), do: Map.get(@known_keys, key)
defp normalize_key(_), do: nil
defp normalize_value(value) when is_integer(value) and value > 0, do: value
defp normalize_value(value) when is_binary(value) do
case Integer.parse(value) do
{parsed, ""} when parsed > 0 -> parsed
_ -> 1
end
end
defp normalize_value(_), do: 1
end
defmodule Snakepit.PythonVersion do
@moduledoc """
Detects the active Python runtime version and recommends worker profiles.
"""
alias Snakepit.PythonRuntime
@type version :: {non_neg_integer(), non_neg_integer(), non_neg_integer()}
@min_supported {3, 8, 0}
@free_threading_min {3, 13, 0}
@spec detect() :: {:ok, version()} | {:error, term()}
def detect do
case PythonRuntime.resolve_executable() do
{:ok, path, _meta} -> detect(path)
{:error, :not_found} -> {:error, :python_not_found}
{:error, :managed_missing} -> {:error, :python_not_found}
end
end
@spec detect(binary()) :: {:ok, version()} | {:error, term()}
def detect(path) when is_binary(path) do
args = ["-c", "import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))"]
case System.cmd(path, args, stderr_to_stdout: true) do
{output, 0} ->
parse_version(output)
{output, _code} ->
{:error, {:python_failed, String.trim(output)}}
end
rescue
_error -> {:error, :python_not_found}
end
@spec supports_free_threading?(version()) :: boolean()
def supports_free_threading?(version) do
version >= @free_threading_min
end
@spec recommend_profile() :: :process | :thread
def recommend_profile do
case detect() do
{:ok, version} -> recommend_profile(version)
_ -> :process
end
end
@spec recommend_profile(version()) :: :process | :thread
def recommend_profile(version) do
if supports_free_threading?(version), do: :thread, else: :process
end
@spec validate() :: :ok | {:error, term()}
def validate do
case detect() do
{:ok, version} ->
if version >= @min_supported do
:ok
else
{:error, {:unsupported_version, version}}
end
{:error, reason} ->
{:error, reason}
end
end
defp parse_version(output) do
version = String.trim(output)
case Regex.run(~r/(\d+)\.(\d+)\.(\d+)/, version) do
[_, major, minor, patch] ->
{:ok, {String.to_integer(major), String.to_integer(minor), String.to_integer(patch)}}
_ ->
{:error, {:invalid_version, version}}
end
end
end
defmodule Snakepit.RetryPolicy do
@moduledoc """
Retry policy with exponential backoff.
Configures retry behavior including max attempts, backoff timing,
and which errors are retriable.
## Usage
policy = RetryPolicy.new(
max_attempts: 3,
backoff_ms: [100, 200, 400],
jitter: true
)
if RetryPolicy.should_retry?(policy, attempt) do
delay = RetryPolicy.backoff_for_attempt(policy, attempt)
Process.sleep(delay)
# retry...
end
"""
@type t :: %__MODULE__{
max_attempts: pos_integer(),
backoff_ms: [non_neg_integer()],
base_backoff_ms: non_neg_integer(),
backoff_multiplier: float(),
max_backoff_ms: non_neg_integer(),
jitter: boolean(),
jitter_factor: float(),
retriable_errors: [atom()] | :all
}
# Note: These struct defaults are compile-time constants.
# Runtime configurable defaults are handled in new/1 via Snakepit.Defaults
defstruct max_attempts: 3,
backoff_ms: [100, 200, 400, 800, 1600],
base_backoff_ms: 100,
backoff_multiplier: 2.0,
max_backoff_ms: 30_000,
jitter: false,
jitter_factor: 0.25,
retriable_errors: [:timeout, :unavailable, :connection_refused, :worker_crash]
alias Snakepit.Defaults
@doc """
Creates a new retry policy.
## Options
- `:max_attempts` - Maximum retry attempts (default: 3)
- `:backoff_ms` - List of backoff delays in ms (default: [100, 200, 400, 800, 1600])
- `:base_backoff_ms` - Base for exponential backoff (default: 100)
- `:backoff_multiplier` - Multiplier for exponential backoff (default: 2.0)
- `:max_backoff_ms` - Maximum backoff delay (default: 30000)
- `:jitter` - Add random jitter to delays (default: false)
- `:jitter_factor` - Jitter range as fraction of delay (default: 0.25)
- `:retriable_errors` - List of error atoms to retry, or :all (default: common errors)
"""
@spec new(keyword()) :: t()
def new(opts) do
# Use runtime-configurable defaults, then apply any user-provided options
defaults = [
max_attempts: Defaults.retry_max_attempts(),
backoff_ms: Defaults.retry_backoff_sequence(),
base_backoff_ms: Defaults.retry_base_backoff_ms(),
backoff_multiplier: Defaults.retry_backoff_multiplier(),
max_backoff_ms: Defaults.retry_max_backoff_ms(),
jitter_factor: Defaults.retry_jitter_factor()
]
merged = Keyword.merge(defaults, opts)
struct(__MODULE__, merged)
end
@doc """
Checks if another retry attempt should be made.
"""
@spec should_retry?(t(), non_neg_integer()) :: boolean()
def should_retry?(%__MODULE__{max_attempts: max}, attempt) do
attempt < max
end
@doc """
Checks if an error is retriable according to the policy.
"""
@spec retry_for_error?(t(), {:error, atom()} | term()) :: boolean()
def retry_for_error?(%__MODULE__{retriable_errors: :all}, _error), do: true
def retry_for_error?(%__MODULE__{retriable_errors: errors}, {:error, reason})
when is_atom(reason) do
reason in errors
end
def retry_for_error?(_policy, _error), do: false
@doc """
Returns the backoff delay for a given attempt.
"""
@spec backoff_for_attempt(t(), pos_integer()) :: non_neg_integer()
def backoff_for_attempt(%__MODULE__{} = policy, attempt) do
base_delay =
case Enum.at(policy.backoff_ms, attempt - 1) do
nil -> List.last(policy.backoff_ms) || policy.base_backoff_ms
delay -> delay
end
# Apply max cap
delay = min(base_delay, policy.max_backoff_ms)
# Apply jitter if enabled
if policy.jitter do
apply_jitter(delay, policy.jitter_factor)
else
delay
end
end
# Private functions
defp apply_jitter(delay, factor) do
jitter_range = trunc(delay * factor)
jitter = :rand.uniform(jitter_range * 2 + 1) - jitter_range - 1
max(0, delay + jitter)
end
end
defmodule Snakepit.RunID do
@moduledoc """
Generates short, unique BEAM run identifiers.
Format: 7 characters, base36-encoded
Components: timestamp (5 chars) + random (2 chars)
Example: "k3x9a2p"
These IDs are embedded in Python process command lines for reliable
identification and cleanup across BEAM restarts.
"""
@doc """
Generates a unique 7-character run ID.
## Examples
iex> run_id = Snakepit.RunID.generate()
iex> String.length(run_id)
7
"""
def generate do
# Use last 5 digits of microsecond timestamp (base36)
# Cycles every ~60 million seconds (~2 years)
timestamp = System.system_time(:microsecond)
time_component = timestamp |> rem(60_466_176) |> Integer.to_string(36) |> String.downcase()
time_part = String.pad_leading(time_component, 5, "0")
# Add 2 random characters for collision resistance
random_part =
:rand.uniform(1296)
|> Kernel.-(1)
|> Integer.to_string(36)
|> String.downcase()
|> String.pad_leading(2, "0")
time_part <> random_part
end
@doc """
Extracts run ID from a process command line.
Supports both --snakepit-run-id and --run-id formats.
## Examples
iex> cmd = "python3 grpc_server.py --snakepit-run-id k3x9a2p --port 50051"
iex> Snakepit.RunID.extract_from_command(cmd)
{:ok, "k3x9a2p"}
iex> Snakepit.RunID.extract_from_command("no run id here")
{:error, :not_found}
"""
def extract_from_command(command) when is_binary(command) do
# Try new format first (--run-id)
case Regex.run(~r/--run-id\s+([0-9a-z]{7})/, command) do
[_, run_id] ->
{:ok, run_id}
nil ->
# Try old format (--snakepit-run-id) with 7-char short IDs
case Regex.run(~r/--snakepit-run-id\s+([0-9a-z]{7})/, command) do
[_, run_id] -> {:ok, run_id}
nil -> {:error, :not_found}
end
end
end
def extract_from_command(_), do: {:error, :not_found}
end
defmodule Snakepit.RuntimeCleanup do
@moduledoc """
Deterministic shutdown cleanup for external worker processes.
This module performs a bounded cleanup pass:
- SIGTERM all known worker processes
- Wait until they exit or timeout
- Escalate to SIGKILL for survivors
"""
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.ProcessRegistry
alias Snakepit.ProcessKiller
@telemetry_start [:snakepit, :cleanup, :start]
@telemetry_success [:snakepit, :cleanup, :success]
@telemetry_timeout [:snakepit, :cleanup, :timeout]
@log_category :shutdown
def cleanup_current_run(opts \\ []) do
run_id = ProcessRegistry.get_beam_run_id()
entries = ProcessRegistry.current_run_entries()
run(entries, Keyword.put_new(opts, :run_id, run_id))
end
def run(entries, opts \\ []) when is_list(entries) do
timeout_ms = Keyword.get(opts, :timeout_ms, cleanup_on_stop_timeout_ms())
poll_interval_ms = Keyword.get(opts, :poll_interval_ms, cleanup_poll_interval_ms())
process_group_kill? = Keyword.get(opts, :process_group_kill, process_group_kill_enabled?())
run_id = Keyword.get(opts, :run_id)
targets =
entries
|> Enum.map(&to_target/1)
|> Enum.filter(fn target -> is_integer(target.process_pid) end)
if targets == [] do
:ok
else
emit_telemetry(@telemetry_start, length(targets), run_id)
kill_targets(targets, :sigterm, process_group_kill?)
remaining = wait_for_targets(targets, poll_interval_ms, timeout_ms)
if remaining == [] do
emit_telemetry(@telemetry_success, length(targets), run_id)
:ok
else
SLog.warning(
@log_category,
"Cleanup timeout after #{timeout_ms}ms; escalating to SIGKILL for #{length(remaining)} processes"
)
kill_targets(remaining, :sigkill, process_group_kill?)
still_alive = wait_for_targets(remaining, poll_interval_ms, poll_interval_ms)
log_incomplete_cleanup(still_alive)
emit_telemetry(@telemetry_timeout, length(remaining), run_id)
{:timeout, still_alive}
end
end
end
defp to_target({worker_id, info}) when is_map(info) do
%{
worker_id: worker_id,
process_pid: Map.get(info, :process_pid),
pgid: Map.get(info, :pgid),
process_group?: Map.get(info, :process_group?, false)
}
end
defp to_target(info) when is_map(info) do
%{
worker_id: Map.get(info, :worker_id),
process_pid: Map.get(info, :process_pid),
pgid: Map.get(info, :pgid),
process_group?: Map.get(info, :process_group?, false)
}
end
defp kill_targets(targets, signal, process_group_kill?) do
Enum.each(targets, fn target ->
case kill_target(target, signal, process_group_kill?) do
:ok ->
:ok
{:error, reason} ->
SLog.warning(
@log_category,
"Failed to send #{signal} to #{format_target(target)}: #{inspect(reason)}"
)
end
end)
end
defp kill_target(%{process_group?: true, pgid: pgid}, signal, true)
when is_integer(pgid) do
ProcessKiller.kill_process_group(pgid, signal)
end
defp kill_target(%{process_pid: pid}, signal, _process_group_kill?) when is_integer(pid) do
ProcessKiller.kill_process(pid, signal)
end
defp kill_target(_target, _signal, _process_group_kill?), do: :ok
defp wait_for_targets(targets, poll_interval_ms, timeout_ms) do
deadline = System.monotonic_time(:millisecond) + timeout_ms
wait_for_targets_loop(targets, poll_interval_ms, deadline)
end
defp wait_for_targets_loop(targets, poll_interval_ms, deadline) do
remaining =
Enum.filter(targets, fn target ->
ProcessKiller.process_alive?(target.process_pid)
end)
cond do
remaining == [] ->
[]
System.monotonic_time(:millisecond) >= deadline ->
remaining
true ->
receive do
after
poll_interval_ms -> :ok
end
wait_for_targets_loop(remaining, poll_interval_ms, deadline)
end
end
defp emit_telemetry(event, count, run_id) do
:telemetry.execute(
event,
%{count: count},
%{beam_run_id: run_id, timestamp: System.system_time(:second)}
)
end
defp cleanup_on_stop_timeout_ms do
Application.get_env(:snakepit, :cleanup_on_stop_timeout_ms, 3_000)
end
defp cleanup_poll_interval_ms do
Application.get_env(:snakepit, :cleanup_poll_interval_ms, 50)
end
defp process_group_kill_enabled? do
Application.get_env(:snakepit, :process_group_kill, true)
end
defp log_incomplete_cleanup([]), do: :ok
defp log_incomplete_cleanup(still_alive) do
SLog.warning(
@log_category,
"Cleanup incomplete: #{length(still_alive)} processes still alive after SIGKILL"
)
end
defp format_target(%{worker_id: worker_id, process_pid: pid})
when is_binary(worker_id) and is_integer(pid) do
"#{worker_id} (pid #{pid})"
end
defp format_target(%{process_pid: pid}) when is_integer(pid), do: "pid #{pid}"
defp format_target(_), do: "unknown target"
end
# Snakepit Supervision Tree
The following Mermaid diagram captures the runtime supervision tree that
`Snakepit.Application` boots. Two sets of children exist:
- **Base services** always start (session store + tool registry).
- **Pooling branch** activates when `:pooling_enabled` is true, which is the
typical production configuration.
```mermaid
flowchart TD
A["Snakepit.Supervisor\n(strategy: one_for_one)"]
A --> B["Snakepit.Bridge.SessionStore\n(GenServer)"]
A --> C["Snakepit.Bridge.ToolRegistry\n(GenServer)"]
subgraph "Pooling Enabled (:pooling_enabled == true)"
direction TB
A --> D["GRPC.Server.Supervisor\nSnakepit.GRPC.Endpoint"]
D --> D1["Cowboy/Ranch Acceptors"]
A --> E["Task.Supervisor\n(Snakepit.TaskSupervisor)"]
A --> F["Snakepit.Pool.Registry\n(Registry)"]
A --> G["Snakepit.Pool.Worker.StarterRegistry\n(Registry)"]
A --> H["Snakepit.Pool.ProcessRegistry\n(GenServer + ETS)"]
A --> I["Snakepit.Pool.WorkerSupervisor\n(DynamicSupervisor)"]
A --> J["Snakepit.Worker.LifecycleManager\n(GenServer)"]
A --> K["Snakepit.Pool\n(GenServer)\nrequest routing"]
A --> L["Snakepit.Pool.ApplicationCleanup\n(GenServer)\nfirst to stop"]
end
subgraph "Pooling Disabled (:pooling_enabled == false)"
direction TB
A --> M["(no additional children)\nuseful for tests"]
end
```
**Notes**
- `GRPC.Server.Supervisor` encapsulates `Snakepit.GRPC.Endpoint` which exposes
the Elixir bridge services.
- `Snakepit.Pool.WorkerSupervisor` dynamically supervises worker GenServers
(`Snakepit.GRPCWorker` et al.) under a `:one_for_one` strategy.
- `Snakepit.Pool.ProcessRegistry` tracks external OS PIDs and run IDs to ensure
cleanup routines know which processes belong to the current BEAM instance.
- `Snakepit.Pool.ApplicationCleanup` is intentionally listed last so it
terminates first and can reap external processes before other supervisors
shut down.
defmodule Snakepit.TelemetryMetrics do
@moduledoc """
Telemetry metric definitions and reporters for Snakepit.
Metrics focus on heartbeat and worker lifecycle events. Reporters are opt-in
via configuration under `:snakepit, :telemetry_metrics`.
"""
import Telemetry.Metrics
@type reporter_child_spec :: Supervisor.child_spec()
@default_config %{
prometheus: %{
enabled: false,
port: 9568,
name: :snakepit_prometheus_metrics,
protocol: :http
}
}
@doc """
Returns the metric definitions for Snakepit telemetry.
"""
@spec metrics() :: [Telemetry.Metrics.t()]
def metrics do
[
counter("snakepit.heartbeat.pings",
event_name: [:snakepit, :heartbeat, :ping_sent],
measurement: :pings,
tags: [:worker_id]
),
counter("snakepit.heartbeat.pongs",
event_name: [:snakepit, :heartbeat, :pong_received],
measurement: :pongs,
tags: [:worker_id]
),
counter("snakepit.heartbeat.failures",
event_name: [:snakepit, :heartbeat, :monitor_failure],
measurement: :failures,
tags: [:worker_id, :failure_reason]
),
summary("snakepit.heartbeat.latency",
event_name: [:snakepit, :heartbeat, :pong_received],
measurement: :latency_ms,
unit: :millisecond,
tags: [:worker_id]
),
last_value("snakepit.heartbeat.missed",
event_name: [:snakepit, :heartbeat, :heartbeat_timeout],
measurement: :missed,
tags: [:worker_id]
),
counter("snakepit.grpc.worker.executions",
event_name: [:snakepit, :grpc_worker, :execute, :stop],
measurement: :executions,
tags: [:worker_id, :command]
),
counter("snakepit.grpc.worker.errors",
event_name: [:snakepit, :grpc_worker, :execute, :stop],
measurement: :errors,
tags: [:worker_id, :command, :error]
),
summary("snakepit.grpc.worker.duration",
event_name: [:snakepit, :grpc_worker, :execute, :stop],
measurement: :duration_ms,
unit: :millisecond,
tags: [:worker_id, :command]
),
counter("snakepit.worker.recycled",
event_name: [:snakepit, :worker, :recycled],
measurement: :count,
tags: [:pool, :reason]
),
summary("snakepit.worker.memory_mb",
event_name: [:snakepit, :worker, :recycled],
measurement: &Map.get(&1, :memory_mb),
unit: :megabyte,
tags: [:pool]
)
]
end
@doc """
Returns reporter child specs enabled via configuration.
"""
@spec reporter_children() :: [reporter_child_spec()]
def reporter_children do
Enum.flat_map([prometheus_child_spec()], fn
nil -> []
spec -> [spec]
end)
end
defp prometheus_child_spec do
%{prometheus: prometheus_config} = load_config()
if truthy?(prometheus_config[:enabled]) do
TelemetryMetricsPrometheus.child_spec(
Keyword.merge(
[
metrics: metrics(),
port: fetch_integer(prometheus_config[:port], 9568),
name: prometheus_config[:name] || :snakepit_prometheus_metrics,
protocol: prometheus_config[:protocol] || :http
],
prometheus_extra_options(prometheus_config)
)
)
end
end
defp prometheus_extra_options(config) do
opts = []
opts =
case config[:ip] do
nil -> opts
ip when is_tuple(ip) -> Keyword.put(opts, :plug_cowboy_opts, ip: ip)
ip when is_binary(ip) -> Keyword.put(opts, :plug_cowboy_opts, ip: parse_ip(ip))
_ -> opts
end
opts
end
defp parse_ip(ip_string) do
case :inet.parse_address(String.to_charlist(ip_string)) do
{:ok, tuple} -> tuple
{:error, _reason} -> {0, 0, 0, 0}
end
end
defp truthy?(value) when is_boolean(value), do: value
defp truthy?(value) when is_binary(value) do
normalized = String.downcase(String.trim(value))
normalized in ["true", "1", "yes", "on"]
end
defp truthy?(value) when is_integer(value), do: value != 0
defp truthy?(true), do: true
defp truthy?(_), do: false
defp fetch_integer(nil, default), do: default
defp fetch_integer(value, default) when is_binary(value) do
case Integer.parse(value) do
{int, _rest} -> int
:error -> default
end
end
defp fetch_integer(value, _default) when is_integer(value), do: value
defp fetch_integer(_, default), do: default
defp load_config do
base_config =
Application.get_env(:snakepit, :telemetry_metrics, %{})
|> to_map()
Map.merge(@default_config, base_config, fn
_key, default_value, user_value when is_map(default_value) and is_map(user_value) ->
Map.merge(default_value, user_value)
_key, _default_value, user_value ->
user_value
end)
end
defp to_map(value) when is_map(value), do: value
defp to_map(value) when is_list(value) do
Enum.into(value, %{}, fn
{key, val} when is_atom(key) -> {key, to_map(val)}
{key, val} -> {String.to_atom(to_string(key)), to_map(val)}
end)
end
defp to_map(other), do: other
end
defmodule Snakepit.Telemetry do
@moduledoc """
Telemetry event definitions for Snakepit.
This module provides:
- Complete event catalog (Layer 1: Infrastructure, Layer 2: Python, Layer 3: gRPC)
- Event handler management
- Integration with the distributed telemetry system
See `Snakepit.Telemetry.Naming` for event name validation and atom safety.
See `Snakepit.Telemetry.GrpcStream` for Python telemetry folding.
## Usage
# Attach handlers to specific events
:telemetry.attach(
"my-handler",
[:snakepit, :python, :call, :stop],
&MyApp.Telemetry.handle_python_call/4,
nil
)
# Emit a pool event
:telemetry.execute(
[:snakepit, :pool, :worker, :spawned],
%{duration: 1000, system_time: System.system_time()},
%{node: node(), worker_id: "worker_1", pool_name: :default}
)
"""
alias Snakepit.Logger, as: SLog
@log_category :telemetry
@doc """
Lists all telemetry events used by Snakepit.
"""
def events do
session_events() ++
program_events() ++
heartbeat_events() ++
pool_events() ++ python_events() ++ grpc_events() ++ runtime_events()
end
## Layer 0: Session Store & Heartbeat (Legacy)
@doc """
Session-related telemetry events (session store).
"""
def session_events do
[
[:snakepit, :session_store, :session, :created],
[:snakepit, :session_store, :session, :accessed],
[:snakepit, :session_store, :session, :deleted],
[:snakepit, :session_store, :session, :expired]
]
end
@doc """
Program-related telemetry events (session store).
"""
def program_events do
[
[:snakepit, :session_store, :program, :stored],
[:snakepit, :session_store, :program, :retrieved],
[:snakepit, :session_store, :program, :deleted]
]
end
@doc """
Heartbeat and monitor telemetry events.
"""
def heartbeat_events do
[
[:snakepit, :heartbeat, :monitor_started],
[:snakepit, :heartbeat, :monitor_stopped],
[:snakepit, :heartbeat, :monitor_failure],
[:snakepit, :heartbeat, :ping_sent],
[:snakepit, :heartbeat, :pong_received],
[:snakepit, :heartbeat, :heartbeat_timeout]
]
end
## Layer 1: Infrastructure Events (Pool, Worker, Session)
@doc """
Pool and worker lifecycle events.
"""
def pool_events do
[
[:snakepit, :pool, :initialized],
[:snakepit, :pool, :status],
[:snakepit, :pool, :queue, :enqueued],
[:snakepit, :pool, :queue, :dequeued],
[:snakepit, :pool, :queue, :timeout],
[:snakepit, :pool, :worker, :spawn_started],
[:snakepit, :pool, :worker, :spawned],
[:snakepit, :pool, :worker, :spawn_failed],
[:snakepit, :pool, :worker, :terminated],
[:snakepit, :pool, :worker, :restarted],
[:snakepit, :worker, :recycled],
[:snakepit, :session, :created],
[:snakepit, :session, :destroyed],
[:snakepit, :session, :affinity, :assigned],
[:snakepit, :session, :affinity, :broken]
]
end
## Layer 2: Python Execution Events (Folded from Python)
@doc """
Python worker telemetry events (folded back from Python workers).
"""
def python_events do
[
[:snakepit, :python, :call, :start],
[:snakepit, :python, :call, :stop],
[:snakepit, :python, :call, :exception],
[:snakepit, :python, :memory, :sampled],
[:snakepit, :python, :cpu, :sampled],
[:snakepit, :python, :gc, :completed],
[:snakepit, :python, :error, :occurred],
[:snakepit, :python, :tool, :execution, :start],
[:snakepit, :python, :tool, :execution, :stop],
[:snakepit, :python, :tool, :execution, :exception],
[:snakepit, :python, :tool, :result_size]
]
end
## Layer 3: gRPC Bridge Events
@doc """
gRPC communication events.
"""
def grpc_events do
[
[:snakepit, :grpc, :call, :start],
[:snakepit, :grpc, :call, :stop],
[:snakepit, :grpc, :call, :exception],
[:snakepit, :grpc, :stream, :opened],
[:snakepit, :grpc, :stream, :message],
[:snakepit, :grpc, :stream, :closed],
[:snakepit, :grpc, :connection, :established],
[:snakepit, :grpc, :connection, :lost],
[:snakepit, :grpc, :connection, :reconnected]
]
end
@doc """
Runtime enhancement events (zero-copy, crash barrier, exception translation).
"""
def runtime_events do
[
[:snakepit, :zero_copy, :export],
[:snakepit, :zero_copy, :import],
[:snakepit, :zero_copy, :fallback],
[:snakepit, :worker, :crash],
[:snakepit, :worker, :tainted],
[:snakepit, :worker, :restarted],
[:snakepit, :python, :exception, :mapped],
[:snakepit, :python, :exception, :unmapped]
]
end
@doc """
Attaches default handlers for all events.
"""
def attach_handlers do
attach_session_handlers()
attach_program_handlers()
attach_heartbeat_handlers()
end
@doc """
Attaches default handlers for session events.
"""
def attach_session_handlers do
:telemetry.attach_many(
"snakepit-session-logger",
session_events(),
&handle_event/4,
nil
)
end
@doc """
Attaches default handlers for program events.
"""
def attach_program_handlers do
:telemetry.attach_many(
"snakepit-program-logger",
program_events(),
&handle_event/4,
nil
)
end
@doc """
Attaches default handlers for heartbeat events.
"""
def attach_heartbeat_handlers do
:telemetry.attach_many(
"snakepit-heartbeat-logger",
heartbeat_events(),
&handle_event/4,
nil
)
end
# Event handlers
# Session event handlers
defp handle_event([:snakepit, :session_store, :session, :created], _measurements, metadata, _) do
SLog.info(@log_category, "Session created: #{metadata.session_id}")
end
defp handle_event([:snakepit, :session_store, :session, :accessed], _measurements, metadata, _) do
SLog.debug(@log_category, "Session accessed: #{metadata.session_id}")
end
defp handle_event([:snakepit, :session_store, :session, :deleted], _measurements, metadata, _) do
SLog.info(@log_category, "Session deleted: #{metadata.session_id}")
end
defp handle_event([:snakepit, :session_store, :session, :expired], measurements, _metadata, _) do
SLog.info(@log_category, "Sessions expired: count=#{measurements.count}")
end
# Program event handlers
defp handle_event([:snakepit, :session_store, :program, :stored], _measurements, metadata, _) do
SLog.debug(
@log_category,
"Program stored: #{metadata.program_id} in session #{metadata.session_id}"
)
end
defp handle_event([:snakepit, :session_store, :program, :retrieved], _measurements, metadata, _) do
SLog.debug(
@log_category,
"Program retrieved: #{metadata.program_id} from session #{metadata.session_id}"
)
end
defp handle_event([:snakepit, :session_store, :program, :deleted], _measurements, metadata, _) do
SLog.debug(
@log_category,
"Program deleted: #{metadata.program_id} from session #{metadata.session_id}"
)
end
# Heartbeat events
defp handle_event([:snakepit, :heartbeat, :monitor_started], _measurements, metadata, _) do
SLog.debug(@log_category, "Heartbeat monitor started for #{metadata.worker_id}")
end
defp handle_event([:snakepit, :heartbeat, :monitor_stopped], _measurements, metadata, _) do
SLog.debug(
@log_category,
"Heartbeat monitor stopped for #{metadata.worker_id} reason=#{inspect(metadata.reason)}"
)
end
defp handle_event([:snakepit, :heartbeat, :monitor_failure], _measurements, metadata, _) do
SLog.warning(
@log_category,
"Heartbeat monitor triggered failure for #{metadata.worker_id}: #{inspect(metadata.failure_reason)}"
)
end
defp handle_event([:snakepit, :heartbeat, :ping_sent], measurements, metadata, _) do
SLog.debug(
@log_category,
"Heartbeat ping sent for #{metadata.worker_id} (count=#{measurements[:count]})"
)
end
defp handle_event([:snakepit, :heartbeat, :pong_received], measurements, metadata, _) do
SLog.debug(
@log_category,
"Heartbeat pong received for #{metadata.worker_id} latency=#{measurements[:latency_ms]}ms"
)
end
defp handle_event([:snakepit, :heartbeat, :heartbeat_timeout], measurements, metadata, _) do
SLog.warning(
@log_category,
"Heartbeat timeout for #{metadata.worker_id} missed=#{measurements[:count]}"
)
end
# Catch-all handler for any unhandled events
defp handle_event(event, measurements, metadata, _) do
SLog.debug(
@log_category,
"Telemetry event: #{inspect(event)} measurements=#{inspect(measurements)} metadata=#{inspect(metadata)}"
)
end
end
defmodule Snakepit.WorkerProfile do
@moduledoc """
Behaviour for worker profiles (process vs thread).
A worker profile defines how workers are created, managed, and utilized.
Snakepit v0.6.0 introduces dual-mode parallelism:
## Process Profile (`:process`)
- Many single-threaded Python processes
- Process isolation and GIL compatibility
- Optimal for: I/O-bound workloads, high concurrency, legacy Python
## Thread Profile (`:thread`)
- Few multi-threaded Python processes
- Shared memory and CPU parallelism
- Optimal for: CPU-bound workloads, Python 3.13+, large data
## Implementing a Profile
Profiles control the full worker lifecycle:
defmodule MyProfile do
@behaviour Snakepit.WorkerProfile
def start_worker(config) do
# Start worker according to profile
{:ok, worker_handle}
end
def get_capacity(worker_handle) do
# Return concurrent request capacity
1 # or N for multi-threaded
end
end
See `Snakepit.WorkerProfile.Process` and `Snakepit.WorkerProfile.Thread` for reference implementations.
"""
@type worker_handle :: pid() | reference()
@type config :: map()
@type capacity :: pos_integer()
@doc """
Start a worker with the given configuration.
Returns `{:ok, worker_handle}` where worker_handle is typically a GenServer PID,
or `{:error, reason}` if startup fails.
The config map contains all pool and adapter configuration for this worker.
"""
@callback start_worker(config) :: {:ok, worker_handle} | {:error, term()}
@doc """
Stop a worker gracefully.
Should perform cleanup and shutdown the worker process.
"""
@callback stop_worker(worker_handle) :: :ok
@doc """
Execute a request on a worker.
For process-based workers, this typically blocks until the request completes.
For thread-based workers, this may execute concurrently with other requests
on the same worker.
The timeout is in milliseconds.
"""
@callback execute_request(worker_handle, request :: map(), timeout :: timeout()) ::
{:ok, term()} | {:error, term()}
@doc """
Get the maximum capacity of a worker (how many concurrent requests it can handle).
- Process profile: returns 1 (single-threaded)
- Thread profile: returns N (thread pool size)
This is used by the pool for load balancing decisions.
"""
@callback get_capacity(worker_handle) :: capacity()
@doc """
Get the current load of a worker (how many requests are currently in-flight).
Returns 0 if no requests are active, up to the worker's capacity.
"""
@callback get_load(worker_handle) :: non_neg_integer()
@doc """
Check if a worker is healthy and responsive.
Returns `:ok` if healthy, `{:error, reason}` if unhealthy.
"""
@callback health_check(worker_handle) :: :ok | {:error, term()}
@doc """
Get profile-specific metadata about a worker.
Optional callback. Returns a map with profile-specific information.
"""
@callback get_metadata(worker_handle) :: {:ok, map()} | {:error, term()}
@optional_callbacks [get_metadata: 1]
end
defmodule Snakepit.ZeroCopyRef do
@moduledoc """
Opaque handle for zero-copy payloads.
The handle metadata travels through the runtime so adapters can resolve
DLPack or Arrow buffers without copying.
"""
@type kind :: :dlpack | :arrow
@type device :: :cpu | :cuda | :mps
@enforce_keys [:kind, :ref]
defstruct [
:kind,
:device,
:dtype,
:shape,
:owner,
:ref,
:copy,
:bytes,
:metadata
]
@type t :: %__MODULE__{
kind: kind(),
device: device() | nil,
dtype: atom() | String.t() | nil,
shape: tuple() | list() | nil,
owner: :elixir | :python | nil,
ref: reference(),
copy: boolean() | nil,
bytes: non_neg_integer() | nil,
metadata: map() | nil
}
@doc false
def to_map(%__MODULE__{} = ref) do
%{
"__snakepit_zero_copy__" => true,
"kind" => Atom.to_string(ref.kind),
"device" => maybe_string(ref.device),
"dtype" => maybe_string(ref.dtype),
"shape" => ref.shape,
"owner" => maybe_string(ref.owner),
"ref" => encode_ref(ref.ref),
"copy" => ref.copy,
"bytes" => ref.bytes,
"metadata" => ref.metadata || %{}
}
|> Enum.reject(fn {_key, value} -> is_nil(value) end)
|> Map.new()
end
@doc false
def maybe_from_map(%{"__snakepit_zero_copy__" => true} = map), do: from_map(map)
def maybe_from_map(%{__snakepit_zero_copy__: true} = map), do: from_map(map)
def maybe_from_map(other), do: other
@doc false
def from_map(map) when is_map(map) do
ref =
map
|> fetch_value([:ref, "ref"])
|> decode_ref()
%__MODULE__{
kind: map |> fetch_value([:kind, "kind"]) |> to_kind(),
device: map |> fetch_value([:device, "device"]) |> to_device(),
dtype: fetch_value(map, [:dtype, "dtype"]),
shape: fetch_value(map, [:shape, "shape"]),
owner: map |> fetch_value([:owner, "owner"]) |> to_owner(),
ref: ref,
copy: fetch_value(map, [:copy, "copy"]),
bytes: fetch_value(map, [:bytes, "bytes"]),
metadata: fetch_value(map, [:metadata, "metadata"])
}
end
defp fetch_value(map, keys) when is_list(keys) do
Enum.find_value(keys, fn key -> Map.get(map, key) end)
end
defp encode_ref(ref) when is_reference(ref) do
ref
|> :erlang.term_to_binary()
|> Base.encode64()
end
defp decode_ref(ref) when is_reference(ref), do: ref
defp decode_ref(ref) when is_binary(ref) do
ref
|> Base.decode64!()
|> :erlang.binary_to_term()
end
defp decode_ref(_), do: make_ref()
defp to_kind(kind) when is_atom(kind), do: kind
defp to_kind(kind) when is_binary(kind) do
case String.downcase(kind) do
"dlpack" -> :dlpack
"arrow" -> :arrow
_ -> :dlpack
end
end
defp to_kind(_), do: :dlpack
defp to_device(device) when is_atom(device), do: device
defp to_device(device) when is_binary(device) do
case String.downcase(device) do
"cpu" -> :cpu
"cuda" -> :cuda
"mps" -> :mps
_ -> nil
end
end
defp to_device(_), do: nil
defp to_owner(owner) when is_atom(owner), do: owner
defp to_owner(owner) when is_binary(owner) do
case String.downcase(owner) do
"elixir" -> :elixir
"python" -> :python
_ -> nil
end
end
defp to_owner(_), do: nil
defp maybe_string(nil), do: nil
defp maybe_string(value) when is_atom(value), do: Atom.to_string(value)
defp maybe_string(value), do: to_string(value)
end
defimpl Jason.Encoder, for: Snakepit.ZeroCopyRef do
def encode(ref, opts) do
ref
|> Snakepit.ZeroCopyRef.to_map()
|> Jason.Encode.map(opts)
end
end
defmodule Snakepit.ZeroCopy do
@moduledoc """
Zero-copy interop helpers for DLPack and Arrow.
Handles create/import lifecycle for zero-copy handles, with copy-based
fallbacks when unavailable.
"""
alias Snakepit.Logger, as: SLog
alias Snakepit.ZeroCopyRef
@table :snakepit_zero_copy_handles
@log_category :bridge
@default_config %{
enabled: false,
dlpack: true,
arrow: true,
allow_fallback: true,
max_bytes: 64 * 1024 * 1024,
strict: false
}
@type export_opts ::
[
device: ZeroCopyRef.device(),
dtype: atom() | String.t(),
shape: tuple() | list(),
owner: :elixir | :python
]
def to_dlpack(term, opts \\ []), do: export(:dlpack, term, opts)
def to_arrow(term, opts \\ []), do: export(:arrow, term, opts)
def from_dlpack(ref, opts \\ []), do: import(:dlpack, ref, opts)
def from_arrow(ref, opts \\ []), do: import(:arrow, ref, opts)
def close(ref) do
ref = normalize_ref(ref)
ensure_table()
:ets.delete(@table, ref.ref)
:ok
end
defp export(kind, term, opts) when kind in [:dlpack, :arrow] do
start = System.monotonic_time(:microsecond)
config = config()
meta = build_metadata(kind, term, opts)
cond do
zero_copy_enabled?(config, kind) ->
ref = store_handle(kind, term, meta, copy?: false)
emit(:export, start, ref)
{:ok, ref}
config.allow_fallback ->
SLog.warning(
@log_category,
"Zero-copy #{kind} unavailable; falling back to copy (reason: :zero_copy_unavailable)"
)
ref = store_handle(kind, copy_term(term), meta, copy?: true)
emit(:fallback, start, ref)
{:ok, ref}
config.strict ->
{:error, :zero_copy_unavailable}
true ->
{:error, :zero_copy_unavailable}
end
end
defp import(kind, ref, _opts) when kind in [:dlpack, :arrow] do
start = System.monotonic_time(:microsecond)
ref = normalize_ref(ref)
ensure_table()
case :ets.lookup(@table, ref.ref) do
[{_ref, stored_kind, value, metadata}] when stored_kind == kind ->
emit(:import, start, Map.merge(ref, metadata))
{:ok, value}
[{_ref, stored_kind, _value, _metadata}] ->
{:error, {:invalid_zero_copy_kind, stored_kind}}
_ ->
{:error, :zero_copy_handle_not_found}
end
end
defp store_handle(kind, term, metadata, opts) do
ensure_table()
ref = make_ref()
entry = {ref, kind, term, metadata}
true = :ets.insert(@table, entry)
%ZeroCopyRef{
kind: kind,
device: metadata.device,
dtype: metadata.dtype,
shape: metadata.shape,
owner: metadata.owner,
ref: ref,
copy: Keyword.get(opts, :copy?, false),
bytes: metadata.bytes,
metadata: metadata.extra
}
end
defp build_metadata(kind, term, opts) do
%{
kind: kind,
device: Keyword.get(opts, :device, :cpu),
dtype: Keyword.get(opts, :dtype),
shape: Keyword.get(opts, :shape),
owner: Keyword.get(opts, :owner, :elixir),
bytes: term_bytes(term),
extra: Keyword.get(opts, :metadata, %{})
}
end
defp emit(event, start_us, %ZeroCopyRef{} = ref) do
duration_us = System.monotonic_time(:microsecond) - start_us
measurements =
%{duration_ms: duration_us / 1000}
|> maybe_put(:bytes, ref.bytes)
metadata = %{
kind: ref.kind,
device: ref.device,
dtype: ref.dtype,
shape: ref.shape
}
:telemetry.execute([:snakepit, :zero_copy, event], measurements, metadata)
end
defp maybe_put(map, _key, nil), do: map
defp maybe_put(map, key, value), do: Map.put(map, key, value)
defp copy_term(term) do
term
|> :erlang.term_to_binary()
|> :erlang.binary_to_term()
end
defp term_bytes(term) when is_binary(term), do: byte_size(term)
defp term_bytes(term) do
if function_exported?(:erlang, :external_size, 1) do
:erlang.external_size(term)
else
nil
end
end
defp zero_copy_enabled?(config, :dlpack),
do: config.enabled && config.dlpack
defp zero_copy_enabled?(config, :arrow),
do: config.enabled && config.arrow
defp config do
config =
:snakepit
|> Application.get_env(:zero_copy, [])
|> Map.new()
Map.merge(@default_config, config)
end
defp normalize_ref(%ZeroCopyRef{} = ref), do: ref
defp normalize_ref(map) when is_map(map), do: ZeroCopyRef.from_map(map)
defp normalize_ref(_), do: %ZeroCopyRef{kind: :dlpack, ref: make_ref()}
defp ensure_table do
case :ets.whereis(@table) do
:undefined ->
:ets.new(@table, [:named_table, :set, :public, {:read_concurrency, true}])
_ ->
@table
end
end
end
defmodule Snakepit do
@moduledoc """
Snakepit - A generalized high-performance pooler and session manager.
Extracted from DSPex V3 pool implementation, Snakepit provides:
- Concurrent worker initialization and management
- Stateless pool system with session affinity
- Generalized adapter pattern for any external process
- High-performance OTP-based process management
## Basic Usage
# Configure in config/config.exs
config :snakepit,
pooling_enabled: true,
adapter_module: YourAdapter
# Execute commands on any available worker
{:ok, result} = Snakepit.execute("ping", %{test: true})
# Session-based execution with worker affinity
{:ok, result} = Snakepit.execute_in_session("my_session", "command", %{})
"""
alias Snakepit.Logger, as: SLog
alias Snakepit.Pool.ProcessRegistry
# Type definitions
@type command :: String.t()
@type args :: map()
@type result :: term()
@type session_id :: String.t()
@type callback_fn :: (term() -> any())
@type pool_name :: atom() | pid()
@doc """
Convenience function to execute commands on the pool.
## Examples
{:ok, result} = Snakepit.execute("ping", %{test: true})
## Options
* `:pool` - The pool to use (default: `Snakepit.Pool`)
* `:timeout` - Request timeout in ms (default: 60000)
* `:session_id` - Execute with session affinity
"""
@spec execute(command(), args(), keyword()) :: {:ok, result()} | {:error, Snakepit.Error.t()}
def execute(command, args, opts \\ []) do
Snakepit.Pool.execute(command, args, opts)
end
@doc """
Executes a command in session context with worker affinity.
This function executes commands with session-based worker affinity,
ensuring that subsequent calls with the same session_id prefer
the same worker when possible for state continuity.
Args are passed through unchanged - no domain-specific enhancement.
"""
@spec execute_in_session(session_id(), command(), args(), keyword()) ::
{:ok, result()} | {:error, Snakepit.Error.t()}
def execute_in_session(session_id, command, args, opts \\ []) do
# Add session_id to opts for session affinity
opts_with_session = Keyword.put(opts, :session_id, session_id)
# Execute command with session affinity (no args enhancement)
execute(command, args, opts_with_session)
end
@doc """
Get pool statistics.
Returns aggregate stats across all pools or stats for a specific pool.
"""
@spec get_stats(pool_name()) :: map()
def get_stats(pool \\ Snakepit.Pool) do
Snakepit.Pool.get_stats(pool)
end
@doc """
List workers from the pool.
Returns a list of worker IDs.
"""
@spec list_workers(pool_name()) :: [String.t()]
def list_workers(pool \\ Snakepit.Pool) do
Snakepit.Pool.list_workers(pool)
end
@doc """
Executes a streaming command with a callback function.
## Examples
Snakepit.execute_stream("batch_inference", %{items: [...]}, fn chunk ->
handle_chunk(chunk)
end)
## Options
* `:pool` - The pool to use (default: `Snakepit.Pool`)
* `:timeout` - Request timeout in ms (default: 300000)
* `:session_id` - Run in a specific session
## Returns
Returns `:ok` on success or `{:error, %Snakepit.Error{}}` on failure.
Note: Streaming is only supported with gRPC adapters.
"""
@spec execute_stream(command(), args(), callback_fn(), keyword()) ::
:ok | {:error, Snakepit.Error.t()}
def execute_stream(command, args \\ %{}, callback_fn, opts \\ []) do
ensure_started!()
adapter = Application.get_env(:snakepit, :adapter_module)
if function_exported?(adapter, :uses_grpc?, 0) and adapter.uses_grpc?() do
Snakepit.Pool.execute_stream(command, args, callback_fn, opts)
else
{:error,
Snakepit.Error.validation_error("Streaming not supported by adapter", %{
adapter: adapter
})}
end
end
@doc """
Manually trigger cleanup of external worker processes for the current run.
Useful for library embedding or scripts that control the lifecycle directly.
"""
@spec cleanup() :: :ok | {:timeout, list()}
def cleanup do
if Process.whereis(Snakepit.Pool.ProcessRegistry) do
Snakepit.RuntimeCleanup.cleanup_current_run()
else
:ok
end
end
@doc """
Executes a command in a session with a callback function.
"""
@spec execute_in_session_stream(session_id(), command(), args(), callback_fn(), keyword()) ::
:ok | {:error, Snakepit.Error.t()}
def execute_in_session_stream(session_id, command, args \\ %{}, callback_fn, opts \\ []) do
ensure_started!()
adapter = Application.get_env(:snakepit, :adapter_module)
if function_exported?(adapter, :uses_grpc?, 0) and adapter.uses_grpc?() do
opts_with_session = Keyword.put(opts, :session_id, session_id)
Snakepit.Pool.execute_stream(command, args, callback_fn, opts_with_session)
else
{:error,
Snakepit.Error.validation_error("Streaming not supported by adapter", %{
adapter: adapter
})}
end
end
defp ensure_started! do
case Application.ensure_all_started(:snakepit) do
{:ok, _} -> :ok
{:error, _} -> raise "Snakepit application not started"
end
end
@doc """
Starts the Snakepit application, executes a given function,
and ensures graceful shutdown.
This is the recommended way to use Snakepit for short-lived scripts or
Mix tasks to prevent orphaned processes.
It handles the full OTP application lifecycle (start, run, stop)
automatically.
## Examples
# In a Mix task
Snakepit.run_as_script(fn ->
{:ok, result} = Snakepit.execute("my_command", %{data: "value"})
handle_result(result)
end)
# For demos or scripts
Snakepit.run_as_script(fn ->
MyApp.run_load_test()
end)
## Options
* `:timeout` - Maximum time to wait for pool initialization (default: 15000ms)
* `:shutdown_timeout` - Time to wait for supervisor shutdown confirmation (default: 15000ms)
* `:cleanup_timeout` - Time to wait for worker process cleanup before forcing cleanup (default: 5000ms)
(cleanup is bounded; if it exceeds `cleanup_timeout + 1000` ms the script continues)
* `:restart` - Restart Snakepit if already started to apply script config (`:auto` | true | false)
* `:await_pool` - Wait for pool readiness (default: `pooling_enabled` setting)
* `:halt` - Force `System.halt/1` after cleanup for scripts that must exit (default: false,
or set `SNAKEPIT_SCRIPT_HALT=true`)
## Returns
Returns the result of the provided function, or `{:error, reason}` if
the pool fails to initialize.
"""
@spec run_as_script((-> any()), keyword()) :: any() | {:error, term()}
def run_as_script(fun, opts \\ []) when is_function(fun, 0) do
startup_timeout = Keyword.get(opts, :timeout, 15_000)
shutdown_timeout = Keyword.get(opts, :shutdown_timeout, 15_000)
cleanup_timeout = Keyword.get(opts, :cleanup_timeout, 5_000)
restart = Keyword.get(opts, :restart, :auto)
await_pool = Keyword.get(opts, :await_pool, pooling_enabled?())
halt = Keyword.get(opts, :halt, env_truthy?("SNAKEPIT_SCRIPT_HALT"))
# Ensure all dependencies are started, including Snakepit itself
maybe_restart_snakepit(restart, shutdown_timeout, cleanup_timeout)
{:ok, _apps} = Application.ensure_all_started(:snakepit)
# Deterministically wait for the pool to be fully initialized
startup_result =
if await_pool do
Snakepit.Pool.await_ready(Snakepit.Pool, startup_timeout)
else
:ok
end
case startup_result do
:ok ->
beam_run_id = safe_beam_run_id()
result =
try do
{:ok, fun.()}
catch
kind, reason ->
{:error, {kind, reason, __STACKTRACE__}}
after
SLog.info(:shutdown, "Script execution finished. Shutting down gracefully.")
stop_snakepit(shutdown_timeout, label: "Shutdown")
run_cleanup_with_timeout(beam_run_id, cleanup_timeout)
end
case result do
{:ok, value} ->
maybe_halt(halt, 0)
value
{:error, {kind, reason, stacktrace}} ->
maybe_halt(halt, 1)
:erlang.raise(kind, reason, stacktrace)
end
{:error, %Snakepit.Error{category: :timeout}} ->
SLog.error(:startup, "Pool failed to initialize within #{startup_timeout}ms",
timeout_ms: startup_timeout
)
Application.stop(:snakepit)
maybe_halt(halt, 1)
{:error, :pool_initialization_timeout}
end
end
defp pooling_enabled? do
Application.get_env(:snakepit, :pooling_enabled, false)
end
defp maybe_restart_snakepit(restart, shutdown_timeout, cleanup_timeout) do
if should_restart?(restart) and snakepit_started?() do
SLog.info(:startup, "Restarting to apply script configuration")
beam_run_id = safe_beam_run_id()
stop_snakepit(shutdown_timeout, label: "Restart cleanup")
maybe_cleanup_orphaned_workers(beam_run_id, cleanup_timeout)
end
end
defp should_restart?(true), do: true
defp should_restart?(false), do: false
defp should_restart?(:auto), do: mix_project_loaded?()
defp should_restart?(_), do: false
defp mix_project_loaded? do
mix_started? =
Enum.any?(Application.started_applications(), fn {app, _desc, _vsn} ->
app == :mix
end)
if mix_started? and Code.ensure_loaded?(Mix.Project) and
function_exported?(Mix.Project, :get, 0) do
try do
Mix.Project.get() != nil
catch
_, _ -> false
end
else
false
end
end
defp snakepit_started? do
Enum.any?(Application.started_applications(), fn {app, _desc, _vsn} ->
app == :snakepit
end)
end
defp stop_snakepit(shutdown_timeout, opts) do
label = Keyword.get(opts, :label, "Shutdown")
# Monitor the supervisor to wait for actual shutdown signal
case Process.whereis(Snakepit.Supervisor) do
nil ->
Application.stop(:snakepit)
# Already shut down
SLog.info(:shutdown, "#{label} complete (supervisor already terminated).")
supervisor_pid ->
ref = Process.monitor(supervisor_pid)
Application.stop(:snakepit)
# Wait for :DOWN signal from BEAM - no guessing with sleep
receive do
{:DOWN, ^ref, :process, ^supervisor_pid, _reason} ->
SLog.info(:shutdown, "#{label} complete (confirmed via :DOWN signal).")
after
shutdown_timeout ->
SLog.warning(
:shutdown,
"#{label} confirmation timeout after #{shutdown_timeout}ms. Proceeding anyway.",
shutdown_timeout_ms: shutdown_timeout
)
end
end
end
defp safe_beam_run_id do
ProcessRegistry.get_beam_run_id()
catch
_, _ -> nil
end
defp maybe_cleanup_orphaned_workers(nil, _timeout_ms), do: :ok
defp maybe_cleanup_orphaned_workers(_run_id, timeout_ms) when timeout_ms <= 0, do: :ok
defp maybe_cleanup_orphaned_workers(run_id, timeout_ms) do
if wait_for_run_id_shutdown(run_id, timeout_ms) do
:ok
else
SLog.warning(
:shutdown,
"Worker processes still running after #{timeout_ms}ms. Forcing cleanup...",
timeout_ms: timeout_ms
)
Snakepit.ProcessKiller.kill_by_run_id(run_id)
if not wait_for_run_id_shutdown(run_id, timeout_ms) do
SLog.warning(:shutdown, "Worker processes still running after forced cleanup.")
end
end
end
defp wait_for_run_id_shutdown(run_id, timeout_ms) do
deadline = System.monotonic_time(:millisecond) + timeout_ms
wait_for_run_id_shutdown_loop(run_id, deadline)
end
defp wait_for_run_id_shutdown_loop(run_id, deadline) do
if System.monotonic_time(:millisecond) >= deadline do
false
else
if run_id_processes?(run_id) do
receive do
after
100 -> :ok
end
wait_for_run_id_shutdown_loop(run_id, deadline)
else
true
end
end
end
defp run_id_processes?(run_id) do
Snakepit.ProcessKiller.find_python_processes()
|> Enum.any?(fn pid ->
case Snakepit.ProcessKiller.get_process_command(pid) do
{:ok, cmd} -> run_id_in_command?(cmd, run_id)
_ -> false
end
end)
end
defp run_id_in_command?(command, run_id) do
has_script =
String.contains?(command, "grpc_server.py") or
String.contains?(command, "grpc_server_threaded.py")
has_run_id =
String.contains?(command, "--snakepit-run-id #{run_id}") or
String.contains?(command, "--run-id #{run_id}")
has_script and has_run_id
end
defp run_cleanup_with_timeout(run_id, cleanup_timeout) do
if cleanup_timeout <= 0 or is_nil(run_id) do
maybe_cleanup_orphaned_workers(run_id, cleanup_timeout)
else
task_timeout = cleanup_timeout + 1_000
task = Task.async(fn -> maybe_cleanup_orphaned_workers(run_id, cleanup_timeout) end)
case Task.yield(task, task_timeout) || Task.shutdown(task, :brutal_kill) do
{:ok, _} ->
:ok
nil ->
SLog.warning(
:shutdown,
"Cleanup exceeded #{task_timeout}ms. Skipping remaining cleanup.",
timeout_ms: task_timeout
)
{:exit, reason} ->
SLog.warning(:shutdown, "Cleanup crashed: #{inspect(reason)}", reason: reason)
end
end
end
defp maybe_halt(true, status) do
if status != 0 do
SLog.error(:shutdown, "Halting BEAM with status #{status}.", status: status)
end
# Flush all IO before halting to ensure output is visible
:ok = :io.put_chars(:standard_io, [])
:ok = :io.put_chars(:standard_error, [])
System.halt(status)
end
defp maybe_halt(_, _status), do: :ok
defp env_truthy?(name) do
case System.get_env(name) do
nil -> false
value -> String.downcase(String.trim(value)) in ["1", "true", "yes", "y", "on"]
end
end
end