Make the value in the label optional (#2465)

Signed-off-by: Kevin Su <[email protected]>

flytekit/clis/sdk_in_container/run.py

@@ -32,7 +32,7 @@
 from flytekit.core.type_engine import TypeEngine
 from flytekit.core.workflow import PythonFunctionWorkflow, WorkflowBase
 from flytekit.exceptions.system import FlyteSystemException
-from flytekit.interaction.click_types import FlyteLiteralConverter, key_value_callback
+from flytekit.interaction.click_types import FlyteLiteralConverter, key_value_callback, labels_callback
 from flytekit.interaction.string_literals import literal_string_repr
 from flytekit.loggers import logger
 from flytekit.models import security
@@ -174,7 +174,7 @@ class RunLevelParams(PyFlyteParams):
             multiple=True,
             type=str,
             show_default=True,
-            callback=key_value_callback,
+            callback=labels_callback,
             help="Labels to be attached to the execution of the format `label_key=label_value`.",
         )
     )

flytekit/interaction/click_types.py

@@ -57,6 +57,22 @@ def key_value_callback(_: typing.Any, param: str, values: typing.List[str]) -> t
     return result
 
 
+def labels_callback(_: typing.Any, param: str, values: typing.List[str]) -> typing.Optional[typing.Dict[str, str]]:
+    """
+    Callback for click to parse labels.
+    """
+    if not values:
+        return None
+    result = {}
+    for v in values:
+        if "=" not in v:
+            result[v.strip()] = ""
+        else:
+            k, v = v.split("=", 1)
+            result[k.strip()] = v.strip()
+    return result
+
+
 class DirParamType(click.ParamType):
     name = "directory path"
 

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -76,6 +76,16 @@ def test_pyflyte_run_wf(remote, remote_flag, workflow_file):
         assert result.exit_code == 0
 
 
+def test_pyflyte_run_with_labels():
+    workflow_file = pathlib.Path(__file__).parent / "workflow.py"
+    with mock.patch("flytekit.configuration.plugin.FlyteRemote"):
+        runner = CliRunner()
+        result = runner.invoke(
+            pyflyte.main, ["run", "--remote",  str(workflow_file), "my_wf", "--help"], catch_exceptions=False
+        )
+        assert result.exit_code == 0
+
+
 def test_imperative_wf():
     runner = CliRunner()
     result = runner.invoke(