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

zenml/VERSION

@@ -1 +1 @@
-0.83.1.dev20250710
+0.83.1.dev20250711

zenml/artifact_stores/base_artifact_store.py

@@ -33,11 +33,15 @@ from typing import (
     cast,
 )
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
-from zenml.constants import ENV_ZENML_SERVER
+from zenml.constants import (
+    ENV_ZENML_SERVER,
+    ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS,
+    handle_bool_env_var,
+)
 from zenml.enums import StackComponentType
-from zenml.exceptions import ArtifactStoreInterfaceError
+from zenml.exceptions import ArtifactStoreInterfaceError, IllegalOperationError
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent, StackComponentConfig
@@ -73,6 +77,12 @@ class _sanitize_paths:
         """
         self.func = func
         self.fixed_root_path = fixed_root_path
+        if ENV_ZENML_SERVER in os.environ:
+            self.allow_local_file_access = handle_bool_env_var(
+                ENV_ZENML_SERVER_ALLOW_LOCAL_FILE_ACCESS, False
+            )
+        else:
+            self.allow_local_file_access = True
 
         self.path_args: List[int] = []
         self.path_kwargs: List[str] = []
@@ -93,7 +103,15 @@ class _sanitize_paths:
         Raises:
             FileNotFoundError: If the path is outside of the artifact store
                 bounds.
+            IllegalOperationError: If the path is a local file and the server
+                is not configured to allow local file access.
         """
+        if not self.allow_local_file_access and not io_utils.is_remote(path):
+            raise IllegalOperationError(
+                "Files in a local artifact store cannot be accessed from the "
+                "server."
+            )
+
         if not path.startswith(self.fixed_root_path):
             raise FileNotFoundError(
                 f"File `{path}` is outside of "
@@ -169,9 +187,18 @@ class _sanitize_paths:
 
 
 class BaseArtifactStoreConfig(StackComponentConfig):
-    """Config class for `BaseArtifactStore`."""
+    """Config class for `BaseArtifactStore`.
 
-    path: str
+    Base configuration for artifact storage backends.
+    Field descriptions are defined inline using Field() descriptors.
+    """
+
+    path: str = Field(
+        description="Root path for artifact storage. Must be a valid URI supported by the "
+        "specific artifact store implementation. Examples: 's3://my-bucket/artifacts', "
+        "'/local/storage/path', 'gs://bucket-name/zenml-artifacts', 'azure://container/path'. "
+        "Path must be accessible with the configured credentials and permissions"
+    )
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]]
     IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
@@ -435,40 +462,41 @@ class BaseArtifactStore(StackComponent):
             **kwargs: The keyword arguments to pass to the Pydantic object.
         """
         super(BaseArtifactStore, self).__init__(*args, **kwargs)
+        self._add_path_sanitization()
 
         # If running in a ZenML server environment, we don't register
         # the filesystems. We always use the artifact stores directly.
         if ENV_ZENML_SERVER not in os.environ:
             self._register()
 
+    def _add_path_sanitization(self) -> None:
+        """Add path sanitization to the artifact store."""
+        for method_name, method in inspect.getmembers(BaseArtifactStore):
+            if getattr(method, "__isabstractmethod__", False):
+                method_implementation = getattr(self, method_name)
+                sanitized_method = _sanitize_paths(
+                    method_implementation, self.path
+                )
+                setattr(self, method_name, sanitized_method)
+
     def _register(self) -> None:
         """Create and register a filesystem within the filesystem registry."""
         from zenml.io.filesystem import BaseFilesystem
         from zenml.io.filesystem_registry import default_filesystem_registry
         from zenml.io.local_filesystem import LocalFilesystem
 
-        overloads: Dict[str, Any] = {
-            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
-        }
-        for abc_method in inspect.getmembers(BaseArtifactStore):
-            if getattr(abc_method[1], "__isabstractmethod__", False):
-                sanitized_method = _sanitize_paths(
-                    getattr(self, abc_method[0]), self.path
-                )
-                # prepare overloads for filesystem methods
-                overloads[abc_method[0]] = staticmethod(sanitized_method)
-
-                # decorate artifact store methods
-                setattr(
-                    self,
-                    abc_method[0],
-                    sanitized_method,
-                )
-
         # Local filesystem is always registered, no point in doing it again.
         if isinstance(self, LocalFilesystem):
             return
 
+        overloads: Dict[str, Any] = {
+            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
+        }
+        for method_name, method in inspect.getmembers(BaseArtifactStore):
+            if getattr(method, "__isabstractmethod__", False):
+                method_implementation = getattr(self, method_name)
+                overloads[method_name] = staticmethod(method_implementation)
+
         filesystem_class = type(
             self.__class__.__name__, (BaseFilesystem,), overloads
         )

zenml/artifacts/utils.py

@@ -49,6 +49,7 @@ from zenml.enums import (
 )
 from zenml.exceptions import (
     DoesNotExistException,
+    IllegalOperationError,
     StepContextError,
 )
 from zenml.io import fileio
@@ -925,6 +926,7 @@ def _load_file_from_artifact_store(
         DoesNotExistException: If the file does not exist in the artifact store.
         NotImplementedError: If the artifact store cannot open the file.
         IOError: If the artifact store rejects the request.
+        IllegalOperationError: If the artifact store rejects the request.
     """
     try:
         with artifact_store.open(uri, mode) as text_file:
@@ -946,7 +948,7 @@ def _load_file_from_artifact_store(
             f"File '{uri}' does not exist in artifact store "
             f"'{artifact_store.name}'."
         )
-    except IOError as e:
+    except (IOError, IllegalOperationError) as e:
         raise e
     except Exception as e:
         logger.exception(e)

zenml/cli/pipeline.py

@@ -604,16 +604,27 @@ def delete_pipeline_run(
 
 @runs.command("refresh")
 @click.argument("run_name_or_id", type=str, required=True)
-def refresh_pipeline_run(run_name_or_id: str) -> None:
+@click.option(
+    "--include-steps",
+    is_flag=True,
+    default=False,
+    help="Also refresh the status of individual steps.",
+)
+def refresh_pipeline_run(
+    run_name_or_id: str, include_steps: bool = False
+) -> None:
     """Refresh the status of a pipeline run.
 
     Args:
         run_name_or_id: The name or ID of the pipeline run to refresh.
+        include_steps: If True, also refresh the status of individual steps.
     """
     try:
         # Fetch and update the run
         run = Client().get_pipeline_run(name_id_or_prefix=run_name_or_id)
-        run.refresh_run_status()
+        run_utils.refresh_run_status(
+            run=run, include_step_updates=include_steps
+        )
 
     except KeyError as e:
         cli_utils.error(str(e))

zenml/constants.py

@@ -166,6 +166,9 @@ 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"

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/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: