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!