mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
feat(controller): Implement PodSpecOverride API (#2614)
* feat(controller): Implement PodSpecOverride API Signed-off-by: Andrey Velichkevich <[email protected]> * Update API comments Signed-off-by: Andrey Velichkevich <[email protected]> * Fix InitContainers override Signed-off-by: Andrey Velichkevich <[email protected]> * Fix generate Signed-off-by: Andrey Velichkevich <[email protected]> * Update pkg/apis/trainer/v1alpha1/trainjob_types.go Co-authored-by: Shao Wang <[email protected]> Signed-off-by: Andrey Velichkevich <[email protected]> * Change API to targetJob and add validation Signed-off-by: Andrey Velichkevich <[email protected]> * Add unit test for the JobSet builder Signed-off-by: Andrey Velichkevich <[email protected]> * Add integration tests Signed-off-by: Andrey Velichkevich <[email protected]> * Update pkg/apis/trainer/v1alpha1/trainjob_types.go Co-authored-by: Shao Wang <[email protected]> Signed-off-by: Andrey Velichkevich <[email protected]> * Update pkg/apis/trainer/v1alpha1/trainjob_types.go Co-authored-by: Shao Wang <[email protected]> Signed-off-by: Andrey Velichkevich <[email protected]> * Update pkg/apis/trainer/v1alpha1/trainjob_types.go Co-authored-by: Antonin Stefanutti <[email protected]> Signed-off-by: Andrey Velichkevich <[email protected]> * Merge PodSpecOverrides before the build phase Signed-off-by: Antonin Stefanutti <[email protected]> * Fix make generate Signed-off-by: Antonin Stefanutti <[email protected]> * Fix PodSpec SMT patch unmarshalling into JobSet template Signed-off-by: Antonin Stefanutti <[email protected]> * Add UT for volumes PodSpecOverrides Signed-off-by: Antonin Stefanutti <[email protected]> * Add UT for tolerations PodSpecOverrides Signed-off-by: Antonin Stefanutti <[email protected]> * Add UT for node selector PodSpecOverrides Signed-off-by: Antonin Stefanutti <[email protected]> * Revert back to using TargetJobs []PodSpecOverrideTargetJob Signed-off-by: Antonin Stefanutti <[email protected]> * Run make fmt Signed-off-by: Antonin Stefanutti <[email protected]> * Add target jobs unicity validation for PodSpecOverrides Signed-off-by: Antonin Stefanutti <[email protected]> * Update with code review from @tenzen-y Signed-off-by: Antonin Stefanutti <[email protected]> --------- Signed-off-by: Andrey Velichkevich <[email protected]> Signed-off-by: Antonin Stefanutti <[email protected]> Co-authored-by: Shao Wang <[email protected]> Co-authored-by: Antonin Stefanutti <[email protected]> Co-authored-by: Antonin Stefanutti <[email protected]>
This commit is contained in:
Andrey Velichkevich
committed by
GitHub
GitHub
co-authored by
Shao Wang
Antonin Stefanutti
Antonin Stefanutti
Shao Wang
Antonin Stefanutti
Antonin Stefanutti
parent
2a9911b927
commit
434d953004
@@ -19,7 +19,6 @@ 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_env_from_source import IoK8sApiCoreV1EnvFromSource
|
||||
from kubeflow_trainer_api.models.io_k8s_api_core_v1_env_var import IoK8sApiCoreV1EnvVar
|
||||
from kubeflow_trainer_api.models.io_k8s_api_core_v1_volume_mount import IoK8sApiCoreV1VolumeMount
|
||||
from typing import Optional, Set
|
||||
@@ -27,15 +26,12 @@ from typing_extensions import Self
|
||||
|
||||
class TrainerV1alpha1ContainerOverride(BaseModel):
|
||||
"""
|
||||
ContainerOverride represents parameters that can be overridden using PodSpecOverrides. Parameters from the Trainer, DatasetConfig, and ModelConfig will take precedence.
|
||||
ContainerOverride represents parameters that can be overridden using PodSpecOverrides.
|
||||
""" # noqa: E501
|
||||
args: Optional[List[StrictStr]] = Field(default=None, description="Arguments to the entrypoint for the training container.")
|
||||
command: Optional[List[StrictStr]] = Field(default=None, description="Entrypoint commands for the training container.")
|
||||
env: Optional[List[IoK8sApiCoreV1EnvVar]] = Field(default=None, description="List of environment variables to set in the container. These values will be merged with the TrainingRuntime's environments.")
|
||||
env_from: Optional[List[IoK8sApiCoreV1EnvFromSource]] = Field(default=None, description="List of sources to populate environment variables in the container. These values will be merged with the TrainingRuntime's environments.", alias="envFrom")
|
||||
env: Optional[List[IoK8sApiCoreV1EnvVar]] = Field(default=None, description="List of environment variables to set in the container. These values will be merged with the TrainingRuntime's environments. These values can't be set for container with the name: `node`, `dataset-initializer`, or `model-initializer`. For those containers the envs can only be set via Trainer or Initializer APIs.")
|
||||
name: StrictStr = Field(description="Name for the container. TrainingRuntime must have this container.")
|
||||
volume_mounts: Optional[List[IoK8sApiCoreV1VolumeMount]] = Field(default=None, description="Pod volumes to mount into the container's filesystem.", alias="volumeMounts")
|
||||
__properties: ClassVar[List[str]] = ["args", "command", "env", "envFrom", "name", "volumeMounts"]
|
||||
__properties: ClassVar[List[str]] = ["env", "name", "volumeMounts"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
populate_by_name=True,
|
||||
@@ -83,13 +79,6 @@ class TrainerV1alpha1ContainerOverride(BaseModel):
|
||||
if _item_env:
|
||||
_items.append(_item_env.to_dict())
|
||||
_dict['env'] = _items
|
||||
# override the default output from pydantic by calling `to_dict()` of each item in env_from (list)
|
||||
_items = []
|
||||
if self.env_from:
|
||||
for _item_env_from in self.env_from:
|
||||
if _item_env_from:
|
||||
_items.append(_item_env_from.to_dict())
|
||||
_dict['envFrom'] = _items
|
||||
# override the default output from pydantic by calling `to_dict()` of each item in volume_mounts (list)
|
||||
_items = []
|
||||
if self.volume_mounts:
|
||||
@@ -109,10 +98,7 @@ class TrainerV1alpha1ContainerOverride(BaseModel):
|
||||
return cls.model_validate(obj)
|
||||
|
||||
_obj = cls.model_validate({
|
||||
"args": obj.get("args"),
|
||||
"command": obj.get("command"),
|
||||
"env": [IoK8sApiCoreV1EnvVar.from_dict(_item) for _item in obj["env"]] if obj.get("env") is not None else None,
|
||||
"envFrom": [IoK8sApiCoreV1EnvFromSource.from_dict(_item) for _item in obj["envFrom"]] if obj.get("envFrom") is not None else None,
|
||||
"name": obj.get("name") if obj.get("name") is not None else '',
|
||||
"volumeMounts": [IoK8sApiCoreV1VolumeMount.from_dict(_item) for _item in obj["volumeMounts"]] if obj.get("volumeMounts") is not None else None
|
||||
})
|
||||
|
||||
@@ -32,11 +32,11 @@ class TrainerV1alpha1PodSpecOverride(BaseModel):
|
||||
""" # 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")
|
||||
node_selector: Optional[Dict[str, StrictStr]] = Field(default=None, description="Override for the node selector to place Pod on the specific mode.", alias="nodeSelector")
|
||||
node_selector: Optional[Dict[str, StrictStr]] = Field(default=None, description="Override for the node selector to place Pod on the specific node.", alias="nodeSelector")
|
||||
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 configuration.")
|
||||
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"]
|
||||
|
||||
model_config = ConfigDict(
|
||||
|
||||
Reference in New Issue
Block a user