[FEAT] add driverPod/executorPod in Spark #6085

machichima · 2024-12-05T13:53:30Z

Tracking issue

Why are the changes needed?

Enable setting K8sPod separately for Spark Driver and Executor pods.

What changes were proposed in this pull request?

Add driverPod and executorPod field with type K8sPod in SparkJob. Uses existing mergePodSpecs to merge default podSpec with our driverPod or executorPod.

How was this patch tested?

Unit tests

I extended the existing Spark unit test TestBuildResourceContainer and TestBuildResourcePodTemplate and create a new test named TestBuildResourceCustomK8SPod for testing.

Test with my_spark example

See flyteorg/flytekit#3016

Setup process

Screenshots

Check all the applicable boxes

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

Related PRs

flyteorg/flytekit#3016

Docs link

Summary by Bito

This PR enhances Spark task configuration by introducing separate driver and executor pod specifications through new SparkJob message fields. The implementation adds support for customizing Kubernetes pod configurations independently for Spark driver and executor components, including tolerations, environment variables, and other pod settings. The changes encompass protobuf definitions with generated code for multiple languages and improvements to pod helper functions.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 4

Add driverPod/executorPod field in SparkJob class and use them as Spark driver and executor Signed-off-by: machichima <[email protected]>

codecov · 2024-12-05T13:56:57Z

Codecov Report

Attention: Patch coverage is 55.88235% with 30 lines in your changes missing coverage. Please review.

Project coverage is 36.87%. Comparing base (ea00864) to head (c7dbcfb).
Report is 9 commits behind head on master.

Files with missing lines	Patch %	Lines
flyteplugins/go/tasks/plugins/k8s/spark/spark.go	64.91%	14 Missing and 6 partials ⚠️
flyteidl/gen/pb-go/flyteidl/plugins/spark.pb.go	0.00%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6085      +/-   ##
==========================================
+ Coverage   36.86%   36.87%   +0.01%     
==========================================
  Files        1318     1318              
  Lines      134530   134703     +173     
==========================================
+ Hits        49591    49674      +83     
- Misses      80619    80702      +83     
- Partials     4320     4327       +7

Flag	Coverage Δ
unittests-datacatalog	`51.58% <ø> (ø)`
unittests-flyteadmin	`51.96% <ø> (+0.01%)`	⬆️
unittests-flytecopilot	`30.99% <ø> (ø)`
unittests-flytectl	`62.29% <ø> (ø)`
unittests-flyteidl	`7.22% <0.00%> (-0.02%)`	⬇️
unittests-flyteplugins	`54.03% <65.51%> (+0.11%)`	⬆️
unittests-flytepropeller	`42.78% <ø> (-0.01%)`	⬇️
unittests-flytestdlib	`55.35% <ø> (ø)`

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

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

fix protobuf number mismatch pass K8sPod instead of annotation and label separately Signed-off-by: machichima <[email protected]>

successfully apply pods specify in SparkJob Signed-off-by: machichima <[email protected]>

Signed-off-by: machichima <[email protected]>

flyteplugins/go/tasks/plugins/k8s/spark/spark.go

pingsutw · 2024-12-27T15:28:21Z

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go

+		// of c.Name
+		if val, ok := taskTemplate.GetConfig()[PrimaryContainerKey]; ok {
+			primaryContainerName = val
+			c.Name = primaryContainerName


could we add a small unit test for it

I am wondering whether this part is needed. As I know, to set the primary container name, we need to define in pod_template=PodTemplate(primary_container_name="primary") . But if we set this, we will get into case *core.TaskTemplate_K8SPod here. Therefore, this if val, ok := taskTemplate.GetConfig()[PrimaryContainerKey] will never be True.

I am thinking of removing this and use TaskTemplate_K8SPod in spark_test.go instead.

Signed-off-by: machichima <[email protected]>

flyte-bot · 2024-12-28T11:44:24Z

Code Review Agent Run #ebad2b

Actionable Suggestions - 8

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go - 2
- Consider validating primary container name value · Line 285-290
- Consider impact of function visibility change · Line 573-573
flyteplugins/go/tasks/plugins/k8s/spark/spark.go - 3
- Consider using appropriate unmarshal function · Line 69-69
- Consider extracting annotation merging logic · Line 152-165
- Consider nil check for PodSpec · Line 252-252
flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go - 3
- Consider extracting common task template code · Line 354-375
- Consider consolidating env var assertions · Line 1051-1055
- Fix typo in SecurityContext property name · Line 1100-1100

Additional Suggestions - 1

flyteplugins/go/tasks/plugins/k8s/spark/spark.go - 1
- Consider more specific utils import alias · Line 28-29

Review Details

Files reviewed - 10 · Commit Range: d847a63..48902c9
- flyteidl/gen/pb-es/flyteidl/plugins/spark_pb.ts
- flyteidl/gen/pb-go/flyteidl/plugins/spark.pb.go
- flyteidl/gen/pb_python/flyteidl/plugins/spark_pb2.py
- flyteidl/gen/pb_python/flyteidl/plugins/spark_pb2.pyi
- flyteidl/gen/pb_rust/flyteidl.plugins.rs
- flyteidl/protos/flyteidl/plugins/spark.proto
- flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go
- flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go
- flyteplugins/go/tasks/plugins/k8s/spark/spark.go
- flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go
Files skipped - 0
Tools
- Golangci-lint (Linter) - ✖︎ Failed
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2024-12-28T11:52:54Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Enhanced Spark Pod Configuration Support	- `spark_pb.ts` - Added driver and executor pod configuration support in SparkJob - `spark.pb.go` - Implemented driver and executor pod configuration in Go protobuf - `spark_pb2.py` - Added Python bindings for driver and executor pod configuration - `spark_pb2.pyi` - Updated Python type hints for new pod configuration fields - `flyteidl.plugins.rs` - Added Rust support for driver and executor pod configuration - `spark.proto` - Added K8sPod fields for driver and executor pod configuration
Other Improvements - Pod Helper Function Improvements	- `pod_helper.go` - Enhanced pod helper functions with primary container name configuration and exposed MergePodSpecs
Feature Improvement - Enhanced Spark Pod Configuration Support	- `spark_pb.ts` - Added driver and executor pod configuration support in SparkJob - `spark.pb.go` - Implemented driver and executor pod configuration in Go protobuf - `spark_pb2.py` - Added Python bindings for driver and executor pod configuration - `spark_pb2.pyi` - Updated Python type hints for new pod configuration fields - `flyteidl.plugins.rs` - Added Rust support for driver and executor pod configuration - `spark.proto` - Added K8sPod fields for driver and executor pod configuration - `spark.go` - Implemented support for custom driver and executor pod configurations - `spark_test.go` - Added tests for custom driver and executor pod configurations
Other Improvements - Pod Helper Function Improvements	- `pod_helper.go` - Enhanced pod helper functions with primary container name configuration and exposed MergePodSpecs - `pod_helper_test.go` - Updated tests to use exposed MergePodSpecs function

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go

flyteplugins/go/tasks/plugins/k8s/spark/spark.go

flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go

Signed-off-by: machichima <[email protected]>

As for spark driver/executor podTemplate, flytekit will only pass the container match the primary_container_name Signed-off-by: machichima <[email protected]>

Signed-off-by: machichima <[email protected]>

flyte-bot · 2025-01-19T09:03:46Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Agent didn't review this pull request automatically because it exceeded the size limit. No action is needed if you didn't intend for the agent to review it. Otherwise, you can initiate the review by typing /review in a comment below.

Signed-off-by: machichima <[email protected]>

flyte-bot · 2025-01-19T12:43:46Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Agent didn't review this pull request automatically because it exceeded the size limit. No action is needed if you didn't intend for the agent to review it. Otherwise, you can initiate the review by typing /review in a comment below.

Signed-off-by: machichima <[email protected]>

flyte-bot · 2025-01-19T14:43:57Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Agent didn't review this pull request automatically because it exceeded the size limit. No action is needed if you didn't intend for the agent to review it. Otherwise, you can initiate the review by typing /review in a comment below.

Signed-off-by: machichima <[email protected]>

flyte-bot · 2025-02-05T14:59:09Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Agent didn't review this pull request automatically because it exceeded the size limit. No action is needed if you didn't intend for the agent to review it. Otherwise, you can initiate the review by typing /review in a comment below.

Signed-off-by: machichima <[email protected]>

flyte-bot · 2025-02-12T13:51:18Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Agent didn't review this pull request automatically because it exceeded the size limit. No action is needed if you didn't intend for the agent to review it. Otherwise, you can initiate the review by typing /review in a comment below.

pingsutw · 2025-02-12T21:56:48Z

LGTM, thank you!

eapolinario

LGTM. Thank you.

feat: add driverPod/executorPod in Spark

d847a63

Add driverPod/executorPod field in SparkJob class and use them as Spark driver and executor Signed-off-by: machichima <[email protected]>

machichima added 4 commits December 6, 2024 23:29

fix+refactor: fix error+make code more readable

241dc46

fix protobuf number mismatch pass K8sPod instead of annotation and label separately Signed-off-by: machichima <[email protected]>

feat: successfully add driver/executor pod

f64e171

successfully apply pods specify in SparkJob Signed-off-by: machichima <[email protected]>

test: update test for SparkJob K8sPod (unfinished)

1faaa00

Signed-off-by: machichima <[email protected]>

feat: merge default and custom podSpec

88e3b6c

Signed-off-by: machichima <[email protected]>

machichima force-pushed the 4105-spark-driver-executor-podtemplate branch from ae39e8f to 394c269 Compare December 15, 2024 15:21

test: test custom driver/executor pod

da4199b

Signed-off-by: machichima <[email protected]>

machichima force-pushed the 4105-spark-driver-executor-podtemplate branch from 394c269 to da4199b Compare December 20, 2024 14:59

clean up unwanted comments

b7daee1

Signed-off-by: machichima <[email protected]>

machichima mentioned this pull request Dec 20, 2024

[FEAT] add driver/executor pod in Spark flyteorg/flytekit#3016

Open

3 tasks

machichima changed the title ~~[WIP] feat: add driverPod/executorPod in Spark~~ [FEAT] add driverPod/executorPod in Spark Dec 20, 2024

fix: fix lint issue

70cfdff

Signed-off-by: machichima <[email protected]>

machichima force-pushed the 4105-spark-driver-executor-podtemplate branch from c3eed97 to 70cfdff Compare December 21, 2024 03:14

pingsutw reviewed Dec 27, 2024

View reviewed changes

fix: based on suggested change

48902c9

Signed-off-by: machichima <[email protected]>