Skip to content

Commit 4cc748c

Browse files
viktorhtviktorht
committed
Installs, cmdstan and compile package during build. However user should still be able to use pseudobatch without cmdstanpy installed
1 parent 55a70ba commit 4cc748c

File tree

1 file changed

+156
-160
lines changed

1 file changed

+156
-160
lines changed

‎setup.py

Lines changed: 156 additions & 160 deletions
Original file line numberDiff line numberDiff line change
@@ -8,175 +8,171 @@
88
from pathlib import Path
99
from shutil import copy, copytree, rmtree
1010
from typing import Tuple
11-
from importlib.metadata import distribution
11+
import cmdstanpy
1212
from setuptools import Extension, setup
1313

1414

15-
try:
16-
distribution("cmdstanpy")
17-
import cmdstanpy
18-
19-
cmdstanpy_installed = True
20-
except:
21-
print("cmdstanpy not found, skipping model compilation.")
22-
cmdstanpy_installed = False
15+
from setuptools.command.build_ext import build_ext
16+
from distutils.command.clean import clean
17+
from wheel.bdist_wheel import bdist_wheel
18+
19+
MODEL_DIR = "pseudobatch/error_propagation/stan"
20+
21+
MODELS = [
22+
"error_propagation",
23+
]
24+
25+
CMDSTAN_VERSION = "2.33.0"
26+
BINARIES_DIR = "bin"
27+
BINARIES = ["diagnose", "print", "stanc", "stansummary"]
28+
MATH_LIB = "stan/lib/stan_math/lib"
29+
TBB_DIRS = ["tbb", "tbb_2020.3"]
30+
31+
32+
def prune_cmdstan(cmdstan_dir: os.PathLike) -> None:
33+
"""
34+
Keep only the cmdstan executables and tbb files
35+
(minimum required to run a cmdstanpy commands on a pre-compiled model).
36+
"""
37+
original_dir = Path(cmdstan_dir).resolve()
38+
parent_dir = original_dir.parent
39+
temp_dir = parent_dir / "temp"
40+
if temp_dir.is_dir():
41+
rmtree(temp_dir)
42+
temp_dir.mkdir()
43+
44+
print("Copying ", original_dir, " to ", temp_dir, " for pruning")
45+
copytree(original_dir / BINARIES_DIR, temp_dir / BINARIES_DIR)
46+
copy(original_dir / "makefile", temp_dir / "makefile")
47+
for f in (temp_dir / BINARIES_DIR).iterdir():
48+
if f.is_dir():
49+
rmtree(f)
50+
elif f.is_file() and f.stem not in BINARIES:
51+
os.remove(f)
52+
for tbb_dir in TBB_DIRS:
53+
copytree(
54+
original_dir / MATH_LIB / tbb_dir, temp_dir / MATH_LIB / tbb_dir
55+
)
56+
57+
rmtree(original_dir)
58+
temp_dir.rename(original_dir)
59+
60+
61+
def repackage_cmdstan() -> bool:
62+
return os.environ.get("PSEUDOBATCH_REPACKAGE_CMDSTAN", "").lower() in [
63+
"true",
64+
"1",
65+
]
2366

24-
if not cmdstanpy_installed:
25-
setup()
26-
else:
27-
from setuptools.command.build_ext import build_ext
28-
from distutils.command.clean import clean
29-
from wheel.bdist_wheel import bdist_wheel
3067

31-
MODEL_DIR = "pseudobatch/error_propagation/stan"
68+
def maybe_install_cmdstan_toolchain() -> None:
69+
"""Install C++ compilers required to build stan models on Windows machines."""
3270

33-
MODELS = [
34-
"error_propagation",
35-
]
71+
try:
72+
cmdstanpy.utils.cxx_toolchain_path()
73+
except Exception:
74+
from cmdstanpy.install_cxx_toolchain import run_rtools_install
3675

37-
CMDSTAN_VERSION = "2.31.0"
38-
BINARIES_DIR = "bin"
39-
BINARIES = ["diagnose", "print", "stanc", "stansummary"]
40-
MATH_LIB = "stan/lib/stan_math/lib"
41-
TBB_DIRS = ["tbb", "tbb_2020.3"]
42-
43-
def prune_cmdstan(cmdstan_dir: os.PathLike) -> None:
44-
"""
45-
Keep only the cmdstan executables and tbb files
46-
(minimum required to run a cmdstanpy commands on a pre-compiled model).
47-
"""
48-
original_dir = Path(cmdstan_dir).resolve()
49-
parent_dir = original_dir.parent
50-
temp_dir = parent_dir / "temp"
51-
if temp_dir.is_dir():
52-
rmtree(temp_dir)
53-
temp_dir.mkdir()
54-
55-
print("Copying ", original_dir, " to ", temp_dir, " for pruning")
56-
copytree(original_dir / BINARIES_DIR, temp_dir / BINARIES_DIR)
57-
copy(original_dir / "makefile", temp_dir / "makefile")
58-
for f in (temp_dir / BINARIES_DIR).iterdir():
59-
if f.is_dir():
60-
rmtree(f)
61-
elif f.is_file() and f.stem not in BINARIES:
62-
os.remove(f)
63-
for tbb_dir in TBB_DIRS:
64-
copytree(
65-
original_dir / MATH_LIB / tbb_dir, temp_dir / MATH_LIB / tbb_dir
66-
)
76+
run_rtools_install({"version": None, "dir": None, "verbose": True})
77+
cmdstanpy.utils.cxx_toolchain_path()
6778

68-
rmtree(original_dir)
69-
temp_dir.rename(original_dir)
7079

71-
def repackage_cmdstan() -> bool:
72-
return os.environ.get("PSEUDOBATCH_REPACKAGE_CMDSTAN", "").lower() in [
73-
"true",
74-
"1",
75-
]
80+
def install_cmdstan_deps(cmdstan_dir: Path) -> None:
81+
from multiprocessing import cpu_count
7682

77-
def maybe_install_cmdstan_toolchain() -> None:
78-
"""Install C++ compilers required to build stan models on Windows machines."""
83+
if repackage_cmdstan():
84+
if platform.platform().startswith("Win"):
85+
maybe_install_cmdstan_toolchain()
86+
print("Installing cmdstan to", cmdstan_dir)
87+
if os.path.isdir(cmdstan_dir):
88+
print("Removing existing dir", cmdstan_dir)
89+
rmtree(cmdstan_dir)
7990

80-
try:
81-
cmdstanpy.utils.cxx_toolchain_path()
82-
except Exception:
83-
from cmdstanpy.install_cxx_toolchain import run_rtools_install
84-
85-
run_rtools_install({"version": None, "dir": None, "verbose": True})
86-
cmdstanpy.utils.cxx_toolchain_path()
87-
88-
def install_cmdstan_deps(cmdstan_dir: Path) -> None:
89-
from multiprocessing import cpu_count
90-
91-
if repackage_cmdstan():
92-
if platform.platform().startswith("Win"):
93-
maybe_install_cmdstan_toolchain()
94-
print("Installing cmdstan to", cmdstan_dir)
95-
if os.path.isdir(cmdstan_dir):
96-
print("Removing existing dir", cmdstan_dir)
97-
rmtree(cmdstan_dir)
98-
99-
if not cmdstanpy.install_cmdstan(
100-
version=CMDSTAN_VERSION,
101-
dir=cmdstan_dir.parent,
102-
overwrite=True,
103-
verbose=True,
104-
cores=cpu_count(),
105-
):
106-
raise RuntimeError(
107-
"CmdStan failed to install in repackaged directory"
108-
)
109-
else:
110-
try:
111-
cmdstanpy.cmdstan_path()
112-
except ValueError as e:
113-
raise SystemExit(
114-
"CmdStan not installed, but the package is building from source"
115-
) from e
116-
117-
def build_models(target_dir: str) -> None:
118-
cmdstan_dir = (
119-
Path(target_dir) / f"cmdstan-{CMDSTAN_VERSION}"
120-
).resolve()
121-
install_cmdstan_deps(cmdstan_dir)
122-
for model in MODELS:
123-
sm = cmdstanpy.CmdStanModel(
124-
stan_file=os.path.join(MODEL_DIR, model + ".stan"),
125-
stanc_options={"O1": True},
91+
if not cmdstanpy.install_cmdstan(
92+
version=CMDSTAN_VERSION,
93+
dir=cmdstan_dir.parent,
94+
overwrite=True,
95+
verbose=True,
96+
cores=cpu_count(),
97+
):
98+
raise RuntimeError(
99+
"CmdStan failed to install in repackaged directory"
126100
)
127-
copy(sm.exe_file, os.path.join(target_dir, model + ".exe"))
128-
129-
if repackage_cmdstan():
130-
prune_cmdstan(cmdstan_dir)
131-
132-
class BuildModels(build_ext):
133-
"""Custom build command to pre-compile Stan models."""
134-
135-
def run(self) -> None:
136-
if not self.dry_run:
137-
target_dir = os.path.join(self.build_lib, MODEL_DIR)
138-
self.mkpath(target_dir)
139-
build_models(target_dir)
140-
# don't call build_ext.run, since we're not really building c files
141-
142-
def clean_models(target_dir: str) -> None:
143-
# Remove compiled stan files
144-
for model in MODELS:
145-
for filename in [model, f"{model}.hpp", f"{model}.exe"]:
146-
stan_file = Path(target_dir) / filename
147-
if stan_file.exists():
148-
stan_file.unlink()
149-
150-
class CleanModels(clean):
151-
"""Custom clean command to remove pre-compile Stan models."""
152-
153-
def run(self) -> None:
154-
if not self.dry_run:
155-
target_dir = os.path.join(self.build_lib, MODEL_DIR)
156-
clean_models(target_dir)
157-
clean_models(MODEL_DIR)
158-
super().run()
159-
160-
# this is taken from the cibuildwheel example https://github.com/joerick/python-ctypes-package-sample
161-
# it marks the wheel as not specific to the Python API version.
162-
# This means the wheel will only be built once per platform, rather than per-Python-per-platform.
163-
# If you are combining with any actual C extensions, you will most likely want to remove this.
164-
class WheelABINone(bdist_wheel):
165-
def finalize_options(self) -> None:
166-
bdist_wheel.finalize_options(self)
167-
self.root_is_pure = False
168-
169-
def get_tag(self) -> Tuple[str, str, str]:
170-
_, _, plat = bdist_wheel.get_tag(self)
171-
return "py3", "none", plat
172-
173-
setup(
174-
# Extension marks this as platform-specific
175-
ext_modules=[Extension("pseudobatch.error_propagation.stan", [])],
176-
# override the build and bdist commands
177-
cmdclass={
178-
"build_ext": BuildModels,
179-
"bdist_wheel": WheelABINone,
180-
"clean": CleanModels,
181-
},
182-
)
101+
else:
102+
try:
103+
cmdstanpy.cmdstan_path()
104+
except ValueError as e:
105+
raise SystemExit(
106+
"CmdStan not installed, but the package is building from source"
107+
) from e
108+
109+
110+
def build_models(target_dir: str) -> None:
111+
cmdstan_dir = (Path(target_dir) / f"cmdstan-{CMDSTAN_VERSION}").resolve()
112+
install_cmdstan_deps(cmdstan_dir)
113+
for model in MODELS:
114+
sm = cmdstanpy.CmdStanModel(
115+
stan_file=os.path.join(MODEL_DIR, model + ".stan"),
116+
stanc_options={"O1": True},
117+
)
118+
copy(sm.exe_file, os.path.join(target_dir, model + ".exe"))
119+
120+
if repackage_cmdstan():
121+
prune_cmdstan(cmdstan_dir)
122+
123+
124+
class BuildModels(build_ext):
125+
"""Custom build command to pre-compile Stan models."""
126+
127+
def run(self) -> None:
128+
if not self.dry_run:
129+
target_dir = os.path.join(self.build_lib, MODEL_DIR)
130+
self.mkpath(target_dir)
131+
build_models(target_dir)
132+
# don't call build_ext.run, since we're not really building c files
133+
134+
135+
def clean_models(target_dir: str) -> None:
136+
# Remove compiled stan files
137+
for model in MODELS:
138+
for filename in [model, f"{model}.hpp", f"{model}.exe"]:
139+
stan_file = Path(target_dir) / filename
140+
if stan_file.exists():
141+
stan_file.unlink()
142+
143+
144+
class CleanModels(clean):
145+
"""Custom clean command to remove pre-compile Stan models."""
146+
147+
def run(self) -> None:
148+
if not self.dry_run:
149+
target_dir = os.path.join(self.build_lib, MODEL_DIR)
150+
clean_models(target_dir)
151+
clean_models(MODEL_DIR)
152+
super().run()
153+
154+
155+
# this is taken from the cibuildwheel example https://github.com/joerick/python-ctypes-package-sample
156+
# it marks the wheel as not specific to the Python API version.
157+
# This means the wheel will only be built once per platform, rather than per-Python-per-platform.
158+
# If you are combining with any actual C extensions, you will most likely want to remove this.
159+
class WheelABINone(bdist_wheel):
160+
def finalize_options(self) -> None:
161+
bdist_wheel.finalize_options(self)
162+
self.root_is_pure = False
163+
164+
def get_tag(self) -> Tuple[str, str, str]:
165+
_, _, plat = bdist_wheel.get_tag(self)
166+
return "py3", "none", plat
167+
168+
169+
setup(
170+
# Extension marks this as platform-specific
171+
ext_modules=[Extension("pseudobatch.error_propagation.stan", [])],
172+
# override the build and bdist commands
173+
cmdclass={
174+
"build_ext": BuildModels,
175+
"bdist_wheel": WheelABINone,
176+
"clean": CleanModels,
177+
},
178+
)

0 commit comments

Comments
 (0)