precommit

Author: ishandhanani

examples/sglang/components/embedding_worker.py

@@ -17,18 +17,14 @@
 Using SGLang and Dynamo to serve embedding models!
 """
 
-import asyncio
 import logging
-import random
-import socket
-from typing import Any
 
 import sglang as sgl
 from utils.protocol import EmbeddingRequest
 from utils.sglang import parse_sglang_args
 
 from dynamo.llm import ModelType, register_llm
-from dynamo.sdk import async_on_start, depends, dynamo_context, endpoint, service
+from dynamo.sdk import async_on_start, dynamo_context, endpoint, service
 
 logger = logging.getLogger(__name__)
 
@@ -41,7 +37,6 @@
     workers=1,
 )
 class SGLangEmbeddingWorker:
-
     def __init__(self):
         class_name = self.__class__.__name__
         self.engine_args = parse_sglang_args(class_name, "")
@@ -70,11 +65,11 @@ async def generate(self, request: EmbeddingRequest):
             input = [i for i in request.input]
         else:
             raise ValueError(f"Invalid input type: {type(request.input)}")
-        
+
         g = await self.engine.async_encode(
             prompt=input,
         )
-        
+
         # Transform response to match OpenAI embedding format
         response = self._transform_response(g, request.model)
         yield response
@@ -83,16 +78,18 @@ def _transform_response(self, ret, model_name):
         """Transform SGLang response to OpenAI embedding format"""
         if not isinstance(ret, list):
             ret = [ret]
-        
+
         embedding_objects = []
         prompt_tokens = 0
-        
+
         for idx, ret_item in enumerate(ret):
-            embedding_objects.append({
-                "object": "embedding",
-                "embedding": ret_item["embedding"],
-                "index": idx,
-            })
+            embedding_objects.append(
+                {
+                    "object": "embedding",
+                    "embedding": ret_item["embedding"],
+                    "index": idx,
+                }
+            )
             prompt_tokens += ret_item["meta_info"]["prompt_tokens"]
 
         return {
@@ -103,4 +100,4 @@ def _transform_response(self, ret, model_name):
                 "prompt_tokens": prompt_tokens,
                 "total_tokens": prompt_tokens,
             },
-        }
+        }

examples/sglang/components/frontend.py

@@ -17,8 +17,8 @@
 import subprocess
 from pathlib import Path
 
-from components.worker import SGLangWorker
 from components.embedding_worker import SGLangEmbeddingWorker
+from components.worker import SGLangWorker
 from fastapi import FastAPI
 from pydantic import BaseModel
 

examples/sglang/configs/embedding.yaml

@@ -5,7 +5,6 @@ Frontend:
 SGLangEmbeddingWorker:
   model-path: intfloat/e5-base-v2
   served-model-name: e5
-  is-embedding: true
   tp: 1
   trust-remote-code: true
   is-embedding: true

examples/sglang/graphs/embedding.py

@@ -1,4 +1,3 @@
-
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -14,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from components.frontend import Frontend
 from components.embedding_worker import SGLangEmbeddingWorker
+from components.frontend import Frontend
 
 Frontend.link(SGLangEmbeddingWorker)

examples/sglang/utils/protocol.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Union, Literal
+from typing import List, Literal, Optional, Union
 
 from pydantic import BaseModel, Field
 
@@ -61,18 +61,17 @@ class DisaggPreprocessedRequest(BaseModel):
     bootstrap_port: int
     bootstrap_room: int
 
-EmbeddingInput = Union[
-    str,
-    List[str], 
-    List[int],  
-    List[List[int]]  
-]
+
+EmbeddingInput = Union[str, List[str], List[int], List[List[int]]]
 
 EncodingFormat = Literal["float", "base64"]
 
+
 class EmbeddingRequest(BaseModel):
     model: str
     input: EmbeddingInput
     encoding_format: Optional[EncodingFormat] = None
     user: Optional[str] = None
-    dimensions: Optional[int] = None # only supported in text-embedding-3 and later models from OpenAI
+    dimensions: Optional[
+        int
+    ] = None  # only supported in text-embedding-3 and later models from OpenAI