Pytorch-LightningのCheckpointでLAN内のminio(S3互換ストレージ)へモデルを保存する方法

概要

現在のpytorch-lightning(ver.1.2.3)では、S3への保存を行う際にendpointを指定できないため、LAN内にあるS3互換ストレージなどへ保存を行うことが出来ない。
以下のsetup_endpointを実行することで、fsspecモジュールへ強引にモンキーパッチを当ててエンドポイントを指定する。

方法

• pytorch-lightning がクラウドストレージへのIOで使用している関数へモンキーパッチを当てEndpointを強引に指定する。
• 以下のコードをコピペしてsetup_endpointを実行するだけで可能。

コード

from typing import Any
import gorilla
import fsspec

def apply_gorrila(function: Any, module: Any):
    patch = gorilla.Patch(
        module, 
        function.__name__, 
        function, 
        settings=gorilla.Settings(allow_hit=True))
    gorilla.apply(patch)

def setup_endpoint(endpoint_url: str):
    def filesystem(protocol, **storage_options):
        if protocol == "s3":
            storage_options["client_kwargs"] = {
                "endpoint_url": endpoint_url,
            }
        original = gorilla.get_original_attribute(fsspec, 'filesystem')
        return original(protocol, **storage_options)
    def open_files(*args, **kwargs):
        kwargs["client_kwargs"] = {
                "endpoint_url": endpoint_url,
        }
        original = gorilla.get_original_attribute(fsspec.core, 'open_files')
        return original(*args, **kwargs)
    apply_gorrila(filesystem, fsspec)
    apply_gorrila(open_files, fsspec.core)