# CelerisLab/utils.py import pycuda.driver as cuda import subprocess import json import os from typing import NamedTuple, Optional, List, Tuple, Union class CudaDeviceInfo(NamedTuple): name: str compute_capability: str multiprocessors: int total_global_memory: int max_shared_memory_per_block: int max_threads_per_block: int max_blocks_per_multiprocessor: int device_interconnect: Optional[str] = None class FlowFieldConfig(NamedTuple): data_type: str dimensionality: int lattice: int field_dim_in_U: Tuple[int, int, int] viscosity: float velocity: float boundary_conditions: Tuple[str, str, str, str, str, str] class CudaConfig(NamedTuple): multi_gpu: bool gpu_connection: str required_cuda_capability: str threads_per_block: int unit_dimensions: Tuple[int, int, int] def check_cuda_device_availability(device_id=0): if cuda.Device.count() == 0: raise RuntimeError("No CUDA device is available.") if device_id < 0 or device_id >= cuda.Device.count(): raise ValueError( f"Invalid device_id {device_id}. Must be between 0 and {cuda.Device.count() - 1}." ) try: subprocess.check_output(["nvidia-smi", "--version"]) except subprocess.CalledProcessError: raise RuntimeError("nvidia-smi is not available or not installed correctly.") def query_cuda_device_info(device_id=0) -> CudaDeviceInfo: check_cuda_device_availability(device_id) try: output = subprocess.check_output( ["nvidia-smi", "-q", "-d", "TOPOLOGY", "-i", str(device_id)], text=True ) if "NVLink" in output: device_interconnect = "NVLink" elif "PCIe" in output: device_interconnect = "PCIe" else: device_interconnect = "Unknown" except Exception as e: device_interconnect = None device = cuda.Device(device_id) return CudaDeviceInfo( name=device.name(), compute_capability=f"{device.compute_capability()[0]}.{device.compute_capability()[1]}", multiprocessors=device.get_attribute( cuda.device_attribute.MULTIPROCESSOR_COUNT ), total_global_memory=device.total_memory(), max_shared_memory_per_block=device.get_attribute( cuda.device_attribute.MAX_SHARED_MEMORY_PER_BLOCK ), max_threads_per_block=device.get_attribute( cuda.device_attribute.MAX_THREADS_PER_BLOCK ), max_blocks_per_multiprocessor=device.get_attribute( cuda.device_attribute.MAX_BLOCKS_PER_MULTIPROCESSOR ), device_interconnect=device_interconnect, ) def find_config_file(config_filename: str, config_path: Optional[str] = None) -> str: """ Find configuration file by searching in multiple locations. Search priority: 1. Provided config_path (if given) 2. Environment variable CELERISLAB_CONFIG_DIR 3. Current working directory ./configs/ 4. Package installation location (relative to this utils.py file) Args: config_filename: Name of the config file (e.g., 'config_cuda.json') config_path: Optional explicit path to config file Returns: Absolute path to the config file Raises: FileNotFoundError: If config file cannot be found in any location """ search_paths = [] # Priority 1: Explicit path provided if config_path: search_paths.append(config_path) # Priority 2: Environment variable env_config_dir = os.environ.get('CELERISLAB_CONFIG_DIR') if env_config_dir: search_paths.append(os.path.join(env_config_dir, config_filename)) # Priority 3: Current working directory search_paths.append(os.path.join(os.getcwd(), 'configs', config_filename)) # Priority 4: Package installation location (relative to this utils.py) # configs are in lbm/configs/ package_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) search_paths.append(os.path.join(package_root, 'lbm', 'configs', config_filename)) # Search for the file for path in search_paths: if os.path.isfile(path): return os.path.abspath(path) # File not found, provide helpful error message error_msg = f"Configuration file '{config_filename}' not found. Searched in:\n" for path in search_paths: error_msg += f" - {path}\n" error_msg += "\nTo fix this, you can:\n" error_msg += " 1. Set CELERISLAB_CONFIG_DIR environment variable\n" error_msg += " 2. Place config files in ./configs/ directory\n" error_msg += " 3. Provide explicit config_path parameter" raise FileNotFoundError(error_msg) def load_flow_field_config(config_path: Optional[str] = None) -> FlowFieldConfig: """ Load flow field configuration from JSON file. Args: config_path: Optional path to config file. If None, searches in standard locations. Can be relative path like 'configs/config_flowfield.json' or just filename. Returns: FlowFieldConfig object """ # Determine config filename and full path if config_path: # Check if it's just a filename or a path if os.path.basename(config_path) == config_path: # Just a filename, search for it config_file = find_config_file(config_path, None) else: # It's a path, use it if exists, otherwise try to find the basename if os.path.isfile(config_path): config_file = config_path else: config_file = find_config_file(os.path.basename(config_path), None) else: # No path provided, search for default filename config_file = find_config_file('config_flowfield.json', None) try: with open(config_file, "r") as file: config = json.load(file) required_keys = [ "data_type", "dimensionality", "lattice", "field_dim_in_U", "viscosity", "boundary_conditions", ] if not all(key in config for key in required_keys): raise ValueError("Missing required configuration items.") if config["data_type"] not in ["FP32", "FP64"]: raise ValueError("Data type must be either FP32 or FP64.") if config["dimensionality"] not in [2, 3]: raise ValueError("Dimensionality must be either 2 or 3.") if config["dimensionality"] == 2 and config["field_dim_in_U"][2] != 1: raise ValueError( "Field dimensions must be 1 in the third dimension for 2D simulations." ) if config["lattice"] not in [9]: raise ValueError("Lattice must be either 9 or 19.") boundary_conditions = tuple( condition for key in ["x", "y", "z"] for condition in config["boundary_conditions"].get(key, []) ) if len(boundary_conditions) != 6: raise ValueError("Boundary conditions must contain exactly six elements.") return FlowFieldConfig( data_type=config["data_type"], dimensionality=config["dimensionality"], lattice=config["lattice"], field_dim_in_U=tuple(config["field_dim_in_U"]), viscosity=config["viscosity"], velocity=config["velocity"], boundary_conditions=boundary_conditions, ) except Exception as e: raise RuntimeError(f"Failed to load or parse the flow field configuration: {e}") def load_cuda_config(config_path: Optional[str] = None) -> CudaConfig: """ Load CUDA configuration from JSON file. Args: config_path: Optional path to config file. If None, searches in standard locations. Can be relative path like 'configs/config_cuda.json' or just filename. Returns: CudaConfig object """ # Determine config filename and full path if config_path: # Check if it's just a filename or a path if os.path.basename(config_path) == config_path: # Just a filename, search for it config_file = find_config_file(config_path, None) else: # It's a path, use it if exists, otherwise try to find the basename if os.path.isfile(config_path): config_file = config_path else: config_file = find_config_file(os.path.basename(config_path), None) else: # No path provided, search for default filename config_file = find_config_file('config_cuda.json', None) try: with open(config_file, "r") as file: config = json.load(file) required_keys = [ "multi_gpu", "gpu_connection", "required_cuda_capability", "threads_per_block", "X_1U", "Y_1U", "Z_1U", ] if not all(key in config for key in required_keys): raise ValueError("Missing required configuration items.") return CudaConfig( multi_gpu=config["multi_gpu"], gpu_connection=config["gpu_connection"], required_cuda_capability=config["required_cuda_capability"], threads_per_block=config["threads_per_block"], unit_dimensions=(config["X_1U"], config["Y_1U"], config["Z_1U"]), ) except Exception as e: raise RuntimeError(f"Failed to load or parse the CUDA configuration: {e}") def check_cuda_capability( field_config: FlowFieldConfig, cuda_config: CudaConfig, device_id: Union[int, List[int]] = None, ): SAFE_FACTOR = 0.8 if cuda_config.multi_gpu: if device_id is None or isinstance(device_id, int): raise ValueError("Multi-GPU support requires a list of device IDs.") raise NotImplementedError("Multi-GPU support is not implemented yet.") else: if isinstance(device_id, list): if len(device_id) > 1: raise ValueError( "Single-GPU mode does not support multiple device IDs." ) device_id = device_id[0] elif device_id is None: device_id = 0 device_info = query_cuda_device_info(device_id) if device_info.compute_capability != cuda_config.required_cuda_capability: raise ValueError( f"Device {device_info.name} has compute capability {device_info.compute_capability}, but {cuda_config.required_cuda_capability} is required." ) field_size = sum( size * unit for size, unit in zip( field_config.field_dim_in_U, cuda_config.unit_dimensions ) ) if ( device_info.total_global_memory * SAFE_FACTOR < calc_field_memory_consumption( field_size, field_config.dimensionality, field_config.lattice, field_config.data_type, ) ): raise ValueError( f"Device {device_info.name} does not have enough memory to store the flow field." ) if ( device_info.max_threads_per_block * SAFE_FACTOR < cuda_config.threads_per_block ): raise ValueError( f"Device {device_info.name} does not have enough threads per block to run the simulation." ) block_size = cuda_config.threads_per_block if ( device_info.max_shared_memory_per_block * SAFE_FACTOR < 2 * calc_field_memory_consumption( block_size, field_config.dimensionality, field_config.lattice, field_config.data_type, ) ): raise ValueError( f"Device {device_info.name} does not have enough shared memory per block to run the simulation." ) def calc_field_memory_consumption( field_size: int, dimensionality: int, directions: int, data_type: str ) -> int: if data_type == "FP32": data_size = 4 elif data_type == "FP64": data_size = 8 else: raise ValueError(f"Unsupported data type {data_type}.") return ( field_size * directions * data_size * 2 + field_size * dimensionality * data_size + field_size )