AWS SageMaker Serverless

Mix.install([
  {:aws, "~> 1.0"},
  {:ex_aws, "~> 2.5"},
  {:ex_aws_sts, "~> 2.3"},
  {:ex_aws_s3, "~> 2.5"},
  {:req, "~> 0.5"},
  {:nx, "~> 0.9"},
  {:stb_image, "~> 0.6"},
  {:hackney, "~> 1.20"},
  {:sweet_xml, "~> 0.7"},
  {:kino, "~> 0.14"},
  {:flow, "~> 1.2"}
])

認証

access_key_id_input = Kino.Input.password("ACCESS_KEY_ID")
secret_access_key_input = Kino.Input.password("SECRET_ACCESS_KEY")
region_input = Kino.Input.text("REGION")

[
  access_key_id_input,
  secret_access_key_input,
  region_input
]
|> Kino.Layout.grid(columns: 3)

client =
  AWS.Client.create(
    Kino.Input.read(access_key_id_input),
    Kino.Input.read(secret_access_key_input),
    Kino.Input.read(region_input)
  )

auth_config = [
  access_key_id: Kino.Input.read(access_key_id_input),
  secret_access_key: Kino.Input.read(secret_access_key_input),
  region: Kino.Input.read(region_input)
]

Kino.nothing()

アカウントID取得

account_id =
  ExAws.STS.get_caller_identity()
  |> ExAws.request!(auth_config)
  |> then(& &1.body.account)

ECRリポジトリー作成

region = Kino.Input.read(region_input)

image = "sagemaker-phoenix"

fullname = "#{account_id}.dkr.ecr.#{region}.amazonaws.com/#{image}:latest"

client
|> AWS.ECR.create_repository(%{
  "repositoryName" => image
})

client
|> AWS.ECR.describe_repositories(%{})

client
|> AWS.ECR.set_repository_policy(%{
  "repositoryName" => image,
  "policyText" =>
    Jason.encode!(%{
      "Statement" => [
        %{
          "Sid" => "ECR202201051440",
          "Effect" => "Allow",
          "Principal" => %{
            "AWS" => ["*"]
          },
          "Action" => "ecr:*"
        }
      ]
    })
})

ECR リポジトリーへのイメージプッシュ

token =
  client
  |> AWS.ECR.get_authorization_token(%{})
  |> elem(1)
  |> Map.get("authorizationData")
  |> Enum.at(0)
  |> Map.get("authorizationToken")

[username, password] =
  token
  |> Base.decode64!()
  |> String.split(":")

System.cmd(
  "docker",
  [
    "login",
    "--username",
    username,
    "--password",
    password,
    fullname
  ]
)

{result, 0} =
  System.cmd(
    "docker",
    [
      "build",
      "--platform",
      "linux/amd64",
      "-t",
      image,
      "/sagemaker/serve"
    ]
  )

Kino.Markdown.new(result)

{result, 0} =
  System.cmd(
    "docker",
    [
      "tag",
      image,
      fullname
    ]
  )

{result, 0} =
  System.cmd(
    "docker",
    [
      "push",
      fullname
    ]
  )

Kino.Markdown.new(result)

client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)

S3へのモデルファイルアップロード

アップロード先の指定

ExAws.S3.list_buckets()
|> ExAws.request!(auth_config)
|> then(& &1.body.buckets)
|> Kino.DataTable.new()

bucket_name_input = Kino.Input.text("BUCKET_NAME")

bucket_name = Kino.Input.read(bucket_name_input)
model_prefix = "models/"

モデルファイルの圧縮

models_path = "/sagemaker/serve/models"

filenames =
  models_path
  |> File.ls!()
  |> Enum.map(fn filename ->
    {
      to_charlist(filename),
      to_charlist(Path.join(models_path, filename))
    }
  end)

tar_filename = "models.tar.gz"

:erl_tar.create(tar_filename, filenames, [:compressed])

アップロードの実行

tar_filename
|> ExAws.S3.Upload.stream_file()
|> ExAws.S3.upload(bucket_name, model_prefix <> tar_filename)
|> ExAws.request!(auth_config)

bucket_name
|> ExAws.S3.list_objects(prefix: model_prefix)
|> ExAws.request!(auth_config)
|> then(& &1.body.contents)
|> Kino.DataTable.new()

IAM ロールの定義

client
|> AWS.IAM.create_role(%{
  "RoleName" => "sagemaker-phoenix-role",
  "AssumeRolePolicyDocument" =>
    Jason.encode!(%{
      "Statement" => [
        %{
          "Sid" => "STS202201051440",
          "Effect" => "Allow",
          "Principal" => %{
            "Service" => ["sagemaker.amazonaws.com"]
          },
          "Action" => "sts:AssumeRole"
        }
      ]
    })
})

client
|> AWS.IAM.create_policy(%{
  "PolicyName" => "sagemaker-phoenix-role-policy",
  "PolicyDocument" =>
    Jason.encode!(%{
      "Version" => "2012-10-17",
      "Statement" => [
        %{
          "Effect" => "Allow",
          "Action" => [
            "cloudwatch:PutMetricData",
            "logs:CreateLogStream",
            "logs:PutLogEvents",
            "logs:CreateLogGroup",
            "logs:DescribeLogStreams",
            "ecr:GetAuthorizationToken"
          ],
          "Resource" => ["*"]
        },
        %{
          "Effect" => "Allow",
          "Action" => [
            "s3:GetObject"
          ],
          "Resource" => [
            "arn:aws:s3:::#{bucket_name}/*"
          ]
        },
        %{
          "Effect" => "Allow",
          "Action" => [
            "ecr:BatchCheckLayerAvailability",
            "ecr:GetDownloadUrlForLayer",
            "ecr:BatchGetImage"
          ],
          "Resource" => [
            "arn:aws:ecr:::repository/#{image}"
          ]
        }
      ]
    })
})

client
|> AWS.IAM.get_policy_version(%{
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy",
  "VersionId" => "v1"
})
|> elem(1)
|> then(& &1["GetPolicyVersionResponse"]["GetPolicyVersionResult"]["PolicyVersion"]["Document"])
|> URI.decode()
|> Jason.decode!()

client
|> AWS.IAM.attach_role_policy(%{
  "RoleName" => "sagemaker-phoenix-role",
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

SageMaker モデル作成

client
|> AWS.SageMaker.create_model(%{
  "ModelName" => "sagemaker-phoenix-model",
  "ExecutionRoleArn" => "arn:aws:iam::#{account_id}:role/sagemaker-phoenix-role",
  "PrimaryContainer" => %{
    "Image" => fullname,
    "ModelDataUrl" => "s3://#{bucket_name}/#{model_prefix}#{tar_filename}"
  }
})

SageMaker エンドポイント設定作成

client
|> AWS.SageMaker.create_endpoint_config(%{
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config",
  "ProductionVariants" => [
    %{
      "VariantName" => "variant-1",
      "ModelName" => "sagemaker-phoenix-model",
      "ServerlessConfig" => %{ 
        "MaxConcurrency" => 4,
        "MemorySizeInMB" => 6144,
        "ProvisionedConcurrency" => 1
      }
    }
  ]
})

SageMaker エンドポイント作成

client
|> AWS.SageMaker.create_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint",
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config"
})

client
|> AWS.SageMaker.describe_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})
|> elem(1)

client
|> AWS.SageMaker.describe_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})
|> elem(1)
|> Map.get("EndpointStatus")

SageMaker エンドポイント呼出

binary =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
  |> Req.get!()
  |> Map.get(:body)

binary
|> StbImage.read_binary!()
|> StbImage.to_nx()
|> Kino.Image.new()

client
|> AWS.SageMakerRuntime.invoke_endpoint(
  "sagemaker-phoenix-endpoint",
  %{
    "Body" => binary,
    "ContentType" => "image/jpeg",
    "Accept" => "application/json"
  },
  recv_timeout: 60_000
)
|> elem(1)
|> Map.get("Body")
|> Jason.decode!()
|> Map.get("predictions")
|> Kino.DataTable.new()

0..10
|> Flow.from_enumerable(max_demand: 1)
|> Flow.map(fn _index ->
  client
  |> AWS.SageMakerRuntime.invoke_endpoint(
    "sagemaker-phoenix-endpoint",
    %{
      "Body" => binary,
      "ContentType" => "image/jpeg",
      "Accept" => "application/json"
    },
    recv_timeout: 60_000
  )
end)
|> Enum.to_list()

image_input = Kino.Input.image("INPUT_IMAGE", format: :jpeg)

binary =
  image_input
  |> Kino.Input.read()
  |> Map.get(:file_ref)
  |> Kino.Input.file_path()
  |> File.read!()

client
|> AWS.SageMakerRuntime.invoke_endpoint(
  "sagemaker-phoenix-endpoint",
  %{
    "Body" => binary,
    "ContentType" => "image/jpeg",
    "Accept" => "application/json"
  },
  recv_timeout: 60_000
)
|> elem(1)
|> Map.get("Body")
|> Jason.decode!()
|> Map.get("predictions")
|> Kino.DataTable.new()

SageMaker エンドポイント削除

client
|> AWS.SageMaker.delete_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})

client
|> AWS.SageMaker.delete_endpoint_config(%{
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config"
})

client
|> AWS.SageMaker.delete_model(%{
  "ModelName" => "sagemaker-phoenix-model"
})

IAM ロール削除

client
|> AWS.IAM.detach_role_policy(%{
  "RoleName" => "sagemaker-phoenix-role",
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

client
|> AWS.IAM.delete_policy(%{
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

client
|> AWS.IAM.delete_role(%{
  "RoleName" => "sagemaker-phoenix-role"
})

ECR リポジトリー削除

image_digest_list =
  client
  |> AWS.ECR.list_images(%{
    "repositoryName" => image
  })
  |> elem(1)
  |> Map.get("imageIds")
  |> Enum.map(&Map.take(&1,["imageDigest"]))

client
|> AWS.ECR.batch_delete_image(%{
  "repositoryName" => image,
  "imageIds" => image_digest_list
})

client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)
|> Map.get("imageIds")
|> Kino.DataTable.new()

client
|> AWS.ECR.delete_repository(%{
  "repositoryName" => image
})

モデルファイル削除

ExAws.S3.delete_object(bucket_name, model_prefix <> tar_filename)
|> ExAws.request!(auth_config)