[BUG] MlflowStorage skips queued WAITING trials

[!WARNING] Before submitting a PR, please make sure that:

• A maintainer has triaged this issue and applied the ready label
• This issue has no assignee
• No duplicate PR exists

PRs not meeting these requirements may be automatically closed.

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local machine, using a remote MLflow tracking server over HTTP.

MLflow version

• Client: source checkout at a192cbfa6
• Tracking server: source checkout at a192cbfa6
• Optuna: 5.0.0

System information

• OS Platform and Distribution: Linux
• Python version: 3.10

Describe the problem

MlflowStorage.set_trial_state_values updates an MLflow run from SCHEDULED (Optuna WAITING) to RUNNING, but returns False. Optuna interprets False as an unsuccessful claim of the waiting trial, so Study.optimize skips the queued trial and creates a different trial instead. The queued run remains RUNNING and is never evaluated.

The expected behavior is that a WAITING -> RUNNING transition returns True, while a repeated RUNNING -> RUNNING transition returns False.

Tracking information

The tracking server was started with:

mlflow server --backend-store-uri sqlite:////tmp/mlflow.db --no-serve-artifacts --host 127.0.0.1 --port 5135 --workers 1

Code to reproduce issue

import mlflow
import optuna
from mlflow.optuna import MlflowStorage

mlflow.set_tracking_uri("http://127.0.0.1:5135")
experiment_id = mlflow.create_experiment("optuna_waiting_trial_bug")
storage = MlflowStorage(experiment_id=experiment_id)
study = optuna.create_study(
    storage=storage,
    study_name="waiting-trial-regression",
    direction="minimize",
)
study.enqueue_trial({"x": 0.25})
executed = []

def objective(trial):
    x = trial.suggest_float("x", 0.0, 1.0)
    executed.append((trial.number, x))
    return x

study.optimize(objective, n_trials=1)
print(executed)
print([(trial.number, trial.state.name) for trial in study.trials])

On the affected code, the queued trial is left RUNNING, a new trial is executed, and the output contains trial number 1 instead of the enqueued trial number 0.

Stack trace

There is no exception. The incorrect return value causes the queued trial to be silently skipped.

Other info / logs

The issue is in mlflow/optuna/storage.py in MlflowStorage.set_trial_state_values. The method currently checks the requested state rather than the run's existing state, making the condition for RUNNING always return False.

Willingness to contribute

Yes. I can contribute a fix for this bug independently.

What component(s) does this bug affect?

• area/tracking: Tracking Service, tracking client APIs, autologging