Entrypoints
zenml.entrypoints
special
Initializations for ZenML entrypoints module.
step_entrypoint
Functionality to run ZenML steps.
main()
Runs the ZenML step defined by the command line arguments.
This main logic for running the step is not implemented in this file,
instead it simply creates an object of a
zenml.entrypoints.StepEntrypointConfiguration subclass (the concrete
implementation can be specified using the command line arguments) and calls
its run() method.
Exceptions:
| Type | Description |
|---|---|
TypeError |
If the command line arguments are invalid. |
Source code in zenml/entrypoints/step_entrypoint.py
def main() -> None:
"""Runs the ZenML step defined by the command line arguments.
This main logic for running the step is not implemented in this file,
instead it simply creates an object of a
`zenml.entrypoints.StepEntrypointConfiguration` subclass (the concrete
implementation can be specified using the command line arguments) and calls
its `run()` method.
Raises:
TypeError: If the command line arguments are invalid.
"""
# Read the source for the entrypoint configuration class from the command
# line arguments
parser = argparse.ArgumentParser()
parser.add_argument(f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}", required=True)
args, remaining_args = parser.parse_known_args()
# Create an instance of the entrypoint configuration and pass it the
# remaining command line arguments
entrypoint_config_class = source_utils.load_source_path_class(
args.entrypoint_config_source
)
if not issubclass(entrypoint_config_class, StepEntrypointConfiguration):
raise TypeError(
f"The entrypoint config source `{args.entrypoint_config_source}` "
f"passed to the entrypoint is not pointing to a "
f"`{StepEntrypointConfiguration}` subclass."
)
entrypoint_config = entrypoint_config_class(arguments=remaining_args)
# Run the entrypoint configuration
entrypoint_config.run()
step_entrypoint_configuration
Abstract base class for entrypoint configurations that run a single step.
StepEntrypointConfiguration (ABC)
Abstract base class for entrypoint configurations that run a single step.
If an orchestrator needs to run steps in a separate process or environment (e.g. a docker container), you should create a custom entrypoint configuration class that inherits from this class and use it to implement your custom entrypoint logic.
How to subclass:
There is only one mandatory method get_run_name(...) that you need to
implement in order to get a functioning entrypoint. Inside this method you
need to return a string which has to be the same for all steps that are
executed as part of the same pipeline run.
Passing additional arguments to the entrypoint:
If you need to pass additional arguments to the entrypoint, there are
two methods that you need to implement:
* get_custom_entrypoint_options(): This method should return all
the additional options that you require in the entrypoint.
* `get_custom_entrypoint_arguments(...)`: This method should return
a list of arguments that should be passed to the entrypoint.
The arguments need to provide values for all options defined
in the `custom_entrypoint_options()` method mentioned above.
You'll be able to access the argument values from `self.entrypoint_args`
inside your `StepEntrypointConfiguration` subclass.
Running custom code inside the entrypoint:
If you need to run custom code in the entrypoint, you can overwrite
the setup(...) and post_run(...) methods which allow you to run
code before and after the step execution respectively.
How to use:
After you created your StepEntrypointConfiguration subclass, you only
have to run the entrypoint somewhere. To do this, you should execute the
command returned by the get_entrypoint_command() method with the
arguments returned by the get_entrypoint_arguments(...) method.
Examples:
class MyStepEntrypointConfiguration(StepEntrypointConfiguration):
...
class MyOrchestrator(BaseOrchestrator):
def prepare_or_run_pipeline(
self,
sorted_list_of_steps: List[BaseStep],
pipeline: "BasePipeline",
pb2_pipeline: Pb2Pipeline,
stack: "Stack",
runtime_configuration: "RuntimeConfiguration",
) -> Any:
...
cmd = MyStepEntrypointConfiguration.get_entrypoint_command()
for step in sorted_list_of_steps:
...
args = MyStepEntrypointConfiguration.get_entrypoint_arguments(
step=step, pb2_pipeline=pb2_pipeline
)
# Run the command and pass it the arguments. Our example
# orchestrator here executes the entrypoint in a separate
# process, but in a real-world scenario you would probably run
# it inside a docker container or a different environment.
import subprocess
subprocess.check_call(cmd + args)
Source code in zenml/entrypoints/step_entrypoint_configuration.py
class StepEntrypointConfiguration(ABC):
"""Abstract base class for entrypoint configurations that run a single step.
If an orchestrator needs to run steps in a separate process or environment
(e.g. a docker container), you should create a custom entrypoint
configuration class that inherits from this class and use it to implement
your custom entrypoint logic.
How to subclass:
----------------
There is only one mandatory method `get_run_name(...)` that you need to
implement in order to get a functioning entrypoint. Inside this method you
need to return a string which **has** to be the same for all steps that are
executed as part of the same pipeline run.
Passing additional arguments to the entrypoint:
If you need to pass additional arguments to the entrypoint, there are
two methods that you need to implement:
* `get_custom_entrypoint_options()`: This method should return all
the additional options that you require in the entrypoint.
* `get_custom_entrypoint_arguments(...)`: This method should return
a list of arguments that should be passed to the entrypoint.
The arguments need to provide values for all options defined
in the `custom_entrypoint_options()` method mentioned above.
You'll be able to access the argument values from `self.entrypoint_args`
inside your `StepEntrypointConfiguration` subclass.
Running custom code inside the entrypoint:
If you need to run custom code in the entrypoint, you can overwrite
the `setup(...)` and `post_run(...)` methods which allow you to run
code before and after the step execution respectively.
How to use:
-----------
After you created your `StepEntrypointConfiguration` subclass, you only
have to run the entrypoint somewhere. To do this, you should execute the
command returned by the `get_entrypoint_command()` method with the
arguments returned by the `get_entrypoint_arguments(...)` method.
Example:
```python
class MyStepEntrypointConfiguration(StepEntrypointConfiguration):
...
class MyOrchestrator(BaseOrchestrator):
def prepare_or_run_pipeline(
self,
sorted_list_of_steps: List[BaseStep],
pipeline: "BasePipeline",
pb2_pipeline: Pb2Pipeline,
stack: "Stack",
runtime_configuration: "RuntimeConfiguration",
) -> Any:
...
cmd = MyStepEntrypointConfiguration.get_entrypoint_command()
for step in sorted_list_of_steps:
...
args = MyStepEntrypointConfiguration.get_entrypoint_arguments(
step=step, pb2_pipeline=pb2_pipeline
)
# Run the command and pass it the arguments. Our example
# orchestrator here executes the entrypoint in a separate
# process, but in a real-world scenario you would probably run
# it inside a docker container or a different environment.
import subprocess
subprocess.check_call(cmd + args)
```
"""
def __init__(self, arguments: List[str]):
"""Initializes the entrypoint configuration.
Args:
arguments: Command line arguments to configure this object.
"""
self.entrypoint_args = self._parse_arguments(arguments)
@classmethod
def get_custom_entrypoint_options(cls) -> Set[str]:
"""Custom options for this entrypoint in addition to what's needed by default.
The options returned by this method should be strings like "my_option"
(no "--" prefix). When the entrypoint is executed, it will parse the
command line arguments and expect all options to be passed in the form
"--my_option my_value". You'll be able to retrieve the argument
value by calling `self.entrypoint_args["my_option"]`.
Returns:
A set of strings with the custom options.
"""
return set()
@classmethod
def get_custom_entrypoint_arguments(
cls, step: BaseStep, **kwargs: Any
) -> List[str]:
"""Custom arguments that the entrypoint command should be called with.
The argument list should be something that
`argparse.ArgumentParser.parse_args(...)` can handle (e.g.
`["--some_option", "some_value"]` or `["--some_option=some_value"]`).
It needs to provide values for all options returned by the
`get_custom_entrypoint_options()` method of this class.
Args:
step: The step that the entrypoint should run using the arguments
returned by this method.
**kwargs: Additional keyword arguments that are passed to the
`get_custom_entrypoint_options()` method.
Returns:
A list of strings with the arguments.
"""
return []
@abstractmethod
def get_run_name(self, pipeline_name: str) -> str:
"""Returns the run name.
Subclasses must implement this and return a run name that is the same
for all steps of this pipeline run.
Examples:
* If you're in an orchestrator environment which has an equivalent
concept to a pipeline run, you can use that as the run name. E.g.
Kubeflow Pipelines sets a run id as environment variable in all
pods which we can reuse: `return os.environ["KFP_RUN_ID"]`
* If that isn't possible, you could pass a unique value as an argument
to the entrypoint (make sure to also return it from
`get_custom_entrypoint_options()`) and use it like this:
`return self.entrypoint_args["run_name_option"]`
Args:
pipeline_name: The name of the pipeline.
Returns:
The run name.
"""
def setup(self, pipeline_name: str, step_name: str) -> None:
"""Code that needs to run before user code is imported or the step is executed.
Subclasses should overwrite this method if they need to perform any
additional setup, but should in most cases include a
`super().setup(...)` call so the ZenML setup code will still be
executed.
Args:
pipeline_name: Name of the parent pipeline of the step that will
be executed.
step_name: Name of the step that will be executed.
"""
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
def post_run(
self,
pipeline_name: str,
step_name: str,
pipeline_node: Pb2PipelineNode,
execution_info: Optional[data_types.ExecutionInfo] = None,
) -> None:
"""Does cleanup or post-processing after the step finished running.
Subclasses should overwrite this method if they need to run any
additional code after the step execution.
Args:
pipeline_name: Name of the parent pipeline of the step that was
executed.
step_name: Name of the step that was executed.
pipeline_node: Protobuf node corresponding to the step that was
executed.
execution_info: Info about the finished step execution.
"""
@classmethod
def get_entrypoint_command(cls) -> List[str]:
"""Returns a command that runs the entrypoint module.
This entrypoint module must execute a ZenML step when called. If
subclasses don't overwrite this method, it will default to running the
`zenml.entrypoints.step_entrypoint` module.
**Note**: This command won't work on its own but needs to be called with
the arguments returned by the `get_entrypoint_arguments(...)`
method of this class.
Returns:
A list of strings with the command.
"""
return DEFAULT_SINGLE_STEP_CONTAINER_ENTRYPOINT_COMMAND
@classmethod
def get_entrypoint_options(cls) -> Set[str]:
"""Gets all options required for running an entrypoint with this configuration.
**Note**: Subclasses should implement the
`get_custom_entrypoint_options()` class method instead of this
one if they require custom options.
Returns:
A set of strings with all required options.
"""
zenml_options = {
# Importable source pointing to the entrypoint configuration class
# that should be used inside the entrypoint.
ENTRYPOINT_CONFIG_SOURCE_OPTION,
# Importable source pointing to the python module that was executed
# to run a pipeline. This will be imported inside the entrypoint to
# make sure all custom materializer/artifact types are registered.
MAIN_MODULE_SOURCE_OPTION,
# Importable source pointing to the ZenML step class that should be
# run inside the entrypoint. This will be used to recreate the tfx
# executor class that is required to run the step function.
STEP_SOURCE_OPTION,
# Base64 encoded json representation of the parent pipeline of
# the step that will be executed. This is needed in order to create
# the tfx launcher in the entrypoint that will run the ZenML step.
PIPELINE_JSON_OPTION,
# Base64 encoded json dictionary mapping the step input names to
# importable sources pointing to
# `zenml.artifacts.base_artifact.BaseArtifact` subclasses.
# These classes are needed to recreate the tfx executor
# class and can't be inferred as we do not have access to the
# output artifacts from previous steps.
INPUT_ARTIFACT_SOURCES_OPTION,
# Base64 encoded json dictionary mapping the step output names to
# importable sources pointing to
# `zenml.materializers.base_materializer.BaseMaterializer`
# subclasses. These classes are needed to recreate the tfx executor
# class if custom materializers were specified for the step.
MATERIALIZER_SOURCES_OPTION,
}
custom_options = cls.get_custom_entrypoint_options()
return zenml_options.union(custom_options)
@classmethod
def get_entrypoint_arguments(
cls,
step: BaseStep,
pb2_pipeline: Pb2Pipeline,
**kwargs: Any,
) -> List[str]:
"""Gets all arguments that the entrypoint command should be called with.
The argument list should be something that
`argparse.ArgumentParser.parse_args(...)` can handle (e.g.
`["--some_option", "some_value"]` or `["--some_option=some_value"]`).
It needs to provide values for all options returned by the
`get_entrypoint_options()` method of this class.
**Note**: Subclasses should implement the
`get_custom_entrypoint_arguments(...)` class method instead of
this one if they require custom arguments.
Args:
step: The step that the entrypoint should run using the arguments
returned by this method.
pb2_pipeline: The protobuf representation of the pipeline to which
the `step` belongs.
**kwargs: Custom options that will be passed to
`get_custom_entrypoint_arguments()`.
Returns:
A list of strings with the arguments.
"""
# Get an importable source of the user main module. If the `__main__`
# module is not the actual module that a user executed (e.g. if we're
# in some docker container entrypoint) use the ZenML constant that
# points to the actual user main module. If not resolve the `__main__`
# module to something that can be imported during entrypoint execution.
main_module_source = (
constants.USER_MAIN_MODULE
or source_utils.get_module_source_from_module(
sys.modules["__main__"]
)
)
# TODO [ENG-887]: Move this method to source utils
def _resolve_class(class_: Type[Any]) -> str:
"""Resolves the input class so it is importable inside the entrypoint.
Args:
class_: The class to resolve.
Returns:
The importable source of the class.
"""
source = source_utils.resolve_class(class_)
module_source, class_source = source.rsplit(".", 1)
if module_source == "__main__":
# If the class is defined inside the `__main__` module it will
# not be importable in the entrypoint (The `__main__` module
# there will be the actual entrypoint module). We therefore
# replace the module source by the importable main module
# source computed above.
module_source = main_module_source
return f"{module_source}.{class_source}"
# Resolve the entrypoint config source
entrypoint_config_source = _resolve_class(cls)
# Resolve the step class
step_source = _resolve_class(step.__class__)
# Resolve the input artifact classes
input_artifact_sources = {
input_name: _resolve_class(input_type)
for input_name, input_type in step.INPUT_SPEC.items()
}
# Resolve the output materializer classes
materializer_classes = step.get_materializers(ensure_complete=True)
materializer_sources = {
output_name: _resolve_class(materializer_class)
for output_name, materializer_class in materializer_classes.items()
}
# See `get_entrypoint_options()` for an in -depth explanation of all
# these arguments.
zenml_arguments = [
f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}",
entrypoint_config_source,
f"--{MAIN_MODULE_SOURCE_OPTION}",
main_module_source,
f"--{STEP_SOURCE_OPTION}",
step_source,
# Base64 encode the json strings to make sure there are no issues
# when passing these arguments
f"--{PIPELINE_JSON_OPTION}",
string_utils.b64_encode(json_format.MessageToJson(pb2_pipeline)),
f"--{INPUT_ARTIFACT_SOURCES_OPTION}",
string_utils.b64_encode(json.dumps(input_artifact_sources)),
f"--{MATERIALIZER_SOURCES_OPTION}",
string_utils.b64_encode(json.dumps(materializer_sources)),
]
custom_arguments = cls.get_custom_entrypoint_arguments(
step=step, **kwargs
)
all_arguments = zenml_arguments + custom_arguments
# Try to parse the arguments here to validate the format and that
# values for all options are provided.
cls._parse_arguments(all_arguments)
return all_arguments
@classmethod
def _parse_arguments(cls, arguments: List[str]) -> Dict[str, Any]:
"""Parses command line arguments.
This method will create an `argparse.ArgumentParser` and add required
arguments for all the options specified in the
`get_entrypoint_options()` method of this class. Subclasses that
require additional arguments during entrypoint execution should provide
them by implementing the `get_custom_entrypoint_options()` class method.
Args:
arguments: Arguments to parse. The format should be something that
`argparse.ArgumentParser.parse_args(...)` can handle (e.g.
`["--some_option", "some_value"]` or
`["--some_option=some_value"]`).
Returns:
Dictionary of the parsed arguments.
# noqa: DAR402
Raises:
ValueError: If the arguments are not valid.
"""
# Argument parser subclass that suppresses some argparse logs and
# raises an exception instead of the `sys.exit()` call
class _CustomParser(argparse.ArgumentParser):
def error(self, message: str) -> NoReturn:
raise ValueError(
f"Failed to parse entrypoint arguments: {message}"
)
parser = _CustomParser()
for option_name in cls.get_entrypoint_options():
if option_name == ENTRYPOINT_CONFIG_SOURCE_OPTION:
# This option is already used by
# `zenml.entrypoints.step_entrypoint` to read which config
# class to use
continue
parser.add_argument(f"--{option_name}", required=True)
result, _ = parser.parse_known_args(arguments)
return vars(result)
@staticmethod
def _create_executor_class(
step: BaseStep,
input_artifact_sources: Dict[str, str],
materializer_sources: Dict[str, str],
) -> None:
"""Creates an executor class for the given step.
Args:
step: The step for which the executor class should be created.
input_artifact_sources: Dictionary mapping step input names to a
source strings of the artifact class to use for that input.
materializer_sources: Dictionary mapping step output names to a
source string of the materializer class to use for that output.
Raises:
TypeError: If the step does not have a valid input or output
specification.
"""
# Import the input artifact classes from the given sources
input_spec = {}
for input_name, source in input_artifact_sources.items():
artifact_class = source_utils.load_source_path_class(source)
if not issubclass(artifact_class, BaseArtifact):
raise TypeError(
f"The artifact source `{source}` passed to the entrypoint "
f"for the step input '{input_name}' is not pointing to a "
f"`{BaseArtifact}` subclass."
)
input_spec[input_name] = artifact_class
output_spec = {
key: type_registry.get_artifact_type(value)[0]
for key, value in step.OUTPUT_SIGNATURE.items()
}
execution_parameters = {
**step.PARAM_SPEC,
**step._internal_execution_parameters,
}
materializers = {}
for output_name, source in materializer_sources.items():
materializer_class = source_utils.load_source_path_class(source)
if not issubclass(materializer_class, BaseMaterializer):
raise TypeError(
f"The materializer source `{source}` passed to the "
f"entrypoint for the step output '{output_name}' is not "
f"pointing to a `{BaseMaterializer}` subclass."
)
materializers[output_name] = materializer_class
step_utils.generate_component_class(
step_name=step.name,
step_module=step.__module__,
input_spec=input_spec,
output_spec=output_spec,
execution_parameter_names=set(execution_parameters),
step_function=step.entrypoint,
materializers=materializers,
)
def run(self) -> None:
"""Runs a single ZenML step.
Subclasses should in most cases not need to overwrite this method and
implement their custom logic in the `setup(...)` and `post_run(...)`
methods instead. If you still need to customize the functionality of
this method, make sure to still include all the existing logic as your
step won't be executed properly otherwise.
Raises:
TypeError: If the arguments passed to the entrypoint are invalid.
"""
# Make sure this entrypoint does not run an entire pipeline when
# importing user modules. This could happen if the `pipeline.run()` call
# is not wrapped in a function or an `if __name__== "__main__":` check)
constants.SHOULD_PREVENT_PIPELINE_EXECUTION = True
# Extract and parse all the entrypoint arguments required to execute
# the step. See `get_entrypoint_options()` for an in-depth explanation
# of all these arguments.
pb2_pipeline = Pb2Pipeline()
pb2_pipeline_json = string_utils.b64_decode(
self.entrypoint_args[PIPELINE_JSON_OPTION]
)
json_format.Parse(pb2_pipeline_json, pb2_pipeline)
main_module_source = self.entrypoint_args[MAIN_MODULE_SOURCE_OPTION]
step_source = self.entrypoint_args[STEP_SOURCE_OPTION]
input_artifact_sources = json.loads(
string_utils.b64_decode(
self.entrypoint_args[INPUT_ARTIFACT_SOURCES_OPTION]
)
)
materializer_sources = json.loads(
string_utils.b64_decode(
self.entrypoint_args[MATERIALIZER_SOURCES_OPTION]
)
)
# Get some common values that will be used throughout the remainder of
# this method
pipeline_name = pb2_pipeline.pipeline_info.id
if "@" in step_source:
# Get rid of potential ZenML version pins if the source looks like
# this `zenml.some_module.SomeClass@zenml_0.9.9`
step_source, _ = step_source.split("@", 1)
step_module, step_name = step_source.rsplit(".", 1)
# Allow subclasses to run custom code before user code is imported and
# the step is executed.
self.setup(pipeline_name=pipeline_name, step_name=step_name)
# Activate all the integrations. This makes sure that all materializers
# and stack component flavors are registered.
integration_registry.activate_integrations()
# Import the main module that was executed to run the pipeline to which
# the step getting executed belongs. Even if the step class is not
# defined in this module, we need to import it in case the user
# defined/imported custom materializers/artifacts here that ZenML needs
# to execute the step.
importlib.import_module(main_module_source)
# The `__main__` module when running a step inside a docker container
# will always be the entrypoint file
# (e.g. `zenml.entrypoints.single_step_entrypoint`). If this entrypoint
# now needs to execute something in a different environment,
# `sys.modules["__main__"]` would not point to the original user main
# module anymore. That's why we set this constant here so other
# components (e.g. step operators) can use it to correctly import it
# in the containers they create.
constants.USER_MAIN_MODULE = main_module_source
# Create an instance of the ZenML step that this entrypoint should run.
step_class = source_utils.load_source_path_class(step_source)
if not issubclass(step_class, BaseStep):
raise TypeError(
f"The step source `{step_source}` passed to the "
f"entrypoint is not pointing to a `{BaseStep}` subclass."
)
step = step_class()
# Compute the name of the module in which the step was in when the
# pipeline run was started. This could potentially differ from the
# module part of `STEP_SOURCE_OPTION` if the step was defined in
# the `__main__` module (e.g. the step was defined in `run.py` and the
# pipeline run was started by the call `python run.py`).
# Why we need this: The executor source is stored inside the protobuf
# pipeline representation object `pb2_pipeline` which was passed to
# this entrypoint. Tfx will try to load the executor class from that
# source, so we need to recreate the executor class in the same module.
original_step_module = (
"__main__" if step_module == main_module_source else step_module
)
# Make sure the `__module__` attribute of our step instance points to
# this original module so the executor class gets created in the
# correct module.
step.__module__ = original_step_module
# Create the executor class that is responsible for running the step
# function. This method call dynamically creates an executor class
# which will later be used when `orchestrator.run_step(...)` is running
# the step.
self._create_executor_class(
step=step,
input_artifact_sources=input_artifact_sources,
materializer_sources=materializer_sources,
)
# Execute the actual step code.
run_name = self.get_run_name(pipeline_name=pipeline_name)
orchestrator = Repository().active_stack.orchestrator
execution_info = orchestrator.run_step(
step=step, run_name=run_name, pb2_pipeline=pb2_pipeline
)
# Allow subclasses to run custom code after the step finished executing.
pipeline_node = orchestrator._get_node_with_step_name(
step_name=step_name, pb2_pipeline=pb2_pipeline
)
self.post_run(
pipeline_name=pipeline_name,
step_name=step_name,
pipeline_node=pipeline_node,
execution_info=execution_info,
)
__init__(self, arguments)
special
Initializes the entrypoint configuration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
arguments |
List[str] |
Command line arguments to configure this object. |
required |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
def __init__(self, arguments: List[str]):
"""Initializes the entrypoint configuration.
Args:
arguments: Command line arguments to configure this object.
"""
self.entrypoint_args = self._parse_arguments(arguments)
get_custom_entrypoint_arguments(step, **kwargs)
classmethod
Custom arguments that the entrypoint command should be called with.
The argument list should be something that
argparse.ArgumentParser.parse_args(...) can handle (e.g.
["--some_option", "some_value"] or ["--some_option=some_value"]).
It needs to provide values for all options returned by the
get_custom_entrypoint_options() method of this class.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
step |
BaseStep |
The step that the entrypoint should run using the arguments returned by this method. |
required |
**kwargs |
Any |
Additional keyword arguments that are passed to the
|
{} |
Returns:
| Type | Description |
|---|---|
List[str] |
A list of strings with the arguments. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@classmethod
def get_custom_entrypoint_arguments(
cls, step: BaseStep, **kwargs: Any
) -> List[str]:
"""Custom arguments that the entrypoint command should be called with.
The argument list should be something that
`argparse.ArgumentParser.parse_args(...)` can handle (e.g.
`["--some_option", "some_value"]` or `["--some_option=some_value"]`).
It needs to provide values for all options returned by the
`get_custom_entrypoint_options()` method of this class.
Args:
step: The step that the entrypoint should run using the arguments
returned by this method.
**kwargs: Additional keyword arguments that are passed to the
`get_custom_entrypoint_options()` method.
Returns:
A list of strings with the arguments.
"""
return []
get_custom_entrypoint_options()
classmethod
Custom options for this entrypoint in addition to what's needed by default.
The options returned by this method should be strings like "my_option"
(no "--" prefix). When the entrypoint is executed, it will parse the
command line arguments and expect all options to be passed in the form
"--my_option my_value". You'll be able to retrieve the argument
value by calling self.entrypoint_args["my_option"].
Returns:
| Type | Description |
|---|---|
Set[str] |
A set of strings with the custom options. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@classmethod
def get_custom_entrypoint_options(cls) -> Set[str]:
"""Custom options for this entrypoint in addition to what's needed by default.
The options returned by this method should be strings like "my_option"
(no "--" prefix). When the entrypoint is executed, it will parse the
command line arguments and expect all options to be passed in the form
"--my_option my_value". You'll be able to retrieve the argument
value by calling `self.entrypoint_args["my_option"]`.
Returns:
A set of strings with the custom options.
"""
return set()
get_entrypoint_arguments(step, pb2_pipeline, **kwargs)
classmethod
Gets all arguments that the entrypoint command should be called with.
The argument list should be something that
argparse.ArgumentParser.parse_args(...) can handle (e.g.
["--some_option", "some_value"] or ["--some_option=some_value"]).
It needs to provide values for all options returned by the
get_entrypoint_options() method of this class.
Note: Subclasses should implement the
get_custom_entrypoint_arguments(...) class method instead of
this one if they require custom arguments.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
step |
BaseStep |
The step that the entrypoint should run using the arguments returned by this method. |
required |
pb2_pipeline |
Pipeline |
The protobuf representation of the pipeline to which
the |
required |
**kwargs |
Any |
Custom options that will be passed to
|
{} |
Returns:
| Type | Description |
|---|---|
List[str] |
A list of strings with the arguments. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@classmethod
def get_entrypoint_arguments(
cls,
step: BaseStep,
pb2_pipeline: Pb2Pipeline,
**kwargs: Any,
) -> List[str]:
"""Gets all arguments that the entrypoint command should be called with.
The argument list should be something that
`argparse.ArgumentParser.parse_args(...)` can handle (e.g.
`["--some_option", "some_value"]` or `["--some_option=some_value"]`).
It needs to provide values for all options returned by the
`get_entrypoint_options()` method of this class.
**Note**: Subclasses should implement the
`get_custom_entrypoint_arguments(...)` class method instead of
this one if they require custom arguments.
Args:
step: The step that the entrypoint should run using the arguments
returned by this method.
pb2_pipeline: The protobuf representation of the pipeline to which
the `step` belongs.
**kwargs: Custom options that will be passed to
`get_custom_entrypoint_arguments()`.
Returns:
A list of strings with the arguments.
"""
# Get an importable source of the user main module. If the `__main__`
# module is not the actual module that a user executed (e.g. if we're
# in some docker container entrypoint) use the ZenML constant that
# points to the actual user main module. If not resolve the `__main__`
# module to something that can be imported during entrypoint execution.
main_module_source = (
constants.USER_MAIN_MODULE
or source_utils.get_module_source_from_module(
sys.modules["__main__"]
)
)
# TODO [ENG-887]: Move this method to source utils
def _resolve_class(class_: Type[Any]) -> str:
"""Resolves the input class so it is importable inside the entrypoint.
Args:
class_: The class to resolve.
Returns:
The importable source of the class.
"""
source = source_utils.resolve_class(class_)
module_source, class_source = source.rsplit(".", 1)
if module_source == "__main__":
# If the class is defined inside the `__main__` module it will
# not be importable in the entrypoint (The `__main__` module
# there will be the actual entrypoint module). We therefore
# replace the module source by the importable main module
# source computed above.
module_source = main_module_source
return f"{module_source}.{class_source}"
# Resolve the entrypoint config source
entrypoint_config_source = _resolve_class(cls)
# Resolve the step class
step_source = _resolve_class(step.__class__)
# Resolve the input artifact classes
input_artifact_sources = {
input_name: _resolve_class(input_type)
for input_name, input_type in step.INPUT_SPEC.items()
}
# Resolve the output materializer classes
materializer_classes = step.get_materializers(ensure_complete=True)
materializer_sources = {
output_name: _resolve_class(materializer_class)
for output_name, materializer_class in materializer_classes.items()
}
# See `get_entrypoint_options()` for an in -depth explanation of all
# these arguments.
zenml_arguments = [
f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}",
entrypoint_config_source,
f"--{MAIN_MODULE_SOURCE_OPTION}",
main_module_source,
f"--{STEP_SOURCE_OPTION}",
step_source,
# Base64 encode the json strings to make sure there are no issues
# when passing these arguments
f"--{PIPELINE_JSON_OPTION}",
string_utils.b64_encode(json_format.MessageToJson(pb2_pipeline)),
f"--{INPUT_ARTIFACT_SOURCES_OPTION}",
string_utils.b64_encode(json.dumps(input_artifact_sources)),
f"--{MATERIALIZER_SOURCES_OPTION}",
string_utils.b64_encode(json.dumps(materializer_sources)),
]
custom_arguments = cls.get_custom_entrypoint_arguments(
step=step, **kwargs
)
all_arguments = zenml_arguments + custom_arguments
# Try to parse the arguments here to validate the format and that
# values for all options are provided.
cls._parse_arguments(all_arguments)
return all_arguments
get_entrypoint_command()
classmethod
Returns a command that runs the entrypoint module.
This entrypoint module must execute a ZenML step when called. If
subclasses don't overwrite this method, it will default to running the
zenml.entrypoints.step_entrypoint module.
Note: This command won't work on its own but needs to be called with
the arguments returned by the get_entrypoint_arguments(...)
method of this class.
Returns:
| Type | Description |
|---|---|
List[str] |
A list of strings with the command. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@classmethod
def get_entrypoint_command(cls) -> List[str]:
"""Returns a command that runs the entrypoint module.
This entrypoint module must execute a ZenML step when called. If
subclasses don't overwrite this method, it will default to running the
`zenml.entrypoints.step_entrypoint` module.
**Note**: This command won't work on its own but needs to be called with
the arguments returned by the `get_entrypoint_arguments(...)`
method of this class.
Returns:
A list of strings with the command.
"""
return DEFAULT_SINGLE_STEP_CONTAINER_ENTRYPOINT_COMMAND
get_entrypoint_options()
classmethod
Gets all options required for running an entrypoint with this configuration.
Note: Subclasses should implement the
get_custom_entrypoint_options() class method instead of this
one if they require custom options.
Returns:
| Type | Description |
|---|---|
Set[str] |
A set of strings with all required options. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@classmethod
def get_entrypoint_options(cls) -> Set[str]:
"""Gets all options required for running an entrypoint with this configuration.
**Note**: Subclasses should implement the
`get_custom_entrypoint_options()` class method instead of this
one if they require custom options.
Returns:
A set of strings with all required options.
"""
zenml_options = {
# Importable source pointing to the entrypoint configuration class
# that should be used inside the entrypoint.
ENTRYPOINT_CONFIG_SOURCE_OPTION,
# Importable source pointing to the python module that was executed
# to run a pipeline. This will be imported inside the entrypoint to
# make sure all custom materializer/artifact types are registered.
MAIN_MODULE_SOURCE_OPTION,
# Importable source pointing to the ZenML step class that should be
# run inside the entrypoint. This will be used to recreate the tfx
# executor class that is required to run the step function.
STEP_SOURCE_OPTION,
# Base64 encoded json representation of the parent pipeline of
# the step that will be executed. This is needed in order to create
# the tfx launcher in the entrypoint that will run the ZenML step.
PIPELINE_JSON_OPTION,
# Base64 encoded json dictionary mapping the step input names to
# importable sources pointing to
# `zenml.artifacts.base_artifact.BaseArtifact` subclasses.
# These classes are needed to recreate the tfx executor
# class and can't be inferred as we do not have access to the
# output artifacts from previous steps.
INPUT_ARTIFACT_SOURCES_OPTION,
# Base64 encoded json dictionary mapping the step output names to
# importable sources pointing to
# `zenml.materializers.base_materializer.BaseMaterializer`
# subclasses. These classes are needed to recreate the tfx executor
# class if custom materializers were specified for the step.
MATERIALIZER_SOURCES_OPTION,
}
custom_options = cls.get_custom_entrypoint_options()
return zenml_options.union(custom_options)
get_run_name(self, pipeline_name)
Returns the run name.
Subclasses must implement this and return a run name that is the same for all steps of this pipeline run.
Examples:
- If you're in an orchestrator environment which has an equivalent
concept to a pipeline run, you can use that as the run name. E.g.
Kubeflow Pipelines sets a run id as environment variable in all
pods which we can reuse:
return os.environ["KFP_RUN_ID"] - If that isn't possible, you could pass a unique value as an argument
to the entrypoint (make sure to also return it from
get_custom_entrypoint_options()) and use it like this:return self.entrypoint_args["run_name_option"]
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
pipeline_name |
str |
The name of the pipeline. |
required |
Returns:
| Type | Description |
|---|---|
str |
The run name. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
@abstractmethod
def get_run_name(self, pipeline_name: str) -> str:
"""Returns the run name.
Subclasses must implement this and return a run name that is the same
for all steps of this pipeline run.
Examples:
* If you're in an orchestrator environment which has an equivalent
concept to a pipeline run, you can use that as the run name. E.g.
Kubeflow Pipelines sets a run id as environment variable in all
pods which we can reuse: `return os.environ["KFP_RUN_ID"]`
* If that isn't possible, you could pass a unique value as an argument
to the entrypoint (make sure to also return it from
`get_custom_entrypoint_options()`) and use it like this:
`return self.entrypoint_args["run_name_option"]`
Args:
pipeline_name: The name of the pipeline.
Returns:
The run name.
"""
post_run(self, pipeline_name, step_name, pipeline_node, execution_info=None)
Does cleanup or post-processing after the step finished running.
Subclasses should overwrite this method if they need to run any additional code after the step execution.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
pipeline_name |
str |
Name of the parent pipeline of the step that was executed. |
required |
step_name |
str |
Name of the step that was executed. |
required |
pipeline_node |
PipelineNode |
Protobuf node corresponding to the step that was executed. |
required |
execution_info |
Optional[tfx.orchestration.portable.data_types.ExecutionInfo] |
Info about the finished step execution. |
None |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
def post_run(
self,
pipeline_name: str,
step_name: str,
pipeline_node: Pb2PipelineNode,
execution_info: Optional[data_types.ExecutionInfo] = None,
) -> None:
"""Does cleanup or post-processing after the step finished running.
Subclasses should overwrite this method if they need to run any
additional code after the step execution.
Args:
pipeline_name: Name of the parent pipeline of the step that was
executed.
step_name: Name of the step that was executed.
pipeline_node: Protobuf node corresponding to the step that was
executed.
execution_info: Info about the finished step execution.
"""
run(self)
Runs a single ZenML step.
Subclasses should in most cases not need to overwrite this method and
implement their custom logic in the setup(...) and post_run(...)
methods instead. If you still need to customize the functionality of
this method, make sure to still include all the existing logic as your
step won't be executed properly otherwise.
Exceptions:
| Type | Description |
|---|---|
TypeError |
If the arguments passed to the entrypoint are invalid. |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
def run(self) -> None:
"""Runs a single ZenML step.
Subclasses should in most cases not need to overwrite this method and
implement their custom logic in the `setup(...)` and `post_run(...)`
methods instead. If you still need to customize the functionality of
this method, make sure to still include all the existing logic as your
step won't be executed properly otherwise.
Raises:
TypeError: If the arguments passed to the entrypoint are invalid.
"""
# Make sure this entrypoint does not run an entire pipeline when
# importing user modules. This could happen if the `pipeline.run()` call
# is not wrapped in a function or an `if __name__== "__main__":` check)
constants.SHOULD_PREVENT_PIPELINE_EXECUTION = True
# Extract and parse all the entrypoint arguments required to execute
# the step. See `get_entrypoint_options()` for an in-depth explanation
# of all these arguments.
pb2_pipeline = Pb2Pipeline()
pb2_pipeline_json = string_utils.b64_decode(
self.entrypoint_args[PIPELINE_JSON_OPTION]
)
json_format.Parse(pb2_pipeline_json, pb2_pipeline)
main_module_source = self.entrypoint_args[MAIN_MODULE_SOURCE_OPTION]
step_source = self.entrypoint_args[STEP_SOURCE_OPTION]
input_artifact_sources = json.loads(
string_utils.b64_decode(
self.entrypoint_args[INPUT_ARTIFACT_SOURCES_OPTION]
)
)
materializer_sources = json.loads(
string_utils.b64_decode(
self.entrypoint_args[MATERIALIZER_SOURCES_OPTION]
)
)
# Get some common values that will be used throughout the remainder of
# this method
pipeline_name = pb2_pipeline.pipeline_info.id
if "@" in step_source:
# Get rid of potential ZenML version pins if the source looks like
# this `zenml.some_module.SomeClass@zenml_0.9.9`
step_source, _ = step_source.split("@", 1)
step_module, step_name = step_source.rsplit(".", 1)
# Allow subclasses to run custom code before user code is imported and
# the step is executed.
self.setup(pipeline_name=pipeline_name, step_name=step_name)
# Activate all the integrations. This makes sure that all materializers
# and stack component flavors are registered.
integration_registry.activate_integrations()
# Import the main module that was executed to run the pipeline to which
# the step getting executed belongs. Even if the step class is not
# defined in this module, we need to import it in case the user
# defined/imported custom materializers/artifacts here that ZenML needs
# to execute the step.
importlib.import_module(main_module_source)
# The `__main__` module when running a step inside a docker container
# will always be the entrypoint file
# (e.g. `zenml.entrypoints.single_step_entrypoint`). If this entrypoint
# now needs to execute something in a different environment,
# `sys.modules["__main__"]` would not point to the original user main
# module anymore. That's why we set this constant here so other
# components (e.g. step operators) can use it to correctly import it
# in the containers they create.
constants.USER_MAIN_MODULE = main_module_source
# Create an instance of the ZenML step that this entrypoint should run.
step_class = source_utils.load_source_path_class(step_source)
if not issubclass(step_class, BaseStep):
raise TypeError(
f"The step source `{step_source}` passed to the "
f"entrypoint is not pointing to a `{BaseStep}` subclass."
)
step = step_class()
# Compute the name of the module in which the step was in when the
# pipeline run was started. This could potentially differ from the
# module part of `STEP_SOURCE_OPTION` if the step was defined in
# the `__main__` module (e.g. the step was defined in `run.py` and the
# pipeline run was started by the call `python run.py`).
# Why we need this: The executor source is stored inside the protobuf
# pipeline representation object `pb2_pipeline` which was passed to
# this entrypoint. Tfx will try to load the executor class from that
# source, so we need to recreate the executor class in the same module.
original_step_module = (
"__main__" if step_module == main_module_source else step_module
)
# Make sure the `__module__` attribute of our step instance points to
# this original module so the executor class gets created in the
# correct module.
step.__module__ = original_step_module
# Create the executor class that is responsible for running the step
# function. This method call dynamically creates an executor class
# which will later be used when `orchestrator.run_step(...)` is running
# the step.
self._create_executor_class(
step=step,
input_artifact_sources=input_artifact_sources,
materializer_sources=materializer_sources,
)
# Execute the actual step code.
run_name = self.get_run_name(pipeline_name=pipeline_name)
orchestrator = Repository().active_stack.orchestrator
execution_info = orchestrator.run_step(
step=step, run_name=run_name, pb2_pipeline=pb2_pipeline
)
# Allow subclasses to run custom code after the step finished executing.
pipeline_node = orchestrator._get_node_with_step_name(
step_name=step_name, pb2_pipeline=pb2_pipeline
)
self.post_run(
pipeline_name=pipeline_name,
step_name=step_name,
pipeline_node=pipeline_node,
execution_info=execution_info,
)
setup(self, pipeline_name, step_name)
Code that needs to run before user code is imported or the step is executed.
Subclasses should overwrite this method if they need to perform any
additional setup, but should in most cases include a
super().setup(...) call so the ZenML setup code will still be
executed.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
pipeline_name |
str |
Name of the parent pipeline of the step that will be executed. |
required |
step_name |
str |
Name of the step that will be executed. |
required |
Source code in zenml/entrypoints/step_entrypoint_configuration.py
def setup(self, pipeline_name: str, step_name: str) -> None:
"""Code that needs to run before user code is imported or the step is executed.
Subclasses should overwrite this method if they need to perform any
additional setup, but should in most cases include a
`super().setup(...)` call so the ZenML setup code will still be
executed.
Args:
pipeline_name: Name of the parent pipeline of the step that will
be executed.
step_name: Name of the step that will be executed.
"""
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)