[Python] Python モジュールを遅延ロードするにはどうすればよいですか? - MLflow から LazyLoader を分析する

(画像出典：https://www.irasutoya.com/2019/03/blog-post_72.html)

イントロ

ある日、私は MLflow など、Python で人気のある ML ライブラリをいくつか調べていました。ソース コードを眺めていると、__init__.py の LazyLoader というクラスに興味を惹かれました (これは実際には wandb プロジェクトを反映していますが、ご覧のとおり、元のコードは MLflow が現在使用しているものから変更されています)。

遅延読み込みの概念については、Web フロントエンドの画像読み込み、キャッシュ戦略など、さまざまな場面で聞いたことがあるでしょう。これらすべての遅延読み込みの概念の本質は、「私は怠け者でロード 今」ということだと思います - そう、「今」という隠された言葉です。 。つまり、アプリケーションは、必要な場合にのみそのリソースをロードして使用します。したがって、この MLflow ライブラリでは、その中のリソース (変数、関数、クラス) がアクセスされた場合にのみモジュールがロードされます。

でもどうやって？これが私の主な興味でした。そこでソースコードを読んでみましたが、一見すると非常に単純そうに見えました。しかし、驚くことに、その仕組みを理解するのに少し時間がかかり、コードを読むことで多くのことを学びました。この記事は、Python 言語のさまざまなテクニックを使用してこのような遅延読み込みがどのように機能するかを理解するために、MLflow のソース コードを分析することについてです。

LazyLoader を使って遊んでみる

分析の目的で、ローカル マシン上に lazyloading という単純なパッケージを作成し、次のようにモジュールを配置しました。

lazyloading/
├─ __init__.py
├─ __main__.py
├─ lazy_load.py
├─ heavy_module.py

• __init__.py: このファイルはディレクトリ全体をパッケージにします。
• __main__.py: このファイルは、次のようにパッケージ全体を実行する場合のエントリ ポイントです: python -m Lazyloading.
• Lazy_load.py: LazyLoader はこのファイル内にあります。
• Heavy_module.py: これは、シミュレーションのためにロードされる重いパッケージ (PyTorch など) を含むモジュールを表します:

import time

for i in range(5):
    time.sleep(1)
    print(5 - i, " seconds left before loading")

print("I am heavier than Pytorch!")

HEAVY_ATTRIBUTE = "heavy”

次に、この Heavy_module を __main__.py 内にインポートします。

if __name__ == "__main__":
    from lazyloading import heavy_module 

このパッケージを実行して結果を見てみましょう:

python -m lazyloading
5  seconds left before loading
4  seconds left before loading
3  seconds left before loading
2  seconds left before loading
1  seconds left before loading
I am heavier than pytorch!

ここでは、PyTorch などの重いパッケージを単純にインポートすると、アプリケーション全体のオーバーヘッドになる可能性があることが明確にわかります。だからこそ、ここで遅延読み込みが必要なのです。 __main__.py を次のように変更しましょう:

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("lazyloading.heavy_module", globals(), "lazyloading.heavy_module")
    print("nothing happens yet")
    print(heavy_module.HEAVY_ATTRIBUTE)

結果は次のようになります:

python -m lazyloading
nothing happens yet
5  seconds left before loading
4  seconds left before loading
3  seconds left before loading
2  seconds left before loading
1  seconds left before loading
heavy

はい、LazyLoader によってインポートされたモジュールは、スクリプトを実行したり、他のパッケージをインポートしたりする必要はありません。これは、モジュールのいずれかの属性がアクセスされた場合にのみ発生します。これが遅延読み込みの威力です!

LazyLoader は MLflow でどのように動作しますか? - ソースコード分析

コード自体は短くてシンプルです。説明のために型の注釈といくつかのコメント (<, > で囲まれた行) を追加しました。他のすべてのコメントは、元のソース コード内のコメントです。

"""Utility to lazy load modules."""
import importlib
import sys
import types

from typing import Any, TypeVar

T = TypeVar("T") # <this is added by me>

class LazyLoader(types.ModuleType):
    """Class for module lazy loading.

    This class helps lazily load modules at package level, which avoids pulling in large
    dependencies like `tensorflow` or `torch`. This class is mirrored from wandb's LazyLoader:
    https://github.com/wandb/wandb/blob/79b2d4b73e3a9e4488e503c3131ff74d151df689/wandb/sdk/lib/lazyloader.py#L9
    """

    _local_name: str # <the name of the package that is used inside code>
    _parent_module_globals: dict[str, types.ModuleType] # <importing module's namespace, accessible by calling globals()>
    _module: types.ModuleType | None # <actual module>

    def __init__(
        self, 
        local_name: str, 
        parent_module_globals: dict[str, types.ModuleType], 
        name: Any # <to be used in types.ModuleType(name=str(name)), the full package name (such as pkg.subpkg.subsubpkg)>
    ):
        self._local_name = local_name
        self._parent_module_globals = parent_module_globals
        self._module = None

        super().__init__(str(name)) 

    def _load(self) -> types.ModuleType:
        """Load the module and insert it into the parent's globals."""
        if self._module:
            # If already loaded, return the loaded module.
            return self._module

        # Import the target module and insert it into the parent's namespace

        # <see https://docs.python.org/3/library/importlib.html#importlib.import_module>
        # <absolute import, importing the module itself from a package rather than the top-level package only(like __import__)>
        # <here, self.__name__ is the variable `name` in __init__>
        # <this is why that `name` in __init__ must be the full module path>
        module = importlib.import_module(self.__name__) # this automatically updates sys.modules

        # <add the name of the module to the importing module(=parent module)'s namespace>
        # <so that you can use this module's name as a variable inside the importing module, even if it is called inside a function defined in the importing module>
        self._parent_module_globals[self._local_name] = module

        # <add the module to the list of loaded modules for caching>
        # <see https://docs.python.org/3/reference/import.html#the-module-cache>
        # <this makes possible to import cached module with the variable _local_name
        sys.modules[self._local_name] = module

        # Update this object's dict so that if someone keeps a reference to the `LazyLoader`,
        # lookups are efficient (`__getattr__` is only called on lookups that fail).
        self.__dict__.update(module.__dict__)

        return module

    def __getattr__(self, item: T) -> T:
        module = self._load()
        return getattr(module, item)

    def __dir__(self):
        module = self._load()
        return dir(module)

    def __repr__(self):
        if not self._module:
            return f"<module '{self.__name__} (Not loaded yet)'>"
        return repr(self._module)

ここで、heavy_module を遅延ロードしながらコードを調べてみましょう。モジュールの重さをシミュレートする必要がなくなったので、time.sleep(1) ループ部分を削除しましょう。

1. LazyLoader のインスタンスを作成し、元のモジュールをプロキシします

LazyLoader の __init__() を見てみましょう。

class LazyLoader(types.ModuleType):
    # …
    # code omitted
    # …

    def __init__(
        self, 
        local_name: str, 
        parent_module_globals: dict[str, types.ModuleType], 
        name: Any # <to be used in types.ModuleType(name=str(name)); the full package name(such as pkg.subpkg.subsubpkg)>
    ):
        self._local_name = local_name
        self._parent_module_globals = parent_module_globals
        self._module = None

        super().__init__(str(name)) 

local_name、parent_module_globals、および name をコンストラクター __init__() に提供します。現時点では、これらすべてが何を意味するのかはわかりませんが、LazyLoader は types.ModuleType を継承しているため、少なくとも最後の行は実際にモジュール super().__init__(str(name)) を生成していることを示しています。変数名を指定すると、LazyLoaderで作成したモジュールがname(heavy_module.__name__と同じ)という名前のモジュールとして認識されます。

モジュール自体を出力すると、次のことが証明されます:

# __main__.py
# run python -m lazyloading

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("lazyloading.heavy_module", globals(), "lazyloading.heavy_module")

    print(heavy_module.__name__)

これは端末上で次のようになります:

lazyloading.heavy_module

ただし、コンストラクターでは、インスタンス変数に値を代入し、この プロキシ モジュールにモジュールの名前を指定しただけです。さて、モジュールの属性にアクセスしようとすると何が起こるでしょうか?

2. 属性へのアクセス - __getattribute__、__getattr__、getattr

これはこのクラスの楽しい部分の 1 つです。一般に、Python オブジェクトの属性にアクセスすると何が起こるでしょうか? heavy_module.HEAVY_ATTRIBUTE を呼び出して、heavy_module の HEAVY_ATTRIBUTE にアクセスするとします。ここのコード、またはいくつかの Python プロジェクトでのあなた自身の経験から、__getattr__() が呼び出されると推測するかもしれませんが、それは部分的には正しいです。公式ドキュメントを参照してください:

デフォルトの属性アクセスが AttributeError で失敗した場合に呼び出されます（name がインスタンス属性または自分自身のクラス ツリー内の属性ではないため、getattribute() が AttributeError を発生させるか、get name プロパティの場合、AttributeError が発生します)。

(Please ignore __get__ because it is out of scope of this post, and our LazyLoader doesn’t implement __get__ either).

So __getattribute__() the key method here is __getattribute__. According to the docs, when we try to access an attribute, __getattribute__ will be called first, and if the attribute we’re looking for cannot be found by __getattribute__, AttributeError will be raised, which will in turn invoke our __getattr__ in the code. To verify this, let’s override __getattribute__ of the LazyLoader class, and change __getattr__() a little bit as follows:

def __getattribute__(self, name: str) -> Any:
    try:
        print(f"__getattribute__ is called when accessing attribute '{name}'")
        return super().__getattribute__(name)

    except Exception as error:
        print(f"an error has occurred when __getattribute__() is invoked as accessing '{name}': {error}")
        raise

def __getattr__(self, item: T) -> T:
    print(f"__getattr__ is called when accessing attribute '{item}'")
    module = self._load()
    return getattr(module, item)

When we access HEAVY_ATTRIBUTE that exists in heavy_module, the result is:

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("lazyloading.heavy_module", globals(), "lazyloading.heavy_module")

    print(heavy_module.HEAVY_ATTRIBUTE)

python -m lazyloading
__getattribute__ is called when accessing attribute 'HEAVY_ATTRIBUTE'
an error has occurred when __getattribute__() is invoked as accessing 'HEAVY_ATTRIBUTE': module 'lazyloading.heavy_module' has no attribute 'HEAVY_ATTRIBUTE'
__getattr__ is called when accessing attribute 'HEAVY_ATTRIBUTE'
__getattribute__ is called when accessing attribute '_load'
__getattribute__ is called when accessing attribute '_module'
__getattribute__ is called when accessing attribute '__name__'
I am heavier than Pytorch!
__getattribute__ is called when accessing attribute '_parent_module_globals'
__getattribute__ is called when accessing attribute '_local_name'
__getattribute__ is called when accessing attribute '__dict__'
heavy

So __getattr__ is actually not called directly, but __getattribute__ is called first, and it raises AttributeError because our LazyLoader instance doesn’t have attribute HEAVY_ATTRIBUTE. Now __getattr__() is called as a failover. Then we meet getattr(), but this code line getattr(module, item) is equivalent to code module.item in Python. So eventually, we access the HEAVY_ATTRIBUTE in the actual module heavy_module, if module variable in __getattr__() is correctly imported and returned by self._load().

But before we move on to investigating _load() method, let’s call HEAVY_ATTRIBUTE once again in __main__.py and run the package:

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("lazyloading.heavy_module", globals(), "lazyloading.heavy_module")

    print(heavy_module.HEAVY_ATTRIBUTE)
    print(heavy_module.HEAVY_ATTRIBUTE)

Now we see the additional logs on the terminal:

# … the same log as above
__getattribute__ is called when accessing attribute 'HEAVY_ATTRIBUTE'
heavy

It seems that __getattribute__ can access HEAVY_ATTRIBUTE now inside the proxy module(our LazyLoader instance). This is because(!!!spoiler alert!!!) _load caches the accessed attribute in __dict__ attribute of the LazyLoader instance. We’ll get back to this in the next section.

3. Loading and caching the actual module

This section covers the core part the post - loading the actual module in the function _load().

3-1. Module caching at the level of LazyLoader class

First, it checks whether our LazyLoader instance has already imported the module before (which reminds us of the Singleton pattern).

if self._module:
    # If already loaded, return the loaded module.
    return self._module

3-2. Importing the actual module with importlib.import_module

Otherwise, the method tries to import the module named __name__, which we saw in the __init__ constructor:

# <see https://docs.python.org/3/library/importlib.html#importlib.import_module>
# <absolute import, importing the module itself from a package rather than the top-level package only(like __import__)>
# <here, self.__name__ is the variable `name` in __init__>
# <this is why that `name` in __init__ must be the full module path>
module = importlib.import_module(self.__name__) # this automatically updates sys.modules

According to the docs of importlib.import_module, when we don’t provide the pkg argument and only the path string, the function tries to import the package in the absolute manner. Therefore, when we create a LazyLoader instance, the name argument should be the absolute term. You can run your own experiment to see it raises ModuleNotFoundError:

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("heavy_module", globals(), "heavy_module")

    print(heavy_module.HEAVY_ATTRIBUTE)

# logs omitted
ModuleNotFoundError: No module named 'heavy_module'

Notably, invoking importlib.import_module(self.__name__) caches the module with name self.__name__ in the global scope. If you run the following lines in __main__.py

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader
    heavy_module = LazyLoader("heavy_module", globals(), "lazyloading.heavy_module")

    # check whether the module is cached at the global scope
    import sys
    print("lazyloading.heavy_module" in sys.modules)

    # accessing any attribute to load the module
    heavy_module.HEAVY_ATTRIBUTE

    print("lazyloading.heavy_module" in sys.modules)

and run the package, then the logs should be:

python -m lazyloading
False
I am heavier than Pytorch!
True

This way of caching using sys.modules is related to the next two lines that also cache the module in different ways.

3-3. Caching the module with given local_name

# <add the name of the module to the importing module(=parent module)'s namespace>
# <so that you can use this module's name as a variable inside the importing module, even if it is called inside a function defined in the importing module>
self._parent_module_globals[self._local_name] = module

# <add the module to the list of loaded modules for caching>
# <see https://docs.python.org/3/reference/import.html#the-module-cache>
# <this makes possible to import cached module with the variable _local_name
sys.modules[self._local_name] = module

Both lines cache the module in the dictionaries self._parent_module_globals and sys.modules respectively, but with the key self._local_name(not self.__name__). This is the variable we provided as local_name when creating this proxy module instance with __init__(). But what does this caching accomplish?

First, we can use the module with the given _local_name in the "parent module"’s globals(from the parameter’s name and seeing how MLflow uses in its uppermost __init__.py, we can infer that here the word globals means (globals()). This means that importing the module inside a function doesn’t limit the module to be used outside the function’s scope:

if __name__ == "__main__":
    from lazyloading.lazy_load import LazyLoader

    def load_heavy_module() -> None:
        # import the module inside a function
        heavy_module = LazyLoader("heavy_module", globals(), "lazyloading.heavy_module")
        print(heavy_module.HEAVY_ATTRIBUTE)

    # loads the heavy_module inside the function's scope
    load_heavy_module()

    # the module is now in the scope of this module
    print(heavy_module)

Running the package gives:

python -m lazyloading
I am heavier than Pytorch!
heavy
<module 'lazyloading.heavy_module' from ‘…’> # the path of the heavy_module(a Python file)

Of course, if you provide the second argument locals(), then you’ll get NameError(give it a try!).

Second, we can also import the module in any other place inside the whole package with the given local name. Let’s create another module heavy_module_loader.py inside the current package lazyloading :

lazyloading/
├─ __init__.py
├─ __main__.py
├─ lazy_load.py
├─ heavy_module.py
├─ heavy_module_loader.py

Note that I used a custom name heavy_module_local for the local variable name of the proxy module.

# heavy_module_loader.py

from lazyloading.lazy_load import LazyLoader

heavy_module = LazyLoader("heavy_module_local", globals(), "lazyloading.heavy_module")
heavy_module.HEAVY_ATTRIBUTE

Now let __main__.py be simpler:

from lazyloading import heavy_module_loader

if __name__ == "__main__":
    import heavy_module_local
    print(heavy_module_local)

Your IDE will probably alert this line as having a syntax error, but actually running it will give us the expected result:

python -m lazyloading
I am heavier than Pytorch!
<module 'lazyloading.heavy_module' from ‘…’> # the path of the heavy_module(a Python file)

Although MLflow seems to use the same string value for both local_name and name when creating LazyLoader instances, we can use the local_name as an alias for the actual package name, thanks to this caching mechanism.

3-4. Caching the attributes of the actual module in __dict__

# Update this object's dict so that if someone keeps a reference to the `LazyLoader`,
# lookups are efficient (`__getattr__` is only called on lookups that fail).
self.__dict__.update(module.__dict__)

In Python, the attribute __dict__ gives the dictionary of attributes of the given object. Updating this proxy module’s attributes with the actual module’s ones makes the user easier to access the attributes of the real one. As we discussed in section 2(2. Accessing an attribute - __getattribute__, __getattr__, and getattr) and noted in the comments of the original source code, this allows __getattribute__ and __getattr__ to directly access the target attributes.

In my view, this part is somewhat unnecessary, as we already cache modules and use them whenever their attributes are accessed. However, this could be useful when we need to debug and inspect __dict__.

4. __dir__ and __repr__

Similar to __dict__, these two dunder functions might not be strictly necessary when using LazyLoader modules. However, they could be useful for debugging. __repr__ is particularly helpful as it indicates whether the module has been loaded.

<p>if not self.<em>module</em>:<br>
    return f"<module '{self.<em>name</em>_} (Not loaded yet)'>"<br>
return repr(self._module)</p>

Conclusion

Although the source code itself is quite short, we covered several advanced topics, including importing modules, module scopes, and accessing object attributes in Python. Also, the concept of lazyloading is very common in computer science, but we rarely get the chance to examine how it is implemented in detail. By investigating how LazyLoader works, we learned more than we expected. Our biggest takeaway is that short code doesn’t necessarily mean easy code to analyze!

以上が[Python] Python モジュールを遅延ロードするにはどうすればよいですか? - MLflow から LazyLoader を分析するの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。