Pyflyte --env Option for Register and Serialize (flyteorg#1880)

Use with `pyflyte register --env key=val ...`

Signed-off-by: Minura Punchihewa <[email protected]>

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,
         )

flytekit/clis/sdk_in_container/serialize.py

@@ -9,13 +9,15 @@
 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
 
 CTX_IMAGE = "image"
 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),
     )
 
 

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: