Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move external allocators into rmm.allocators module to defer imports #1221

Merged
merged 14 commits into from
Feb 27, 2023

Commits on Feb 23, 2023

  1. Move third-party allocators to submodule

    Preparation for removing by-default import of non-essential packages.
    wence- committed Feb 23, 2023
    Configuration menu
    Copy the full SHA
    172bbb5 View commit details
    Browse the repository at this point in the history
  2. Move import of allocators to __init__.py

    This removes an indirection through rmm.rmm, to simplify removal of
    the top-level imports by default.
    wence- committed Feb 23, 2023
    Configuration menu
    Copy the full SHA
    e20d071 View commit details
    Browse the repository at this point in the history
  3. Configuration menu
    Copy the full SHA
    6f204e0 View commit details
    Browse the repository at this point in the history
  4. Don't import allocators by default

    This will make importing RMM faster if you don't want the hooks for
    cupy, numba, or pytorch.
    
    Before, a sampling trace of `import rmm` with pyinstrument shows:
    
        $ pyinstrument -i 0.01 importrmm.py
    
          _     ._   __/__   _ _  _  _ _/_   Recorded: 10:19:56  Samples:  67
         /_//_/// /_\ / //_// / //_'/ //     Duration: 0.839     CPU time: 0.837
        /   _/                      v4.4.0
    
        Program: importrmm.py
    
        0.839 <module>  importrmm.py:1
        └─ 0.839 <module>  rmm/__init__.py:1
           ├─ 0.315 <module>  rmm/allocators/torch.py:1
           │  └─ 0.315 <module>  torch/__init__.py:1
           │        [96 frames hidden]  torch, <built-in>, enum, inspect, tok...
           ├─ 0.297 <module>  rmm/mr.py:1
           │  └─ 0.297 <module>  rmm/_lib/__init__.py:1
           │     ├─ 0.216 <module>  numba/__init__.py:1
           │     │     [140 frames hidden]  numba, abc, <built-in>, importlib, em...
           │     ├─ 0.040 <module>  numba/cuda/__init__.py:1
           │     │     [34 frames hidden]  numba, asyncio, ssl, <built-in>, re, ...
           │     ├─ 0.030 __new__  enum.py:180
           │     │     [5 frames hidden]  enum, <built-in>
           │     └─ 0.011 [self]  None
           └─ 0.227 <module>  rmm/allocators/cupy.py:1
              └─ 0.227 <module>  cupy/__init__.py:1
                    [123 frames hidden]  cupy, pytest, _pytest, attr, <built-i...
    
    After:
    
        $ pyinstrument -i 0.01 importrmm.py
    
          _     ._   __/__   _ _  _  _ _/_   Recorded: 10:20:10  Samples:  28
         /_//_/// /_\ / //_// / //_'/ //     Duration: 0.297     CPU time: 0.297
        /   _/                      v4.4.0
    
        Program: importrmm.py
    
        0.296 <module>  importrmm.py:1
        └─ 0.296 <module>  rmm/__init__.py:1
           └─ 0.296 <module>  rmm/mr.py:1
              └─ 0.296 <module>  rmm/_lib/__init__.py:1
                 ├─ 0.216 <module>  numba/__init__.py:1
                 │     [141 frames hidden]  numba, <built-in>, importlib, email, ...
                 ├─ 0.040 <module>  numba/cuda/__init__.py:1
                 │     [19 frames hidden]  numba, asyncio, ssl, <built-in>, unit...
                 ├─ 0.031 [self]  None
                 └─ 0.010 __new__  enum.py:180
                       [4 frames hidden]  enum, <built-in>
    wence- committed Feb 23, 2023
    Configuration menu
    Copy the full SHA
    359fcc0 View commit details
    Browse the repository at this point in the history
  5. Expunge final top-level imports of numba

    After this change, a sampling trace of `import rmm` with pyinstrument
    shows:
    
        $ pyinstrument -i 0.01 importrmm.py
    
          _     ._   __/__   _ _  _  _ _/_   Recorded: 10:37:40  Samples:  9
         /_//_/// /_\ / //_// / //_'/ //     Duration: 0.099     CPU time: 0.099
        /   _/                      v4.4.0
    
        Program: importrmm.py
    
        0.099 <module>  importrmm.py:1
        └─ 0.099 <module>  rmm/__init__.py:1
           └─ 0.099 <module>  rmm/mr.py:1
              └─ 0.099 <module>  rmm/_lib/__init__.py:1
                 ├─ 0.059 <module>  numpy/__init__.py:1
                 │     [31 frames hidden]  numpy, re, sre_compile, <built-in>, s...
                 ├─ 0.020 __new__  enum.py:180
                 │     [2 frames hidden]  enum
                 ├─ 0.010 <module>  ctypes/__init__.py:1
                 │     [3 frames hidden]  ctypes, <built-in>
                 └─ 0.010 _EnumDict.__setitem__  enum.py:89
                       [3 frames hidden]  enum
    
    Prior to this, when we had expunged cupy and pytorch imports from the
    top-level, but not yet numba, an import would take 0.3, rather than
    0.1s.
    wence- committed Feb 23, 2023
    Configuration menu
    Copy the full SHA
    7fd885b View commit details
    Browse the repository at this point in the history
  6. Configuration menu
    Copy the full SHA
    73074de View commit details
    Browse the repository at this point in the history
  7. Clean up duplicate import

    wence- committed Feb 23, 2023
    Configuration menu
    Copy the full SHA
    47e50d4 View commit details
    Browse the repository at this point in the history

Commits on Feb 24, 2023

  1. Absolute import of rmm.mr

    wence- committed Feb 24, 2023
    Configuration menu
    Copy the full SHA
    8dd78e6 View commit details
    Browse the repository at this point in the history
  2. Fix docstring example

    wence- committed Feb 24, 2023
    Configuration menu
    Copy the full SHA
    6f0900e View commit details
    Browse the repository at this point in the history
  3. Fix header level

    wence- committed Feb 24, 2023
    Configuration menu
    Copy the full SHA
    04df09f View commit details
    Browse the repository at this point in the history
  4. Configuration menu
    Copy the full SHA
    caafde5 View commit details
    Browse the repository at this point in the history
  5. Configuration menu
    Copy the full SHA
    8ad1acb View commit details
    Browse the repository at this point in the history

Commits on Feb 27, 2023

  1. Simplify sentence.

    bdice committed Feb 27, 2023
    Configuration menu
    Copy the full SHA
    c436826 View commit details
    Browse the repository at this point in the history
  2. Configuration menu
    Copy the full SHA
    d40f71e View commit details
    Browse the repository at this point in the history