Adds secret env_var #3048

thomasjpfan · 2025-01-10T20:01:02Z

Tracking issue

Related to flyteorg/flyte#6141 (comment)
Requires flyteorg/flyte#6160

Why are the changes needed?

This PR adds an env_var to the Secrets IDL, which has the follow behavior:

If mount_requirement is ENV_VAR, then we set an environment variable named env_name to the value of the secret.
If mount_requirement is FILE, then we set an environment variable named env_name to the path of the mounted secret.

What changes were proposed in this pull request?

This PR adds env_name to the Secrets IDL. This makes it easy to configure a secret in a Flyte task. For example, one can easily set a hugging face secret:

@task(secret_requests=[Secret(..., env_var="HF_TOKEN", mount_requirement=Secret.MountType.ENV_VAR)
def hello():
    ...

Or for secrets that require a file:

@task(secret_requests=[Secret(..., env_var="HF_TOKEN_PATH", mount_requirement=Secret.MountType.FILE)
def hello():
    ...

How was this patch tested?

I ran the following to try the two different modes:

from flytekit import task, Secret, ImageSpec
from typing import Optional
import os

image = ImageSpec(
    apt_packages=["git"],
    packages=[
        "git+https://github.com/thomasjpfan/flytekit.git@8814da92a8578dbad07550486a5aaced5785d7a5",
    ],
    registry="localhost:30000",
    commands=[
        "uv pip install git+https://github.com/thomasjpfan/flyte.git@25a7d9c100f571a9ce394e802d50ef8861eae7a4#subdirectory=flyteidl"
    ],
)



@task(
    container_image=image,
    secret_requests=[
        Secret(
            key="token",
            group="my-fun-group",
            mount_requirement=Secret.MountType.ENV_VAR,
            env_var="HELLO_WORLD",
        )
    ],
)
def get_secret_env_var() -> Optional[str]:
    return os.getenv("HELLO_WORLD")


@task(
    container_image=image,
    secret_requests=[
        Secret(
            key="token",
            group="my-fun-group",
            mount_requirement=Secret.MountType.FILE,
            env_var="HELLO_WORLD",
        )
    ],
)
def get_secret_file() -> str:
    with open(os.getenv("HELLO_WORLD"), "r") as f:
        return f.read()

Docs link

Summary by Bito

Enhanced Secret class in Flytekit implementing dual access methods (ENV_VAR and FILE) for secret management. Changes include renamed env_var parameter, improved functionality for direct secret value access, and core improvements to error handling. Added integration tests for Kubernetes secrets management with test fixtures and workflow demonstrations. This pull request also introduces an environment variable to the Secrets IDL, allowing secrets to be accessed as environment variables or file paths. It includes updates to integration tests for Kubernetes secrets with added checks for kubectl availability.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 5

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-10T20:01:18Z

Code Review Agent Run #f2a7e1

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 615499e..615499e
- flytekit/models/security.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-10T20:03:38Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Enhance Secret Class with env_var	- `security.py` - Added env_var parameter to Secret class for flexible secret management.
Other Improvements - Update Dependency Version	- `pyproject.toml` - Updated flyteidl dependency version from 1.14.1 to 1.14.2.
Testing - Add Integration Tests for Secret Management	- `test_remote.py` - Introduced tests for Kubernetes secret management with fixtures and parametrized tests. - `get_secret.py` - Added new workflow tasks to test secret retrieval via environment variable and file.

codecov · 2025-01-10T20:22:59Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 83.58%. Comparing base (a86048a) to head (b2007d9).
Report is 6 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3048      +/-   ##
==========================================
+ Coverage   77.44%   83.58%   +6.14%     
==========================================
  Files         238        3     -235     
  Lines       23158      195   -22963     
  Branches     2760        0    -2760     
==========================================
- Hits        17935      163   -17772     
+ Misses       4412       32    -4380     
+ Partials      811        0     -811

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

cosmicBboy · 2025-01-13T20:05:03Z

lgtm, tests will pass in this PR once this is merged right? flyteorg/flyte#6160

thomasjpfan · 2025-01-13T21:01:54Z

lgtm, tests will pass in this PR once this is merged right? flyteorg/flyte#6160

Yea and when a new flyteidl version is released.

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-16T18:47:37Z

Code Review Agent Run #7c9e36

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 615499e..8814da9
- flytekit/models/security.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

…t_env_name Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

eapolinario · 2025-01-22T20:16:13Z

Does it make sense to add an integration test (and leave it disabled while until we get a new Flyte sandbox version up)? At least you'll be able to validate that this works locally.

flyte-bot · 2025-01-22T20:53:34Z

Code Review Agent Run #4296d4

Actionable Suggestions - 0

Additional Suggestions - 10

flytekit/exceptions/user.py - 1
- Consider using classmethod for error code · Line 32-34
flytekit/remote/remote.py - 1
- Consider consolidating duplicated signal handling logic · Line 635-721
plugins/flytekit-optuna/flytekitplugins/optuna/optimizer.py - 1
- Consider simplifying validation logic conditions · Line 84-87
flytekit/core/worker_queue.py - 4
- Consider using dataclass for Update class · Line 80-88
- Make polling parameters configurable · Line 287-305
- Consider configurable sleep duration value · Line 400-400
- Consider encapsulating enum state checks · Line 114-115
tests/flytekit/clis/sdk_in_container/test_serve.py - 2
- Consider consistent CLI argument naming convention · Line 19-20
- Consider more robust mocking approach · Line 14-14
flytekit/core/type_engine.py - 1
- Consider adding runtime type info to error · Line 1528-1535

Review Details

Files reviewed - 48 · Commit Range: 8814da9..7e0bf7e
- flytekit/bin/entrypoint.py
- flytekit/clis/sdk_in_container/run.py
- flytekit/clis/sdk_in_container/serve.py
- flytekit/core/array_node_map_task.py
- flytekit/core/context_manager.py
- flytekit/core/node.py
- flytekit/core/promise.py
- flytekit/core/python_function_task.py
- flytekit/core/type_engine.py
- flytekit/core/worker_queue.py
- flytekit/exceptions/user.py
- flytekit/image_spec/default_builder.py
- flytekit/image_spec/image_spec.py
- flytekit/remote/remote.py
- flytekit/tools/translator.py
- plugins/flytekit-envd/flytekitplugins/envd/image_builder.py
- plugins/flytekit-envd/tests/test_image_spec.py
- plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py
- plugins/flytekit-kf-pytorch/tests/test_elastic_task.py
- plugins/flytekit-optuna/flytekitplugins/optuna/__init__.py
- plugins/flytekit-optuna/flytekitplugins/optuna/optimizer.py
- plugins/flytekit-optuna/setup.py
- plugins/flytekit-optuna/tests/test_callback.py
- plugins/flytekit-optuna/tests/test_decorator.py
- plugins/flytekit-optuna/tests/test_imperative.py
- plugins/flytekit-optuna/tests/test_optimizer.py
- plugins/flytekit-optuna/tests/test_validation.py
- plugins/flytekit-pydantic/flytekitplugins/pydantic/__init__.py
- plugins/flytekit-pydantic/flytekitplugins/pydantic/basemodel_transformer.py
- plugins/flytekit-pydantic/flytekitplugins/pydantic/commons.py
- plugins/flytekit-pydantic/flytekitplugins/pydantic/deserialization.py
- plugins/flytekit-pydantic/flytekitplugins/pydantic/serialization.py
- plugins/flytekit-pydantic/setup.py
- plugins/flytekit-pydantic/tests/folder/test_file1.txt
- plugins/flytekit-pydantic/tests/folder/test_file2.txt
- plugins/flytekit-pydantic/tests/test_type_transformer.py
- pyproject.toml
- tests/flytekit/clis/sdk_in_container/test_serve.py
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/integration/remote/workflows/basic/signal_test.py
- tests/flytekit/unit/bin/test_python_entrypoint.py
- tests/flytekit/unit/core/image_spec/test_default_builder.py
- tests/flytekit/unit/core/test_array_node_map_task.py
- tests/flytekit/unit/core/test_type_engine.py
- tests/flytekit/unit/core/test_type_hints.py
- tests/flytekit/unit/core/test_worker_queue.py
- tests/flytekit/unit/exceptions/test_user.py
- tests/flytekit/unit/remote/test_remote.py
Files skipped - 2
- plugins/flytekit-optuna/README.md - Reason: Filter setting
- plugins/flytekit-pydantic/README.md - Reason: Filter setting
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

thomasjpfan · 2025-01-22T22:40:12Z

Does it make sense to add an integration test (and leave it disabled while until we get a new Flyte sandbox version up)?

Added integration test in bc9b040 (Which is skipped now)

For it to be enabled, kubectl must be usable in the CI to create the secret.

flyte-bot · 2025-01-22T23:43:33Z

Code Review Agent Run #7b0105

Actionable Suggestions - 3

tests/flytekit/integration/remote/workflows/basic/get_secret.py - 1
- Consider adding error handling for secrets · Line 10-11
tests/flytekit/integration/remote/test_remote.py - 2
- Use secure secret generation method · Line 906-907
- Add error handling for kubectl commands · Line 908-928

Review Details

Files reviewed - 2 · Commit Range: 7e0bf7e..27a0d3e
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/integration/remote/workflows/basic/get_secret.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-22T23:46:50Z

tests/flytekit/integration/remote/workflows/basic/get_secret.py

+def get_secret() -> str:
+    return current_context().secrets.get(group="my-group", key="token")


Consider adding error handling for secrets

Consider adding error handling around current_context().secrets.get() call to handle cases where the secret may not exist or be accessible.

Code suggestion

Check the AI-generated fix before applying

Suggested change

def get_secret() -> str:

return current_context().secrets.get(group="my-group", key="token")

def get_secret() -> str:

try:

return current_context().secrets.get(group="my-group", key="token")

except ValueError as e:

raise ValueError(f"Failed to retrieve secret: {str(e)}")

Code Review Run #7b0105

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

tests/flytekit/integration/remote/test_remote.py

flyte-bot · 2025-01-22T23:46:52Z

tests/flytekit/integration/remote/test_remote.py

+    subprocess.run([
+        "kubectl",
+        "create",
+        "secret",
+        "-n",
+        "flytesnacks-development",
+        "generic",
+        "my-group",
+        f"--from-literal=token={secret}",
+    ],  capture_output=True, text=True)
+    yield secret
+
+    # Remove secret
+    subprocess.run([
+        "kubectl",
+        "delete",
+        "secrets",
+        "-n",
+        "flytesnacks-development",
+        "my-group",
+    ],  capture_output=True, text=True)


Add error handling for kubectl commands

Consider handling potential errors from subprocess.run() calls when creating and deleting the secret. The command could fail if kubectl is not installed or user lacks permissions.

Code suggestion

Check the AI-generated fix before applying

@@ -908,11 +908,14 @@ result = subprocess.run([ "kubectl", "create", "secret", "-n", "flytesnacks-development", "generic", "my-group", f"--from-literal=token={secret}", ], capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"Failed to create secret: {result.stderr}") yield secret # Remove secret result = subprocess.run([ "kubectl", "delete", "secrets", "-n", "flytesnacks-development", "my-group", ], capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"Failed to delete secret: {result.stderr}")

Code Review Run #7b0105

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-23T03:43:33Z

Code Review Agent Run #aa41e1

Actionable Suggestions - 3

tests/flytekit/integration/remote/workflows/basic/get_secret.py - 3
- Consider more descriptive env var names · Line 7-13
- Consider adding env var error handling · Line 20-20
- Consider adding error handling for file · Line 25-26

Review Details

Files reviewed - 2 · Commit Range: 27a0d3e..867bfd1
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/integration/remote/workflows/basic/get_secret.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

tests/flytekit/integration/remote/workflows/basic/get_secret.py

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-23T19:59:33Z

Code Review Agent Run #dcba26

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 867bfd1..b2007d9
- tests/flytekit/integration/remote/test_remote.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Adds secret env name

615499e

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

thomasjpfan requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla and Future-Outlier as code owners January 10, 2025 20:01

Use env_var

8814da9

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

thomasjpfan changed the title ~~Adds secret env_name~~ Adds secret env_var Jan 16, 2025

thomasjpfan mentioned this pull request Jan 16, 2025

Rename Secret.env_name to Secret.env_var flyteorg/flyte#6176

Merged

thomasjpfan added 3 commits January 22, 2025 14:47

Merge remote-tracking branch 'upstream/master' into thomasjpfan/secre…

c147535

…t_env_name Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Bump flyteidl

771ba14

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Smaller change

7e0bf7e

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

thomasjpfan added 2 commits January 22, 2025 17:39

Add integration test

bc9b040

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Add integration test

27a0d3e

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot reviewed Jan 22, 2025

View reviewed changes

tests/flytekit/integration/remote/test_remote.py Show resolved Hide resolved

flyte-bot reviewed Jan 22, 2025

View reviewed changes

Use env_var

daaf43e

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Check for file and env_var

867bfd1

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot reviewed Jan 23, 2025

View reviewed changes

tests/flytekit/integration/remote/workflows/basic/get_secret.py Show resolved Hide resolved

flyte-bot reviewed Jan 23, 2025

View reviewed changes

tests/flytekit/integration/remote/workflows/basic/get_secret.py Show resolved Hide resolved

flyte-bot reviewed Jan 23, 2025

View reviewed changes

tests/flytekit/integration/remote/workflows/basic/get_secret.py Show resolved Hide resolved

Add check for kubectl

b2007d9

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

eapolinario approved these changes Jan 23, 2025

View reviewed changes

thomasjpfan merged commit f2a1742 into flyteorg:master Jan 23, 2025
106 of 108 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adds secret env_var #3048

Adds secret env_var #3048

thomasjpfan commented Jan 10, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 10, 2025 •

edited

Loading

Code Review Agent Run #f2a7e1

flyte-bot commented Jan 10, 2025 •

edited

Loading

Changelist by Bito

codecov bot commented Jan 10, 2025 •

edited

Loading

cosmicBboy commented Jan 13, 2025

thomasjpfan commented Jan 13, 2025

flyte-bot commented Jan 16, 2025 •

edited

Loading

Code Review Agent Run #7c9e36

eapolinario commented Jan 22, 2025

flyte-bot commented Jan 22, 2025 •

edited

Loading

Code Review Agent Run #4296d4

thomasjpfan commented Jan 22, 2025

flyte-bot commented Jan 22, 2025 •

edited

Loading

Code Review Agent Run #7b0105

flyte-bot Jan 22, 2025

flyte-bot Jan 22, 2025

flyte-bot commented Jan 23, 2025 •

edited

Loading

Code Review Agent Run #aa41e1

flyte-bot commented Jan 23, 2025 •

edited

Loading

Code Review Agent Run #dcba26

		def get_secret() -> str:
		return current_context().secrets.get(group="my-group", key="token")

Adds secret env_var #3048

Adds secret env_var #3048

Conversation

thomasjpfan commented Jan 10, 2025 • edited by flyte-bot Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Docs link

Summary by Bito

flyte-bot commented Jan 10, 2025 • edited Loading

Code Review Agent Run #f2a7e1

flyte-bot commented Jan 10, 2025 • edited Loading

Changelist by Bito

codecov bot commented Jan 10, 2025 • edited Loading

Codecov Report

cosmicBboy commented Jan 13, 2025

thomasjpfan commented Jan 13, 2025

flyte-bot commented Jan 16, 2025 • edited Loading

Code Review Agent Run #7c9e36

eapolinario commented Jan 22, 2025

flyte-bot commented Jan 22, 2025 • edited Loading

Code Review Agent Run #4296d4

thomasjpfan commented Jan 22, 2025

flyte-bot commented Jan 22, 2025 • edited Loading

Code Review Agent Run #7b0105

flyte-bot Jan 22, 2025

Choose a reason for hiding this comment

flyte-bot Jan 22, 2025

Choose a reason for hiding this comment

flyte-bot commented Jan 23, 2025 • edited Loading

Code Review Agent Run #aa41e1

flyte-bot commented Jan 23, 2025 • edited Loading

Code Review Agent Run #dcba26

thomasjpfan commented Jan 10, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 10, 2025 •

edited

Loading

flyte-bot commented Jan 10, 2025 •

edited

Loading

codecov bot commented Jan 10, 2025 •

edited

Loading

flyte-bot commented Jan 16, 2025 •

edited

Loading

flyte-bot commented Jan 22, 2025 •

edited

Loading

flyte-bot commented Jan 22, 2025 •

edited

Loading

flyte-bot commented Jan 23, 2025 •

edited

Loading

flyte-bot commented Jan 23, 2025 •

edited

Loading