mirror of
https://github.com/saymrwulf/prophet.git
synced 2026-05-15 20:50:48 +00:00
269 lines
8.4 KiB
Python
269 lines
8.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import os
|
|
import sys
|
|
import platform
|
|
from pathlib import Path
|
|
from shutil import copy, copytree, rmtree
|
|
from typing import List
|
|
|
|
from pkg_resources import add_activation_listener, normalize_path, require, working_set
|
|
from setuptools import find_packages, setup, Extension
|
|
from setuptools.command.build_ext import build_ext
|
|
from setuptools.command.build_py import build_py
|
|
from setuptools.command.develop import develop
|
|
from setuptools.command.test import test as test_command
|
|
|
|
|
|
MODEL_DIR = "stan"
|
|
MODEL_TARGET_DIR = os.path.join("prophet", "stan_model")
|
|
|
|
CMDSTAN_VERSION = "2.26.1"
|
|
BINARIES_DIR = "bin"
|
|
BINARIES = ["diagnose", "print", "stanc", "stansummary"]
|
|
TBB_PARENT = "stan/lib/stan_math/lib"
|
|
TBB_DIRS = ["tbb", "tbb_2019_U8"]
|
|
|
|
|
|
def prune_cmdstan(cmdstan_dir: str) -> None:
|
|
"""
|
|
Keep only the cmdstan executables and tbb files (minimum required to run a cmdstanpy commands on a pre-compiled model).
|
|
"""
|
|
original_dir = Path(cmdstan_dir).resolve()
|
|
parent_dir = original_dir.parent
|
|
temp_dir = parent_dir / "temp"
|
|
if temp_dir.is_dir():
|
|
rmtree(temp_dir)
|
|
temp_dir.mkdir()
|
|
|
|
print("Copying ", original_dir, " to ", temp_dir, " for pruning")
|
|
copytree(original_dir / BINARIES_DIR, temp_dir / BINARIES_DIR)
|
|
for f in (temp_dir / BINARIES_DIR).iterdir():
|
|
if f.is_dir():
|
|
rmtree(f)
|
|
elif f.is_file() and f.stem not in BINARIES:
|
|
os.remove(f)
|
|
for tbb_dir in TBB_DIRS:
|
|
copytree(original_dir / TBB_PARENT / tbb_dir, temp_dir / TBB_PARENT / tbb_dir)
|
|
|
|
rmtree(original_dir)
|
|
temp_dir.rename(original_dir)
|
|
|
|
def repackage_cmdstan():
|
|
return os.environ.get("PROPHET_REPACKAGE_CMDSTAN", "").lower() not in ["false", "0"]
|
|
|
|
|
|
def maybe_install_cmdstan_toolchain():
|
|
"""Install C++ compilers required to build stan models on Windows machines."""
|
|
import cmdstanpy
|
|
try:
|
|
cmdstanpy.utils.cxx_toolchain_path()
|
|
except Exception:
|
|
try:
|
|
from cmdstanpy.install_cxx_toolchain import run_rtools_install
|
|
except ImportError:
|
|
# older versions
|
|
from cmdstanpy.install_cxx_toolchain import main as run_rtools_install
|
|
|
|
run_rtools_install({"version": None, "dir": None, "verbose": True})
|
|
cmdstanpy.utils.cxx_toolchain_path()
|
|
|
|
|
|
def install_cmdstan_deps(cmdstan_dir: Path):
|
|
import cmdstanpy
|
|
from multiprocessing import cpu_count
|
|
|
|
if repackage_cmdstan():
|
|
if platform.platform().startswith("Win"):
|
|
maybe_install_cmdstan_toolchain()
|
|
print("Installing cmdstan to", cmdstan_dir)
|
|
if os.path.isdir(cmdstan_dir):
|
|
rmtree(cmdstan_dir)
|
|
|
|
if not cmdstanpy.install_cmdstan(
|
|
version=CMDSTAN_VERSION,
|
|
dir=cmdstan_dir.parent,
|
|
overwrite=True,
|
|
verbose=True,
|
|
cores=cpu_count(),
|
|
progress=True,
|
|
):
|
|
raise RuntimeError("CmdStan failed to install in repackaged directory")
|
|
|
|
|
|
def build_cmdstan_model(target_dir):
|
|
"""
|
|
Rebuild cmdstan in the build environment, then use this installation to compile the stan model.
|
|
The stan model is copied to {target_dir}/prophet_model.bin
|
|
The cmdstan files required to run cmdstanpy commands are copied to {target_dir}/cmdstan-{version}.
|
|
|
|
Parameters
|
|
----------
|
|
target_dir: Directory to copy the compiled model executable and core cmdstan files to.
|
|
"""
|
|
import cmdstanpy
|
|
|
|
cmdstan_dir = (Path(target_dir) / f"cmdstan-{CMDSTAN_VERSION}").resolve()
|
|
install_cmdstan_deps(cmdstan_dir)
|
|
model_name = "prophet.stan"
|
|
target_name = "prophet_model.bin"
|
|
sm = cmdstanpy.CmdStanModel(stan_file=os.path.join(MODEL_DIR, model_name))
|
|
copy(sm.exe_file, os.path.join(target_dir, target_name))
|
|
|
|
# Clean up
|
|
for f in Path(MODEL_DIR).iterdir():
|
|
if f.is_file() and f.name != model_name:
|
|
os.remove(f)
|
|
|
|
if repackage_cmdstan():
|
|
prune_cmdstan(cmdstan_dir)
|
|
|
|
|
|
|
|
def get_backends_from_env() -> List[str]:
|
|
return os.environ.get("STAN_BACKEND", "CMDSTANPY").split(",")
|
|
|
|
|
|
def build_models(target_dir):
|
|
print("Compiling cmdstanpy model")
|
|
build_cmdstan_model(target_dir)
|
|
|
|
if 'PYSTAN' in get_backends_from_env():
|
|
raise ValueError("PyStan backend is not supported for Prophet >= 1.1")
|
|
|
|
class BuildPyCommand(build_py):
|
|
"""Custom build command to pre-compile Stan models."""
|
|
|
|
def run(self):
|
|
if not self.dry_run:
|
|
target_dir = os.path.join(self.build_lib, MODEL_TARGET_DIR)
|
|
self.mkpath(target_dir)
|
|
build_models(target_dir)
|
|
|
|
build_py.run(self)
|
|
|
|
|
|
class BuildExtCommand(build_ext):
|
|
"""Ensure built extensions are added to the correct path in the wheel."""
|
|
|
|
def run(self):
|
|
pass
|
|
|
|
|
|
class DevelopCommand(develop):
|
|
"""Custom develop command to pre-compile Stan models in-place."""
|
|
|
|
def run(self):
|
|
if not self.dry_run:
|
|
target_dir = os.path.join(self.setup_path, MODEL_TARGET_DIR)
|
|
self.mkpath(target_dir)
|
|
build_models(target_dir)
|
|
|
|
develop.run(self)
|
|
|
|
|
|
class TestCommand(test_command):
|
|
user_options = [
|
|
("test-module=", "m", "Run 'test_suite' in specified module"),
|
|
(
|
|
"test-suite=",
|
|
"s",
|
|
"Run single test, case or suite (e.g. 'module.test_suite')",
|
|
),
|
|
("test-runner=", "r", "Test runner to use"),
|
|
("test-slow", "w", "Test slow suites (default off)"),
|
|
]
|
|
test_slow = None
|
|
|
|
def initialize_options(self):
|
|
super(TestCommand, self).initialize_options()
|
|
self.test_slow = False
|
|
|
|
def finalize_options(self):
|
|
super(TestCommand, self).finalize_options()
|
|
if self.test_slow is None:
|
|
self.test_slow = getattr(self.distribution, "test_slow", False)
|
|
|
|
"""We must run tests on the build directory, not source."""
|
|
|
|
def with_project_on_sys_path(self, func):
|
|
# Ensure metadata is up-to-date
|
|
self.reinitialize_command("build_py", inplace=0)
|
|
self.run_command("build_py")
|
|
bpy_cmd = self.get_finalized_command("build_py")
|
|
build_path = normalize_path(bpy_cmd.build_lib)
|
|
|
|
# Build extensions
|
|
self.reinitialize_command("egg_info", egg_base=build_path)
|
|
self.run_command("egg_info")
|
|
|
|
self.reinitialize_command("build_ext", inplace=0)
|
|
self.run_command("build_ext")
|
|
|
|
ei_cmd = self.get_finalized_command("egg_info")
|
|
|
|
old_path = sys.path[:]
|
|
old_modules = sys.modules.copy()
|
|
|
|
try:
|
|
sys.path.insert(0, normalize_path(ei_cmd.egg_base))
|
|
working_set.__init__()
|
|
add_activation_listener(lambda dist: dist.activate())
|
|
require("%s==%s" % (ei_cmd.egg_name, ei_cmd.egg_version))
|
|
func()
|
|
finally:
|
|
sys.path[:] = old_path
|
|
sys.modules.clear()
|
|
sys.modules.update(old_modules)
|
|
working_set.__init__()
|
|
|
|
|
|
with open("README.md", "r", encoding="utf-8") as f:
|
|
long_description = f.read()
|
|
|
|
with open("requirements.txt", "r") as f:
|
|
install_requires = f.read().splitlines()
|
|
|
|
about = {}
|
|
here = Path(__file__).parent.resolve()
|
|
with open(here / "prophet" / "__version__.py", "r") as f:
|
|
exec(f.read(), about)
|
|
|
|
setup(
|
|
name=about["__title__"],
|
|
version=about["__version__"],
|
|
description=about["__description__"],
|
|
url=about["__url__"],
|
|
project_urls={
|
|
"Source": "https://github.com/facebook/prophet",
|
|
},
|
|
author=about["__author__"],
|
|
author_email=about["__author_email__"],
|
|
license=about["__license__"],
|
|
packages=find_packages(),
|
|
install_requires=install_requires,
|
|
python_requires=">=3.7",
|
|
zip_safe=False,
|
|
include_package_data=True,
|
|
ext_modules=[Extension("prophet.stan_model", [])],
|
|
cmdclass={
|
|
"build_ext": BuildExtCommand,
|
|
"build_py": BuildPyCommand,
|
|
"develop": DevelopCommand,
|
|
"test": TestCommand,
|
|
},
|
|
test_suite="prophet.tests",
|
|
classifiers=[
|
|
"Programming Language :: Python",
|
|
"Programming Language :: Python :: 3",
|
|
"Programming Language :: Python :: 3.7",
|
|
"Programming Language :: Python :: 3.8",
|
|
"Programming Language :: Python :: 3.9",
|
|
"Programming Language :: Python :: 3.10",
|
|
],
|
|
long_description=long_description,
|
|
long_description_content_type="text/markdown",
|
|
)
|