Amazon SageMakerでS3のどこに何が出力されるかをちゃんと制御したい

こんにちは、小澤です。

Amazon SageMaker(以下SageMaker)を使う際、入出力はS3であることが前提となっています。 デフォルトバケットというものがあったり、出力のデフォルト値が用意されている場合があったりと、 意識しなくてもそのまま使えてしまうケースも多くあります。

しかし、私のようなズボラな人間がそこに頼っていると、いつの日かどこに何があるのかさっぱりわからなくなる日が訪れます。 そこで、今回はSageMakerがどんな時にどんなファイルをS3に出力するのか整理ていきます。

まずはデフォルトバケット

SageMaker Python SDKを使っている場合、 sagemaker.Sessiondefault_bucket 関数にてデフォルトで利用するバケットの名前が取得できます。 これは、 sagemaker-- となっています。

このデフォルトバケットは任意のものに変更可能です。 Sessionのインスタンスを作成する際に default_bucket 引数があるためそこに設定したバケットをデフォルトバケットとして利用するように変更できます。

sess = sagemaker.Session(default_bucket='<任意のバケット名>')

出力されるもの

さて、実際にファイルを出力するジョブを見ていきます。 Processing Job, Training Job, Batch Transform Jobを考えていきます。

Processing Jobはoutputsで明示的に指定したもののみがS3への出力対象となるため、 どこに出力されたかわからなくなる、ということは少ないでしょう。

Training Jobの出力先はEstimatorのインスタンス作成時にoutput_pathで明示的に指定が可能です。 これによって任意のバケット、プレフィックスを出力先として選択できます。

ただし、指定したパスにそのまま出力されるわけではなく、末尾にジョブ名が付与されたものが実際に出力されるプレフィックスとなります。 ジョブ名は重複できないため、出力が上書きされることはありませんが常に同じパスに出力することもできません。

また、Training Jobでは最終的な出力先のほかにモデルのチェックポイントもS3に保存しておけます。 こちらは、 checkpoint_s3_uri 引数と checkpoint_local_path 引数で対応関係を明示することで利用可能になるため、 指定してない時にどのに出力されたかわからなくなることは少ないでしょう。

最後のBatch Transform Jobも同様にインスタンス作成時のoutput_pathで出力先を明示的に指定できます。 こちらはTraining Jobのようなジョブ名を末尾に付与するような仕組みは無いため、指定したパスがそのまま出力先になります。

各ジョブの出力に関しては上記の通り、output_pathを指定するものはデフォルト値がありそれ以外は明示的に指定するものとなっています。 また、入力に関してはTraining Jobから直接TransformerやPredictorような暗黙的に特定できる場合を除いて、常に明示的に指定するようになっているため、 こちらで迷子になることは少ないでしょう。

それ以外でS3に出力されるもの

さて、SageMaker Python SDKを使っていると実はこのほかにもS3に出力されているファイルがあります。 SageMakerではジョブ実行時に利用するスクリプトをコンテナ内に含めるのではなく実行時に指定する仕組みがいくつか用意されています。

entry_pointやcodeとしてローカルのファイルが指定できますが、 それらがジョブが実行可能な仕組みの実現方法として

  • SDK側で一度S3にスクリプトをアップロード
  • コンテナ内の処理でS3から取得したものを実行する

という動きになっています。 そのため、実行するスクリプトが暗黙的に一度S3にアップロードされています。 この時の動きに関してはEstimator, ModelとProcessorとで少し状況が異なります。

EstimatorやModel(Predictor, Transformer共通で利用するモデルの設定)はentry_pointやsrc_dirで指定したファイルが対象となりますが、 こちらは厳密にはsagemaker-training-toolkitなどのライブラリ側でのS3からファイルを取得する仕組みを実現しており、 sagemaker.estimator.Framework を継承したクラスが entry_point , source_dir 引数として受け取ったものをS3にアップロードする仕組みになっています。

entry_point引数単体での指定はローカルのパスまたはGitからの取得のみとなります。 そのため、こちらは実行するたびにジョブ名を含むパスでS3にスクリプトのファイルが出力されます。 S3のパスを指定する場合は、 code_location 引数で s3:// で始まるS3のパスを指定できます。 source_dirを利世する場合は個の引数に s3:// で始まるパスを指定できます。

Processorについてですが、こちらはScriptProcessorとそのサブクラスを利用している際にcode引数で指定するものになります。 こちらは、実際にはS3にスクリプトをアップロードしたのち、Processing Jobのinputsにそれを追加しています。 アップロード先はジョブ名が含まれるパスとなるため、実行するたびにアップロードされます。

こちらは、 s3:// で始まるパスも指定可能ですし、通常のProcessorでinputsの中で同じようにしてすることも可能です。 そのため、source_dirに相当する引数はProcessorにはありませんがこの指定方法で同様の動きを実現できます。

おわりに

SageMakerの各ジョブがどこで何をS3に出力していて、その出力先をどう制御するかを見ていきました。 複数人でバケットを使い分けて利用したい場合など、どこに出力されるか意識したい時にご参照いただければと思います。