BudgetOptimizer self._compile_objective_and_grad() recompiles using different CPU architectures

[AlfredoJF]

I'm wondering if there's a way to create a "generic CPU" compilation for the BudgetOptimizer in PyTensor to compile the objective and grad functions.

pymc-marketing/pymc_marketing/mmm/budget_optimizer.py

Lines 209 to 211 in a87f1fa

	# 6. Compile objective & gradient
	self._compiled_functions = {}
	self._compile_objective_and_grad()

I trained my model on a MacBook Pro M1 and created a Docker container with a precompilation snippet to reduce cold starts at runtime. Locally, this precompilation reduces the run time from 3 minutes to about 15 seconds, which works perfectly for my use case.

However, when deploying to the cloud using Google Cloud Build and serving via Google Cloud Run, it seems the CPU architecture mismatch between the build environment (where the precomplie snippet runs) and the Cloud Run instance forces PyTensor to recompile the model at runtime. This results in a 3-minute delay before the first optimized response is returned from the self.optimize_budget function, and unfortunately, this recompile is not persistent.

Here is how my Dockerfile looks like:

# Set PyTensor flags to use g++ and BLAS, with a persistent compiledir
ENV PYTENSOR_FLAGS="blas__ldflags=-lopenblas,cxx=/usr/bin/g++,compiledir=/app/pytensor_cache"

# Precompile model during the Docker build
RUN PYTHONPATH="/app" python -c "\
from pymc_marketing.mmm import MMM; \
mmm = MMM.load('/app/my_model.nc'); \
_ , _ = mmm.optimize_budget(budget=1000, num_periods=1); \
"

Hope that anyone with more experience with PyTensor than me has a workaround.

cc: @cetagostini @ricardoV94

[AlfredoJF]

Seems like adding this to the PYTENSOR_FLAGS environment variable has significantly improved the process! The precompilation now runs flawlessly in both my local setup and the production container, completing in just a few seconds.

gcc__cxxflags=-march=x86-64 -mtune=generic