hanzo-ml/trainer
1
0
Fork 0
mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support affinity in TrainJob pod spec overrides (#2796)

Browse Source 
Signed-off-by: Tsubasa Nagasawa <[email protected]>
This commit is contained in:
Tsubasa Nagasawa
2025-08-28 15:52:22 +00:00
committed by GitHub
parent 023ca6ce32
commit 187a5de400
11 changed files with 2063 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_spec_override.py
+7 -1
View File
@@ -19,6 +19,7 @@ import json
from pydantic import BaseModel, ConfigDict, Field, StrictStr
from typing import Any, ClassVar, Dict, List, Optional
from kubeflow_trainer_api.models.io_k8s_api_core_v1_affinity import IoK8sApiCoreV1Affinity
from kubeflow_trainer_api.models.io_k8s_api_core_v1_pod_scheduling_gate import IoK8sApiCoreV1PodSchedulingGate
from kubeflow_trainer_api.models.io_k8s_api_core_v1_toleration import IoK8sApiCoreV1Toleration
from kubeflow_trainer_api.models.io_k8s_api_core_v1_volume import IoK8sApiCoreV1Volume
@@ -31,6 +32,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
"""
PodSpecOverride represents the custom overrides that will be applied for the TrainJob's resources.
""" # noqa: E501
affinity: Optional[IoK8sApiCoreV1Affinity] = Field(default=None, description="Override for the Pod's affinity.")
containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the containers in the target job templates.")
init_containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the init container in the target job templates.", alias="initContainers")
node_selector: Optional[Dict[str, StrictStr]] = Field(default=None, description="Override for the node selector to place Pod on the specific node.", alias="nodeSelector")
@@ -39,7 +41,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
target_jobs: List[TrainerV1alpha1PodSpecOverrideTargetJob] = Field(description="TrainJobs is the training job replicas in the training runtime template to apply the overrides.", alias="targetJobs")
tolerations: Optional[List[IoK8sApiCoreV1Toleration]] = Field(default=None, description="Override for the Pod's tolerations.")
volumes: Optional[List[IoK8sApiCoreV1Volume]] = Field(default=None, description="Overrides for the Pod volume configurations.")
__properties: ClassVar[List[str]] = ["containers", "initContainers", "nodeSelector", "schedulingGates", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
__properties: ClassVar[List[str]] = ["affinity", "containers", "initContainers", "nodeSelector", "schedulingGates", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
model_config = ConfigDict(
populate_by_name=True,
@@ -80,6 +82,9 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
exclude=excluded_fields,
exclude_none=True,
)
# override the default output from pydantic by calling `to_dict()` of affinity
if self.affinity:
_dict['affinity'] = self.affinity.to_dict()
# override the default output from pydantic by calling `to_dict()` of each item in containers (list)
_items = []
if self.containers:
@@ -134,6 +139,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
return cls.model_validate(obj)
_obj = cls.model_validate({
"affinity": IoK8sApiCoreV1Affinity.from_dict(obj["affinity"]) if obj.get("affinity") is not None else None,
"containers": [TrainerV1alpha1ContainerOverride.from_dict(_item) for _item in obj["containers"]] if obj.get("containers") is not None else None,
"initContainers": [TrainerV1alpha1ContainerOverride.from_dict(_item) for _item in obj["initContainers"]] if obj.get("initContainers") is not None else None,
"nodeSelector": obj.get("nodeSelector"),