zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/config/compiler.py

@@ -39,6 +39,7 @@ from zenml.config.step_configurations import (
 from zenml.environment import get_run_environment_dict
 from zenml.exceptions import StackValidationError
 from zenml.models import PipelineDeploymentBase
+from zenml.new.pipelines.run_utils import get_default_run_name
 from zenml.utils import pydantic_utils, settings_utils
 
 if TYPE_CHECKING:
@@ -74,7 +75,7 @@ class Compiler:
         pipeline: "Pipeline",
         stack: "Stack",
         run_configuration: PipelineRunConfiguration,
-    ) -> Tuple[PipelineDeploymentBase, PipelineSpec]:
+    ) -> PipelineDeploymentBase:
         """Compiles a ZenML pipeline to a serializable representation.
 
         Args:
@@ -83,7 +84,7 @@ class Compiler:
             run_configuration: The run configuration for this pipeline.
 
         Returns:
-            The compiled pipeline deployment and spec
+            The compiled pipeline deployment.
         """
         logger.debug("Compiling pipeline `%s`.", pipeline.name)
         # Copy the pipeline before we apply any run-level configurations, so
@@ -127,12 +128,17 @@ class Compiler:
 
         self._ensure_required_stack_components_exist(stack=stack, steps=steps)
 
-        run_name = run_configuration.run_name or self._get_default_run_name(
+        run_name = run_configuration.run_name or get_default_run_name(
             pipeline_name=pipeline.name
         )
 
         client_version, server_version = get_zenml_versions()
 
+        step_specs = [step.spec for step in steps.values()]
+        pipeline_spec = self._compute_pipeline_spec(
+            pipeline=pipeline, step_specs=step_specs
+        )
+
         deployment = PipelineDeploymentBase(
             run_name_template=run_name,
             pipeline_configuration=pipeline.configuration,
@@ -140,17 +146,15 @@ class Compiler:
             client_environment=get_run_environment_dict(),
             client_version=client_version,
             server_version=server_version,
-        )
-
-        step_specs = [step.spec for step in steps.values()]
-        pipeline_spec = self._compute_pipeline_spec(
-            pipeline=pipeline, step_specs=step_specs
+            pipeline_version_hash=pipeline._compute_unique_identifier(
+                pipeline_spec=pipeline_spec
+            ),
+            pipeline_spec=pipeline_spec,
         )
 
         logger.debug("Compiled pipeline deployment: %s", deployment)
-        logger.debug("Compiled pipeline spec: %s", pipeline_spec)
 
-        return deployment, pipeline_spec
+        return deployment
 
     def compile_spec(self, pipeline: "Pipeline") -> PipelineSpec:
         """Compiles a ZenML pipeline to a pipeline spec.
@@ -474,18 +478,6 @@ class Compiler:
         )
         return Step(spec=step_spec, config=complete_step_configuration)
 
-    @staticmethod
-    def _get_default_run_name(pipeline_name: str) -> str:
-        """Gets the default name for a pipeline run.
-
-        Args:
-            pipeline_name: Name of the pipeline which will be run.
-
-        Returns:
-            Run name.
-        """
-        return f"{pipeline_name}-{{date}}-{{time}}"
-
     def _get_sorted_invocations(
         self,
         pipeline: "Pipeline",

zenml/config/docker_settings.py

@@ -16,8 +16,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Field, model_validator
-from pydantic_settings import SettingsConfigDict
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
@@ -49,15 +48,6 @@ class PythonEnvironmentExportMethod(Enum):
         }[self]
 
 
-class SourceFileMode(Enum):
-    """Different methods to handle source files in Docker images."""
-
-    INCLUDE = "include"
-    DOWNLOAD_OR_INCLUDE = "download_or_include"
-    DOWNLOAD = "download"
-    IGNORE = "ignore"
-
-
 class PythonPackageInstaller(Enum):
     """Different installers for python packages."""
 
@@ -102,7 +92,6 @@ class DockerSettings(BaseSettings):
     Depending on the configuration of this object, requirements will be
     installed in the following order (each step optional):
     - The packages installed in your local python environment
-    - The packages specified via the `required_hub_plugins` attribute
     - The packages required by the stack unless this is disabled by setting
       `install_stack_requirements=False`.
     - The packages specified via the `required_integrations`
@@ -135,10 +124,9 @@ class DockerSettings(BaseSettings):
             when the `dockerfile` attribute is set. If this is left empty, the
             build context will only contain the Dockerfile.
         parent_image_build_config: Configuration for the parent image build.
-        build_options: DEPRECATED, use parent_image_build_config.build_options
-            instead.
         skip_build: If set to `True`, the parent image will be used directly to
             run the steps of your pipeline.
+        prevent_build_reuse: Prevent the reuse of an existing build.
         target_repository: Name of the Docker repository to which the
             image should be pushed. This repository will be appended to the
             registry URI of the container registry of your stack and should
@@ -161,11 +149,7 @@ class DockerSettings(BaseSettings):
         required_integrations: List of ZenML integrations that should be
             installed. All requirements for the specified integrations will
             be installed inside the Docker image.
-        required_hub_plugins: List of ZenML Hub plugins to install.
-            Expected format: '(<author_username>/)<plugin_name>==<version>'.
-            If no version is specified, the latest version is taken. The
-            packages of required plugins and all their dependencies will be
-            installed inside the Docker image.
+        required_hub_plugins: DEPRECATED/UNUSED.
         install_stack_requirements: If `True`, ZenML will automatically detect
             if components of your active stack are part of a ZenML integration
             and install the corresponding requirements and apt packages.
@@ -176,33 +160,32 @@ class DockerSettings(BaseSettings):
         environment: Dictionary of environment variables to set inside the
             Docker image.
         build_config: Configuration for the main image build.
-        dockerignore: DEPRECATED, use build_config.dockerignore instead.
-        copy_files: DEPRECATED, use the `source_files` attribute instead.
-        copy_global_config: DEPRECATED/UNUSED.
         user: If not `None`, will set the user, make it owner of the `/app`
             directory which contains all the user code and run the container
             entrypoint as this user.
-        source_files: Defines how the user source files will be handled when
-            building the Docker image.
-            * INCLUDE: The files will be included in the Docker image.
-            * DOWNLOAD: The files will be downloaded when running the image. If
-              this is specified, the files must be inside a registered code
-              repository and the repository must have no local changes,
-              otherwise the build will fail.
-            * DOWNLOAD_OR_INCLUDE: The files will be downloaded if they're
-              inside a registered code repository and the repository has no
-              local changes, otherwise they will be included in the image.
-            * IGNORE: The files will not be included or downloaded in the image.
-              If you use this option, you're responsible that all the files
-              to run your steps exist in the right place.
+        allow_including_files_in_images: If `True`, code can be included in the
+            Docker images if code download from a code repository or artifact
+            store is disabled or not possible.
+        allow_download_from_code_repository: If `True`, code can be downloaded
+            from a code repository if possible.
+        allow_download_from_artifact_store: If `True`, code can be downloaded
+            from the artifact store.
+        build_options: DEPRECATED, use parent_image_build_config.build_options
+            instead.
+        dockerignore: DEPRECATED, use build_config.dockerignore instead.
+        copy_files: DEPRECATED/UNUSED.
+        copy_global_config: DEPRECATED/UNUSED.
+        source_files: DEPRECATED. Use allow_including_files_in_images,
+            allow_download_from_code_repository and
+            allow_download_from_artifact_store instead.
     """
 
     parent_image: Optional[str] = None
     dockerfile: Optional[str] = None
     build_context_root: Optional[str] = None
-    build_options: Dict[str, Any] = {}
     parent_image_build_config: Optional[DockerBuildConfig] = None
     skip_build: bool = False
+    prevent_build_reuse: bool = False
     target_repository: Optional[str] = None
     python_package_installer: PythonPackageInstaller = (
         PythonPackageInstaller.PIP
@@ -215,49 +198,89 @@ class DockerSettings(BaseSettings):
         default=None, union_mode="left_to_right"
     )
     required_integrations: List[str] = []
-    required_hub_plugins: List[str] = []
     install_stack_requirements: bool = True
     apt_packages: List[str] = []
     environment: Dict[str, Any] = {}
-    dockerignore: Optional[str] = None
-    copy_files: bool = True
-    copy_global_config: bool = True
     user: Optional[str] = None
     build_config: Optional[DockerBuildConfig] = None
 
-    source_files: SourceFileMode = SourceFileMode.DOWNLOAD_OR_INCLUDE
+    allow_including_files_in_images: bool = True
+    allow_download_from_code_repository: bool = True
+    allow_download_from_artifact_store: bool = True
+
+    # Deprecated attributes
+    build_options: Dict[str, Any] = {}
+    dockerignore: Optional[str] = None
+    copy_files: bool = True
+    copy_global_config: bool = True
+    source_files: Optional[str] = None
+    required_hub_plugins: List[str] = []
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
-        "copy_files", "copy_global_config"
+        "copy_files",
+        "copy_global_config",
+        "source_files",
+        "required_hub_plugins",
     )
 
     @model_validator(mode="before")
     @classmethod
     @before_validator_handler
-    def _migrate_copy_files(cls, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Migrates the value from the old copy_files attribute.
+    def _migrate_source_files(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Migrate old source_files values.
 
         Args:
-            data: The settings values.
+            data: The model data.
+
+        Raises:
+            ValueError: If an invalid source file mode is specified.
 
         Returns:
-            The migrated settings values.
+            The migrated data.
         """
-        copy_files = data.get("copy_files", None)
+        source_files = data.get("source_files", None)
 
-        if copy_files is None:
+        if source_files is None:
             return data
 
-        if data.get("source_files", None):
-            # Ignore the copy files value in favor of the new source files
+        replacement_attributes = [
+            "allow_including_files_in_images",
+            "allow_download_from_code_repository",
+            "allow_download_from_artifact_store",
+        ]
+        if any(v in data for v in replacement_attributes):
             logger.warning(
-                "Both `copy_files` and `source_files` specified for the "
-                "DockerSettings, ignoring the `copy_files` value."
+                "Both `source_files` and one of %s specified for the "
+                "DockerSettings, ignoring the `source_files` value.",
+                replacement_attributes,
             )
-        elif copy_files is True:
-            data["source_files"] = SourceFileMode.INCLUDE
-        elif copy_files is False:
-            data["source_files"] = SourceFileMode.IGNORE
+            return data
+
+        allow_including_files_in_images = False
+        allow_download_from_code_repository = False
+        allow_download_from_artifact_store = False
+
+        if source_files == "download":
+            allow_download_from_code_repository = True
+        elif source_files == "include":
+            allow_including_files_in_images = True
+        elif source_files == "download_or_include":
+            allow_including_files_in_images = True
+            allow_download_from_code_repository = True
+        elif source_files == "ignore":
+            pass
+        else:
+            raise ValueError(f"Invalid source file mode `{source_files}`.")
+
+        data["allow_including_files_in_images"] = (
+            allow_including_files_in_images
+        )
+        data["allow_download_from_code_repository"] = (
+            allow_download_from_code_repository
+        )
+        data["allow_download_from_artifact_store"] = (
+            allow_download_from_artifact_store
+        )
 
         return data
 
@@ -282,7 +305,7 @@ class DockerSettings(BaseSettings):
 
         return self
 
-    model_config = SettingsConfigDict(
+    model_config = ConfigDict(
         # public attributes are immutable
         frozen=True,
         # prevent extra attributes during model initialization

zenml/config/pipeline_run_configuration.py

@@ -21,6 +21,7 @@ from pydantic import Field, SerializeAsAny
 from zenml.config.base_settings import BaseSettings
 from zenml.config.retry_config import StepRetryConfig
 from zenml.config.schedule import Schedule
+from zenml.config.source import SourceWithValidator
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
@@ -48,3 +49,5 @@ class PipelineRunConfiguration(
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
     retry: Optional[StepRetryConfig] = None
+    failure_hook_source: Optional[SourceWithValidator] = None
+    success_hook_source: Optional[SourceWithValidator] = None

zenml/config/server_config.py

@@ -65,6 +65,9 @@ def generate_jwt_secret_key() -> str:
 class ServerConfiguration(BaseModel):
     """ZenML Server configuration attributes.
 
+    All these attributes can be set through the environment with the `ZENML_SERVER_`-Prefix.
+    The value of the `ZENML_SERVER_DEPLOYMENT_TYPE` environment variable will be extracted to deployment_type.
+
     Attributes:
         deployment_type: The type of ZenML server deployment that is running.
         server_url: The URL where the ZenML server API is reachable. Must be

zenml/config/source.py

@@ -42,6 +42,7 @@ class SourceType(Enum):
     INTERNAL = "internal"
     DISTRIBUTION_PACKAGE = "distribution_package"
     CODE_REPOSITORY = "code_repository"
+    NOTEBOOK = "notebook"
     UNKNOWN = "unknown"
 
 
@@ -229,6 +230,63 @@ class CodeRepositorySource(Source):
         return value
 
 
+class NotebookSource(Source):
+    """Source representing an object defined in a notebook.
+
+    Attributes:
+        code_path: Path where the notebook cell code for this source is
+            uploaded.
+        replacement_module: Name of the module from which this source should
+            be loaded in case the code is not running in a notebook.
+    """
+
+    code_path: Optional[str] = None
+    replacement_module: Optional[str] = None
+    type: SourceType = SourceType.NOTEBOOK
+
+    # Private attribute that is used to store the code but should not be
+    # serialized
+    _cell_code: Optional[str] = None
+
+    @field_validator("type")
+    @classmethod
+    def _validate_type(cls, value: SourceType) -> SourceType:
+        """Validate the source type.
+
+        Args:
+            value: The source type.
+
+        Raises:
+            ValueError: If the source type is not `NOTEBOOK`.
+
+        Returns:
+            The source type.
+        """
+        if value != SourceType.NOTEBOOK:
+            raise ValueError("Invalid source type.")
+
+        return value
+
+    @field_validator("module")
+    @classmethod
+    def _validate_module(cls, value: str) -> str:
+        """Validate the module.
+
+        Args:
+            value: The module.
+
+        Raises:
+            ValueError: If the module is not `__main__`.
+
+        Returns:
+            The module.
+        """
+        if value != "__main__":
+            raise ValueError("Invalid module for notebook source.")
+
+        return value
+
+
 def convert_source(source: Any) -> Any:
     """Converts an old source string to a source object.
 
@@ -245,5 +303,6 @@ def convert_source(source: Any) -> Any:
 
 
 SourceWithValidator = Annotated[
-    SerializeAsAny[Source], BeforeValidator(convert_source)
+    SerializeAsAny[Source],
+    BeforeValidator(convert_source),
 ]

zenml/constants.py

@@ -165,7 +165,6 @@ ENV_ZENML_SKIP_IMAGE_BUILDER_DEFAULT = "ZENML_SKIP_IMAGE_BUILDER_DEFAULT"
 ENV_ZENML_REQUIRES_CODE_DOWNLOAD = "ZENML_REQUIRES_CODE_DOWNLOAD"
 ENV_ZENML_SERVER = "ZENML_SERVER"
 ENV_ZENML_LOCAL_SERVER = "ZENML_LOCAL_SERVER"
-ENV_ZENML_HUB_URL = "ZENML_HUB_URL"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
@@ -279,7 +278,6 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT = "nosniff"
 _csp_script_src_urls = ["https://widgets-v3.featureos.app"]
 _csp_connect_src_urls = [
     "https://sdkdocs.zenml.io",
-    "https://hubapi.zenml.io",
     "https://analytics.zenml.io",
 ]
 _csp_img_src_urls = [
@@ -315,6 +313,7 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS = (
 )
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
 DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD = False
+DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS = 30
 
 # Configurations to decide which resources report their usage and check for
 # entitlement in the case of a cloud deployment. Expected Format is this:
@@ -364,6 +363,7 @@ PIPELINES = "/pipelines"
 PIPELINE_SPEC = "/pipeline-spec"
 PLUGIN_FLAVORS = "/plugin-flavors"
 RUNS = "/runs"
+RUN_TEMPLATES = "/run_templates"
 RUN_METADATA = "/run-metadata"
 SCHEDULES = "/schedules"
 SECRETS = "/secrets"
@@ -424,6 +424,7 @@ PAGE_SIZE_MAXIMUM: int = handle_int_env_var(
     ENV_ZENML_PAGINATION_DEFAULT_LIMIT, default=10000
 )
 FILTERING_DATETIME_FORMAT: str = "%Y-%m-%d %H:%M:%S"
+SORT_PIPELINES_BY_LATEST_RUN_KEY = "latest_run"
 
 # Metadata constants
 METADATA_ORCHESTRATOR_URL = "orchestrator_url"
@@ -492,3 +493,11 @@ BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
 
 
 STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
+
+# ZenML Pro
+ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT = (
+    "\nHINT: Since you are trying to communicate with the ZenML Pro Tenant, "
+    "please make sure that your tenant is in RUNNING state on your "
+    "Organization page. If the tenant is PAUSED you can `Resume` it via UI "
+    "and try again."
+)

zenml/entrypoints/base_entrypoint_configuration.py

@@ -27,10 +27,15 @@ from zenml.constants import (
     handle_bool_env_var,
 )
 from zenml.logger import get_logger
-from zenml.utils import code_repository_utils, source_utils, uuid_utils
+from zenml.utils import (
+    code_repository_utils,
+    code_utils,
+    source_utils,
+    uuid_utils,
+)
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import CodeReferenceResponse, PipelineDeploymentResponse
 
 logger = get_logger(__name__)
 DEFAULT_ENTRYPOINT_COMMAND = [
@@ -198,7 +203,7 @@ class BaseEntrypointConfiguration(ABC):
 
         Raises:
             RuntimeError: If the current environment requires code download
-                but the deployment does not have an associated code reference.
+                but the deployment does not have a reference to any code.
         """
         requires_code_download = handle_bool_env_var(
             ENV_ZENML_REQUIRES_CODE_DOWNLOAD
@@ -207,17 +212,33 @@ class BaseEntrypointConfiguration(ABC):
         if not requires_code_download:
             return
 
-        code_reference = deployment.code_reference
-        if not code_reference:
+        if code_reference := deployment.code_reference:
+            self.download_code_from_code_repository(
+                code_reference=code_reference
+            )
+        elif code_path := deployment.code_path:
+            self.download_code_from_artifact_store(code_path=code_path)
+        else:
             raise RuntimeError(
-                "Code download required but no code reference provided."
+                "Code download required but no code reference or path provided."
             )
 
+        logger.info("Code download finished.")
+
+    def download_code_from_code_repository(
+        self, code_reference: "CodeReferenceResponse"
+    ) -> None:
+        """Download code from a code repository.
+
+        Args:
+            code_reference: The reference to the code.
+        """
         logger.info(
             "Downloading code from code repository `%s` (commit `%s`).",
             code_reference.code_repository.name,
             code_reference.commit,
         )
+
         model = Client().get_code_repository(code_reference.code_repository.id)
         repo = BaseCodeRepository.from_model(model)
         code_repo_root = os.path.abspath("code")
@@ -234,10 +255,34 @@ class BaseEntrypointConfiguration(ABC):
         code_repository_utils.set_custom_local_repository(
             root=code_repo_root, commit=code_reference.commit, repo=repo
         )
-        # Add downloaded file directory to python path
+
         sys.path.insert(0, download_dir)
+        os.chdir(download_dir)
 
-        logger.info("Code download finished.")
+    def download_code_from_artifact_store(self, code_path: str) -> None:
+        """Download code from the artifact store.
+
+        Args:
+            code_path: Path where the code is stored.
+        """
+        logger.info(
+            "Downloading code from artifact store path `%s`.", code_path
+        )
+
+        # Do not remove this line, we need to instantiate the artifact store to
+        # register the filesystem needed for the file download
+        _ = Client().active_stack.artifact_store
+
+        extract_dir = os.path.abspath("code")
+        os.makedirs(extract_dir)
+
+        code_utils.download_and_extract_code(
+            code_path=code_path, extract_dir=extract_dir
+        )
+
+        source_utils.set_custom_source_root(extract_dir)
+        sys.path.insert(0, extract_dir)
+        os.chdir(extract_dir)
 
     @abstractmethod
     def run(self) -> None:

zenml/enums.py

@@ -171,7 +171,6 @@ class CliCategories(StrEnum):
 
     STACK_COMPONENTS = "Stack Components"
     MODEL_DEPLOYMENT = "Model Deployment"
-    HUB = "ZenML Hub"
     INTEGRATIONS = "Integrations"
     MANAGEMENT_TOOLS = "Management Tools"
     MODEL_CONTROL_PLANE = "Model Control Plane"
@@ -303,6 +302,7 @@ class EnvironmentType(StrEnum):
     NOTEBOOK = "notebook"
     PAPERSPACE = "paperspace"
     WSL = "wsl"
+    LIGHTNING_AI_STUDIO = "lightning_ai_studio"
 
 
 class ModelStages(StrEnum):
@@ -338,6 +338,9 @@ class TaggableResourceTypes(StrEnum):
     ARTIFACT_VERSION = "artifact_version"
     MODEL = "model"
     MODEL_VERSION = "model_version"
+    PIPELINE = "pipeline"
+    PIPELINE_RUN = "pipeline_run"
+    RUN_TEMPLATE = "run_template"
 
 
 class ResponseUpdateStrategy(StrEnum):

zenml/environment.py

@@ -15,7 +15,6 @@
 
 import os
 import platform
-from importlib.util import find_spec
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, cast
 
@@ -70,6 +69,8 @@ def get_environment() -> str:
         return EnvironmentType.BITBUCKET_CI
     elif Environment.in_ci():
         return EnvironmentType.GENERIC_CI
+    elif Environment.in_lightning_ai_studio():
+        return EnvironmentType.LIGHTNING_AI_STUDIO
     elif Environment.in_docker():
         return EnvironmentType.DOCKER
     elif Environment.in_container():
@@ -257,15 +258,17 @@ class Environment(metaclass=SingletonMetaClass):
         if Environment.in_google_colab():
             return True
 
-        if find_spec("IPython") is not None:
-            from IPython import get_ipython
+        try:
+            ipython = get_ipython()  # type: ignore[name-defined]
+        except NameError:
+            return False
 
-            if get_ipython().__class__.__name__ in [
-                "TerminalInteractiveShell",
-                "ZMQInteractiveShell",
-                "DatabricksShell",
-            ]:
-                return True
+        if ipython.__class__.__name__ in [
+            "TerminalInteractiveShell",
+            "ZMQInteractiveShell",
+            "DatabricksShell",
+        ]:
+            return True
         return False
 
     @staticmethod
@@ -339,6 +342,19 @@ class Environment(metaclass=SingletonMetaClass):
         """
         return "microsoft-standard" in platform.uname().release
 
+    @staticmethod
+    def in_lightning_ai_studio() -> bool:
+        """If the current Python process is running in Lightning.ai studios.
+
+        Returns:
+            `True` if the current Python process is running in Lightning.ai studios,
+            `False` otherwise.
+        """
+        return (
+            "LIGHTNING_CLOUD_URL" in os.environ
+            and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
+        )
+
     def register_component(
         self, component: "BaseEnvironmentComponent"
     ) -> "BaseEnvironmentComponent":

zenml/image_builders/base_image_builder.py

@@ -119,7 +119,7 @@ class BaseImageBuilder(StackComponent, ABC):
 
         hash_ = hashlib.sha1()  # nosec
         with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as f:
-            build_context.write_archive(f, gzip=True)
+            build_context.write_archive(f, use_gzip=True)
 
             while True:
                 data = f.read(64 * 1024)