モデルをトレーニングし、MLflow で結果を保存する pandas_udf 関数に各モデルが渡される複数の Prophet モデルを構築しています。
@pandas_udf(result_schema, PandasUDFType.GROUPED_MAP)
def forecast(data):
......
with mlflow.start_run() as run:
......
次に、各 KPI のモデルをトレーニングするこの UDF を呼び出します。
df.groupBy('KPI').apply(forecast)
各 KPI について、モデルが複数のハイパーパラメーターでトレーニングされ、各モデルに最適なパラメーターが MLflow に格納されるという考え方です。Hyperopt を使用して、検索をより効率的にしたいと考えています。
この場合、目的関数をどこに配置すればよいでしょうか。データはモデルごとに UDF に渡されるため、実行ごとにデータを使用する UDF 内に内部関数を作成することを考えました。これは理にかなっていますか？