Skip to content

Commit 579e0af

Browse files
allnesallnes
committed
Format MPI extension Python helpers
1 parent 4cd1142 commit 579e0af

2 files changed

Lines changed: 98 additions & 25 deletions

File tree

‎.github/actions/setup-mpi-extensions/setup_mpi_extensions.py‎

Lines changed: 35 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424

2525

2626
def parse_args() -> argparse.Namespace:
27-
parser = argparse.ArgumentParser(description="Set up mpi-extensions Open MPI for GitHub Actions.")
27+
parser = argparse.ArgumentParser(
28+
description="Set up mpi-extensions Open MPI for GitHub Actions."
29+
)
2830
parser.add_argument("--repo-root", required=True, type=Path)
2931
parser.add_argument("--prefix", required=True, type=Path)
3032
parser.add_argument("--platform", required=True)
@@ -35,7 +37,9 @@ def parse_args() -> argparse.Namespace:
3537
return parser.parse_args()
3638

3739

38-
def run(command: list[str], *, env: dict[str, str] | None = None, check: bool = True) -> subprocess.CompletedProcess[str]:
40+
def run(
41+
command: list[str], *, env: dict[str, str] | None = None, check: bool = True
42+
) -> subprocess.CompletedProcess[str]:
3943
print("+ " + " ".join(command), flush=True)
4044
return subprocess.run(command, check=check, env=env)
4145

@@ -68,7 +72,12 @@ def purge_macos_mpi() -> None:
6872
if not shutil.which("brew"):
6973
return
7074
for package in MPI_PACKAGES_MACOS:
71-
if subprocess.run(["brew", "list", "--versions", package], check=False).returncode == 0:
75+
if (
76+
subprocess.run(
77+
["brew", "list", "--versions", package], check=False
78+
).returncode
79+
== 0
80+
):
7281
run(["brew", "uninstall", "--ignore-dependencies", "--force", package])
7382

7483

@@ -91,9 +100,26 @@ def prepare_python(runner_temp: Path) -> Path:
91100
run([str(python_bin), "-m", "pip", "install", "requests>=2.31,<3"])
92101
return python_bin
93102
except subprocess.CalledProcessError:
94-
print("::warning::Python venv setup failed; trying the current interpreter with pip")
95-
if subprocess.run([sys.executable, "-m", "pip", "install", "--user", "requests>=2.31,<3"], check=False).returncode != 0:
96-
run([sys.executable, "-m", "pip", "install", "--break-system-packages", "requests>=2.31,<3"])
103+
print(
104+
"::warning::Python venv setup failed; trying the current interpreter with pip"
105+
)
106+
if (
107+
subprocess.run(
108+
[sys.executable, "-m", "pip", "install", "--user", "requests>=2.31,<3"],
109+
check=False,
110+
).returncode
111+
!= 0
112+
):
113+
run(
114+
[
115+
sys.executable,
116+
"-m",
117+
"pip",
118+
"install",
119+
"--break-system-packages",
120+
"requests>=2.31,<3",
121+
]
122+
)
97123
return Path(sys.executable)
98124

99125

@@ -114,7 +140,9 @@ def install_openmpi(args: argparse.Namespace, python_bin: Path) -> None:
114140
)
115141

116142

117-
def append_github_environment(prefix: Path, github_env: Path, github_path: Path) -> None:
143+
def append_github_environment(
144+
prefix: Path, github_env: Path, github_path: Path
145+
) -> None:
118146
lib_path = prefix / "lib"
119147
env_lines = [
120148
f"MPI_EXTENSIONS_HOME={prefix}",

‎scripts/install_mpi_extensions.py‎

Lines changed: 63 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -99,19 +99,29 @@ def select_asset(release: dict[str, Any], asset_name: str) -> dict[str, Any]:
9999
for asset in release.get("assets", []):
100100
if asset.get("name") == asset_name:
101101
return asset
102-
available = ", ".join(asset.get("name", "<unnamed>") for asset in release.get("assets", []))
103-
raise SystemExit(f"Release asset not found: {asset_name}. Available assets: {available}")
102+
available = ", ".join(
103+
asset.get("name", "<unnamed>") for asset in release.get("assets", [])
104+
)
105+
raise SystemExit(
106+
f"Release asset not found: {asset_name}. Available assets: {available}"
107+
)
104108

105109

106-
def select_archive_asset(release: dict[str, Any], target_platform: str) -> dict[str, Any]:
107-
pattern = re.compile(rf"^mpi-extensions-openmpi-.+-{re.escape(target_platform)}\.tar\.gz$")
110+
def select_archive_asset(
111+
release: dict[str, Any], target_platform: str
112+
) -> dict[str, Any]:
113+
pattern = re.compile(
114+
rf"^mpi-extensions-openmpi-.+-{re.escape(target_platform)}\.tar\.gz$"
115+
)
108116
matches = [
109117
asset
110118
for asset in release.get("assets", [])
111119
if pattern.match(str(asset.get("name", "")))
112120
]
113121
if len(matches) != 1:
114-
available = ", ".join(asset.get("name", "<unnamed>") for asset in release.get("assets", []))
122+
available = ", ".join(
123+
asset.get("name", "<unnamed>") for asset in release.get("assets", [])
124+
)
115125
raise SystemExit(
116126
f"Expected exactly one nightly archive for {target_platform}, "
117127
f"found {len(matches)}. Available assets: {available}"
@@ -148,30 +158,43 @@ def safe_extract(archive: Path, destination: Path) -> Path:
148158
continue
149159
roots.add(parts[0])
150160
target = (destination / member.name).resolve()
151-
if target != resolved_destination and resolved_destination not in target.parents:
161+
if (
162+
target != resolved_destination
163+
and resolved_destination not in target.parents
164+
):
152165
raise SystemExit(f"Archive contains unsafe path: {member.name}")
153166
if len(roots) != 1:
154-
raise SystemExit(f"Expected archive to contain one top-level directory, found: {roots}")
167+
raise SystemExit(
168+
f"Expected archive to contain one top-level directory, found: {roots}"
169+
)
155170
tar.extractall(destination, members=members)
156171
return destination / roots.pop()
157172

158173

159174
def install_tree(source_root: Path, prefix: Path, force: bool) -> None:
160175
required_tools = ["mpicc", "mpicxx", "mpirun", "mpiexec"]
161-
missing_tools = [tool for tool in required_tools if not (source_root / "bin" / tool).exists()]
176+
missing_tools = [
177+
tool for tool in required_tools if not (source_root / "bin" / tool).exists()
178+
]
162179
if missing_tools:
163-
raise SystemExit(f"Archive is missing required MPI tools: {', '.join(missing_tools)}")
180+
raise SystemExit(
181+
f"Archive is missing required MPI tools: {', '.join(missing_tools)}"
182+
)
164183

165184
if prefix.exists():
166185
if not force and any(prefix.iterdir()):
167-
raise SystemExit(f"Installation prefix is not empty: {prefix}. Use --force to replace it.")
186+
raise SystemExit(
187+
f"Installation prefix is not empty: {prefix}. Use --force to replace it."
188+
)
168189
shutil.rmtree(prefix)
169190
prefix.parent.mkdir(parents=True, exist_ok=True)
170191
shutil.copytree(source_root, prefix, symlinks=True)
171192

172193

173194
def run_tool(command: list[str]) -> subprocess.CompletedProcess[str]:
174-
return subprocess.run(command, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
195+
return subprocess.run(
196+
command, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
197+
)
175198

176199

177200
def otool_lines(path: Path) -> list[str]:
@@ -185,7 +208,9 @@ def patch_macos_install_names(prefix: Path) -> None:
185208
if platform.system() != "Darwin":
186209
return
187210
if not shutil.which("install_name_tool") or not shutil.which("otool"):
188-
raise SystemExit("macOS install_name_tool and otool are required to install mpi-extensions")
211+
raise SystemExit(
212+
"macOS install_name_tool and otool are required to install mpi-extensions"
213+
)
189214

190215
lib_dir = prefix / "lib"
191216
if not lib_dir.exists():
@@ -208,12 +233,19 @@ def patch_macos_install_names(prefix: Path) -> None:
208233
result = run_tool(["otool", "-D", str(path)])
209234
if result.returncode != 0:
210235
continue
211-
install_names = [line.strip() for line in result.stdout.splitlines()[1:] if line.strip()]
236+
install_names = [
237+
line.strip() for line in result.stdout.splitlines()[1:] if line.strip()
238+
]
212239
if install_names:
213240
local_name = local_libs.get(Path(install_names[0]).name)
214241
if local_name:
215242
subprocess.run(
216-
["install_name_tool", "-id", f"@rpath/{local_name.name}", str(path)],
243+
[
244+
"install_name_tool",
245+
"-id",
246+
f"@rpath/{local_name.name}",
247+
str(path),
248+
],
217249
check=True,
218250
)
219251

@@ -231,7 +263,13 @@ def patch_macos_install_names(prefix: Path) -> None:
231263
else:
232264
replacement = f"@rpath/{local_dependency.name}"
233265
subprocess.run(
234-
["install_name_tool", "-change", dependency, replacement, str(path)],
266+
[
267+
"install_name_tool",
268+
"-change",
269+
dependency,
270+
replacement,
271+
str(path),
272+
],
235273
check=True,
236274
)
237275

@@ -258,10 +296,14 @@ def main() -> int:
258296

259297
session = requests.Session()
260298
session.headers.update(github_headers())
261-
release_url = f"https://api.github.com/repos/{args.repo}/releases/tags/{NIGHTLY_TAG}"
299+
release_url = (
300+
f"https://api.github.com/repos/{args.repo}/releases/tags/{NIGHTLY_TAG}"
301+
)
262302
release = get_json(session, release_url)
263303
if release.get("tag_name") != NIGHTLY_TAG or not release.get("prerelease"):
264-
raise SystemExit(f"Release {args.repo}@{NIGHTLY_TAG} is not the nightly prerelease")
304+
raise SystemExit(
305+
f"Release {args.repo}@{NIGHTLY_TAG} is not the nightly prerelease"
306+
)
265307

266308
archive_asset = select_archive_asset(release, target_platform)
267309
archive_name = archive_asset["name"]
@@ -271,7 +313,10 @@ def main() -> int:
271313
temp_dir = Path(temp_dir_text)
272314
archive_path = temp_dir / archive_name
273315
checksum_path = temp_dir / f"{archive_name}.sha256"
274-
print(f"Downloading {archive_name} from {args.repo}@{NIGHTLY_TAG}", file=sys.stderr)
316+
print(
317+
f"Downloading {archive_name} from {args.repo}@{NIGHTLY_TAG}",
318+
file=sys.stderr,
319+
)
275320
download(session, archive_asset["browser_download_url"], archive_path)
276321
download(session, checksum_asset["browser_download_url"], checksum_path)
277322

0 commit comments

Comments
 (0)