feat: Implement Training Options pattern for flexible TrainJob customization

[abhijeet-dhumal]

What this PR does / why we need it:

Fixes #87, #92, #116

QuickStart :

from kubeflow.trainer import TrainerClient, CustomTrainer
from kubeflow.trainer.options import Name, Labels

client = TrainerClient()
client.train(
    runtime=client.get_runtime("torch-distributed"),
    trainer=CustomTrainer(func=my_function),
    options=[
        Name(name="my-job"),
        Labels(labels={"team": "ml"})
    ]
)

Common Options:

from kubeflow.trainer.options import Name

Kubernetes Options

from kubeflow.trainer.options.kubernetes import (
    Name,        #JobName
    Labels,              # Job metadata labels
    Annotations,         # Job metadata annotations
    SpecLabels,          # Labels for pods/jobsets
    SpecAnnotations,     # Annotations for pods/jobsets
    PodTemplateOverrides,  # Pod customization
    TrainerCommand,      # Override container command (CustomTrainerContainer only)
    TrainerArgs,         # Override container args (CustomTrainerContainer only)
)

Basic config:

from kubeflow.trainer.options import Name, Labels

options=[
    Name(name="my-job"),
    Labels(labels={"team": "ml-team", "project": "nlp"})
]

Custom container specific options:

from kubeflow.trainer.options import TrainerCommand, TrainerArgs

options=[
    TrainerCommand(["python", "-m", "torch.distributed.run"]),
    TrainerArgs(["train.py", "--epochs=10"])
]

Pod customisation :

from kubeflow.trainer.options import (
    PodTemplateOverride,
    PodTemplateOverrides,
    PodSpecOverride
)

options=[
    PodTemplateOverrides(
        PodTemplateOverride(
            target_jobs=["node"],
            spec=PodSpecOverride(
                node_selector={"gpu-type": "a100"},
                volumes=[{
                    "name": "data",
                    "persistentVolumeClaim": {"claimName": "my-data"}
                }]
            )
        )
    )
]

Creating custom options:

from dataclasses import dataclass
from typing import Any, Optional

@dataclass
class MyCustomOption:
    """My custom option for Kubernetes jobs."""
    
    timeout_seconds: int
    
    def __call__(self, job_spec: dict[str, Any], trainer, backend) -> None:
        # Validate backend compatibility
        from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
        if not isinstance(backend, KubernetesBackend):
            raise ValueError("MyCustomOption only works with Kubernetes")
        
        # Apply your changes
        spec = job_spec.setdefault("spec", {})
        spec["activeDeadlineSeconds"] = self.timeout_seconds

client.train(
    trainer=my_trainer,
    options=[MyCustomOption(timeout_seconds=3600)]
)

Overrides:

PodTemplateOverride Fields

PodTemplateOverride(
    target_jobs=["node", "launcher"],  # Which pods to apply to
    metadata={...},                     # Pod labels/annotations
    spec=PodSpecOverride(...)          # Pod spec customizations
)

PodSpecOverride Fields

PodSpecOverride(
    service_account_name="...",
    node_selector={...},
    tolerations=[...],
    affinity={...},
    volumes=[...],
    init_containers=[...],
    containers=[...],
    scheduling_gates=[...],
    image_pull_secrets=[...]
)

ContainerOverride Fields

ContainerOverride(
    name="trainer",
    env=[...],
    volume_mounts=[...]
)

Checklist:

• Docs included if any changes are user facing

[szaher]

Thanks @abhijeet-dhumal

I believe we need to handle how options will be applied for while backends. either we ignore options for localprocess or make options targeted towards specific backend.

[coveralls]

Pull Request Test Coverage Report for Build 19096824715

Details

• 390 of 463 (84.23%) changed or added relevant lines in 11 files are covered.
• 1 unchanged line in 1 file lost coverage.
• Overall coverage decreased (-11.6%) to 68.007%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
kubeflow/trainer/api/trainer_client.py	0	1	0.0%
kubeflow/trainer/options/kubernetes_test.py	85	86	98.84%
kubeflow/trainer/backends/localprocess/backend.py	17	19	89.47%
kubeflow/trainer/backends/localprocess/backend_test.py	111	114	97.37%
kubeflow/trainer/backends/kubernetes/backend.py	25	32	78.13%
kubeflow/trainer/options/kubernetes.py	117	176	66.48%

Files with Coverage Reduction	New Missed Lines	%
kubeflow/trainer/backends/kubernetes/backend.py	1	79.92%

Totals
Change from base Build 19075264540:	-11.6%
Covered Lines:	2487
Relevant Lines:	3657

---

💛 - Coveralls

[astefanutti]

Thanks @abhijeet-dhumal for this outstanding contribution (and your patience 😅)!

/lgtm

/assign @andreyvelich @kramaranya

[kramaranya]

Thank you for this amazing work, @abhijeet-dhumal! 🎉
/lgtm

Also, we just beat the record for the most commented PR :)

[andreyvelich]

Thank you for this amazing contribution @abhijeet-dhumal!
/lgtm
/approve
/hold cancel

[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

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

[abhijeet-dhumal]

Thanks @andreyvelich @astefanutti @kramaranya , delighted to have made this contribution. Thank you to all reviewers for their help.
For the next course of action, I will be adding PR in https://github.com/kubeflow/website/tree/master/content/en/docs/components/trainer/user-guides too, for documenting the training options available via SDK 🏁
Thanks 🥳