ミニマルな機械学習の実験管理を目指して

import luigi

class HelloWorld(luigi.Task):
    def requires(self):
        # requires no tasks
        return None

    def output(self):
        # output to a local file
        # `luigi.LocalTarget` を置き換えることでローカル以外にもファイルを保存することができます。
        return luigi.LocalTarget('helloworld.txt')

    def run(self):
        with self.output().open('w') as outfile:
            outfile.write('Hello World!\n')

class MyTask(luigi.Task):
    date = luigi.DataParameter()
    learning_rate = luigi.FloatParameter(default=.005)

[MyTask]
date = 2021-10-01
learning_rate = 0.1

luigi.build([MyTask(), ]) 

import mlflow

with mlflow.start_run():
    mlflow.log_param('lr', .005)
    mlflow.log_metric('accuracy', .5)

import hashlib
import json
from pathlib import Path
from typing import Dict, Any, Optional, final, TypeVar, Callable, List, Tuple

import luigi
import mlflow
from luigi import LocalTarget
from mlflow.entities import Run
from mlflow.protos.service_pb2 import ACTIVE_ONLY

from .utils import get_tmp_filepath  # dummy


T = TypeVar('T')


class MlflowTask(luigi.Task):

    class Meta:
        # 対応する MLflow の experiment の名前 
        mlflow_experiment_name = None
        # {出力ファイルの識別名: 出力ファイルのファイル名}
        # 例 `{"ids": "ids.csv", "embeddings": "embeddings.pickle"}`
        mlflow_artifact_fnames = dict()
        
    @abstractmethod
    def requires(self) -> Dict[str, luigi.Task]:
        """
        :return: {タスク名: タスク} からなる dict。例: `{"texts": FetchText(), "embeddings": EmbeddingTexts()}`
        """
        raise NotImplementedError()

    @abstractmethod
    def to_mlflow_tags(self) -> Dict[str, Any]:
        """
        パラメータを MLflow のタグにシリアライズする。
        """
        raise NotImplementedError()
    
    @abstractmethod
    def _run(self):
        """
        メインの処理 (入力から出力まで書く)
        """
        raise NotImplementedError()

    @final
    def run(self):
        mlflow.set_experiment(self.Meta.mlflow_experiment_name)
        self._run()

    def search_for_mlflow_run(self) -> Optional[Run]:
        """
        :return: 同じタグを持つ run が MLflow に既に存在する場合はそれを返す。無ければ None を返す。
        """
        experiment = mlflow.get_experiment_by_name(self.Meta.mlflow_experiment_name)
        # MLflow 用のクエリ文字列を生成する
        query_items = [
            f'tag.{pname} = "{pval}"'
            for pname, pval in self.to_mlflow_tags().items()
        ]
        query = ' and '.join(query_items)  # Convert into a query string used in MLflow
        res = mlflow.search_runs(
            experiment_ids=[experiment.experiment_id, ],
            filter_string=query,
            max_results=1,
            run_view_type=ACTIVE_ONLY,
            output_format='list',
        )
        if len(res) > 0:
            return res[0]
        else:
            return None

    def complete(self):
        # run が MLflow に存在すれば完了済みとみなす。
        is_complete = (self.search_for_mlflow_run() is not None)
        self.logger.info(f'is_complete: {is_complete}')
        return is_complete

    def output(self) -> Dict[str, LocalTarget]:
        """
        試用
        - タスクが既に完了してる (= MLflow に run がある) 場合は、その artifact path(s) を返す。
        - タスクがまだ完了していない  (= MLflow に run がない) 場合は、一時ディレクトリ path を返す。
          (MLflow run を作るまでの一時的な待避先)
        """
        maybe_mlf_run = self.search_for_mlflow_run()
        if maybe_mlf_run is None:
            paths = {
                key: get_tmp_filepath(self, fname)
                for key, fname in self.Meta.mlflow_artifact_fnames.items()
            }
        else:
            paths = {
                key: Path(maybe_mlf_run.info.artifact_uri) / fname
                for key, fname in self.Meta.mlflow_artifact_fnames.items()
            }
        return {
            key: LocalTarget(str(p))
            for key, p in paths.items()
        }

    @final
    def save_to_mlflow(
        self,
        artifacts_and_save_funcs: Dict[str, Tuple[T, Callable[[T, str], None]]] = None,
        metrics: Dict[str, float] = None,
    ):
        """
        出力 (artifact), メトリクスを MLflow に登録する。`_run()` から呼び出されることを想定。
        
        :param artifacts_and_save_funcs:
            `(artifact_obj, artifact_saver)`. のリスト。
           `artifact_saver` は artifact とファイルパスを受け取り、artifact をそのパスに保存する関数。
           `artifact_saver` の一例:  `lambda df, path: df.to_csv(path)`
        """
        artifact_paths: List[str] = []
        artifacts_and_save_funcs = artifacts_and_save_funcs or []
        # Save artifacts
        for name, (artifact, save_fn) in artifacts_and_save_funcs.items():
            out_path = self.output()[name].path
            save_fn(artifact, out_path)
            artifact_paths.append(out_path)
        # Log into MLflow
        with mlflow.start_run():
            mlflow.set_tags(self.to_mlflow_tags())
            if metrics is not None:
                mlflow.log_metrics(metrics)
            for path in artifact_paths:
                mlflow.log_artifact(path)

from typing import Any, Dict

import luigi
import pandas as pd

from .tasks.base import MlflowTask
from .tasks.extract_ids import FetchArticleIds  # dummy
from .utils import get_title_from_id  # dummy


class FetchArticleTitles(MlflowTask):

    class Meta:
        mlflow_experiment_name = 'fetch_article_titles'
        mlflow_artifact_fnames = {
            'title': 'titles.csv'
        }

    def requires(self) -> Dict[str, luigi.Task]:
        return {
            'ids': FetchArticleIds(),
        }

    def to_mlflow_tags(self) -> Dict[str, Any]:
        return {
            'ids_path': self.input()['ids']['ids'].path,
        }


    def _run(self):
        articles_df = pd.read_csv(self.input()['ids']['ids'].path)
        # get title
        articles_df['title'] = articles_df['id'].apply(
            lambda x: get_title_from_id(x)
        )
        # Save
        self.save_to_mlflow(
            {
                'title': (articles_df, lambda df, path: df.to_csv(path)),
            }
        )

luigi.build([FetchArticleTitles(), ])