Coverage for cuda / pathfinder / _utils / env_vars.py: 94.29%

1# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4"""Centralized CUDA environment variable handling. 

5 

6This module defines the canonical search order for CUDA Toolkit environment variables 

7used throughout cuda-python packages (cuda.pathfinder, cuda.core, cuda.bindings). 

8 

9Search Order Priority: 

10 1. CUDA_PATH (higher priority) 

11 2. CUDA_HOME (lower priority) 

12 

13If both are set and differ, CUDA_PATH takes precedence and a warning is issued. 

14 

15Important Note on Caching: 

16 The result of get_cuda_path_or_home() is cached for the process lifetime. The first 

17 call determines the CUDA Toolkit path, and all subsequent calls return the cached 

18 value, even if environment variables change later. This ensures consistent behavior 

19 throughout the application lifecycle. 

20""" 

21 

22import functools (empty)

23import os (empty)

24import warnings (empty)

25 

26_CUDA_PATH_ENV_VARS_ORDERED = ("CUDA_PATH", "CUDA_HOME") (empty)

27 

28 

29def _paths_differ(a: str, b: str) -> bool: (empty)

30 """ 

31 Return True if paths are observably different. 

32 

33 Strategy: 

34 1) Compare os.path.normcase(os.path.normpath(...)) for quick, robust textual equality. 

35 - Handles trailing slashes and case-insensitivity on Windows. 

36 2) If still different AND both exist, use os.path.samefile to resolve symlinks/junctions. 

37 3) Otherwise (nonexistent paths or samefile unavailable), treat as different. 

38 """ 

39 norm_a = os.path.normcase(os.path.normpath(a)) 22 ctx1aghijklmnopqrstbNMcfde

40 norm_b = os.path.normcase(os.path.normpath(b)) 22 ctx1aghijklmnopqrstbNMcfde

41 if norm_a == norm_b: 22 ctx1aghijklmnopqrstbNMcfde

42 return False 16 ctx1aghijklmnopqrstM

43 

44 try: 7 ctx1bNMcfde

45 if os.path.exists(a) and os.path.exists(b): 7 ctx1bNMcfde

46 # samefile raises on non-existent paths; only call when both exist. 

47 return not os.path.samefile(a, b) 5 ctx1Ncfde

48 except OSError: 

49 # Fall through to "different" if samefile isn't applicable/available. 

50 pass 

51 

52 # If normalized strings differ and we couldn't prove they're the same entry, treat as different. 

53 return True 2 ctx1bM

54 

55 

56@functools.cache (empty)

57def get_cuda_path_or_home() -> str | None: (empty)

58 """Get CUDA Toolkit path from environment variables. 

59 

60 Returns the value of CUDA_PATH or CUDA_HOME. If both are set and differ, 

61 CUDA_PATH takes precedence and a warning is issued. 

62 

63 The result is cached for the process lifetime. The first call determines the CUDA 

64 Toolkit path, and subsequent calls return the cached value. 

65 

66 Returns: 

67 Path to CUDA Toolkit, or None if neither variable is set or all are empty. 

68 

69 Warnings: 

70 UserWarning: If multiple CUDA environment variables are set but point to 

71 different locations (only on the first call). 

72 

73 """ 

74 # Collect non-empty environment variables in priority order. 

75 # Empty strings are treated as undefined — no valid CUDA path is empty. 

76 set_vars = {} 46 ctx1auvwxyzABOCPghiDEFjklmnoQRGSpqHTrIUJKstbcVfdLe

77 for var in _CUDA_PATH_ENV_VARS_ORDERED: 46 ctx1auvwxyzABOCPghiDEFjklmnoQRGSpqHTrIUJKstbcVfdLe

78 val = os.environ.get(var) 46 ctx1auvwxyzABOCPghiDEFjklmnoQRGSpqHTrIUJKstbcVfdLe

79 if val: 46 ctx1auvwxyzABOCPghiDEFjklmnoQRGSpqHTrIUJKstbcVfdLe

80 set_vars[var] = val 38 ctx1auvwxyzABCghiDEFjklmnoGpqHrIJKstbcfdLe

81 

82 if not set_vars: 46 ctx1auvwxyzABOCPghiDEFjklmnoQRGSpqHTrIUJKstbcVfdLe

83 return None 8 ctx1OPQRSTUV

84 

85 # If multiple variables are set, check if they differ and warn 

86 if len(set_vars) > 1: 38 ctx1auvwxyzABCghiDEFjklmnoGpqHrIJKstbcfdLe

87 values = list(set_vars.items()) 20 ctx1aghijklmnopqrstbcfde

88 values_differ = False 20 ctx1aghijklmnopqrstbcfde

89 for i in range(len(values) - 1): 20 ctx1aghijklmnopqrstbcfde

90 if _paths_differ(values[i][1], values[i + 1][1]): 20 ctx1aghijklmnopqrstbcfde

91 values_differ = True 4 ctx1bcde

92 break 4 ctx1bcde

93 

94 if values_differ: 20 ctx1aghijklmnopqrstbcfde

95 var_list = "\n".join(f" {var}={val}" for var, val in set_vars.items()) 4 ctx1bcde

96 warnings.warn( 4 ctx1bcde

97 f"Multiple CUDA environment variables are set but differ:\n" 

98 f"{var_list}\n" 

99 f"Using {_CUDA_PATH_ENV_VARS_ORDERED[0]} (highest priority).", 

100 UserWarning, 

101 stacklevel=2, 

102 ) 

103 

104 # Return the first (highest priority) set variable 

105 return next(iter(set_vars.values())) 38 ctx1auvwxyzABCghiDEFjklmnoGpqHrIJKstbcfdLe