Coverage for biobb_gromacs/gromacs/mdrun_base.py: 55%

1#!/usr/bin/env python3 

2 

3"""Module containing the MdrunBase class shared by the GROMACS mdrun building blocks.""" 

4import re 

5import pathlib 

6from biobb_common.generic.biobb_object import BiobbObject 

7from biobb_common.tools import file_utils as fu 

8from biobb_gromacs.gromacs.common import get_gromacs_version 

9 

10 

11class MdrunBase(BiobbObject): 

12 """Shared behaviour for the GROMACS ``mdrun`` building blocks 

13 (:class:`Mdrun <gromacs.mdrun.Mdrun>`, 

14 :class:`MdrunPlumed <gromacs.mdrun_plumed.MdrunPlumed>` and 

15 :class:`MdrunMultidir <gromacs.mdrun_multidir.MdrunMultidir>`). 

16 

17 Concrete subclasses declare their own ``self.io_dict`` and assemble the 

18 building-block-specific part of ``self.cmd`` in ``launch()``. This base 

19 class centralises the common GROMACS runtime properties, the MPI runner 

20 prefix, the CPU/GPU/thread command-line flags and the ``-noappend`` output 

21 renaming logic. 

22 """ 

23 

24 def _init_common_properties(self, properties: dict) -> None: 

25 """Parse the GROMACS runtime properties shared by all mdrun wrappers. 

26 

27 Must be called from the subclass ``__init__`` after 

28 ``super().__init__(properties)`` (it relies on ``self.container_path``). 

29 """ 

30 # general mpi properties 

31 self.mpi_bin = properties.get('mpi_bin') 

32 self.mpi_np = properties.get('mpi_np') 

33 self.mpi_flags = properties.get('mpi_flags') 

34 # gromacs cpu mpi/openmp properties 

35 self.num_threads = str(properties.get('num_threads', '')) 

36 self.num_threads_mpi = str(properties.get('num_threads_mpi', '')) 

37 self.num_threads_omp = str(properties.get('num_threads_omp', '')) 

38 self.num_threads_omp_pme = str(properties.get('num_threads_omp_pme', '')) 

39 # gromacs gpus 

40 self.use_gpu = properties.get('use_gpu', False) # Adds: -nb gpu -pme gpu 

41 self.gpu_id = str(properties.get('gpu_id', '')) 

42 self.gpu_tasks = str(properties.get('gpu_tasks', '')) 

43 # gromacs 

44 self.checkpoint_time = properties.get('checkpoint_time') 

45 self.noappend = properties.get('noappend', False) 

46 

47 # Properties common in all GROMACS BB 

48 self.gmx_lib = properties.get('gmx_lib', None) 

49 self.binary_path: str = properties.get('binary_path', 'gmx') 

50 self.gmx_nobackup = properties.get('gmx_nobackup', True) 

51 self.gmx_nocopyright = properties.get('gmx_nocopyright', True) 

52 if self.gmx_nobackup: 

53 self.binary_path += ' -nobackup' 

54 if self.gmx_nocopyright: 

55 self.binary_path += ' -nocopyright' 

56 if (not self.mpi_bin) and (not self.container_path): 

57 self.gmx_version = get_gromacs_version(self.binary_path) 

58 

59 def _get_working_dir(self) -> str: 

60 """Return the working directory for the command (container or sandbox).""" 

61 if self.container_path: 

62 return self.container_volume_path if self.container_volume_path else "/data" 

63 return self.stage_io_dict.get('unique_dir', '') 

64 

65 def _prepend_mpi_runner(self) -> None: 

66 """Prepend the MPI runner (e.g. mpirun/srun) to ``self.cmd`` if configured.""" 

67 if self.mpi_bin: 

68 mpi_cmd = [self.mpi_bin] 

69 if self.mpi_np: 

70 mpi_cmd.append('-n') 

71 mpi_cmd.append(str(self.mpi_np)) 

72 if self.mpi_flags: 

73 mpi_cmd.append(self.mpi_flags) 

74 self.cmd = mpi_cmd + self.cmd 

75 

76 def _append_gmx_runtime_flags(self) -> None: 

77 """Append the shared CPU/GPU/thread flags to ``self.cmd`` and set GMXLIB.""" 

78 # gromacs cpu mpi/openmp properties 

79 if self.num_threads: 

80 fu.log(f'User added number of gmx threads: {self.num_threads}', self.out_log) 

81 self.cmd.append('-nt') 

82 self.cmd.append(self.num_threads) 

83 if self.num_threads_mpi: 

84 fu.log(f'User added number of gmx mpi threads: {self.num_threads_mpi}', self.out_log) 

85 self.cmd.append('-ntmpi') 

86 self.cmd.append(self.num_threads_mpi) 

87 if self.num_threads_omp: 

88 fu.log(f'User added number of gmx omp threads: {self.num_threads_omp}', self.out_log) 

89 self.cmd.append('-ntomp') 

90 self.cmd.append(self.num_threads_omp) 

91 if self.num_threads_omp_pme: 

92 fu.log(f'User added number of gmx omp_pme threads: {self.num_threads_omp_pme}', self.out_log) 

93 self.cmd.append('-ntomp_pme') 

94 self.cmd.append(self.num_threads_omp_pme) 

95 # GMX gpu properties 

96 if self.use_gpu: 

97 fu.log('Adding GPU specific settings adds: -nb gpu -pme gpu', self.out_log) 

98 self.cmd += ["-nb", "gpu", "-pme", "gpu"] 

99 if self.gpu_id: 

100 fu.log(f'list of unique GPU device IDs available to use: {self.gpu_id}', self.out_log) 

101 self.cmd.append('-gpu_id') 

102 self.cmd.append(self.gpu_id) 

103 if self.gpu_tasks: 

104 fu.log(f'list of GPU device IDs, mapping each PP task on each node to a device: {self.gpu_tasks}', self.out_log) 

105 self.cmd.append('-gputasks') 

106 self.cmd.append(self.gpu_tasks) 

107 

108 if self.noappend: 

109 self.cmd.append('-noappend') 

110 

111 if self.gmx_lib: 

112 self.env_vars_dict['GMXLIB'] = self.gmx_lib 

113 

114 def _apply_noappend_renaming(self) -> None: 

115 """Update expected output paths in the sandbox to catch -noappend renames. 

116 

117 GROMACS mdrun changes output file names from md.gro to md.part0001.gro 

118 when the ``-noappend`` flag is used. 

119 """ 

120 if not self.noappend: 

121 return 

122 

123 def capture_part_pattern(filename): 

124 """Capture the 'part' pattern followed by digits from a string.""" 

125 match = re.search(r'part\d+', filename) 

126 return match.group(0) if match else None 

127 

128 # List files in the staging directory and find the part000x pattern 

129 staging_path = self.stage_io_dict["unique_dir"] 

130 files_in_staging = list(pathlib.Path(staging_path).glob('*')) 

131 part_pattern = None 

132 for file in files_in_staging: 

133 part_pattern = capture_part_pattern(file.name) 

134 if part_pattern: 

135 break 

136 

137 # Update expected output files 

138 for file_ref, stage_file_path in self.stage_io_dict["out"].items(): 

139 if stage_file_path: 

140 parent_path = pathlib.Path(stage_file_path).parent 

141 file_stem = pathlib.Path(stage_file_path).stem 

142 file_suffix = pathlib.Path(stage_file_path).suffix 

143 # Rename all output files except checkpoint files 

144 if file_suffix != '.cpt' and part_pattern: 

145 new_file_name = f"{file_stem}.{part_pattern}{file_suffix}" 

146 self.stage_io_dict["out"][file_ref] = str(parent_path / new_file_name) 

147 

148 def _extra_copy_to_host(self) -> None: 

149 """Hook for subclasses to copy building-block-specific outputs. No-op by default.""" 

150 

151 def copy_to_host(self): 

152 """Copy output files back to the host, accounting for -noappend renames.""" 

153 self._apply_noappend_renaming() 

154 super().copy_to_host() 

155 self._extra_copy_to_host()