Hyperopt で SparkTrials が有効になっている場合に実行が入れ子にならないRuns are not nested when SparkTrials is enabled in Hyperopt

問題Problem

SparkTrials は Hyperoptの拡張機能であり、実行を Spark worker に配布できます。SparkTrials is an extension of Hyperopt, which allows runs to be distributed to Spark workers.

ワーカー関数でを使用して MLflow 実行を開始すると、 nested=True 結果は親の実行の下に入れ子になっているはずです。When you start an MLflow run with nested=True in the worker function, the results are supposed to be nested under the parent run.

ワーカー関数でを使用して SparkTrials を実行した場合でも、結果が親の実行の下に正しく入れ子になって いない ことがあり nested=True ます。Sometimes the results are not correctly nested under the parent run, even though you ran SparkTrials with nested=True in the worker function.

次に例を示します。For example:

from hyperopt import fmin, tpe, hp, Trials, STATUS_OK

def train(params):
  """
  An example train method that computes the square of the input.
  This method will be passed to `hyperopt.fmin()`.

  :param params: hyperparameters. Its structure is consistent with how search space is defined. See below.
  :return: dict with fields 'loss' (scalar loss) and 'status' (success/failure status of run)
  """
  with mlflow.start_run(run_name='inner_run', nested=True) as run:

    x, = params
  return {'loss': x ** 2, 'status': STATUS_OK}

with mlflow.start_run(run_name='outer_run_with_sparktrials'):
  spark_trials_run_id = mlflow.active_run().info.run_id
  argmin = fmin(
    fn=train,
    space=search_space,
    algo=algo,
    max_evals=16,
    trials=spark_trials
  )

期待される結果:Expected results:

SparkTrials からの出力が必要です

実際の結果:Actual results:

SparkTrials からの実際の出力

原因Cause

オープンソースバージョンの Hyperopt は、Azure Databricks で SparkTrials MLflow の実行を適切に入れ子にするために必要な機能をサポートしていません。The open source version of Hyperopt does not support the required features necessary to properly nest SparkTrials MLflow runs on Azure Databricks.

解決策Solution

Machine Learning の Databricks Runtime には、Hyperopt の内部フォークと追加機能が含まれています。Databricks Runtime for Machine Learning includes an internal fork of Hyperopt with additional features. SparkTrials を使用する場合は、オープンソースリポジトリから Hyperopt を手動でインストールするのではなく、Machine Learning に Databricks Runtime を使用する必要があります。If you want to use SparkTrials, you should use Databricks Runtime for Machine Learning instead of installing Hyperopt manually from open-source repositories.