Access intermediate storage URI

[amarrella]

Hi :)

I know that the intermediate storage API is being refactored, but I wanted to ask you if it's possible somehow (even piercing the internal APIs if necessary, I'm writing an internal framework) to access the intermediate storage uri and not just the file object / handle.

My use case is the following:

• execute some code that generates training set and a bunch of other files for a ML job
• run training on SageMaker (it currently needs to be on SageMaker as we are using one of Amazon's proprietary containers). SageMaker wants the S3 URI of the file, it would be amazing if at the same time I could be able to specify both the data dependency using the Dagster solid input def and also somehow derive the S3 URI for the input(s). For SageMaker i'm using a custom solid that invokes the AWS APIs.

I'm fine that the only allowed execution mode for that pipeline requires having S3 as intermediate storage.

Thanks for your help and for this amazing project :)

[schrockn]

Yes you should be able to do this but it takes a little doing. The internal implementation is a bit overabstracted as you'll see.

To get access to the intermediate storage during computation you can get here:

@solid
def a_solid(context):
    print(context.get_system_context().intermediate_storage)

@pipeline
def a_pipeline():
    a_solid()

And then to fish out the s3 uri it is something like this. Note that these are internal APIs and liable to change at any time. The future structure will be dramatically more simple to interact with.

from dagster import pipeline, solid, execute_pipeline
from dagster.core.execution.plan.objects import StepOutputHandle

@solid
def a_solid(context):
    system_context = context.get_system_context()
    intermediate_storage = system_context.intermediate_storage
    step_key = system_context.step.key
    step_output_handle = StepOutputHandle(step_key, "result")
    paths = intermediate_storage._get_paths(step_output_handle)
    obj_store = intermediate_storage.object_store
    key = obj_store.key_for_paths([intermediate_storage.root] + paths) 
    uri = obj_store.uri_for_key(key)
    print(uri)

@pipeline
def a_pipeline():
    a_solid()

execute_pipeline(a_pipeline)

[amarrella]

Thanks a lot for this :)

I'm looking forward to the new structure but this will get me going for now!

Thanks!