Add driver and executor pod template fields to SparkJob #4179

andrewwdye · 2023-10-06T00:18:43Z

Tracking issue

Describe your changes

This change defines a common pattern for per-role pod configuration and adds RoleSpec fields for driver and executor pods, accordingly.

Adds a RoleSpec to common.proto with initial fields for pod and pod_template_name. This mirrors the interface we have for container tasks spread across TaskTemplate and TaskMetadata. Fields are optional.
Adds driverSpec and executorSpec to SparkJob

This PR was manually migrated from flyteorg/flyteidl#450. See initial comment discussions there.

How should a plugin handle TaskTemplate/TaskMetadata pod templates vs. these RoleSpecs?
For Spark, the task-level pod template is not consumed today and can be treated as an error if set (not supported). The plugin can inspect optional per-role specs for configuration.

For existing plugins that consume the task-level pod template (Dask, MPI, PyTorch, Ray), they can continue doing so. If/when we are ready to introduce per-role pod configuration, we can add RoleSpec fields to the plugin's config, deprecate the task-level pod template, and rely on a runtime check to validate that only task-level pod template OR per-role pod templates are specified.

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Signed-off-by: Andrew Dye <[email protected]>

wild-endeavor · 2023-10-06T04:04:16Z

@eapolinario i see that one +1 was enough to merge this pr. should we continue to make idl a +2? I think we should.

not sure if there was resolution to @hamersaw's comment.

eapolinario · 2023-10-06T17:43:52Z

@eapolinario i see that one +1 was enough to merge this pr. should we continue to make idl a +2? I think we should.

not sure if there was resolution to @hamersaw's comment.

As far as I can tell CODEOWNERS does not support that feature (as per https://github.com/orgs/community/discussions/22522). I opened #4182 to track this.

pingsutw

Should we follow the same pattern as tensorflow.proto

like

message SparkJob {
    ...
    sparkReplicaSpec driver = 11;
    sparkReplicaSpec worker = 12;
}

message sparkReplicaSpec {
    oneof PodValue {
        core.K8sPod podTemplate = 1;
    }
    string PodTemplateName = 2;
    string image = 3;
     ...
}

Maybe we should put repliaSpec to common.proto, all other plugins can import this proto in the future, like Ray.

# common.proto
message ReplicaSpec {
    oneof PodValue {
        core.K8sPod podTemplate = 1;
    }
    string PodTemplateName = 2;
    string image = 3;
     ...
}

cc @eapolinario @wild-endeavor @hamersaw wdyt

andrewwdye · 2023-10-06T20:02:30Z

Wanted to also cross reference this comment thread from the original PR, proposing using the TaskTemplate.PodTemplate for all roles, unless specifically overridden with a role-specific template/spec.

@task(
    task_config=Spark(
        spark_conf={
           "spark.driver.memory": "4G",
           "spark.executor.memory": "8G",
        },
        # driver only pod template overrides (takes precedence over top level pod_template)
        driver_pod_template=...
    ),
    # default pod template overrides for BOTH driver and executor
    pod_template=...
)

andrewwdye · 2023-10-06T20:18:13Z

Pulling together references from existing leader/worker patterns for various plugins

Dask uses a common spec for scheduler/worker w/ scheduler treated as non-interruptible
MPI uses a common pod spec with some per role args
PyTorch uses a common pod spec with some per role args
Tensorflow uses separate replica specs for each of chief/worker/ps
Ray uses a common pod spec for headgroup and worker

It does seem like a good idea to establish a pattern (and shared definitions if appropriate)

Signed-off-by: Andrew Dye <[email protected]>

codecov · 2023-10-08T05:31:56Z

Codecov Report

All modified lines are covered by tests ✅

Comparison is base (4ee73f5) 59.44% compared to head (8d12d1b) 78.48%.
Report is 4 commits behind head on master.

❗ Current head 8d12d1b differs from pull request most recent head 6fefc0b. Consider uploading reports for the commit 6fefc0b to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #4179       +/-   ##
===========================================
+ Coverage   59.44%   78.48%   +19.03%     
===========================================
  Files         637       18      -619     
  Lines       54262     1250    -53012     
===========================================
- Hits        32254      981    -31273     
+ Misses      19467      212    -19255     
+ Partials     2541       57     -2484

Flag	Coverage Δ
unittests	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

see 637 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

andrewwdye · 2023-10-08T05:32:20Z

Based on comments from the initial PR thread and above, I've moved pod and pod template name behind a RoleSpec definition in common.proto. See latest PR summary for the pattern this establishes and expected handling of TaskTemplate.K8sPod.

wild-endeavor

+1 but is this the right repo @eapolinario ?

hamersaw · 2023-10-09T14:07:35Z

flyteidl/protos/flyteidl/plugins/spark.proto

+    RoleSpec driverSpec = 10;
+    // The executor spec, used in place of the task's pod template
+    RoleSpec executorSpec = 11;


Is there ever a scenario where both of these differ from the values set in the task decorator?

Curious what you mean here. With this change we will fail

@task( pod_template=... )

in favor of specifying one or both values in

@task( task_config=Spark( executor_pod_template=... driver_pod_template=... ) )

hamersaw · 2023-10-09T14:10:53Z

Maybe we should put repliaSpec to common.proto, all other plugins can import this proto in the future, like Ray.

In theory I think this is a great idea. But in practice I would bet there will be corner case plugins where not all options are supported.

andrewwdye

In theory I think this is a great idea. But in practice I would bet there will be corner case plugins where not all options are supported

Are you thinking of initial fields in RoleSpec or future ones? If we believe that K8sPod is the correct abstraction for overriding parts of the default pod template, then including it as part of the role spec seems appropriate. I could imagine, however, additional role-specific configuration desired over time that wouldn't fit into a common RoleSpec definition. Thoughts?

andrewwdye · 2023-10-09T15:03:30Z

flyteidl/protos/flyteidl/plugins/spark.proto

+    RoleSpec driverSpec = 10;
+    // The executor spec, used in place of the task's pod template
+    RoleSpec executorSpec = 11;


Curious what you mean here. With this change we will fail

@task( pod_template=... )

in favor of specifying one or both values in

@task( task_config=Spark( executor_pod_template=... driver_pod_template=... ) )

andrewwdye · 2023-10-09T19:28:19Z

Given some of the outstanding questions for Spark, the pattern for multi-role plugins more broadly, and the deprecation/migration story, let me close this out until we get alignment. For now Spark plugin can use the existing TaskTemplate K8sPod field to override tolerations for both driver and executor pods, which is complimentary to the existing interface for node selectors: spark.kubernetes.node.selector... (same values for both roles). This is also consistent with how other plugins use pod templates for multiple roles today.

Add driver and executor pod template fields to SparkJob

6e45b86

Signed-off-by: Andrew Dye <[email protected]>

andrewwdye mentioned this pull request Oct 6, 2023

Add driver and executor pod template fields to SparkJob flyteorg/flyteidl#450

Closed

8 tasks

andrewwdye requested review from wild-endeavor, hamersaw and katrogan October 6, 2023 00:19

wild-endeavor previously approved these changes Oct 6, 2023

View reviewed changes

andrewwdye mentioned this pull request Oct 6, 2023

Implement Spark pod template overrides #4183

Merged

3 tasks

pingsutw reviewed Oct 6, 2023

View reviewed changes

Use RoleSpec for driver and executor

6fefc0b

Signed-off-by: Andrew Dye <[email protected]>

andrewwdye dismissed wild-endeavor’s stale review via 6fefc0b October 8, 2023 05:27

andrewwdye requested review from wild-endeavor and pingsutw October 8, 2023 05:45

pingsutw approved these changes Oct 8, 2023

View reviewed changes

wild-endeavor approved these changes Oct 9, 2023

View reviewed changes

hamersaw reviewed Oct 9, 2023

View reviewed changes

andrewwdye commented Oct 9, 2023

View reviewed changes

andrewwdye closed this Oct 9, 2023

andrewwdye mentioned this pull request Oct 9, 2023

Add driver and executor pod template params to Spark config flyteorg/flytekit#1863

Closed

8 tasks

troychiu mentioned this pull request Nov 12, 2023

[Core feature] Refactor distributed job using common ReplicaSpec #4408

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add driver and executor pod template fields to SparkJob #4179

Add driver and executor pod template fields to SparkJob #4179

andrewwdye commented Oct 6, 2023 •

edited

Loading

wild-endeavor commented Oct 6, 2023

eapolinario commented Oct 6, 2023

pingsutw left a comment

andrewwdye commented Oct 6, 2023 •

edited

Loading

andrewwdye commented Oct 6, 2023 •

edited

Loading

codecov bot commented Oct 8, 2023

andrewwdye commented Oct 8, 2023

wild-endeavor left a comment

hamersaw Oct 9, 2023 •

edited

Loading

andrewwdye Oct 9, 2023

hamersaw commented Oct 9, 2023

andrewwdye left a comment

andrewwdye Oct 9, 2023

andrewwdye commented Oct 9, 2023

Add driver and executor pod template fields to SparkJob #4179

Add driver and executor pod template fields to SparkJob #4179

Conversation

andrewwdye commented Oct 6, 2023 • edited Loading

Tracking issue

Describe your changes

Check all the applicable boxes

wild-endeavor commented Oct 6, 2023

eapolinario commented Oct 6, 2023

pingsutw left a comment

Choose a reason for hiding this comment

andrewwdye commented Oct 6, 2023 • edited Loading

andrewwdye commented Oct 6, 2023 • edited Loading

codecov bot commented Oct 8, 2023

Codecov Report

andrewwdye commented Oct 8, 2023

wild-endeavor left a comment

Choose a reason for hiding this comment

hamersaw Oct 9, 2023 • edited Loading

Choose a reason for hiding this comment

andrewwdye Oct 9, 2023

Choose a reason for hiding this comment

hamersaw commented Oct 9, 2023

andrewwdye left a comment

Choose a reason for hiding this comment

andrewwdye Oct 9, 2023

Choose a reason for hiding this comment

andrewwdye commented Oct 9, 2023

andrewwdye commented Oct 6, 2023 •

edited

Loading

andrewwdye commented Oct 6, 2023 •

edited

Loading

andrewwdye commented Oct 6, 2023 •

edited

Loading

hamersaw Oct 9, 2023 •

edited

Loading