Запуски не вкладываются, если в Hyperopt включено расширение SparkTrials

Проблема

Спарктриалс — это расширение хиперопт, которое позволяет распределять запуски между рабочими процессами Spark.

При запуске Млфлов с помощью nested=True в рабочей функции результаты должны быть вложенными в родительский запуск.

Иногда результаты не будут правильно вложены в родительский запуск, даже если вы выполнили спарктриалс с nested=True в рабочей функции.

Пример.

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
  )

Ожидаемый результат:

Ожидаемые выходные данные из Спарктриалс

Фактические результаты:

Фактические выходные данные из Спарктриалс

Причина

Версия Хиперопт с открытым исходным кодом не поддерживает необходимые компоненты, необходимые для правильного вложения Спарктриалс Млфлов на Azure Databricks.

Решение

Databricks Runtime для Машинное обучение включает внутреннюю вилку Хиперопт с дополнительными функциями. Если вы хотите использовать Спарктриалс, следует использовать Databricks Runtime для Машинное обучение вместо установки Хиперопт вручную из репозиториев с открытым кодом.