Skip to main content

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 @environment to pass and then print the contents of STEP_ENV_VAR in read_locally.
  • Use @environment to pass and then print the contents of STEP_ENV_VAR in read_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 @batch decorator with @kubernetes.
set_environment_variables_decorator.py
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
...
[414/read_locally/2469 (pid 88407)] Task is starting.
[414/read_locally/2469 (pid 88407)] secret message: hi from local environment
[414/read_locally/2469 (pid 88407)] Task finished successfully.
...
[414/read_in_container/2470 (pid 88410)] Task is starting.
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task is starting (status SUBMITTED)...
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task is starting (status RUNNABLE)...
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task is starting (status STARTING)...
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task is starting (status RUNNING)...
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Setting up task environment.
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Downloading code package...
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Code package downloaded.
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task is starting.
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] secret message: hi from a container
[414/read_in_container/2470 (pid 88410)] [241ee31e-ca72-4f3e-9553-3d0791b16f0b] Task finished with exit code 0.
[414/read_in_container/2470 (pid 88410)] Task finished successfully.
...

Further Reading

For help or feedback, please join Metaflow Slack. To suggest an article, you may open an issue on GitHub.

Join Slack