Introduce FirstAvailableStrategyI base class encapsulating Redis coordination and monitoring logic; update FirstAvailableStrategy and FirstAvailableOptimStrategy to inherit from it, remove duplicated scripts, buffer handling and monitor setup, and add configurable strategy_prefix for Redis key prefixes.

Author: Paweł Kędzia

llm_router_api/base/lb/first_available.py

@@ -36,11 +36,11 @@
 from typing import List, Dict, Optional, Any
 
 from llm_router_api.base.constants import REDIS_PORT, REDIS_HOST
-from llm_router_api.base.lb.strategy import ChooseProviderStrategyI
 from llm_router_api.base.lb.provider_monitor import RedisProviderMonitor
+from llm_router_api.base.lb.first_available_i import FirstAvailableStrategyI
 
 
-class FirstAvailableStrategy(ChooseProviderStrategyI):
+class FirstAvailableStrategy(FirstAvailableStrategyI):
  """
  Strategy that selects the first free provider for a model using Redis.
 
@@ -87,53 +87,16 @@ def __init__(
  clear_buffers:
  Whether to clear all buffers when starting. Default is ``True``.
  """
- if not REDIS_IS_AVAILABLE:
- raise RuntimeError("Redis is not available. Please install it first.")
-
- super().__init__(models_config_path=models_config_path, logger=logger)
-
- self.redis_client = redis.Redis(
- host=redis_host, port=redis_port, db=redis_db, decode_responses=True
- )
- self.timeout = timeout
- self.check_interval = check_interval
-
- # Atomic acquire script – treat missing field as “available”
- self._acquire_script = self.redis_client.register_script(
- """
- local redis_key = KEYS[1]
- local field = ARGV[1]
- local v = redis.call('HGET', redis_key, field)
- -- v == false -> field does not exist (nil)
- -- v == 'false' -> explicitly marked as free
- if v == false or v == 'false' then
- redis.call('HSET', redis_key, field, 'true')
- return 1
- end
- return 0
- """
- )
-
- # Atomic release script – simply delete the field (no race condition)
- self._release_script = self.redis_client.register_script(
- """
- local redis_key = KEYS[1]
- local field = ARGV[1]
- -- Delete the field; returns 1 if field existed, 0 otherwise
- redis.call('HDEL', redis_key, field)
- return 1
- """
- )
-
- if clear_buffers:
- self._clear_buffers()
-
- # Start providers monitor
- self._monitor = RedisProviderMonitor(
- redis_client=self.redis_client,
- check_interval=30,
+ super().__init__(
+ models_config_path=models_config_path,
+ redis_host=redis_host,
+ redis_port=redis_port,
+ redis_db=redis_db,
+ timeout=timeout,
+ check_interval=check_interval,
  clear_buffers=clear_buffers,
- logger=self.logger,
+ logger=logger,
+ strategy_prefix="fa_",
  )
 
  def get_provider(
@@ -185,24 +148,14 @@ def get_provider(
  * Call :meth:`put_provider` to release the lock once the provider is no
  longer needed.
  """
- if not providers:
- return None
 
- # Register providers for monitoring (only once per model)
- self._monitor.add_providers(model_name, providers)
+ redis_key, is_random = self.init_provider(
+ model_name=model_name, providers=providers, options=options
+ )
+ if not redis_key:
+ return None
 
- redis_key = self._get_redis_key(model_name)
  start_time = time.time()
-
- # Ensure fields exist; if someone removed the hash, recreate it
- if not self.redis_client.exists(redis_key):
- for p in providers:
- self.redis_client.hset(redis_key, self._provider_field(p), "false")
-
- # self._print_provider_status(redis_key, providers)
-
- is_random = options and options.get("random_choice", False)
-
  while True:
  _providers = self._get_active_providers(
  model_name=model_name, providers=providers
@@ -348,145 +301,3 @@ def _get_active_providers(
  model_name=model_name, only_active=True
  )
  return active_providers
-
- def _get_redis_key(self, model_name: str) -> str:
- """
- Return Redis key prefix for a given model.
- """
- for ch in self.REPLACE_PROVIDER_KEY:
- model_name = model_name.replace(ch, "_")
- return f"model:{model_name}"
-
- def _provider_field(self, provider: dict) -> str:
- """
- Build the Redis hash field name that stores the chosen flag
- for a given provider.
-
- Parameters
- ----------
- provider : dict
- Provider configuration dictionary.
-
- Returns
- -------
- str
- Field name in the format ``{provider_id}:is_chosen``.
- """
- provider_id = self._provider_key(provider)
- return f"{provider_id}:is_chosen"
-
- def _init_flag(self, model_name: str) -> str:
- """
- Build the Redis key used as an initialization flag for a model.
-
- Parameters
- ----------
- model_name : str
- Name of the model.
-
- Returns
- -------
- str
- Flag key in the format ``model:{model_name}:initialized``.
- """
- return f"{self._get_redis_key(model_name)}:initialized"
-
- def _initialize_providers(self, model_name: str, providers: List[Dict]) -> None:
- """
- Ensure that the provider lock fields for *model_name* exist in Redis.
-
- This method is idempotent – it will create the hash fields only if the
- model has not been initialized before. An auxiliary flag key
- ``model:{model_name}:initialized`` is used to guard against repeated
- initialization, which could otherwise overwrite the current lock state
- of providers that are already in use.
-
- Parameters
- ----------
- model_name : str
- The name of the model whose providers are being prepared.
- providers : List[Dict]
- A list of provider configuration dictionaries. Each dictionary must
- contain enough information for :meth:`_provider_field` to generate a
- unique field name.
-
- Notes
- -----
- * The provider fields are stored in a Redis hash whose key is
- ``model:{model_name}``. Each field is set to the string ``'false'``
- to indicate that the provider is currently free.
- * The initialization flag is a simple Redis key with value ``'1'``.
- Its existence signals that the hash has already been populated.
- """
- redis_key = self._get_redis_key(model_name)
-
- # Check if already initialized using a flag
- init_flag = self._init_flag(model_name)
- if self.redis_client.exists(init_flag):
- return
-
- # Initialize all providers as available
- for provider in providers:
- provider_field = self._provider_field(provider)
- self.redis_client.hset(redis_key, provider_field, "false")
-
- # Set initialization flag
- self.redis_client.set(init_flag, "1")
-
- def _clear_buffers(self) -> None:
- """
- Reset the Redis state for all active models.
-
- This method removes any existing initialization flags and provider
- lock fields, then re‑initialises the providers as available. It is
- typically invoked during strategy start‑up to ensure a clean slate.
- """
- active_models = self._api_model_config.active_models
- models_configs = self._api_model_config.models_configs
- for _, models_names in active_models.items():
- for model_name in models_names:
- redis_key = self._get_redis_key(model_name)
- providers = models_configs[model_name]["providers"]
- if len(providers) > 0:
- model_path = providers[0].get("model_path", "").strip()
- if model_path:
- model_name = model_path
-
- init_flag = self._init_flag(model_name)
- self.redis_client.delete(init_flag)
-
- for provider in providers:
- provider_field = self._provider_field(provider)
- self.redis_client.hset(redis_key, provider_field, "false")
-
- self._initialize_providers(
- model_name=model_name, providers=providers
- )
-
- def _print_provider_status(self, redis_key: str, providers: List[Dict]) -> None:
- """
- Print the lock status of each provider stored in the Redis hash
- ``redis_key``. Uses emojis for a quick visual cue:
-
- * 🟢 – provider is free (`'false'` or missing)
- * 🔴 – provider is currently taken (`'true'`)
-
- The output is formatted in a table‑like layout for readability.
- """
- try:
- # Retrieve the entire hash; missing fields default to None
- hash_data = self.redis_client.hgetall(redis_key)
- except Exception as exc:
- print(f"[⚠️] Could not read Redis key '{redis_key}': {exc}")
- return
-
- print("\nProvider lock status:")
- print("-" * 40)
- for provider in providers:
- field = self._provider_field(provider)
- status = hash_data.get(field, "false")
- icon = "🔴" if status == "true" else "🟢"
- # Show a short identifier for the provider (fallback to field)
- provider_id = provider.get("id") or provider.get("name") or field
- print(f"{icon} {provider_id:<30} [{field}]")
- print("-" * 40)