Skip to content

Commit

Permalink
Update eager: make sure client secret can be specified as env var (#2720
Browse files Browse the repository at this point in the history
)

* fix eager mode

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* bind secret to env var

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* add remote creation error handling

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update new arg to client_secret_env_var

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix bug

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix kwargs

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* try creating secret

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* add event loop if needed

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* debug

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* debug

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update error

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* pass default domain and project to new remote

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

---------

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
  • Loading branch information
cosmicBboy committed Sep 5, 2024
1 parent a9c9c46 commit 9f9f197
Showing 1 changed file with 36 additions and 7 deletions.
43 changes: 36 additions & 7 deletions flytekit/experimental/eager_function.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import asyncio
import inspect
import os
import signal
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
Expand Down Expand Up @@ -382,6 +383,7 @@ def eager(
timeout: Optional[timedelta] = None,
poll_interval: Optional[timedelta] = None,
local_entrypoint: bool = False,
client_secret_env_var: Optional[str] = None,
**kwargs,
):
"""Eager workflow decorator.
Expand All @@ -396,6 +398,8 @@ def eager(
:param local_entrypoint: If True, the eager workflow will can be executed locally but use the provided
:py:func:`~flytekit.remote.FlyteRemote` object to create task/workflow executions. This is useful for local
testing against a remote Flyte cluster.
:param client_secret_env_var: if specified, binds the client secret to the specified environment variable for
remote authentication.
:param kwargs: keyword-arguments forwarded to :py:func:`~flytekit.task`.
This type of workflow will execute all flyte entities within it eagerly, meaning that all python constructs can be
Expand Down Expand Up @@ -488,7 +492,10 @@ async def eager_workflow(x: int) -> int:
remote=remote,
client_secret_group=client_secret_group,
client_secret_key=client_secret_key,
timeout=timeout,
poll_interval=poll_interval,
local_entrypoint=local_entrypoint,
client_secret_env_var=client_secret_env_var,
**kwargs,
)

Expand All @@ -510,7 +517,9 @@ async def wrapper(*args, **kws):
execution_id = exec_params.execution_id

async_stack = AsyncStack(task_id, execution_id)
_remote = _prepare_remote(_remote, ctx, client_secret_group, client_secret_key, local_entrypoint)
_remote = _prepare_remote(
_remote, ctx, client_secret_group, client_secret_key, local_entrypoint, client_secret_env_var
)

# make sure sub-nodes as cleaned up on termination signal
loop = asyncio.get_event_loop()
Expand All @@ -533,8 +542,10 @@ async def wrapper(*args, **kws):
await cleanup_fn()

secret_requests = kwargs.pop("secret_requests", None) or []
if client_secret_group is not None and client_secret_key is not None:
try:
secret_requests.append(Secret(group=client_secret_group, key=client_secret_key))
except ValueError:
pass

return task(
wrapper,
Expand All @@ -551,6 +562,7 @@ def _prepare_remote(
client_secret_group: Optional[str] = None,
client_secret_key: Optional[str] = None,
local_entrypoint: bool = False,
client_secret_env_var: Optional[str] = None,
) -> Optional[FlyteRemote]:
"""Prepare FlyteRemote object for accessing Flyte cluster in a task running on the same cluster."""

Expand All @@ -576,7 +588,7 @@ def _prepare_remote(
if remote.config.platform.endpoint.startswith("localhost"):
# replace sandbox endpoints with internal dns, since localhost won't exist within the Flyte cluster
return _internal_demo_remote(remote)
return _internal_remote(remote, client_secret_group, client_secret_key)
return _internal_remote(remote, client_secret_group, client_secret_key, client_secret_env_var)


def _internal_demo_remote(remote: FlyteRemote) -> FlyteRemote:
Expand Down Expand Up @@ -605,16 +617,33 @@ def _internal_demo_remote(remote: FlyteRemote) -> FlyteRemote:

def _internal_remote(
remote: FlyteRemote,
client_secret_group: str,
client_secret_key: str,
client_secret_group: Optional[str],
client_secret_key: Optional[str],
client_secret_env_var: Optional[str],
) -> FlyteRemote:
"""Derives a FlyteRemote object from a yaml configuration file, modifying parts to make it work internally."""
assert client_secret_group is not None, "secret_group must be defined when using a remote cluster"
assert client_secret_key is not None, "secret_key must be defined a remote cluster"
secrets_manager = current_context().secrets

assert (
client_secret_group is not None or client_secret_key is not None
), "One of client_secret_group or client_secret_key must be defined when using a remote cluster"

client_secret = secrets_manager.get(client_secret_group, client_secret_key)
# get the raw output prefix from the context that's set from the pyflyte-execute entrypoint
# (see flytekit/bin/entrypoint.py)

if client_secret_env_var is not None:
# this creates a remote client where the env var client secret is sufficient for authentication
os.environ[client_secret_env_var] = client_secret
try:
remote_cls = type(remote)
return remote_cls(
default_domain=remote.default_domain,
default_project=remote.default_project,
)
except Exception as exc:
raise TypeError(f"Unable to authenticate remote class {remote_cls} with client secret") from exc

ctx = FlyteContextManager.current_context()
return FlyteRemote(
config=remote.config.with_params(
Expand Down

0 comments on commit 9f9f197

Please sign in to comment.