zenml-nightly 0.83.0.dev20250603__tar.gz → 0.83.0.dev20250605__tar.gz

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.83.0.dev20250603
+Version: 0.83.0.dev20250605
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.83.0.dev20250603"
+version = "0.83.0.dev20250605"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.83.0.dev20250605/src/zenml/VERSION

@@ -0,0 +1 @@
+0.83.0.dev20250605

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/artifacts/utils.py

@@ -37,6 +37,7 @@ from zenml.artifacts.preexisting_data_materializer import (
 )
 from zenml.client import Client
 from zenml.constants import (
+    ENV_ZENML_SERVER,
     MODEL_METADATA_YAML_FILE_NAME,
 )
 from zenml.enums import (
@@ -820,11 +821,18 @@ def _load_artifact_store(
             an artifact store.
         NotImplementedError: If the artifact store could not be loaded.
     """
+    client = Client()
     if isinstance(artifact_store_id, str):
         artifact_store_id = UUID(artifact_store_id)
 
+    if (
+        ENV_ZENML_SERVER not in os.environ
+        and client.active_stack.artifact_store.id == artifact_store_id
+    ):
+        return client.active_stack.artifact_store
+
     if zen_store is None:
-        zen_store = Client().zen_store
+        zen_store = client.zen_store
 
     try:
         artifact_store_model = zen_store.get_stack_component(artifact_store_id)
@@ -857,9 +865,18 @@ def _load_artifact_store(
 def _get_artifact_store_from_response_or_from_active_stack(
     artifact: ArtifactVersionResponse,
 ) -> "BaseArtifactStore":
+    client = Client()
+
     if artifact.artifact_store_id:
+        if (
+            ENV_ZENML_SERVER not in os.environ
+            and client.active_stack.artifact_store.id
+            == artifact.artifact_store_id
+        ):
+            return client.active_stack.artifact_store
+
         try:
-            artifact_store_model = Client().get_stack_component(
+            artifact_store_model = client.get_stack_component(
                 component_type=StackComponentType.ARTIFACT_STORE,
                 name_id_or_prefix=artifact.artifact_store_id,
             )
@@ -881,7 +898,8 @@ def _get_artifact_store_from_response_or_from_active_stack(
                 "the environment that you are loading this artifact from "
                 "has the right dependencies."
             )
-    return Client().active_stack.artifact_store
+
+    return client.active_stack.artifact_store
 
 
 def _load_file_from_artifact_store(

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/client.py

@@ -1480,9 +1480,6 @@ class Client(metaclass=ClientMetaClass):
 
         Returns:
             The model of the active stack for this client.
-
-        Raises:
-            RuntimeError: If the active stack is not set.
         """
         if env_stack_id := os.environ.get(ENV_ZENML_ACTIVE_STACK_ID):
             if not self._active_stack or self._active_stack.id != UUID(
@@ -1509,12 +1506,6 @@ class Client(metaclass=ClientMetaClass):
 
             stack_id = GlobalConfiguration().get_active_stack_id()
 
-        if not stack_id:
-            raise RuntimeError(
-                "No active stack is configured. Run "
-                "`zenml stack set STACK_NAME` to set the active stack."
-            )
-
         return self.get_stack(stack_id)
 
     def activate_stack(
@@ -1550,9 +1541,6 @@ class Client(metaclass=ClientMetaClass):
 
         Args:
             stack: The stack to validate.
-
-        Raises:
-            ValidationError: If the stack configuration is invalid.
         """
         local_components: List[str] = []
         remote_components: List[str] = []
@@ -1617,13 +1605,6 @@ class Client(metaclass=ClientMetaClass):
                 f"local resources."
             )
 
-        if not stack.is_valid:
-            raise ValidationError(
-                "Stack configuration is invalid. A valid"
-                "stack must contain an Artifact Store and "
-                "an Orchestrator."
-            )
-
     # ----------------------------- Services -----------------------------------
 
     def create_service(

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/config/global_config.py

@@ -751,13 +751,15 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         Returns:
             The model of the active project.
         """
-        project_id = self.get_active_project_id()
+        # Initialize the store first. This might create a local database if it
+        # doesn't exist yet, which will then refresh our active project ID.
+        _ = self.zen_store
 
         if self._active_project is not None:
             return self._active_project
 
         project = self.zen_store.get_project(
-            project_name_or_id=project_id,
+            project_name_or_id=self.get_active_project_id(),
         )
         return self.set_active_project(project)
 

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -20,24 +20,35 @@ from typing import (
     Any,
     Dict,
     Iterator,
+    List,
     Optional,
     Tuple,
     Type,
+    Union,
     cast,
 )
 from uuid import UUID
 
 import boto3
-import sagemaker
 from botocore.exceptions import WaiterError
+from sagemaker.estimator import Estimator
+from sagemaker.inputs import TrainingInput
 from sagemaker.network import NetworkConfig
-from sagemaker.processing import ProcessingInput, ProcessingOutput
+from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
+from sagemaker.session import Session
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.execution_variables import (
-    ExecutionVariable,
     ExecutionVariables,
 )
 from sagemaker.workflow.pipeline import Pipeline
-from sagemaker.workflow.steps import ProcessingStep, TrainingStep
+from sagemaker.workflow.step_collections import (
+    StepCollection,
+)
+from sagemaker.workflow.steps import (
+    ProcessingStep,
+    Step,
+    TrainingStep,
+)
 from sagemaker.workflow.triggers import PipelineSchedule
 
 from zenml.client import Client
@@ -215,7 +226,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         """
         return SagemakerOrchestratorSettings
 
-    def _get_sagemaker_session(self) -> sagemaker.Session:
+    def _get_sagemaker_session(self) -> Session:
         """Method to create the sagemaker session with proper authentication.
 
         Returns:
@@ -258,7 +269,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                     aws_session_token=credentials["SessionToken"],
                     region_name=self.config.region,
                 )
-        return sagemaker.Session(
+        return Session(
             boto_session=boto_session, default_bucket=self.config.bucket
         )
 
@@ -324,12 +335,8 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 SagemakerOrchestratorSettings, self.get_settings(step)
             )
 
-            environment[ENV_ZENML_SAGEMAKER_RUN_ID] = (
-                ExecutionVariables.PIPELINE_EXECUTION_ARN
-            )
-
             if step_settings.environment:
-                step_environment = step_settings.environment.copy()
+                split_environment = step_settings.environment.copy()
                 # Sagemaker does not allow environment variables longer than 256
                 # characters to be passed to Processor steps. If an environment variable
                 # is longer than 256 characters, we split it into multiple environment
@@ -337,9 +344,9 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 # custom entrypoint configuration.
                 split_environment_variables(
                     size_limit=SAGEMAKER_PROCESSOR_STEP_ENV_VAR_SIZE_LIMIT,
-                    env=step_environment,
+                    env=split_environment,
                 )
-                environment.update(step_environment)
+                environment.update(split_environment)
 
             use_training_step = (
                 step_settings.use_training_step
@@ -421,28 +428,51 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                     )
 
             # Construct S3 inputs to container for step
-            inputs = None
+            training_inputs: Optional[
+                Union[TrainingInput, Dict[str, TrainingInput]]
+            ] = None
+            processing_inputs: Optional[List[ProcessingInput]] = None
 
             if step_settings.input_data_s3_uri is None:
                 pass
             elif isinstance(step_settings.input_data_s3_uri, str):
-                inputs = [
-                    ProcessingInput(
-                        source=step_settings.input_data_s3_uri,
-                        destination="/opt/ml/processing/input/data",
-                        s3_input_mode=step_settings.input_data_s3_mode,
+                if use_training_step:
+                    training_inputs = TrainingInput(
+                        s3_data=step_settings.input_data_s3_uri,
+                        input_mode=step_settings.input_data_s3_mode,
                     )
-                ]
-            elif isinstance(step_settings.input_data_s3_uri, dict):
-                inputs = []
-                for channel, s3_uri in step_settings.input_data_s3_uri.items():
-                    inputs.append(
+                else:
+                    processing_inputs = [
                         ProcessingInput(
-                            source=s3_uri,
-                            destination=f"/opt/ml/processing/input/data/{channel}",
+                            source=step_settings.input_data_s3_uri,
+                            destination="/opt/ml/processing/input/data",
                             s3_input_mode=step_settings.input_data_s3_mode,
                         )
-                    )
+                    ]
+            elif isinstance(step_settings.input_data_s3_uri, dict):
+                if use_training_step:
+                    training_inputs = {}
+                    for (
+                        channel,
+                        s3_uri,
+                    ) in step_settings.input_data_s3_uri.items():
+                        training_inputs[channel] = TrainingInput(
+                            s3_data=s3_uri,
+                            input_mode=step_settings.input_data_s3_mode,
+                        )
+                else:
+                    processing_inputs = []
+                    for (
+                        channel,
+                        s3_uri,
+                    ) in step_settings.input_data_s3_uri.items():
+                        processing_inputs.append(
+                            ProcessingInput(
+                                source=s3_uri,
+                                destination=f"/opt/ml/processing/input/data/{channel}",
+                                s3_input_mode=step_settings.input_data_s3_mode,
+                            )
+                        )
 
             # Construct S3 outputs from container for step
             outputs = None
@@ -475,42 +505,56 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                         )
                     )
 
-            # Convert environment to a dict of strings
-            environment = {
+            step_environment: Dict[str, Union[str, PipelineVariable]] = {
                 key: str(value)
-                if not isinstance(value, ExecutionVariable)
+                if not isinstance(value, PipelineVariable)
                 else value
                 for key, value in environment.items()
             }
 
+            step_environment[ENV_ZENML_SAGEMAKER_RUN_ID] = (
+                ExecutionVariables.PIPELINE_EXECUTION_ARN
+            )
+
+            sagemaker_step: Step
             if use_training_step:
                 # Create Estimator and TrainingStep
-                estimator = sagemaker.estimator.Estimator(
+                estimator = Estimator(
                     keep_alive_period_in_seconds=step_settings.keep_alive_period_in_seconds,
                     output_path=output_path,
-                    environment=environment,
+                    environment=step_environment,
                     container_entry_point=entrypoint,
                     **args_for_step_executor,
                 )
+
                 sagemaker_step = TrainingStep(
                     name=step_name,
-                    depends_on=step.spec.upstream_steps,
-                    inputs=inputs,
+                    depends_on=cast(
+                        Optional[List[Union[str, Step, StepCollection]]],
+                        step.spec.upstream_steps,
+                    ),
+                    inputs=training_inputs,
                     estimator=estimator,
                 )
             else:
                 # Create Processor and ProcessingStep
-                processor = sagemaker.processing.Processor(
-                    entrypoint=entrypoint,
-                    env=environment,
+                processor = Processor(
+                    entrypoint=cast(
+                        Optional[List[Union[str, PipelineVariable]]],
+                        entrypoint,
+                    ),
+                    env=step_environment,
                     **args_for_step_executor,
                 )
 
                 sagemaker_step = ProcessingStep(
                     name=step_name,
                     processor=processor,
-                    depends_on=step.spec.upstream_steps,
-                    inputs=inputs,
+                    depends_on=cast(
+                        Optional[List[Union[str, Step, StepCollection]]],
+                        step.spec.upstream_steps,
+                    ),
+                    inputs=processing_inputs,
                     outputs=outputs,
                 )
 

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py

@@ -13,10 +13,21 @@
 #  permissions and limitations under the License.
 """Implementation of the Sagemaker Step Operator."""
 
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 import boto3
-import sagemaker
+from sagemaker.estimator import Estimator
+from sagemaker.inputs import TrainingInput
+from sagemaker.session import Session
 
 from zenml.client import Client
 from zenml.config.build_configuration import BuildConfiguration
@@ -216,7 +227,7 @@ class SagemakerStepOperator(BaseStepOperator):
         else:
             boto_session = boto3.Session()
 
-        session = sagemaker.Session(
+        session = Session(
             boto_session=boto_session, default_bucket=self.config.bucket
         )
 
@@ -232,9 +243,7 @@ class SagemakerStepOperator(BaseStepOperator):
         estimator_args["sagemaker_session"] = session
 
         # Create Estimator
-        estimator = sagemaker.estimator.Estimator(
-            image_name, self.config.role, **estimator_args
-        )
+        estimator = Estimator(image_name, self.config.role, **estimator_args)
 
         # SageMaker allows 63 characters at maximum for job name - ZenML uses 60 for safety margin.
         step_name = Client().get_run_step(info.step_run_id).name
@@ -248,18 +257,14 @@ class SagemakerStepOperator(BaseStepOperator):
         )
 
         # Construct training input object, if necessary
-        inputs = None
+        inputs: Optional[Union[TrainingInput, Dict[str, TrainingInput]]] = None
 
         if isinstance(settings.input_data_s3_uri, str):
-            inputs = sagemaker.inputs.TrainingInput(
-                s3_data=settings.input_data_s3_uri
-            )
+            inputs = TrainingInput(s3_data=settings.input_data_s3_uri)
         elif isinstance(settings.input_data_s3_uri, dict):
             inputs = {}
             for channel, s3_uri in settings.input_data_s3_uri.items():
-                inputs[channel] = sagemaker.inputs.TrainingInput(
-                    s3_data=s3_uri
-                )
+                inputs[channel] = TrainingInput(s3_data=s3_uri)
 
         experiment_config = {}
         if settings.experiment_name:

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/logging/step_logging.py

@@ -549,8 +549,14 @@ class PipelineLogsStorageContext:
                 except Exception:
                     pass
 
-                if step_context and args[0] != "\n":
-                    message = f"[{step_context.step_name}] " + message
+                if step_context and args[0] not in ["\n", ""]:
+                    # For progress bar updates (with \r), inject the step name after the \r
+                    if "\r" in message:
+                        message = message.replace(
+                            "\r", f"\r[{step_context.step_name}] "
+                        )
+                    else:
+                        message = f"[{step_context.step_name}] {message}"
 
                 output = method(message, *args[1:], **kwargs)
 

{zenml_nightly-0.83.0.dev20250603 → zenml_nightly-0.83.0.dev20250605}/src/zenml/models/v2/core/stack.py

@@ -27,7 +27,7 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 from sqlmodel import and_
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
@@ -94,18 +94,31 @@ class StackRequest(UserScopedRequest):
         "scratch.",
     )
 
-    @property
-    def is_valid(self) -> bool:
-        """Check if the stack is valid.
+    @field_validator("components")
+    def _validate_components(
+        cls, value: Dict[StackComponentType, List[Union[UUID, ComponentInfo]]]
+    ) -> Dict[StackComponentType, List[Union[UUID, ComponentInfo]]]:
+        """Validate the components of the stack.
+
+        Args:
+            value: The components of the stack.
+
+        Raises:
+            ValueError: If the stack does not contain an orchestrator and
+                artifact store.
 
         Returns:
-            True if the stack is valid, False otherwise.
+            The components of the stack.
         """
-        if not self.components:
-            return False
-        return (
-            StackComponentType.ARTIFACT_STORE in self.components
-            and StackComponentType.ORCHESTRATOR in self.components
+        if value:
+            artifact_stores = value.get(StackComponentType.ARTIFACT_STORE, [])
+            orchestrators = value.get(StackComponentType.ORCHESTRATOR, [])
+
+            if orchestrators and artifact_stores:
+                return value
+
+        raise ValueError(
+            "Stack must contain at least an orchestrator and artifact store."
         )
 
     @model_validator(mode="after")
@@ -163,6 +176,39 @@ class StackUpdate(BaseUpdate):
         title="The stack labels.",
     )
 
+    @field_validator("components")
+    def _validate_components(
+        cls,
+        value: Optional[
+            Dict[StackComponentType, List[Union[UUID, ComponentInfo]]]
+        ],
+    ) -> Optional[Dict[StackComponentType, List[Union[UUID, ComponentInfo]]]]:
+        """Validate the components of the stack.
+
+        Args:
+            value: The components of the stack.
+
+        Raises:
+            ValueError: If the stack does not contain an orchestrator and
+                artifact store.
+
+        Returns:
+            The components of the stack.
+        """
+        if value is None:
+            return None
+
+        if value:
+            artifact_stores = value.get(StackComponentType.ARTIFACT_STORE, [])
+            orchestrators = value.get(StackComponentType.ORCHESTRATOR, [])
+
+            if orchestrators and artifact_stores:
+                return value
+
+        raise ValueError(
+            "Stack must contain at least an orchestrator and artifact store."
+        )
+
 
 # ------------------ Response Model ------------------