Add Modal orchestrator with step operator and orchestrator flavors #3733
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the 📝 WalkthroughWalkthroughThis update introduces a new Modal orchestrator integration to ZenML, enabling orchestration of entire pipelines or individual steps on Modal's serverless cloud. It adds new configuration, settings, and utility modules, updates the Modal integration to support both orchestrator and step operator flavors, enhances documentation, and refactors related code for modularity and clarity. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ZenML
participant ModalOrchestrator
participant ModalSandboxExecutor
participant Modal Cloud
User->>ZenML: Trigger pipeline run
ZenML->>ModalOrchestrator: Prepare deployment
ModalOrchestrator->>ModalSandboxExecutor: Setup execution (pipeline/step mode)
ModalSandboxExecutor->>Modal Cloud: Launch sandbox(es) with resources and image
Modal Cloud-->>ModalSandboxExecutor: Run pipeline/steps and stream logs
ModalSandboxExecutor->>ModalOrchestrator: Report completion or errors
ModalOrchestrator->>ZenML: Finalize run status
ZenML-->>User: Pipeline run results
Suggested reviewers
Poem
✨ Finishing Touches🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
github-actions
bot
added
internal
- Adds new Modal orchestrator flavor for serverless pipeline execution - Implements optimized execution modes: pipeline (default) and per_step - Supports GPU/CPU resource configuration with intelligent defaults - Features persistent apps with warm containers for fast execution - Includes comprehensive documentation and examples - Simplifies execution model by removing redundant single_function mode 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Documentation Link Check Results❌ Absolute links check failed |
|
✅ Branch tenant has been deployed! Access it at: https://staging.cloud.zenml.io/workspaces/feature-modal-orchestrator/projects |
| elif ( | ||
| log_age < 300 | ||
| ): # Only show logs from last 5 minutes | ||
| # This log is recent enough to likely be ours | ||
| logger.info(f"{log_msg}") | ||
| # Else: skip old logs |
There was a problem hiding this comment.
This bit seems maybe like it could use more attention. Also are you sure that we wouldn't have time zone mismatches here etc?
| else: | ||
| # Fallback to first step's resource settings if no pipeline-level resources | ||
| if deployment.step_configurations: | ||
| first_step = list(deployment.step_configurations.values())[0] |
There was a problem hiding this comment.
what is the first step is unrepresentative, though?
There was a problem hiding this comment.
We should check / just make this really explicit in the docs that this is how we assume this
| # Register the orchestrator with explicit credentials | ||
| zenml orchestrator register <ORCHESTRATOR_NAME> \ | ||
| --flavor=modal \ | ||
| --token=<MODAL_TOKEN> \ | ||
| --workspace=<MODAL_WORKSPACE> \ | ||
| --synchronous=true |
There was a problem hiding this comment.
I think this should use --token-id and --token-secret separately as per the code?
| ### Authentication with different environments | ||
|
|
||
| For production deployments, you can specify different Modal environments: |
There was a problem hiding this comment.
Maybe could have a little info box in this section (or maybe even above, linking down here) to say that you might want to have two different stacks, each associated with a different modal environment, one for prod and the other for development etc etc.
| log_stream_active.set() | ||
| start_time = time.time() | ||
|
|
||
| def stream_logs() -> None: |
There was a problem hiding this comment.
Function in function smells a bit wrong and also wondering if we should instead use their Python SDK to stream the logs? https://github.com/modal-labs/modal-client/blob/4177d0b994ac69e01ada7d7a96655c9dcaae570e/modal/cli/utils.py#L24
Possibly something for down the line, though the func-in-func seems off.
| if TYPE_CHECKING: | ||
| from zenml.integrations.modal.flavors.modal_orchestrator_flavor import ( | ||
| ModalOrchestratorConfig, | ||
| ModalOrchestratorSettings, | ||
| ) | ||
|
|
||
| from zenml.integrations.modal.flavors.modal_orchestrator_flavor import ( | ||
| ModalExecutionMode, | ||
| ) | ||
|
|
||
| if TYPE_CHECKING: | ||
| from zenml.models import PipelineDeploymentResponse, PipelineRunResponse | ||
| from zenml.models.v2.core.pipeline_deployment import PipelineDeploymentBase |
There was a problem hiding this comment.
combine the 'if TYPE_CHECKING' parts?
|
|
||
| The Modal orchestrator supports two execution modes: | ||
|
|
||
| 1. **`pipeline` (default)**: Runs the entire pipeline in a single Modal function for maximum speed and cost efficiency |
There was a problem hiding this comment.
Not sure I understand why this pipeline option is max speed. Isn't it running everything sequentially in the same container? Wouldn't running things in parallel in separate Modal function calls run faster?
src/zenml/integrations/modal/flavors/modal_orchestrator_flavor.py
Outdated
Show resolved
Hide resolved
| Using the ZenML `modal` integration, you can orchestrate and scale your ML pipelines on [Modal's](https://modal.com/) serverless cloud platform with minimal setup and maximum efficiency. | ||
|
|
||
| The Modal orchestrator is designed for speed and cost-effectiveness, running entire pipelines in single serverless functions to minimize cold starts and optimize resource utilization. | ||
|
|
There was a problem hiding this comment.
Maybe some representative screenshot of the Modal UI in here to make the docs a bit friendlier?
There was a problem hiding this comment.
i think its fine without
- Extract nested log streaming function into ModalLogStreamer class for better code organization - Remove unreliable timezone-based log filtering that could miss logs due to clock skew - Implement smarter resource fallback: use highest requirements across all steps instead of potentially unrepresentative first step - Add logging for resource selection decisions to improve debugging - Fix function-in-function code smell identified in PR review 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Combine duplicate TYPE_CHECKING blocks into single import section - Improve import organization and reduce redundancy - Maintain all existing functionality while improving code structure 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
- Import MODAL_ORCHESTRATOR_FLAVOR constant from central location to avoid duplication - Update requirements to modal>=1 after testing compatibility with both orchestrator and step operator - Remove unnecessary utils import that was only for mypy discovery - Maintain consistent import patterns across Modal integration files 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
Based on PR review feedback: - Fix token authentication examples to use --token-id and --token-secret - Add "When NOT to use it" section with clear tradeoffs and alternatives - Add info boxes for environment separation best practices and cost implications - Document Modal vs Step Operator differences with usage recommendations - Add GPU base image requirements and CUDA compatibility warnings - Clarify execution modes: "pipeline" mode reduces overhead vs enables parallelism - Document resource fallback behavior and warming window defaults - Add container warming cost implications with specific guidance - Remove tracking pixel per review request - Improve overall documentation clarity and completeness 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (2)
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py (1)
159-159: Setting pipeline run ID in environment works correctly.Yes, this works. Setting the pipeline run ID as an environment variable is a standard pattern for passing runtime information to the execution environment.
src/zenml/integrations/modal/utils.py (1)
320-321: Remove unused attribute assignment.Based on the codebase context, setting this attribute on the app object appears to be unused. The caching is already handled via Modal's Dict.
- # Store the image in the app for future use - setattr(app, image_name_key, zenml_image)
🧹 Nitpick comments (4)
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py (1)
144-149: Consider supporting scheduled pipelines in the future.The warning about schedules not being supported is helpful for users. Consider adding this as a feature request for future iterations.
Would you like me to create an issue to track adding schedule support for the Modal orchestrator?
src/zenml/integrations/modal/orchestrators/modal_sandbox_executor.py (1)
525-527: Consider using debug level for expected scenarios.The warning "Image not hydrated after sandbox creation" might occur in normal scenarios. Consider using
logger.debuginstead oflogger.warningfor this case.else: - logger.warning("Image not hydrated after sandbox creation") + logger.debug("Image not hydrated after sandbox creation")docs/book/component-guide/step-operators/modal.md (1)
30-48: Fix unordered list style for consistency.Use dashes instead of asterisks for unordered lists to maintain consistency with the documentation style guide.
-* You want to run only specific compute-intensive steps (like training or inference) on Modal while keeping other steps local. -* You need different hardware requirements for different steps in your pipeline. -* You want to leverage Modal's fast execution and GPU access for select steps without moving your entire pipeline to the cloud. -* You have a hybrid workflow where some steps need to access local resources or data. +- You want to run only specific compute-intensive steps (like training or inference) on Modal while keeping other steps local. +- You need different hardware requirements for different steps in your pipeline. +- You want to leverage Modal's fast execution and GPU access for select steps without moving your entire pipeline to the cloud. +- You have a hybrid workflow where some steps need to access local resources or data. ### When NOT to use it The Modal step operator may not be the best choice if: -* **You want to run entire pipelines on Modal**: Use the [Modal orchestrator](../orchestrators/modal.md) instead for complete pipeline execution with better cost efficiency and reduced overhead. +- **You want to run entire pipelines on Modal**: Use the [Modal orchestrator](../orchestrators/modal.md) instead for complete pipeline execution with better cost efficiency and reduced overhead. -* **You have simple, lightweight steps**: For steps that don't require significant compute resources, the overhead of running them on Modal may not be worth it. +- **You have simple, lightweight steps**: For steps that don't require significant compute resources, the overhead of running them on Modal may not be worth it. -* **You need very low latency**: The step operator introduces some overhead for individual step execution compared to running steps locally. +- **You need very low latency**: The step operator introduces some overhead for individual step execution compared to running steps locally. -* **You have tight data locality requirements**: If your steps need to access large amounts of local data, transferring it to Modal for each step execution may be inefficient. +- **You have tight data locality requirements**: If your steps need to access large amounts of local data, transferring it to Modal for each step execution may be inefficient.docs/book/component-guide/orchestrators/modal.md (1)
583-619: Valuable cost optimization guidance.The detailed cost examples and optimization strategies provide essential information for users to make informed decisions about execution modes and resource allocation. Consider adding a link to Modal's pricing page for the most up-to-date pricing information.
Add a link to Modal's pricing page:
**Cost Examples (approximate)**: +For current pricing, see [Modal's pricing page](https://modal.com/pricing). + ```python
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
docs/book/component-guide/orchestrators/README.md(1 hunks)docs/book/component-guide/orchestrators/modal.md(1 hunks)docs/book/component-guide/step-operators/modal.md(7 hunks)docs/book/component-guide/toc.md(1 hunks)src/zenml/integrations/modal/__init__.py(2 hunks)src/zenml/integrations/modal/flavors/__init__.py(1 hunks)src/zenml/integrations/modal/flavors/modal_orchestrator_flavor.py(1 hunks)src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py(2 hunks)src/zenml/integrations/modal/orchestrators/__init__.py(1 hunks)src/zenml/integrations/modal/orchestrators/modal_orchestrator.py(1 hunks)src/zenml/integrations/modal/orchestrators/modal_orchestrator_entrypoint.py(1 hunks)src/zenml/integrations/modal/orchestrators/modal_orchestrator_entrypoint_configuration.py(1 hunks)src/zenml/integrations/modal/orchestrators/modal_sandbox_executor.py(1 hunks)src/zenml/integrations/modal/step_operators/modal_step_operator.py(4 hunks)src/zenml/integrations/modal/utils.py(1 hunks)tests/integration/integrations/modal/step_operators/test_modal_step_operator.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
docs/**/*.md
Instructions used from:
Sources:
⚙️ CodeRabbit Configuration File
tests/**/*.py
Instructions used from:
Sources:
⚙️ CodeRabbit Configuration File
src/zenml/**/*.py
Instructions used from:
Sources:
⚙️ CodeRabbit Configuration File
🧠 Learnings (8)
docs/book/component-guide/toc.md (2)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/component-guide/**/* : Add documentation in /docs/book/component-guide/ when adding new integrations
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Update documentation for user-facing changes (or ensure that nothing was broken)
src/zenml/integrations/modal/orchestrators/__init__.py (2)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/integrations/**/* : Create integration package in /src/zenml/integrations/ when adding new integrations
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/**/*.py : Document public APIs thoroughly when implementing features
src/zenml/integrations/modal/flavors/__init__.py (2)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/integrations/**/* : Create integration package in /src/zenml/integrations/ when adding new integrations
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/**/*.py : Document public APIs thoroughly when implementing features
src/zenml/integrations/modal/__init__.py (2)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/**/*.py : Document public APIs thoroughly when implementing features
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/integrations/**/* : Create integration package in /src/zenml/integrations/ when adding new integrations
src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py (1)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Use environment variables or ZenML's secret management for sensitive data
docs/book/component-guide/step-operators/modal.md (2)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Update documentation for user-facing changes (or ensure that nothing was broken)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Use environment variables or ZenML's secret management for sensitive data
docs/book/component-guide/orchestrators/modal.md (7)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/component-guide/**/* : Add documentation in /docs/book/component-guide/ when adding new integrations
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/**/*.py : Document public APIs thoroughly when implementing features
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Update documentation for user-facing changes (or ensure that nothing was broken)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Explanation of concepts in documentation
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/toc.md : Multiple sections each have their own toc.md file for navigation
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Usage patterns and best practices in documentation
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Include metadata fields at the top of documentation pages (follow existing patterns)
src/zenml/integrations/modal/utils.py (6)
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/**/*.py : Document public APIs thoroughly when implementing features
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to src/zenml/integrations/**/* : Create integration package in /src/zenml/integrations/ when adding new integrations
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Clear, concise language in documentation
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Documentation should be readable and conversational
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Explanation of concepts in documentation
Learnt from: CR
PR: zenml-io/zenml#0
File: CLAUDE.md:0-0
Timestamp: 2025-07-01T17:52:34.101Z
Learning: Applies to docs/book/**/* : Match tone and style with existing documentation
🧬 Code Graph Analysis (5)
src/zenml/integrations/modal/orchestrators/__init__.py (1)
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py (1)
ModalOrchestrator(55-207)
src/zenml/integrations/modal/orchestrators/modal_orchestrator_entrypoint_configuration.py (2)
src/zenml/config/docker_settings.py (1)
command(36-48)src/zenml/models/v2/core/pipeline_run.py (1)
orchestrator_run_id(491-497)
src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py (2)
src/zenml/utils/secret_utils.py (1)
SecretField(91-105)src/zenml/step_operators/base_step_operator.py (1)
BaseStepOperatorConfig(33-34)
src/zenml/integrations/modal/step_operators/modal_step_operator.py (1)
src/zenml/integrations/modal/utils.py (4)
build_modal_image(325-384)create_modal_stack_validator(422-461)get_gpu_values(149-178)setup_modal_client(62-146)
src/zenml/integrations/modal/flavors/modal_orchestrator_flavor.py (4)
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py (2)
config(63-69)ModalOrchestrator(55-207)src/zenml/orchestrators/base_orchestrator.py (2)
BaseOrchestratorConfig(78-139)BaseOrchestratorFlavor(526-554)src/zenml/utils/secret_utils.py (1)
SecretField(91-105)src/zenml/stack/flavor.py (2)
generate_default_docs_url(206-230)generate_default_sdk_docs_url(232-264)
🪛 LanguageTool
docs/book/component-guide/step-operators/modal.md
[style] ~33-~33: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...odal while keeping other steps local. * You need different hardware requirements fo...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~34-~34: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...for different steps in your pipeline. * You want to leverage Modal's fast execution...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~35-~35: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ng your entire pipeline to the cloud. * You have a hybrid workflow where some steps...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~45-~45: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...them on Modal may not be worth it. * You need very low latency: The step opera...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~47-~47: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...compared to running steps locally. * You have tight data locality requirements...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
docs/book/component-guide/orchestrators/modal.md
[style] ~21-~21: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...ion with minimal cold start overhead. * you want cost-effective ML pipeline orchest...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~22-~22: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...tion without managing infrastructure. * you need easy access to GPUs and high-perfo...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~23-~23: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...high-performance computing resources. * you prefer a simple setup process without c...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~35-~35: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...er suited for other orchestrators. * You need complex workflow patterns: Modal...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[style] ~114-~114: Consider using an alternative to strengthen your wording.
Context: ...ml.io/concepts/containerization) if you want to learn more about how ZenML builds th...
(WANT_KEEN)
[grammar] ~236-~236: Ensure spelling is correct
Context: ...figuration! You can mix different GPUs, CPU, and memory settings across steps. Pipe...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
🪛 markdownlint-cli2 (0.17.2)
docs/book/component-guide/step-operators/modal.md
32-32: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
33-33: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
34-34: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
35-35: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
41-41: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
43-43: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
45-45: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
47-47: Unordered list style
Expected: dash; Actual: asterisk
(MD004, ul-style)
96-96: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
104-104: Emphasis used instead of a heading
(MD036, no-emphasis-as-heading)
🔇 Additional comments (36)
docs/book/component-guide/toc.md (1)
22-22: Documentation entry added correctly.The Modal Orchestrator entry follows the established pattern and is properly placed in the orchestrators section.
docs/book/component-guide/orchestrators/README.md (1)
37-37: Table entry added correctly.The Modal orchestrator entry follows the established table format and provides accurate flavor, integration, and description information.
tests/integration/integrations/modal/step_operators/test_modal_step_operator.py (1)
45-45: Test updated correctly for refactored function signature.The change correctly updates the function call to pass
settings.gpudirectly instead of the entire settings object, aligning with the refactoredget_gpu_valuesutility function.src/zenml/integrations/modal/flavors/__init__.py (2)
16-20: Orchestrator flavor imports added correctly.The imports follow the established pattern for Modal integration flavors and properly expose the orchestrator configuration, flavor, and settings classes.
28-30: Public API exports updated appropriately.The new orchestrator classes are correctly added to the
__all__list, maintaining consistency with the existing step operator exports.src/zenml/integrations/modal/orchestrators/__init__.py (1)
1-20: New orchestrator package properly structured.The init file correctly serves as the entry point for the Modal orchestrator package, following ZenML conventions and Python best practices with proper imports and exports.
src/zenml/integrations/modal/__init__.py (4)
14-18: LGTM!The docstring update accurately reflects the addition of orchestrator functionality alongside the existing step operator.
25-26: LGTM!The new orchestrator flavor constant follows the same naming pattern as the step operator flavor.
42-48: LGTM!The flavors method correctly imports and returns both the orchestrator and step operator flavors, properly registering them with the integration.
33-33: Verify Modal v1.0+ Breaking ChangesWe’ve upgraded the Modal requirement from
“modal>=0.64.49,<1”to“modal>=1”, which introduces several breaking changes. Before merging, please audit your integration code (step operator, orchestrator, anymodal.Imageusage) for:
Automounting removal
• v1.0 no longer auto‐include local Python packages.
• Add explicit calls tomodal.Image.add_local_python_source(...)for any local modules you rely on.Removed
automountconfig/ENV
• Remove anyautomountflags orMODAL_AUTOMOUNTreferences.Deprecated copy methods enforced
• ReplaceImage.copy_local_dir/copy_local_filewithImage.add_local_dir/add_local_file(usecopy=Trueif you need copy semantics).@modal.cls initialization changes
• Migrate any explicit__init__in@modal.clsclasses to usemodal.parameter()and move setup logic into@modal.enter()methods.Migration guide: https://modal.com/docs/guide/modal-1-0-migration
If you haven’t already addressed these deprecations, consider pinning tomodal~=1.0.0first to avoid unexpected breakages.src/zenml/integrations/modal/orchestrators/modal_orchestrator_entrypoint_configuration.py (1)
26-85: Well-structured entrypoint configuration!The
ModalOrchestratorEntrypointConfigurationclass provides a clean interface for command-line argument handling with proper separation of concerns:
- Clear distinction between required and optional options
- Proper command construction using Python module execution
- Type hints with forward references for better code clarity
src/zenml/integrations/modal/flavors/modal_step_operator_flavor.py (1)
21-21: Excellent security and configuration enhancements!The updates properly implement:
- Secure credential handling using
SecretFieldfor sensitive data (following best practices)- Clear documentation for GPU usage and authentication fallback
- Sensible timeout default matching Modal's 24-hour maximum
- Consistent configuration pattern with the new orchestrator flavor
Also applies to: 40-52, 58-73
src/zenml/integrations/modal/flavors/modal_orchestrator_flavor.py (4)
29-40: Well-designed execution modes!The
ModalExecutionModeenum clearly defines two execution strategies with appropriate documentation explaining the trade-offs between speed (PIPELINE mode) and granular control (PER_STEP mode).
42-71: Comprehensive orchestrator settings!The settings class provides all necessary configuration options:
- Resource configuration (GPU, region, cloud)
- Execution control (mode, parallelism, synchronous)
- Sensible defaults (24h timeout, pipeline mode, synchronous execution)
73-110: Secure and well-structured configuration!The configuration class properly:
- Uses
SecretFieldfor sensitive credentials- Correctly implements
is_synchronousproperty (not hardcoded)- Documents authentication fallback behavior
- Maintains consistency with other orchestrator implementations
112-170: Complete flavor implementation!The orchestrator flavor correctly implements all required properties and follows ZenML patterns for documentation URLs, logo, and lazy loading of the implementation class.
src/zenml/integrations/modal/orchestrators/modal_orchestrator_entrypoint.py (4)
92-178: Excellent optimization with image pre-building!The
prepare_shared_image_cachefunction implements a smart optimization by:
- Pre-building all unique images before step execution to avoid redundant builds
- Creating a shared Modal app for resource efficiency
- Properly handling the case where deployment.build is None
- Generating deterministic cache keys using build ID and image hash
This will significantly improve performance for pipelines with multiple steps.
194-264: Well-structured per-step execution with proper parallelism control!The implementation correctly:
- Leverages the pre-built image cache for efficiency
- Creates a shared executor to avoid redundant setup
- Uses
ThreadedDagRunnerwith configurablemax_parallelism- Maintains proper step dependencies through the DAG
The past review comment about settings persistence across steps appears to have been properly addressed.
266-323: Robust finalization with proper error handling!The
finalize_runfunction demonstrates good practices:
- Updates both step and pipeline run statuses based on execution results
- Only marks runs as failed if they're still in running/initializing state
- Gracefully handles
AuthorizationExceptionwhich may occur after token expirationThis ensures accurate status reporting even in edge cases.
325-383: Clean orchestration with proper execution mode dispatch!The main function effectively:
- Sets up the Modal client with proper authentication
- Retrieves deployment and settings from the active stack
- Dispatches to the appropriate execution mode (pipeline vs per-step)
- Provides comprehensive error logging throughout
The architecture allows for easy extension of execution modes in the future.
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py (2)
84-102: Good use of shared utilities for authentication and validation.The refactoring to use
setup_modal_clientandcreate_modal_stack_validatorfrom the utils module promotes code reuse and consistency across the Modal integration.
103-120: Proper error handling for orchestrator run ID retrieval.The method correctly handles the missing environment variable case with a clear error message.
src/zenml/integrations/modal/orchestrators/modal_sandbox_executor.py (4)
113-166: Excellent robust implementation for resource settings conversion.This method demonstrates excellent defensive programming by gracefully handling various input formats (ResourceSettings instance, pydantic models, dicts, or even objects with
__dict__). The fallback logic and logging ensure reliability even with historical or malformed data.
301-312: Smart default GPU type handling.Good UX decision to default to "T4" when
gpu_count > 0but no GPU type is specified. This prevents confusing errors while still allowing users to override with specific GPU types.
325-410: Well-structured parameter validation for Modal API.The method properly validates required parameters and conditionally includes optional ones, preventing invalid API calls. The parameter summary logging (excluding complex objects) is helpful for debugging.
624-689: Clean implementation of execution methods.The separation between
execute_pipelineandexecute_stepis well-designed, with each method building the appropriate entrypoint commands for their respective execution modes.docs/book/component-guide/step-operators/modal.md (3)
11-29: Excellent comparison between Modal components.The comparison table and decision guide provide clear guidance for users to choose between the step operator and orchestrator based on their needs. This will significantly improve the user experience.
58-77: Great addition of a quick start guide.The 5-minute quick start provides an excellent onboarding experience for new users with clear, actionable steps.
204-240: Valuable best practices for environment separation.The recommendation to use separate step operators for different environments (dev, staging, prod) is an important security and operational best practice that will help users avoid configuration mistakes.
src/zenml/integrations/modal/step_operators/modal_step_operator.py (2)
28-33: Excellent refactoring to use shared utilities.The migration to use
setup_modal_client,create_modal_stack_validator,build_modal_image, andget_gpu_valuesfrom the utils module improves code maintainability and ensures consistency between the step operator and orchestrator implementations.Also applies to: 80-81, 125-130, 132-136
162-162: Good improvement: configurable timeout.Making the timeout configurable via
settings.timeoutinstead of hardcoding it provides better flexibility for users with different execution time requirements.docs/book/component-guide/orchestrators/modal.md (3)
11-14: Important warning about deployment requirements.The warning about requiring a remote ZenML deployment is crucial for setting proper expectations and preventing user frustration with local deployments.
279-374: Outstanding documentation of execution modes.The detailed explanation of pipeline vs per-step execution modes with concrete examples is excellent. The per-step resource configuration example (lines 320-367) clearly demonstrates the power and flexibility of this feature for cost optimization.
529-559: Clear architectural explanation.The "Persistent Apps + Dynamic Sandboxes" architecture explanation effectively communicates the design decisions and benefits, helping users understand the performance characteristics and isolation guarantees.
src/zenml/integrations/modal/utils.py (2)
181-232: Well-implemented resource validation.The function provides comprehensive validation with helpful warnings for edge cases. Good defensive programming with proper error handling.
422-462: Well-structured validation logic.The stack validator properly ensures remote components are configured with clear, actionable error messages.
| def build_modal_image( | ||
| image_name: str, | ||
| stack: "Stack", | ||
| environment: Optional[Dict[str, str]] = None, | ||
| ) -> Any: | ||
| """Build a Modal image from a ZenML-built Docker image. | ||
|
|
||
| Args: | ||
| image_name: The name of the Docker image to use as base. | ||
| stack: The ZenML stack containing container registry. | ||
| environment: The environment variables to pass to the image. | ||
|
|
||
| Returns: | ||
| The configured Modal image. | ||
|
|
||
| Raises: | ||
| RuntimeError: If no Docker credentials are found. | ||
| ValueError: If no container registry is found. | ||
| """ | ||
| if not stack.container_registry: | ||
| raise ValueError( | ||
| "No Container registry found in the stack. " | ||
| "Please add a container registry and ensure " | ||
| "it is correctly configured." | ||
| ) | ||
|
|
||
| logger.info("Building Modal image") | ||
| logger.info(f"Base image: {image_name}") | ||
|
|
||
| if docker_creds := stack.container_registry.credentials: | ||
| docker_username, docker_password = docker_creds | ||
| else: | ||
| raise RuntimeError( | ||
| "No Docker credentials found for the container registry." | ||
| ) | ||
|
|
||
| # Create Modal secret for registry authentication | ||
| registry_secret = modal.Secret.from_dict( | ||
| { | ||
| "REGISTRY_USERNAME": docker_username, | ||
| "REGISTRY_PASSWORD": docker_password, | ||
| } | ||
| ) | ||
|
|
||
| # Build new Modal image and register it with consistent name | ||
| logger.info(f"🔨 Building Modal image from base: {image_name}") | ||
|
|
||
| # Build Modal image from the ZenML-built image | ||
| # Modal will automatically cache layers and reuse when possible | ||
| logger.info(f"Creating Modal image from base: {image_name}") | ||
| zenml_image = ( | ||
| modal.Image.from_registry( | ||
| image_name, secret=registry_secret | ||
| ).pip_install("modal") # Install Modal in the container | ||
| ) | ||
|
|
||
| if environment: | ||
| zenml_image = zenml_image.env(environment) | ||
|
|
||
| return zenml_image | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Refactor to eliminate code duplication.
This function duplicates most of the logic from get_or_build_modal_image. Consider extracting the common image building logic into a shared helper function.
Create a helper function:
def _build_modal_image_from_registry(
image_name: str,
stack: "Stack",
environment: Optional[Dict[str, str]] = None,
) -> Any:
"""Build Modal image from registry with authentication."""
if not stack.container_registry:
raise ValueError(
"No Container registry found in the stack. "
"Please add a container registry and ensure "
"it is correctly configured."
)
if docker_creds := stack.container_registry.credentials:
docker_username, docker_password = docker_creds
else:
raise RuntimeError(
"No Docker credentials found for the container registry."
)
registry_secret = modal.Secret.from_dict({
"REGISTRY_USERNAME": docker_username,
"REGISTRY_PASSWORD": docker_password,
})
logger.info(f"🔨 Building Modal image from base: {image_name}")
zenml_image = (
modal.Image.from_registry(image_name, secret=registry_secret)
.pip_install("modal")
)
if environment:
zenml_image = zenml_image.env(environment)
return zenml_imageThen simplify both functions to use this helper.
🤖 Prompt for AI Agents
In src/zenml/integrations/modal/utils.py between lines 325 and 385, the
build_modal_image function duplicates logic found in get_or_build_modal_image
related to building a Modal image from a Docker registry with authentication. To
fix this, extract the shared logic that checks for container registry and
credentials, creates the Modal secret, builds the image with
pip_install("modal"), and applies environment variables into a new helper
function named _build_modal_image_from_registry. Then refactor build_modal_image
and get_or_build_modal_image to call this helper, removing the duplicated code
and centralizing the image building process.
| if gpu_count is None or gpu_count == 0: | ||
| return gpu_type |
There was a problem hiding this comment.
Clarify behavior when gpu_count is 0.
When gpu_count is explicitly set to 0, the function still returns the GPU type. This could be confusing - users setting gpu_count=0 likely intend to disable GPU usage entirely.
Consider this fix:
- # GPU type specified but no count, default to 1
- if gpu_count is None or gpu_count == 0:
+ # No GPU requested if count is explicitly 0
+ if gpu_count == 0:
+ return None
+
+ # GPU type specified but no count, default to 1
+ if gpu_count is None:
return gpu_type📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if gpu_count is None or gpu_count == 0: | |
| return gpu_type | |
| # No GPU requested if count is explicitly 0 | |
| if gpu_count == 0: | |
| return None | |
| # GPU type specified but no count, default to 1 | |
| if gpu_count is None: | |
| return gpu_type |
🤖 Prompt for AI Agents
In src/zenml/integrations/modal/utils.py around lines 174 to 175, the function
returns the GPU type even when gpu_count is explicitly set to 0, which can
confuse users who intend to disable GPU usage. Modify the condition to return
None or an appropriate value indicating no GPU should be used when gpu_count is
0, ensuring the function behavior aligns with user intent to disable GPU usage.
| # If lookup fails, we'll build a new image | ||
| pass | ||
|
|
There was a problem hiding this comment.
Remove redundant exception handling.
This outer try/except block appears redundant and uses a bare except clause.
- except Exception:
- # If lookup fails, we'll build a new image
- pass📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # If lookup fails, we'll build a new image | |
| pass |
🤖 Prompt for AI Agents
In src/zenml/integrations/modal/utils.py around lines 286 to 288, remove the
outer try/except block that uses a bare except clause as it is redundant.
Instead, allow exceptions to propagate naturally or handle them more
specifically where necessary, eliminating the unnecessary pass statement inside
the except block.
| # Dict doesn't exist or image not found, will build new one | ||
| pass | ||
| except Exception: |
There was a problem hiding this comment.
Avoid bare except clauses.
The bare except clause could hide real errors. Be more specific about which exceptions you're catching.
- except Exception:
- # Dict doesn't exist or image not found, will build new one
- pass
+ except (modal.exceptions.NotFoundError, KeyError) as e:
+ # Dict doesn't exist or image not found, will build new one
+ logger.debug(f"Cache lookup failed, will build new image: {e}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Dict doesn't exist or image not found, will build new one | |
| pass | |
| except Exception: | |
| except (modal.exceptions.NotFoundError, KeyError) as e: | |
| # Dict doesn't exist or image not found, will build new one | |
| logger.debug(f"Cache lookup failed, will build new image: {e}") |
🤖 Prompt for AI Agents
In src/zenml/integrations/modal/utils.py around lines 283 to 285, replace the
bare except clause with a more specific exception type that you expect could
occur in this block. Identify the likely exceptions (e.g., KeyError,
FileNotFoundError) and catch only those to avoid hiding unexpected errors. This
will make error handling clearer and safer.
| max_parallelism: Maximum number of parallel sandboxes (for PER_STEP mode). | ||
| synchronous: Wait for completion (True) or fire-and-forget (False). | ||
| """ | ||
|
|
There was a problem hiding this comment.
I think having an option to specify the sandbox name would be nice?
| cloud: The cloud provider to use for the pipeline execution. | ||
| modal_environment: The Modal environment to use for the pipeline execution. | ||
| timeout: Maximum execution time in seconds (default 24h). | ||
| execution_mode: Execution mode - PIPELINE (fastest) or PER_STEP (granular). |
There was a problem hiding this comment.
Needs more explanation IMO, also the naming I'm not sure about.
Both modes execute your pipeline, but PIPELINE uses a single sandbox and runs sequentially, while PER_STEP spins up a separate sandbox per step.
There was a problem hiding this comment.
sandboxes dont have names and the app names are equal to the pipeline names
There was a problem hiding this comment.
I think this reply is on the wrong comment. What I meant was actually the app name. I would equate this with specifying a kubernetes namespace or pod name, both of which we allow users to specify but fallback to defaults like the pipeline name if no value is given.
src/zenml/integrations/modal/step_operators/modal_step_operator.py
Outdated
Show resolved
Hide resolved
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py
Outdated
Show resolved
Hide resolved
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py
Outdated
Show resolved
Hide resolved
| execution_mode = getattr( | ||
| settings, "execution_mode", ModalExecutionMode.PIPELINE | ||
| ) |
There was a problem hiding this comment.
| execution_mode = getattr( | |
| settings, "execution_mode", ModalExecutionMode.PIPELINE | |
| ) | |
| execution_mode = settings.execution_mode |
| logger.error(f"Pipeline execution failed: {e}") | ||
| raise | ||
|
|
||
| logger.info("✅ Pipeline execution completed successfully") |
There was a problem hiding this comment.
This is only true if the orchestrator is running in sync mode
| orchestrator_run_id = str(uuid4()) | ||
| environment[ENV_ZENML_MODAL_ORCHESTRATOR_RUN_ID] = orchestrator_run_id |
There was a problem hiding this comment.
There is no need to pass this to the entrypoint if it is simply a random UUID. In that case, we can simply generate it in the entrypoint itself.
| logger.info( | ||
| f"🚀 Executing pipeline with Modal ({execution_mode.lower()} mode)" | ||
| ) | ||
|
|
There was a problem hiding this comment.
If the execution mode is pipeline, we do not allow per-step images (as they cannot be used). I suggest you overwrite the get_docker_builds(...) method as follows:
- Call the super implementation
- Check if there is any build configuration in the result that has
key="orchestrator"and a non-Nonestep_name. In that case, raise an error.
…r.py Co-authored-by: Michael Schuster <[email protected]>
Co-authored-by: Michael Schuster <[email protected]>
…l into feature/modal-orchestrator
| ) | ||
|
|
||
| # Build Modal image | ||
| zenml_image = self._build_modal_image(deployment, stack, environment) |
There was a problem hiding this comment.
Comment: Dont build each time, use metadata of build id to track whether this build is already in modal or not
| logger.error(f"Pipeline execution failed: {e}") | ||
| raise | ||
|
|
||
| logger.info("✅ Pipeline execution completed successfully") |
| execution_mode = getattr( | ||
| settings, "execution_mode", ModalExecutionMode.PIPELINE | ||
| ) |
| logger.info( | ||
| f"🚀 Executing pipeline with Modal ({execution_mode.lower()} mode)" | ||
| ) | ||
|
|
| orchestrator_run_id = str(uuid4()) | ||
| environment[ENV_ZENML_MODAL_ORCHESTRATOR_RUN_ID] = orchestrator_run_id |
src/zenml/integrations/modal/orchestrators/modal_orchestrator.py
Outdated
Show resolved
Hide resolved
| cloud: The cloud provider to use for the pipeline execution. | ||
| modal_environment: The Modal environment to use for the pipeline execution. | ||
| timeout: Maximum execution time in seconds (default 24h). | ||
| execution_mode: Execution mode - PIPELINE (fastest) or PER_STEP (granular). |
There was a problem hiding this comment.
sandboxes dont have names and the app names are equal to the pipeline names
| raise ValueError( | ||
| f"Per-step Docker settings are not supported in PIPELINE " | ||
| f"execution mode. Step '{build.step_name}' has custom Docker " | ||
| f"settings but will be ignored since the entire pipeline runs " |
There was a problem hiding this comment.
It will not be ignored, but instead the run will fail right here with this exception
| ) | ||
|
|
||
| # Determine if we should wait for completion | ||
| synchronous = ( |
There was a problem hiding this comment.
synchronous = settings.synchronous
| @@ -231,7 +158,7 @@ async def run_sandbox() -> asyncio.Future[None]: | |||
| cloud=settings.cloud, | |||
| region=settings.region, | |||
| app=app, | |||
| timeout=86400, # 24h, the max Modal allows | |||
| timeout=settings.timeout, | |||
There was a problem hiding this comment.
Feel free to ignore as part of this PR as well, but seems like this could also be refactored to use the same code that runs a step of a pipeline in a new sandbox?
| token_id=orchestrator.config.token_id, | ||
| token_secret=orchestrator.config.token_secret, | ||
| workspace=orchestrator.config.workspace, | ||
| environment=orchestrator.config.modal_environment, |
There was a problem hiding this comment.
The modal_environment is not a config option and can therefore be overwritten using the settings. So I'm guessing we should probably fetch this from the pipeline settings instead?
| token_id=self.config.token_id, | ||
| token_secret=self.config.token_secret, | ||
| workspace=self.config.workspace, | ||
| environment=self.config.modal_environment, |
There was a problem hiding this comment.
Same as for the entrypoint, this should be fetched from the settings instead.
| The configured Modal image. | ||
| """ | ||
| # Try to get existing image from the app | ||
| image_name_key = f"zenml_image_{build_id}" |
There was a problem hiding this comment.
How is this even supposed to work? This key is the same for all images in the build, which are certainly not referring to the same docker image.
| logger.warning(f"Invalid memory {memory_mb}MB, ignoring.") | ||
| memory_mb = None | ||
| elif memory_mb < 128: # Less than 128MB seems too low | ||
| logger.warning( |
There was a problem hiding this comment.
Why would it not be intentional, the user explicitly specified this?
| ) | ||
|
|
||
| # Build Modal image from the ZenML-built image | ||
| logger.debug(f"Building Modal image from base: {image_name}") |
| zenml_image = ( | ||
| modal.Image.from_registry( | ||
| image_name, secret=registry_secret | ||
| ).pip_install("modal") # Install Modal in the container |
There was a problem hiding this comment.
modal is installed in the container, as part of the ZenML image build that happened for the stack.
| def _build_modal_image_from_registry( | ||
| image_name: str, | ||
| stack: "Stack", | ||
| environment: Optional[Dict[str, str]] = None, |
There was a problem hiding this comment.
Didn't we say we don't build these into the image anymore?
If you do, your image reuse doesn't work, as these environment variables contain e.g. credentials to access the ZenML server. If they get reused, they point to potentially a wrong ZenML user, are credentials for the wrong user, ...
| environment: Dict[str, str], | ||
| settings: ModalOrchestratorSettings, | ||
| shared_image_cache: Optional[Dict[str, Any]] = None, | ||
| shared_app: Optional[Any] = None, |
| run_id: The pipeline run ID. | ||
| synchronous: Whether to wait for completion. | ||
| """ | ||
| logger.debug("Executing entire pipeline in single sandbox") |
There was a problem hiding this comment.
This is called no matter which mode the orchestrator is running in.
| command = ( | ||
| ModalOrchestratorEntrypointConfiguration.get_entrypoint_command() | ||
| ) | ||
| from uuid import UUID |
| ) | ||
|
|
||
| # Submit the pipeline | ||
| run_id = str(placeholder_run.id) if placeholder_run else None |
There was a problem hiding this comment.
You're converting this to a string here, only to convert it back to a UUID in the executor.execute_pipeline method.
| environment_name=settings.modal_environment, | ||
| ) | ||
|
|
||
| def _build_entrypoint_command( |
There was a problem hiding this comment.
10 lines of code for a +? Can be removed I thjink
| else: | ||
| return f"{build_id}_pipeline_{image_hash}" | ||
|
|
||
| async def execute_pipeline( |
There was a problem hiding this comment.
This is more like start_pipeline_sandbox right?
| def _prepare_modal_api_params( | ||
| self, | ||
| entrypoint_command: List[str], | ||
| image: Any, |
|
|
||
| # Store the image ID for future caching after sandbox creation | ||
| # The image should be hydrated after being used in sandbox creation | ||
| await self._store_image_id(zenml_image) |
There was a problem hiding this comment.
Should probably happen before the sync waiting?
| """ | ||
| try: | ||
| # After sandbox creation, the image should be hydrated | ||
| zenml_image.hydrate() |
There was a problem hiding this comment.
No zenml image has this method, probably because this is a modal image instead. No Any type annotations.
| # Execute pipeline sandbox | ||
| await self._execute_sandbox( | ||
| entrypoint_command=entrypoint_command, | ||
| mode="PIPELINE", |
There was a problem hiding this comment.
Actually the mode here can be either pipeline or step, because this launches a sandbox that either
- runs the entire pipeline
- only orchestrator the steps
So we should rather pass the actual mode here, or nothing at all. This is anyway only used for tags.
New Features
Documentation
Refactor
Tests
Breaking Changes
None - this is a new integration that doesn't affect existing functionality.
Dependencies
Note: This orchestrator follows the same patterns as other ZenML orchestrators
(GCP Vertex, Kubernetes) and integrates seamlessly with the existing ZenML
stack architecture.
Note: I also updated the step operator logic to unify it
Pre-requisites
Please ensure you have done the following:
developand the open PR is targetingdevelop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.Types of changes
Summary by CodeRabbit
New Features
Enhancements
Bug Fixes
Refactor
Documentation