diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 57c7168f36..cef13f0014 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -148,7 +148,7 @@ class RunLevelParams(PyFlyteParams): ) wait_execution: bool = make_click_option_field( click.Option( - param_decls=["--wait-execution", "wait_execution"], + param_decls=["--wait", "--wait-execution", "wait_execution"], required=False, is_flag=True, default=False, @@ -156,6 +156,16 @@ class RunLevelParams(PyFlyteParams): help="Whether to wait for the execution to finish", ) ) + poll_interval: int = make_click_option_field( + click.Option( + param_decls=["-i", "--poll-interval", "poll_interval"], + required=False, + type=int, + default=None, + show_default=True, + help="Poll interval in seconds to check the status of the execution", + ) + ) dump_snippet: bool = make_click_option_field( click.Option( param_decls=["--dump-snippet", "dump_snippet"], @@ -532,7 +542,6 @@ def run_remote( project=project, domain=domain, execution_name=run_level_params.name, - wait=run_level_params.wait_execution, options=options_from_run_params(run_level_params), type_hints=type_hints, overwrite_cache=run_level_params.overwrite_cache, @@ -541,15 +550,16 @@ def run_remote( cluster_pool=run_level_params.cluster_pool, execution_cluster_label=run_level_params.execution_cluster_label, ) + s = ( + click.style("\n[✔] ", fg="green") + + "Go to " + + click.style(execution.execution_url, fg="cyan") + + " to see execution in the console." + ) + click.echo(s) - console_url = remote.generate_console_url(execution) - s = ( - click.style("\n[✔] ", fg="green") - + "Go to " - + click.style(console_url, fg="cyan") - + " to see execution in the console." - ) - click.echo(s) + if run_level_params.wait_execution: + execution = remote.wait(execution, poll_interval=run_level_params.poll_interval) if run_level_params.wait_execution: if execution.closure.phase != WorkflowExecutionPhase.SUCCEEDED: