feat: KEP 2841 Flux Policy to support Flux Framework

[vsoch]

What this PR does / why we need it:

This KEP proposes adding a policy to support Flux Framework that provides more traditional HPC features. Using an HPC workload manager like Flux to bootstrap MPI will empower users to run MPI-based and other distributed workloads with advanced scheduling, topology awareness, and a more robust bootstrapping mechanism than traditional SSH-based methods. The proposal introduces a new flux policy in the TrainJob API, allowing users to select and configure the HPC workload manager, Flux.

The WIP implementation for the design discussed.

Authors

Myself and @milroy

Which issue(s) this PR fixes This will fix #2841.

Ping @andreyvelich @astefanutti

Checklist:

• Docs included if any changes are user facing

[andreyvelich]

Thanks for driving this great feature @vsoch, and sorry for the delay, got swamped with the KubeCon. I left my initial thoughts.

cc @kubeflow/kubeflow-trainer-team @kubeflow/kubeflow-sdk-team

[coveralls]

Pull Request Test Coverage Report for Build 20256291728

Details

• 0 of 0 changed or added relevant lines in 0 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage remained the same at 51.435%

---

Totals
Change from base Build 20255809727:	0.0%
Covered Lines:	1237
Relevant Lines:	2405

---

💛 - Coveralls

[andreyvelich]

/ok-to-test

[vsoch]

I think the error in CI is a flaky test? Note that I'm currently pushing for a more generic HPCPolicy that can support multiple plugin backends with a flexible Settings field. This means not using any hard coded variables (akin to the current MPI plugin).

[vsoch]

For the FluxMLPolicySource, we define the minimum required parameters needed for Flux and installing the view, along with the most highly used parameters in HPC. This largely includes the network device and view for compatibility. If you get it wrong, it cannot work. My first thinking was to create clean separation between fields for different components that will likely emerge:

type MLPolicySource struct {
    Torch *TorchMLPolicySource `json:"torch,omitempty"`
    MPI   *MPIMLPolicySource   `json:"mpi,omitempty"`

    // FluxMLPolicy defines policy only for Flux
	// +optional
    Flux  *FluxMLPolicySource  `json:"flux,omitempty"`
}

// FluxMLPolicySource represents a Flux HPC runtime configuration.
type FluxMLPolicySource struct {

	// numNodes is the number of physical nodes for the job.
	// This is defined a level up on the Trainer

	// numProcPerNode is the number of processes per node.
	// Defaults to 1.
	// +kubebuilder:default=1
	// +optional
	NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`	

	// fluxInstall describes how to install Flux
	// +optional
	FluxInstall FluxInstall `json:"install,omitempty"`	
}

// FluxInstall describes the install, network, and scheduling policy
// This is more modular for the Flux operator, and squashed here.
type FluxInstall struct {

    // Container image to use for Flux view that installs Flux
	// This must be compatible with the application container
	// Get the flux view container (these are choices)
	// ghcr.io/converged-computing/flux-view-rocky:arm-9
	// ghcr.io/converged-computing/flux-view-rocky:arn-8
	// ghcr.io/converged-computing/flux-view-rocky:tag-9
	// ghcr.io/converged-computing/flux-view-rocky:tag-8
	// ghcr.io/converged-computing/flux-view-ubuntu:tag-noble
	// ghcr.io/converged-computing/flux-view-ubuntu:tag-jammy
	// ghcr.io/converged-computing/flux-view-ubuntu:tag-focal
	// ghcr.io/converged-computing/flux-view-ubuntu:arm-jammy
	// ghcr.io/converged-computing/flux-view-ubuntu:arm-focal
	// We use an ubuntu (more recent) default since it is common
    // +kubebuilder:default="ghcr.io/converged-computing/flux-view-ubuntu:arm-jammy"
    Image string `json:"image,omitempty"`

    // Network device for flux to use
    // +kubebuilder:default="eth0"
    NetworkDevice string `json:"networkDevice,omitempty"`

    // Queue policy for Flux to use
    // +kubebuilder:default="fcfs"
    QueuePolicy string `json:"queuePolicy,omitempty"`
}

That said, the design of the others adheres to a flat structure, so I have refactored to reflect that - no FluxInstall but everything under one group.

// FluxMLPolicySource represents a Flux HPC runtime configuration.
type FluxMLPolicySource struct {

    // numNodes is the number of physical nodes for the job.
    // This is defined a level up on the Trainer

    // numProcPerNode is the number of processes per node.
    // Defaults to 1.
    // +kubebuilder:default=1
    // +optional
    NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`

    // Container image to use for Flux view that installs Flux
    // This must be compatible with the application container
    // Get the flux view container (these are choices)
    // ghcr.io/converged-computing/flux-view-rocky:arm-9
    // ghcr.io/converged-computing/flux-view-rocky:arn-8
    // ghcr.io/converged-computing/flux-view-rocky:tag-9
    // ghcr.io/converged-computing/flux-view-rocky:tag-8
    // ghcr.io/converged-computing/flux-view-ubuntu:tag-noble
    // ghcr.io/converged-computing/flux-view-ubuntu:tag-jammy
    // ghcr.io/converged-computing/flux-view-ubuntu:tag-focal
    // ghcr.io/converged-computing/flux-view-ubuntu:arm-jammy
    // ghcr.io/converged-computing/flux-view-ubuntu:arm-focal
    // We use an ubuntu (more recent) default since it is common
    // +kubebuilder:default="ghcr.io/converged-computing/flux-view-ubuntu:arm-jammy"
    Image string `json:"image,omitempty"`

    // Network device for flux to use
    // +kubebuilder:default="eth0"
    NetworkDevice string `json:"networkDevice,omitempty"`

    // Queue policy for Flux to use
    // +kubebuilder:default="fcfs"
    QueuePolicy string `json:"queuePolicy,omitempty"`
}

After staring at it, I think I prefer it! I will update the PR as needed for tests to pass (or mostly pass) and then we can do another review pass.

[vsoch]

@andreyvelich ready for another look!

[andreyvelich]

Thanks for the updates @vsoch!
Overall lgtm, I just left a few questions.

[vsoch]

Updates applied. Thanks for the review @andreyvelich ! Let me know if you have follow up questions in the discussion above.

[vsoch]

Updates:

• network policy and queue policy are removed. Defaults will be eth0 and fcfs
• initContainers is added as an example to the flux-runtime yaml spec.

[andreyvelich]

Thank you for this @vsoch! Exciting to see this moving forward!
/lgtm
/assign @astefanutti @tenzen-y @Electronic-Waste

[vsoch]

Thank you @andreyvelich for the speedy follow up reviews! I am also pumped to add Flux (and continue working on the implementation, which is the next step after the KEP).

Thanks to the other reviewers in advance for their feedback.

[vsoch]

Hi @astefanutti @tenzen-y @Electronic-Waste what do you need from our side to make progress here? Thanks!

[astefanutti]

/lgtm

Thanks @vsoch @milroy!

[astefanutti]

Maybe you've PoC'ed this already, but I don't think the runtime framework currently handles multiple plugins contributing to the same resource (in that case the JobSet resource). We may need to improve the framework machinery to handle this properly which is going to be useful generally.

[tenzen-y]

Yes, the order of plugin execution is not guaranteed.
So, we might need to introduce single truth of information across plugins.

We probably can make PodSetInfo as a truth cache or introduce another single truth cache.

[andreyvelich]

I guess, the single source of truth across plugins is Info object.
Like in MPI plugin, we apply volumes to it, so once the JobSet plugin is called, we sync those changes: https://github.com/kubeflow/trainer/blob/master/pkg/runtime/framework/plugins/mpi/mpi.go#L142

[tenzen-y]

This looks great, thank you for this effort, @vsoch !
/lgtm

[andreyvelich]

We should be good to move this forward.
Thanks again for this effort @vsoch!
/lgtm
/approve

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

• OWNERS [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[vsoch]

Thank you @andreyvelich @tenzen-y - the review was excellent, and I'm glad to see this moving through. It might be a bit early, but Happy New Year! I appreciate everything you do for the Kubeflow (and larger Kubernetes) communities.

[andreyvelich]

Thank you, excited to see this finally move forward!
Happy holidays @vsoch!