機械学習のために EKS クラスターで AWS Inferentia インスタンスを使用する

Inf1 Amazon EC2 インスタンスを実行するノードを持つクラスターを作成します。inf1.2xlarge を任意の Inf1 インスタンスタイプ に置き換えることができます。eksctl ユーティリティは、Inf1 インスタンスタイプでノードグループを起動していることを検出し、アクセラレーションされた Amazon EKS 最適化 Amazon Linux AMI のいずれかを使用してノードを起動します。

eksctl create cluster \
    --name inferentia \
    --region region-code \
    --nodegroup-name ng-inf1 \
    --node-type inf1.2xlarge \
    --nodes 2 \
    --nodes-min 1 \
    --nodes-max 4 \
    --ssh-access \
    --ssh-public-key your-key \
    --with-oidc

[9]  adding identity "arn:aws:iam::111122223333:role/eksctl-inferentia-nodegroup-ng-in-NodeInstanceRole-FI7HIYS3BS09" to auth ConfigMap

Inf1 インスタンスでノードグループを起動すると、eksctl によって AWS Neuron Kubernetes デバイスプラグインが自動的にインストールされます。このプラグインは、Neuron デバイスをシステムリソースとして Kubernetes スケジューラにアドバタイズします。このスケジューラはコンテナによってリクエストできます。デフォルトの Amazon EKS ノードの IAM ポリシーに加えて、Amazon S3 読み取り専用のアクセスポリシーが追加されます。これにより、後のステップで説明するサンプルアプリケーションで Amazon S3 からトレーニングされたモデルをロードできるようになります。

すべての Pods が正しく起動していることを確認します。

kubectl get pods -n kube-system

簡略化された出力:

NAME                                   READY   STATUS    RESTARTS   AGE
[...]
neuron-device-plugin-daemonset-6djhp   1/1     Running   0          5m
neuron-device-plugin-daemonset-hwjsj   1/1     Running   0          5m

「クラスターを作成する」のステップ 1 で作成したノードのインスタンスロールに AmazonS3ReadOnlyAccess IAM ポリシーを追加します。これは、サンプルアプリケーションが Amazon S3 からトレーニングされたモデルをロードできるようにするために必要です。

aws iam attach-role-policy \
    --policy-arn arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess \
    --role-name eksctl-inferentia-nodegroup-ng-in-NodeInstanceRole-FI7HIYS3BS09

次の内容で、rn50_deployment.yaml という名前のファイルを作成します。希望の設定に合わせて、リージョンコードとモデルパスを更新します。モデル名は、クライアントが TensorFlow サーバーにリクエストを行う際の識別を目的としています。この例では、モデル名を使用してサンプルの ResNet50 のクライアントスクリプトと照合します。このスクリプトは、後のステップで予測リクエストを送信する際に使用されます。

aws ecr list-images --repository-name neuron-rtd --registry-id 790709498068 --region us-west-2

kind: Deployment
apiVersion: apps/v1
metadata:
  name: eks-neuron-test
  labels:
    app: eks-neuron-test
    role: master
spec:
  replicas: 2
  selector:
    matchLabels:
      app: eks-neuron-test
      role: master
  template:
    metadata:
      labels:
        app: eks-neuron-test
        role: master
    spec:
      containers:
        - name: eks-neuron-test
          image: 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-inference-neuron:1.15.4-neuron-py37-ubuntu18.04
          command:
            - /usr/local/bin/entrypoint.sh
          args:
            - --port=8500
            - --rest_api_port=9000
            - --model_name=resnet50_neuron
            - --model_base_path=s3://your-bucket-of-models/resnet50_neuron/
          ports:
            - containerPort: 8500
            - containerPort: 9000
          imagePullPolicy: IfNotPresent
          env:
            - name: AWS_REGION
              value: "us-east-1"
            - name: S3_USE_HTTPS
              value: "1"
            - name: S3_VERIFY_SSL
              value: "0"
            - name: S3_ENDPOINT
              value: s3.us-east-1.amazonaws.com
            - name: AWS_LOG_LEVEL
              value: "3"
          resources:
            limits:
              cpu: 4
              memory: 4Gi
              aws.amazon.com/neuron: 1
            requests:
              cpu: "1"
              memory: 1Gi
          securityContext:
            capabilities:
              add:
                - IPC_LOCK

モデルをデプロイします。

kubectl apply -f rn50_deployment.yaml

次の内容で、rn50_service.yaml というファイルを作成します。予測リクエストを受け入れるために HTTP ポートと gRPC ポートが開かれます。

kind: Service
apiVersion: v1
metadata:
  name: eks-neuron-test
  labels:
    app: eks-neuron-test
spec:
  type: ClusterIP
  ports:
    - name: http-tf-serving
      port: 8500
      targetPort: 8500
    - name: grpc-tf-serving
      port: 9000
      targetPort: 9000
  selector:
    app: eks-neuron-test
    role: master

TensorFlow モデルサービスアプリケーション用の Kubernetes サービスを作成します。

kubectl apply -f rn50_service.yaml

ローカルでテストするには、gRPC ポートを eks-neuron-test サービスに転送します。

kubectl port-forward service/eks-neuron-test 8500:8500 &

次の内容で tensorflow-model-server-infer.py という Python スクリプトを作成します。このスクリプトは、サービスフレームワークである gRPC を介して推論を実行します。

import numpy as np
   import grpc
   import tensorflow as tf
   from tensorflow.keras.preprocessing import image
   from tensorflow.keras.applications.resnet50 import preprocess_input
   from tensorflow_serving.apis import predict_pb2
   from tensorflow_serving.apis import prediction_service_pb2_grpc
   from tensorflow.keras.applications.resnet50 import decode_predictions

   if __name__ == '__main__':
       channel = grpc.insecure_channel('localhost:8500')
       stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
       img_file = tf.keras.utils.get_file(
           "./kitten_small.jpg",
           "https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg")
       img = image.load_img(img_file, target_size=(224, 224))
       img_array = preprocess_input(image.img_to_array(img)[None, ...])
       request = predict_pb2.PredictRequest()
       request.model_spec.name = 'resnet50_inf1'
       request.inputs['input'].CopyFrom(
           tf.make_tensor_proto(img_array, shape=img_array.shape))
       result = stub.Predict(request)
       prediction = tf.make_ndarray(result.outputs['output'])
       print(decode_predictions(prediction))

スクリプトを実行して、予測をサービスに送信します。

python3 tensorflow-model-server-infer.py

出力例は次のとおりです。

[[(u'n02123045', u'tabby', 0.68817204), (u'n02127052', u'lynx', 0.12701613), (u'n02123159', u'tiger_cat', 0.08736559), (u'n02124075', u'Egyptian_cat', 0.063844085), (u'n02128757', u'snow_leopard', 0.009240591)]]