zenml-nightly 0.83.1.dev20250709__py3-none-any.whl → 0.83.1.dev20250711__py3-none-any.whl

zenml/config/retry_config.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Retry configuration for a step."""
 
+from pydantic import NonNegativeInt, PositiveInt
+
 from zenml.config.strict_base_model import StrictBaseModel
 
 
@@ -22,6 +24,6 @@ class StepRetryConfig(StrictBaseModel):
     Delay is an integer (specified in seconds).
     """
 
-    max_retries: int = 1
-    delay: int = 0  # in seconds
-    backoff: int = 0
+    max_retries: PositiveInt = 1
+    delay: NonNegativeInt = 0  # in seconds
+    backoff: NonNegativeInt = 0

zenml/config/step_configurations.py

@@ -289,6 +289,7 @@ class StepConfiguration(PartialStepConfiguration):
                 "extra",
                 "failure_hook_source",
                 "success_hook_source",
+                "retry",
                 "substitutions",
             },
             exclude_none=True,
@@ -300,12 +301,17 @@ class StepConfiguration(PartialStepConfiguration):
                     "extra",
                     "failure_hook_source",
                     "success_hook_source",
+                    "retry",
                     "substitutions",
                 },
                 exclude_none=True,
             )
 
-            updated_config = self.model_copy(update=pipeline_values, deep=True)
+            updated_config_dict = {
+                **self.model_dump(),
+                **pipeline_values,
+            }
+            updated_config = self.model_validate(updated_config_dict)
             return update_model(updated_config, original_values)
         else:
             return self.model_copy(deep=True)

zenml/console.py

@@ -18,11 +18,14 @@ from rich.style import Style
 from rich.theme import Theme
 
 zenml_style_defaults = {
-    "info": Style(color="cyan", dim=True),
+    "info": Style(color="white", dim=True),
     "warning": Style(color="yellow"),
     "danger": Style(color="red", bold=True),
     "title": Style(color="cyan", bold=True, underline=True),
     "error": Style(color="red"),
+    "success": Style(color="green"),
+    "repr.str": Style(color="white", dim=False),
+    "repr.uuid": Style(color="magenta", dim=False),
 }
 
 zenml_custom_theme = Theme(zenml_style_defaults)

zenml/constants.py

@@ -166,12 +166,14 @@ ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING = (
 ENV_ZENML_SKIP_IMAGE_BUILDER_DEFAULT = "ZENML_SKIP_IMAGE_BUILDER_DEFAULT"
 ENV_ZENML_SKIP_STACK_VALIDATION = "ZENML_SKIP_STACK_VALIDATION"
 ENV_ZENML_SERVER = "ZENML_SERVER"
+ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS = (
+    "ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS"
+)
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE = "ZENML_DISABLE_PIPELINE_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS = "ZENML_DISABLE_STEP_NAMES_IN_LOGS"
-ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"

zenml/container_registries/base_container_registry.py

@@ -16,7 +16,7 @@
 import re
 from typing import TYPE_CHECKING, Optional, Tuple, Type, cast
 
-from pydantic import field_validator
+from pydantic import Field, field_validator
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.enums import StackComponentType
@@ -36,12 +36,24 @@ if TYPE_CHECKING:
 class BaseContainerRegistryConfig(AuthenticationConfigMixin):
     """Base config for a container registry.
 
-    Attributes:
-        uri: The URI of the container registry.
+    Configuration for connecting to container image registries.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    uri: str
-    default_repository: Optional[str] = None
+    uri: str = Field(
+        description="Container registry URI (e.g., 'gcr.io' for Google Container "
+        "Registry, 'docker.io' for Docker Hub, 'registry.gitlab.com' for GitLab "
+        "Container Registry, 'ghcr.io' for GitHub Container Registry). This is "
+        "the base URL where container images will be pushed to and pulled from."
+    )
+    default_repository: Optional[str] = Field(
+        default=None,
+        description="Default repository namespace for image storage (e.g., "
+        "'username' for Docker Hub, 'project-id' for GCR, 'organization' for "
+        "GitHub Container Registry). If not specified, images will be stored at "
+        "the registry root. For Docker Hub this would mean only official images "
+        "can be pushed.",
+    )
 
     @field_validator("uri")
     @classmethod

zenml/enums.py

@@ -78,6 +78,8 @@ class ExecutionStatus(StrEnum):
     COMPLETED = "completed"
     RUNNING = "running"
     CACHED = "cached"
+    RETRYING = "retrying"
+    RETRIED = "retried"
     STOPPED = "stopped"
     STOPPING = "stopping"
 
@@ -92,9 +94,19 @@ class ExecutionStatus(StrEnum):
             ExecutionStatus.FAILED,
             ExecutionStatus.COMPLETED,
             ExecutionStatus.CACHED,
+            ExecutionStatus.RETRIED,
             ExecutionStatus.STOPPED,
         }
 
+    @property
+    def is_successful(self) -> bool:
+        """Whether the execution status refers to a successful execution.
+
+        Returns:
+            Whether the execution status refers to a successful execution.
+        """
+        return self in {ExecutionStatus.COMPLETED, ExecutionStatus.CACHED}
+
 
 class LoggingLevels(Enum):
     """Enum for logging levels."""
@@ -411,9 +423,6 @@ class OnboardingStep(StrEnum):
     DEVICE_VERIFIED = "device_verified"
     PROJECT_CREATED = "project_created"
     PIPELINE_RUN = "pipeline_run"
-    SECOND_PIPELINE_RUN = "second_pipeline_run"
-    THIRD_PIPELINE_RUN = "third_pipeline_run"
-    STARTER_SETUP_COMPLETED = "starter_setup_completed"
     STACK_WITH_REMOTE_ORCHESTRATOR_CREATED = (
         "stack_with_remote_orchestrator_created"
     )
@@ -426,7 +435,7 @@ class OnboardingStep(StrEnum):
     PIPELINE_RUN_WITH_REMOTE_ARTIFACT_STORE = (
         "pipeline_run_with_remote_artifact_store"
     )
-    PRODUCTION_SETUP_COMPLETED = "production_setup_completed"
+    OSS_ONBOARDING_COMPLETED = "oss_onboarding_completed"
     PRO_ONBOARDING_COMPLETED = "pro_onboarding_completed"
 
 

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -37,100 +37,123 @@ DEFAULT_OUTPUT_DATA_S3_MODE = "EndOfJob"
 
 
 class SagemakerOrchestratorSettings(BaseSettings):
-    """Settings for the Sagemaker orchestrator.
-
-    Attributes:
-        synchronous: If `True`, the client running a pipeline using this
-            orchestrator waits until all steps finish running. If `False`,
-            the client returns immediately and the pipeline is executed
-            asynchronously. Defaults to `True`.
-        instance_type: The instance type to use for the processing job.
-        execution_role: The IAM role to use for the step execution.
-        processor_role: DEPRECATED: use `execution_role` instead.
-        volume_size_in_gb: The size of the EBS volume to use for the processing
-            job.
-        max_runtime_in_seconds: The maximum runtime in seconds for the
-            processing job.
-        tags: Tags to apply to the Processor/Estimator assigned to the step.
-        pipeline_tags: Tags to apply to the pipeline via the
-            sagemaker.workflow.pipeline.Pipeline.create method.
-        processor_tags: DEPRECATED: use `tags` instead.
-        keep_alive_period_in_seconds: The time in seconds after which the
-            provisioned instance will be terminated if not used. This is only
-            applicable for TrainingStep type and it is not possible to use
-            TrainingStep type if the `output_data_s3_uri` is set to Dict[str, str].
-        use_training_step: Whether to use the TrainingStep type.
-            It is not possible to use TrainingStep type
-            if the `output_data_s3_uri` is set to Dict[str, str] or if the
-            `output_data_s3_mode` != "EndOfJob".
-        processor_args: Arguments that are directly passed to the SageMaker
-            Processor for a specific step, allowing for overriding the default
-            settings provided when configuring the component. See
-            https://sagemaker.readthedocs.io/en/stable/api/training/processing.html#sagemaker.processing.Processor
-            for a full list of arguments.
-            For processor_args.instance_type, check
-            https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html
-            for a list of available instance types.
-        environment: Environment variables to pass to the container.
-        estimator_args: Arguments that are directly passed to the SageMaker
-            Estimator for a specific step, allowing for overriding the default
-            settings provided when configuring the component. See
-            https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html#sagemaker.estimator.Estimator
-            for a full list of arguments.
-            For a list of available instance types, check
-            https://docs.aws.amazon.com/sagemaker/latest/dg/cmn-info-instance-types.html.
-        input_data_s3_mode: How data is made available to the container.
-            Two possible input modes: File, Pipe.
-        input_data_s3_uri: S3 URI where data is located if not locally,
-            e.g. s3://my-bucket/my-data/train. How data will be made available
-            to the container is configured with input_data_s3_mode. Two possible
-            input types:
-                - str: S3 location where training data is saved.
-                - Dict[str, str]: (ChannelName, S3Location) which represent
-                - Dict[str, str]: (ChannelName, S3Location) which represent
-                    channels (e.g. training, validation, testing) where
-                    specific parts of the data are saved in S3.
-        output_data_s3_mode: How data is uploaded to the S3 bucket.
-            Two possible output modes: EndOfJob, Continuous.
-        output_data_s3_uri: S3 URI where data is uploaded after or during processing run.
-            e.g. s3://my-bucket/my-data/output. How data will be made available
-            to the container is configured with output_data_s3_mode. Two possible
-            input types:
-                - str: S3 location where data will be uploaded from a local folder
-                    named /opt/ml/processing/output/data.
-                - Dict[str, str]: (ChannelName, S3Location) which represent
-                    channels (e.g. output_one, output_two) where
-                    specific parts of the data are stored locally for S3 upload.
-                    Data must be available locally in /opt/ml/processing/output/data/<ChannelName>.
-    """
-
-    synchronous: bool = True
+    """Settings for the Sagemaker orchestrator."""
+
+    synchronous: bool = Field(
+        True,
+        description="Controls whether pipeline execution blocks the client. If True, "
+        "the client waits until all steps complete before returning. If False, "
+        "returns immediately and executes asynchronously. Useful for long-running "
+        "production pipelines where you don't want to maintain a connection",
+    )
 
-    instance_type: Optional[str] = None
-    execution_role: Optional[str] = None
-    volume_size_in_gb: int = 30
-    max_runtime_in_seconds: int = 86400
-    tags: Dict[str, str] = {}
-    pipeline_tags: Dict[str, str] = {}
-    keep_alive_period_in_seconds: Optional[int] = 300  # 5 minutes
-    use_training_step: Optional[bool] = None
+    instance_type: Optional[str] = Field(
+        None,
+        description="AWS EC2 instance type for step execution. Must be a valid "
+        "SageMaker-supported instance type. Examples: 'ml.t3.medium' (2 vCPU, 4GB RAM), "
+        "'ml.m5.xlarge' (4 vCPU, 16GB RAM), 'ml.p3.2xlarge' (8 vCPU, 61GB RAM, 1 GPU). "
+        "Defaults to ml.m5.xlarge for training steps or ml.t3.medium for processing steps",
+    )
+    execution_role: Optional[str] = Field(
+        None,
+        description="IAM role ARN for SageMaker step execution permissions. Must have "
+        "necessary policies attached (SageMakerFullAccess, S3 access, etc.). "
+        "Example: 'arn:aws:iam::123456789012:role/SageMakerExecutionRole'. "
+        "If not provided, uses the default SageMaker execution role",
+    )
+    volume_size_in_gb: int = Field(
+        30,
+        description="EBS volume size in GB for step execution storage. Must be between "
+        "1-16384 GB. Used for temporary files, model artifacts, and data processing. "
+        "Larger volumes needed for big datasets or model training. Example: 30 for "
+        "small jobs, 100+ for large ML training jobs",
+    )
+    max_runtime_in_seconds: int = Field(
+        86400,  # 24 hours
+        description="Maximum execution time in seconds before job termination. Must be "
+        "between 1-432000 seconds (5 days). Used to prevent runaway jobs and control costs. "
+        "Examples: 3600 (1 hour), 86400 (24 hours), 259200 (3 days). "
+        "Consider your longest expected step duration",
+    )
+    tags: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Tags to apply to the Processor/Estimator assigned to the step. "
+        "Example: {'Environment': 'Production', 'Project': 'MLOps'}",
+    )
+    pipeline_tags: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Tags to apply to the pipeline via the "
+        "sagemaker.workflow.pipeline.Pipeline.create method. Example: "
+        "{'Environment': 'Production', 'Project': 'MLOps'}",
+    )
+    keep_alive_period_in_seconds: Optional[int] = Field(
+        300,  # 5 minutes
+        description="The time in seconds after which the provisioned instance "
+        "will be terminated if not used. This is only applicable for "
+        "TrainingStep type.",
+    )
+    use_training_step: Optional[bool] = Field(
+        None,
+        description="Whether to use the TrainingStep type. It is not possible "
+        "to use TrainingStep type if the `output_data_s3_uri` is set to "
+        "Dict[str, str] or if the `output_data_s3_mode` != 'EndOfJob'.",
+    )
 
-    processor_args: Dict[str, Any] = {}
-    estimator_args: Dict[str, Any] = {}
-    environment: Dict[str, str] = {}
+    processor_args: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments that are directly passed to the SageMaker "
+        "Processor for a specific step, allowing for overriding the default "
+        "settings provided when configuring the component. Example: "
+        "{'instance_count': 2, 'base_job_name': 'my-processing-job'}",
+    )
+    estimator_args: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments that are directly passed to the SageMaker "
+        "Estimator for a specific step, allowing for overriding the default "
+        "settings provided when configuring the component. Example: "
+        "{'train_instance_count': 2, 'train_max_run': 3600}",
+    )
+    environment: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Environment variables to pass to the container. "
+        "Example: {'LOG_LEVEL': 'INFO', 'DEBUG_MODE': 'False'}",
+    )
 
-    input_data_s3_mode: str = "File"
+    input_data_s3_mode: str = Field(
+        "File",
+        description="How data is made available to the container. "
+        "Two possible input modes: File, Pipe.",
+    )
     input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
-        default=None, union_mode="left_to_right"
+        default=None,
+        union_mode="left_to_right",
+        description="S3 URI where data is located if not locally. Example string: "
+        "'s3://my-bucket/my-data/train'. Example dict: "
+        "{'training': 's3://bucket/train', 'validation': 's3://bucket/val'}",
     )
 
-    output_data_s3_mode: str = DEFAULT_OUTPUT_DATA_S3_MODE
+    output_data_s3_mode: str = Field(
+        DEFAULT_OUTPUT_DATA_S3_MODE,
+        description="How data is uploaded to the S3 bucket. "
+        "Two possible output modes: EndOfJob, Continuous.",
+    )
     output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
-        default=None, union_mode="left_to_right"
+        default=None,
+        union_mode="left_to_right",
+        description="S3 URI where data is uploaded after or during processing run. "
+        "Example string: 's3://my-bucket/my-data/output'. Example dict: "
+        "{'output_one': 's3://bucket/out1', 'output_two': 's3://bucket/out2'}",
+    )
+    processor_role: Optional[str] = Field(
+        None,
+        description="DEPRECATED: use `execution_role` instead. "
+        "The IAM role to use for the step execution.",
+    )
+    processor_tags: Optional[Dict[str, str]] = Field(
+        None,
+        description="DEPRECATED: use `tags` instead. "
+        "Tags to apply to the Processor assigned to the step.",
     )
-
-    processor_role: Optional[str] = None
-    processor_tags: Optional[Dict[str, str]] = None
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
         ("processor_role", "execution_role"), ("processor_tags", "tags")
     )
@@ -184,39 +207,49 @@ class SagemakerOrchestratorConfig(
         `aws_secret_access_key`, and optional `aws_auth_role_arn`,
     - If none of the above are provided, unspecified credentials will be
         loaded from the default AWS config.
-
-    Attributes:
-        execution_role: The IAM role ARN to use for the pipeline.
-        scheduler_role: The ARN of the IAM role that will be assumed by
-            the EventBridge service to launch Sagemaker pipelines
-            (For more details regarding the required permissions, please check:
-            https://docs.zenml.io/stacks/stack-components/orchestrators/sagemaker#required-iam-permissions-for-schedules)
-        aws_access_key_id: The AWS access key ID to use to authenticate to AWS.
-            If not provided, the value from the default AWS config will be used.
-        aws_secret_access_key: The AWS secret access key to use to authenticate
-            to AWS. If not provided, the value from the default AWS config will
-            be used.
-        aws_profile: The AWS profile to use for authentication if not using
-            service connectors or explicit credentials. If not provided, the
-            default profile will be used.
-        aws_auth_role_arn: The ARN of an intermediate IAM role to assume when
-            authenticating to AWS.
-        region: The AWS region where the processing job will be run. If not
-            provided, the value from the default AWS config will be used.
-        bucket: Name of the S3 bucket to use for storing artifacts
-            from the job run. If not provided, a default bucket will be created
-            based on the following format:
-            "sagemaker-{region}-{aws-account-id}".
     """
 
-    execution_role: str
-    scheduler_role: Optional[str] = None
-    aws_access_key_id: Optional[str] = SecretField(default=None)
-    aws_secret_access_key: Optional[str] = SecretField(default=None)
-    aws_profile: Optional[str] = None
-    aws_auth_role_arn: Optional[str] = None
-    region: Optional[str] = None
-    bucket: Optional[str] = None
+    execution_role: str = Field(
+        ..., description="The IAM role ARN to use for the pipeline."
+    )
+    scheduler_role: Optional[str] = Field(
+        None,
+        description="The ARN of the IAM role that will be assumed by "
+        "the EventBridge service to launch Sagemaker pipelines. "
+        "Required for scheduled pipelines.",
+    )
+    aws_access_key_id: Optional[str] = SecretField(
+        default=None,
+        description="The AWS access key ID to use to authenticate to AWS. "
+        "If not provided, the value from the default AWS config will be used.",
+    )
+    aws_secret_access_key: Optional[str] = SecretField(
+        default=None,
+        description="The AWS secret access key to use to authenticate to AWS. "
+        "If not provided, the value from the default AWS config will be used.",
+    )
+    aws_profile: Optional[str] = Field(
+        None,
+        description="The AWS profile to use for authentication if not using "
+        "service connectors or explicit credentials. If not provided, the "
+        "default profile will be used.",
+    )
+    aws_auth_role_arn: Optional[str] = Field(
+        None,
+        description="The ARN of an intermediate IAM role to assume when "
+        "authenticating to AWS.",
+    )
+    region: Optional[str] = Field(
+        None,
+        description="The AWS region where the processing job will be run. "
+        "If not provided, the value from the default AWS config will be used.",
+    )
+    bucket: Optional[str] = Field(
+        None,
+        description="Name of the S3 bucket to use for storing artifacts "
+        "from the job run. If not provided, a default bucket will be created "
+        "based on the following format: 'sagemaker-{region}-{aws-account-id}'.",
+    )
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py

@@ -34,38 +34,37 @@ if TYPE_CHECKING:
 
 
 class SagemakerStepOperatorSettings(BaseSettings):
-    """Settings for the Sagemaker step operator.
-
-    Attributes:
-        experiment_name: The name for the experiment to which the job
-            will be associated. If not provided, the job runs would be
-            independent.
-        input_data_s3_uri: S3 URI where training data is located if not locally,
-            e.g. s3://my-bucket/my-data/train. How data will be made available
-            to the container is configured with estimator_args.input_mode. Two possible
-            input types:
-                - str: S3 location where training data is saved.
-                - Dict[str, str]: (ChannelName, S3Location) which represent
-                    channels (e.g. training, validation, testing) where
-                    specific parts of the data are saved in S3.
-        estimator_args: Arguments that are directly passed to the SageMaker
-            Estimator. See
-            https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html#sagemaker.estimator.Estimator
-            for a full list of arguments.
-            For estimator_args.instance_type, check
-            https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html
-            for a list of available instance types.
-        environment: Environment variables to pass to the container.
-
-    """
-
-    instance_type: Optional[str] = None
-    experiment_name: Optional[str] = None
+    """Settings for the Sagemaker step operator."""
+
+    instance_type: Optional[str] = Field(
+        None,
+        description="DEPRECATED: The instance type to use for the step execution. "
+        "Use estimator_args instead. Example: 'ml.m5.xlarge'",
+    )
+    experiment_name: Optional[str] = Field(
+        None,
+        description="The name for the experiment to which the job will be associated. "
+        "If not provided, the job runs would be independent. Example: 'my-training-experiment'",
+    )
     input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
-        default=None, union_mode="left_to_right"
+        default=None,
+        union_mode="left_to_right",
+        description="S3 URI where training data is located if not locally. "
+        "Example string: 's3://my-bucket/my-data/train'. Example dict: "
+        "{'training': 's3://bucket/train', 'validation': 's3://bucket/val'}",
+    )
+    estimator_args: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments that are directly passed to the SageMaker Estimator. "
+        "See SageMaker documentation for available arguments and instance types. Example: "
+        "{'instance_type': 'ml.m5.xlarge', 'instance_count': 1, "
+        "'train_max_run': 3600, 'input_mode': 'File'}",
+    )
+    environment: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Environment variables to pass to the container during execution. "
+        "Example: {'LOG_LEVEL': 'INFO', 'DEBUG_MODE': 'False'}",
     )
-    estimator_args: Dict[str, Any] = {}
-    environment: Dict[str, str] = {}
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
         "instance_type"
@@ -75,18 +74,20 @@ class SagemakerStepOperatorSettings(BaseSettings):
 class SagemakerStepOperatorConfig(
     BaseStepOperatorConfig, SagemakerStepOperatorSettings
 ):
-    """Config for the Sagemaker step operator.
-
-    Attributes:
-        role: The role that has to be assigned to the jobs which are
-            running in Sagemaker.
-        bucket: Name of the S3 bucket to use for storing artifacts
-            from the job run. If not provided, a default bucket will be created
-            based on the following format: "sagemaker-{region}-{aws-account-id}".
-    """
-
-    role: str
-    bucket: Optional[str] = None
+    """Config for the Sagemaker step operator."""
+
+    role: str = Field(
+        ...,
+        description="The IAM role ARN that has to be assigned to the jobs "
+        "running in SageMaker. This role must have the necessary permissions "
+        "to access SageMaker and S3 resources.",
+    )
+    bucket: Optional[str] = Field(
+        None,
+        description="Name of the S3 bucket to use for storing artifacts from the job run. "
+        "If not provided, a default bucket will be created based on the format: "
+        "'sagemaker-{region}-{aws-account-id}'.",
+    )
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -807,14 +807,20 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             settings=settings,
         )
 
-    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+    def fetch_status(
+        self, run: "PipelineRunResponse", include_steps: bool = False
+    ) -> Tuple[
+        Optional[ExecutionStatus], Optional[Dict[str, ExecutionStatus]]
+    ]:
         """Refreshes the status of a specific pipeline run.
 
         Args:
             run: The run that was executed by this orchestrator.
+            include_steps: Whether to fetch steps
 
         Returns:
-            the actual status of the pipeline job.
+            A tuple of (pipeline_status, step_statuses_dict).
+            Step statuses are not supported for SageMaker, so step_statuses_dict will always be None.
 
         Raises:
             AssertionError: If the run was not executed by to this orchestrator.
@@ -855,18 +861,21 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         # Map the potential outputs to ZenML ExecutionStatus. Potential values:
         # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribePipelineExecution.html
         if status == "Executing":
-            return ExecutionStatus.RUNNING
+            pipeline_status = ExecutionStatus.RUNNING
         elif status == "Stopping":
-            return ExecutionStatus.STOPPING
+            pipeline_status = ExecutionStatus.STOPPING
         elif status == "Stopped":
-            return ExecutionStatus.STOPPED
+            pipeline_status = ExecutionStatus.STOPPED
         elif status == "Failed":
-            return ExecutionStatus.FAILED
+            pipeline_status = ExecutionStatus.FAILED
         elif status == "Succeeded":
-            return ExecutionStatus.COMPLETED
+            pipeline_status = ExecutionStatus.COMPLETED
         else:
             raise ValueError("Unknown status for the pipeline execution.")
 
+        # SageMaker doesn't support step-level status fetching yet
+        return pipeline_status, None
+
     def compute_metadata(
         self,
         execution_arn: str,