Pull secrets from environment when running locally

Signed-off-by: Thomas J. Fan <[email protected]>

flytekit/core/context_manager.py

@@ -346,9 +346,14 @@ def __getattr__(self, item: str) -> str:
     def __init__(self, secrets_cfg: typing.Optional[SecretsConfig] = None):
         if secrets_cfg is None:
             secrets_cfg = SecretsConfig.auto()
+            is_local_execution = os.getenv("FLYTE_INTERNAL_EXECUTION_ID") is None
+            if is_local_execution:
+                self._env_prefix = ""
+            else:
+                self._env_prefix = secrets_cfg.env_prefix.strip()
+
         self._base_dir = secrets_cfg.default_dir.strip()
         self._file_prefix = secrets_cfg.file_prefix.strip()
-        self._env_prefix = secrets_cfg.env_prefix.strip()
 
     def __getattr__(self, item: str) -> _GroupSecrets:
         """

flytekit/core/task.py

@@ -302,6 +302,13 @@ def foo2():
                      Possible options for secret stores are - Vault, Confidant, Kube secrets, AWS KMS etc
                      Refer to :py:class:`Secret` to understand how to specify the request for a secret. It
                      may change based on the backend provider.
+
+                     .. note::
+
+                         During local execution, the secrets will be pulled from the local environment variables
+                         with the format `{GROUP}_{GROUP_VERSION}_{KEY}`, where all the characters are capitalized
+                         and the prefix is not used.
+
     :param execution_mode: This is mainly for internal use. Please ignore. It is filled in automatically.
     :param node_dependency_hints: A list of tasks, launchplans, or workflows that this task depends on. This is only
         for dynamic tasks/workflows, where flyte cannot automatically determine the dependencies prior to runtime.

tests/flytekit/unit/core/test_context_manager.py

@@ -239,6 +239,46 @@ def test_secrets_manager_env():
     assert sec.get(group="group", key="key") == "value"
 
 
+@pytest.mark.parametrize("is_local_execution", [True, False])
+def test_secrets_manager_execution(monkeypatch, is_local_execution):
+    if not is_local_execution:
+        monkeypatch.setenv("FLYTE_INTERNAL_EXECUTION_ID", "my-execution-id")
+
+    sec = SecretsManager()
+    prefix = sec._env_prefix
+
+    if not is_local_execution:
+        assert prefix == "_FSEC_"
+    else:
+        assert prefix == ""
+
+    monkeypatch.setenv(f"{prefix}ABC_XYZ", "my-abc-secret")
+    assert sec.get(group="ABC", key="XYZ") == "my-abc-secret"
+
+
+@pytest.mark.parametrize("is_local_execution", [True, False])
+def test_secrets_manager_execution_no_group_required(monkeypatch, is_local_execution):
+    # Remove group requirements
+    plugin_mock = Mock()
+    plugin_mock.secret_requires_group.return_value = False
+    mock_global_plugin = {"plugin": plugin_mock}
+    monkeypatch.setattr(flytekit.configuration.plugin, "_GLOBAL_CONFIG", mock_global_plugin)
+
+    if not is_local_execution:
+        monkeypatch.setenv("FLYTE_INTERNAL_EXECUTION_ID", "my-execution-id")
+
+    sec = SecretsManager()
+    prefix = sec._env_prefix
+
+    if not is_local_execution:
+        assert prefix == "_FSEC_"
+    else:
+        assert prefix == ""
+
+    monkeypatch.setenv(f"{prefix}XYZ", "my-abc-secret")
+    assert sec.get(key="XYZ") == "my-abc-secret"
+
+
 def test_serialization_settings_transport():
     default_img = Image(name="default", fqn="test", tag="tag")
     serialization_settings = SerializationSettings(