Set Environment Variables with Metaflow Decorator
Question
How can I use environment variables in a step of my Metaflow flow?
Solution
There are several ways this page will show you how to do this. The most straightforward way is to pass a dictionary to the vars argument of Metaflow's @environment decorator. This will work locally and with remote compute instances running your steps.
These methods demonstrate how to set arbitrary environment variables in the execution environment of steps in your flow. If you are looking to understand the environment variables specific to your Metaflow configuration - which by default lives at $HOME/.metaflowconfig/config.json - you can see them here.
You can pass a dictionary containing variables you want to use in a step to the vars argument of Metaflow's @environment step decorator.
1Set Environment Variable
You can pass a dictionary containing variables you want to use in a step to the vars argument of Metaflow's @environment step decorator.
Before running the flow, let's define an environment variable to read:
export MSG="hi"
2Run Flow
This flow shows how to:
- Use
@environmentto pass and then print the contents ofSTEP_ENV_VARinread_locally. - Use
@environmentto pass and then print the contents ofSTEP_ENV_VARinread_in_container.- This step demonstrates reading the environment variable into a customizable Docker image.
- This step is run on AWS batch. You could replace the
@batchdecorator with@kubernetes.
from metaflow import FlowSpec, step, batch, environment
import os
IMAGE = 'public.ecr.aws/docker/library/' + \
'python:3.9.12-buster'
class EnvVarFlow(FlowSpec):
@step
def start(self):
self.next(self.read_locally,
self.read_in_container)
@environment(vars={
"STEP_VAR": f"{os.getenv('MSG')} from local"})
@step
def read_locally(self):
print(f"secret message: {os.getenv('STEP_VAR')}")
self.next(self.join)
@environment(vars =
{"STEP_VAR": f"{os.getenv('MSG')} from container"})
@batch(image=IMAGE,
cpu=1)
@step
def read_in_container(self):
print(f"secret message: {os.getenv('STEP_VAR')}")
self.next(self.join)
@step
def join(self, inputs):
self.next(self.end)
@step
def end(self):
pass
if __name__ == "__main__":
EnvVarFlow()
python set_environment_variables_decorator.py run