feat: Add GPU-specific Ray worker images with CI/CD
- Add Dockerfiles for nvidia, rdna2, strixhalo, and intel GPU targets - Add ray-serve modules (embeddings, whisper, tts, llm, reranker) - Add Gitea Actions workflow for automated builds - Add Makefile for local development - Update README with comprehensive documentation
This commit is contained in:
142
ray-serve/serve_reranker.py
Normal file
142
ray-serve/serve_reranker.py
Normal file
@@ -0,0 +1,142 @@
|
||||
"""
|
||||
Ray Serve deployment for sentence-transformers CrossEncoder reranking.
|
||||
Runs on: drizzt (Radeon 680M iGPU, ROCm) or danilo (Intel i915 iGPU, OpenVINO/IPEX)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from ray import serve
|
||||
|
||||
|
||||
@serve.deployment(name="RerankerDeployment", num_replicas=1)
|
||||
class RerankerDeployment:
|
||||
def __init__(self):
|
||||
from sentence_transformers import CrossEncoder
|
||||
import torch
|
||||
|
||||
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"
|
||||
print("Intel GPU detected via IPEX, using XPU device")
|
||||
else:
|
||||
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")
|
||||
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]:
|
||||
"""
|
||||
Handle reranking requests.
|
||||
|
||||
Expected request format:
|
||||
{
|
||||
"query": "search query",
|
||||
"documents": ["doc1", "doc2", "doc3"],
|
||||
"top_k": 3,
|
||||
"return_documents": true
|
||||
}
|
||||
|
||||
Alternative format (pairs):
|
||||
{
|
||||
"pairs": [["query", "doc1"], ["query", "doc2"]]
|
||||
}
|
||||
"""
|
||||
# Handle pairs format
|
||||
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),
|
||||
})
|
||||
|
||||
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)):
|
||||
result = {
|
||||
"index": i,
|
||||
"score": float(score),
|
||||
}
|
||||
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,
|
||||
"model": self.model_id,
|
||||
"usage": {
|
||||
"total_pairs": len(pairs),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
app = RerankerDeployment.bind()
|
||||
Reference in New Issue
Block a user