Zakaria
75 lines
3.2 KiB
Python
75 lines
3.2 KiB
Python
import os
|
|
import time
|
|
import uuid
|
|
|
|
import boto3
|
|
import pytest
|
|
from botocore.config import Config
|
|
|
|
|
|
def client(service):
|
|
return boto3.client(
|
|
service,
|
|
# Every other test file gets this from the conftest.py fixtures, which read
|
|
# FLOCI_ENDPOINT (the test runner container's compat-net hostname for floci, e.g.
|
|
# http://floci:4566). This file has no fixture of its own, so it must read the same
|
|
# variable directly instead of hardcoding localhost, which only resolves to floci
|
|
# when running the suite straight on the host rather than in its own CI container.
|
|
endpoint_url=os.environ.get("FLOCI_ENDPOINT", "http://localhost:4566"),
|
|
region_name="us-east-1",
|
|
aws_access_key_id="test",
|
|
aws_secret_access_key="test",
|
|
config=Config(retries={"max_attempts": 1}),
|
|
)
|
|
|
|
|
|
def test_sagemaker_control_plane_and_training():
|
|
sm = client("sagemaker")
|
|
s3 = client("s3")
|
|
suffix = uuid.uuid4().hex[:8]
|
|
model = f"compat-sm-{suffix}"
|
|
cfg = f"compat-sm-cfg-{suffix}"
|
|
bucket = f"compat-sm-{suffix}"
|
|
|
|
sm.create_model(ModelName=model, PrimaryContainer={"Image": "public.ecr.aws/docker/library/busybox:stable"}, ExecutionRoleArn="arn:aws:iam::000000000000:role/r")
|
|
assert sm.describe_model(ModelName=model)["ModelName"] == model
|
|
sm.create_endpoint_config(EndpointConfigName=cfg, ProductionVariants=[{"VariantName": "AllTraffic", "ModelName": model, "InitialInstanceCount": 1, "InstanceType": "ml.t2.medium"}])
|
|
assert any(c["EndpointConfigName"] == cfg for c in sm.list_endpoint_configs()["EndpointConfigs"])
|
|
|
|
s3.create_bucket(Bucket=bucket)
|
|
s3.put_object(Bucket=bucket, Key="input/data.txt", Body=b"hello")
|
|
job = f"compat-train-{suffix}"
|
|
sm.create_training_job(
|
|
TrainingJobName=job,
|
|
RoleArn="arn:aws:iam::000000000000:role/r",
|
|
AlgorithmSpecification={
|
|
"TrainingImage": "public.ecr.aws/docker/library/busybox:stable",
|
|
"TrainingInputMode": "File",
|
|
"ContainerEntrypoint": ["/bin/sh", "-c"],
|
|
"ContainerArguments": ["mkdir -p /opt/ml/model && echo ok > /opt/ml/model/model.txt"],
|
|
},
|
|
InputDataConfig=[{
|
|
"ChannelName": "train",
|
|
"DataSource": {"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": f"s3://{bucket}/input"}},
|
|
}],
|
|
OutputDataConfig={"S3OutputPath": f"s3://{bucket}/output"},
|
|
ResourceConfig={"InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 1},
|
|
StoppingCondition={"MaxRuntimeInSeconds": 60},
|
|
)
|
|
deadline = time.time() + 120
|
|
while time.time() < deadline:
|
|
desc = sm.describe_training_job(TrainingJobName=job)
|
|
if desc["TrainingJobStatus"] == "Completed":
|
|
artifact = desc["ModelArtifacts"]["S3ModelArtifacts"]
|
|
key = artifact.split(f"s3://{bucket}/", 1)[1]
|
|
assert s3.get_object(Bucket=bucket, Key=key)["Body"].read()
|
|
break
|
|
if desc["TrainingJobStatus"] == "Failed":
|
|
pytest.fail(desc.get("FailureReason", "training failed"))
|
|
time.sleep(1)
|
|
else:
|
|
pytest.fail("Timed out waiting for SageMaker training job")
|
|
|
|
sm.delete_endpoint_config(EndpointConfigName=cfg)
|
|
sm.delete_model(ModelName=model)
|