如何在Luigi中使用早期任务的输出?

时间:2017-06-04 22:22:40

标签: python luigi

我正在编写一个管道,后来的任务需要读取早期任务的输出,这样他们才能知道需要通过哪些参数来满足他们的要求。

我在下面创建了一个简化的设置示例。

import random
import pickle
import luigi


class WriteNumbers(luigi.Task):

    def requires(self):
        pass

    def run(self):
        # pickle a random list of ints 1-100
        numbers = [random.randint(1, 100) for _ in range(100)]
        pickle.dump(numbers, open("./numbers.pkl", 'wb'))

    def output(self):
        return luigi.LocalTarget("./numbers.pkl")


class SquareNumber(luigi.Task):
    number = luigi.IntParameter()

    def requires(self):
        pass

    def run(self):
        # given a number as the parameter, write a file containing its square
        with open("./squared_{}".format(self.number), 'w') as f:
            f.write(str(self.number ** 2))

    def output(self):
        return luigi.LocalTarget("./squared_{}".format(self.number))


class SquareAll(luigi.WrapperTask):

    def requires(self):
        yield WriteNumbers()  # require the number list to be pickled first
        numbers = pickle.load("./numbers.pkl")  # load the number list
        for n in numbers:  # square each number in the number list
            yield SquareNumber(number=n)

class CubeNumber(luigi.Task):
    number = luigi.IntParameter()

    def requires(self):
        pass

    def run(self):
        # given a number as the parameter, write a file containing its cube
        with open("./cubed_{}".format(self.number), 'w') as f:
            f.write(str(self.number ** 3))

    def output(self):
        return luigi.LocalTarget("./cubed_{}".format(self.number))


class CubeAll(luigi.WrapperTask):

    def requires(self):
        yield WriteNumbers()  # require the number list to be pickled first
        numbers = pickle.load("./numbers.pkl")  # load the number list
        for n in numbers:  # square each number in the number list
            yield CubeNumber(number=n)

class CrunchNumbers(luigi.WrapperTask):
    def requires(self):
        yield SquareAll()
        yield CubeAll()

if __name__ == '__main__':
    luigi.run()

当通过python luigi_example.py CrunchNumbers运行时,将创建100个随机数字,并且列表会被腌制。转储到磁盘。 SquareAll加载该pickle列表并使用它来要求SquareNumber任务包含所需的参数。 CubeAll为其类似任务引用相同的结果文件。

问题是,运行时会抛出异常,因为numbers.pkl文件尚不存在。

如何允许以后的任务根据早期任务的输出生成依赖项?我在这里使用随机数表示输出无法提前知道:我的实际应用程序是从API中处理数据。

1 个答案:

答案 0 :(得分:0)

您正在使用Dynamic Dependencies,因此需要通过run方法来调用它们(当requires的结果可作为input使用时),因此{{1 }}和CubeAll的结构应如下:

SquareAll