Skip to content
27 changes: 13 additions & 14 deletions tests/integ/sagemaker/serve/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,24 +12,23 @@
# language governing permissions and limitations under the License.
from __future__ import absolute_import

# import os
import os
import platform

# from tests.integ import DATA_DIR
from tests.integ import DATA_DIR

# SERVE_IN_PROCESS_TIMEOUT = 5
# SERVE_MODEL_PACKAGE_TIMEOUT = 10
# SERVE_LOCAL_CONTAINER_TIMEOUT = 10
SERVE_IN_PROCESS_TIMEOUT = 5
SERVE_MODEL_PACKAGE_TIMEOUT = 10
SERVE_LOCAL_CONTAINER_TIMEOUT = 10
SERVE_SAGEMAKER_ENDPOINT_TIMEOUT = 15
# SERVE_SAVE_TIMEOUT = 2
SERVE_SAVE_TIMEOUT = 2

# NOT_RUNNING_ON_PY38 = platform.python_version_tuple()[1] != "8"
NOT_RUNNING_ON_PY310 = platform.python_version_tuple()[1] != "10"
# NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE = os.getenv("TEST_OWNER") != "INF_EXP_DEV"
PYTHON_VERSION_IS_NOT_38 = platform.python_version_tuple()[1] != "8"
PYTHON_VERSION_IS_NOT_310 = platform.python_version_tuple()[1] != "10"

# XGB_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "xgboost")
# PYTORCH_SQUEEZENET_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "pytorch")
# TF_EFFICIENT_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "tensorflow")
# HF_DIR = os.path.join(DATA_DIR, "serve_resources", "hf")
XGB_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "xgboost")
PYTORCH_SQUEEZENET_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "pytorch")
TF_EFFICIENT_RESOURCE_DIR = os.path.join(DATA_DIR, "serve_resources", "tensorflow")
HF_DIR = os.path.join(DATA_DIR, "serve_resources", "hf")

# BYOC_IMAGE_URI_TEMPLATE = "661407751302.dkr.ecr.{}.amazonaws.com/byoc-integ-test-images:{}"
BYOC_IMAGE_URI_TEMPLATE = "661407751302.dkr.ecr.{}.amazonaws.com/byoc-integ-test-images:{}"
5 changes: 3 additions & 2 deletions tests/integ/sagemaker/serve/test_serve_js_happy.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from sagemaker.serve.builder.schema_builder import SchemaBuilder
from tests.integ.sagemaker.serve.constants import (
SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
NOT_RUNNING_ON_PY310,
PYTHON_VERSION_IS_NOT_310,
)

from tests.integ.timeout import timeout
Expand Down Expand Up @@ -47,9 +47,10 @@ def happy_model_builder(sagemaker_session):


@pytest.mark.skipif(
NOT_RUNNING_ON_PY310,
PYTHON_VERSION_IS_NOT_310,
reason="The goal of these test are to test the serving components of our feature",
)
@pytest.mark.slow_test
def test_happy_tgi_sagemaker_endpoint(happy_model_builder, gpu_instance_type):
logger.info("Running in SAGEMAKER_ENDPOINT mode...")
caught_ex = None
Expand Down
Loading