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: Add schedulingGates to PodSpecOverrides (#2700)

Browse Source 
* feat: Add schedulingGates to PodSpecOverrides

Signed-off-by: Antonin Stefanutti <[email protected]>

* Change desired job to target job in PodSpecOverrides comments

Signed-off-by: Antonin Stefanutti <[email protected]>

---------

Signed-off-by: Antonin Stefanutti <[email protected]>
This commit is contained in:
Antonin Stefanutti
2025-07-03 11:03:20 +00:00
committed by GitHub
parent c8d6a203ed
commit 8dc6334ccb
10 changed files with 198 additions and 14 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
+13 -3
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_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
from kubeflow_trainer_api.models.trainer_v1alpha1_container_override import TrainerV1alpha1ContainerOverride
@@ -30,14 +31,15 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
"""
PodSpecOverride represents the custom overrides that will be applied for the TrainJob's resources.
""" # noqa: E501
containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the containers in the desired job templates.")
init_containers: Optional[List[TrainerV1alpha1ContainerOverride]] = Field(default=None, description="Overrides for the init container in the desired job templates.", alias="initContainers")
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")
scheduling_gates: Optional[List[IoK8sApiCoreV1PodSchedulingGate]] = Field(default=None, description="SchedulingGates overrides the scheduling gates of the Pods in the target job templates. More info: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/", alias="schedulingGates")
service_account_name: Optional[StrictStr] = Field(default=None, description="Override for the service account.", alias="serviceAccountName")
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", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
__properties: ClassVar[List[str]] = ["containers", "initContainers", "nodeSelector", "schedulingGates", "serviceAccountName", "targetJobs", "tolerations", "volumes"]
model_config = ConfigDict(
populate_by_name=True,
@@ -92,6 +94,13 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
if _item_init_containers:
_items.append(_item_init_containers.to_dict())
_dict['initContainers'] = _items
# override the default output from pydantic by calling `to_dict()` of each item in scheduling_gates (list)
_items = []
if self.scheduling_gates:
for _item_scheduling_gates in self.scheduling_gates:
if _item_scheduling_gates:
_items.append(_item_scheduling_gates.to_dict())
_dict['schedulingGates'] = _items
# override the default output from pydantic by calling `to_dict()` of each item in target_jobs (list)
_items = []
if self.target_jobs:
@@ -128,6 +137,7 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
"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"),
"schedulingGates": [IoK8sApiCoreV1PodSchedulingGate.from_dict(_item) for _item in obj["schedulingGates"]] if obj.get("schedulingGates") is not None else None,
"serviceAccountName": obj.get("serviceAccountName"),
"targetJobs": [TrainerV1alpha1PodSpecOverrideTargetJob.from_dict(_item) for _item in obj["targetJobs"]] if obj.get("targetJobs") is not None else None,
"tolerations": [IoK8sApiCoreV1Toleration.from_dict(_item) for _item in obj["tolerations"]] if obj.get("tolerations") is not None else None,