zenml-nightly 0.66.0.dev20240923__py3-none-any.whl → 0.66.0.dev20240927__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.66.0.dev20240923
+0.66.0.dev20240927

zenml/cli/__init__.py

@@ -1715,6 +1715,13 @@ To delete a pipeline run, use:
 zenml pipeline runs delete <PIPELINE_RUN_NAME_OR_ID>
 ```
 
+To refresh the status of a pipeline run, you can use the `refresh` command (
+only supported for pipelines executed on Vertex, Sagemaker or AzureML).
+
+```bash
+zenml pipeline runs refresh <PIPELINE_RUN_NAME_OR_ID>
+```
+
 If you run any of your pipelines with `pipeline.run(schedule=...)`, ZenML keeps
 track of the schedule and you can list all schedules via:
 

zenml/cli/base.py

@@ -83,11 +83,11 @@ ZENML_PROJECT_TEMPLATES = dict(
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",
-        github_tag="2024.08.28",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.09.23",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="2024.08.29",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.09.23",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     llm_finetuning=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-llm-finetuning",

zenml/cli/pipeline.py

@@ -500,6 +500,27 @@ def delete_pipeline_run(
         cli_utils.declare(f"Deleted pipeline run '{run_name_or_id}'.")
 
 
+@runs.command("refresh")
+@click.argument("run_name_or_id", type=str, required=True)
+def refresh_pipeline_run(run_name_or_id: str) -> None:
+    """Refresh the status of a pipeline run.
+
+    Args:
+        run_name_or_id: The name or ID of the pipeline run to refresh.
+    """
+    try:
+        # Fetch and update the run
+        run = Client().get_pipeline_run(name_id_or_prefix=run_name_or_id)
+        run.refresh_run_status()
+
+    except KeyError as e:
+        cli_utils.error(str(e))
+    else:
+        cli_utils.declare(
+            f"Refreshed the status of pipeline run '{run.name}'."
+        )
+
+
 @pipeline.group()
 def builds() -> None:
     """Commands for pipeline builds."""

zenml/cli/utils.py

@@ -76,6 +76,7 @@ from zenml.models import (
     StrFilter,
     UUIDFilter,
 )
+from zenml.models.v2.base.filter import FilterGenerator
 from zenml.services import BaseService, ServiceState
 from zenml.stack import StackComponent
 from zenml.stack.stack_component import StackComponentConfig
@@ -2477,12 +2478,13 @@ def create_filter_help_text(filter_model: Type[BaseFilter], field: str) -> str:
     Returns:
         The help text.
     """
-    if filter_model.is_sort_by_field(field):
+    filter_generator = FilterGenerator(filter_model)
+    if filter_generator.is_sort_by_field(field):
         return (
             "[STRING] Example: --sort_by='desc:name' to sort by name in "
             "descending order. "
         )
-    if filter_model.is_datetime_field(field):
+    if filter_generator.is_datetime_field(field):
         return (
             f"[DATETIME] The following datetime format is supported: "
             f"'{FILTERING_DATETIME_FORMAT}'. Make sure to keep it in "
@@ -2491,23 +2493,23 @@ def create_filter_help_text(filter_model: Type[BaseFilter], field: str) -> str:
             f"'{GenericFilterOps.GTE}:{FILTERING_DATETIME_FORMAT}' to "
             f"filter for everything created on or after the given date."
         )
-    elif filter_model.is_uuid_field(field):
+    elif filter_generator.is_uuid_field(field):
         return (
             f"[UUID] Example: --{field}='{GenericFilterOps.STARTSWITH}:ab53ca' "
             f"to filter for all UUIDs starting with that prefix."
         )
-    elif filter_model.is_int_field(field):
+    elif filter_generator.is_int_field(field):
         return (
             f"[INTEGER] Example: --{field}='{GenericFilterOps.GTE}:25' to "
             f"filter for all entities where this field has a value greater than "
             f"or equal to the value."
         )
-    elif filter_model.is_bool_field(field):
+    elif filter_generator.is_bool_field(field):
         return (
             f"[BOOL] Example: --{field}='True' to "
             f"filter for all instances where this field is true."
         )
-    elif filter_model.is_str_field(field):
+    elif filter_generator.is_str_field(field):
         return (
             f"[STRING] Example: --{field}='{GenericFilterOps.CONTAINS}:example' "
             f"to filter everything that contains the query string somewhere in "
@@ -2529,27 +2531,28 @@ def create_data_type_help_text(
     Returns:
         The help text.
     """
-    if filter_model.is_datetime_field(field):
+    filter_generator = FilterGenerator(filter_model)
+    if filter_generator.is_datetime_field(field):
         return (
             f"[DATETIME] supported filter operators: "
             f"{[str(op) for op in NumericFilter.ALLOWED_OPS]}"
         )
-    elif filter_model.is_uuid_field(field):
+    elif filter_generator.is_uuid_field(field):
         return (
             f"[UUID] supported filter operators: "
             f"{[str(op) for op in UUIDFilter.ALLOWED_OPS]}"
         )
-    elif filter_model.is_int_field(field):
+    elif filter_generator.is_int_field(field):
         return (
             f"[INTEGER] supported filter operators: "
             f"{[str(op) for op in NumericFilter.ALLOWED_OPS]}"
         )
-    elif filter_model.is_bool_field(field):
+    elif filter_generator.is_bool_field(field):
         return (
             f"[BOOL] supported filter operators: "
             f"{[str(op) for op in BoolFilter.ALLOWED_OPS]}"
         )
-    elif filter_model.is_str_field(field):
+    elif filter_generator.is_str_field(field):
         return (
             f"[STRING] supported filter operators: "
             f"{[str(op) for op in StrFilter.ALLOWED_OPS]}"

zenml/client.py

@@ -1236,6 +1236,8 @@ class Client(metaclass=ClientMetaClass):
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
         component_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
+        component: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[StackResponse]:
         """Lists all stacks.
@@ -1252,6 +1254,8 @@ class Client(metaclass=ClientMetaClass):
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
             component_id: The id of the component to filter by.
+            user: The name/ID of the user to filter by.
+            component: The name/ID of the component to filter by.
             name: The name of the stack to filter by.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
@@ -1267,6 +1271,8 @@ class Client(metaclass=ClientMetaClass):
             workspace_id=workspace_id,
             user_id=user_id,
             component_id=component_id,
+            user=user,
+            component=component,
             name=name,
             description=description,
             id=id,
@@ -2348,8 +2354,10 @@ class Client(metaclass=ClientMetaClass):
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
         name: Optional[str] = None,
+        latest_run_status: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         tag: Optional[str] = None,
         hydrate: bool = False,
     ) -> Page[PipelineResponse]:
@@ -2364,8 +2372,11 @@ class Client(metaclass=ClientMetaClass):
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
             name: The name of the pipeline to filter by.
+            latest_run_status: Filter by the status of the latest run of a
+                pipeline.
             workspace_id: The id of the workspace to filter by.
             user_id: The id of the user to filter by.
+            user: The name/ID of the user to filter by.
             tag: Tag to filter by.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
@@ -2382,8 +2393,10 @@ class Client(metaclass=ClientMetaClass):
             created=created,
             updated=updated,
             name=name,
+            latest_run_status=latest_run_status,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             tag=tag,
         )
         pipeline_filter_model.set_scope_workspace(self.active_workspace.id)
@@ -3467,6 +3480,9 @@ class Client(metaclass=ClientMetaClass):
         build_id: Optional[Union[str, UUID]] = None,
         stack_id: Optional[Union[str, UUID]] = None,
         code_repository_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
+        pipeline: Optional[Union[UUID, str]] = None,
+        stack: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[RunTemplateResponse]:
         """Get a page of run templates.
@@ -3486,6 +3502,9 @@ class Client(metaclass=ClientMetaClass):
             build_id: Filter by build ID.
             stack_id: Filter by stack ID.
             code_repository_id: Filter by code repository ID.
+            user: Filter by user name/ID.
+            pipeline: Filter by pipeline name/ID.
+            stack: Filter by stack name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -3507,6 +3526,9 @@ class Client(metaclass=ClientMetaClass):
             build_id=build_id,
             stack_id=stack_id,
             code_repository_id=code_repository_id,
+            user=user,
+            pipeline=pipeline,
+            stack=stack,
         )
 
         return self.zen_store.list_run_templates(
@@ -3742,6 +3764,7 @@ class Client(metaclass=ClientMetaClass):
         deployment_id: Optional[Union[str, UUID]] = None,
         code_repository_id: Optional[Union[str, UUID]] = None,
         template_id: Optional[Union[str, UUID]] = None,
+        model_version_id: Optional[Union[str, UUID]] = None,
         orchestrator_run_id: Optional[str] = None,
         status: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
@@ -3750,6 +3773,11 @@ class Client(metaclass=ClientMetaClass):
         unlisted: Optional[bool] = None,
         templatable: Optional[bool] = None,
         tag: Optional[str] = None,
+        user: Optional[Union[UUID, str]] = None,
+        pipeline: Optional[Union[UUID, str]] = None,
+        code_repository: Optional[Union[UUID, str]] = None,
+        model: Optional[Union[UUID, str]] = None,
+        stack: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[PipelineRunResponse]:
         """List all pipeline runs.
@@ -3764,7 +3792,8 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             pipeline_id: The id of the pipeline to filter by.
-            pipeline_name: The name of the pipeline to filter by.
+            pipeline_name: DEPRECATED. Use `pipeline` instead to filter by
+                pipeline name.
             user_id: The id of the user to filter by.
             stack_id: The id of the stack to filter by.
             schedule_id: The id of the schedule to filter by.
@@ -3772,6 +3801,7 @@ class Client(metaclass=ClientMetaClass):
             deployment_id: The id of the deployment to filter by.
             code_repository_id: The id of the code repository to filter by.
             template_id: The ID of the template to filter by.
+            model_version_id: The ID of the model version to filter by.
             orchestrator_run_id: The run id of the orchestrator to filter by.
             name: The name of the run to filter by.
             status: The status of the pipeline run
@@ -3781,6 +3811,11 @@ class Client(metaclass=ClientMetaClass):
             unlisted: If the runs should be unlisted or not.
             templatable: If the runs should be templatable or not.
             tag: Tag to filter by.
+            user: The name/ID of the user to filter by.
+            pipeline: The name/ID of the pipeline to filter by.
+            code_repository: Filter by code repository name/ID.
+            model: Filter by model name/ID.
+            stack: Filter by stack name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -3804,6 +3839,7 @@ class Client(metaclass=ClientMetaClass):
             deployment_id=deployment_id,
             code_repository_id=code_repository_id,
             template_id=template_id,
+            model_version_id=model_version_id,
             orchestrator_run_id=orchestrator_run_id,
             user_id=user_id,
             stack_id=stack_id,
@@ -3813,6 +3849,11 @@ class Client(metaclass=ClientMetaClass):
             num_steps=num_steps,
             tag=tag,
             unlisted=unlisted,
+            user=user,
+            pipeline=pipeline,
+            code_repository=code_repository,
+            stack=stack,
+            model=model,
             templatable=templatable,
         )
         runs_filter_model.set_scope_workspace(self.active_workspace.id)
@@ -3892,6 +3933,7 @@ class Client(metaclass=ClientMetaClass):
         original_step_run_id: Optional[Union[str, UUID]] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        model_version_id: Optional[Union[str, UUID]] = None,
         num_outputs: Optional[Union[int, str]] = None,
         hydrate: bool = False,
     ) -> Page[StepRunResponse]:
@@ -3911,6 +3953,7 @@ class Client(metaclass=ClientMetaClass):
             user_id: The  id of the user to filter by.
             pipeline_run_id: The  id of the pipeline run to filter by.
             original_step_run_id: The  id of the pipeline run to filter by.
+            model_version_id: The ID of the model version to filter by.
             name: The name of the run to filter by.
             entrypoint_name: The entrypoint_name of the run to filter by.
             code_hash: The code_hash of the run to filter by.
@@ -3942,6 +3985,7 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             workspace_id=workspace_id,
             user_id=user_id,
+            model_version_id=model_version_id,
             num_outputs=num_outputs,
         )
         step_run_filter_model.set_scope_workspace(self.active_workspace.id)
@@ -4150,8 +4194,11 @@ class Client(metaclass=ClientMetaClass):
         user_id: Optional[Union[str, UUID]] = None,
         only_unused: Optional[bool] = False,
         has_custom_name: Optional[bool] = None,
-        hydrate: bool = False,
+        user: Optional[Union[UUID, str]] = None,
+        model: Optional[Union[UUID, str]] = None,
+        pipeline_run: Optional[Union[UUID, str]] = None,
         tag: Optional[str] = None,
+        hydrate: bool = False,
     ) -> Page[ArtifactVersionResponse]:
         """Get a list of artifact versions.
 
@@ -4177,9 +4224,12 @@ class Client(metaclass=ClientMetaClass):
             only_unused: Only return artifact versions that are not used in
                 any pipeline runs.
             has_custom_name: Filter artifacts with/without custom names.
+            tag: A tag to filter by.
+            user: Filter by user name or ID.
+            model: Filter by model name or ID.
+            pipeline_run: Filter by pipeline run name or ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
-            tag: A tag to filter by.
 
         Returns:
             A list of artifact versions.
@@ -4206,6 +4256,9 @@ class Client(metaclass=ClientMetaClass):
             only_unused=only_unused,
             has_custom_name=has_custom_name,
             tag=tag,
+            user=user,
+            model=model,
+            pipeline_run=pipeline_run,
         )
         artifact_version_filter_model.set_scope_workspace(
             self.active_workspace.id
@@ -6102,6 +6155,7 @@ class Client(metaclass=ClientMetaClass):
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
         name: Optional[str] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
         tag: Optional[str] = None,
     ) -> Page[ModelResponse]:
@@ -6115,6 +6169,7 @@ class Client(metaclass=ClientMetaClass):
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
             name: The name of the model to filter by.
+            user: Filter by user name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
             tag: The tag of the model to filter by.
@@ -6131,6 +6186,7 @@ class Client(metaclass=ClientMetaClass):
             created=created,
             updated=updated,
             tag=tag,
+            user=user,
         )
 
         return self.zen_store.list_models(
@@ -6310,6 +6366,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         number: Optional[int] = None,
         stage: Optional[Union[str, ModelStages]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
         tag: Optional[str] = None,
     ) -> Page[ModelVersionResponse]:
@@ -6327,6 +6384,7 @@ class Client(metaclass=ClientMetaClass):
             name: name or id of the model version.
             number: number of the model version.
             stage: stage of the model version.
+            user: Filter by user name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
             tag: The tag to filter by.
@@ -6345,6 +6403,7 @@ class Client(metaclass=ClientMetaClass):
             number=number,
             stage=stage,
             tag=tag,
+            user=user,
         )
 
         return self.zen_store.list_model_versions(
@@ -6422,6 +6481,7 @@ class Client(metaclass=ClientMetaClass):
         only_model_artifacts: Optional[bool] = None,
         only_deployment_artifacts: Optional[bool] = None,
         has_custom_name: Optional[bool] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[ModelVersionArtifactResponse]:
         """Get model version to artifact links by filter in Model Control Plane.
@@ -6443,6 +6503,7 @@ class Client(metaclass=ClientMetaClass):
             only_model_artifacts: Use to filter by model artifacts
             only_deployment_artifacts: Use to filter by deployment artifacts
             has_custom_name: Filter artifacts with/without custom names.
+            user: Filter by user name/ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -6467,6 +6528,7 @@ class Client(metaclass=ClientMetaClass):
                 only_model_artifacts=only_model_artifacts,
                 only_deployment_artifacts=only_deployment_artifacts,
                 has_custom_name=has_custom_name,
+                user=user,
             ),
             hydrate=hydrate,
         )
@@ -6536,6 +6598,7 @@ class Client(metaclass=ClientMetaClass):
         model_version_id: Optional[Union[UUID, str]] = None,
         pipeline_run_id: Optional[Union[UUID, str]] = None,
         pipeline_run_name: Optional[str] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[ModelVersionPipelineRunResponse]:
         """Get all model version to pipeline run links by filter.
@@ -6553,6 +6616,7 @@ class Client(metaclass=ClientMetaClass):
             model_version_id: Use the model version id for filtering
             pipeline_run_id: Use the pipeline run id for filtering
             pipeline_run_name: Use the pipeline run name for filtering
+            user: Filter by user name or ID.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response
 
@@ -6573,6 +6637,7 @@ class Client(metaclass=ClientMetaClass):
                 model_version_id=model_version_id,
                 pipeline_run_id=pipeline_run_id,
                 pipeline_run_name=pipeline_run_name,
+                user=user,
             ),
             hydrate=hydrate,
         )

zenml/config/step_configurations.py

@@ -137,7 +137,6 @@ class ArtifactConfiguration(PartialArtifactConfiguration):
 class StepConfigurationUpdate(StrictBaseModel):
     """Class for step configuration updates."""
 
-    name: Optional[str] = None
     enable_cache: Optional[bool] = None
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
@@ -154,10 +153,6 @@ class StepConfigurationUpdate(StrictBaseModel):
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
-    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
-        "name"
-    )
-
 
 class PartialStepConfiguration(StepConfigurationUpdate):
     """Class representing a partial step configuration."""

zenml/constants.py

@@ -364,6 +364,7 @@ PIPELINE_DEPLOYMENTS = "/pipeline_deployments"
 PIPELINES = "/pipelines"
 PIPELINE_SPEC = "/pipeline-spec"
 PLUGIN_FLAVORS = "/plugin-flavors"
+REFRESH = "/refresh"
 RUNS = "/runs"
 RUN_TEMPLATES = "/run_templates"
 RUN_METADATA = "/run-metadata"
@@ -430,6 +431,8 @@ SORT_PIPELINES_BY_LATEST_RUN_KEY = "latest_run"
 
 # Metadata constants
 METADATA_ORCHESTRATOR_URL = "orchestrator_url"
+METADATA_ORCHESTRATOR_LOGS_URL = "orchestrator_logs_url"
+METADATA_ORCHESTRATOR_RUN_ID = "orchestrator_run_id"
 METADATA_EXPERIMENT_TRACKER_URL = "experiment_tracker_url"
 METADATA_DEPLOYED_MODEL_URL = "deployed_model_url"
 

zenml/enums.py

@@ -244,6 +244,7 @@ class GenericFilterOps(StrEnum):
     """Ops for all filters for string values on list methods."""
 
     EQUALS = "equals"
+    NOT_EQUALS = "notequals"
     CONTAINS = "contains"
     STARTSWITH = "startswith"
     ENDSWITH = "endswith"
@@ -251,6 +252,7 @@ class GenericFilterOps(StrEnum):
     GT = "gt"
     LTE = "lte"
     LT = "lt"
+    IN = "in"
 
 
 class SorterOps(StrEnum):

zenml/integrations/__init__.py

@@ -69,6 +69,7 @@ from zenml.integrations.skypilot_aws import SkypilotAWSIntegration  # noqa
 from zenml.integrations.skypilot_gcp import SkypilotGCPIntegration  # noqa
 from zenml.integrations.skypilot_azure import SkypilotAzureIntegration  # noqa
 from zenml.integrations.skypilot_lambda import SkypilotLambdaIntegration  # noqa
+from zenml.integrations.skypilot_kubernetes import SkypilotKubernetesIntegration  # noqa
 from zenml.integrations.slack import SlackIntegration  # noqa
 from zenml.integrations.spark import SparkIntegration  # noqa
 from zenml.integrations.tekton import TektonIntegration  # noqa

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.aws import (
@@ -25,23 +25,38 @@ from zenml.integrations.aws import (
 from zenml.models import ServiceConnectorRequirements
 from zenml.orchestrators import BaseOrchestratorConfig
 from zenml.orchestrators.base_orchestrator import BaseOrchestratorFlavor
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
     from zenml.integrations.aws.orchestrators import SagemakerOrchestrator
 
+DEFAULT_TRAINING_INSTANCE_TYPE = "ml.m5.xlarge"
+DEFAULT_PROCESSING_INSTANCE_TYPE = "ml.t3.medium"
+DEFAULT_OUTPUT_DATA_S3_MODE = "EndOfJob"
+
 
 class SagemakerOrchestratorSettings(BaseSettings):
     """Settings for the Sagemaker orchestrator.
 
     Attributes:
         instance_type: The instance type to use for the processing job.
-        processor_role: The IAM role to use for the step execution on a Processor.
+        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.
-        processor_tags: Tags to apply to the Processor assigned to the step.
+        tags: Tags to apply to the Processor/Estimator assigned to the step.
+        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
@@ -50,6 +65,13 @@ class SagemakerOrchestratorSettings(BaseSettings):
             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.
+        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,
@@ -74,23 +96,70 @@ class SagemakerOrchestratorSettings(BaseSettings):
                     Data must be available locally in /opt/ml/processing/output/data/<ChannelName>.
     """
 
-    instance_type: str = "ml.t3.medium"
-    processor_role: Optional[str] = None
+    instance_type: Optional[str] = None
+    execution_role: Optional[str] = None
     volume_size_in_gb: int = 30
     max_runtime_in_seconds: int = 86400
-    processor_tags: Dict[str, str] = {}
+    tags: Dict[str, str] = {}
+    keep_alive_period_in_seconds: Optional[int] = 300  # 5 minutes
+    use_training_step: Optional[bool] = None
 
     processor_args: Dict[str, Any] = {}
+    estimator_args: Dict[str, Any] = {}
+
     input_data_s3_mode: str = "File"
     input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
         default=None, union_mode="left_to_right"
     )
 
-    output_data_s3_mode: str = "EndOfJob"
+    output_data_s3_mode: str = DEFAULT_OUTPUT_DATA_S3_MODE
     output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
         default=None, union_mode="left_to_right"
     )
 
+    processor_role: Optional[str] = None
+    processor_tags: Dict[str, str] = {}
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("processor_role", "execution_role"), ("processor_tags", "tags")
+    )
+
+    @model_validator(mode="before")
+    def validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Check if model is configured correctly.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The validated model data.
+
+        Raises:
+            ValueError: If the model is configured incorrectly.
+        """
+        use_training_step = data.get("use_training_step", True)
+        output_data_s3_uri = data.get("output_data_s3_uri", None)
+        output_data_s3_mode = data.get(
+            "output_data_s3_mode", DEFAULT_OUTPUT_DATA_S3_MODE
+        )
+        if use_training_step and (
+            isinstance(output_data_s3_uri, dict)
+            or (
+                isinstance(output_data_s3_uri, str)
+                and (output_data_s3_mode != DEFAULT_OUTPUT_DATA_S3_MODE)
+            )
+        ):
+            raise ValueError(
+                "`use_training_step=True` is not supported when `output_data_s3_uri` is a dict or "
+                f"when `output_data_s3_mode` is not '{DEFAULT_OUTPUT_DATA_S3_MODE}'."
+            )
+        instance_type = data.get("instance_type", None)
+        if instance_type is None:
+            if use_training_step:
+                data["instance_type"] = DEFAULT_TRAINING_INSTANCE_TYPE
+            else:
+                data["instance_type"] = DEFAULT_PROCESSING_INSTANCE_TYPE
+        return data
+
 
 class SagemakerOrchestratorConfig(
     BaseOrchestratorConfig, SagemakerOrchestratorSettings