|
14 | 14 |
|
15 | 15 | import pytest
|
16 | 16 |
|
17 |
| -from sagemaker.serve import Mode |
18 | 17 | from sagemaker.serve.builder.model_builder import ModelBuilder
|
19 | 18 | from sagemaker.serve.builder.schema_builder import SchemaBuilder
|
20 | 19 | from tests.integ.sagemaker.serve.constants import (
|
21 | 20 | SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
|
22 |
| - SERVE_LOCAL_CONTAINER_TUNE_TIMEOUT, |
23 | 21 | PYTHON_VERSION_IS_NOT_310,
|
24 | 22 | )
|
25 | 23 |
|
@@ -77,32 +75,3 @@ def test_happy_tgi_sagemaker_endpoint(happy_model_builder, gpu_instance_type):
|
77 | 75 | )
|
78 | 76 | if caught_ex:
|
79 | 77 | raise caught_ex
|
80 |
| - |
81 |
| - |
82 |
| -@pytest.mark.skipif( |
83 |
| - PYTHON_VERSION_IS_NOT_310, |
84 |
| - reason="The goal of these tests are to test the serving components of our feature", |
85 |
| -) |
86 |
| -@pytest.mark.local_mode |
87 |
| -def test_happy_tune_tgi_local_mode(sagemaker_local_session): |
88 |
| - logger.info("Running in LOCAL_CONTAINER mode...") |
89 |
| - caught_ex = None |
90 |
| - |
91 |
| - model_builder = ModelBuilder( |
92 |
| - model="huggingface-llm-bilingual-rinna-4b-instruction-ppo-bf16", |
93 |
| - schema_builder=SchemaBuilder(SAMPLE_PROMPT, SAMPLE_RESPONSE), |
94 |
| - mode=Mode.LOCAL_CONTAINER, |
95 |
| - sagemaker_session=sagemaker_local_session, |
96 |
| - ) |
97 |
| - |
98 |
| - model = model_builder.build() |
99 |
| - |
100 |
| - with timeout(minutes=SERVE_LOCAL_CONTAINER_TUNE_TIMEOUT): |
101 |
| - try: |
102 |
| - tuned_model = model.tune() |
103 |
| - assert tuned_model.env is not None |
104 |
| - except Exception as e: |
105 |
| - caught_ex = e |
106 |
| - finally: |
107 |
| - if caught_ex: |
108 |
| - raise caught_ex |
0 commit comments