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_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 708b7d651fe..6e95d6559b7 100644 --- a/src/zenml/integrations/modal/sandbox_utils.py +++ b/src/zenml/integrations/modal/sandbox_utils.py @@ -14,6 +14,8 @@ """Shared Modal Sandbox helpers.""" import math +import os +import subprocess import time from typing import Any, Dict, List, Optional, Protocol, Tuple @@ -209,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, @@ -398,3 +413,66 @@ 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. + + 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 or times out. + """ + 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] + + 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}': {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 402abe979eb..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, ) @@ -267,3 +270,47 @@ 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. + + Args: + step_run: The finished step run. + """ + settings = cast(ModalStepOperatorSettings, self.get_settings(step_run)) + if not settings.stop_app: + return + + 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, + 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, + ) 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 +