[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
242ec57305
commit
c40ed3766a
1 changed files with 63 additions and 42 deletions
|
|
@ -43,7 +43,11 @@ DEFAULT_VARIANT = "UD-Q4_K_XL"
|
|||
PORT = 18222 # high port unlikely to collide
|
||||
HOST = "127.0.0.1"
|
||||
STARTUP_TIMEOUT = 120 # seconds to wait for banner
|
||||
LOG_FILE = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "test_studio_run.log"
|
||||
LOG_FILE = (
|
||||
Path(__file__).resolve().parent.parent.parent.parent
|
||||
/ "temp"
|
||||
/ "test_studio_run.log"
|
||||
)
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────────────────
|
||||
|
|
@ -59,14 +63,14 @@ def _http(
|
|||
) -> tuple[int, str]:
|
||||
"""Minimal stdlib HTTP helper. Returns (status_code, body_text)."""
|
||||
data = json.dumps(body).encode() if body else None
|
||||
req = urllib.request.Request(url, data=data, headers=headers or {}, method=method)
|
||||
req = urllib.request.Request(url, data = data, headers = headers or {}, method = method)
|
||||
if body:
|
||||
req.add_header("Content-Type", "application/json")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
||||
return resp.status, resp.read().decode()
|
||||
except urllib.error.HTTPError as exc:
|
||||
return exc.code, exc.read().decode(errors="replace")
|
||||
return exc.code, exc.read().decode(errors = "replace")
|
||||
|
||||
|
||||
def _stream_http(
|
||||
|
|
@ -78,11 +82,11 @@ def _stream_http(
|
|||
) -> tuple[int, list[dict]]:
|
||||
"""POST a streaming request and collect SSE chunks."""
|
||||
data = json.dumps(body).encode()
|
||||
req = urllib.request.Request(url, data=data, headers=headers, method="POST")
|
||||
req = urllib.request.Request(url, data = data, headers = headers, method = "POST")
|
||||
req.add_header("Content-Type", "application/json")
|
||||
chunks: list[dict] = []
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
||||
status = resp.status
|
||||
for raw_line in resp:
|
||||
line = raw_line.decode().strip()
|
||||
|
|
@ -103,9 +107,9 @@ def test_help_output():
|
|||
"""``unsloth studio run --help`` should show all documented options."""
|
||||
result = subprocess.run(
|
||||
["unsloth", "studio", "run", "--help"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=15,
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 15,
|
||||
)
|
||||
out = result.stdout
|
||||
assert result.returncode == 0, f"--help exited with {result.returncode}"
|
||||
|
|
@ -130,11 +134,11 @@ def test_curl_basic(base_url: str, api_key: str):
|
|||
status, text = _http(
|
||||
"POST",
|
||||
f"{base_url}/v1/chat/completions",
|
||||
body={
|
||||
body = {
|
||||
"messages": [{"role": "user", "content": "Say just the word hello"}],
|
||||
"stream": False,
|
||||
},
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
headers = {"Authorization": f"Bearer {api_key}"},
|
||||
)
|
||||
assert status == 200, f"Expected 200, got {status}: {text[:300]}"
|
||||
data = json.loads(text)
|
||||
|
|
@ -162,11 +166,11 @@ def test_curl_streaming(base_url: str, api_key: str):
|
|||
"""Example 2: streaming chat completion via HTTP SSE."""
|
||||
status, chunks = _stream_http(
|
||||
f"{base_url}/v1/chat/completions",
|
||||
body={
|
||||
body = {
|
||||
"messages": [{"role": "user", "content": "Count from 1 to 3"}],
|
||||
"stream": True,
|
||||
},
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
headers = {"Authorization": f"Bearer {api_key}"},
|
||||
)
|
||||
assert status == 200, f"Expected 200, got {status}"
|
||||
assert len(chunks) > 0, "No SSE chunks received"
|
||||
|
|
@ -183,11 +187,13 @@ def test_openai_sdk(base_url: str, api_key: str):
|
|||
print(" SKIP openai SDK not installed")
|
||||
return
|
||||
|
||||
client = OpenAI(base_url=f"{base_url}/v1", api_key=api_key)
|
||||
client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key)
|
||||
response = client.chat.completions.create(
|
||||
model="current",
|
||||
messages=[{"role": "user", "content": "What is 2+2? Answer with just the number."}],
|
||||
stream=True,
|
||||
model = "current",
|
||||
messages = [
|
||||
{"role": "user", "content": "What is 2+2? Answer with just the number."}
|
||||
],
|
||||
stream = True,
|
||||
)
|
||||
content_parts = []
|
||||
for chunk in response:
|
||||
|
|
@ -211,25 +217,26 @@ def test_curl_with_tools(base_url: str, api_key: str):
|
|||
"""
|
||||
status, chunks = _stream_http(
|
||||
f"{base_url}/v1/chat/completions",
|
||||
body={
|
||||
body = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 123 * 456? Use code to compute it."}
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 123 * 456? Use code to compute it.",
|
||||
}
|
||||
],
|
||||
"stream": True,
|
||||
"enable_tools": True,
|
||||
"enabled_tools": ["python"],
|
||||
"session_id": "test-session",
|
||||
},
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=120,
|
||||
headers = {"Authorization": f"Bearer {api_key}"},
|
||||
timeout = 120,
|
||||
)
|
||||
assert status == 200, f"Expected 200, got {status}"
|
||||
assert len(chunks) > 0, "No SSE chunks received for tools request"
|
||||
|
||||
# Check that at least one chunk has the expected shape
|
||||
has_valid_chunk = any(
|
||||
"choices" in c or "type" in c for c in chunks
|
||||
)
|
||||
has_valid_chunk = any("choices" in c or "type" in c for c in chunks)
|
||||
assert has_valid_chunk, "No valid chunks in tools response"
|
||||
full = _collect_streamed_content(chunks)
|
||||
print(f" PASS curl with tools: {len(chunks)} chunks, {len(full)} chars content")
|
||||
|
|
@ -240,11 +247,11 @@ def test_invalid_key_rejected(base_url: str):
|
|||
status, _text = _http(
|
||||
"POST",
|
||||
f"{base_url}/v1/chat/completions",
|
||||
body={
|
||||
body = {
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"stream": False,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-unsloth-boguskey123"},
|
||||
headers = {"Authorization": "Bearer sk-unsloth-boguskey123"},
|
||||
)
|
||||
assert status == 401, f"Expected 401 for invalid key, got {status}"
|
||||
print(" PASS invalid API key rejected (401)")
|
||||
|
|
@ -255,7 +262,7 @@ def test_no_key_rejected(base_url: str):
|
|||
status, _text = _http(
|
||||
"POST",
|
||||
f"{base_url}/v1/chat/completions",
|
||||
body={
|
||||
body = {
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"stream": False,
|
||||
},
|
||||
|
|
@ -273,22 +280,28 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st
|
|||
Returns (process, api_key).
|
||||
"""
|
||||
cmd = [
|
||||
"unsloth", "studio", "run",
|
||||
"--model", model,
|
||||
"--port", str(PORT),
|
||||
"--host", HOST,
|
||||
"--api-key-name", "test",
|
||||
"unsloth",
|
||||
"studio",
|
||||
"run",
|
||||
"--model",
|
||||
model,
|
||||
"--port",
|
||||
str(PORT),
|
||||
"--host",
|
||||
HOST,
|
||||
"--api-key-name",
|
||||
"test",
|
||||
]
|
||||
if variant:
|
||||
cmd.extend(["--gguf-variant", variant])
|
||||
|
||||
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
LOG_FILE.parent.mkdir(parents = True, exist_ok = True)
|
||||
log_fh = open(LOG_FILE, "w")
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=log_fh,
|
||||
stderr=subprocess.STDOUT,
|
||||
preexec_fn=os.setsid,
|
||||
stdout = log_fh,
|
||||
stderr = subprocess.STDOUT,
|
||||
preexec_fn = os.setsid,
|
||||
)
|
||||
|
||||
# Wait for the banner containing the API key
|
||||
|
|
@ -327,25 +340,31 @@ def _kill_server(proc: subprocess.Popen):
|
|||
except (ProcessLookupError, PermissionError):
|
||||
pass
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
proc.wait(timeout = 10)
|
||||
except subprocess.TimeoutExpired:
|
||||
try:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
except (ProcessLookupError, PermissionError):
|
||||
pass
|
||||
proc.wait(timeout=5)
|
||||
proc.wait(timeout = 5)
|
||||
|
||||
|
||||
# ── Main ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="End-to-end tests for unsloth studio run")
|
||||
parser.add_argument(
|
||||
"--model", default=DEFAULT_MODEL, help=f"Model to test with (default: {DEFAULT_MODEL})"
|
||||
parser = argparse.ArgumentParser(
|
||||
description = "End-to-end tests for unsloth studio run"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gguf-variant", default=DEFAULT_VARIANT, help=f"GGUF variant (default: {DEFAULT_VARIANT})"
|
||||
"--model",
|
||||
default = DEFAULT_MODEL,
|
||||
help = f"Model to test with (default: {DEFAULT_MODEL})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gguf-variant",
|
||||
default = DEFAULT_VARIANT,
|
||||
help = f"GGUF variant (default: {DEFAULT_VARIANT})",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
|
@ -370,7 +389,9 @@ def main():
|
|||
run_test(test_help_output)
|
||||
|
||||
# ── 2-7. Start server and run API tests ──────────────────────────
|
||||
print(f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}...")
|
||||
print(
|
||||
f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..."
|
||||
)
|
||||
proc = None
|
||||
try:
|
||||
proc, api_key = _start_server(args.model, args.gguf_variant)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue