Hello! I was wondering how can I acquire AWS credentials with metaflow. I have a code that runs fine on it's own, but in the flow, it does not.

# working standalone code
import boto3

AWS_PROFILE = "dev"


def get_vault_session(aws_profile: str) -> boto3.Session:
    import json
    import subprocess

    aws_env_json = subprocess.check_output(
        ["aws-vault", "export", "--format=json", aws_profile]
    )
    aws_env = json.loads(aws_env_json)

    return boto3.Session(
        aws_access_key_id=aws_env["AccessKeyId"],
        aws_secret_access_key=aws_env["SecretAccessKey"],
        aws_session_token=aws_env["SessionToken"],
    )

if __name__ == "__main__":
    session = get_vault_session(AWS_PROFILE)
    print(session)

And while wrapped in metaflow, it returns an error.

# not working metaflow
import boto3
from metaflow import FlowSpec, step, project

AWS_PROFILE = "dev"


def get_vault_session(aws_profile: str) -> boto3.Session:
    import json
    import subprocess

    aws_env_json = subprocess.check_output(
        ["aws-vault", "export", "--format=json", aws_profile]
    )
    aws_env = json.loads(aws_env_json)

    return boto3.Session(
        aws_access_key_id=aws_env["AccessKeyId"],
        aws_secret_access_key=aws_env["SecretAccessKey"],
        aws_session_token=aws_env["SessionToken"],
    )

@project(name="test_flow")
class TestFlow(FlowSpec):
    @step
    def start(self) -> None:
        print("Running Test")
        self.next(self.get_credentials)

    @step
    def get_credentials(self) -> None:
        self.session = get_vault_session(AWS_PROFILE)
        self.next(self.end)

    @step
    def end(self) -> None:
        print("Test Completed")


if __name__ == "__main__":
    TestFlow()

Error:

AttributeError: Can't pickle local object 'lazy_call.<locals>._handler'

Any ideas how to overcome this? I want to use the credentials to invoke aws services. Thanks!

Generally you don't want to store AWS (or any other) credentials as an artifact, i.e. assign them to self. First, its not very secure since they get serialized to S3. Second, if those task run on AWS Batch or Kubernetes, they'd typically acquire the credentials independently from each other, typically from some mechansim like AWS Instance Profile or IRSA, or some secret store. Remember also you may want to use resume functionality, in this case the resumed execution may happen days later and it will only run the flow from the failed step. So if you get credentials in the first step, they'll be long expired. For those reasons typically each step should acquire credentials independently.