mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
Tariq HasanandGitHub
1f57df4484
* ci(boilerplate): enforce copyright headers on YAML files Signed-off-by: tariq-hasan <[email protected]> * ci(boilerplate): enforce year-less header for new files Signed-off-by: tariq-hasan <[email protected]> --------- Signed-off-by: tariq-hasan <[email protected]>
107 lines
3.4 KiB
YAML
107 lines
3.4 KiB
YAML
# Copyright The Kubeflow Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Example 2: Runtime Patches
|
|
#
|
|
# Pod customization with the runtimePatches API. Each patch is keyed by a
|
|
# "manager" so users, Kueue, and admission webhooks can each own one entry
|
|
# without conflicting.
|
|
#
|
|
# Notes:
|
|
# - Resources for training nodes are set via spec.trainer.resourcesPerNode.
|
|
# - Env vars for the "node" container go in spec.trainer.env
|
|
# (RuntimePatches container env is forbidden for the "node",
|
|
# "dataset-initializer", and "model-initializer" containers).
|
|
# - runtimePatches expose JobSet metadata, Job metadata, Pod metadata, and
|
|
# a restricted set of Pod spec fields (serviceAccountName, volumes,
|
|
# nodeSelector, affinity, tolerations, etc.).
|
|
#
|
|
# Apply: kubectl apply -f 02-runtime-patches.yaml
|
|
# Check: kubectl describe trainjob runtime-patches-example
|
|
# Pods: kubectl get pods -l jobset.sigs.k8s.io/jobset-name=runtime-patches-example
|
|
# Clean: kubectl delete trainjob runtime-patches-example
|
|
|
|
apiVersion: trainer.kubeflow.org/v1alpha1
|
|
kind: TrainJob
|
|
metadata:
|
|
name: runtime-patches-example
|
|
spec:
|
|
runtimeRef:
|
|
apiGroup: trainer.kubeflow.org
|
|
kind: ClusterTrainingRuntime
|
|
name: torch-distributed
|
|
|
|
trainer:
|
|
numNodes: 2
|
|
command:
|
|
- torchrun
|
|
- --no-python
|
|
- python3
|
|
- -c
|
|
- |
|
|
import os
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
device, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")
|
|
dist.init_process_group(backend=backend)
|
|
|
|
print("PyTorch Distributed Environment")
|
|
print(f"Using device: {device}")
|
|
print(f"WORLD_SIZE: {dist.get_world_size()}")
|
|
print(f"RANK: {dist.get_rank()}")
|
|
print(f"LOCAL_RANK: {os.environ['LOCAL_RANK']}")
|
|
print(f"CUSTOM_VAR: {os.getenv('CUSTOM_VAR')}")
|
|
|
|
env:
|
|
- name: CUSTOM_VAR
|
|
value: "custom-value"
|
|
- name: PYTHONUNBUFFERED
|
|
value: "1"
|
|
|
|
resourcesPerNode:
|
|
requests:
|
|
cpu: "2"
|
|
memory: "4Gi"
|
|
limits:
|
|
cpu: "4"
|
|
memory: "8Gi"
|
|
|
|
runtimePatches:
|
|
- manager: example.io/platform-team
|
|
trainingRuntimeSpec:
|
|
template:
|
|
metadata:
|
|
labels:
|
|
team: ml-platform
|
|
cost-center: research
|
|
annotations:
|
|
prometheus.io/scrape: "true"
|
|
prometheus.io/port: "8080"
|
|
spec:
|
|
replicatedJobs:
|
|
- name: node
|
|
template:
|
|
spec:
|
|
template:
|
|
spec:
|
|
serviceAccountName: training-sa
|
|
nodeSelector:
|
|
node-type: training
|
|
tolerations:
|
|
- key: training-workload
|
|
operator: Equal
|
|
value: "true"
|
|
effect: NoSchedule
|