second first commit

2026-09-22 08:57:00 -04:00
parent e88d7e2ff3
commit 0aef7fb2d8
4766 changed files with 1236986 additions and 0 deletions
@@ -0,0 +1,74 @@
+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)