82 lines
3.5 KiB
Python
82 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
from django.core.management.base import BaseCommand, CommandError
|
|
|
|
from control_plane.resources.models import Resource
|
|
from model_router.providers import QwenProvider
|
|
from model_router.router import ModelCapability, ModelRequestContract
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Probe Qwen/local_inference health and bounded concurrent completion capacity."
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument("--requests", type=int, default=1)
|
|
parser.add_argument("--concurrency", type=int, default=1)
|
|
parser.add_argument("--prompt", default="Return only valid JSON: {\"ok\": true, \"model\": \"qwen\"}.")
|
|
parser.add_argument("--token-budget", type=int, default=256)
|
|
|
|
def handle(self, *args, **options):
|
|
resource = Resource.objects.filter(provider="local_inference", is_active=True).first()
|
|
if resource is None:
|
|
raise CommandError("No Qwen/local_inference resource configured. Run seed_spark_resources first.")
|
|
provider = QwenProvider(resource)
|
|
before_health = provider.health()
|
|
total = max(1, int(options["requests"]))
|
|
concurrency = max(1, int(options["concurrency"]))
|
|
started = time.monotonic()
|
|
|
|
def call(index: int) -> dict[str, object]:
|
|
call_started = time.monotonic()
|
|
try:
|
|
response = provider.complete(
|
|
ModelRequestContract(
|
|
purpose=ModelCapability.REASONING,
|
|
prompt=f"{options['prompt']} Request index: {index}",
|
|
token_budget=int(options["token_budget"]),
|
|
)
|
|
)
|
|
return {
|
|
"index": index,
|
|
"status": "COMPLETE",
|
|
"latency_ms": int((time.monotonic() - call_started) * 1000),
|
|
"content_chars": len(response.content),
|
|
"model": response.model,
|
|
"attempts": response.metadata.get("attempts"),
|
|
}
|
|
except Exception as exc:
|
|
return {
|
|
"index": index,
|
|
"status": "FAILED",
|
|
"latency_ms": int((time.monotonic() - call_started) * 1000),
|
|
"failure": str(exc),
|
|
}
|
|
|
|
results = []
|
|
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
|
futures = {executor.submit(call, index): index for index in range(1, total + 1)}
|
|
for future in as_completed(futures):
|
|
results.append(future.result())
|
|
results.sort(key=lambda row: row["index"])
|
|
after_health = provider.health()
|
|
completed = sum(1 for row in results if row["status"] == "COMPLETE")
|
|
failed = total - completed
|
|
latencies = [int(row["latency_ms"]) for row in results]
|
|
output = {
|
|
"before_health": before_health,
|
|
"after_health": after_health,
|
|
"requests": total,
|
|
"concurrency": concurrency,
|
|
"completed": completed,
|
|
"failed": failed,
|
|
"runtime_seconds": round(time.monotonic() - started, 2),
|
|
"latency_ms": {"min": min(latencies), "max": max(latencies), "avg": round(sum(latencies) / len(latencies), 1)},
|
|
"results": results,
|
|
}
|
|
self.stdout.write(json.dumps(output, indent=2))
|
|
if failed:
|
|
raise CommandError(f"Qwen probe failed {failed}/{total} requests")
|