container_image_uri = sagemaker.image_uris.retrieve(
    "pytorch",  # PyTorch のマネージドコンテナを利用
    sagemaker.session.Session().boto_region_name, # ECR のリージョンを指定
    version='1.12', # PyTorch のバージョンを指定
    instance_type = 'ml.g5.xlarge', # インスタンスタイプを指定
    image_scope = 'inference' # 推論コンテナを指定
)

print(container_image_uri)
# >> 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:1.12-gpu-py38

dockerfile_txt = f'''FROM {container_image_uri}

RUN apt update && apt install -y poppler-utils && pip install einops kornia yacs pdf2image
'''
print(dockerfile_txt)
with open('./docker/Dockerfile','wt') as f:
    f.write(dockerfile_txt)

ecr_repository = 'techbot-sagemaker-inference'
!aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin $container_image_uri
%cd docker
!docker build -t {ecr_repository} .
%cd ../

# boto3の機能を使ってリポジトリ名に必要な情報を取得する
account_id = boto3.client('sts').get_caller_identity().get('Account')
region = boto3.session.Session().region_name
tag = ':latest'

image_uri = f'{account_id}.dkr.ecr.{region}.amazonaws.com/{ecr_repository+tag}'

!$(aws ecr get-login --region $region --registry-ids $account_id --no-include-email)

# リポジトリの作成
# すでにある場合はこのコマンドは必要ない
!aws ecr create-repository --repository-name $ecr_repository

!docker tag {ecr_repository + tag} $image_uri
!docker push $image_uri

print(f'コンテナは {image_uri} へ登録されています。')
# >> コンテナは XXXXXXXXXXXX.dkr.ecr.ap-northeast-1.amazonaws.com/techbot-sagemaker-inference:latest へ登録されています

%%writefile src/inference.py
# 標準モジュール
import os
from glob import glob
import json
import shutil
import logging
from uuid import uuid4
import sys
# 3rd モジュール
import torch
# 自作モジュール
import classes.ModelA.eval_ModelA as eval_ModelA
import classes.ModelB.eval_ModelB as eval_ModelB
from pdf2png import pdf2png
from png2feat import png2feat

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logger.addHandler(logging.StreamHandler(sys.stdout))

# モデルをロードする関数
# エンドポイント立ち上げ時に一度だけ呼ばれる
def model_fn(model_dir):
    logger.info('loading ModelA...')
    ModelA = eval_ModelA.load_model(compute='gpu' if torch.cuda.is_available() else 'cpu')
    logger.info('loading PCR...')
    ModelB = eval_ModelB.load_model()
    logger.info('completed model_fn')
    return {'ModelA':ModelA, 'ModelB':ModelB}

# データの前処理を行う関数
# エンドポイントにリクエストされた時に最初呼ばれる
# input_fn が存在しない場合は default_input_fn が呼ばれる
def input_fn(input_data, content_type):
    if content_type == 'application/pdf':
        name = str(uuid4())
        logger.debug(f'{name=}')
        pdf_path = f'/tmp/{name}.pdf'
        logger.debug(f'{pdf_path=}')
        png_dir = f'/tmp/{name}/'
        logger.debug(f'{png_dir=}')
        with open(pdf_path, 'wb') as f:
            f.write(input_data)
        os.makedirs(png_dir)
        logger.info('pdf2png start...')
        is_success = pdf2png(pdf_path, png_dir)
        logger.info('pdf2png end')
        pngs = glob(os.path.join(png_dir,'*.png'))
        logger.debug(f'{pngs=}')
        os.remove(pdf_path)
        return pngs
    else:
        raise TypeError('ContentType is only allowed application/pdf.')

# 推論する関数
def predict_fn(pngs, model):
    logger.info('png2feat start...')
    feat = png2feat(model['ModelA'], model['ModelB'], pngs)
    logger.info('png2feat end')
    png_dir = pngs[0].replace(pngs[0].split('/')[-1],'')
    logger.info('delete {png_dir}')
    shutil.rmtree(png_dir)
    return feat

# 推論結果の後処理を行う関数
# output_fn が存在しない場合は default_output_fn が呼ばれる
def output_fn(feat, accept):
    if accept == 'application/json':
        response = json.dumps({'result_array':feat.tolist()})
    elif accept == 'text/csv':
        response = ','.join([str(element) for element in feat.tolist()])
    else:
        raise TypeError('Accept is only allowed application/json or text/csv')
    return response

with tarfile.open('model.tar.gz', mode='w:gz') as tar:
    tar.add('src/','code/') # コードを追加
    tar.add('./wait.ckpt')
    tar.add('./ModelB_model_320x240_4096')

# model.tar.gz check
with tarfile.open('./model.tar.gz','r:gz') as tar:
    tar.list()

source_s3_uri = sagemaker.session.Session().upload_data(
    f'./model.tar.gz',
    key_prefix = 'techbot-gpu'
)
print(source_s3_uri)

base_name = 'techbot-gpu'
model_name = f'{base_name}-model'
sm_client.create_model(
    ModelName=model_name,
    PrimaryContainer={
        'Image': image_uri,
        'Mode': 'SingleModel',
        'ModelDataUrl': source_s3_uri,
        'Environment':{
            'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
            'SAGEMAKER_PROGRAM': 'inference.py',
            'SAGEMAKER_REGION': 'region',
            'SAGEMAKER_SUBMIT_DIRECTORY': '/opt/ml/model/code'
        }
    },
    ExecutionRoleArn=role
)

endpoint_config_name = f'{base_name}-endpoint-config'
sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[{
        'VariantName': 'AllTraffic',
        'ModelName': model_name,
        'InitialInstanceCount': 1,
        'InstanceType': 'ml.g5.xlarge',
        'InitialVariantWeight': 1.0
    }]
)

endpoint_name = f'{base_name}-endpoint'
sm_client.create_endpoint(
    EndpointName=endpoint_name,
    EndpointConfigName = endpoint_config_name
)

endpoint_inservice_waiter.wait(
    EndpointName=endpoint_name,
    WaiterConfig={'Delay': 5,}
)

%%time
response = smr_client.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType='application/pdf',
    Accept='application/json',
    Body=content
)

result = json.loads(response['Body'].read().decode('utf-8'))
pprint(result)

sm_client.delete_endpoint(EndpointName=endpoint_name)
sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
sm_client.delete_model(ModelName=model_name)