feat: Extra Node Selector and Tolerations - #65
Conversation
📝 WalkthroughWalkthroughAdded optional node role scheduling via a Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Fix all issues with AI Agents 🤖
In @snakemake_executor_plugin_kubernetes/__init__.py:
- Around line 220-251: Wrap the parsing of resources_dict["node_taint"] (the
toleration = str(...), key, rest = toleration.split("=", 1), value, effect =
rest.split(":", 1)) in a try/except that catches ValueError and raises a
WorkflowError with an informative message about the expected "key=value:effect"
format; keep the existing checks for empty parts and effect validity, and when
constructing the kubernetes.client.V1Toleration use consistent keyword spacing
(e.g., key=key, operator="Equal", value=value, effect=effect) and retain the
self.logger.debug call to log the added toleration.
🧹 Nitpick comments (2)
snakemake_executor_plugin_kubernetes/__init__.py (2)
206-209: Consider documenting or making the node selector label key configurable.The implementation hardcodes
"role"as the node selector key. While this may work for specific cluster setups, Kubernetes doesn't enforce a standard "role" label. Consider:
- Documenting that users must label their nodes with
role=<value>for this feature to work- Alternatively, making the label key configurable (e.g.,
node_role_labelsetting) to support different cluster labeling schemesExample of configurable approach
If configurability is desired, you could allow users to specify the full label in the resource:
# Example: node_selector_label="custom-pool=snakemake-exec" if "node_selector_label" in resources_dict: label_spec = str(resources_dict["node_selector_label"]) if "=" in label_spec: key, value = label_spec.split("=", 1) node_selector[key] = value else: node_selector[label_spec] = "true"However, the current implementation is acceptable if the
rolelabel is standard in your environment.
229-238: Optional: Consider the static analysis hint about exception messages.The static analysis tool suggests avoiding long messages outside the exception class (TRY003). However, for this use case, the contextual error messages are appropriate and clear. You may safely ignore this hint, or if you prefer to address it, you could define custom exception subclasses.
This is a low-priority style suggestion and doesn't affect functionality. Based on coding guidelines for this project.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
snakemake_executor_plugin_kubernetes/__init__.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
**/*.py: Do not try to improve formatting.
Do not suggest type annotations for functions that are defined inside of functions or methods.
Do not suggest type annotation of theselfargument of methods.
Do not suggest type annotation of theclsargument of classmethods.
Do not suggest return type annotation if a function or method does not contain areturnstatement.
Files:
snakemake_executor_plugin_kubernetes/__init__.py
🪛 Ruff (0.14.10)
snakemake_executor_plugin_kubernetes/__init__.py
230-233: Avoid specifying long messages outside the exception class
(TRY003)
235-238: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (1)
snakemake_executor_plugin_kubernetes/__init__.py (1)
220-251: Integration with existing GPU tolerations is correct.The node_taint toleration handling integrates well with the existing GPU toleration logic (lines 252-298). Both properly initialize
pod_spec.tolerationsif needed and can coexist, allowing pods to tolerate both custom taints and GPU-specific taints.
MessereN
changed the title
…s missing characters in the format
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
snakemake_executor_plugin_kubernetes/__init__.py (1)
220-258: LGTM! Error handling and validation are solid.The implementation correctly handles the
node_taintresource with proper error handling for parsing failures and comprehensive validation. The try/except block at lines 227-234 addresses the parsing concern from the previous review.Optionally, you can improve the exception chaining by using
raise ... from Noneto indicate that the ValueError is intentionally being replaced:🔎 Optional improvement for exception chaining
try: key, rest = toleration.split("=", 1) value, effect = rest.split(":", 1) except ValueError: - raise WorkflowError( + raise WorkflowError( f"Invalid node_taint format: '{toleration}'. " "Expected format: key=value:effect" - ) + ) from NoneThis makes it clearer that the ValueError is expected and being intentionally replaced with a more user-friendly WorkflowError.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
snakemake_executor_plugin_kubernetes/__init__.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
⚙️ CodeRabbit configuration file
**/*.py: Do not try to improve formatting.
Do not suggest type annotations for functions that are defined inside of functions or methods.
Do not suggest type annotation of theselfargument of methods.
Do not suggest type annotation of theclsargument of classmethods.
Do not suggest return type annotation if a function or method does not contain areturnstatement.
Files:
snakemake_executor_plugin_kubernetes/__init__.py
🪛 Ruff (0.14.10)
snakemake_executor_plugin_kubernetes/__init__.py
231-234: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
231-234: Avoid specifying long messages outside the exception class
(TRY003)
237-240: Avoid specifying long messages outside the exception class
(TRY003)
242-245: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (1)
snakemake_executor_plugin_kubernetes/__init__.py (1)
206-209: No changes needed — the hardcoded "role" label key is intentional.The code comment explicitly documents that this feature is for scheduling on GKE burst-pool nodes using the
rolelabel key (e.g.,role=snakemake-exec). The hardcoding is not a flexibility issue but rather the correct implementation for this specific use case. Therolelabel is the standard convention for GKE burst pools, and the design is properly documented in the code.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@snakemake_executor_plugin_kubernetes/__init__.py`:
- Around line 254-257: The log message in the logger.debug call incorrectly
hardcodes the prefix "workload=" while emitting the value from
resources_dict['node_taint']; update the message to use a generic or accurate
label (e.g., "taint=" or "node_taint=") and/or interpolate the actual key name
instead of "workload" so the output matches the value; modify the call to
logger.debug(...) that references resources_dict['node_taint'] to something like
"Added taint toleration: taint={resources_dict['node_taint']}" (or
"node_taint={...}") to remove the misleading "workload=" prefix.
🧹 Nitpick comments (2)
snakemake_executor_plugin_kubernetes/__init__.py (2)
230-234: Chain the exception for better debugging.When re-raising in an except block, use
from errto preserve the original traceback, which aids debugging.♻️ Proposed fix
- except ValueError: + except ValueError as err: raise WorkflowError( f"Invalid node_taint format: '{toleration}'. " "Expected format: key=value:effect" - ) + ) from err
306-312: Duplicate privileged container handling.The security context is set at lines 306-312, and then again at lines 400-405 when
self.privilegedis True. This creates redundant code - whenself.privileged=True, the context is assigned twice.Consider consolidating this logic. The block at lines 400-405 could be removed entirely, or the condition at line 306 could be simplified to only handle the APPTAINER case (since the self.privileged case is handled later with logging).
♻️ Option 1: Remove the redundant block at lines 400-405
Add logging to the first block and remove the second:
# capabilities if self.privileged or ( DeploymentMethod.APPTAINER in self.workflow.deployment_settings.deployment_method ): container.security_context = kubernetes.client.V1SecurityContext( privileged=True ) + self.logger.debug("Container set to run in privileged mode.")Then delete lines 400-405.
♻️ Option 2: Keep blocks separate but avoid duplicate assignment
# capabilities -if self.privileged or ( - DeploymentMethod.APPTAINER +if ( + DeploymentMethod.APPTAINER in self.workflow.deployment_settings.deployment_method ): container.security_context = kubernetes.client.V1SecurityContext( privileged=True ) + self.logger.debug("Container set to run in privileged mode (APPTAINER).")This keeps the APPTAINER-specific handling separate from the explicit privileged flag handling at lines 400-405.
Also applies to: 400-405
| self.logger.debug( | ||
| "Added workload taint toleration: " | ||
| f"workload={resources_dict['node_taint']}" | ||
| ) |
There was a problem hiding this comment.
Misleading log message prefix.
The log hardcodes "workload=" but the toleration key can be any value, not just "workload". If node_taint="foo=bar:NoSchedule", this logs "workload=foo=bar:NoSchedule" which is confusing.
🐛 Proposed fix
self.logger.debug(
- "Added workload taint toleration: "
- f"workload={resources_dict['node_taint']}"
+ f"Added taint toleration: {resources_dict['node_taint']}"
)📝 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.
| self.logger.debug( | |
| "Added workload taint toleration: " | |
| f"workload={resources_dict['node_taint']}" | |
| ) | |
| self.logger.debug( | |
| f"Added taint toleration: {resources_dict['node_taint']}" | |
| ) |
🤖 Prompt for AI Agents
In `@snakemake_executor_plugin_kubernetes/__init__.py` around lines 254 - 257, The
log message in the logger.debug call incorrectly hardcodes the prefix
"workload=" while emitting the value from resources_dict['node_taint']; update
the message to use a generic or accurate label (e.g., "taint=" or "node_taint=")
and/or interpolate the actual key name instead of "workload" so the output
matches the value; modify the call to logger.debug(...) that references
resources_dict['node_taint'] to something like "Added taint toleration:
taint={resources_dict['node_taint']}" (or "node_taint={...}") to remove the
misleading "workload=" prefix.
1 participant
Added additional functionality in the plugin to account for new resource keys: node_role and node_taint.
node_role: Allows users to define in either default-resources or in rule specific resources to schedule jobs on labeled nodes.
cluster configuration: ie. node-labels=role=snakemake-exec
snakemake resource usage: ie. node_role=snakemake-exec
node_taint: Introduces a new resource key that enables users to specify custom taints on their nodes in order to craft the necessary tolerations for the pod to be scheduled on the tainted node(s). This can be initialized in either default-resources or within rule specific resources. The expected format of this resource is key=value:effect. It is essential to note that the only supported operator is "=" but perhaps in the future this can be expanded upon. Also, there are error mechanisms in place to account for proper formatting and ensuring that the effect is one of the supported NoSchedule, PreferNoSchedule, NoExecute.
cluster configuration: ie. node-taints=workload=snakemake:NoSchedule
snakemake resource usage: ie. node_taint=workload=snakemake:NoSchedule
These extra resources are useful for custom node taints and selectors on a GKE cluster so jobs can land properly on specific nodes/pools and have the associated tolerations.
Summary by CodeRabbit
New Features
Bug Fixes / Behavior
✏️ Tip: You can customize this high-level summary in your review settings.