Skip to content

Commit 89a4c30

Browse files
test sm endpoint
1 parent c001e68 commit 89a4c30

File tree

1 file changed

+15
-1
lines changed

1 file changed

+15
-1
lines changed

test/vllm/sagemaker/test_sm_endpoint.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,28 @@
55
from sagemaker.model import Model
66
from sagemaker.predictor import Predictor
77
from sagemaker import serializers
8-
from ec2.test_artifacts.test_ec2 import get_secret_hf_token
98

109
# Fixed parameters
1110
AWS_REGION = "us-west-2"
1211
INSTANCE_TYPE = "ml.g5.12xlarge"
1312
ROLE = "SageMakerRole"
1413

1514

15+
def get_secret_hf_token():
16+
secret_name = "test/hf_token"
17+
region_name = "us-west-2"
18+
19+
session = boto3.session.Session()
20+
client = session.client(service_name="secretsmanager", region_name=region_name)
21+
try:
22+
get_secret_value_response = client.get_secret_value(SecretId=secret_name)
23+
except ClientError as e:
24+
raise e
25+
26+
response = json.loads(get_secret_value_response["SecretString"])
27+
return response
28+
29+
1630
def deploy_endpoint(name, image_uri, role, instance_type):
1731
try:
1832
model = Model(

0 commit comments

Comments
 (0)