Class RemoteEmbeddings

tests/modal_remote_vllm/remote_cache_dense.py:202–283

class RemoteEmbeddings:
    def __init__(
        self,
        base_url: str,
        api_key: str = "EMPTY",
        max_concurrent: int = CONCURRENT_BATCHES,
    ):
        self.client = OpenAI(base_url=base_url, api_key=api_key)
        self.max_concurrent = max_concurrent

    def embed(self, texts: List[str], max_retries: int = 3) -> List[np.ndarray]:
        # Validate all texts before sending
        validated_texts = []
        for text in texts:
            if not validate_token_length(text, MAX_TOKENS):
                # If validation fails, re-clip the text more aggressively
                text = clip_to_max_tokens(text, MAX_TOKENS - 10)  # Extra buffer
                if not validate_token_length(text, MAX_TOKENS):
                    # Last resort: truncate very aggressively
                    tokens = get_tokenizer().encode(text, add_special_tokens=False)
                    safe_tokens = tokens[: MAX_TOKENS - 10]
                    text = get_tokenizer().decode(safe_tokens, skip_special_tokens=True)
            validated_texts.append(text)

        # Retry logic for network resilience under high load
        for attempt in range(max_retries):
            try:
                resp = self.client.embeddings.create(
                    model="Snowflake/snowflake-arctic-embed-m-v1.5",
                    input=validated_texts,
                    encoding_format="float",
                )
                # Convert to float32 for downstream usage
                return [np.array(e.embedding, dtype=np.float32) for e in resp.data]
            except Exception as e:
                if attempt == max_retries - 1:
                    raise e
                # Exponential backoff on retry
                wait_time = (2**attempt) * 0.5  # 0.5s, 1s, 2s
                time.sleep(wait_time)

    def embed_concurrent(self, texts: List[str]) -> List[np.ndarray]:
        """
        Process embeddings with concurrent batches to maximize vLLM server utilization.
        """
        if len(texts) <= EMB_BATCH:
            return self.embed(texts)

        # Split texts into batches
        batches = [texts[i : i + EMB_BATCH] for i in range(0, len(texts), EMB_BATCH)]

        embeddings = []

        # Process batches concurrently with ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=self.max_concurrent) as executor:
            # Submit all batches
            future_to_batch = {
                executor.submit(self.embed, batch): i for i, batch in enumerate(batches)