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(runtimes): Add JAX training runtime (#3151)

Browse Source 
Signed-off-by: mahdikhashan <[email protected]>
Signed-off-by: Tomas Tormo <[email protected]>
This commit is contained in:
Tomas Tormo
2026-02-06 18:15:22 +00:00
committed by GitHub
parent 19dd9e4ae7
commit 62812f9ce3
28 changed files with 1624 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_ml_policy.py
Generated
+3 -1
View File
@@ -28,10 +28,11 @@ class TrainerV1alpha1MLPolicy(BaseModel):
"""
MLPolicy represents configuration for the model training with ML-specific parameters.
""" # noqa: E501
jax: Optional[Dict[str, Any]] = Field(default=None, description="jax defines the configuration for the JAX Runtime")
mpi: Optional[TrainerV1alpha1MPIMLPolicySource] = Field(default=None, description="mpi defines the configuration for the MPI Runtime.")
num_nodes: Optional[StrictInt] = Field(default=None, description="numNodes is the number of training nodes. Defaults to 1.", alias="numNodes")
torch: Optional[TrainerV1alpha1TorchMLPolicySource] = Field(default=None, description="torch defines the configuration for the PyTorch runtime.")
__properties: ClassVar[List[str]] = ["mpi", "numNodes", "torch"]
__properties: ClassVar[List[str]] = ["jax", "mpi", "numNodes", "torch"]
model_config = ConfigDict(
populate_by_name=True,
@@ -90,6 +91,7 @@ class TrainerV1alpha1MLPolicy(BaseModel):
return cls.model_validate(obj)
_obj = cls.model_validate({
"jax": obj.get("jax"),
"mpi": TrainerV1alpha1MPIMLPolicySource.from_dict(obj["mpi"]) if obj.get("mpi") is not None else None,
"numNodes": obj.get("numNodes"),
"torch": TrainerV1alpha1TorchMLPolicySource.from_dict(obj["torch"]) if obj.get("torch") is not None else None
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_ml_policy_source.py
Generated
+3 -1
View File
@@ -28,9 +28,10 @@ class TrainerV1alpha1MLPolicySource(BaseModel):
"""
MLPolicySource represents the runtime-specific configuration for various technologies. One of the following specs can be set.
""" # noqa: E501
jax: Optional[Dict[str, Any]] = Field(default=None, description="jax defines the configuration for the JAX Runtime")
mpi: Optional[TrainerV1alpha1MPIMLPolicySource] = Field(default=None, description="mpi defines the configuration for the MPI Runtime.")
torch: Optional[TrainerV1alpha1TorchMLPolicySource] = Field(default=None, description="torch defines the configuration for the PyTorch runtime.")
__properties: ClassVar[List[str]] = ["mpi", "torch"]
__properties: ClassVar[List[str]] = ["jax", "mpi", "torch"]
model_config = ConfigDict(
populate_by_name=True,
@@ -89,6 +90,7 @@ class TrainerV1alpha1MLPolicySource(BaseModel):
return cls.model_validate(obj)
_obj = cls.model_validate({
"jax": obj.get("jax"),
"mpi": TrainerV1alpha1MPIMLPolicySource.from_dict(obj["mpi"]) if obj.get("mpi") is not None else None,
"torch": TrainerV1alpha1TorchMLPolicySource.from_dict(obj["torch"]) if obj.get("torch") is not None else None
})