feat(sdk) Add Input Parameter support for node affinity #12028
Conversation
VaniHaripriya
force-pushed
the
add-node-affinity
branch
4 times, most recently
from
229ed04 to
251291b
Compare
|
@VaniHaripriya Possible for you to provide an example of pipeline (sdk component) to verify this? |
VaniHaripriya
force-pushed
the
add-node-affinity
branch
8 times, most recently
from
b09e168 to
1c59ce6
Compare
|
/lgtm |
|
@VaniHaripriya Thanks for adding the test description, its very clear now as to how someone can add node affinity. However I am still not sure about the validation to see if this actually worked or not 🤔 |
@nsingla I confirmed that the node affinity settings are correctly reflected in the generated pipeline YAML and pod specs during runtime. I used a multi-node cluster and set node affinity criteria that do not match the node labels — this caused the pods to remain unscheduled, which confirms that the affinity rules were enforced. |
There was a problem hiding this comment.
Looks good in general @VaniHaripriya.
I just left some suggestions to validate the input at compile time (SDK) for a better experience.
| key=expr.get("key"), | ||
| operator=expr.get("operator"), | ||
| values=expr.get("values", []), | ||
| ) | ||
| if match_fields: | ||
| for field in match_fields: | ||
| affinity_term.match_fields.add( | ||
| key=field.get("key"), | ||
| operator=field.get("operator"), | ||
| values=field.get("values", []), |
There was a problem hiding this comment.
We could add some validations here to catch errors early.
- Validate that
keyandoperatorare non-empty - Validate that
operatorhas a valid value likeIn,NotIn,Exists...
Also, add tests for those validations.
| if weight is not None: | ||
| affinity_term.weight = weight |
There was a problem hiding this comment.
We can validate that weight is 1-100.
|
|
||
| def add_node_affinity_json( | ||
| task: PipelineTask, | ||
| node_affinity_json: Union[pipeline_channel.PipelineParameterChannel, dict], |
There was a problem hiding this comment.
We could validate node_affinity_json against the schema.
You probably can use something like the following (AI-generated):
from kubernetes import client, config
def validate_node_affinity(node_affinity_json: dict):
"""
Validates a JSON object against the V1NodeAffinity model.
Args:
node_affinity_json: A Python dictionary loaded from your JSON.
Returns:
True if the JSON is valid, False otherwise.
"""
try:
# Create an empty V1NodeAffinity object to merge into.
v1_node_affinity = client.V1NodeAffinity()
# The ApiMerger is used to deserialize Python dicts into Kubernetes models.
merger = client.api_client.ApiMerger()
merger.merge(v1_node_affinity, node_affinity_json)
print("✅ JSON is a valid V1NodeAffinity.")
return True
except Exception as e:
print(f"❌ Invalid V1NodeAffinity JSON: {e}")
return FalseAnd add a test for an invalid JSON.
VaniHaripriya
force-pushed
the
add-node-affinity
branch
2 times, most recently
from
73082f4 to
6d466f0
Compare
| if match_expressions: | ||
| for expr in match_expressions: | ||
| key = expr.get("key") | ||
| operator = expr.get("operator") | ||
| if not key: | ||
| raise ValueError("Each match_expression must have a non-empty 'key'.") | ||
| if not operator: | ||
| raise ValueError(f"Each match_expression for key '{key}' must have a non-empty 'operator'.") | ||
| if operator not in VALID_OPERATORS: | ||
| raise ValueError(f"Invalid operator '{operator}' for key '{key}'. Must be one of {sorted(VALID_OPERATORS)}.") | ||
| affinity_term.match_expressions.add( | ||
| key=key, | ||
| operator=operator, | ||
| values=expr.get("values", []), | ||
| ) | ||
| if match_fields: | ||
| for field in match_fields: | ||
| key = field.get("key") | ||
| operator = field.get("operator") | ||
| if not key: | ||
| raise ValueError("Each match_field must have a non-empty 'key'.") | ||
| if not operator: | ||
| raise ValueError(f"Each match_field for key '{key}' must have a non-empty 'operator'.") | ||
| if operator not in VALID_OPERATORS: | ||
| raise ValueError(f"Invalid operator '{operator}' for key '{key}'. Must be one of {sorted(VALID_OPERATORS)}.") | ||
| affinity_term.match_fields.add( | ||
| key=key, | ||
| operator=operator, | ||
| values=field.get("values", []), | ||
| ) |
There was a problem hiding this comment.
Sorry for not noticing earlier, but we could create a helper method to eliminate duplicate code here. Both loops do almost the same thing.
| try: | ||
| k8s_model_dict = common.deserialize_dict_to_k8s_model_keys(node_affinity_json) | ||
| client.V1NodeAffinity(**k8s_model_dict) | ||
| return True |
There was a problem hiding this comment.
nit: It doesn't need to return anything (can be void).
backend/src/v2/driver/k8s.go
Outdated
| k8snodeAffinity, err := json.Marshal(nodeAffinityJSON) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to marshal node affinity json: %w", err) | ||
| } | ||
| var nodeAffinity k8score.NodeAffinity | ||
| if err := json.Unmarshal(k8snodeAffinity, &nodeAffinity); err != nil { | ||
| return fmt.Errorf("failed to unmarshal node affinity json: %w", err) | ||
| } |
There was a problem hiding this comment.
I may be missing something, but it doesn't seem right to me. Marshal and then unmarshal?
There was a problem hiding this comment.
Updated the code to avoid both marshal and then unmarshal.
backend/src/v2/driver/k8s.go
Outdated
| var requiredTerms []k8score.NodeSelectorTerm | ||
| var preferredTerms []k8score.PreferredSchedulingTerm | ||
|
|
||
| for _, nodeAffinityTerm := range nodeAffinityTerms { |
There was a problem hiding this comment.
It might be worth adding a quick check at the start of the loop to skip empty terms early on:
for i, nodeAffinityTerm := range nodeAffinityTerms {
if nodeAffinityTerm.GetNodeAffinityJson() == nil &&
len(nodeAffinityTerm.GetMatchExpressions()) == 0 &&
len(nodeAffinityTerm.GetMatchFields()) == 0 {
glog.Warningf("NodeAffinityTerm %d is empty, skipping", i)
continue
}
// rest of processing...
}Right now, we only catch empty explicit terms later in the function, but JSON terms could also be empty and we'd still try to process them. This would save some unnecessary work and give clearer debugging info.
| def test_component_pipeline_input_required_scheduling(self): | ||
| """Test JSON-based node affinity with pipeline input for required scheduling.""" | ||
| @dsl.pipeline | ||
| def my_pipeline(affinity_input: dict): | ||
| task = comp() | ||
| kubernetes.add_node_affinity_json( | ||
| task, | ||
| node_affinity_json=affinity_input, | ||
| ) | ||
|
|
||
| assert json_format.MessageToDict(my_pipeline.platform_spec) == { | ||
| 'platforms': { | ||
| 'kubernetes': { | ||
| 'deploymentSpec': { | ||
| 'executors': { | ||
| 'exec-comp': { | ||
| 'nodeAffinity': [{ | ||
| 'nodeAffinityJson': { | ||
| 'componentInputParameter': 'affinity_input' | ||
| } | ||
| }] | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| def test_component_pipeline_input_preferred_scheduling(self): | ||
| """Test JSON-based node affinity with pipeline input for preferred scheduling.""" | ||
| @dsl.pipeline | ||
| def my_pipeline(affinity_input: dict): | ||
| task = comp() | ||
| kubernetes.add_node_affinity_json( | ||
| task, | ||
| node_affinity_json=affinity_input, | ||
| ) | ||
|
|
||
| assert json_format.MessageToDict(my_pipeline.platform_spec) == { | ||
| 'platforms': { | ||
| 'kubernetes': { | ||
| 'deploymentSpec': { | ||
| 'executors': { | ||
| 'exec-comp': { | ||
| 'nodeAffinity': [{ | ||
| 'nodeAffinityJson': { | ||
| 'componentInputParameter': 'affinity_input' | ||
| } | ||
| }] | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
These two tests look exactly the same.
| task, | ||
| node_affinity_json=affinity_input_2, | ||
| ) | ||
| print(json_format.MessageToDict(my_pipeline.platform_spec)) |
There was a problem hiding this comment.
Do we really want to keep this print, or is it a leftover?
There was a problem hiding this comment.
Its a leftover :-)
VaniHaripriya
force-pushed
the
add-node-affinity
branch
from
134c5b5 to
73d7f2c
Compare
Signed-off-by: VaniHaripriya <[email protected]>
VaniHaripriya
force-pushed
the
add-node-affinity
branch
from
73d7f2c to
3bd6ec5
Compare
Signed-off-by: VaniHaripriya <[email protected]>
VaniHaripriya
force-pushed
the
add-node-affinity
branch
from
3bd6ec5 to
8dcb764
Compare
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: hbelmiro The full list of commands accepted by this bot can be found here. The pull request process is described here
Needs approval from an approver in each of these files:
Approvers can indicate their approval by writing |
4 participants
Description of your changes:
This change enables the specification of node affinity through input parameters, in alignment with the enhancement discussed in #9682.
Testing Instructions
SDK
Use the example code to compile
You should be able to compile and find the following snippet in the main.yaml file:
Checklist: