zenml-nightly 0.73.0.dev20250123__py3-none-any.whl → 0.73.0.dev20250125__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.73.0.dev20250123
+0.73.0.dev20250125

zenml/analytics/context.py

@@ -17,7 +17,6 @@ This module is based on the 'analytics-python' package created by Segment.
 The base functionalities are adapted to work with the ZenML analytics server.
 """
 
-import datetime
 import locale
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
@@ -32,6 +31,7 @@ from zenml.constants import (
 )
 from zenml.environment import Environment, get_environment
 from zenml.logger import get_logger
+from zenml.utils.time_utils import utc_now_tz_aware
 
 if TYPE_CHECKING:
     from zenml.analytics.enums import AnalyticsEvent
@@ -284,11 +284,7 @@ class AnalyticsContext:
 
         try:
             # Timezone as tzdata
-            tz = (
-                datetime.datetime.now(datetime.timezone.utc)
-                .astimezone()
-                .tzname()
-            )
+            tz = utc_now_tz_aware().astimezone().tzname()
             if tz is not None:
                 properties.update({"timezone": tz})
 

zenml/cli/annotator.py

@@ -182,7 +182,7 @@ def register_annotator_subcommands() -> None:
         kwargs_dict = {}
         for arg in kwargs:
             if arg.startswith("--"):
-                key, value = arg.lstrip("--").split("=", 1)
+                key, value = arg.removeprefix("--").split("=", 1)
                 kwargs_dict[key] = value
 
         if annotator.flavor == "prodigy":

zenml/cli/login.py

@@ -249,6 +249,7 @@ def connect_to_pro_server(
     from zenml.login.pro.tenant.models import TenantStatus
 
     pro_api_url = pro_api_url or ZENML_PRO_API_URL
+    pro_api_url = pro_api_url.rstrip("/")
 
     server_id, server_url, server_name = None, None, None
     login = False
@@ -434,6 +435,7 @@ def is_pro_server(
     from zenml.login.credentials_store import get_credentials_store
     from zenml.login.server_info import get_server_info
 
+    url = url.rstrip("/")
     # First, check the credentials store
     credentials_store = get_credentials_store()
     credentials = credentials_store.get_credentials(url)
@@ -790,15 +792,16 @@ def login(
     )
 
     if server is not None:
-        if not re.match(r"^(https?|mysql)://", server):
-            # The server argument is a ZenML Pro server name or UUID
-            connect_to_pro_server(
-                pro_server=server,
+        if re.match(r"^mysql://", server):
+            # The server argument is a MySQL URL, we can directly connect to it
+
+            connect_to_server(
+                url=server,
                 api_key=api_key_value,
+                verify_ssl=verify_ssl,
                 refresh=refresh,
-                pro_api_url=pro_api_url,
             )
-        else:
+        elif re.match(r"^https?://", server):
             # The server argument is a server URL
 
             # First, try to discover if the server is a ZenML Pro server or not
@@ -819,6 +822,14 @@ def login(
                     verify_ssl=verify_ssl,
                     refresh=refresh,
                 )
+        else:
+            # The server argument is a ZenML Pro server name or UUID
+            connect_to_pro_server(
+                pro_server=server,
+                api_key=api_key_value,
+                refresh=refresh,
+                pro_api_url=pro_api_url,
+            )
 
     elif current_non_local_server:
         # The server argument is not provided, so we default to

zenml/cli/server.py

@@ -578,6 +578,7 @@ def server_list(
     from zenml.login.pro.tenant.models import TenantRead, TenantStatus
 
     pro_api_url = pro_api_url or ZENML_PRO_API_URL
+    pro_api_url = pro_api_url.rstrip("/")
 
     credentials_store = get_credentials_store()
     pro_token = credentials_store.get_pro_token(

zenml/cli/service_connectors.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Service connector CLI commands."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Any, Dict, List, Optional, Union, cast
 from uuid import UUID
 
@@ -25,7 +25,6 @@ from zenml.cli.utils import (
     is_sorted_or_filtered,
     list_options,
     print_page_info,
-    seconds_to_human_readable,
 )
 from zenml.client import Client
 from zenml.console import console
@@ -37,6 +36,7 @@ from zenml.models import (
     ServiceConnectorResourcesModel,
     ServiceConnectorResponse,
 )
+from zenml.utils.time_utils import seconds_to_human_readable, utc_now
 
 
 # Service connectors
@@ -292,7 +292,7 @@ def prompt_expires_at(
         default_str = ""
         if default is not None:
             seconds = int(
-                (default - datetime.now(timezone.utc)).total_seconds()
+                (default - utc_now(tz_aware=default)).total_seconds()
             )
             default_str = (
                 f" [{str(default)} i.e. in "
@@ -309,7 +309,7 @@ def prompt_expires_at(
 
         assert expires_at is not None
         assert isinstance(expires_at, datetime)
-        if expires_at < datetime.now(timezone.utc):
+        if expires_at < utc_now(tz_aware=expires_at):
             cli_utils.warning(
                 "The expiration time must be in the future. Please enter a "
                 "later date and time."
@@ -317,7 +317,7 @@ def prompt_expires_at(
             continue
 
         seconds = int(
-            (expires_at - datetime.now(timezone.utc)).total_seconds()
+            (expires_at - utc_now(tz_aware=expires_at)).total_seconds()
         )
 
         confirm = click.confirm(

zenml/cli/stack.py

@@ -17,7 +17,6 @@ import getpass
 import re
 import time
 import webbrowser
-from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -77,6 +76,7 @@ from zenml.service_connectors.service_connector_utils import (
 )
 from zenml.utils import requirements_utils
 from zenml.utils.dashboard_utils import get_component_url, get_stack_url
+from zenml.utils.time_utils import utc_now_tz_aware
 from zenml.utils.yaml_utils import read_yaml, write_yaml
 
 if TYPE_CHECKING:
@@ -1575,7 +1575,7 @@ def deploy(
         ):
             raise click.Abort()
 
-        date_start = datetime.now(timezone.utc)
+        date_start = utc_now_tz_aware()
 
         webbrowser.open(deployment_config.deployment_url)
         console.print(

zenml/cli/utils.py

@@ -14,7 +14,6 @@
 """Utility functions for the CLI."""
 
 import contextlib
-import datetime
 import json
 import os
 import platform
@@ -79,6 +78,7 @@ from zenml.services import BaseService, ServiceState
 from zenml.stack import StackComponent
 from zenml.stack.stack_component import StackComponentConfig
 from zenml.utils import secret_utils
+from zenml.utils.time_utils import expires_in
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -1581,58 +1581,6 @@ def print_components_table(
     print_table(configurations)
 
 
-def seconds_to_human_readable(time_seconds: int) -> str:
-    """Converts seconds to human-readable format.
-
-    Args:
-        time_seconds: Seconds to convert.
-
-    Returns:
-        Human readable string.
-    """
-    seconds = time_seconds % 60
-    minutes = (time_seconds // 60) % 60
-    hours = (time_seconds // 3600) % 24
-    days = time_seconds // 86400
-    tokens = []
-    if days:
-        tokens.append(f"{days}d")
-    if hours:
-        tokens.append(f"{hours}h")
-    if minutes:
-        tokens.append(f"{minutes}m")
-    if seconds:
-        tokens.append(f"{seconds}s")
-
-    return "".join(tokens)
-
-
-def expires_in(
-    expires_at: datetime.datetime,
-    expired_str: str,
-    skew_tolerance: Optional[int] = None,
-) -> str:
-    """Returns a human-readable string of the time until the token expires.
-
-    Args:
-        expires_at: Datetime object of the token expiration.
-        expired_str: String to return if the token is expired.
-        skew_tolerance: Seconds of skew tolerance to subtract from the
-            expiration time. If the token expires within this time, it will be
-            considered expired.
-
-    Returns:
-        Human readable string.
-    """
-    now = datetime.datetime.now(datetime.timezone.utc)
-    expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
-    if skew_tolerance:
-        expires_at -= datetime.timedelta(seconds=skew_tolerance)
-    if expires_at < now:
-        return expired_str
-    return seconds_to_human_readable(int((expires_at - now).total_seconds()))
-
-
 def print_service_connectors_table(
     client: "Client",
     connectors: Sequence["ServiceConnectorResponse"],
@@ -2660,7 +2608,7 @@ def print_model_url(url: Optional[str]) -> None:
         warning(
             "You can display various ZenML entities including pipelines, "
             "runs, stacks and much more on the ZenML Dashboard. "
-            "You can try it locally, by running `zenml local --local`, or "
+            "You can try it locally, by running `zenml login --local`, or "
             "remotely, by deploying ZenML on the infrastructure of your choice."
         )
 

zenml/config/pipeline_configurations.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from pydantic import SerializeAsAny, field_validator
@@ -23,6 +23,7 @@ from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import SourceWithValidator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings
@@ -61,7 +62,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
             The full substitutions dict including date and time.
         """
         if start_time is None:
-            start_time = datetime.now(timezone.utc)
+            start_time = utc_now()
         ret = self.substitutions.copy()
         ret.setdefault("date", start_time.strftime("%Y_%m_%d"))
         ret.setdefault("time", start_time.strftime("%H_%M_%S_%f"))

zenml/config/schedule.py

@@ -98,27 +98,3 @@ class Schedule(BaseModel):
                 "or a run once start time "
                 "need to be set for a valid schedule."
             )
-
-    @property
-    def utc_start_time(self) -> Optional[str]:
-        """Optional ISO-formatted string of the UTC start time.
-
-        Returns:
-            Optional ISO-formatted string of the UTC start time.
-        """
-        if not self.start_time:
-            return None
-
-        return self.start_time.astimezone(datetime.timezone.utc).isoformat()
-
-    @property
-    def utc_end_time(self) -> Optional[str]:
-        """Optional ISO-formatted string of the UTC end time.
-
-        Returns:
-            Optional ISO-formatted string of the UTC end time.
-        """
-        if not self.end_time:
-            return None
-
-        return self.end_time.astimezone(datetime.timezone.utc).isoformat()

zenml/enums.py

@@ -376,6 +376,7 @@ class MetadataResourceTypes(StrEnum):
     STEP_RUN = "step_run"
     ARTIFACT_VERSION = "artifact_version"
     MODEL_VERSION = "model_version"
+    SCHEDULE = "schedule"
 
 
 class DatabaseBackupStrategy(StrEnum):

zenml/event_hub/base_event_hub.py

@@ -14,7 +14,7 @@
 """Base class for event hub implementations."""
 
 from abc import ABC, abstractmethod
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 from zenml import EventSourceResponse
@@ -28,6 +28,7 @@ from zenml.models import (
     TriggerExecutionResponse,
     TriggerResponse,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.jwt import JWTToken
 
@@ -134,9 +135,7 @@ class BaseEventHub(ABC):
         )
         expires: Optional[datetime] = None
         if trigger.action.auth_window:
-            expires = datetime.now(timezone.utc) + timedelta(
-                minutes=trigger.action.auth_window
-            )
+            expires = utc_now() + timedelta(minutes=trigger.action.auth_window)
         encoded_token = token.encode(expires=expires)
         auth_context = AuthContext(
             user=trigger.action.service_account,

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -42,6 +42,7 @@ from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils import io_utils
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.config import ResourceSettings
@@ -408,8 +409,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         if schedule:
             if schedule.cron_expression:
                 start_time = schedule.start_time or (
-                    datetime.datetime.now(datetime.timezone.utc)
-                    - datetime.timedelta(7)
+                    utc_now() - datetime.timedelta(7)
                 )
                 return {
                     "schedule": schedule.cron_expression,
@@ -429,7 +429,6 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
             "schedule": "@once",
             # set a start time in the past and disable catchup so airflow
             # runs the dag immediately
-            "start_date": datetime.datetime.now(datetime.timezone.utc)
-            - datetime.timedelta(7),
+            "start_date": utc_now() - datetime.timedelta(7),
             "catchup": False,
         }

zenml/integrations/aws/__init__.py

@@ -35,12 +35,13 @@ AWS_RESOURCE_TYPE = "aws-generic"
 S3_RESOURCE_TYPE = "s3-bucket"
 AWS_IMAGE_BUILDER_FLAVOR = "aws"
 
+
 class AWSIntegration(Integration):
     """Definition of AWS integration for ZenML."""
 
     NAME = AWS
     REQUIREMENTS = [
-        "sagemaker>=2.117.0",
+        "sagemaker>=2.199.0",
         "kubernetes",
         "aws-profile-manager",
     ]

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -85,6 +85,7 @@ class SagemakerOrchestratorSettings(BaseSettings):
             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.
@@ -184,6 +185,10 @@ class SagemakerOrchestratorConfig(
 
     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/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
@@ -203,6 +208,7 @@ class SagemakerOrchestratorConfig(
     """
 
     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
@@ -232,6 +238,15 @@ class SagemakerOrchestratorConfig(
         """
         return self.synchronous
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return True
+
 
 class SagemakerOrchestratorFlavor(BaseOrchestratorFlavor):
     """Flavor for the Sagemaker orchestrator."""