fix: apply ruff fixes to ray_serve package
Some checks failed
Build and Publish ray-serve-apps / lint (push) Successful in 1m30s
Build and Publish ray-serve-apps / publish (push) Failing after 2m44s

[ray-serve only]

- Fix whitespace in docstrings
- Add strict=True to zip() calls
- Use ternary operators where appropriate
- Rename unused loop variables
This commit is contained in:
2026-02-02 11:09:35 -05:00
parent 16f6199534
commit 12987c6adc
6 changed files with 128 additions and 131 deletions

View File

@@ -4,9 +4,7 @@ Runs on: drizzt (Radeon 680M iGPU, ROCm) or danilo (Intel i915 iGPU, OpenVINO/IP
"""
import os
import time
import uuid
from typing import Any, Dict, List, Tuple
from typing import Any
from ray import serve
@@ -14,16 +12,17 @@ from ray import serve
@serve.deployment(name="RerankerDeployment", num_replicas=1)
class RerankerDeployment:
def __init__(self):
from sentence_transformers import CrossEncoder
import torch
from sentence_transformers import CrossEncoder
self.model_id = os.environ.get("MODEL_ID", "BAAI/bge-reranker-v2-m3")
self.use_ipex = False
self.device = "cpu"
# Detect device - check for Intel GPU first via IPEX
try:
import intel_extension_for_pytorch as ipex
self.use_ipex = True
if hasattr(torch, "xpu") and torch.xpu.is_available():
self.device = "xpu"
@@ -32,36 +31,37 @@ class RerankerDeployment:
print("IPEX available, will use CPU optimization")
except ImportError:
print("IPEX not available, checking for other GPUs")
# Check for CUDA/ROCm if not using Intel
if not self.use_ipex:
if torch.cuda.is_available():
self.device = "cuda"
print(f"Using CUDA/ROCm device")
print("Using CUDA/ROCm device")
else:
print("No GPU detected, using CPU")
print(f"Loading reranker model: {self.model_id}")
print(f"Using device: {self.device}")
# Load model
self.model = CrossEncoder(self.model_id, device=self.device)
# Apply IPEX optimization if available
if self.use_ipex and self.device == "cpu":
try:
import intel_extension_for_pytorch as ipex
self.model.model = ipex.optimize(self.model.model)
print("IPEX CPU optimization applied")
except Exception as e:
print(f"IPEX optimization failed: {e}")
print(f"Reranker model loaded successfully")
async def __call__(self, request: Dict[str, Any]) -> Dict[str, Any]:
print("Reranker model loaded successfully")
async def __call__(self, request: dict[str, Any]) -> dict[str, Any]:
"""
Handle reranking requests.
Expected request format:
{
"query": "search query",
@@ -69,7 +69,7 @@ class RerankerDeployment:
"top_k": 3,
"return_documents": true
}
Alternative format (pairs):
{
"pairs": [["query", "doc1"], ["query", "doc2"]]
@@ -79,42 +79,44 @@ class RerankerDeployment:
if "pairs" in request:
pairs = request["pairs"]
scores = self.model.predict(pairs)
results = []
for i, (pair, score) in enumerate(zip(pairs, scores)):
results.append({
"index": i,
"score": float(score),
})
for i, (_pair, score) in enumerate(zip(pairs, scores, strict=False)):
results.append(
{
"index": i,
"score": float(score),
}
)
return {
"object": "list",
"results": results,
"model": self.model_id,
}
# Handle query + documents format
query = request.get("query", "")
documents = request.get("documents", [])
top_k = request.get("top_k", len(documents))
return_documents = request.get("return_documents", True)
if not documents:
return {
"object": "list",
"results": [],
"model": self.model_id,
}
# Create query-document pairs
pairs = [[query, doc] for doc in documents]
# Get scores
scores = self.model.predict(pairs)
# Create results with indices and scores
results = []
for i, (doc, score) in enumerate(zip(documents, scores)):
for i, (doc, score) in enumerate(zip(documents, scores, strict=False)):
result = {
"index": i,
"score": float(score),
@@ -122,13 +124,13 @@ class RerankerDeployment:
if return_documents:
result["document"] = doc
results.append(result)
# Sort by score descending
results.sort(key=lambda x: x["score"], reverse=True)
# Apply top_k
results = results[:top_k]
return {
"object": "list",
"results": results,