zenml-nightly 0.70.0.dev20241201__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/client.py

@@ -1702,6 +1702,7 @@ class Client(metaclass=ClientMetaClass):
         updated: Optional[datetime] = None,
         type: Optional[str] = None,
         flavor: Optional[str] = None,
+        user: Optional[Union[UUID, str]] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
         hydrate: bool = False,
@@ -1727,6 +1728,7 @@ class Client(metaclass=ClientMetaClass):
             flavor: Use the service flavor for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The id of the user 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.
             running: Use the running status for filtering
@@ -1753,6 +1755,7 @@ class Client(metaclass=ClientMetaClass):
             flavor=flavor,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             running=running,
             name=service_name,
             pipeline_name=pipeline_name,
@@ -2249,6 +2252,7 @@ class Client(metaclass=ClientMetaClass):
         type: Optional[str] = None,
         integration: Optional[str] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[FlavorResponse]:
         """Fetches all the flavor models.
@@ -2262,6 +2266,7 @@ class Client(metaclass=ClientMetaClass):
             created: Use to flavors by time of creation
             updated: Use the last updated date for filtering
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             name: The name of the flavor to filter by.
             type: The type of the flavor to filter by.
             integration: The integration of the flavor to filter by.
@@ -2277,6 +2282,7 @@ class Client(metaclass=ClientMetaClass):
             sort_by=sort_by,
             logical_operator=logical_operator,
             user_id=user_id,
+            user=user,
             name=name,
             type=type,
             integration=integration,
@@ -2661,13 +2667,16 @@ class Client(metaclass=ClientMetaClass):
         updated: Optional[Union[datetime, str]] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         pipeline_id: Optional[Union[str, UUID]] = None,
         stack_id: Optional[Union[str, UUID]] = None,
+        container_registry_id: Optional[Union[UUID, str]] = None,
         is_local: Optional[bool] = None,
         contains_code: Optional[bool] = None,
         zenml_version: Optional[str] = None,
         python_version: Optional[str] = None,
         checksum: Optional[str] = None,
+        stack_checksum: Optional[str] = None,
         hydrate: bool = False,
     ) -> Page[PipelineBuildResponse]:
         """List all builds.
@@ -2682,13 +2691,17 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             pipeline_id: The id of the pipeline to filter by.
             stack_id: The id of the stack to filter by.
+            container_registry_id: The id of the container registry to
+                filter by.
             is_local: Use to filter local builds.
             contains_code: Use to filter builds that contain code.
             zenml_version: The version of ZenML to filter by.
             python_version: The Python version to filter by.
             checksum: The build checksum to filter by.
+            stack_checksum: The stack checksum to filter by.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -2705,13 +2718,16 @@ class Client(metaclass=ClientMetaClass):
             updated=updated,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             pipeline_id=pipeline_id,
             stack_id=stack_id,
+            container_registry_id=container_registry_id,
             is_local=is_local,
             contains_code=contains_code,
             zenml_version=zenml_version,
             python_version=python_version,
             checksum=checksum,
+            stack_checksum=stack_checksum,
         )
         build_filter_model.set_scope_workspace(self.active_workspace.id)
         return self.zen_store.list_builds(
@@ -2771,7 +2787,7 @@ class Client(metaclass=ClientMetaClass):
         allow_name_prefix_match: bool = True,
         hydrate: bool = True,
     ) -> EventSourceResponse:
-        """Get a event source by name, ID or prefix.
+        """Get an event source by name, ID or prefix.
 
         Args:
             name_id_or_prefix: The name, ID or prefix of the stack.
@@ -2804,6 +2820,7 @@ class Client(metaclass=ClientMetaClass):
         event_source_type: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[EventSourceResponse]:
         """Lists all event_sources.
@@ -2818,6 +2835,7 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             name: The name of the event_source to filter by.
             flavor: The flavor of the event_source to filter by.
             event_source_type: The subtype of the event_source to filter by.
@@ -2834,6 +2852,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             name=name,
             flavor=flavor,
             plugin_subtype=event_source_type,
@@ -3001,6 +3020,7 @@ class Client(metaclass=ClientMetaClass):
         action_type: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[ActionResponse]:
         """List actions.
@@ -3015,6 +3035,7 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The id of the user to filter by.
+            user: Filter by user name/ID.
             name: The name of the action to filter by.
             flavor: The flavor of the action to filter by.
             action_type: The type of the action to filter by.
@@ -3031,6 +3052,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             name=name,
             id=id,
             flavor=flavor,
@@ -3179,6 +3201,7 @@ class Client(metaclass=ClientMetaClass):
         action_subtype: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[TriggerResponse]:
         """Lists all triggers.
@@ -3193,6 +3216,7 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             name: The name of the trigger to filter by.
             event_source_id: The event source associated with the trigger.
             action_id: The action associated with the trigger.
@@ -3215,6 +3239,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             name=name,
             event_source_id=event_source_id,
             action_id=action_id,
@@ -3365,6 +3390,7 @@ class Client(metaclass=ClientMetaClass):
         updated: Optional[Union[datetime, str]] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         pipeline_id: Optional[Union[str, UUID]] = None,
         stack_id: Optional[Union[str, UUID]] = None,
         build_id: Optional[Union[str, UUID]] = None,
@@ -3383,6 +3409,7 @@ class Client(metaclass=ClientMetaClass):
             updated: Use the last updated date for filtering
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             pipeline_id: The id of the pipeline to filter by.
             stack_id: The id of the stack to filter by.
             build_id: The id of the build to filter by.
@@ -3403,6 +3430,7 @@ class Client(metaclass=ClientMetaClass):
             updated=updated,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             pipeline_id=pipeline_id,
             stack_id=stack_id,
             build_id=build_id,
@@ -3488,6 +3516,7 @@ class Client(metaclass=ClientMetaClass):
         logical_operator: LogicalOperators = LogicalOperators.AND,
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
+        id: Optional[Union[UUID, str]] = None,
         name: Optional[str] = None,
         tag: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
@@ -3510,6 +3539,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator: Which logical operator to use [and, or].
             created: Filter by the creation date.
             updated: Filter by the last updated date.
+            id: Filter by run template ID.
             name: Filter by run template name.
             tag: Filter by run template tags.
             workspace_id: Filter by workspace ID.
@@ -3534,6 +3564,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             created=created,
             updated=updated,
+            id=id,
             name=name,
             tag=tag,
             workspace_id=workspace_id,
@@ -3650,6 +3681,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         pipeline_id: Optional[Union[str, UUID]] = None,
         orchestrator_id: Optional[Union[str, UUID]] = None,
         active: Optional[Union[str, bool]] = None,
@@ -3674,6 +3706,7 @@ class Client(metaclass=ClientMetaClass):
             name: The name of the stack to filter by.
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             pipeline_id: The id of the pipeline to filter by.
             orchestrator_id: The id of the orchestrator to filter by.
             active: Use to filter by active status.
@@ -3700,6 +3733,7 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             pipeline_id=pipeline_id,
             orchestrator_id=orchestrator_id,
             active=active,
@@ -3940,6 +3974,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,
+        user: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[str, UUID]] = None,
         model: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
@@ -3958,6 +3993,7 @@ class Client(metaclass=ClientMetaClass):
             end_time: Use to filter by the time when the step finished running
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
+            user: Filter by user name/ID.
             pipeline_run_id: The id of the pipeline run to filter by.
             deployment_id: The id of the deployment to filter by.
             original_step_run_id: The id of the original step run to filter by.
@@ -3992,6 +4028,7 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
             model_version_id=model_version_id,
             model=model,
         )
@@ -4664,6 +4701,7 @@ class Client(metaclass=ClientMetaClass):
         scope: Optional[SecretScope] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[SecretResponse]:
         """Fetches all the secret models.
@@ -4683,6 +4721,7 @@ class Client(metaclass=ClientMetaClass):
             scope: The scope of the secret to filter by.
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user 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.
 
@@ -4699,6 +4738,7 @@ class Client(metaclass=ClientMetaClass):
             sort_by=sort_by,
             logical_operator=logical_operator,
             user_id=user_id,
+            user=user,
             workspace_id=workspace_id,
             name=name,
             scope=scope,
@@ -5013,6 +5053,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[CodeRepositoryResponse]:
         """List all code repositories.
@@ -5028,6 +5069,7 @@ class Client(metaclass=ClientMetaClass):
             name: The name of the code repository to filter by.
             workspace_id: The id of the workspace to filter by.
             user_id: The id of the user 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.
 
@@ -5045,6 +5087,7 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             workspace_id=workspace_id,
             user_id=user_id,
+            user=user,
         )
         filter_model.set_scope_workspace(self.active_workspace.id)
         return self.zen_store.list_code_repositories(
@@ -5405,6 +5448,7 @@ class Client(metaclass=ClientMetaClass):
         resource_id: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        user: Optional[Union[UUID, str]] = None,
         labels: Optional[Dict[str, Optional[str]]] = None,
         secret_id: Optional[Union[str, UUID]] = None,
         hydrate: bool = False,
@@ -5427,6 +5471,7 @@ class Client(metaclass=ClientMetaClass):
                 they can give access to.
             workspace_id: The id of the workspace to filter by.
             user_id: The id of the user to filter by.
+            user: Filter by user name/ID.
             name: The name of the service connector to filter by.
             labels: The labels of the service connector to filter by.
             secret_id: Filter by the id of the secret that is referenced by the
@@ -5444,6 +5489,7 @@ class Client(metaclass=ClientMetaClass):
             logical_operator=logical_operator,
             workspace_id=workspace_id or self.active_workspace.id,
             user_id=user_id,
+            user=user,
             name=name,
             connector_type=connector_type,
             auth_method=auth_method,
@@ -6596,6 +6642,7 @@ class Client(metaclass=ClientMetaClass):
         client_id: Union[UUID, str, None] = None,
         status: Union[OAuthDeviceStatus, str, None] = None,
         trusted_device: Union[bool, str, None] = None,
+        user: Optional[Union[UUID, str]] = None,
         failed_auth_attempts: Union[int, str, None] = None,
         last_login: Optional[Union[datetime, str, None]] = None,
         hydrate: bool = False,
@@ -6613,6 +6660,7 @@ class Client(metaclass=ClientMetaClass):
             expires: Use the expiration date for filtering.
             client_id: Use the client id for filtering.
             status: Use the status for filtering.
+            user: Filter by user name/ID.
             trusted_device: Use the trusted device flag for filtering.
             failed_auth_attempts: Use the failed auth attempts for filtering.
             last_login: Use the last login date for filtering.
@@ -6632,6 +6680,7 @@ class Client(metaclass=ClientMetaClass):
             updated=updated,
             expires=expires,
             client_id=client_id,
+            user=user,
             status=status,
             trusted_device=trusted_device,
             failed_auth_attempts=failed_auth_attempts,
@@ -6730,7 +6779,7 @@ class Client(metaclass=ClientMetaClass):
         trigger_execution_id: UUID,
         hydrate: bool = True,
     ) -> TriggerExecutionResponse:
-        """Get an trigger execution by ID.
+        """Get a trigger execution by ID.
 
         Args:
             trigger_execution_id: The ID of the trigger execution to get.
@@ -6751,6 +6800,7 @@ class Client(metaclass=ClientMetaClass):
         size: int = PAGE_SIZE_DEFAULT,
         logical_operator: LogicalOperators = LogicalOperators.AND,
         trigger_id: Optional[UUID] = None,
+        user: Optional[Union[UUID, str]] = None,
         hydrate: bool = False,
     ) -> Page[TriggerExecutionResponse]:
         """List all trigger executions matching the given filter criteria.
@@ -6761,6 +6811,7 @@ class Client(metaclass=ClientMetaClass):
             size: The maximum size of all pages.
             logical_operator: Which logical operator to use [and, or].
             trigger_id: ID of the trigger 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.
 
@@ -6772,6 +6823,7 @@ class Client(metaclass=ClientMetaClass):
             sort_by=sort_by,
             page=page,
             size=size,
+            user=user,
             logical_operator=logical_operator,
         )
         filter_model.set_scope_workspace(self.active_workspace.id)

zenml/config/secret_reference_mixin.py

@@ -66,7 +66,7 @@ class SecretReferenceMixin(BaseModel):
                         "but future versions of ZenML will require you to pass "
                         "in sensitive information as secrets. Check out the "
                         "documentation on how to configure values with secrets "
-                        "here: https://docs.zenml.io/getting-started/deploying-zenml/manage-the-deployed-services/secret-management"
+                        "here: https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                     )
                 continue
 

zenml/config/server_config.py

@@ -27,6 +27,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
     DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
@@ -269,6 +270,9 @@ class ServerConfiguration(BaseModel):
     generic_api_token_lifetime: PositiveInt = (
         DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME
     )
+    generic_api_token_max_lifetime: PositiveInt = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME
+    )
 
     external_login_url: Optional[str] = None
     external_user_info_url: Optional[str] = None

zenml/constants.py

@@ -170,6 +170,9 @@ ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 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 Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
@@ -268,6 +271,9 @@ DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY = 1000
 DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME = 60 * 60  # 1 hour
+DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME = (
+    60 * 60 * 24 * 7
+)  # 7 days
 
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS = (
     "max-age=63072000; includeSubdomains"
@@ -466,3 +472,7 @@ BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
 
 
 STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
+
+ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = handle_int_env_var(
+    ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION, default=0
+)

zenml/image_builders/base_image_builder.py

@@ -25,6 +25,7 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent
 from zenml.stack.stack_component import StackComponentConfig
+from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
     from zenml.container_registries import BaseContainerRegistry
@@ -100,6 +101,7 @@ class BaseImageBuilder(StackComponent, ABC):
     def _upload_build_context(
         build_context: "BuildContext",
         parent_path_directory_name: str,
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> str:
         """Uploads a Docker image build context to a remote location.
 
@@ -109,6 +111,7 @@ class BaseImageBuilder(StackComponent, ABC):
                 the build context to. It will be appended to the artifact
                 store path to create the parent path where the build context
                 will be uploaded to.
+            archive_type: The type of archive to create.
 
         Returns:
             The path to the uploaded build context.
@@ -119,7 +122,7 @@ class BaseImageBuilder(StackComponent, ABC):
 
         hash_ = hashlib.sha1()  # nosec
         with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as f:
-            build_context.write_archive(f, use_gzip=True)
+            build_context.write_archive(f, archive_type)
 
             while True:
                 data = f.read(64 * 1024)
@@ -127,7 +130,7 @@ class BaseImageBuilder(StackComponent, ABC):
                     break
                 hash_.update(data)
 
-            filename = f"{hash_.hexdigest()}.tar.gz"
+            filename = f"{hash_.hexdigest()}.{archive_type.value}"
             filepath = f"{parent_path}/{filename}"
             if not fileio.exists(filepath):
                 logger.info("Uploading build context to `%s`.", filepath)

zenml/image_builders/build_context.py

@@ -20,7 +20,7 @@ from zenml.constants import REPOSITORY_DIRECTORY_NAME
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import io_utils, string_utils
-from zenml.utils.archivable import Archivable
+from zenml.utils.archivable import Archivable, ArchiveType
 
 logger = get_logger(__name__)
 
@@ -69,28 +69,19 @@ class BuildContext(Archivable):
         return None
 
     def write_archive(
-        self, output_file: IO[bytes], use_gzip: bool = True
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            use_gzip: Whether to use `gzip` to compress the file.
+            archive_type: The type of archive to create.
         """
-        from docker.utils import build as docker_build_utils
-
-        files = self.get_files()
-        extra_files = self.get_extra_files()
-
-        context_archive = docker_build_utils.create_archive(
-            fileobj=output_file,
-            root=self._root,
-            files=sorted(files.keys()),
-            gzip=use_gzip,
-            extra_files=list(extra_files.items()),
-        )
+        super().write_archive(output_file, archive_type)
 
-        build_context_size = os.path.getsize(context_archive.name)
+        build_context_size = os.path.getsize(output_file.name)
         if (
             self._root
             and build_context_size > 50 * 1024 * 1024

zenml/integrations/aws/__init__.py

@@ -33,6 +33,7 @@ AWS_SAGEMAKER_ORCHESTRATOR_FLAVOR = "sagemaker"
 AWS_CONNECTOR_TYPE = "aws"
 AWS_RESOURCE_TYPE = "aws-generic"
 S3_RESOURCE_TYPE = "s3-bucket"
+AWS_IMAGE_BUILDER_FLAVOR = "aws"
 
 class AWSIntegration(Integration):
     """Definition of AWS integration for ZenML."""
@@ -59,12 +60,14 @@ class AWSIntegration(Integration):
         """
         from zenml.integrations.aws.flavors import (
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerOrchestratorFlavor,
             SagemakerStepOperatorFlavor,
         )
 
         return [
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerStepOperatorFlavor,
             SagemakerOrchestratorFlavor,
         ]

zenml/integrations/aws/flavors/__init__.py

@@ -17,6 +17,10 @@ from zenml.integrations.aws.flavors.aws_container_registry_flavor import (
     AWSContainerRegistryConfig,
     AWSContainerRegistryFlavor,
 )
+from zenml.integrations.aws.flavors.aws_image_builder_flavor import (
+    AWSImageBuilderConfig,
+    AWSImageBuilderFlavor,
+)
 from zenml.integrations.aws.flavors.sagemaker_orchestrator_flavor import (
     SagemakerOrchestratorConfig,
     SagemakerOrchestratorFlavor,
@@ -29,6 +33,8 @@ from zenml.integrations.aws.flavors.sagemaker_step_operator_flavor import (
 __all__ = [
     "AWSContainerRegistryFlavor",
     "AWSContainerRegistryConfig",
+    "AWSImageBuilderConfig",
+    "AWSImageBuilderFlavor",
     "SagemakerStepOperatorFlavor",
     "SagemakerStepOperatorConfig",
     "SagemakerOrchestratorFlavor",