Streamline fast register options with new flag (#2690)

Signed-off-by: Yee Hing Tong <[email protected]>

flytekit/clis/sdk_in_container/helpers.py

@@ -7,6 +7,7 @@
 from flytekit.configuration import ImageConfig
 from flytekit.configuration.plugin import get_plugin
 from flytekit.remote.remote import FlyteRemote
+from flytekit.tools.fast_registration import CopyFileDetection
 
 FLYTE_REMOTE_INSTANCE_KEY = "flyte_remote"
 
@@ -61,3 +62,17 @@ def patch_image_config(config_file: Optional[str], image_config: ImageConfig) ->
             if addl.name not in additional_image_names:
                 new_additional_images.append(addl)
     return replace(image_config, default_image=new_default, images=new_additional_images)
+
+
+def parse_copy(ctx, param, value) -> Optional[CopyFileDetection]:
+    """Helper function to parse cmd line args into enum"""
+    if value == "auto":
+        copy_style = CopyFileDetection.LOADED_MODULES
+    elif value == "all":
+        copy_style = CopyFileDetection.ALL
+    elif value == "none":
+        copy_style = CopyFileDetection.NO_COPY
+    else:
+        copy_style = None
+
+    return copy_style

flytekit/clis/sdk_in_container/package.py

@@ -1,16 +1,19 @@
 import os
+import typing
 
 import rich_click as click
 
 from flytekit.clis.helpers import display_help_with_error
 from flytekit.clis.sdk_in_container import constants
+from flytekit.clis.sdk_in_container.helpers import parse_copy
 from flytekit.configuration import (
     DEFAULT_RUNTIME_PYTHON_INTERPRETER,
     FastSerializationSettings,
     ImageConfig,
     SerializationSettings,
 )
 from flytekit.interaction.click_types import key_value_callback
+from flytekit.tools.fast_registration import CopyFileDetection, FastPackageOptions
 from flytekit.tools.repo import NoSerializableEntitiesError, serialize_and_package
 
 
@@ -50,8 +53,18 @@
     is_flag=True,
     default=False,
     required=False,
-    help="This flag enables fast packaging, that allows `no container build` deploys of flyte workflows and tasks. "
-    "Note this needs additional configuration, refer to the docs.",
+    help="[Will be deprecated, see --copy] This flag enables fast packaging, that allows `no container build`"
+    " deploys of flyte workflows and tasks. You can specify --copy all/auto instead"
+    " Note this needs additional configuration, refer to the docs.",
+)
+@click.option(
+    "--copy",
+    required=False,
+    type=click.Choice(["all", "auto", "none"], case_sensitive=False),
+    default=None,  # this will be changed to "none" after removing fast option
+    callback=parse_copy,
+    help="[Beta] Specify whether local files should be copied and uploaded so task containers have up-to-date code"
+    " 'all' will behave as the current 'fast' flag, copying all files, 'auto' copies only loaded Python modules",
 )
 @click.option(
     "-f",
@@ -100,6 +113,7 @@ def package(
     source,
     output,
     force,
+    copy: typing.Optional[CopyFileDetection],
     fast,
     in_container_source_path,
     python_interpreter,
@@ -113,6 +127,12 @@ def package(
     object contains the WorkflowTemplate, along with the relevant tasks for that workflow.
     This serialization step will set the name of the tasks to the fully qualified name of the task function.
     """
+    if copy is not None and fast:
+        raise ValueError("--fast and --copy cannot be used together. Please use --copy all instead.")
+    elif copy == CopyFileDetection.ALL or copy == CopyFileDetection.LOADED_MODULES:
+        # for those migrating, who only set --copy all/auto but don't have --fast set.
+        fast = True
+
     if os.path.exists(output) and not force:
         raise click.BadParameter(
             click.style(
@@ -136,6 +156,12 @@ def package(
         display_help_with_error(ctx, "No packages to scan for flyte entities. Aborting!")
 
     try:
-        serialize_and_package(pkgs, serialization_settings, source, output, fast, deref_symlinks)
+        # verbosity greater than 0 means to print the files
+        show_files = ctx.obj[constants.CTX_VERBOSE] > 0
+
+        fast_options = FastPackageOptions([], copy_style=copy, show_files=show_files)
+        serialize_and_package(
+            pkgs, serialization_settings, source, output, fast, deref_symlinks, fast_options=fast_options
+        )
     except NoSerializableEntitiesError:
         click.secho(f"No flyte objects found in packages {pkgs}", fg="yellow")

flytekit/clis/sdk_in_container/register.py

@@ -5,13 +5,18 @@
 
 from flytekit.clis.helpers import display_help_with_error
 from flytekit.clis.sdk_in_container import constants
-from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context, patch_image_config
+from flytekit.clis.sdk_in_container.helpers import (
+    get_and_save_remote_with_click_context,
+    parse_copy,
+    patch_image_config,
+)
 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 logger
 from flytekit.tools import repo
+from flytekit.tools.fast_registration import CopyFileDetection
 
 _register_help = """
 This command is similar to ``package`` but instead of producing a zip file, all your Flyte entities are compiled,
@@ -93,7 +98,17 @@
     "--non-fast",
     default=False,
     is_flag=True,
-    help="Skip zipping and uploading the package",
+    help="[Will be deprecated, see --copy] Skip zipping and uploading the package. You can specify --copy none instead",
+)
+@click.option(
+    "--copy",
+    required=False,
+    type=click.Choice(["all", "auto", "none"], case_sensitive=False),
+    default=None,  # this will be changed to "all" after removing non-fast option
+    callback=parse_copy,
+    help="[Beta] Specify how and whether to use fast register"
+    " 'all' is the current behavior copying all files from root, 'auto' copies only loaded Python modules"
+    " 'none' means no files are copied, i.e. don't use fast register",
 )
 @click.option(
     "--dry-run",
@@ -139,6 +154,7 @@ def register(
     version: typing.Optional[str],
     deref_symlinks: bool,
     non_fast: bool,
+    copy: typing.Optional[CopyFileDetection],
     package_or_module: typing.Tuple[str],
     dry_run: bool,
     activate_launchplans: bool,
@@ -148,14 +164,24 @@ def register(
     """
     see help
     """
+    if copy is not None and non_fast:
+        raise ValueError("--non-fast and --copy cannot be used together. Use --copy none instead.")
+
+    # Handle the new case where the copy flag is used instead of non-fast
+    if copy == CopyFileDetection.NO_COPY:
+        non_fast = True
+        # Set this to None because downstream logic currently detects None to mean old logic.
+        copy = None
+    show_files = ctx.obj[constants.CTX_VERBOSE] > 0
+
     pkgs = ctx.obj[constants.CTX_PACKAGES]
     if not pkgs:
         logger.debug("No pkgs")
     if pkgs:
         raise ValueError("Unimplemented, just specify pkgs like folder/files as args at the end of the command")
 
     if non_fast and not version:
-        raise ValueError("Version is a required parameter in case --non-fast is specified.")
+        raise ValueError("Version is a required parameter in case --non-fast/--copy none is specified.")
 
     if len(package_or_module) == 0:
         display_help_with_error(
@@ -190,10 +216,12 @@ def register(
         version,
         deref_symlinks,
         fast=not non_fast,
+        copy_style=copy,
         package_or_module=package_or_module,
         remote=remote,
         env=env,
         dry_run=dry_run,
         activate_launchplans=activate_launchplans,
         skip_errors=skip_errors,
+        show_files=show_files,
     )

flytekit/clis/sdk_in_container/run.py

@@ -19,7 +19,10 @@
 from typing_extensions import get_origin
 
 from flytekit import Annotations, FlyteContext, FlyteContextManager, Labels, Literal
-from flytekit.clis.sdk_in_container.helpers import patch_image_config
+from flytekit.clis.sdk_in_container.helpers import (
+    parse_copy,
+    patch_image_config,
+)
 from flytekit.clis.sdk_in_container.utils import (
     PyFlyteParams,
     domain_option,
@@ -63,6 +66,7 @@
 )
 from flytekit.remote.executions import FlyteWorkflowExecution
 from flytekit.tools import module_loader
+from flytekit.tools.fast_registration import CopyFileDetection, FastPackageOptions
 from flytekit.tools.script_mode import _find_project_root, compress_scripts, get_all_modules
 from flytekit.tools.translator import Options
 
@@ -104,7 +108,20 @@ class RunLevelParams(PyFlyteParams):
             is_flag=True,
             default=False,
             show_default=True,
-            help="Copy all files in the source root directory to the destination directory",
+            help="[Will be deprecated, see --copy] Copy all files in the source root directory to"
+            " the destination directory. You can specify --copy all instead",
+        )
+    )
+    copy: typing.Optional[CopyFileDetection] = make_click_option_field(
+        click.Option(
+            param_decls=["--copy"],
+            required=False,
+            default=None,  # this will change to "auto" after removing copy_all option
+            type=click.Choice(["all", "auto"], case_sensitive=False),
+            show_default=True,
+            callback=parse_copy,
+            help="[Beta] Specifies how to detect which files to copy into image."
+            " 'all' will behave as the current copy-all flag, 'auto' copies only loaded Python modules",
         )
     )
     image_config: ImageConfig = make_click_option_field(
@@ -626,6 +643,12 @@ def _run(*args, **kwargs):
             image_config = patch_image_config(config_file, image_config)
 
             with context_manager.FlyteContextManager.with_context(remote.context.new_builder()):
+                show_files = run_level_params.verbose > 0
+                fast_package_options = FastPackageOptions(
+                    [],
+                    copy_style=run_level_params.copy,
+                    show_files=show_files,
+                )
                 remote_entity = remote.register_script(
                     entity,
                     project=run_level_params.project,
@@ -635,6 +658,7 @@ def _run(*args, **kwargs):
                     source_path=run_level_params.computed_params.project_root,
                     module_name=run_level_params.computed_params.module,
                     copy_all=run_level_params.copy_all,
+                    fast_package_options=fast_package_options,
                 )
 
                 run_remote(

flytekit/remote/remote.py

@@ -1062,7 +1062,7 @@ def register_script(
             image_config = ImageConfig.auto_default_image()
 
         with tempfile.TemporaryDirectory() as tmp_dir:
-            if copy_all:
+            if copy_all or (fast_package_options and fast_package_options.copy_style):
                 md5_bytes, upload_native_url = self.fast_package(
                     pathlib.Path(source_path), False, tmp_dir, fast_package_options
                 )