zenml-nightly 0.83.1.dev20250624__py3-none-any.whl → 0.83.1.dev20250625__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.83.1.dev20250624
+0.83.1.dev20250625

zenml/cli/login.py

@@ -229,6 +229,7 @@ def connect_to_pro_server(
     api_key: Optional[str] = None,
     refresh: bool = False,
     pro_api_url: Optional[str] = None,
+    verify_ssl: Union[str, bool] = True,
 ) -> None:
     """Connect the client to a ZenML Pro server.
 
@@ -238,6 +239,8 @@ def connect_to_pro_server(
         api_key: The API key to use to authenticate with the ZenML Pro server.
         refresh: Whether to force a new login flow with the ZenML Pro server.
         pro_api_url: The URL for the ZenML Pro API.
+        verify_ssl: Whether to verify the server's TLS certificate. If a string
+            is passed, it is interpreted as the path to a CA bundle file.
 
     Raises:
         ValueError: If incorrect parameters are provided.
@@ -280,7 +283,12 @@ def connect_to_pro_server(
         # server to connect to.
         if api_key:
             if server_url:
-                connect_to_server(server_url, api_key=api_key, pro_server=True)
+                connect_to_server(
+                    server_url,
+                    api_key=api_key,
+                    pro_server=True,
+                    verify_ssl=verify_ssl,
+                )
                 return
             else:
                 raise ValueError(
@@ -299,6 +307,7 @@ def connect_to_pro_server(
         try:
             token = web_login(
                 pro_api_url=pro_api_url,
+                verify_ssl=verify_ssl,
             )
         except AuthorizationException as e:
             cli_utils.error(f"Authorization error: {e}")
@@ -418,7 +427,9 @@ def connect_to_pro_server(
         f"Connecting to ZenML Pro server: {server.name} [{str(server.id)}] "
     )
 
-    connect_to_server(server.url, api_key=api_key, pro_server=True)
+    connect_to_server(
+        server.url, api_key=api_key, pro_server=True, verify_ssl=verify_ssl
+    )
 
     # Update the stored server info with more accurate data taken from the
     # ZenML Pro workspace object.
@@ -555,7 +566,7 @@ def _fail_if_authentication_environment_variables_set() -> None:
         not return until the server exits or is stopped with CTRL+C
 
       * `--docker`: start the local ZenML server as a Docker container instead
-        of a local process
+        of a local background process.
 
       * `--port`: use a custom TCP port value for the local ZenML server
 
@@ -775,6 +786,9 @@ def login(
             pro_server=server,
             refresh=True,
             pro_api_url=pro_api_url,
+            verify_ssl=ssl_ca_cert
+            if ssl_ca_cert is not None
+            else not no_verify_ssl,
         )
         return
 
@@ -822,6 +836,7 @@ def login(
                     # Prefer the pro API URL extracted from the server info if
                     # available
                     pro_api_url=server_pro_api_url or pro_api_url,
+                    verify_ssl=verify_ssl,
                 )
             else:
                 connect_to_server(
@@ -837,6 +852,7 @@ def login(
                 api_key=api_key_value,
                 refresh=refresh,
                 pro_api_url=pro_api_url,
+                verify_ssl=verify_ssl,
             )
 
     elif current_non_local_server and not refresh:
@@ -861,6 +877,7 @@ def login(
                 # Prefer the pro API URL extracted from the server info if
                 # available
                 pro_api_url=server_pro_api_url or pro_api_url,
+                verify_ssl=verify_ssl,
             )
         else:
             cli_utils.declare(
@@ -890,6 +907,7 @@ def login(
         connect_to_pro_server(
             api_key=api_key_value,
             pro_api_url=pro_api_url,
+            verify_ssl=verify_ssl,
         )
 
 

zenml/config/__init__.py

@@ -23,12 +23,23 @@ configuration. This ``GlobalConfiguration`` object handles the serialization and
 deserialization of the configuration options that are stored in the file in
 order to persist the configuration across sessions.
 """
-from zenml.config.docker_settings import DockerSettings
-from zenml.config.resource_settings import ResourceSettings
+from zenml.config.docker_settings import (
+    DockerSettings,
+    PythonPackageInstaller,
+    PythonEnvironmentExportMethod,
+)
+from zenml.config.resource_settings import ResourceSettings, ByteUnit
 from zenml.config.retry_config import StepRetryConfig
+from zenml.config.schedule import Schedule
+from zenml.config.store_config import StoreConfiguration
 
 __all__ = [
     "DockerSettings",
+    "PythonPackageInstaller",
+    "PythonEnvironmentExportMethod",
     "ResourceSettings",
+    "ByteUnit",
     "StepRetryConfig",
+    "Schedule",
+    "StoreConfiguration",
 ]

zenml/constants.py

@@ -173,7 +173,6 @@ 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_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
 )

zenml/exceptions.py

@@ -220,3 +220,19 @@ class CustomFlavorImportError(ImportError):
 
 class MaxConcurrentTasksError(ZenMLBaseException):
     """Raised when the maximum number of concurrent tasks is reached."""
+
+
+class RunMonitoringError(ZenMLBaseException):
+    """Raised when an error occurs while monitoring a pipeline run."""
+
+    def __init__(
+        self,
+        original_exception: BaseException,
+    ) -> None:
+        """Initializes the error.
+
+        Args:
+            original_exception: The original exception that occurred while
+                monitoring the pipeline run.
+        """
+        self.original_exception = original_exception

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -38,7 +38,7 @@ from zenml.integrations.airflow.flavors.airflow_orchestrator_flavor import (
 )
 from zenml.io import fileio
 from zenml.logger import get_logger
-from zenml.orchestrators import ContainerizedOrchestrator
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils import io_utils
@@ -191,21 +191,29 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         if self.config.local:
             stack.check_local_paths()
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Creates and writes an Airflow DAG zip file.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
+
+        Returns:
+            Optional submission result.
         """
         pipeline_settings = cast(
             AirflowOrchestratorSettings, self.get_settings(deployment)
@@ -277,6 +285,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
             dag_generator_values=dag_generator_values,
             output_dir=pipeline_settings.dag_output_dir or self.dags_directory,
         )
+        return None
 
     def _apply_resource_settings(
         self,

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -19,7 +19,6 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
-    Iterator,
     List,
     Optional,
     Tuple,
@@ -60,7 +59,6 @@ from zenml.constants import (
 )
 from zenml.enums import (
     ExecutionStatus,
-    MetadataResourceTypes,
     StackComponentType,
 )
 from zenml.integrations.aws.flavors.sagemaker_orchestrator_flavor import (
@@ -73,7 +71,7 @@ from zenml.integrations.aws.orchestrators.sagemaker_orchestrator_entrypoint_conf
 )
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType, Uri
-from zenml.orchestrators import ContainerizedOrchestrator
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils.env_utils import split_environment_variables
@@ -273,20 +271,25 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             boto_session=boto_session, default_bucket=self.config.bucket
         )
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Iterator[Dict[str, MetadataType]]:
-        """Prepares or runs a pipeline on Sagemaker.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The deployment to prepare or run.
-            stack: The stack to run on.
+            deployment: The pipeline deployment to submit.
+            stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
@@ -296,8 +299,8 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 AWS SageMaker NetworkConfig class.
             ValueError: If the schedule is not valid.
 
-        Yields:
-            A dictionary of metadata related to the pipeline run.
+        Returns:
+            Optional submission result.
         """
         # sagemaker requires pipelineName to use alphanum and hyphens only
         unsanitized_orchestrator_run_name = get_orchestrator_run_name(
@@ -705,26 +708,14 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             )
             logger.info(f"The schedule ARN is: {triggers[0]}")
 
+            schedule_metadata = {}
             try:
-                from zenml.models import RunMetadataResource
-
                 schedule_metadata = self.generate_schedule_metadata(
                     schedule_arn=triggers[0]
                 )
-
-                Client().create_run_metadata(
-                    metadata=schedule_metadata,  # type: ignore[arg-type]
-                    resources=[
-                        RunMetadataResource(
-                            id=deployment.schedule.id,
-                            type=MetadataResourceTypes.SCHEDULE,
-                        )
-                    ],
-                )
             except Exception as e:
                 logger.debug(
-                    "There was an error attaching metadata to the "
-                    f"schedule: {e}"
+                    "There was an error generating schedule metadata: %s", e
                 )
 
             logger.info(
@@ -749,6 +740,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             logger.info(
                 f"`aws scheduler delete-schedule --name {schedule_name}`"
             )
+            return SubmissionResult(metadata=schedule_metadata)
         else:
             # Execute the pipeline immediately if no schedule is specified
             execution = pipeline.start()
@@ -757,33 +749,40 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 "when using the Sagemaker Orchestrator."
             )
 
-            # Yield metadata based on the generated execution object
-            yield from self.compute_metadata(
+            run_metadata = self.compute_metadata(
                 execution_arn=execution.arn, settings=settings
             )
 
-            # mainly for testing purposes, we wait for the pipeline to finish
+            _wait_for_completion = None
             if settings.synchronous:
-                logger.info(
-                    "Executing synchronously. Waiting for pipeline to "
-                    "finish... \n"
-                    "At this point you can `Ctrl-C` out without cancelling the "
-                    "execution."
-                )
-                try:
-                    execution.wait(
-                        delay=POLLING_DELAY, max_attempts=MAX_POLLING_ATTEMPTS
-                    )
-                    logger.info("Pipeline completed successfully.")
-                except WaiterError:
-                    raise RuntimeError(
-                        "Timed out while waiting for pipeline execution to "
-                        "finish. For long-running pipelines we recommend "
-                        "configuring your orchestrator for asynchronous "
-                        "execution. The following command does this for you: \n"
-                        f"`zenml orchestrator update {self.name} "
-                        f"--synchronous=False`"
+
+                def _wait_for_completion() -> None:
+                    logger.info(
+                        "Executing synchronously. Waiting for pipeline to "
+                        "finish... \n"
+                        "At this point you can `Ctrl-C` out without cancelling the "
+                        "execution."
                     )
+                    try:
+                        execution.wait(
+                            delay=POLLING_DELAY,
+                            max_attempts=MAX_POLLING_ATTEMPTS,
+                        )
+                        logger.info("Pipeline completed successfully.")
+                    except WaiterError:
+                        raise RuntimeError(
+                            "Timed out while waiting for pipeline execution to "
+                            "finish. For long-running pipelines we recommend "
+                            "configuring your orchestrator for asynchronous "
+                            "execution. The following command does this for you: \n"
+                            f"`zenml orchestrator update {self.name} "
+                            f"--synchronous=False`"
+                        )
+
+            return SubmissionResult(
+                wait_for_completion=_wait_for_completion,
+                metadata=run_metadata,
+            )
 
     def get_pipeline_run_metadata(
         self, run_id: UUID
@@ -798,20 +797,15 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         """
         execution_arn = os.environ[ENV_ZENML_SAGEMAKER_RUN_ID]
 
-        run_metadata: Dict[str, "MetadataType"] = {}
-
         settings = cast(
             SagemakerOrchestratorSettings,
             self.get_settings(Client().get_pipeline_run(run_id)),
         )
 
-        for metadata in self.compute_metadata(
+        return self.compute_metadata(
             execution_arn=execution_arn,
             settings=settings,
-        ):
-            run_metadata.update(metadata)
-
-        return run_metadata
+        )
 
     def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
         """Refreshes the status of a specific pipeline run.
@@ -873,14 +867,14 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         self,
         execution_arn: str,
         settings: SagemakerOrchestratorSettings,
-    ) -> Iterator[Dict[str, MetadataType]]:
+    ) -> Dict[str, MetadataType]:
         """Generate run metadata based on the generated Sagemaker Execution.
 
         Args:
             execution_arn: The ARN of the pipeline execution.
             settings: The Sagemaker orchestrator settings.
 
-        Yields:
+        Returns:
             A dictionary of metadata related to the pipeline run.
         """
         # Orchestrator Run ID
@@ -901,7 +895,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         ):
             metadata[METADATA_ORCHESTRATOR_LOGS_URL] = Uri(logs_url)
 
-        yield metadata
+        return metadata
 
     def _compute_orchestrator_url(
         self,
@@ -979,7 +973,9 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             return None
 
     @staticmethod
-    def generate_schedule_metadata(schedule_arn: str) -> Dict[str, str]:
+    def generate_schedule_metadata(
+        schedule_arn: str,
+    ) -> Dict[str, MetadataType]:
         """Attaches metadata to the ZenML Schedules.
 
         Args:

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -19,7 +19,6 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
-    Iterator,
     List,
     Optional,
     Tuple,
@@ -63,7 +62,7 @@ from zenml.integrations.azure.orchestrators.azureml_orchestrator_entrypoint_conf
 )
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType, Uri
-from zenml.orchestrators import ContainerizedOrchestrator
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils.string_utils import b64_encode
@@ -198,27 +197,32 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             command=" ".join(command + arguments),
         )
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Iterator[Dict[str, MetadataType]]:
-        """Prepares or runs a pipeline on AzureML.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The deployment to prepare or run.
-            stack: The stack to run on.
+            deployment: The pipeline deployment to submit.
+            stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             RuntimeError: If the creation of the schedule fails.
 
-        Yields:
-            A dictionary of metadata related to the pipeline run.
+        Returns:
+            Optional submission result.
         """
         # Authentication
         if connector := self.get_connector():
@@ -384,14 +388,11 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
                     "Failed to create schedule for the pipeline "
                     f"'{run_name}': {str(e)}"
                 )
-
+            return None
         else:
             job = ml_client.jobs.create_or_update(pipeline_job)
             logger.info(f"Pipeline {run_name} has been started.")
 
-            # Yield metadata based on the generated job object
-            yield from self.compute_metadata(job)
-
             assert job.services is not None
             assert job.name is not None
 
@@ -401,9 +402,17 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
                 f"{job.services['Studio'].endpoint}"
             )
 
+            _wait_for_completion = None
             if settings.synchronous:
-                logger.info("Waiting for pipeline to finish...")
-                ml_client.jobs.stream(job.name)
+
+                def _wait_for_completion() -> None:
+                    logger.info("Waiting for pipeline to finish...")
+                    ml_client.jobs.stream(job.name)
+
+            return SubmissionResult(
+                metadata=self.compute_metadata(job),
+                wait_for_completion=_wait_for_completion,
+            )
 
     def get_pipeline_run_metadata(
         self, run_id: UUID
@@ -518,13 +527,13 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         else:
             raise ValueError("Unknown status for the pipeline job.")
 
-    def compute_metadata(self, job: Any) -> Iterator[Dict[str, MetadataType]]:
+    def compute_metadata(self, job: Any) -> Dict[str, MetadataType]:
         """Generate run metadata based on the generated AzureML PipelineJob.
 
         Args:
             job: The corresponding PipelineJob object.
 
-        Yields:
+        Returns:
             A dictionary of metadata related to the pipeline run.
         """
         # Metadata
@@ -538,7 +547,7 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         if orchestrator_url := self._compute_orchestrator_url(job):
             metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_url)
 
-        yield metadata
+        return metadata
 
     @staticmethod
     def _compute_orchestrator_url(job: Any) -> Optional[str]: