diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py index a5e45979fe2..fc0bea47180 100644 --- a/python/nx-cugraph/_nx_cugraph/__init__.py +++ b/python/nx-cugraph/_nx_cugraph/__init__.py @@ -301,6 +301,45 @@ def get_info(): .lower() == "true", } + + # Enable zero-code change usage with a simple environment variable + # by setting or updating other NETWORKX environment variables. + if os.environ.get("NX_CUGRAPH_AUTOCONFIG", "").strip().lower() == "true": + from itertools import chain + + def update_env_var(varname): + """Add "cugraph" to a list of backend names environment variable.""" + if varname not in os.environ: + os.environ[varname] = "cugraph" + return + string = os.environ[varname] + vals = [ + stripped for x in string.strip().split(",") if (stripped := x.strip()) + ] + if "cugraph" not in vals: + # Should we append or prepend? Let's be first! + os.environ[varname] = ",".join(chain(["cugraph"], vals)) + + # Automatically convert NetworkX Graphs to nx-cugraph for algorithms + if (varname := "NETWORKX_BACKEND_PRIORITY_ALGOS") in os.environ: + # "*_ALGOS" is given priority in NetworkX >=3.4 + update_env_var(varname) + # But update this too to "just work" if users mix env vars and nx versions + os.environ["NETWORKX_BACKEND_PRIORITY"] = os.environ[varname] + else: + update_env_var("NETWORKX_BACKEND_PRIORITY") + # And for older NetworkX versions + update_env_var("NETWORKX_AUTOMATIC_BACKENDS") # For NetworkX 3.2 + update_env_var("NETWORKX_GRAPH_CONVERT") # For NetworkX 3.0 and 3.1 + # Automatically create nx-cugraph Graph from graph generators + update_env_var("NETWORKX_BACKEND_PRIORITY_GENERATORS") + # Run default NetworkX implementation (in >=3.4) if not implemented by nx-cugraph + if (varname := "NETWORKX_FALLBACK_TO_NX") not in os.environ: + os.environ[varname] = "true" + # Cache graph conversions (default is False in NetworkX 3.2 + if (varname := "NETWORKX_CACHE_CONVERTED_GRAPHS") not in os.environ: + os.environ[varname] = "true" + return d