Use module-based k2 import guard

[artbataev]

Signed-off-by: Vladimir Bataev [email protected]

What does this PR do ?

Simplify k2 library import guard: use module-based approach.

Advantages:

• less code is required for using guard
• no need to suppress linters (isort, etc.)

Collection: [ASR]

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

from nemo.core.utils.k2_guard import k2  # import k2 with performing all required checks

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[GNroy]

LGTM

[SeanNaren]

LGTM!

[titu1994]

How is it an optional dependency if merely importing it will raise an error ? K2 is not installed by default, the check should be don't at init inside of where it will be used rather than during simply import of a file.

[titu1994]

Currently, k2 is not included in the import path of init.py files for this reason, but at least the check is done at init of the class. This instead puts it in the import path so if user simply imports the file it will crash with helpful error.

Numba checks specifically defers this crash to the moment user actually tries to use the module rather than simply import the module

[artbataev]

Importing k2 raises an error only when importing k2_guard of k2 from k2_guard
k2_guard will not be imported and cached.

So, every module that tries to import k2 from k2_guard will raise an error (if no k2 / incorrect version) and will not be initialized (no classes in this module).
As a result, importing from any of the modules in nemo, that depend on k2 (maybe not directly) will result in error with the necessary message, which is a correct behavior.

But importing not at the top level will allow to raise an error in runtime, we do not need to call k2_import_guard():

if backend == "k2":
    if criterion_type == "ml":
        from nemo.collections.asr.parts.k2.ml_loss import MLLoss as K2Loss  # maybe error here
    elif criterion_type == "map":
        from nemo.collections.asr.parts.k2.map_loss import MAPLoss as K2Loss  # or here

This doesn't change the previous behavior: in modules k2_import_guard() call was used at the top level, so incorrect k2 version resulted in immediate crash before importing any defined class.

[artbataev]

You can set __K2_MINIMUM_MINOR_VERSION = 111 and see that the behavior didn't change and everything works.

[titu1994]

Approved for now, but K2 should switch to dynamic Numba style runtime check rather than static file import check.