Make all nvtx usage go through safe imports (#5424)

Authors:
  - Dante Gama Dessavre (https://github.com/dantegd)

Approvers:
  - William Hicks (https://github.com/wphicks)

URL: #5424

python/cuml/datasets/blobs.py

@@ -14,14 +14,19 @@
 # limitations under the License.
 #
 
-
-from cuml.datasets.utils import _create_rs_generator
-import cuml.internals
-from cuml.internals.safe_imports import cpu_only_import
-import nvtx
 import numbers
+import cuml.internals
+
 from collections.abc import Iterable
 from cuml.internals.safe_imports import gpu_only_import
+from cuml.datasets.utils import _create_rs_generator
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import_from,
+    null_decorator,
+)
+
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 
 cp = gpu_only_import("cupy")
 np = cpu_only_import("numpy")
@@ -83,7 +88,7 @@ def _get_centers(rs, centers, center_box, n_samples, n_features, dtype):
     return centers, n_centers
 
 
-@nvtx.annotate(message="datasets.make_blobs", domain="cuml_python")
+@nvtx_annotate(message="datasets.make_blobs", domain="cuml_python")
 @cuml.internals.api_return_generic()
 def make_blobs(
     n_samples=100,

python/cuml/datasets/classification.py

@@ -13,14 +13,20 @@
 # limitations under the License.
 #
 
-import nvtx
-from cuml.internals.safe_imports import cpu_only_import
 import cuml.internals
 from cuml.internals.import_utils import has_sklearn
 from cuml.datasets.utils import _create_rs_generator
 
 from cuml.internals.safe_imports import gpu_only_import
 
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import_from,
+    null_decorator,
+)
+
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
+
 cp = gpu_only_import("cupy")
 np = cpu_only_import("numpy")
 
@@ -50,7 +56,7 @@ def _generate_hypercube(samples, dimensions, rng):
     return out
 
 
-@nvtx.annotate(message="datasets.make_classification", domain="cuml_python")
+@nvtx_annotate(message="datasets.make_classification", domain="cuml_python")
 @cuml.internals.api_return_generic()
 def make_classification(
     n_samples=100,

python/cuml/datasets/regression.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2022, NVIDIA CORPORATION.
+# Copyright (c) 2019-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,9 +17,13 @@
 # distutils: language = c++
 
 import typing
-import nvtx
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import_from,
+    null_decorator
+)
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 
-from cuml.internals.safe_imports import cpu_only_import
 np = cpu_only_import('numpy')
 
 import cuml.internals
@@ -74,7 +78,7 @@ inp_to_dtype = {
 }
 
 
-@nvtx.annotate(message="datasets.make_regression", domain="cuml_python")
+@nvtx_annotate(message="datasets.make_regression", domain="cuml_python")
 @cuml.internals.api_return_generic()
 def make_regression(
     n_samples=100,

python/cuml/ensemble/randomforestclassifier.pyx

@@ -16,10 +16,15 @@
 #
 
 # distutils: language = c++
-from cuml.internals.safe_imports import cpu_only_import
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import,
+    gpu_only_import_from,
+    null_decorator
+)
+
 np = cpu_only_import('numpy')
-import nvtx
-from cuml.internals.safe_imports import gpu_only_import
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 rmm = gpu_only_import('rmm')
 import warnings
 
@@ -413,7 +418,7 @@ class RandomForestClassifier(BaseRandomForestModel,
                                  algo=algo,
                                  fil_sparse_format=fil_sparse_format)
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="fit RF-Classifier @randomforestclassifier.pyx",
         domain="cuml_python")
     @generate_docstring(skip_parameters_heading=True,
@@ -554,7 +559,7 @@ class RandomForestClassifier(BaseRandomForestModel,
         del(X_m)
         return preds
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="predict RF-Classifier @randomforestclassifier.pyx",
         domain="cuml_python")
     @insert_into_docstring(parameters=[('dense', '(n_samples, n_features)')],
@@ -676,7 +681,7 @@ class RandomForestClassifier(BaseRandomForestModel,
 
         return preds_proba
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="score RF-Classifier @randomforestclassifier.pyx",
         domain="cuml_python")
     @insert_into_docstring(parameters=[('dense', '(n_samples, n_features)'),

python/cuml/ensemble/randomforestregressor.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2022, NVIDIA CORPORATION.
+# Copyright (c) 2019-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,10 +16,14 @@
 
 # distutils: language = c++
 
-from cuml.internals.safe_imports import cpu_only_import
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import,
+    gpu_only_import_from,
+    null_decorator
+)
 np = cpu_only_import('numpy')
-import nvtx
-from cuml.internals.safe_imports import gpu_only_import
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 rmm = gpu_only_import('rmm')
 import warnings
 
@@ -411,7 +415,7 @@ class RandomForestRegressor(BaseRandomForestModel,
                                  algo=algo,
                                  fil_sparse_format=fil_sparse_format)
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="fit RF-Regressor @randomforestregressor.pyx",
         domain="cuml_python")
     @generate_docstring()
@@ -533,7 +537,7 @@ class RandomForestRegressor(BaseRandomForestModel,
         del(X_m)
         return preds
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="predict RF-Regressor @randomforestclassifier.pyx",
         domain="cuml_python")
     @insert_into_docstring(parameters=[('dense', '(n_samples, n_features)')],
@@ -593,7 +597,7 @@ class RandomForestRegressor(BaseRandomForestModel,
 
         return preds
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="score RF-Regressor @randomforestclassifier.pyx",
         domain="cuml_python")
     @insert_into_docstring(parameters=[('dense', '(n_samples, n_features)'),

python/cuml/experimental/fil/fil.pyx

@@ -16,7 +16,6 @@
 import cupy as cp
 import functools
 import numpy as np
-import nvtx
 import pathlib
 import treelite.sklearn
 import warnings
@@ -41,6 +40,13 @@ from cuml.internals.global_settings import GlobalSettings
 from cuml.internals.mem_type import MemoryType
 from pylibraft.common.handle cimport handle_t as raft_handle_t
 
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import_from,
+    null_decorator
+)
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
+
 cdef extern from "treelite/c_api.h":
     ctypedef void* ModelHandle
 
@@ -1047,7 +1053,7 @@ class ForestInference(UniversalBase, CMajorInputTagMixin):
             device_id=device_id
         )
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message='ForestInference.predict_proba',
         domain='cuml_python'
     )
@@ -1095,7 +1101,7 @@ class ForestInference(UniversalBase, CMajorInputTagMixin):
             )
         return self.forest.predict(X, preds=preds, chunk_size=chunk_size)
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message='ForestInference.predict',
         domain='cuml_python'
     )
@@ -1182,7 +1188,7 @@ class ForestInference(UniversalBase, CMajorInputTagMixin):
                 X, predict_type="default", preds=preds, chunk_size=chunk_size
             )
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message='ForestInference.predict_per_tree',
         domain='cuml_python'
     )
@@ -1235,7 +1241,7 @@ class ForestInference(UniversalBase, CMajorInputTagMixin):
             X, predict_type="per_tree", preds=preds, chunk_size=chunk_size
         )
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message='ForestInference.apply',
         domain='cuml_python'
     )

python/cuml/internals/base.pyx

@@ -20,9 +20,13 @@ import os
 import inspect
 import numbers
 from importlib import import_module
-from cuml.internals.safe_imports import cpu_only_import
+from cuml.internals.safe_imports import (
+    cpu_only_import,
+    gpu_only_import_from,
+    null_decorator
+)
 np = cpu_only_import('numpy')
-import nvtx
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 import typing
 
 import cuml
@@ -444,7 +448,7 @@ class Base(TagsMixin,
                                  addr=hex(id(self)))
                 msg = msg[5:]  # remove cuml.
                 func = getattr(self, func_name)
-                func = nvtx.annotate(message=msg, domain="cuml_python")(func)
+                func = nvtx_annotate(message=msg, domain="cuml_python")(func)
                 setattr(self, func_name, func)
 
 

python/cuml/internals/safe_imports.py

@@ -17,6 +17,8 @@
 
 import importlib
 import traceback
+
+from contextlib import contextmanager
 from cuml.internals.device_support import CPU_ENABLED, GPU_ENABLED
 from cuml.internals import logger
 
@@ -30,6 +32,7 @@ def return_false(*args, **kwargs):
     return False
 
 
+@contextmanager
 def null_decorator(*args, **kwargs):
     if len(kwargs) == 0 and len(args) == 1 and callable(args[0]):
         return args[0]

python/cuml/naive_bayes/naive_bayes.py

@@ -27,10 +27,13 @@
 from cuml.common import CumlArray
 import math
 import warnings
-import nvtx
-
-from cuml.internals.safe_imports import gpu_only_import
+from cuml.internals.safe_imports import (
+    gpu_only_import,
+    gpu_only_import_from,
+    null_decorator,
+)
 
+nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator)
 cp = gpu_only_import("cupy")
 cupyx = gpu_only_import("cupyx")
 
@@ -386,7 +389,7 @@ def fit(self, X, y, sample_weight=None) -> "GaussianNB":
             sample_weight=sample_weight,
         )
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="naive_bayes.GaussianNB._partial_fit", domain="cuml_python"
     )
     def _partial_fit(
@@ -816,7 +819,7 @@ def partial_fit(
             X, y, sample_weight=sample_weight, _classes=classes
         )
 
-    @nvtx.annotate(
+    @nvtx_annotate(
         message="naive_bayes._BaseDiscreteNB._partial_fit",
         domain="cuml_python",
     )