|
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