[Test] Run SageMaker handler-override tests in-process via TestClient (#47250)
Signed-off-by: Jyothirmai Kottu <jkottu@amazon.com> Co-authored-by: Claude <noreply@anthropic.com>
Jyothirmai Kottu committed
Jul 1, 2026 at 02:14 UTC
dee5da1dec8c508cba430fd8536510afd7b80f60
1 file changed
+227
-333
tests/entrypoints/serve/sagemaker/test_sagemaker_handler_overrides.py
+227
-333
@@ -12,15 +12,21 @@ Tests real customer usage scenarios:
12
- Priority: env vars > decorators > customer script files > framework
13
defaults
14
15
-Note: These tests focus on validating server responses rather than directly calling
16
-get_ping_handler() and get_invoke_handler() to ensure full integration testing.
15
+The handler-override scenarios exercise the real vLLM SageMaker router and
16
+bootstrap path via an in-process FastAPI ``TestClient`` instead of launching a
17
+model server. These scenarios fully replace the ``/ping`` and ``/invocations``
18
+endpoints with customer handlers, so no inference engine is required to
19
+validate override behavior. Avoiding the model server also keeps the tests
20
+fast and deterministic rather than depending on the FastAPI version resolved
21
+into the test environment at runtime.
22
"""
23
24
import os
20
-import tempfile
25
26
import pytest
27
import requests
28
+from fastapi import FastAPI
29
+from fastapi.testclient import TestClient
30
31
from tests.utils import RemoteOpenAIServer
32
@@ -29,6 +35,23 @@ from .conftest import (
35
)
36
37
38
+def _build_sagemaker_test_client() -> TestClient:
39
+ """Build a TestClient over the real SageMaker router and bootstrap path.
40
+
41
+ ``attach_router`` is called with empty supported tasks because the override
42
+ tests replace the endpoints with customer handlers, so no framework
43
+ invocation handler (and therefore no engine) is exercised.
44
+ """
45
+ from vllm.entrypoints.serve.sagemaker.api_router import (
46
+ attach_router,
47
+ sagemaker_standards_bootstrap,
48
+ )
49
+
50
+ app = FastAPI()
51
+ attach_router(app, ())
52
+ return TestClient(sagemaker_standards_bootstrap(app))
53
+
54
+
55
class TestHandlerOverrideIntegration:
56
"""Integration tests simulating real customer usage scenarios.
57
@@ -89,8 +112,7 @@ class TestHandlerOverrideIntegration:
112
except ImportError:
113
pass
114
92
- @pytest.mark.asyncio
93
- async def test_customer_script_functions_auto_loaded(self):
115
+ def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path):
116
"""Test customer scenario: script functions automatically override
117
framework defaults."""
118
try:
@@ -101,15 +123,15 @@ class TestHandlerOverrideIntegration:
123
pytest.skip("model-hosting-container-standards not available")
124
125
# Customer writes a script file with ping() and invoke() functions
104
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
105
- f.write(
106
- """
126
+ script_path = tmp_path / "model.py"
127
+ script_path.write_text(
128
+ """
129
from fastapi import Request
130
131
async def custom_sagemaker_ping_handler():
132
return {
133
"status": "healthy",
112
- "source": "customer_override",
134
+ "source": "customer_override",
135
"message": "Custom ping from customer script"
136
}
137
@@ -119,62 +141,39 @@ async def custom_sagemaker_invocation_handler(request: Request):
141
"source": "customer_override"
142
}
143
"""
122
- )
123
- script_path = f.name
144
+ )
145
125
- try:
126
- script_dir = os.path.dirname(script_path)
127
- script_name = os.path.basename(script_path)
146
+ # Customer sets SageMaker environment variables to point to their script
147
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
148
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
149
129
- # Customer sets SageMaker environment variables to point to their script
130
- env_vars = {
131
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
132
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
133
- }
150
+ with _build_sagemaker_test_client() as client:
151
+ # Customer tests their server and sees their overrides work
152
+ # automatically
153
+ ping_response = client.get("/ping")
154
+ assert ping_response.status_code == 200
155
+ ping_data = ping_response.json()
156
135
- args = [
136
- "--dtype",
137
- "bfloat16",
138
- "--max-model-len",
139
- "2048",
140
- "--enforce-eager",
141
- "--max-num-seqs",
142
- "32",
157
+ invoke_response = client.post(
158
+ "/invocations",
159
+ json={
160
+ "model": MODEL_NAME_SMOLLM,
161
+ "messages": [{"role": "user", "content": "Hello"}],
162
+ "max_tokens": 5,
163
+ },
164
+ )
165
+ assert invoke_response.status_code == 200
166
+ invoke_data = invoke_response.json()
167
+
168
+ # Customer sees their functions are used
169
+ assert ping_data["source"] == "customer_override"
170
+ assert ping_data["message"] == "Custom ping from customer script"
171
+ assert invoke_data["source"] == "customer_override"
172
+ assert invoke_data["predictions"] == [
173
+ "Custom response from customer script"
174
]
175
145
- with RemoteOpenAIServer(
146
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
147
- ) as server:
148
- # Customer tests their server and sees their overrides work
149
- # automatically
150
- ping_response = requests.get(server.url_for("ping"))
151
- assert ping_response.status_code == 200
152
- ping_data = ping_response.json()
153
-
154
- invoke_response = requests.post(
155
- server.url_for("invocations"),
156
- json={
157
- "model": MODEL_NAME_SMOLLM,
158
- "messages": [{"role": "user", "content": "Hello"}],
159
- "max_tokens": 5,
160
- },
161
- )
162
- assert invoke_response.status_code == 200
163
- invoke_data = invoke_response.json()
164
-
165
- # Customer sees their functions are used
166
- assert ping_data["source"] == "customer_override"
167
- assert ping_data["message"] == "Custom ping from customer script"
168
- assert invoke_data["source"] == "customer_override"
169
- assert invoke_data["predictions"] == [
170
- "Custom response from customer script"
171
- ]
172
-
173
- finally:
174
- os.unlink(script_path)
175
-
176
- @pytest.mark.asyncio
177
- async def test_customer_decorator_usage(self):
176
+ def test_customer_decorator_usage(self, monkeypatch, tmp_path):
177
"""Test customer scenario: using @custom_ping_handler and
178
@custom_invocation_handler decorators."""
179
try:
@@ -185,9 +184,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
184
pytest.skip("model-hosting-container-standards not available")
185
186
# Customer writes a script file with decorators
188
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
189
- f.write(
190
- """
187
+ script_path = tmp_path / "model.py"
188
+ script_path.write_text(
189
+ """
190
import model_hosting_container_standards.sagemaker as sagemaker_standards
191
from fastapi import Request
192
@@ -198,62 +197,39 @@ async def my_ping():
197
"source": "customer_decorator"
198
}
199
201
-@sagemaker_standards.custom_invocation_handler
200
+@sagemaker_standards.custom_invocation_handler
201
async def my_invoke(request: Request):
202
return {
204
- "type": "invoke",
203
+ "type": "invoke",
204
"source": "customer_decorator"
205
}
206
"""
208
- )
209
- script_path = f.name
207
+ )
208
211
- try:
212
- script_dir = os.path.dirname(script_path)
213
- script_name = os.path.basename(script_path)
209
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
210
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
211
215
- env_vars = {
216
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
217
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
218
- }
212
+ with _build_sagemaker_test_client() as client:
213
+ ping_response = client.get("/ping")
214
+ assert ping_response.status_code == 200
215
+ ping_data = ping_response.json()
216
220
- args = [
221
- "--dtype",
222
- "bfloat16",
223
- "--max-model-len",
224
- "2048",
225
- "--enforce-eager",
226
- "--max-num-seqs",
227
- "32",
228
- ]
217
+ invoke_response = client.post(
218
+ "/invocations",
219
+ json={
220
+ "model": MODEL_NAME_SMOLLM,
221
+ "messages": [{"role": "user", "content": "Hello"}],
222
+ "max_tokens": 5,
223
+ },
224
+ )
225
+ assert invoke_response.status_code == 200
226
+ invoke_data = invoke_response.json()
227
230
- with RemoteOpenAIServer(
231
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
232
- ) as server:
233
- ping_response = requests.get(server.url_for("ping"))
234
- assert ping_response.status_code == 200
235
- ping_data = ping_response.json()
236
-
237
- invoke_response = requests.post(
238
- server.url_for("invocations"),
239
- json={
240
- "model": MODEL_NAME_SMOLLM,
241
- "messages": [{"role": "user", "content": "Hello"}],
242
- "max_tokens": 5,
243
- },
244
- )
245
- assert invoke_response.status_code == 200
246
- invoke_data = invoke_response.json()
247
-
248
- # Customer sees their handlers are used by the server
249
- assert ping_data["source"] == "customer_decorator"
250
- assert invoke_data["source"] == "customer_decorator"
251
-
252
- finally:
253
- os.unlink(script_path)
228
+ # Customer sees their handlers are used by the server
229
+ assert ping_data["source"] == "customer_decorator"
230
+ assert invoke_data["source"] == "customer_decorator"
231
255
- @pytest.mark.asyncio
256
- async def test_handler_priority_order(self):
232
+ def test_handler_priority_order(self, monkeypatch, tmp_path):
233
"""Test priority: @custom_ping_handler/@custom_invocation_handler
234
decorators vs script functions."""
235
try:
@@ -264,9 +240,9 @@ async def my_invoke(request: Request):
240
pytest.skip("model-hosting-container-standards not available")
241
242
# Customer writes a script with both decorator and regular functions
267
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
268
- f.write(
269
- """
243
+ script_path = tmp_path / "model.py"
244
+ script_path.write_text(
245
+ """
246
import model_hosting_container_standards.sagemaker as sagemaker_standards
247
from fastapi import Request
248
@@ -275,7 +251,7 @@ from fastapi import Request
251
async def decorated_ping():
252
return {
253
"status": "healthy",
278
- "source": "ping_decorator_in_script",
254
+ "source": "ping_decorator_in_script",
255
"priority": "decorator"
256
}
257
@@ -296,60 +272,37 @@ async def custom_sagemaker_invocation_handler(request: Request):
272
"priority": "function"
273
}
274
"""
299
- )
300
- script_path = f.name
275
+ )
276
302
- try:
303
- script_dir = os.path.dirname(script_path)
304
- script_name = os.path.basename(script_path)
277
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
278
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
279
306
- env_vars = {
307
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
308
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
309
- }
280
+ with _build_sagemaker_test_client() as client:
281
+ ping_response = client.get("/ping")
282
+ assert ping_response.status_code == 200
283
+ ping_data = ping_response.json()
284
311
- args = [
312
- "--dtype",
313
- "bfloat16",
314
- "--max-model-len",
315
- "2048",
316
- "--enforce-eager",
317
- "--max-num-seqs",
318
- "32",
319
- ]
285
+ invoke_response = client.post(
286
+ "/invocations",
287
+ json={
288
+ "model": MODEL_NAME_SMOLLM,
289
+ "messages": [{"role": "user", "content": "Hello"}],
290
+ "max_tokens": 5,
291
+ },
292
+ )
293
+ assert invoke_response.status_code == 200
294
+ invoke_data = invoke_response.json()
295
321
- with RemoteOpenAIServer(
322
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
323
- ) as server:
324
- ping_response = requests.get(server.url_for("ping"))
325
- assert ping_response.status_code == 200
326
- ping_data = ping_response.json()
327
-
328
- invoke_response = requests.post(
329
- server.url_for("invocations"),
330
- json={
331
- "model": MODEL_NAME_SMOLLM,
332
- "messages": [{"role": "user", "content": "Hello"}],
333
- "max_tokens": 5,
334
- },
335
- )
336
- assert invoke_response.status_code == 200
337
- invoke_data = invoke_response.json()
338
-
339
- # @custom_ping_handler decorator has higher priority than
340
- # script function
341
- assert ping_data["source"] == "ping_decorator_in_script"
342
- assert ping_data["priority"] == "decorator"
343
-
344
- # Script function is used for invoke
345
- assert invoke_data["source"] == "script_invoke_function"
346
- assert invoke_data["priority"] == "function"
347
-
348
- finally:
349
- os.unlink(script_path)
296
+ # @custom_ping_handler decorator has higher priority than
297
+ # script function
298
+ assert ping_data["source"] == "ping_decorator_in_script"
299
+ assert ping_data["priority"] == "decorator"
300
351
- @pytest.mark.asyncio
352
- async def test_environment_variable_script_loading(self):
301
+ # Script function is used for invoke
302
+ assert invoke_data["source"] == "script_invoke_function"
303
+ assert invoke_data["priority"] == "function"
304
+
305
+ def test_environment_variable_script_loading(self, monkeypatch, tmp_path):
306
"""Test that environment variables correctly specify script location
307
and loading."""
308
try:
@@ -360,9 +313,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
313
pytest.skip("model-hosting-container-standards not available")
314
315
# Customer writes a script in a specific directory
363
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
364
- f.write(
365
- """
316
+ script_path = tmp_path / "model.py"
317
+ script_path.write_text(
318
+ """
319
from fastapi import Request
320
321
async def custom_sagemaker_ping_handler():
@@ -379,60 +332,43 @@ async def custom_sagemaker_invocation_handler(request: Request):
332
"method": "environment_variable_loading"
333
}
334
"""
382
- )
383
- script_path = f.name
335
+ )
336
385
- try:
386
- script_dir = os.path.dirname(script_path)
387
- script_name = os.path.basename(script_path)
337
+ # Test environment variable script loading
338
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
339
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
340
389
- # Test environment variable script loading
390
- env_vars = {
391
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
392
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
393
- }
341
+ with _build_sagemaker_test_client() as client:
342
+ ping_response = client.get("/ping")
343
+ assert ping_response.status_code == 200
344
+ ping_data = ping_response.json()
345
395
- args = [
396
- "--dtype",
397
- "bfloat16",
398
- "--max-model-len",
399
- "2048",
400
- "--enforce-eager",
401
- "--max-num-seqs",
402
- "32",
403
- ]
346
+ invoke_response = client.post(
347
+ "/invocations",
348
+ json={
349
+ "model": MODEL_NAME_SMOLLM,
350
+ "messages": [{"role": "user", "content": "Hello"}],
351
+ "max_tokens": 5,
352
+ },
353
+ )
354
+ assert invoke_response.status_code == 200
355
+ invoke_data = invoke_response.json()
356
405
- with RemoteOpenAIServer(
406
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
407
- ) as server:
408
- ping_response = requests.get(server.url_for("ping"))
409
- assert ping_response.status_code == 200
410
- ping_data = ping_response.json()
411
-
412
- invoke_response = requests.post(
413
- server.url_for("invocations"),
414
- json={
415
- "model": MODEL_NAME_SMOLLM,
416
- "messages": [{"role": "user", "content": "Hello"}],
417
- "max_tokens": 5,
418
- },
419
- )
420
- assert invoke_response.status_code == 200
421
- invoke_data = invoke_response.json()
422
-
423
- # Verify that the script was loaded via environment variables
424
- assert ping_data["source"] == "env_loaded_script"
425
- assert ping_data["method"] == "environment_variable_loading"
426
- assert invoke_data["source"] == "env_loaded_script"
427
- assert invoke_data["method"] == "environment_variable_loading"
428
-
429
- finally:
430
- os.unlink(script_path)
357
+ # Verify that the script was loaded via environment variables
358
+ assert ping_data["source"] == "env_loaded_script"
359
+ assert ping_data["method"] == "environment_variable_loading"
360
+ assert invoke_data["source"] == "env_loaded_script"
361
+ assert invoke_data["method"] == "environment_variable_loading"
362
363
@pytest.mark.asyncio
364
async def test_framework_default_handlers(self):
365
"""Test that framework default handlers work when no customer
435
- overrides exist."""
366
+ overrides exist.
367
+
368
+ This scenario exercises the real inference path (default
369
+ ``/invocations``), so it keeps using a live model server rather than
370
+ the in-process TestClient.
371
+ """
372
args = [
373
"--dtype",
374
"bfloat16",
@@ -478,8 +414,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
414
)
415
assert invoke_response.status_code == 200
416
481
- @pytest.mark.asyncio
482
- async def test_handler_env_var_override(self):
417
+ def test_handler_env_var_override(self, monkeypatch, tmp_path):
418
"""Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
419
environment variable overrides."""
420
try:
@@ -493,9 +428,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
428
pytest.skip("model-hosting-container-standards not available")
429
430
# Create a script with both env var handlers and script functions
496
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
497
- f.write(
498
- """
431
+ script_path = tmp_path / "model.py"
432
+ script_path.write_text(
433
+ """
434
from fastapi import Request, Response
435
import json
436
@@ -533,68 +468,47 @@ async def custom_sagemaker_invocation_handler(request: Request):
468
"method": "script_function"
469
}
470
"""
536
- )
537
- script_path = f.name
471
+ )
472
+
473
+ # Set environment variables to override both handlers
474
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
475
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
476
+ monkeypatch.setenv(
477
+ FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
478
+ f"{script_path.name}:env_var_ping_handler",
479
+ )
480
+ monkeypatch.setenv(
481
+ FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
482
+ f"{script_path.name}:env_var_invoke_handler",
483
+ )
484
+
485
+ with _build_sagemaker_test_client() as client:
486
+ # Test ping handler override
487
+ ping_response = client.get("/ping")
488
+ assert ping_response.status_code == 200
489
+ ping_data = ping_response.json()
490
539
- try:
540
- script_dir = os.path.dirname(script_path)
541
- script_name = os.path.basename(script_path)
542
-
543
- # Set environment variables to override both handlers
544
- env_vars = {
545
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
546
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
547
- FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
548
- f"{script_name}:env_var_ping_handler"
549
- ),
550
- FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
551
- f"{script_name}:env_var_invoke_handler"
552
- ),
553
- }
491
+ # Environment variable should override script function
492
+ assert ping_data["method"] == "environment_variable"
493
+ assert ping_data["source"] == "env_var_ping"
494
555
- args = [
556
- "--dtype",
557
- "bfloat16",
558
- "--max-model-len",
559
- "2048",
560
- "--enforce-eager",
561
- "--max-num-seqs",
562
- "32",
563
- ]
495
+ # Test invocation handler override
496
+ invoke_response = client.post(
497
+ "/invocations",
498
+ json={
499
+ "model": MODEL_NAME_SMOLLM,
500
+ "messages": [{"role": "user", "content": "Hello"}],
501
+ "max_tokens": 5,
502
+ },
503
+ )
504
+ assert invoke_response.status_code == 200
505
+ invoke_data = invoke_response.json()
506
565
- with RemoteOpenAIServer(
566
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
567
- ) as server:
568
- # Test ping handler override
569
- ping_response = requests.get(server.url_for("ping"))
570
- assert ping_response.status_code == 200
571
- ping_data = ping_response.json()
572
-
573
- # Environment variable should override script function
574
- assert ping_data["method"] == "environment_variable"
575
- assert ping_data["source"] == "env_var_ping"
576
-
577
- # Test invocation handler override
578
- invoke_response = requests.post(
579
- server.url_for("invocations"),
580
- json={
581
- "model": MODEL_NAME_SMOLLM,
582
- "messages": [{"role": "user", "content": "Hello"}],
583
- "max_tokens": 5,
584
- },
585
- )
586
- assert invoke_response.status_code == 200
587
- invoke_data = invoke_response.json()
588
-
589
- # Environment variable should override script function
590
- assert invoke_data["method"] == "environment_variable"
591
- assert invoke_data["source"] == "env_var_invoke"
592
-
593
- finally:
594
- os.unlink(script_path)
507
+ # Environment variable should override script function
508
+ assert invoke_data["method"] == "environment_variable"
509
+ assert invoke_data["source"] == "env_var_invoke"
510
596
- @pytest.mark.asyncio
597
- async def test_env_var_priority_over_decorator_and_script(self):
511
+ def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path):
512
"""Test that environment variables have highest priority over decorators
513
and script functions for both ping and invocation handlers."""
514
try:
@@ -608,9 +522,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
522
pytest.skip("model-hosting-container-standards not available")
523
524
# Create a script with all three handler types for both ping and invocation
611
- with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
612
- f.write(
613
- """
525
+ script_path = tmp_path / "model.py"
526
+ script_path.write_text(
527
+ """
528
import model_hosting_container_standards.sagemaker as sagemaker_standards
529
from fastapi import Request, Response
530
import json
@@ -674,62 +588,42 @@ async def custom_sagemaker_invocation_handler(request: Request):
588
"priority": "script_function"
589
}
590
"""
677
- )
678
- script_path = f.name
591
+ )
592
+
593
+ # Set environment variables to specify highest priority handlers
594
+ monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
595
+ monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
596
+ monkeypatch.setenv(
597
+ FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
598
+ f"{script_path.name}:env_priority_ping",
599
+ )
600
+ monkeypatch.setenv(
601
+ FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
602
+ f"{script_path.name}:env_priority_invoke",
603
+ )
604
+
605
+ with _build_sagemaker_test_client() as client:
606
+ # Test ping handler priority
607
+ ping_response = client.get("/ping")
608
+ assert ping_response.status_code == 200
609
+ ping_data = ping_response.json()
610
680
- try:
681
- script_dir = os.path.dirname(script_path)
682
- script_name = os.path.basename(script_path)
683
-
684
- # Set environment variables to specify highest priority handlers
685
- env_vars = {
686
- SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
687
- SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
688
- FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
689
- f"{script_name}:env_priority_ping"
690
- ),
691
- FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
692
- f"{script_name}:env_priority_invoke"
693
- ),
694
- }
611
+ # Environment variable has highest priority and should be used
612
+ assert ping_data["priority"] == "environment_variable"
613
+ assert ping_data["source"] == "env_var"
614
696
- args = [
697
- "--dtype",
698
- "bfloat16",
699
- "--max-model-len",
700
- "2048",
701
- "--enforce-eager",
702
- "--max-num-seqs",
703
- "32",
704
- ]
615
+ # Test invocation handler priority
616
+ invoke_response = client.post(
617
+ "/invocations",
618
+ json={
619
+ "model": MODEL_NAME_SMOLLM,
620
+ "messages": [{"role": "user", "content": "Hello"}],
621
+ "max_tokens": 5,
622
+ },
623
+ )
624
+ assert invoke_response.status_code == 200
625
+ invoke_data = invoke_response.json()
626
706
- with RemoteOpenAIServer(
707
- MODEL_NAME_SMOLLM, args, env_dict=env_vars
708
- ) as server:
709
- # Test ping handler priority
710
- ping_response = requests.get(server.url_for("ping"))
711
- assert ping_response.status_code == 200
712
- ping_data = ping_response.json()
713
-
714
- # Environment variable has highest priority and should be used
715
- assert ping_data["priority"] == "environment_variable"
716
- assert ping_data["source"] == "env_var"
717
-
718
- # Test invocation handler priority
719
- invoke_response = requests.post(
720
- server.url_for("invocations"),
721
- json={
722
- "model": MODEL_NAME_SMOLLM,
723
- "messages": [{"role": "user", "content": "Hello"}],
724
- "max_tokens": 5,
725
- },
726
- )
727
- assert invoke_response.status_code == 200
728
- invoke_data = invoke_response.json()
729
-
730
- # Environment variable has highest priority and should be used
731
- assert invoke_data["priority"] == "environment_variable"
732
- assert invoke_data["source"] == "env_var"
733
-
734
- finally:
735
- os.unlink(script_path)
627
+ # Environment variable has highest priority and should be used
628
+ assert invoke_data["priority"] == "environment_variable"
629
+ assert invoke_data["source"] == "env_var"