Amazon SageMaker Processing -Manifestファイルで必要なファイルだけコピーする-

Amazon SageMaker Processing

Amazon SageMaker Processingは機械学習の前処理などを、お手軽に実行することができるマネージドサービスです。S3バケットのオブジェクトを入力としたり、結果をS3バケットに出力することが簡単にでき、またローカルで作成したPythonスクリプトがそのまま動かせる感覚で使用できるなど、AWS上でデータを処理する環境としてとても便利です。
基本的な使い方はこちらやこちらを参照ください。

入力ファイルでハマったところ

多くの例ではAWSブログにもある通り、ProcessingInputでS3フォルダとコンテナをディレクトリに紐付ける使い方がされています。この方法ですとProcessing Jobの実行時に、指定されたS3フォルダ以下がコンテナ内の入力ディレクトリにコピーされるため、オブジェクトの量が多い場合には実行開始までに多くの時間がかかってしまいます。S3のフォルダ構成が分析とあっている場合は問題ないですが、例えばS3が

s3://bucket-name/purchase/user=A/year=2021/month=1/2021-01-10.csv
s3://bucket-name/purchase/user=A/year=2021/month=2/2021-02-24.csv
s3://bucket-name/purchase/user=B/year=2021/month=1/2021-01-03.csv
s3://bucket-name/purchase/user=B/year=2021/month=2/2021-02-13.csv

と構成されており、「全てのユーザの2月のデータを処理したい」ような場合にはpurchase以降をコピーするしかなく、必要以上にデータをコピーされてしまい実行時間も長くなってしまっていました。

こんな時のためにManifestファイルによる入力ファイルの指定方法が用意されていますが、事例がなく使い方がわかりませんでした。AWSの方に相談したところ、その使い方を教えていただいたので紹介します。

Manifestファイルによる選択的なファイルコピー

Amazon SageMaker Processingの入力として使用するProcessingInputには、s3_data_typeとしてManifestFileを指定することができ、コピーしたいS3オブジェクトを予めManifestファイルに記述しておくことで、必要なファイルのみをコンテナ内にコピーすることができます。以下にその使い方を紹介します。

Manifestファイルの準備

Manifestファイルは配列からなるJSONファイルで、拡張子は.manifestです。共通パスを表すprefixと、prefix以降の相対パスのリストからなり、例えば前述の「全てのユーザの2月のデータを処理したい」場合は以下のようにS3オブジェクトをそれぞれ記述します。

sample.manifest

[ {"prefix": "s3://bucket-name/purchase/"},
"user=A/year=2021/month=2/2021-02-24.csv",
"user=B/year=2021/month=2/2021-02-13.csv",
...
]

manifestファイルはローカルに置くか、S3に置いておきます。

ProcessingInputでの指定

Processing Jobのインプットはprocessor.runの引数inputsにProcessingInputを列挙することで指定します。manifestファイルを使用する場合にはinputsを以下のように記述します。

    inputs=[
        ProcessingInput(source='sample.manifest',
                        s3_data_type='ManifestFile',
                        destination='/opt/ml/processing/input',
                       )
    ]

sourceにローカルパスを指定すると、manifestファイルは一旦S3のデフォルトバケットにアップロードされ、その後インスタンス起動時にSageMakerに読み込まれます。manifestファイルを自分の好きなバケットに置いておきたい場合は、予めmanifestファイルをS3にアップロードしsourceにそのS3パスを指定することで読み込ませることができます。
S3オブジェクトはdestination以下に相対パスが保持される形でコピーさるため、今回の例ですとS3オブジェクトはそれぞれ以下のパスにコピーされます。

/opt/ml/processing/input/user=A/year=2021/month=2/2021-02-24.csv
/opt/ml/processing/input/user=B/year=2021/month=2/2021-02-13.csv
...

こうしておくことで必要なファイルのみをコピーすることができるため、S3フォルダ全体をコピーする場合に比べてジョブの起動が早くすることができます。

まとめ

今回はAmazon SageMaker Processingでmanifestファイルを使用して必要なS3オブジェクトのみをコピーする方法を紹介しました。ジョブの起動時間が短くなるほか、入力ファイルをコンテナの外側で制御できるためコンテナが使い回しやすいと言う利点もあるかと思います。
Amazon SageMaker ProcessingはAWS Batchよりもお手軽で、コンテナをバッチで動かすのにとても便利だと思います。今後、スポットインスタンスに対応するなど、どんどん使いやすくなっていってほしいです。