zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/integrations/github/plugins/event_sources/github_webhook_event_source.py

@@ -17,7 +17,7 @@ import urllib
 from typing import Any, Dict, List, Optional, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.enums import SecretScope
 from zenml.event_sources.base_event import (
@@ -111,14 +111,10 @@ class GithubEvent(BaseEvent):
     after: str
     repository: Repository
     commits: List[Commit]
-    head_commit: Optional[Commit]
-    tags: Optional[List[Tag]]
-    pull_requests: Optional[List[PullRequest]]
-
-    class Config:
-        """Pydantic configuration class."""
-
-        extra = Extra.allow
+    head_commit: Optional[Commit] = None
+    tags: Optional[List[Tag]] = None
+    pull_requests: Optional[List[PullRequest]] = None
+    model_config = ConfigDict(extra="allow")
 
     @property
     def branch(self) -> Optional[str]:
@@ -157,9 +153,9 @@ class GithubEvent(BaseEvent):
 class GithubWebhookEventFilterConfiguration(WebhookEventFilterConfig):
     """Configuration for github event filters."""
 
-    repo: Optional[str]
-    branch: Optional[str]
-    event_type: Optional[GithubEventType]
+    repo: Optional[str] = None
+    branch: Optional[str] = None
+    event_type: Optional[GithubEventType] = None
 
     def event_matches_filter(self, event: BaseEvent) -> bool:
         """Checks the filter against the inbound event.
@@ -442,7 +438,7 @@ class GithubWebhookEventSourceHandler(BaseWebhookEventSourceHandler):
         if config.rotate_secret:
             # In case the secret is being rotated
             secret_key_value = random_str(12)
-            webhook_secret = SecretUpdate(  # type: ignore[call-arg]
+            webhook_secret = SecretUpdate(
                 values={"webhook_secret": secret_key_value}
             )
             self.zen_store.update_secret(

zenml/integrations/great_expectations/__init__.py

@@ -32,7 +32,7 @@ class GreatExpectationsIntegration(Integration):
 
     NAME = GREAT_EXPECTATIONS
     REQUIREMENTS = [
-        "great-expectations>=0.15.0,<=0.15.47",
+        "great-expectations>=0.17.15,<1.0",
     ]
 
     @staticmethod

zenml/integrations/great_expectations/data_validators/ge_data_validator.py

@@ -17,21 +17,25 @@ import os
 from typing import Any, ClassVar, Dict, List, Optional, Sequence, Type, cast
 
 import pandas as pd
-import yaml
 from great_expectations.checkpoint.types.checkpoint_result import (  # type: ignore[import-untyped]
     CheckpointResult,
 )
 from great_expectations.core import (  # type: ignore[import-untyped]
     ExpectationSuite,
 )
-from great_expectations.data_context.data_context import (  # type: ignore[import-untyped]
-    BaseDataContext,
-    DataContext,
+from great_expectations.data_context.data_context.abstract_data_context import (
+    AbstractDataContext,
 )
-from great_expectations.data_context.types.base import (  # type: ignore[import-untyped]
+from great_expectations.data_context.data_context.context_factory import (
+    get_context,
+)
+from great_expectations.data_context.data_context.ephemeral_data_context import (
+    EphemeralDataContext,
+)
+from great_expectations.data_context.types.base import (
     DataContextConfig,
 )
-from great_expectations.data_context.types.resource_identifiers import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.resource_identifiers import (
     ExpectationSuiteIdentifier,
 )
 from great_expectations.profile.user_configurable_profiler import (  # type: ignore[import-untyped]
@@ -65,8 +69,8 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         GreatExpectationsDataValidatorFlavor
     )
 
-    _context: BaseDataContext = None
-    _context_config: Optional[Dict[str, Any]] = None
+    _context: Optional[AbstractDataContext] = None
+    _context_config: Optional[DataContextConfig] = None
 
     @property
     def config(self) -> GreatExpectationsDataValidatorConfig:
@@ -78,7 +82,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         return cast(GreatExpectationsDataValidatorConfig, self._config)
 
     @classmethod
-    def get_data_context(cls) -> BaseDataContext:
+    def get_data_context(cls) -> AbstractDataContext:
         """Get the Great Expectations data context managed by ZenML.
 
         Call this method to retrieve the data context managed by ZenML
@@ -94,15 +98,11 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         return data_validator.data_context
 
     @property
-    def context_config(self) -> Optional[Dict[str, Any]]:
+    def context_config(self) -> Optional[DataContextConfig]:
         """Get the Great Expectations data context configuration.
 
-        The first time the context config is loaded from the stack component
-        config, it is converted from JSON/YAML string format to a dict.
-
         Raises:
-            ValueError: If the context_config value is not a valid JSON/YAML or
-                if the GE configuration extracted from it fails GE validation.
+            ValueError: In case there is an invalid context_config value
 
         Returns:
             A dictionary with the GE data context configuration.
@@ -111,31 +111,18 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         if self._context_config is not None:
             return self._context_config
 
-        # Otherwise, load it from the stack component config
-        context_config = self.config.context_config
-        if context_config is None:
+        # Otherwise, use the configuration from the stack component config, if
+        # set
+        context_config_dict = self.config.context_config
+        if context_config_dict is None:
             return None
-        if isinstance(context_config, dict):
-            self._context_config = context_config
-            return self._context_config
-
-        # If the context config is a string, try to parse it as JSON/YAML
-        try:
-            context_config_dict = yaml.safe_load(context_config)
-        except yaml.parser.ParserError as e:
-            raise ValueError(
-                f"Malformed `context_config` value. Only JSON and YAML "
-                f"formats are supported: {str(e)}"
-            )
 
         # Validate that the context config is a valid GE config
         try:
-            context_config = DataContextConfig(**context_config_dict)
-            BaseDataContext(project_config=context_config)
+            self._context_config = DataContextConfig(**context_config_dict)
         except Exception as e:
             raise ValueError(f"Invalid `context_config` value: {str(e)}")
 
-        self._context_config = cast(Dict[str, Any], context_config_dict)
         return self._context_config
 
     @property
@@ -203,7 +190,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
         }
 
     @property
-    def data_context(self) -> BaseDataContext:
+    def data_context(self) -> AbstractDataContext:
         """Returns the Great Expectations data context configured for this component.
 
         Returns:
@@ -216,7 +203,9 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             profiler_store_name = "zenml_profiler_store"
             evaluation_parameter_store_name = "evaluation_parameter_store"
 
-            zenml_context_config = dict(
+            # Define default configuration options that plug the GX stores
+            # in the active ZenML artifact store
+            zenml_context_config: Dict[str, Any] = dict(
                 stores={
                     expectations_store_name: self.get_store_config(
                         "ExpectationsStore", "expectations"
@@ -250,18 +239,29 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             if self.config.context_root_dir:
                 # initialize the local data context, if a local path was
                 # configured
-                self._context = DataContext(self.config.context_root_dir)
+                self._context = get_context(
+                    context_root_dir=self.config.context_root_dir
+                )
+
             else:
-                # create an in-memory data context configuration that is not
-                # backed by a local YAML file (see https://docs.greatexpectations.io/docs/guides/setup/configuring_data_contexts/how_to_instantiate_a_data_context_without_a_yml_file/).
+                # create an ephemeral in-memory data context that is not
+                # backed by a local YAML file (see https://docs.greatexpectations.io/docs/oss/guides/setup/configuring_data_contexts/instantiating_data_contexts/instantiate_data_context/).
                 if self.context_config:
-                    context_config = DataContextConfig(**self.context_config)
+                    # Use the data context configuration provided in the stack
+                    # component configuration
+                    context_config = self.context_config
                 else:
+                    # Initialize the data context with the default ZenML
+                    # configuration options effectively plugging the GX stores
+                    # into the ZenML artifact store
                     context_config = DataContextConfig(**zenml_context_config)
                     # skip adding the stores after initialization, as they are
                     # already baked in the initial configuration
                     configure_zenml_stores = False
-                self._context = BaseDataContext(project_config=context_config)
+
+                self._context = EphemeralDataContext(
+                    project_config=context_config
+                )
 
             if configure_zenml_stores:
                 self._context.config.expectations_store_name = (
@@ -277,14 +277,14 @@ class GreatExpectationsDataValidator(BaseDataValidator):
                 self._context.config.evaluation_parameter_store_name = (
                     evaluation_parameter_store_name
                 )
-                for store_name, store_config in zenml_context_config[  # type: ignore[attr-defined]
+                for store_name, store_config in zenml_context_config[
                     "stores"
                 ].items():
                     self._context.add_store(
                         store_name=store_name,
                         store_config=store_config,
                     )
-                for site_name, site_config in zenml_context_config[  # type: ignore[attr-defined]
+                for site_name, site_config in zenml_context_config[
                     "data_docs_sites"
                 ].items():
                     self._context.config.data_docs_sites[site_name] = (
@@ -509,7 +509,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             },
         ]
 
-        checkpoint_config = {
+        checkpoint_config: Dict[str, Any] = {
             "name": checkpoint_name,
             "run_name_template": run_name,
             "config_version": 1,
@@ -517,7 +517,7 @@ class GreatExpectationsDataValidator(BaseDataValidator):
             "expectation_suite_name": expectation_suite_name,
             "action_list": action_list,
         }
-        context.add_checkpoint(**checkpoint_config)
+        context.add_checkpoint(**checkpoint_config)  # type: ignore[has-type]
 
         try:
             results = context.run_checkpoint(

zenml/integrations/great_expectations/flavors/great_expectations_data_validator_flavor.py

@@ -16,7 +16,9 @@
 import os
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from pydantic import validator
+import yaml
+from pydantic import field_validator, model_validator
+from yaml.parser import ParserError
 
 from zenml.data_validators.base_data_validator import (
     BaseDataValidatorConfig,
@@ -26,6 +28,7 @@ from zenml.integrations.great_expectations import (
     GREAT_EXPECTATIONS_DATA_VALIDATOR_FLAVOR,
 )
 from zenml.io import fileio
+from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.integrations.great_expectations.data_validators import (
@@ -41,6 +44,8 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
             data context. If configured, the data validator will only be usable
             with local orchestrators.
         context_config: in-line Great Expectations data context configuration.
+            If the `context_root_dir` attribute is also set, this configuration
+            will be ignored.
         configure_zenml_stores: if set, ZenML will automatically configure
             stores that use the Artifact Store as a backend. If neither
             `context_root_dir` nor `context_config` are set, this is the default
@@ -54,7 +59,8 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
     configure_zenml_stores: bool = False
     configure_local_docs: bool = True
 
-    @validator("context_root_dir")
+    @field_validator("context_root_dir")
+    @classmethod
     def _ensure_valid_context_root_dir(
         cls, context_root_dir: Optional[str] = None
     ) -> Optional[str]:
@@ -78,6 +84,33 @@ class GreatExpectationsDataValidatorConfig(BaseDataValidatorConfig):
                 )
         return context_root_dir
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def validate_context_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert the context configuration if given in JSON/YAML format.
+
+        Args:
+            data: The configuration values.
+
+        Returns:
+            The validated configuration values.
+
+        Raises:
+            ValueError: If the context configuration is not a valid
+                JSON/YAML object.
+        """
+        if isinstance(data.get("context_config"), str):
+            try:
+                data["context_config"] = yaml.safe_load(data["context_config"])
+            except ParserError as e:
+                raise ValueError(
+                    f"Malformed `context_config` value. Only JSON and YAML "
+                    f"formats are supported: {str(e)}"
+                )
+
+        return data
+
     @property
     def is_local(self) -> bool:
         """Checks if this stack component is running locally.

zenml/integrations/great_expectations/ge_store_backend.py

@@ -17,14 +17,16 @@ import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, cast
 
-from great_expectations.data_context.store.tuple_store_backend import (  # type: ignore[import-untyped]
+from great_expectations.data_context.store.tuple_store_backend import (
     TupleStoreBackend,
-    filter_properties_dict,
 )
 from great_expectations.exceptions import (  # type: ignore[import-untyped]
     InvalidKeyError,
     StoreBackendError,
 )
+from great_expectations.util import (  # type: ignore[import-untyped]
+    filter_properties_dict,
+)
 
 from zenml.client import Client
 from zenml.io import fileio
@@ -34,7 +36,7 @@ from zenml.utils import io_utils
 logger = get_logger(__name__)
 
 
-class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
+class ZenMLArtifactStoreBackend(TupleStoreBackend):
     """Great Expectations store backend that uses the active ZenML Artifact Store as a store."""
 
     def __init__(
@@ -105,7 +107,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
         if not isinstance(key, tuple):
             key = key.to_tuple()
         if not is_prefix:
-            object_relative_path = self._convert_key_to_filepath(key)
+            object_relative_path = self._convert_key_to_filepath(key)  # type: ignore[no-untyped-call]
         elif key:
             object_relative_path = os.path.join(*key)
         else:
@@ -116,7 +118,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
             object_key = object_relative_path
         return os.path.join(self.root_path, object_key)
 
-    def _get(self, key: Tuple[str, ...]) -> str:
+    def _get(self, key: Tuple[str, ...]) -> str:  # type: ignore[override]
         """Get the value of an object from the store.
 
         Args:
@@ -140,7 +142,18 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
             )
         return contents
 
-    def _set(self, key: Tuple[str, ...], value: str, **kwargs: Any) -> str:
+    def _get_all(self) -> List[Any]:
+        """Get all objects in the store.
+
+        Raises:
+            NotImplementedError: if the method is not implemented for this store
+                backend.
+        """
+        raise NotImplementedError(
+            "Method `_get_all` is not implemented for this store backend."
+        )
+
+    def _set(self, key: Tuple[str, ...], value: str, **kwargs: Any) -> str:  # type: ignore[override]
         """Set the value of an object in the store.
 
         Args:
@@ -212,12 +225,12 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
                     self.filepath_suffix
                 ):
                     continue
-                key = self._convert_filepath_to_key(filepath)
-                if key and not self.is_ignored_key(key):
+                key = self._convert_filepath_to_key(filepath)  # type: ignore[no-untyped-call]
+                if key and not self.is_ignored_key(key):  # type: ignore[no-untyped-call]
                     key_list.append(key)
         return key_list
 
-    def remove_key(self, key: Tuple[str, ...]) -> bool:
+    def remove_key(self, key: Tuple[str, ...]) -> bool:  # type: ignore[override]
         """Delete an object from the store.
 
         Args:
@@ -250,7 +263,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
         result = fileio.exists(filepath)
         return result
 
-    def get_url_for_key(
+    def get_url_for_key(  # type: ignore[override]
         self, key: Tuple[str, ...], protocol: Optional[str] = None
     ) -> str:
         """Get the URL of an object in the store.
@@ -292,7 +305,7 @@ class ZenMLArtifactStoreBackend(TupleStoreBackend):  # type: ignore[misc]
                 f"requested but `base_public_path` was not configured for the "
                 f"{self.__class__.__name__}"
             )
-        filepath = self._convert_key_to_filepath(key)
+        filepath = self._convert_key_to_filepath(key)  # type: ignore[no-untyped-call]
         public_url = self.base_public_path + filepath.replace(self.proto, "")
         return cast(str, public_url)
 

zenml/integrations/great_expectations/materializers/ge_materializer.py

@@ -25,10 +25,10 @@ from great_expectations.core import (  # type: ignore[import-untyped]
 from great_expectations.core.expectation_validation_result import (  # type: ignore[import-untyped]
     ExpectationSuiteValidationResult,
 )
-from great_expectations.data_context.types.base import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.base import (
     CheckpointConfig,
 )
-from great_expectations.data_context.types.resource_identifiers import (  # type: ignore[import-untyped]
+from great_expectations.data_context.types.resource_identifiers import (
     ExpectationSuiteIdentifier,
     ValidationResultIdentifier,
 )
@@ -86,7 +86,7 @@ class GreatExpectationsMaterializer(BaseMaterializer):
         validation_dict = {}
         for result_ident, results in artifact_dict["run_results"].items():
             validation_ident = (
-                ValidationResultIdentifier.from_fixed_length_tuple(
+                ValidationResultIdentifier.from_fixed_length_tuple(  # type: ignore[no-untyped-call]
                     result_ident.split("::")[1].split("/")
                 )
             )

zenml/integrations/great_expectations/utils.py

@@ -13,14 +13,14 @@
 #  permissions and limitations under the License.
 """Great Expectations data profiling standard step."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import pandas as pd
 from great_expectations.core.batch import (  # type: ignore[import-untyped]
     RuntimeBatchRequest,
 )
-from great_expectations.data_context.data_context import (  # type: ignore[import-untyped]
-    BaseDataContext,
+from great_expectations.data_context.data_context.abstract_data_context import (
+    AbstractDataContext,
 )
 
 from zenml import get_step_context
@@ -31,7 +31,7 @@ logger = get_logger(__name__)
 
 
 def create_batch_request(
-    context: BaseDataContext,
+    context: AbstractDataContext,
     dataset: pd.DataFrame,
     data_asset_name: Optional[str],
 ) -> RuntimeBatchRequest:
@@ -62,7 +62,7 @@ def create_batch_request(
     data_asset_name = data_asset_name or f"{pipeline_name}_{step_name}"
     batch_identifier = "default"
 
-    datasource_config = {
+    datasource_config: Dict[str, Any] = {
         "name": datasource_name,
         "class_name": "Datasource",
         "module_name": "great_expectations.datasource",

zenml/integrations/huggingface/__init__.py

@@ -30,6 +30,9 @@ class HuggingfaceIntegration(Integration):
         "transformers<=4.31",
         "datasets",
         "huggingface_hub>0.19.0",
+        "accelerate",
+        "bitsandbytes>=0.41.3",
+        "peft",
         # temporary fix for CI issue similar to:
         # - https://github.com/huggingface/datasets/issues/6737
         # - https://github.com/huggingface/datasets/issues/6697

zenml/integrations/huggingface/flavors/huggingface_model_deployer_flavor.py

@@ -61,7 +61,7 @@ class HuggingFaceModelDeployerConfig(
         namespace: Hugging Face namespace used to list endpoints
     """
 
-    token: Optional[str] = SecretField()
+    token: Optional[str] = SecretField(default=None)
 
     # The namespace to list endpoints for. Set to `"*"` to list all endpoints
     # from all namespaces (i.e. personal namespace and all orgs the user belongs to).

zenml/integrations/huggingface/steps/__init__.py

@@ -16,3 +16,6 @@
 from zenml.integrations.huggingface.steps.huggingface_deployer import (
     huggingface_model_deployer_step,
 )
+from zenml.integrations.huggingface.steps.accelerate_runner import (
+    run_with_accelerate,
+)

zenml/integrations/huggingface/steps/accelerate_runner.py

@@ -0,0 +1,149 @@
+# Apache Software License 2.0
+#
+# Copyright (c) ZenML GmbH 2024. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Step function to run any ZenML step using Accelerate."""
+
+import functools
+from typing import Any, Callable, Optional, TypeVar, cast
+
+import cloudpickle as pickle
+from accelerate.commands.launch import (  # type: ignore[import-untyped]
+    launch_command,
+    launch_command_parser,
+)
+
+from zenml.logger import get_logger
+from zenml.steps import BaseStep
+from zenml.utils.function_utils import _cli_arg_name, create_cli_wrapped_script
+
+logger = get_logger(__name__)
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def run_with_accelerate(
+    step_function: BaseStep,
+    num_processes: Optional[int] = None,
+    use_cpu: bool = False,
+) -> BaseStep:
+    """Run a function with accelerate.
+
+    Accelerate package: https://huggingface.co/docs/accelerate/en/index
+    Example:
+        ```python
+        from zenml import step, pipeline
+        from zenml.integrations.hugginface.steps import run_with_accelerate
+        @step
+        def training_step(some_param: int, ...):
+            # your training code is below
+            ...
+
+        @pipeline
+        def training_pipeline(some_param: int, ...):
+            run_with_accelerate(training_step, num_processes=4)(some_param, ...)
+        ```
+
+    Args:
+        step_function: The step function to run.
+        num_processes: The number of processes to use.
+        use_cpu: Whether to use the CPU.
+
+    Returns:
+        The accelerate-enabled version of the step.
+    """
+
+    def _decorator(entrypoint: F) -> F:
+        @functools.wraps(entrypoint)
+        def inner(*args: Any, **kwargs: Any) -> Any:
+            if args:
+                raise ValueError(
+                    "Accelerated steps do not support positional arguments."
+                )
+
+            if not use_cpu:
+                import torch
+
+                logger.info("Starting accelerate job...")
+
+                device_count = torch.cuda.device_count()
+                if num_processes is None:
+                    _num_processes = device_count
+                else:
+                    if num_processes > device_count:
+                        logger.warning(
+                            f"Number of processes ({num_processes}) is greater than "
+                            f"the number of available GPUs ({device_count}). Using all GPUs."
+                        )
+                        _num_processes = device_count
+                    else:
+                        _num_processes = num_processes
+            else:
+                _num_processes = num_processes or 1
+
+            with create_cli_wrapped_script(
+                entrypoint, flavour="accelerate"
+            ) as (
+                script_path,
+                output_path,
+            ):
+                commands = ["--num_processes", str(_num_processes)]
+                if use_cpu:
+                    commands += [
+                        "--cpu",
+                        "--num_cpu_threads_per_process",
+                        "10",
+                    ]
+                commands.append(str(script_path.absolute()))
+                for k, v in kwargs.items():
+                    k = _cli_arg_name(k)
+                    if isinstance(v, bool):
+                        if v:
+                            commands.append(f"--{k}")
+                    elif isinstance(v, str):
+                        commands += [f"--{k}", '"{v}"']
+                    elif type(v) in (list, tuple, set):
+                        for each in v:
+                            commands.append(f"--{k}")
+                            if isinstance(each, str):
+                                commands.append(f'"{each}"')
+                            else:
+                                commands.append(f"{each}")
+                    else:
+                        commands += [f"--{k}", f"{v}"]
+
+                logger.debug(commands)
+
+                parser = launch_command_parser()
+                args = parser.parse_args(commands)
+                try:
+                    launch_command(args)
+                except Exception as e:
+                    logger.error(
+                        "Accelerate training job failed... See error message for details."
+                    )
+                    raise RuntimeError(
+                        "Accelerate training job failed."
+                    ) from e
+                else:
+                    logger.info(
+                        "Accelerate training job finished successfully."
+                    )
+                    return pickle.load(open(output_path, "rb"))
+
+        return cast(F, inner)
+
+    setattr(step_function, "entrypoint", _decorator(step_function.entrypoint))
+
+    return step_function