def test_to_environment_variables(config):
env = to_environment_variables(
config,
include={"general.x", "general.y", "general.nested.x", "debug"})
assert env == {
"PREFECT__GENERAL__Y": "hi",
# Converts values to strings
"PREFECT__GENERAL__X": "1",
"PREFECT__DEBUG": "False",
# Handles nesting
"PREFECT__GENERAL__NESTED__X": "1",
}
def test_to_environment_variables_roundtrip(config, monkeypatch,
test_config_file_path):
keys = [".".join(k) for k in dict_to_flatdict(config)]
# Note prefix is different to avoid colliding with the `config` fixture env
env = to_environment_variables(config,
include=keys,
prefix="PREFECT_TEST_ROUND")
for k, v in env.items():
monkeypatch.setenv(k, v)
new_config = configuration.load_configuration(
test_config_file_path, env_var_prefix="PREFECT_TEST_ROUND")
assert new_config == config
def subprocess_heartbeat(heartbeat_cmd: List[str],
logger: Logger) -> Iterator[None]:
p = None
try:
# we use Popen + a prefect CLI for a few reasons:
# - using threads would interfere with the task; for example, a task
# which does not release the GIL would prevent the heartbeat thread from
# firing
# - using multiprocessing.Process would release the GIL but a subprocess
# cannot be spawned from a daemonic subprocess, and Dask sometimes will
# submit tasks to run within daemonic subprocesses
current_env = dict(os.environ).copy()
current_env.update(
to_environment_variables(
prefect.context.config,
include={
"cloud.auth_token",
"cloud.api_key",
"cloud.tenant_id",
"cloud.api",
},
))
clean_env = {k: v for k, v in current_env.items() if v is not None}
p = subprocess.Popen(
heartbeat_cmd,
env=clean_env,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
yield
except Exception:
logger.exception(
"Heartbeat process failed to start. This could result in a zombie run."
)
finally:
if p is not None:
exit_code = p.poll()
if exit_code is not None:
logger.error(
f"Heartbeat process died with exit code {exit_code}")
p.kill()
p.wait()
def generate_flow_run_environ(
flow_run_id: str,
flow_id: str,
run_config: RunConfig,
run_api_key: str = None,
include_local_env: bool = False,
) -> Dict[str, str]:
"""
Utility to generate the environment variables required for a flow run
Args:
- flow_run_id: The id for the flow run that will be executed
- flow_id: The id for the flow of the flow run that will be executed
- run_config: The run config for the flow run, contributes environment variables
- run_api_key: An optional API key to pass to the flow run for authenticating
with the backend. If not set, it will be pulled from the Client
- include_local_env: If `True`, the currently available environment variables
will be passed through to the flow run. Defaults to `False` for security.
Returns:
- A dictionary of environment variables
"""
# TODO: Generalize this and use it for all agents
# Local environment
env = cast(Dict[str, Optional[str]], os.environ.copy() if include_local_env else {})
# Pass through config options that can be overriden by run config
env.update(
to_environment_variables(
prefect.config,
include={
"logging.level",
"logging.format",
"logging.datefmt",
"cloud.send_flow_run_logs",
},
)
)
# Update with run config environment
if run_config is not None and run_config.env is not None:
env.update(run_config.env)
# Update with config options that cannot be overriden by the run config
env.update(
to_environment_variables(
prefect.config,
include={"backend", "cloud.api", "cloud.tenant_id"},
)
)
# Pass authentication through
client = prefect.Client() # Instantiate a client to get the current API key
env["PREFECT__CLOUD__API_KEY"] = run_api_key or client.api_key or ""
# Backwards compat for auth tokens
env["PREFECT__CLOUD__AUTH_TOKEN"] = (
run_api_key
or prefect.config.cloud.agent.get("auth_token")
or prefect.config.cloud.get("auth_token")
)
# Add context information for the run
env.update(
{
"PREFECT__CONTEXT__FLOW_RUN_ID": flow_run_id,
"PREFECT__CONTEXT__FLOW_ID": flow_id,
}
)
# Update hardcoded execution variables
env.update(
{
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
}
)
# Filter out `None` values
return {k: v for k, v in env.items() if v is not None}
def test_to_environment_variables_respects_prefix():
env = to_environment_variables(Config({"key": "value"}),
include={"key"},
prefix="FOO")
assert env == {"FOO__KEY": "value"}