Analyze artifacts

analyze_artifacts.py

from metaflow import Parameter, current, Flow, step, FlowSpec
from functools import wraps

class ArtifactProxy:

    def __init__(self, flow):
        params = {x.id for x in Flow(current.flow_name)[current.run_id]['_parameters'].task}
        base = {x for x in dir(flow) if x not in params}
        self.__dict__.update({
             '_artifacts_flow': flow,
             '_artifacts_in': set(),
             '_artifacts_out': set(),
             '_artifacts_base': base})

    def __setattr__(self, key, val):
        self._artifacts_out.add(key)
        return setattr(self._artifacts_flow, key, val)

    def __getattr__(self, key):
        if key not in self._artifacts_base:
            self._artifacts_in.add(key)
        return getattr(self._artifacts_flow, key)

def analyze_artifacts(f):
    @wraps(f)
    def func(self):
        proxy = ArtifactProxy(self)
        ret = f(proxy)
        print('inputs', proxy._artifacts_in)
        print('outputs', proxy._artifacts_out)
    return func

class AnalyzeArtifactFlow(FlowSpec):

    param = Parameter('param')

    @analyze_artifacts
    @step
    def start(self):
        self.first = 'foo'
        self.second = 'bar'
        self.next(self.middle)

    @analyze_artifacts
    @step
    def middle(self):
        print('reading', self.first)
        x = self.param
        self.third = 'xyz'
        self.next(self.end)

    @analyze_artifacts
    @step
    def end(self):
        pass

if __name__ == '__main__':
    AnalyzeArtifactFlow()

Thank you for the demo. I don't fully understand why on line 30 it #prints {'first', third'} instead of just {'third'}?

@zexuan-zhou this example prints all artifacts accessed, both reads and writes (except parameters). The middle step reads first and writes third, hence both of them.

Take a look at this slightly more complex example that distinguishes reads and writes (and handles parameters): https://gist.github.com/tuulos/9a7966ca5ba6f5206738e2625baec8a1