External Datastores Support [Experimental]
- Users can plugin external metaflow compatible datastores to save objects from the
metaflow-checkpoint family of decorators.
- metaflow-checkpoint exposes a
with_artifact_store flow decorator to save/read all objects from the external datastore.
@with_artifact_store(
type="s3",
config= {
"root": MY_CUSTOM_BUCKET,
"role_arn": ROLE,
},
)
class StatefulResumeRetryCase(FlowSpec):
@checkpoint(load_policy="eager")
@step
def start(self):
from os import environ
- Objects saved to external datastores can be accessed from within the
artifact_store_from context manager. Setting the context manager will allow metaflow to perform all read operations from Checkpoint or load_model on the datastore set in the Task/Run. It will essentially allow metaflow to read from the external datastore while also giving users a way to pass and credentials (or configurations).
from metaflow import Task, Run, artifact_store_from, Checkpoint, load_model
task = Task("TorchTuneFlow/8484/train/53673")
with artifact_store_from(task=task, config={"role_arn":"role_foo"}):
load_model(
task.data.model_ref,
"test-models"
)
run = Run("CheckpointsTestsFlow/8992")
with artifact_store_from(run=run, config={"client_params":{"aws_access_key_id":"foo", "aws_secret_access_key":"bar"}}):
with Checkpoint() as cp:
latest = cp.list(
task=run["start"].task
)[0]
print(latest)
cp.load(
latest,
"test-checkpoints"
)