setup_cpp.py

import os

from pybind11.setup_helpers import Pybind11Extension, build_ext
from setuptools import setup

# Check for CUDA
cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
if not cuda_home and os.path.exists("/usr/local/cuda"):
    cuda_home = "/usr/local/cuda"

include_dirs = [
    os.path.abspath("submodules/heidi-kernel/include"),
]
library_dirs = []
libraries = ["z"]  # Always link zlib
macros = []

if cuda_home and os.path.exists(os.path.join(cuda_home, "include/cuda_runtime.h")):
    include_dirs.append(os.path.join(cuda_home, "include"))
    library_dirs.append(os.path.join(cuda_home, "lib64"))
    libraries.append("cudart")
    macros.append(("HAS_CUDA", "1"))
    print(f"[INFO] CUDA detected at {cuda_home}")

libraries.append("crypto")  # Explicitly require openssl crypto for SHA256

ext_modules = [
    Pybind11Extension(
        "heidi_cpp",
        [
            "heidi_engine/cpp/heidi_cpp.cpp",
            "heidi_engine/cpp/core/clock.cpp",
            "heidi_engine/cpp/core/config.cpp",
            "heidi_engine/cpp/core/core.cpp",
            "heidi_engine/cpp/core/journal_writer.cpp",
            "heidi_engine/cpp/core/run_id.cpp",
            "heidi_engine/cpp/core/status_writer.cpp",
            "heidi_engine/cpp/core/subprocess.cpp",
            "heidi_engine/cpp/core/mock_provider.cpp",
            "heidi_engine/cpp/core/async_collector.cpp",
            "submodules/heidi-kernel/src/governor/resource_governor.cpp",
            "submodules/heidi-kernel/src/metrics/sampler.cpp",
        ],
        include_dirs=include_dirs,
        library_dirs=library_dirs,
        libraries=libraries,
        define_macros=macros,
        cxx_std=23,
    ),
]

setup(
    name="heidi_cpp",
    version="0.2.0",
    author="Heidi Team",
    description="C++ performance optimizations for Heidi Engine (V2)",
    ext_modules=ext_modules,
    cmdclass={"build_ext": build_ext},
    zip_safe=False,
    python_requires=">=3.7",
)