zenml-nightly 0.61.0.dev20240711__py3-none-any.whl → 0.61.0.dev20240713__py3-none-any.whl

zenml/logging/step_logging.py

@@ -240,8 +240,6 @@ class StepLogsStorage:
 
         # Immutable filesystems state
         self.last_merge_time = time.time()
-        self.log_files_not_merged: List[str] = []
-        self.next_merged_file_name: str = self._get_timestamped_filename()
 
     @property
     def artifact_store(self) -> "BaseArtifactStore":
@@ -279,13 +277,16 @@ class StepLogsStorage:
             or time.time() - self.last_save_time >= self.time_interval
         )
 
-    def _get_timestamped_filename(self) -> str:
+    def _get_timestamped_filename(self, suffix: str = "") -> str:
         """Returns a timestamped filename.
 
+        Args:
+            suffix: optional suffix for the file name
+
         Returns:
             The timestamped filename.
         """
-        return f"{time.time()}{LOGS_EXTENSION}"
+        return f"{time.time()}{suffix}{LOGS_EXTENSION}"
 
     def save_to_file(self, force: bool = False) -> None:
         """Method to save the buffer to the given URI.
@@ -302,12 +303,7 @@ class StepLogsStorage:
             try:
                 if self.buffer:
                     if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-                        if not self.log_files_not_merged:
-                            self.next_merged_file_name = (
-                                self._get_timestamped_filename()
-                            )
                         _logs_uri = self._get_timestamped_filename()
-                        self.log_files_not_merged.append(_logs_uri)
                         with self.artifact_store.open(
                             os.path.join(
                                 self.logs_uri,
@@ -346,42 +342,40 @@ class StepLogsStorage:
             and time.time() - self.last_merge_time > self.merge_files_interval
         ):
             try:
-                self.merge_log_files(
-                    self.next_merged_file_name, self.log_files_not_merged
-                )
+                self.merge_log_files()
             except (OSError, IOError) as e:
                 logger.error(f"Error while trying to roll up logs: {e}")
-            else:
-                self.log_files_not_merged = []
             finally:
                 self.last_merge_time = time.time()
 
-    def merge_log_files(
-        self,
-        file_name: Optional[str] = None,
-        files: Optional[List[str]] = None,
-    ) -> None:
+    def merge_log_files(self, merge_all_files: bool = False) -> None:
         """Merges all log files into one in the given URI.
 
         Called on the logging context exit.
 
         Args:
-            file_name: The name of the merged log file.
-            files: The list of log files to merge.
+            merge_all_files: whether to merge all files or only raw files
         """
         if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-            files_ = files or self.artifact_store.listdir(self.logs_uri)
-            file_name_ = file_name or self._get_timestamped_filename()
+            merged_file_suffix = "_merged"
+            files_ = self.artifact_store.listdir(self.logs_uri)
+            if not merge_all_files:
+                # already merged files will not be merged again
+                files_ = [f for f in files_ if merged_file_suffix not in f]
+            file_name_ = self._get_timestamped_filename(
+                suffix=merged_file_suffix
+            )
             if len(files_) > 1:
                 files_.sort()
                 logger.debug("Log files count: %s", len(files_))
 
-                try:
-                    # dump all logs to a local file first
-                    with self.artifact_store.open(
-                        os.path.join(self.logs_uri, file_name_), "w"
-                    ) as merged_file:
-                        for file in files_:
+                missing_files = set()
+                # dump all logs to a local file first
+                with self.artifact_store.open(
+                    os.path.join(self.logs_uri, file_name_), "w"
+                ) as merged_file:
+                    for file in files_:
+                        try:
                             merged_file.write(
                                 str(
                                     _load_file_from_artifact_store(
@@ -391,11 +385,12 @@ class StepLogsStorage:
                                     )
                                 )
                             )
-                except Exception as e:
-                    logger.warning(f"Failed to merge log files. {e}")
-                else:
-                    # clean up left over files
-                    for file in files_:
+                        except DoesNotExistException:
+                            missing_files.add(file)
+
+                # clean up left over files
+                for file in files_:
+                    if file not in missing_files:
                         self.artifact_store.remove(
                             os.path.join(self.logs_uri, str(file))
                         )
@@ -452,7 +447,6 @@ class StepLogsStorageContext:
         Restores the `write` method of both stderr and stdout.
         """
         self.storage.save_to_file(force=True)
-        self.storage.merge_log_files()
 
         setattr(sys.stdout, "write", self.stdout_write)
         setattr(sys.stdout, "flush", self.stdout_flush)
@@ -462,6 +456,11 @@ class StepLogsStorageContext:
 
         redirected.set(False)
 
+        try:
+            self.storage.merge_log_files(merge_all_files=True)
+        except (OSError, IOError) as e:
+            logger.warning(f"Step logs roll-up failed: {e}")
+
     def _wrap_write(self, method: Callable[..., Any]) -> Callable[..., Any]:
         """Wrapper function that utilizes the storage object to store logs.
 

zenml/models/__init__.py

@@ -389,6 +389,7 @@ from zenml.models.v2.core.server_settings import (
 )
 from zenml.models.v2.misc.stack_deployment import (
     DeployedStack,
+    StackDeploymentConfig,
     StackDeploymentInfo,
 )
 
@@ -726,6 +727,7 @@ __all__ = [
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",
+    "StackDeploymentConfig",
     "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",
     "OAuthDeviceAuthorizationResponse",

zenml/models/v2/core/server_settings.py

@@ -15,8 +15,6 @@
 
 from datetime import datetime
 from typing import (
-    Any,
-    Dict,
     Optional,
 )
 from uuid import UUID
@@ -57,10 +55,6 @@ class ServerSettingsUpdate(BaseZenModel):
         default=None,
         title="Whether to display notifications about ZenML updates in the dashboard.",
     )
-    onboarding_state: Optional[Dict[str, Any]] = Field(
-        default=None,
-        title="The server's onboarding state.",
-    )
 
 
 # ------------------ Response Model ------------------
@@ -96,11 +90,6 @@ class ServerSettingsResponseBody(BaseResponseBody):
 class ServerSettingsResponseMetadata(BaseResponseMetadata):
     """Response metadata for server settings."""
 
-    onboarding_state: Dict[str, Any] = Field(
-        default={},
-        title="The server's onboarding state.",
-    )
-
 
 class ServerSettingsResponseResources(BaseResponseResources):
     """Response resources for server settings."""
@@ -199,15 +188,6 @@ class ServerSettingsResponse(
         """
         return self.get_body().updated
 
-    @property
-    def onboarding_state(self) -> Dict[str, Any]:
-        """The `onboarding_state` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().onboarding_state
-
 
 # ------------------ Filter Model ------------------
 

zenml/models/v2/misc/stack_deployment.py

@@ -40,6 +40,11 @@ class StackDeploymentInfo(BaseModel):
         title="The instructions for post-deployment.",
         description="The instructions for post-deployment.",
     )
+    integrations: List[str] = Field(
+        title="ZenML integrations required for the stack.",
+        description="The list of ZenML integrations that need to be installed "
+        "for the stack to be usable.",
+    )
     permissions: Dict[str, List[str]] = Field(
         title="The permissions granted to ZenML to access the cloud resources.",
         description="The permissions granted to ZenML to access the cloud "
@@ -52,6 +57,21 @@ class StackDeploymentInfo(BaseModel):
     )
 
 
+class StackDeploymentConfig(BaseModel):
+    """Configuration about a stack deployment."""
+
+    deployment_url: str = Field(
+        title="The cloud provider console URL where the stack will be deployed.",
+    )
+    deployment_url_text: str = Field(
+        title="A textual description for the cloud provider console URL.",
+    )
+    configuration: Optional[str] = Field(
+        title="Configuration for the stack deployment that the user must "
+        "manually configure into the cloud provider console.",
+    )
+
+
 class DeployedStack(BaseModel):
     """Information about a deployed stack."""
 

zenml/orchestrators/step_launcher.py

@@ -291,6 +291,7 @@ class StepLauncher:
                                 logger.error(
                                     f"Failed to run step `{self._step_name}` after {max_retries} retries. Exiting."
                                 )
+                                logger.exception(e)
                                 publish_utils.publish_failed_step_run(
                                     step_run_response.id
                                 )

zenml/stack_deployments/aws_stack_deployment.py

@@ -11,16 +11,12 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""Functionality to deploy a ZenML stack to a cloud provider."""
+"""Functionality to deploy a ZenML stack to AWS."""
 
-import datetime
-from typing import ClassVar, Dict, List, Optional, Tuple
+from typing import ClassVar, Dict, List
 
-from zenml.client import Client
-from zenml.enums import StackComponentType, StackDeploymentProvider
-from zenml.models import (
-    DeployedStack,
-)
+from zenml.enums import StackDeploymentProvider
+from zenml.models import StackDeploymentConfig
 from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
 from zenml.utils.string_utils import random_str
 
@@ -31,6 +27,7 @@ class AWSZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
     """AWS ZenML Cloud Stack Deployment."""
 
     provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AWS
+    deployment: ClassVar[str] = AWS_DEPLOYMENT_TYPE
 
     @classmethod
     def description(cls) -> str:
@@ -65,9 +62,6 @@ to log into your AWS account and create a CloudFormation ZenML stack. The stack
 parameters will be pre-filled with the necessary information to connect ZenML to
 your AWS account, so you should only need to review and confirm the stack.
 
-After the CloudFormation stack is deployed, you can return to the CLI to view
-details about the associated ZenML stack automatically registered with ZenML.
-
 **NOTE**: The CloudFormation stack will create the following new resources in
 your AWS account. Please ensure you have the necessary permissions and are aware
 of any potential costs:
@@ -84,6 +78,18 @@ The CloudFormation stack will automatically create an AWS secret key and
 will share it with ZenML to give it permission to access the resources created
 by the stack. You can revoke these permissions at any time by deleting the
 CloudFormation stack.
+
+**Estimated costs**
+
+A small training job would cost around: $0.60
+
+These are rough estimates and actual costs may vary based on your usage and specific AWS pricing. 
+Some services may be eligible for the AWS Free Tier. Use [the AWS Pricing Calculator](https://calculator.aws)
+for a detailed estimate based on your usage.
+
+💡 **After the CloudFormation stack is deployed, you can return to the CLI to
+view details about the associated ZenML stack automatically registered with
+ZenML.**
 """
 
     @classmethod
@@ -103,6 +109,19 @@ to clean up the resources created by the stack by using the AWS CloudFormation
 console.
 """
 
+    @classmethod
+    def integrations(cls) -> List[str]:
+        """Return the ZenML integrations required for the stack.
+
+        Returns:
+            The list of ZenML integrations that need to be installed for the
+            stack to be usable.
+        """
+        return [
+            "aws",
+            "s3",
+        ]
+
     @classmethod
     def permissions(cls) -> Dict[str, List[str]]:
         """Return the permissions granted to ZenML to access the cloud resources.
@@ -187,32 +206,34 @@ console.
             "South America (São Paulo)": "sa-east-1",
         }
 
-    def deploy_url(
+    def get_deployment_config(
         self,
-        zenml_server_url: str,
-        zenml_server_api_token: str,
-    ) -> Tuple[str, str]:
-        """Return the URL to deploy the ZenML stack to the specified cloud provider.
-
-        The URL should point to a cloud provider console where the user can
-        deploy the ZenML stack and should include as many pre-filled parameters
-        as possible.
-
-        Args:
-            zenml_server_url: The URL of the ZenML server.
-            zenml_server_api_token: The API token to authenticate with the ZenML
-                server.
+    ) -> StackDeploymentConfig:
+        """Return the configuration to deploy the ZenML stack to the specified cloud provider.
+
+        The configuration should include:
+
+        * a cloud provider console URL where the user will be redirected to
+        deploy the ZenML stack. The URL should include as many pre-filled
+        URL query parameters as possible.
+        * a textual description of the URL
+        * some deployment providers may require additional configuration
+        parameters to be passed to the cloud provider in addition to the
+        deployment URL query parameters. Where that is the case, this method
+        should also return a string that the user can copy and paste into the
+        cloud provider console to deploy the ZenML stack (e.g. a set of
+        environment variables, or YAML configuration snippet etc.).
 
         Returns:
-            The URL to deploy the ZenML stack to the specified cloud provider
-            and a text description of the URL.
+            The configuration to deploy the ZenML stack to the specified cloud
+            provider.
         """
         params = dict(
             stackName=self.stack_name,
             templateURL="https://zenml-cf-templates.s3.eu-central-1.amazonaws.com/aws-ecr-s3-sagemaker.yaml",
             param_ResourceName=f"zenml-{random_str(6).lower()}",
-            param_ZenMLServerURL=zenml_server_url,
-            param_ZenMLServerAPIToken=zenml_server_api_token,
+            param_ZenMLServerURL=self.zenml_server_url,
+            param_ZenMLServerAPIToken=self.zenml_server_api_token,
         )
         # Encode the parameters as URL query parameters
         query_params = "&".join([f"{k}={v}" for k, v in params.items()])
@@ -221,69 +242,13 @@ console.
         if self.location:
             region = f"region={self.location}"
 
-        return (
+        url = (
             f"https://console.aws.amazon.com/cloudformation/home?"
-            f"{region}#/stacks/create/review?{query_params}",
-            "AWS CloudFormation Console",
+            f"{region}#/stacks/create/review?{query_params}"
         )
 
-    def get_stack(
-        self, date_start: Optional[datetime.datetime] = None
-    ) -> Optional[DeployedStack]:
-        """Return the ZenML stack that was deployed and registered.
-
-        This method is called to retrieve a ZenML stack matching the deployment
-        provider.
-
-        Args:
-            date_start: The date when the deployment started.
-
-        Returns:
-            The ZenML stack that was deployed and registered or None if a
-            matching stack was not found.
-        """
-        client = Client()
-
-        # It's difficult to find a stack that matches the CloudFormation
-        # deployment 100% because the user can change the stack name before they
-        # deploy the stack in AWS.
-        #
-        # We try to find a full AWS stack that matches the deployment provider
-        # that was registered after this deployment was created.
-
-        # Get all stacks created after the start date
-        stacks = client.list_stacks(
-            created=f"gt:{str(date_start.replace(microsecond=0))}"
-            if date_start
-            else None,
-            sort_by="desc:created",
-            size=50,
+        return StackDeploymentConfig(
+            deployment_url=url,
+            deployment_url_text="AWS CloudFormation Console",
+            configuration=None,
         )
-
-        if not stacks.items:
-            return None
-
-        # Find a stack that best matches the deployment provider
-        for stack in stacks.items:
-            if not stack.labels:
-                continue
-
-            if stack.labels.get("zenml:provider") != self.provider.value:
-                continue
-
-            if stack.labels.get("zenml:deployment") != AWS_DEPLOYMENT_TYPE:
-                continue
-
-            artifact_store = stack.components[
-                StackComponentType.ARTIFACT_STORE
-            ][0]
-
-            if not artifact_store.connector:
-                continue
-
-            return DeployedStack(
-                stack=stack,
-                service_connector=artifact_store.connector,
-            )
-
-        return None