Skip to content

Commit 4ed25f4

Browse files
xylarxylar
authored
Merge pull request #230 from xylar/add-spack-module-env-util
Add functions for creating spack shell scripts from `config_machines.xml`
2 parents 777f035 + 78d8479 commit 4ed25f4

File tree

6 files changed

+356
-7
lines changed

6 files changed

+356
-7
lines changed

‎docs/developers_guide/api.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ documentation.
3939
make_spack_env
4040
get_spack_script
4141
get_modules_env_vars_and_mpi_compilers
42+
extract_machine_config
43+
config_to_shell_script
44+
extract_spack_from_config_machines
45+
list_machine_compiler_mpilib
4246
```
4347

4448
## sync

‎mache/spack/__init__.py‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,29 @@
1-
import os
2-
import subprocess
1+
import os as os
2+
import subprocess as subprocess
33
from importlib import resources as importlib_resources
44

5-
import yaml
6-
from jinja2 import Template
7-
8-
from mache.machine_info import MachineInfo, discover_machine
9-
from mache.version import __version__
5+
import yaml as yaml
6+
from jinja2 import Template as Template
7+
8+
from mache.machine_info import (
9+
MachineInfo as MachineInfo,
10+
)
11+
from mache.machine_info import (
12+
discover_machine as discover_machine,
13+
)
14+
from mache.spack.config_machines import (
15+
config_to_shell_script as config_to_shell_script,
16+
)
17+
from mache.spack.config_machines import (
18+
extract_machine_config as extract_machine_config,
19+
)
20+
from mache.spack.config_machines import (
21+
extract_spack_from_config_machines as extract_spack_from_config_machines,
22+
)
23+
from mache.spack.list import (
24+
list_machine_compiler_mpilib as list_machine_compiler_mpilib,
25+
)
26+
from mache.version import __version__ as __version__
1027

1128

1229
def make_spack_env(

‎mache/spack/config_machines.py‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
import re
2+
from collections import defaultdict
3+
4+
from lxml import etree
5+
6+
7+
def extract_machine_config(xml_file, machine, compiler, mpilib):
8+
"""
9+
Extract the machine configuration from the XML file.
10+
11+
Parameters
12+
----------
13+
xml_file : str
14+
Path to the XML file.
15+
machine : str
16+
Machine name.
17+
compiler : str
18+
Compiler name.
19+
mpilib : str
20+
MPI library name.
21+
22+
Returns
23+
-------
24+
etree.Element or None
25+
The XML element of the machine configuration if found, otherwise None.
26+
"""
27+
tree = etree.parse(xml_file)
28+
root = tree.getroot()
29+
30+
for mach in root.findall('machine'):
31+
if mach.get('MACH') == machine:
32+
for mod_sys in mach.findall('module_system'):
33+
for mod in mod_sys.findall('modules'):
34+
if not (
35+
re.match(mod.get('compiler', '.*'), compiler)
36+
and re.match(mod.get('mpilib', '.*'), mpilib)
37+
and re.match(mod.get('DEBUG', '.*'), 'FALSE')
38+
):
39+
mod_sys.remove(mod)
40+
for env_vars in mach.findall('environment_variables'):
41+
if not (
42+
re.match(env_vars.get('compiler', '.*'), compiler)
43+
and re.match(env_vars.get('mpilib', '.*'), mpilib)
44+
and re.match(env_vars.get('DEBUG', '.*'), 'FALSE')
45+
):
46+
mach.remove(env_vars)
47+
return mach
48+
return None
49+
50+
51+
def config_to_shell_script(config, shell_type):
52+
"""
53+
Convert the machine configuration to a shell script.
54+
55+
Parameters
56+
----------
57+
config : etree.Element
58+
The XML element of the machine configuration.
59+
shell_type : str
60+
The type of shell script to generate ('sh' or 'csh').
61+
62+
Returns
63+
-------
64+
str
65+
The shell script as a string.
66+
"""
67+
script_lines = []
68+
69+
# TODO: possibly replace \: with :
70+
init_path = None
71+
for init in config.findall('.//module_system/init_path'):
72+
if init.get('lang') == shell_type:
73+
init_path = init.text
74+
break
75+
76+
if init_path is not None:
77+
init_path = init_path.replace(';', '\n')
78+
script_lines.append(f'source {init_path}')
79+
script_lines.append('')
80+
81+
module_commands = defaultdict(list)
82+
e3sm_hdf5_netcdf_modules = defaultdict(list)
83+
for module in config.findall('.//module_system/modules'):
84+
for command in module.findall('command'):
85+
name = command.get('name')
86+
value = command.text
87+
if value:
88+
if 'command' != 'unload' and 'python' in value:
89+
# we don't want to load E3SM's python module
90+
continue
91+
elif 'command' != 'unload' and re.search(
92+
r'hdf5|netcdf', value, re.IGNORECASE
93+
):
94+
# we want to remove hdf5 and netcdf in all cases
95+
e3sm_hdf5_netcdf_modules[name].append(value)
96+
else:
97+
module_commands[name].append(value)
98+
elif name not in module_commands:
99+
module_commands[name] = []
100+
101+
script_lines.extend(
102+
_convert_module_commands_to_script_lines(module_commands, shell_type)
103+
)
104+
script_lines.append('')
105+
106+
if e3sm_hdf5_netcdf_modules:
107+
script_lines.append('{% if e3sm_hdf5_netcdf %}')
108+
script_lines.extend(
109+
_convert_module_commands_to_script_lines(
110+
e3sm_hdf5_netcdf_modules, shell_type
111+
)
112+
)
113+
script_lines.append('{% endif %}')
114+
script_lines.append('')
115+
116+
script_lines.extend(_convert_env_vars_to_script_lines(config, shell_type))
117+
118+
script_lines.append('')
119+
120+
return '\n'.join(script_lines)
121+
122+
123+
def extract_spack_from_config_machines(
124+
machine, compiler, mpilib, shell, output
125+
):
126+
"""
127+
Extract machine configuration from XML and write it to a shell script.
128+
129+
Parameters
130+
----------
131+
machine : str
132+
Machine name.
133+
compiler : str
134+
Compiler name.
135+
mpilib : str
136+
MPI library name.
137+
shell : str
138+
Shell script type ('sh' or 'csh').
139+
output : str
140+
Output file to write the shell script.
141+
"""
142+
config_filename = 'mache/cime_machine_config/config_machines.xml'
143+
144+
config = extract_machine_config(config_filename, machine, compiler, mpilib)
145+
if config is None:
146+
raise ValueError(
147+
f'No configuration found for machine={machine}, '
148+
f'compiler={compiler}, mpilib={mpilib}'
149+
)
150+
151+
script = config_to_shell_script(config, shell)
152+
with open(output, 'w') as f:
153+
f.write(script)
154+
155+
156+
def _convert_module_commands_to_script_lines(module_commands, shell_type):
157+
"""
158+
Convert module commands to script lines.
159+
160+
Parameters
161+
----------
162+
module_commands : dict
163+
Dictionary of module commands.
164+
shell_type : str
165+
The type of shell script to generate ('sh' or 'csh').
166+
167+
Returns
168+
-------
169+
list
170+
List of script lines.
171+
"""
172+
script_lines = []
173+
for i, (name, values) in enumerate(module_commands.items()):
174+
if name == 'unload':
175+
name = 'rm'
176+
# the module rm commands produce distracting output so we
177+
# silence them
178+
suffix = ' &> /dev/null'
179+
else:
180+
suffix = ''
181+
if values:
182+
# same code regardless of shell
183+
if name == 'switch':
184+
# switch commands need to be handled one at a time
185+
for value in values:
186+
script_lines.append(f'module {name} {value}')
187+
else:
188+
script_lines.append(f'module {name} \\')
189+
for value in values:
190+
script_lines.append(f' {value} \\')
191+
script_lines[-1] = script_lines[-1].rstrip(' \\') + suffix
192+
else:
193+
script_lines.append(f'module {name}{suffix}')
194+
if i < len(module_commands) - 1:
195+
script_lines.append('')
196+
return script_lines
197+
198+
199+
def _convert_env_vars_to_script_lines(config, shell_type):
200+
"""
201+
Convert environment variables to script lines.
202+
203+
Parameters
204+
----------
205+
config : etree.Element
206+
The XML element of the machine configuration.
207+
shell_type : str
208+
The type of shell script to generate ('sh' or 'csh').
209+
210+
Returns
211+
-------
212+
list
213+
List of script lines.
214+
"""
215+
script_lines = []
216+
e3sm_hdf5_netcdf_env_vars = []
217+
for env_var in config.findall('.//environment_variables/env'):
218+
name = env_var.get('name')
219+
value = env_var.text
220+
if value is None:
221+
value = ''
222+
if '$SHELL' in value:
223+
# at least for now, these $SHELL environment variables are
224+
# E3SM-specific so we'll leave them out
225+
continue
226+
if 'perl' in value:
227+
# we don't want to add E3SM's perl path
228+
continue
229+
if name.startswith('OMP_'):
230+
# OpenMP environment variables cause trouble with ESMF
231+
continue
232+
233+
value = re.sub(r'\$ENV{([^}]+)}', r'${\1}', value)
234+
if 'NETCDF' in name and 'PATH' in name:
235+
e3sm_hdf5_netcdf_env_vars.append((name, value))
236+
if name == 'NETCDF_PATH':
237+
# also set the NETCDF_C_PATH and NETCDF_FORTRAN_PATH, needed
238+
# by MPAS standalone components
239+
e3sm_hdf5_netcdf_env_vars.append(('NETCDF_C_PATH', value))
240+
e3sm_hdf5_netcdf_env_vars.append(
241+
('NETCDF_FORTRAN_PATH', value)
242+
)
243+
else:
244+
if shell_type == 'sh':
245+
script_lines.append(f'export {name}="{value}"')
246+
elif shell_type == 'csh':
247+
script_lines.append(f'setenv {name} "{value}"')
248+
249+
script_lines.append('')
250+
251+
if e3sm_hdf5_netcdf_env_vars:
252+
script_lines.append('{% if e3sm_hdf5_netcdf %}')
253+
for name, value in e3sm_hdf5_netcdf_env_vars:
254+
if shell_type == 'sh':
255+
script_lines.append(f'export {name}="{value}"')
256+
elif shell_type == 'csh':
257+
script_lines.append(f'setenv {name} "{value}"')
258+
script_lines.append('{% endif %}')
259+
260+
return script_lines

‎mache/spack/list.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import importlib.resources
2+
import re
3+
from pathlib import Path
4+
5+
6+
def list_machine_compiler_mpilib():
7+
"""
8+
List tuples of machine, compiler, and MPI library parsed from the name of
9+
YAML files in the `mache.spack.templates` directory.
10+
11+
Returns
12+
-------
13+
list of tuples
14+
Each tuple contains (machine, compiler, mpilib).
15+
"""
16+
machine_compiler_mpi_list = []
17+
pattern = re.compile(r'([\w-]+)_([\w-]+)_([\w-]+)')
18+
19+
files = sorted(
20+
importlib.resources.files('mache.spack.templates').iterdir(), key=str
21+
)
22+
for file in files:
23+
file_path = Path(str(file))
24+
if file_path.suffix == '.yaml':
25+
match = pattern.match(file_path.stem)
26+
if match:
27+
machine, compiler, mpilib = match.groups()
28+
machine_compiler_mpi_list.append((machine, compiler, mpilib))
29+
30+
return machine_compiler_mpi_list

‎utils/README.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,15 @@ A developer can then copy `new_config_machines.xml` into
1414
`mache/cime_machine_config/config_machines.xml` as part of a PR that makes
1515
relevant updates. They should also make the changes associated with the
1616
differences that this utility displays in the appropriate `mache/spack/templates` files.
17+
18+
## extract spack shell scripts from CIME machine config
19+
20+
The `extract_all_spack_from_config_machines.py` produces shell scripts for
21+
each machine, compiler and MPI library supported for spack builds from `mache`.
22+
The scripts are places in `new_spack` and can be moved or copied to
23+
`mache/spack` (after vetting!).
24+
25+
To run the utiltiy, make sure the `mache_dev` environment is active, then run:
26+
```
27+
./utils/extract_all_spack_from_config_machines.py
28+
```

‎utils/extract_all_spack_from_config_machines.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#!/usr/bin/env python3
2+
import os
3+
4+
from mache.spack.config_machines import extract_spack_from_config_machines
5+
from mache.spack.list import list_machine_compiler_mpilib
6+
7+
8+
def main():
9+
"""
10+
Main function to parse arguments and extract machine configuration.
11+
"""
12+
13+
directory = 'spack_scripts'
14+
os.makedirs(directory, exist_ok=True)
15+
16+
for machine, compiler, mpilib in list_machine_compiler_mpilib():
17+
print(f'Extracting {machine}, {compiler}, {mpilib}')
18+
for shell in ['sh', 'csh']:
19+
filename = f'{directory}/{machine}_{compiler}_{mpilib}.{shell}'
20+
extract_spack_from_config_machines(
21+
machine, compiler, mpilib, shell, filename
22+
)
23+
24+
25+
if __name__ == '__main__':
26+
main()

0 commit comments

Comments
 (0)