fix: tune.run() from concurrent threads corrupts shared runner state

flaml/tune/tune.py

@@ -5,6 +5,7 @@
 import datetime
 import os
 import sys
+import threading
 import time
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple, Union
@@ -41,11 +42,27 @@
     internal_mlflow = False
 
 
-_use_ray = True
-_runner = None
-_verbose = 0
-_running_trial = None
-_training_iteration = 0
+class _TuneState(threading.local):
+    """Per-thread run() state.
+
+    report()/run() coordinate purely through this state (use_ray, runner,
+    verbose, running_trial, training_iteration). A plain module global here
+    is shared by every thread, so two threads calling tune.run() concurrently
+    overwrite each other's runner/trial bookkeeping mid-flight (#996).
+    threading.local's __init__ re-runs on each thread's first access, so
+    every thread starts from these same defaults without an explicit
+    per-thread init call.
+    """
+
+    def __init__(self):
+        self.use_ray = True
+        self.runner = None
+        self.verbose = 0
+        self.running_trial = None
+        self.training_iteration = 0
+
+
+_state = _TuneState()
 
 INCUMBENT_RESULT = "__incumbent_result__"
 
@@ -189,11 +206,7 @@ def compute_with_config(config):
         SystemExit (when using ray):
             A SystemExit exception is raised if the trial has been signaled to stop by ray.
     """
-    global _use_ray
-    global _verbose
-    global _running_trial
-    global _training_iteration
-    if _use_ray:
+    if _state.use_ray:
         try:
             from ray import __version__ as ray_version
 
@@ -211,22 +224,22 @@ def compute_with_config(config):
     result = kwargs
     if _metric is not None:
         result[DEFAULT_METRIC] = _metric
-    trial = getattr(_runner, "running_trial", None)
+    trial = getattr(_state.runner, "running_trial", None)
     if not trial:
         return None
-    if _running_trial == trial:
-        _training_iteration += 1
+    if _state.running_trial == trial:
+        _state.training_iteration += 1
     else:
-        _training_iteration = 0
-        _running_trial = trial
-    result["training_iteration"] = _training_iteration
+        _state.training_iteration = 0
+        _state.running_trial = trial
+    result["training_iteration"] = _state.training_iteration
     result["config"] = trial.config
     if INCUMBENT_RESULT in result["config"]:
         del result["config"][INCUMBENT_RESULT]
     for key, value in trial.config.items():
         result["config/" + key] = value
-    _runner.process_trial_result(trial, result)
-    if _verbose > 2:
+    _state.runner.process_trial_result(trial, result)
+    if _state.verbose > 2:
         logger.info(f"result: {result}")
     if trial.is_finished():
         raise StopIteration
@@ -478,15 +491,11 @@ def easy_objective(config):
         **ray_args: keyword arguments to pass to ray.tune.run().
             Only valid when use_ray=True.
     """
-    global _use_ray
-    global _verbose
-    global _running_trial
-    global _training_iteration
     global internal_mlflow
-    old_use_ray = _use_ray
-    old_verbose = _verbose
-    old_running_trial = _running_trial
-    old_training_iteration = _training_iteration
+    old_use_ray = _state.use_ray
+    old_verbose = _state.verbose
+    old_running_trial = _state.running_trial
+    old_training_iteration = _state.training_iteration
 
     if log_file_name:
         dir_name = os.path.dirname(log_file_name)
@@ -498,13 +507,12 @@ def easy_objective(config):
     if use_ray and use_spark:
         raise ValueError("use_ray and use_spark cannot be both True.")
     if not use_ray:
-        _use_ray = False
-        _verbose = verbose
+        _state.use_ray = False
+        _state.verbose = verbose
         old_handlers = logger.handlers
         old_level = logger.getEffectiveLevel()
         logger.handlers = []
-        global _runner
-        old_runner = _runner
+        old_runner = _state.runner
         assert not ray_args, "ray_args is only valid when use_ray=True"
         if (
             old_handlers
@@ -674,7 +682,7 @@ def easy_objective(config):
             from ray import tune
         except ImportError:
             raise ImportError("Failed to import ray tune. " "Please install ray[tune] or set use_ray=False")
-        _use_ray = True
+        _state.use_ray = True
         try:
             analysis = tune.run(
                 evaluation_function,
@@ -695,10 +703,10 @@ def easy_objective(config):
                         f.write(f"result: {trial.last_result}\n")
             return analysis
         finally:
-            _use_ray = old_use_ray
-            _verbose = old_verbose
-            _running_trial = old_running_trial
-            _training_iteration = old_training_iteration
+            _state.use_ray = old_use_ray
+            _state.verbose = old_verbose
+            _state.running_trial = old_running_trial
+            _state.training_iteration = old_training_iteration
 
     if use_spark:
         # parallel run with spark
@@ -759,7 +767,7 @@ def easy_objective(config):
         with parallel_backend("spark"):
             with Parallel(n_jobs=n_concurrent_trials, verbose=max(0, (verbose - 1) * 50)) as parallel:
                 try:
-                    _runner = SparkTrialRunner(
+                    _state.runner = SparkTrialRunner(
                         search_alg=search_alg,
                         scheduler=scheduler,
                         metric=metric,
@@ -779,23 +787,23 @@ def easy_objective(config):
                         if automl_info and automl_info[1] == "all" and automl_info[0] > 0 and time_budget_s < np.inf:
                             time_budget_s -= automl_info[0] * n_concurrent_trials
                             logger.debug(f"Remaining time budget with mlflow log latency: {time_budget_s} seconds.")
-                        while len(_runner.running_trials) < n_concurrent_trials:
+                        while len(_state.runner.running_trials) < n_concurrent_trials:
                             # suggest trials for spark
-                            trial_next = _runner.step()
+                            trial_next = _state.runner.step()
                             if trial_next:
                                 num_trials += 1
                             else:
                                 num_failures += 1  # break with upperbound_num_failures consecutive failures
                                 logger.debug(f"consecutive failures is {num_failures}")
                                 if num_failures >= upperbound_num_failures:
                                     break
-                        trials_to_run = _runner.running_trials
+                        trials_to_run = _state.runner.running_trials
                         if not trials_to_run:
                             logger.warning(f"fail to sample a trial for {max_failure} times in a row, stopping.")
                             break
                         logger.info(
-                            f"Number of trials: {num_trials}/{num_samples}, {len(_runner.running_trials)} RUNNING,"
-                            f" {len(_runner._trials) - len(_runner.running_trials)} TERMINATED"
+                            f"Number of trials: {num_trials}/{num_samples}, {len(_state.runner.running_trials)} RUNNING,"
+                            f" {len(_state.runner._trials) - len(_state.runner.running_trials)} TERMINATED"
                         )
                         logger.debug(
                             f"Configs of Trials to run: {[trial_to_run.config for trial_to_run in trials_to_run]}"
@@ -817,7 +825,7 @@ def easy_objective(config):
                         while results:
                             result = results.pop(0)
                             trial_to_run = trials_to_run[0]
-                            _runner.running_trial = trial_to_run
+                            _state.runner.running_trial = trial_to_run
                             if result is not None:
                                 if _internal_mlflow:
                                     mlflow_integration.record_trial(result, trial_to_run, metric)
@@ -832,10 +840,10 @@ def easy_objective(config):
                                 else:
                                     logger.info("Brief result: {metric: result}")
                                     report(_metric=result)
-                            _runner.stop_trial(trial_to_run)
+                            _state.runner.stop_trial(trial_to_run)
                         num_failures = 0
                     analysis = ExperimentAnalysis(
-                        _runner.get_trials(),
+                        _state.runner.get_trials(),
                         metric=metric,
                         mode=mode,
                         lexico_objectives=lexico_objectives,
@@ -857,26 +865,26 @@ def easy_objective(config):
                     return analysis
                 finally:
                     # recover the global variables in case of nested run
-                    _use_ray = old_use_ray
-                    _verbose = old_verbose
-                    _running_trial = old_running_trial
-                    _training_iteration = old_training_iteration
+                    _state.use_ray = old_use_ray
+                    _state.verbose = old_verbose
+                    _state.running_trial = old_running_trial
+                    _state.training_iteration = old_training_iteration
                     if not use_ray:
-                        _runner = old_runner
+                        _state.runner = old_runner
                         logger.handlers = old_handlers
                         logger.setLevel(old_level)
                     if _internal_mlflow:
                         mlflow_integration.adopt_children()
 
     # simple sequential run without using tune.run() from ray
     time_start = time.time()
-    _use_ray = False
+    _state.use_ray = False
     if scheduler:
         scheduler.set_search_properties(metric=metric, mode=mode)
     from .trial_runner import SequentialTrialRunner
 
     try:
-        _runner = SequentialTrialRunner(
+        _state.runner = SequentialTrialRunner(
             search_alg=search_alg,
             scheduler=scheduler,
             metric=metric,
@@ -892,7 +900,7 @@ def easy_objective(config):
             and (num_samples < 0 or num_trials < num_samples)
             and num_failures < upperbound_num_failures
         ):
-            trial_to_run = _runner.step()
+            trial_to_run = _state.runner.step()
             if trial_to_run:
                 num_trials += 1
                 if verbose:
@@ -913,7 +921,7 @@ def easy_objective(config):
                             trial_to_run.set_status(Trial.ERROR)
                     else:
                         report(_metric=result)
-                _runner.stop_trial(trial_to_run)
+                _state.runner.stop_trial(trial_to_run)
                 num_failures = 0
                 if trial_to_run.last_result is None:
                     # application stops tuning by returning None
@@ -925,7 +933,7 @@ def easy_objective(config):
         if num_failures == upperbound_num_failures:
             logger.warning(f"fail to sample a trial for {max_failure} times in a row, stopping.")
         analysis = ExperimentAnalysis(
-            _runner.get_trials(),
+            _state.runner.get_trials(),
             metric=metric,
             mode=mode,
             lexico_objectives=lexico_objectives,
@@ -946,12 +954,12 @@ def easy_objective(config):
         return analysis
     finally:
         # recover the global variables in case of nested run
-        _use_ray = old_use_ray
-        _verbose = old_verbose
-        _running_trial = old_running_trial
-        _training_iteration = old_training_iteration
+        _state.use_ray = old_use_ray
+        _state.verbose = old_verbose
+        _state.running_trial = old_running_trial
+        _state.training_iteration = old_training_iteration
         if not use_ray:
-            _runner = old_runner
+            _state.runner = old_runner
             logger.handlers = old_handlers
             logger.setLevel(old_level)
         if _internal_mlflow:

test/tune/test_concurrent_run.py

@@ -0,0 +1,102 @@
+"""Regression test for #996: two threads calling flaml.tune.run() concurrently
+must not corrupt each other's trial state.
+
+flaml.tune.run()/tune.report() used to coordinate through module-level globals
+(_runner/_running_trial/_training_iteration/_use_ray/_verbose in flaml/tune/tune.py).
+Those globals are shared by every thread, so whichever thread's run() call
+executes _runner = SequentialTrialRunner(...) most recently "wins" the global
+for every thread's subsequent report()/stop_trial() calls, including resetting
+it to None (or another thread's runner) out from under a still-running call.
+
+This forces the exact interleaving that corrupts state, deterministically (no
+sleep-based timing): thread A steps its own trial and pauses inside its
+evaluation function; thread B starts while A is paused, creates its own runner
+(overwriting the shared state A is still relying on), steps its own trial, and
+also pauses. A is released and runs to completion, all the way through
+report(), stop_trial(), building its analysis, and its own finally-restore,
+entirely while B is still paused. B is then released: on main, B's report()
+silently drops its result
+(the runner it reads has already been reset to whatever A's finally restored),
+and B's stop_trial() call raises AttributeError on that stale/None runner.
+"""
+import threading
+
+from flaml import tune
+
+
+def test_concurrent_tune_run_does_not_corrupt_state():
+    a_paused = threading.Event()
+    b_paused = threading.Event()
+    release_a = threading.Event()
+    release_b = threading.Event()
+    a_done = threading.Event()
+    results = {}
+
+    def eval_a(config):
+        a_paused.set()
+        assert release_a.wait(timeout=5), "thread A was never released"
+        return {"metric": 1.0}
+
+    def eval_b(config):
+        b_paused.set()
+        assert release_b.wait(timeout=5), "thread B was never released"
+        return {"metric": 2.0}
+
+    def run_a():
+        try:
+            analysis = tune.run(
+                eval_a,
+                config={"x": tune.uniform(0, 1)},
+                metric="metric",
+                mode="min",
+                num_samples=1,
+                verbose=0,
+            )
+            results["A"] = ("ok", [t.last_result for t in analysis.trials])
+        except Exception as e:  # noqa: BLE001 (captured for the test assertion below)
+            results["A"] = ("error", f"{type(e).__name__}: {e}")
+        finally:
+            a_done.set()
+
+    def run_b():
+        try:
+            analysis = tune.run(
+                eval_b,
+                config={"x": tune.uniform(0, 1)},
+                metric="metric",
+                mode="min",
+                num_samples=1,
+                verbose=0,
+            )
+            results["B"] = ("ok", [t.last_result for t in analysis.trials])
+        except Exception as e:  # noqa: BLE001 (captured for the test assertion below)
+            results["B"] = ("error", f"{type(e).__name__}: {e}")
+
+    thread_a = threading.Thread(target=run_a)
+    thread_b = threading.Thread(target=run_b)
+
+    thread_a.start()
+    assert a_paused.wait(timeout=5), "thread A never reached its evaluation function"
+
+    thread_b.start()
+    assert b_paused.wait(timeout=5), "thread B never reached its evaluation function"
+
+    # Let A run to full completion, including its own finally-restore, while
+    # B is still paused inside its evaluation function.
+    release_a.set()
+    assert a_done.wait(timeout=5), "thread A never finished"
+    thread_a.join(timeout=5)
+
+    release_b.set()
+    thread_b.join(timeout=5)
+
+    def metric_of(key):
+        status, trials = results.get(key, ("missing", None))
+        if status != "ok" or trials is None or len(trials) != 1:
+            return None
+        return trials[0].get("metric")
+
+    assert results.get("A", ("missing",))[0] == "ok", f"thread A did not complete cleanly: {results.get('A')}"
+    assert results.get("B", ("missing",))[0] == "ok", f"thread B did not complete cleanly: {results.get('B')}"
+    assert metric_of("A") == 1.0, f"thread A's own trial did not receive thread A's own metric: {results}"
+    assert metric_of("B") == 2.0, f"thread B's own trial did not receive thread B's own metric: {results}"