MLflow застревает в состоянии «РАБОТАЕТ» при использовании Ray Tune ASHAScheduler с XGBoostTrainer. ⇐ Python

[Anonymous]

Я использую Ray 2.x с XGBoostTrainer для настройки гиперпараметров. Хотя Ray Tune предлагает MLflowLoggerCallback на стороне драйвера, он не обеспечивает детального контроля, который мне нужен для регистрации метрик уровня итерации и артефактов модели непосредственно от обучающих рабочих.
Проблема: я использую setup_mlflow() внутри train_loop_per_worker, чтобы инициировать ведение журнала на рабочей стороне. Однако когда ASHAScheduler сокращает пробную версию, рабочий процесс завершается. В результате соответствующий запуск MLflow остается в состоянии "РАБОТАЕТ" на неопределенный срок в пользовательском интерфейсе MLflow, а окончательное состояние/теги никогда не обновляются.

Код: Выделить всё

import os
import ray
from ray import train, tune
from ray.tune.schedulers import ASHAScheduler
from ray.air.integrations.mlflow import setup_mlflow
import mlflow
import time

# Simulation of worker-side training with fine-grained logging
def train_fn_per_worker(config):
# setup_mlflow starts a run on the worker
setup_mlflow(
tracking_uri="http://localhost:5000",
experiment_name="asha_orphaned_runs_test"
)

for step in range(100):
time.sleep(0.1) # Simulate work
intermediate_score = config["hw"] * step

# We log here for fine-grained tracking not available at driver level
mlflow.log_metric("iter_score", intermediate_score, step=step)

# Report to ASHA scheduler
train.report({"score": intermediate_score})

# 1. Setup Scheduler (ASHA will prune trials)
scheduler = ASHAScheduler(metric="score", mode="max", grace_period=5)

# 2. Setup Tuner
tuner = tune.Tuner(
train_fn_per_worker,
tune_config=tune.TuneConfig(
scheduler=scheduler,
num_samples=10
),
param_space={"hw": tune.uniform(0.1, 10)}
)

tuner.fit()

Что я пробовал:

• Завершение цикла обучения в try...finally для вызова mlflow.end_run(). Когда ASHA завершает пробную версию, исполнитель уничтожается таким образом, что блокfinally не выполняется последовательно или сервер MLflow не получает обновление статуса.
• Я не могу использовать только MLflowLoggerCallback на стороне драйвера, поскольку мне нужно регистрировать артефакты и пользовательские метрики, которые генерируются в середине итерации в обратном вызове XGBoost.

Среда:
Ray: 2.X
MLflow: 3.X
Python: 3.12+
Вопросы:

• Как синхронизировать сигнал завершения пробной версии от ASHAScheduler сеанс MLflow на рабочей стороне, чтобы гарантировать, что запуски помечаются как KILLED?
• Существует ли рекомендуемый шаблон для «передачи» запуска MLflow, инициированного работником, драйверу Ray Tune, чтобы драйвер мог управлять жизненным циклом после завершения или сокращения пробной версии?
• Как я могу пометить запуск как завершенный_by_asha, когда рабочий остановлен в середине цикла?