Spaces:
Paused
Paused
| # | |
| # SPDX-FileCopyrightText: Hadad <[email protected]> | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| import random # Import random module to enable random selection from a list | |
| from datetime import datetime # Import datetime class to work with current UTC time | |
| from typing import Dict, List # Import type hints for dictionaries and lists (not explicitly used here but imported) | |
| from config import auth # Import authentication configuration, likely a list of host dictionaries with credentials | |
| from src.utils.helper import busy, mark # Import 'busy' dictionary and 'mark' function to track and update host busy status | |
| # Initialize a global dictionary to map session IDs to assigned hosts | |
| mapping = {} # Store session_id to host assignment mapping to maintain consistent host allocation per session | |
| # Define a function to get an available host for a given session, optionally excluding certain hosts | |
| def get_host(session_id: str, exclude_hosts: List[str] = None) -> dict: | |
| """ | |
| Retrieve or assign a host for the given session ID. | |
| Args: | |
| session_id (str): A unique identifier for the current session. | |
| Returns: | |
| dict: The selected host dictionary from the auth configuration. | |
| Raises: | |
| Exception: If no available hosts are found to assign. | |
| Explanation: | |
| Retrieve an available host for the specified session ID, ensuring excluded hosts are not assigned. | |
| This function maintains a mapping of session IDs to hosts to provide consistent host assignment. | |
| It filters out busy hosts and those explicitly excluded, then randomly selects an available host. | |
| If no hosts are available, it raises an exception. | |
| """ | |
| # If no list of hosts to exclude is provided, initialize it as an empty list | |
| if exclude_hosts is None: # Check if exclude_hosts parameter was omitted or set to None | |
| exclude_hosts = [] # Initialize exclude_hosts to an empty list to avoid errors during filtering | |
| # Check if the session ID already has an assigned host in the mapping dictionary | |
| if session_id in mapping: # Verify if a host was previously assigned to this session | |
| assigned_host = mapping[session_id] # Retrieve the assigned host dictionary for this session | |
| # If the assigned host is not in the list of hosts to exclude, return it immediately | |
| if assigned_host["jarvis"] not in exclude_hosts: # Ensure assigned host is allowed for this request | |
| return assigned_host # Return the cached host assignment for session consistency | |
| else: | |
| # If the assigned host is excluded, remove the mapping to allow reassignment | |
| del mapping[session_id] # Delete the existing session-host mapping to find a new host | |
| # Get the current UTC time to compare against host busy status timestamps | |
| now = datetime.utcnow() # Capture current time to filter out hosts that are still busy | |
| # Create a list of hosts that are not currently busy and not in the exclude list | |
| available_hosts = [ | |
| h for h in auth # Iterate over all hosts defined in the authentication configuration | |
| if h["jarvis"] not in busy or busy[h["jarvis"]] <= now # Include hosts not busy or whose busy time has expired | |
| if h["jarvis"] not in exclude_hosts # Exclude hosts specified in the exclude_hosts list | |
| ] | |
| # If no hosts are available after filtering, raise an exception to indicate resource exhaustion | |
| if not available_hosts: # Check if the filtered list of hosts is empty | |
| raise Exception("No available hosts to assign.") # Inform caller that no hosts can be assigned currently | |
| # Randomly select one host from the list of available hosts to distribute load evenly | |
| selected = random.choice(available_hosts) # Choose a host at random to avoid bias in host selection | |
| # Store the selected host in the mapping dictionary for future requests with the same session ID | |
| mapping[session_id] = selected # Cache the selected host to maintain session affinity | |
| # Mark the selected host as busy using the helper function to update its busy status | |
| mark(selected["jarvis"]) # Update the busy dictionary to indicate this host is now in use | |
| # Return the selected host dictionary to the caller for use in processing the session | |
| return selected # Provide the caller with the assigned host details |