diff --git a/flytekit/clis/sdk_in_container/register.py b/flytekit/clis/sdk_in_container/register.py index c69754aac1..8f936ed8e8 100644 --- a/flytekit/clis/sdk_in_container/register.py +++ b/flytekit/clis/sdk_in_container/register.py @@ -9,6 +9,7 @@ from flytekit.clis.sdk_in_container.utils import domain_option_dec, project_option_dec from flytekit.configuration import ImageConfig from flytekit.configuration.default_images import DefaultImages +from flytekit.interaction.click_types import key_value_callback from flytekit.loggers import cli_logger from flytekit.tools import repo @@ -107,6 +108,15 @@ is_flag=True, help="Activate newly registered Launchplans. This operation deactivates previous versions of Launchplans.", ) +@click.option( + "--env", + "--envvars", + required=False, + multiple=True, + type=str, + callback=key_value_callback, + help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`", +) @click.argument("package-or-module", type=click.Path(exists=True, readable=True, resolve_path=True), nargs=-1) @click.pass_context def register( @@ -124,6 +134,7 @@ def register( package_or_module: typing.Tuple[str], dry_run: bool, activate_launchplans: bool, + env: typing.Optional[typing.Dict[str, str]], ): """ see help @@ -173,6 +184,7 @@ def register( fast=not non_fast, package_or_module=package_or_module, remote=remote, + env=env, dry_run=dry_run, activate_launchplans=activate_launchplans, ) diff --git a/flytekit/clis/sdk_in_container/serialize.py b/flytekit/clis/sdk_in_container/serialize.py index 98118c4492..392b57aee0 100644 --- a/flytekit/clis/sdk_in_container/serialize.py +++ b/flytekit/clis/sdk_in_container/serialize.py @@ -9,6 +9,7 @@ from flytekit.clis.sdk_in_container.constants import CTX_PACKAGES from flytekit.configuration import FastSerializationSettings, ImageConfig, SerializationSettings from flytekit.exceptions.scopes import system_entry_point +from flytekit.interaction.click_types import key_value_callback from flytekit.tools.fast_registration import fast_package from flytekit.tools.repo import serialize_to_folder @@ -16,6 +17,7 @@ CTX_LOCAL_SRC_ROOT = "local_source_root" CTX_FLYTEKIT_VIRTUALENV_ROOT = "flytekit_virtualenv_root" CTX_PYTHON_INTERPRETER = "python_interpreter" +CTX_ENV = "env" class SerializationMode(_Enum): @@ -33,6 +35,7 @@ def serialize_all( flytekit_virtualenv_root: typing.Optional[str] = None, python_interpreter: typing.Optional[str] = None, config_file: typing.Optional[str] = None, + env: typing.Optional[typing.Dict[str, str]] = None, ): """ This function will write to the folder specified the following protobuf types :: @@ -64,6 +67,7 @@ def serialize_all( ), flytekit_virtualenv_root=flytekit_virtualenv_root, python_interpreter=python_interpreter, + env=env, ) serialize_to_folder(pkgs, serialization_settings, local_source_root, folder) @@ -106,9 +110,23 @@ def serialize_all( "installed inside your container. Required for running `pyflyte serialize` in out of container mode when " "your container installs the flytekit virtualenv outside of the default `/opt/venv`", ) +@click.option( + "--env", + "--envvars", + required=False, + multiple=True, + type=str, + callback=key_value_callback, + help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`", +) @click.pass_context def serialize( - ctx, image_config: ImageConfig, local_source_root, in_container_config_path, in_container_virtualenv_root + ctx, + image_config: ImageConfig, + local_source_root, + in_container_config_path, + in_container_virtualenv_root, + env: typing.Optional[typing.Dict[str, str]], ): """ This command produces protobufs for tasks and templates. @@ -119,6 +137,7 @@ def serialize( """ ctx.obj[CTX_IMAGE] = image_config ctx.obj[CTX_LOCAL_SRC_ROOT] = local_source_root + ctx.obj[CTX_ENV] = env click.echo(f"Serializing Flyte elements with image {image_config}") if in_container_virtualenv_root: @@ -155,6 +174,7 @@ def workflows(ctx, folder=None): flytekit_virtualenv_root=ctx.obj[CTX_FLYTEKIT_VIRTUALENV_ROOT], python_interpreter=ctx.obj[CTX_PYTHON_INTERPRETER], config_file=ctx.obj.get(constants.CTX_CONFIG_FILE, None), + env=ctx.obj.get(CTX_ENV, None), ) @@ -194,6 +214,7 @@ def fast_workflows(ctx, folder=None, deref_symlinks=False): flytekit_virtualenv_root=ctx.obj[CTX_FLYTEKIT_VIRTUALENV_ROOT], python_interpreter=ctx.obj[CTX_PYTHON_INTERPRETER], config_file=ctx.obj.get(constants.CTX_CONFIG_FILE, None), + env=ctx.obj.get(CTX_ENV, None), ) diff --git a/flytekit/tools/repo.py b/flytekit/tools/repo.py index 9c38b16eab..195e3eea17 100644 --- a/flytekit/tools/repo.py +++ b/flytekit/tools/repo.py @@ -217,6 +217,7 @@ def register( fast: bool, package_or_module: typing.Tuple[str], remote: FlyteRemote, + env: typing.Optional[typing.Dict[str, str]], dry_run: bool = False, activate_launchplans: bool = False, ): @@ -239,6 +240,7 @@ def register( version=version, image_config=image_config, fast_serialization_settings=fast_serialization_settings, + env=env, ) if not version and fast: