last two steps of the task.py

task.py

import luigi
import os
from pathlib import Path

from util import DockerTask

class TrainModel(DockerTask):
    """Task to train random forest regressor with training datasets"""

    in_path = '/usr/share/data/interim/'

    in_train_features_csv = luigi.Parameter(default= in_path + 'train_features.csv')
    in_train_target_csv = luigi.Parameter(default= in_path + 'train_target.csv')
    out_dir = luigi.Parameter(default='/usr/share/data/output/')

    @property
    def image(self):
        return f'code-challenge/train-model:{VERSION}'

    def requires(self):
        return ImputeData()

    @property
    def command(self):
        return [
            'python', 'train_model.py',
            '--in-train-features-csv', self.in_train_features_csv,
            '--in-train-target-csv', self.in_train_target_csv,
            '--out-dir', self.out_dir
        ]

    def output(self):
        return luigi.LocalTarget(
            path=str(Path(self.out_dir) / 'model.sav')
        )

class EvaluateModel(DockerTask):
    """Task to evaluate random forest regressor with test datasets,
    returns mean square error of the predictions and several plots"""

    in_path = '/usr/share/data/interim/'

    in_test_features_csv = luigi.Parameter(default= in_path + 'test_features.csv')
    in_test_target_csv = luigi.Parameter(default= in_path + 'test_target.csv')
    in_trained_model = luigi.Parameter(default= '/usr/share/data/output/model.sav')
    out_dir = luigi.Parameter(default='/usr/share/data/output/')
    flag = luigi.Parameter('.SUCCESS_EvaluateModel')

    @property
    def image(self):
        return f'code-challenge/evaluate-model:{VERSION}'

    def requires(self):
        return TrainModel()

    @property
    def command(self):
        return [
            'python', 'evaluate_model.py',
            '--in-test-features-csv', self.in_test_features_csv,
            '--in-test-target-csv', self.in_test_target_csv,
            '--in-trained-model', self.in_trained_model,
            '--out-dir', self.out_dir,
            '--flag', self.flag
        ]

    def output(self):
        return luigi.LocalTarget(
            path=str(Path(self.out_dir) / self.flag)
        )