Embedding Model Batching Service

Wrap a provided embedding model in a service that supports batched requests, shutdown, and concurrent processing. One variant adds max batch and max-token limits for batches of sequences. Problem Requirements You need to build a server service. This service wraps around an EmbeddingModel class. Its job is to turn text into embeddings (lists of numbers). The server must handle groups of data efficiently and shut down safely when asked. Main Goals: Batch Processing: Handle many text inputs at the same time. Graceful Shutdown: Listen for a shutdown command and stop safely. Result Mapping: Return the correct embedding for each input ID. Scalability: Handle batches containing up to 100 items. Input Format: # Batch request batch_data = [ {'id': '1', 'data': 'example text 1'}, {'id': '2', 'data': 'example text 2'} ] # Shutdown request shutdown_command = 'shutdown' Output Format: # Batch response [ {'id': '1', 'embedding': [0.1, 0.2, ...]}, {'id': '2', 'embedding': [0.3, 0.4, ...]} ] # Shutdown response 'Server shutdown complete.' Example: server = ModelServer() # Process batch result = server.handle_request([ {'id': '1', 'data': 'machine learning'}, {'id': '2', 'data': 'neural networks'} ]) # Output: [{'id': '1', 'embedding': [...]}, {'id': '2', 'embedding': [...]}] # Shutdown result = server.handle_request('shutdown') # Output: 'Server shutdown complete.' Part 1: Solution 1 (Simple Approach) One Thing at a Time The basic version handles one request after another. It does not do things at the same time. class EmbeddingModel: """Base embedding model class.""" def __init__(self): # Real code would load model weights here pass def process(self, text: str) -> list[float]: """ Make an embedding for one text input. Args: text: Input text Returns: List of floats (the embedding) """ # Simple fake embedding (real code uses a neural network) return [0.0] * 384 # Simulating a 384-dimensional embedding class ModelServer: """Server for handling requests.""" def __init__(self): self.model = EmbeddingModel() self.is_running = True def handle_request(self, request_data): """ Handle incoming requests (batch or shutdown). Args: request_data: List of dicts OR 'shutdown' string Returns: List of embeddings OR shutdown message """ # Check for shutdown command if request_data == 'shutdown': self.is_running = False return 'Server shutdown complete.' # Process batch request results = [] for item in request_data: item_id = item['id'] data = item['data'] embedding = self.model.process(data) results.append({ 'id': item_id, 'embedding': embedding }) return results # Usage example server = ModelServer() # Process batch batch = [ {'id': '1', 'data': 'example text 1'}, {'id': '2', 'data': 'example text 2'} ] print(server.handle_request(batch)) # Shutdown print(server.handle_request('shutdown')) Time Complexity: Processing n items takes O(n × T). T is the time it takes to make one embedding. Because it is sequential, each item waits for the one before it. Space Complexity: O(n × d). n is the batch size. d is the embedding dimension. We store all results in memory before returning them. Part 2: Solution 2 (Faster Approach) Using Multiple Threads For real apps, we use parallel processing. This lets us do more work in less time. from concurrent.futures import ThreadPoolExecutor, as_completed from threading import Lock from typing import List, Dict, Union, Any import time class EmbeddingModel: """Thread-safe embedding model.""" def __init__(self): # Fake setup time self.initialization_time = time.time() def process(self, text: str) -> list[float]: """ Make embedding for one text input. Safe for threads. """ # Fake processing time (real code uses GPU/CPU) time.sleep(0.01) # 10ms per item # Make a dummy embedding based on text hash # Real code would use the actual model text_hash = hash(text) embedding = [(text_hash % 1000) / 1000.0] * 384 return embedding class ThreadSafeModelServer: """Server that uses a thread pool for speed.""" def __init__(self, max_workers: int = 10): """ Start server with thread pool. Args: max_workers: Max number of threads at once """ self.model = EmbeddingModel() self.max_workers = max_workers self.is_running = True self.lock = Lock() # Keep data safe self.executor = ThreadPoolExecutor(max_workers=max_workers) def _process_single_item(self, item: Dict[str, str]) -> Dict[str, Any]: """ Process one item (runs inside a thread). Args: item: Dictionary with 'id' and 'data' Returns: Dictionary with 'id' and 'embedding' """ item_id = item['id'] data = item['data'] # Make embedding (thread-safe) embedding = self.model.process(data) return { 'id': item_id, 'embedding': embedding } def handle_request(self, request_data: Union[str, List[Dict[str, str]]]): """ Handle requests using parallel processing. Args: request_data: List of dicts OR 'shutdown' string Returns: List of embeddings OR shutdown message """ # Check for shutdown command if request_data == 'shutdown': with self.lock: self.is_running = False self.shutdown() return 'Server shutdown complete.' # Check batch size limit if len(request_data) > 100: raise ValueError("Batch size exceeds maximum limit of 100") # Send all items to thread pool future_to_item = { self.executor.submit(self._process_single_item, item): item for item in request_data } # Collect results as they finish results = [] for future in as_completed(future_to_item): try: result = future.result() results.append(result) except Exception as e: # Handle errors for single items item = future_to_item[future] results.append({ 'id': item['id'], 'embedding': None, 'error': str(e) }) # Sort by ID to keep order (optional) results.sort(key=lambda x: x['id']) return results def shutdown(self): """Safely stop server and threads.""" print("Shutting down server...") self.executor.shutdown(wait=True) # Wait for tasks to finish print("All tasks completed. Server stopped.") def __enter__(self): """Support for 'with' statement.""" return self def __exit__(self, exc_type, exc_val, exc_tb): """Cleanup when exiting 'with' block.""" self.shutdown() # Usage example with context manager def main(): with ThreadSafeModelServer(max_workers=5) as server: # Process batch batch = [ {'id': '1', 'data': 'machine learning'}, {'id': '2', 'data': 'neural networks'}, {'id': '3', 'data': 'deep learning'}, {'id': '4', 'data': 'natural language processing'} ] start_time = time.time() results = server.handle_request(batch) end_time = time.time() print(f"Processed {len(results)} items in {end_time - start_time:.2f}s") print(f"Results: {results[:2]}...") # Show first 2 results # Shutdown shutdown_msg = server.handle_request('shutdown') print(shutdown_msg) if __name__ == "__main__": main() Performance Comparison: Method Time for 100 items (10ms each) Speed Single-threaded ~1000ms (one by one) 100 items/s Multi-threaded (10 workers) ~100ms (parallel) 1000 items/s Time Complexity: With w workers: O(⌈n/w⌉ × T). If you have 10 workers, it is roughly 10 times faster (assuming no I/O delays). Space Complexity: Extra space: O(w) for the thread pool. Futures: O(n) to track all the tasks we submitted. Important Concepts The following are supplementary notes for preparation, beyond the core question above. 1. Topic: Threading vs. Multiprocessing Threading: ✅ Low memory cost (shares memory). ✅ Good when waiting for things (like API calls). ❌ Limited by GIL for heavy CPU work. Multiprocessing: ✅ True parallel work for CPU tasks (ignores GIL). ✅ Good for heavy math (ML inference). ❌ Uses more memory (each process has its own memory). Which one for embedding models? GPU inference: Use Threading. GPU calls usually release the GIL. CPU inference: Multiprocessing is often better for true parallel speed. External API: Use Async/await or Threading because you are just waiting for a response. 2. Topic: Safe Shutdown What you need to do: Wait for current tasks: Don't stop halfway through a request. Clear the queue: Finish pending items or reject them safely. Clean up: Close connections and free up memory. Confirm: Tell the client "I am done." How to write it: def shutdown(self): # 1. Stop taking new requests self.is_running = False # 2. Wait for active tasks to finish self.executor.shutdown(wait=True) # 3. Cleanup del self.model # Free model memory # 4. Confirm return 'Server shutdown complete.' 3. Topic: Tracking Performance Key things to measure: Latency: How fast are the responses? (p50, p95, p99). Throughput: How many requests per second? Error rate: How many requests fail? Resource usage: How much CPU and RAM are we using? Active threads: How many threads are busy right now? How to write it: import time class MonitoredModelServer: """Server with built-in monitoring.""" def __init__(self): self.model = EmbeddingModel() self.metrics = { 'latency': [], 'success_count': 0, 'error_count': 0 } self.is_running = True def handle_request(self, request_data): start_time = time.time() try: # Check for shutdown command if request_data == 'shutdown': self.is_running = False return 'Server shutdown complete.' # Process batch request results = [] for item in request_data: embedding = self.model.process(item['data']) results.append({ 'id': item['id'], 'embedding': embedding }) # Record success latency = time.time() - start_time self.metrics['latency'].append(latency) self.metrics['success_count'] += 1 return results except Exception as e: # Record failure self.metrics['error_count'] += 1 raise def get_metrics(self): """Return current metrics.""" latencies = self.metrics['latency'] return { 'total_requests': len(latencies), 'success_count': self.metrics['success_count'], 'error_count': self.metrics['error_count'], 'avg_latency': sum(latencies) / len(latencies) if latencies else 0, 'p95_latency': sorted(latencies)[int(len(latencies) * 0.95)] if latencies else 0 }