Optional entrypoint args

Author: schustmi

src/zenml/deployers/server/entrypoint_configuration.py

@@ -14,7 +14,7 @@
 """ZenML Pipeline Deployment Entrypoint Configuration."""
 
 import os
-from typing import Any, List, Set
+from typing import Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
@@ -46,20 +46,20 @@ class DeploymentEntrypointConfiguration(BaseEntrypointConfiguration):
     """
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for the deployment entrypoint.
 
         Returns:
             Set of required option names
         """
         return {
-            DEPLOYMENT_ID_OPTION,
-            HOST_OPTION,
-            PORT_OPTION,
-            WORKERS_OPTION,
-            LOG_LEVEL_OPTION,
-            CREATE_RUNS_OPTION,
-            AUTH_KEY_OPTION,
+            DEPLOYMENT_ID_OPTION: True,
+            HOST_OPTION: True,
+            PORT_OPTION: True,
+            WORKERS_OPTION: True,
+            LOG_LEVEL_OPTION: True,
+            CREATE_RUNS_OPTION: True,
+            AUTH_KEY_OPTION: True,
         }
 
     @classmethod

src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -17,7 +17,7 @@
 import os
 import sys
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Set
+from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional
 from uuid import UUID
 
 from zenml.client import Client
@@ -83,18 +83,18 @@ def get_entrypoint_command(cls) -> List[str]:
         return DEFAULT_ENTRYPOINT_COMMAND
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
-            A set of strings with all required options.
+            A dictionary of options and whether they are required.
         """
         return {
             # Importable source pointing to the entrypoint configuration class
             # that should be used inside the entrypoint.
-            ENTRYPOINT_CONFIG_SOURCE_OPTION,
+            ENTRYPOINT_CONFIG_SOURCE_OPTION: True,
             # ID of the pipeline snapshot to use in this entrypoint
-            SNAPSHOT_ID_OPTION,
+            SNAPSHOT_ID_OPTION: True,
         }
 
     @classmethod
@@ -178,13 +178,13 @@ def error(self, message: str) -> NoReturn:
 
         parser = _CustomParser()
 
-        for option_name in cls.get_entrypoint_options():
+        for option_name, required in cls.get_entrypoint_options().items():
             if option_name == ENTRYPOINT_CONFIG_SOURCE_OPTION:
                 # This option is already used by
                 # `zenml.entrypoints.entrypoint` to read which config
                 # class to use
                 continue
-            parser.add_argument(f"--{option_name}", required=True)
+            parser.add_argument(f"--{option_name}", required=required)
 
         result, _ = parser.parse_known_args(arguments)
         return vars(result)

src/zenml/entrypoints/step_entrypoint_configuration.py

@@ -15,7 +15,7 @@
 
 import os
 import sys
-from typing import TYPE_CHECKING, Any, List, Set
+from typing import TYPE_CHECKING, Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
@@ -115,14 +115,14 @@ def post_run(
         """
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the name of the
             step to run.
         """
-        return super().get_entrypoint_options() | {STEP_NAME_OPTION}
+        return super().get_entrypoint_options() | {STEP_NAME_OPTION: True}
 
     @classmethod
     def get_entrypoint_arguments(

src/zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py

@@ -15,7 +15,7 @@
 
 import json
 import os
-from typing import Any, List, Set
+from typing import Any, Dict, List
 
 from zenml.entrypoints.step_entrypoint_configuration import (
     StepEntrypointConfiguration,
@@ -30,14 +30,14 @@ class AzureMLEntrypointConfiguration(StepEntrypointConfiguration):
     """Entrypoint configuration for ZenML AzureML pipeline steps."""
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the
             environmental variables.
         """
-        return super().get_entrypoint_options() | {ZENML_ENV_VARIABLES}
+        return super().get_entrypoint_options() | {ZENML_ENV_VARIABLES: True}
 
     @classmethod
     def get_entrypoint_arguments(cls, **kwargs: Any) -> List[str]:

src/zenml/integrations/databricks/orchestrators/databricks_orchestrator_entrypoint_config.py

@@ -16,7 +16,7 @@
 import os
 import sys
 from importlib.metadata import distribution
-from typing import Any, List, Set
+from typing import Any, Dict, List
 
 from zenml.entrypoints.step_entrypoint_configuration import (
     StepEntrypointConfiguration,
@@ -38,17 +38,16 @@ class DatabricksEntrypointConfiguration(StepEntrypointConfiguration):
     """
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the wheel package.
         """
-        return (
-            super().get_entrypoint_options()
-            | {WHEEL_PACKAGE_OPTION}
-            | {DATABRICKS_JOB_ID_OPTION}
-        )
+        return super().get_entrypoint_options() | {
+            WHEEL_PACKAGE_OPTION: True,
+            DATABRICKS_JOB_ID_OPTION: True,
+        }
 
     @classmethod
     def get_entrypoint_arguments(

src/zenml/pipelines/dynamic/entrypoint_configuration.py

@@ -27,7 +27,7 @@
 #  permissions and limitations under the License.
 """Abstract base class for entrypoint configurations that run a pipeline."""
 
-from typing import Any, List, Set
+from typing import Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
@@ -44,14 +44,14 @@ class DynamicPipelineEntrypointConfiguration(BaseEntrypointConfiguration):
     """Base class for entrypoint configurations that run an entire pipeline."""
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the name of the
             step to run.
         """
-        return super().get_entrypoint_options() | {RUN_ID_OPTION}
+        return super().get_entrypoint_options() | {RUN_ID_OPTION: False}
 
     @classmethod
     def get_entrypoint_arguments(
@@ -73,10 +73,15 @@ def get_entrypoint_arguments(
             The superclass arguments as well as arguments for the name of the
             step to run.
         """
-        return super().get_entrypoint_arguments(**kwargs) + [
-            f"--{RUN_ID_OPTION}",
-            str(kwargs[RUN_ID_OPTION]),
-        ]
+        args = super().get_entrypoint_arguments(**kwargs)
+        if RUN_ID_OPTION in kwargs:
+            args.extend(
+                [
+                    f"--{RUN_ID_OPTION}",
+                    str(kwargs[RUN_ID_OPTION]),
+                ]
+            )
+        return args
 
     def run(self) -> None:
         """Prepares the environment and runs the configured pipeline."""
@@ -88,9 +93,9 @@ def run(self) -> None:
 
         self.download_code_if_necessary(snapshot=snapshot)
 
-        # TODO: make this optional
-        run_id = UUID(self.entrypoint_args[RUN_ID_OPTION])
-        run = Client().get_pipeline_run(run_id)
+        run = None
+        if run_id := self.entrypoint_args.get(RUN_ID_OPTION, None):
+            run = Client().get_pipeline_run(UUID(run_id))
 
         runner = DynamicPipelineRunner(snapshot=snapshot, run=run)
         runner.run_pipeline()

src/zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -13,7 +13,8 @@
 #  permissions and limitations under the License.
 """Abstract base class for entrypoint configurations that run a single step."""
 
-from typing import TYPE_CHECKING, Any, List, Set
+from ast import Dict
+from typing import TYPE_CHECKING, Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
@@ -36,14 +37,14 @@ class StepOperatorEntrypointConfiguration(StepEntrypointConfiguration):
     """Base class for step operator entrypoint configurations."""
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the step run id.
         """
         return super().get_entrypoint_options() | {
-            STEP_RUN_ID_OPTION,
+            STEP_RUN_ID_OPTION: True,
         }
 
     @classmethod

src/zenml/zen_server/pipeline_execution/runner_entrypoint_configuration.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Runner entrypoint configuration."""
 
-from typing import Any, List, Set
+from typing import Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
@@ -29,14 +29,16 @@ class RunnerEntrypointConfiguration(BaseEntrypointConfiguration):
     """Runner entrypoint configuration."""
 
     @classmethod
-    def get_entrypoint_options(cls) -> Set[str]:
+    def get_entrypoint_options(cls) -> Dict[str, bool]:
         """Gets all options required for running with this configuration.
 
         Returns:
             The superclass options as well as an option for the name of the
             step to run.
         """
-        return super().get_entrypoint_options() | {PLACEHOLDER_RUN_ID_OPTION}
+        return super().get_entrypoint_options() | {
+            PLACEHOLDER_RUN_ID_OPTION: True
+        }
 
     @classmethod
     def get_entrypoint_arguments(