From ec83b0df5fa0ddb3dbcbecb7773a2d00fe7bdb8e Mon Sep 17 00:00:00 2001 From: coffeecoder08 Date: Thu, 25 Jun 2026 23:58:40 +0530 Subject: [PATCH 1/2] Add opt-in Modal app cleanup after sandbox completion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After a step sandbox finishes, the Modal app entry remains in the dashboard as 'deployed, Tasks: 0'. This is purely visual clutter— idle apps incur no compute charges—but accumulates quickly for users running many ZenML pipelines or steps on Modal. Add a new stop_app_after_run: bool = False setting to ModalSettingsMixin (shared by both the step operator and orchestrator). When True, ZenML overrides cleanup_step_submission() to call modal app stop --yes on the per-step app, moving its entry from 'deployed' to 'stopped' in the Modal dashboard. Stopped apps remain visible in the 'recently stopped' section with their full log history intact, so debugging is unaffected. Changes: - modal_base_flavor.py: add stop_app_after_run to ModalSettingsMixin - sandbox_utils.py: add stop_modal_app() helper that delegates to the Modal CLI and forwards token credentials from ZenML config - modal_step_operator.py: override cleanup_step_submission() (the hook added in #4990) to call stop_modal_app() when the setting is enabled; failures are logged as warnings and never re-raised Fixes #4889 --- .../modal/flavors/modal_base_flavor.py | 9 ++++ src/zenml/integrations/modal/sandbox_utils.py | 54 +++++++++++++++++++ .../step_operators/modal_step_operator.py | 42 +++++++++++++++ 3 files changed, 105 insertions(+) diff --git a/src/zenml/integrations/modal/flavors/modal_base_flavor.py b/src/zenml/integrations/modal/flavors/modal_base_flavor.py index e3a3d8b753e..9f5cbd79766 100644 --- a/src/zenml/integrations/modal/flavors/modal_base_flavor.py +++ b/src/zenml/integrations/modal/flavors/modal_base_flavor.py @@ -68,6 +68,15 @@ class ModalSettingsMixin(BaseSettings): f"Examples: 3600 (1 hour), 7200 (2 hours), {DEFAULT_TIMEOUT_SECONDS} (24 hours maximum). " "Execution will be terminated if it exceeds this timeout", ) + stop_app_after_run: bool = Field( + False, + description="Stop the Modal app after the sandbox reaches a terminal " + "state. When True, ZenML calls `modal app stop` on the app created " + "for this run, removing it from the Modal dashboard's deployed-apps " + "list. Stopped apps are still visible in the 'recently stopped' " + "section and their logs remain accessible. Defaults to False to " + "preserve the previous behavior.", + ) class ModalCredentialsMixin(BaseSettings): diff --git a/src/zenml/integrations/modal/sandbox_utils.py b/src/zenml/integrations/modal/sandbox_utils.py index 708b7d651fe..ee3a1a78666 100644 --- a/src/zenml/integrations/modal/sandbox_utils.py +++ b/src/zenml/integrations/modal/sandbox_utils.py @@ -14,6 +14,7 @@ """Shared Modal Sandbox helpers.""" import math +import os import time from typing import Any, Dict, List, Optional, Protocol, Tuple @@ -398,3 +399,56 @@ def terminate_sandbox( """Terminate a Modal sandbox by ID.""" sandbox = get_sandbox_by_id(sandbox_id, modal_client=modal_client) sandbox.terminate() + + +def stop_modal_app( + app_name: str, + *, + modal_environment: Optional[str], + token_id: Optional[str] = None, + token_secret: Optional[str] = None, +) -> None: + """Stop a deployed Modal app by name using the Modal CLI. + + Moves the app from 'deployed' to 'stopped' in the Modal dashboard, + eliminating the visual clutter of idle ``Tasks: 0`` app entries. + Stopped apps remain visible in the 'recently stopped' section with + their logs intact, so debugging is not affected. + + This should only be called after the associated Modal sandbox has + reached a terminal state, because stopping the app does not wait for + running sandboxes to finish. + + Args: + app_name: The Modal app name to stop (e.g. ``zenml--``). + modal_environment: Modal environment name passed as ``--env``. + ``None`` uses the CLI's default environment. + token_id: Modal token ID to forward as ``MODAL_TOKEN_ID``. When + ``None``, the ambient CLI credential is used. + token_secret: Modal token secret to forward as + ``MODAL_TOKEN_SECRET``. When ``None``, the ambient CLI + credential is used. + + Raises: + RuntimeError: If the ``modal app stop`` command exits with a + non-zero status code. + """ + import subprocess + + cmd = ["modal", "app", "stop", "--yes", app_name] + if modal_environment: + cmd += ["--env", modal_environment] + + env = os.environ.copy() + token_pair = resolve_modal_token_pair( + token_id=token_id, token_secret=token_secret + ) + if token_pair: + env[MODAL_TOKEN_ID_ENV_KEY] = token_pair[0] + env[MODAL_TOKEN_SECRET_ENV_KEY] = token_pair[1] + + result = subprocess.run(cmd, capture_output=True, text=True, env=env) + if result.returncode != 0: + raise RuntimeError( + f"Failed to stop Modal app '{app_name}': {result.stderr.strip()}" + ) diff --git a/src/zenml/integrations/modal/step_operators/modal_step_operator.py b/src/zenml/integrations/modal/step_operators/modal_step_operator.py index 402abe979eb..c66cfbc9117 100644 --- a/src/zenml/integrations/modal/step_operators/modal_step_operator.py +++ b/src/zenml/integrations/modal/step_operators/modal_step_operator.py @@ -267,3 +267,45 @@ def cancel(self, step_run: "StepRunResponse") -> None: sandbox_id = str(step_run.run_metadata[STEP_SANDBOX_ID_METADATA_KEY]) modal_client = self._get_modal_client() sandbox_utils.terminate_sandbox(sandbox_id, modal_client=modal_client) + + def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: + """Stop the Modal app after a sandbox run completes. + + When ``stop_app_after_run`` is enabled in the step settings, this + calls ``modal app stop`` on the per-step app, moving its entry from + 'deployed' to 'stopped' in the Modal dashboard. Stopped apps remain + visible in the 'recently stopped' section with their full log history + intact, so debugging is unaffected. Failures are logged as warnings + and never re-raised so that cleanup cannot mask the step result. + + Args: + step_run: The finished step run. + """ + settings = cast(ModalStepOperatorSettings, self.get_settings(step_run)) + if not settings.stop_app_after_run: + return + + app_name = f"zenml-{step_run.id}-{step_run.name}"[:64] + modal_environment = sandbox_utils.normalize_optional_config_value( + settings.modal_environment + ) + try: + sandbox_utils.stop_modal_app( + app_name, + modal_environment=modal_environment, + token_id=self.config.token_id, + token_secret=self.config.token_secret, + ) + logger.debug( + "Stopped Modal app '%s' for step run '%s'.", + app_name, + step_run.id, + ) + except Exception: + logger.warning( + "Failed to stop Modal app '%s' for step run '%s'. " + "The app may still appear in the Modal dashboard.", + app_name, + step_run.id, + exc_info=True, + ) From 29a95620407bc19a90c1b93f76255ca1a5a24633 Mon Sep 17 00:00:00 2001 From: coffeecoder08 Date: Sun, 26 Jul 2026 20:08:04 +0530 Subject: [PATCH 2/2] Address reviewer comments on Modal app cleanup - Move stop_app setting to ModalStepOperatorSettings and rename from stop_app_after_run to stop_app - Move subprocess import to top-level in sandbox_utils.py - Extract get_modal_app_name helper function in sandbox_utils.py - Simplify docstrings in sandbox_utils.py and modal_step_operator.py - Handle timeout and CLI fallback in stop_modal_app - Respect STEP_MODAL_ENVIRONMENT_METADATA_KEY in cleanup_step_submission - Update documentation and add unit tests --- .../component-guide/step-operators/modal.md | 2 + .../modal/flavors/modal_base_flavor.py | 9 -- .../flavors/modal_step_operator_flavor.py | 8 ++ src/zenml/integrations/modal/sandbox_utils.py | 52 ++++++++---- .../step_operators/modal_step_operator.py | 29 ++++--- .../test_modal_step_operator.py | 84 +++++++++++++++++++ 6 files changed, 149 insertions(+), 35 deletions(-) diff --git a/docs/book/component-guide/step-operators/modal.md b/docs/book/component-guide/step-operators/modal.md index d1e6e196e8a..8dc661e5197 100644 --- a/docs/book/component-guide/step-operators/modal.md +++ b/docs/book/component-guide/step-operators/modal.md @@ -113,6 +113,7 @@ modal_settings = ModalStepOperatorSettings( # cloud="aws", # optional, enterprise/team only # modal_environment="main", # optional Modal environment name # timeout=86400, # optional sandbox timeout in seconds + # stop_app=True, # optional; stop the Modal app after the step finishes ) resource_settings = ResourceSettings( @@ -133,6 +134,7 @@ def my_modal_step(): ``` Important: +- Set `stop_app=True` in `ModalStepOperatorSettings` if you want ZenML to stop the per-step Modal app after execution completes. Stopped apps move from 'deployed' to 'stopped' in the Modal dashboard while preserving full log history. - If you request GPUs with `ResourceSettings.gpu_count > 0`, you must also specify a GPU type via `ModalStepOperatorSettings.gpu`; otherwise the run will fail with a validation error. - If a GPU type is set but `gpu_count == 0`, ZenML treats the step as CPU-only and logs a warning that the GPU type is ignored. - If `gpu_count` is omitted and a GPU type is set, Modal uses one GPU of that type. diff --git a/src/zenml/integrations/modal/flavors/modal_base_flavor.py b/src/zenml/integrations/modal/flavors/modal_base_flavor.py index 9f5cbd79766..e3a3d8b753e 100644 --- a/src/zenml/integrations/modal/flavors/modal_base_flavor.py +++ b/src/zenml/integrations/modal/flavors/modal_base_flavor.py @@ -68,15 +68,6 @@ class ModalSettingsMixin(BaseSettings): f"Examples: 3600 (1 hour), 7200 (2 hours), {DEFAULT_TIMEOUT_SECONDS} (24 hours maximum). " "Execution will be terminated if it exceeds this timeout", ) - stop_app_after_run: bool = Field( - False, - description="Stop the Modal app after the sandbox reaches a terminal " - "state. When True, ZenML calls `modal app stop` on the app created " - "for this run, removing it from the Modal dashboard's deployed-apps " - "list. Stopped apps are still visible in the 'recently stopped' " - "section and their logs remain accessible. Defaults to False to " - "preserve the previous behavior.", - ) class ModalCredentialsMixin(BaseSettings): diff --git a/src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py b/src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py index ff17d2bd03f..0769eb01a36 100644 --- a/src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py +++ b/src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py @@ -15,6 +15,8 @@ from typing import TYPE_CHECKING, Optional, Type +from pydantic import Field + from zenml.integrations.modal import MODAL_STEP_OPERATOR_FLAVOR from zenml.integrations.modal.flavors.modal_base_flavor import ( DEFAULT_TIMEOUT_SECONDS, @@ -37,6 +39,12 @@ class ModalStepOperatorSettings(ModalSettingsMixin): """Settings for the Modal step operator.""" + stop_app: bool = Field( + False, + description="Stop the Modal app after the sandbox reaches a terminal " + "state.", + ) + class ModalStepOperatorConfig( BaseStepOperatorConfig, ModalCredentialsMixin, ModalStepOperatorSettings diff --git a/src/zenml/integrations/modal/sandbox_utils.py b/src/zenml/integrations/modal/sandbox_utils.py index ee3a1a78666..6e95d6559b7 100644 --- a/src/zenml/integrations/modal/sandbox_utils.py +++ b/src/zenml/integrations/modal/sandbox_utils.py @@ -15,6 +15,7 @@ import math import os +import subprocess import time from typing import Any, Dict, List, Optional, Protocol, Tuple @@ -210,6 +211,19 @@ def lookup_modal_app( ) +def get_modal_app_name(step_run_id: Any, step_name: str) -> str: + """Generate the Modal app name for a step run. + + Args: + step_run_id: ID of the step run. + step_name: Name of the pipeline step. + + Returns: + Modal app name string truncated to 64 chars. + """ + return f"zenml-{step_run_id}-{step_name}"[:64] + + def get_gpu_values( settings: ModalSandboxSettings, resource_settings: ResourceSettings, @@ -410,15 +424,6 @@ def stop_modal_app( ) -> None: """Stop a deployed Modal app by name using the Modal CLI. - Moves the app from 'deployed' to 'stopped' in the Modal dashboard, - eliminating the visual clutter of idle ``Tasks: 0`` app entries. - Stopped apps remain visible in the 'recently stopped' section with - their logs intact, so debugging is not affected. - - This should only be called after the associated Modal sandbox has - reached a terminal state, because stopping the app does not wait for - running sandboxes to finish. - Args: app_name: The Modal app name to stop (e.g. ``zenml--``). modal_environment: Modal environment name passed as ``--env``. @@ -431,10 +436,8 @@ def stop_modal_app( Raises: RuntimeError: If the ``modal app stop`` command exits with a - non-zero status code. + non-zero status code or times out. """ - import subprocess - cmd = ["modal", "app", "stop", "--yes", app_name] if modal_environment: cmd += ["--env", modal_environment] @@ -447,8 +450,29 @@ def stop_modal_app( env[MODAL_TOKEN_ID_ENV_KEY] = token_pair[0] env[MODAL_TOKEN_SECRET_ENV_KEY] = token_pair[1] - result = subprocess.run(cmd, capture_output=True, text=True, env=env) + try: + result = subprocess.run( + cmd, capture_output=True, text=True, env=env, timeout=30 + ) + if result.returncode != 0 and "--yes" in result.stderr: + fallback_cmd = ["modal", "app", "stop", app_name] + if modal_environment: + fallback_cmd += ["--env", modal_environment] + result = subprocess.run( + fallback_cmd, + input="y\n", + capture_output=True, + text=True, + env=env, + timeout=30, + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"Timed out stopping Modal app '{app_name}' after 30 seconds." + ) from e + if result.returncode != 0: + error_msg = result.stderr.strip() or result.stdout.strip() raise RuntimeError( - f"Failed to stop Modal app '{app_name}': {result.stderr.strip()}" + f"Failed to stop Modal app '{app_name}': {error_msg}" ) diff --git a/src/zenml/integrations/modal/step_operators/modal_step_operator.py b/src/zenml/integrations/modal/step_operators/modal_step_operator.py index c66cfbc9117..28d1e4f66a6 100644 --- a/src/zenml/integrations/modal/step_operators/modal_step_operator.py +++ b/src/zenml/integrations/modal/step_operators/modal_step_operator.py @@ -209,8 +209,11 @@ def submit( ) modal_client = self._get_modal_client() + app_name = sandbox_utils.get_modal_app_name( + info.step_run_id, info.pipeline_step_name + ) app = sandbox_utils.lookup_modal_app( - f"zenml-{info.step_run_id}-{info.pipeline_step_name}"[:64], + app_name, modal_environment=modal_environment, modal_client=modal_client, ) @@ -271,24 +274,26 @@ def cancel(self, step_run: "StepRunResponse") -> None: def cleanup_step_submission(self, step_run: "StepRunResponse") -> None: """Stop the Modal app after a sandbox run completes. - When ``stop_app_after_run`` is enabled in the step settings, this - calls ``modal app stop`` on the per-step app, moving its entry from - 'deployed' to 'stopped' in the Modal dashboard. Stopped apps remain - visible in the 'recently stopped' section with their full log history - intact, so debugging is unaffected. Failures are logged as warnings - and never re-raised so that cleanup cannot mask the step result. - Args: step_run: The finished step run. """ settings = cast(ModalStepOperatorSettings, self.get_settings(step_run)) - if not settings.stop_app_after_run: + if not settings.stop_app: return - app_name = f"zenml-{step_run.id}-{step_run.name}"[:64] - modal_environment = sandbox_utils.normalize_optional_config_value( - settings.modal_environment + app_name = sandbox_utils.get_modal_app_name( + step_run.id, step_run.name ) + modal_environment = step_run.run_metadata.get( + STEP_MODAL_ENVIRONMENT_METADATA_KEY + ) + if modal_environment is not None: + modal_environment = str(modal_environment) + else: + modal_environment = sandbox_utils.normalize_optional_config_value( + settings.modal_environment + ) + try: sandbox_utils.stop_modal_app( app_name, diff --git a/tests/integration/integrations/modal/step_operators/test_modal_step_operator.py b/tests/integration/integrations/modal/step_operators/test_modal_step_operator.py index f8963c8c9f5..91a9d3c84d2 100644 --- a/tests/integration/integrations/modal/step_operators/test_modal_step_operator.py +++ b/tests/integration/integrations/modal/step_operators/test_modal_step_operator.py @@ -797,3 +797,87 @@ def test_gpu_arg_whitespace_type_treated_as_none_behavior() -> None: rs_none = ResourceSettingsStub(gpu_count=None) assert get_gpu_values(settings, rs_none) is None + + +def test_get_modal_app_name() -> None: + app_name = sandbox_utils.get_modal_app_name("run-123", "train_model_step") + assert app_name == "zenml-run-123-train_model_step" + long_name = sandbox_utils.get_modal_app_name("run-123", "a" * 100) + assert len(long_name) == 64 + + +def test_cleanup_step_submission_disabled_by_default(monkeypatch) -> None: + recorded = [] + + def mock_stop_modal_app(*args, **kwargs): + recorded.append((args, kwargs)) + + monkeypatch.setattr(sandbox_utils, "stop_modal_app", mock_stop_modal_app) + + operator = _make_operator(ModalStepOperatorConfig()) + operator.get_settings = lambda _sr: ModalStepOperatorSettings( + stop_app=False + ) + + step_run = SimpleNamespace(id="step-id", name="step-name", run_metadata={}) + operator.cleanup_step_submission(step_run) + + assert len(recorded) == 0 + + +def test_cleanup_step_submission_enabled(monkeypatch) -> None: + recorded = [] + + def mock_stop_modal_app(*args, **kwargs): + recorded.append((args, kwargs)) + + monkeypatch.setattr(sandbox_utils, "stop_modal_app", mock_stop_modal_app) + + config = ModalStepOperatorConfig( + token_id="ak-test", token_secret="as-test" + ) + operator = _make_operator(config) + operator.get_settings = lambda _sr: ModalStepOperatorSettings( + stop_app=True + ) + + step_run = SimpleNamespace( + id="step-id", + name="step-name", + run_metadata={ + modal_step_operator_module.STEP_MODAL_ENVIRONMENT_METADATA_KEY: "staging" + }, + ) + operator.cleanup_step_submission(step_run) + + assert len(recorded) == 1 + assert recorded[0] == ( + ("zenml-step-id-step-name",), + { + "modal_environment": "staging", + "token_id": "ak-test", + "token_secret": "as-test", + }, + ) + + +def test_cleanup_step_submission_logs_warning_on_failure( + monkeypatch, caplog +) -> None: + def mock_stop_modal_app(*args, **kwargs): + raise RuntimeError("CLI failed") + + monkeypatch.setattr(sandbox_utils, "stop_modal_app", mock_stop_modal_app) + + operator = _make_operator(ModalStepOperatorConfig()) + operator.get_settings = lambda _sr: ModalStepOperatorSettings( + stop_app=True + ) + + step_run = SimpleNamespace(id="step-id", name="step-name", run_metadata={}) + + with caplog.at_level("WARNING"): + operator.cleanup_step_submission(step_run) + + assert "Failed to stop Modal app" in caplog.text +