Hi All, I am trying to use the

resume

functionality and the

S3

client in conjunction, but having some difficulty. It seems that to get the artefact from the original run (i.e.,

origin_run_id

) I need to manually make an instance of the

Run

class for this class and then pass that in. This works, but then I have to know a-priori that I want the artefact for the original run and it wasn't subsequently made in a step of the resumed run. I guess I could use some try/except logic to first check if the artefact exists for the resumed run and, if not, then try the original run. Just wondering what other people do when using these two bits of functionality together?

this may be hard but can you provide a minimal reproducible example of this happening? sounds like its more complicated than just the s3 client since you mentioned you're pulling data from a prior run

Sure, here is my toy example

import json
import logging
import os

from metaflow import (
    FlowSpec,
    Run,
    S3,
    current,
    environment,
    project,
    step,
    workflow_specification,
)

from metaflow.plugins.datatools.s3.s3 import MetaflowS3NotFound


logger = logging.getLogger(__name__)

class S3FileFlow(FlowSpec):
    """
    Metaflow pipeline that illustrates using Metaflow's S3 client for reading from & writing to S3.

    See <https://docs.metaflow.org/scaling/data#data-in-s3-metaflows3> for more details.
    """

    @step
    def start(self):
        """Initialise the flow."""
        data = '1,2,3'
        with S3(run=self) as s3:
            res = json.dumps({'data': data})
            url = s3.put('data-test', res)
            <http://logger.info|logger.info>('Key URL: %s', url)
        self.next(self.read_from_batch)

    @step
    def read_from_batch(self):
        """Read the batch data from S3."""
        try:
            <http://logger.info|logger.info>('Trying to obtain key from current run.')
            with S3(run=self) as s3:
                data = s3.get('data-test').text
        except MetaflowS3NotFound:
            if current.origin_run_id:
                <http://logger.info|logger.info>('Failed. Trying to obtain key from original run.')
                run = Run(f'{current.flow_name}/{current.origin_run_id}')
                with S3(run=run) as s3:
                    data = s3.get('data-test').text
            else:
                raise
        <http://logger.info|logger.info>('File contents: %s', json.loads(data))
        self.next(self.end)

    @step
    def end(self):
        """End the flow."""
        print('Finished reading the data!')


if __name__ == '__main__':
    S3FileFlow(

ahh I get it, yeah the

run=self

is probably messing things up with the resume

Sorry, not sure what you mean. You are suggesting removing the

run=self

in the call and replacing it with the prefix?

that way the path is just the path and you can isolate the problem to

run=self

if the option building the path with current works

Sure, thanks. I think I need to check whether the data exists in the S3 bucket though, right?

ok this is the gist of it:

import json
import logging
import os

from metaflow import (
    FlowSpec,
    Run,
    S3,
    current,
    environment,
    project,
    step,
    workflow_specification,
)

from metaflow.plugins.datatools.s3.s3 import MetaflowS3NotFound

logger = logging.getLogger(__name__)


class S3FileFlow(FlowSpec):
    """
    Metaflow pipeline that illustrates using Metaflow's S3 client for reading from & writing to S3.

    See <https://docs.metaflow.org/scaling/data#data-in-s3-metaflows3> for more details.
    """
    bucket = 'mybucket'

    @step
    def start(self):
        """Initialise the flow."""
        data = '1,2,3'

        self.path = f's3://{self.bucket}/metaflow-runs/{current.flow_name}/{current.run_id}'
        with S3(s3root=self.path) as s3:
            res = json.dumps({'data': data})
            self.url = s3.put('data-test', res)
            <http://logger.info|logger.info>('Key URL: %s', self.url)
        self.next(self.read_from_batch)

    @step
    def read_from_batch(self):
        """Read the batch data from S3."""
        try:
            <http://logger.info|logger.info>('Trying to obtain key from current run.')
            with S3(s3root=self.path) as s3:
                data = s3.get('data-test').text
        except MetaflowS3NotFound:
            if current.origin_run_id:
                <http://logger.info|logger.info>('Failed. Trying to obtain key from original run.')
                run = Run(f'{current.flow_name}/{current.origin_run_id}')
                with S3(run=run) as s3:
                    data = s3.get('data-test').text
            else:
                raise
        <http://logger.info|logger.info>('File contents: %s', json.loads(data))
        self.next(self.end)

    @step
    def end(self):
        """End the flow."""
        print('Finished reading the data!')


if __name__ == '__main__':
    S3FileFlow()

OK, let me have a look. You think the try/except logic is still required in your version?

I'd adjust that logic to log the exception but can leave it in

ok, I see what you've done. Basically the path is stored via

self.path = ...

type functionality, so it remains the same on the resume, right? I think it will work. We have a lot of teams using Metaflow in our company now though, so would be good to understand the general expectation on the interaction between these two bits of functionality. Ideally would like a generic solution, but I am guessing that doesn't currently exist?

so I think the root cause is in the exception handling, when it goes down that path I believe current doesn't have a

origin_run_id

associated

def s3_resume_retry(f):
    @wraps(f)
    def wrapper(**kwargs):
        try:
            <http://logger.info|logger.info>('Using S3 client with current run.')
            return f(**kwargs)
        except MetaflowS3NotFound:
            if current.origin_run_id:
                <http://logger.info|logger.info>('Failed. Trying S3 client with original run.')
                run = Run(f'{current.flow_name}/{current.origin_run_id}')
                return f(**{**kwargs, **{'run': run}})
            else:
                raise
    return wrapper


@s3_resume_retry
def download_data(*, run: Union[FlowSpec, "Run"], bucket: Optional[str] = None, prefix: Optional[str] = None,
                  s3root: Optional[str] = None, **kwargs):
    with S3(run=run, bucket=bucket, prefix=prefix, s3root=s3root, **kwargs) as s3:
        return s3.get('data-test').text

OK, so here is what I have as a more generic version than my first pipeline. You have to move the data download into its own function so that you can decorate it, which is not ideal.

think you’re overthinking this one, if you swap out the run=self with the path created from the current object ,it solves the problem (if I’m following said problem that is)

path created from the current object

- Sorry, what do you mean by this? You mean the example you shared?

np - so the relevant snippet from the flow I sent over is here:

bucket = 'mybucket' # <----- added a constant for no real good reason at all but maybe this would be something that would be meaningful as a param

    @step
    def start(self):
        """Initialise the flow."""
        data = '1,2,3'

        self.path = f's3://{self.bucket}/metaflow-runs/{current.flow_name}/{current.run_id}' # <------ path from current 
        with S3(s3root=self.path) as s3:
            res = json.dumps({'data': data})
            self.url = s3.put('data-test', res) # <-------- we're persisting the whole object path here
            logger.info('Key URL: %s', self.url)
        self.next(self.read_from_batch)

I see. Yeah, I understood this bit. I just see most of our internal users naturally using the

run

argument. That's what I have seen internally and was also my natural inclination when using the tool.