From b0c32a95c9cca11025fc03ebb437bec8d1d48174 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 10 Dec 2024 20:07:16 -0600 Subject: [PATCH] Example using nx dispatch hooks (WIP) to print timing --- _nx_cugraph/__init__.py | 14 ++++++++++++++ pyproject.toml | 2 +- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/_nx_cugraph/__init__.py b/_nx_cugraph/__init__.py index a52583d4d..7c681a49c 100644 --- a/_nx_cugraph/__init__.py +++ b/_nx_cugraph/__init__.py @@ -31,6 +31,19 @@ # always be in .. format. (_version_major, _version_minor) = __version__.split(".")[:2] + +def dispatch_hook_impl(hook_name, hook_token, **kwargs): + import time + + if hook_name == "on_call_with_backend_begin": + hook_token["cugraph"]["t0"] = time.time() + elif hook_name == "on_call_with_backend_end": + t = time.time() - hook_token["cugraph"]["t0"] + backend_name = kwargs["backend_name"] + dispatch_name = hook_token.dispatchable.name + print(f"{backend_name!r} backend ran {dispatch_name} in {t:.3g} seconds") + + # Entries between BEGIN and END are automatically generated _info = { "backend_name": "cugraph", @@ -38,6 +51,7 @@ "package": "nx_cugraph", "url": "https://rapids.ai/nx-cugraph", "short_summary": "GPU-accelerated backend.", + "dispatch_hook_impl": dispatch_hook_impl, # "description": "TODO", "functions": { # BEGIN: functions diff --git a/pyproject.toml b/pyproject.toml index 19b7f8299..8876ec6ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -242,7 +242,7 @@ ignore = [ "__init__.py" = ["F401"] # Allow unused imports (w/o defining `__all__`) # Allow assert, print, RNG, and no docstring "nx_cugraph/**/tests/*py" = ["S101", "S311", "T201", "D103", "D100"] -"_nx_cugraph/__init__.py" = ["E501"] +"_nx_cugraph/__init__.py" = ["E501", "T201"] "nx_cugraph/__init__.py" = ["E402"] # Allow module level import not at top of file "nx_cugraph/algorithms/**/*py" = ["D205", "D401"] # Allow flexible docstrings for algorithms "nx_cugraph/generators/**/*py" = ["D205", "D401"] # Allow flexible docstrings for generators