Add a num_threads helper argument to pathfinder()

Author: WardBrian

cmdstanpy/model.py

@@ -1635,6 +1635,7 @@ def pathfinder(
         refresh: Optional[int] = None,
         time_fmt: str = "%Y%m%d%H%M%S",
         timeout: Optional[float] = None,
+        num_threads: Optional[int] = None,
     ) -> CmdStanPathfinder:
         """
         Run CmdStan's Pathfinder variational inference algorithm.
@@ -1737,6 +1738,10 @@ def pathfinder(
         :param timeout: Duration at which Pathfinder times
             out in seconds. Defaults to None.
 
+        :param num_threads: Number of threads to request for parallel execution.
+            A number other than ``1`` requires the model to have been compiled
+            with STAN_THREADS=True.
+
         :return: A :class:`CmdStanPathfinder` object
 
         References
@@ -1763,6 +1768,17 @@ def pathfinder(
                 "available for CmdStan versions 2.34 and later"
             )
 
+        if num_threads is not None:
+            if (
+                num_threads != 1
+                and exe_info.get('STAN_THREADS', '').lower() != 'true'
+            ):
+                raise ValueError(
+                    "Model must be compiled with 'STAN_THREADS=true' to use"
+                    " 'num_threads' argument"
+                )
+            os.environ['STAN_NUM_THREADS'] = str(num_threads)
+
         if num_paths == 1:
             if num_single_draws is None:
                 num_single_draws = draws

test/test_pathfinder.py

@@ -152,3 +152,20 @@ def test_pathfinder_no_lp_calc():
     n_lp_nan = np.sum(np.isnan(pathfinder.method_variables()['lp__']))
     assert n_lp_nan < 4000  # some lp still calculated during pathfinder
     assert n_lp_nan > 3000  # but most are not
+
+
+def test_pathfinder_threads():
+    stan = DATAFILES_PATH / 'bernoulli.stan'
+    bern_model = cmdstanpy.CmdStanModel(stan_file=stan)
+    jdata = str(DATAFILES_PATH / 'bernoulli.data.json')
+
+    bern_model.pathfinder(data=jdata, num_threads=1)
+
+    with pytest.raises(ValueError, match="STAN_THREADS"):
+        bern_model.pathfinder(data=jdata, num_threads=4)
+
+    bern_model = cmdstanpy.CmdStanModel(
+        stan_file=stan, cpp_options={'STAN_THREADS': True}, force_compile=True
+    )
+    pathfinder = bern_model.pathfinder(data=jdata, num_threads=4)
+    assert pathfinder.draws().shape == (1000, 3)