diff --git a/tunix/experimental/distributed/deployment/yaml_generator.py b/tunix/experimental/distributed/deployment/yaml_generator.py index 738746ede..775e9073a 100644 --- a/tunix/experimental/distributed/deployment/yaml_generator.py +++ b/tunix/experimental/distributed/deployment/yaml_generator.py @@ -89,6 +89,16 @@ def main() -> None: default="sleep infinity", help="Command to run on startup", ) + parser.add_argument( + "--otel_exporter_otlp_endpoint", + default="http://$(NODE_IP):4317", + help="OpenTelemetry Collector OTLP gRPC endpoint URL", + ) + parser.add_argument( + "--otel_exporter_otlp_protocol", + default="grpc", + help="OpenTelemetry Collector OTLP protocol (grpc or http/protobuf)", + ) args = parser.parse_args() @@ -164,6 +174,8 @@ def main() -> None: USER_CONTAINER_IMAGE=args.worker_container_image, USER_CONTAINER_PORT=args.worker_container_port, STARTUP_COMMAND=args.worker_startup_command, + OTEL_EXPORTER_OTLP_ENDPOINT=args.otel_exporter_otlp_endpoint, + OTEL_EXPORTER_OTLP_PROTOCOL=args.otel_exporter_otlp_protocol, ) print(content) diff --git a/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml b/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml index 2ab0562b3..f9598b003 100644 --- a/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml +++ b/tunix/experimental/distributed/deployment/yamls/jobset.cpu.yaml @@ -102,6 +102,19 @@ spec: valueFrom: fieldRef: fieldPath: metadata.labels['batch.kubernetes.io/job-completion-index'] + # OpenTelemetry Telemetry Configuration + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OTEL_EXPORTER_OTLP_ENDPOINT + value: "${OTEL_EXPORTER_OTLP_ENDPOINT}" + - name: OTEL_EXPORTER_OTLP_PROTOCOL + value: "${OTEL_EXPORTER_OTLP_PROTOCOL}" + - name: OTEL_SERVICE_NAME + value: "${JOBSET_NAME}-orchestrator" + - name: OTEL_RESOURCE_ATTRIBUTES + value: "service.namespace=tunix,jobset.name=${JOBSET_NAME},container.name=${USER_CONTAINER}" ports: - containerPort: ${USER_CONTAINER_PORT} name: grpc # Naming it 'grpc' tells GKE to optimize for HTTP/2 diff --git a/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml b/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml index 7dcc2a1fa..9eb2717e0 100644 --- a/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml +++ b/tunix/experimental/distributed/deployment/yamls/jobset.pathways.yaml @@ -90,6 +90,19 @@ spec: fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - name: TPU_SKIP_MDS_QUERY value: "true" + # OpenTelemetry Configuration + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OTEL_EXPORTER_OTLP_ENDPOINT + value: "${OTEL_EXPORTER_OTLP_ENDPOINT}" + - name: OTEL_EXPORTER_OTLP_PROTOCOL + value: "${OTEL_EXPORTER_OTLP_PROTOCOL}" + - name: OTEL_SERVICE_NAME + value: "${JOBSET_NAME}-pathways-rm" + - name: OTEL_RESOURCE_ATTRIBUTES + value: "service.namespace=tunix,jobset.name=${JOBSET_NAME},container.name=pathways-rm" ports: - containerPort: 29001 protocol: TCP @@ -113,6 +126,19 @@ spec: valueFrom: fieldRef: fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + # OpenTelemetry Configuration + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OTEL_EXPORTER_OTLP_ENDPOINT + value: "${OTEL_EXPORTER_OTLP_ENDPOINT}" + - name: OTEL_EXPORTER_OTLP_PROTOCOL + value: "${OTEL_EXPORTER_OTLP_PROTOCOL}" + - name: OTEL_SERVICE_NAME + value: "${JOBSET_NAME}-pathways-proxy" + - name: OTEL_RESOURCE_ATTRIBUTES + value: "service.namespace=tunix,jobset.name=${JOBSET_NAME},container.name=pathways-proxy" ports: - containerPort: 29000 protocol: TCP @@ -177,6 +203,19 @@ spec: value: "grpc://localhost:29000" - name: VLLM_ENABLE_V1_MULTIPROCESSING value: "0" + # OpenTelemetry Configuration + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OTEL_EXPORTER_OTLP_ENDPOINT + value: "${OTEL_EXPORTER_OTLP_ENDPOINT}" + - name: OTEL_EXPORTER_OTLP_PROTOCOL + value: "${OTEL_EXPORTER_OTLP_PROTOCOL}" + - name: OTEL_SERVICE_NAME + value: "${JOBSET_NAME}-${USER_CONTAINER}" + - name: OTEL_RESOURCE_ATTRIBUTES + value: "service.namespace=tunix,jobset.name=${JOBSET_NAME},container.name=${USER_CONTAINER}" ports: - containerPort: ${USER_CONTAINER_PORT} name: grpc # Naming it 'grpc' tells GKE to optimize for HTTP/2 @@ -279,6 +318,19 @@ spec: valueFrom: fieldRef: fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + # OpenTelemetry Configuration + - name: NODE_IP + valueFrom: + fieldRef: + fieldPath: status.hostIP + - name: OTEL_EXPORTER_OTLP_ENDPOINT + value: "${OTEL_EXPORTER_OTLP_ENDPOINT}" + - name: OTEL_EXPORTER_OTLP_PROTOCOL + value: "${OTEL_EXPORTER_OTLP_PROTOCOL}" + - name: OTEL_SERVICE_NAME + value: "${JOBSET_NAME}-pathways-worker" + - name: OTEL_RESOURCE_ATTRIBUTES + value: "service.namespace=tunix,jobset.name=${JOBSET_NAME},container.name=pathways-worker" ports: - containerPort: 29005 protocol: TCP diff --git a/tunix/experimental/distributed/examples/rl/launcher.sh b/tunix/experimental/distributed/examples/rl/launcher.sh index 0e825cde4..1973c21e0 100755 --- a/tunix/experimental/distributed/examples/rl/launcher.sh +++ b/tunix/experimental/distributed/examples/rl/launcher.sh @@ -22,6 +22,8 @@ LOCAL_MODE=false ROLE="" # set by --image TUNIX_IMAGE="us-central1-docker.pkg.dev/cloud-tpu-multipod-dev/yangmu/tunix/tunix_base_image:latest" +# set by --otel_endpoint +OTEL_ENDPOINT='http://$(NODE_IP):4317' while [[ $# -gt 0 ]]; do case "$1" in @@ -45,6 +47,14 @@ while [[ $# -gt 0 ]]; do TUNIX_IMAGE="${1#*=}" shift ;; + --otel_endpoint) + OTEL_ENDPOINT="$2" + shift 2 + ;; + --otel_endpoint=*) + OTEL_ENDPOINT="${1#*=}" + shift + ;; *) shift ;; @@ -93,6 +103,7 @@ launch_orchestrator() { --cpu_machine=n2-standard-64 \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=12345 \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_id=orchestrator \ --discovery_port=12345 \ @@ -124,6 +135,7 @@ launch_rollout() { --tpu_slice=tpuv5e:4x4 \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=$((10000+i)) \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_addrs=orchestrator:12345 \ --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \ @@ -152,6 +164,7 @@ launch_trainer() { --tpu_slice=tpuv5e:4x4 \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=20000 \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_addrs=orchestrator:12345 \ --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \ diff --git a/tunix/experimental/distributed/examples/vllm_rollout/launcher.sh b/tunix/experimental/distributed/examples/vllm_rollout/launcher.sh index ad10c4782..7367c41ed 100755 --- a/tunix/experimental/distributed/examples/vllm_rollout/launcher.sh +++ b/tunix/experimental/distributed/examples/vllm_rollout/launcher.sh @@ -24,6 +24,8 @@ ROLE="" WORKER_ID="all" # Docker image URI used for worker containers on Kubernetes. TUNIX_IMAGE="us-central1-docker.pkg.dev/cloud-tpu-multipod-dev/yangmu/tunix/tunix_base_image:latest" +# set by --otel_endpoint +OTEL_ENDPOINT='http://$(NODE_IP):4317' while [[ $# -gt 0 ]]; do case "$1" in @@ -55,6 +57,14 @@ while [[ $# -gt 0 ]]; do TUNIX_IMAGE="${1#*=}" shift ;; + --otel_endpoint) + OTEL_ENDPOINT="$2" + shift 2 + ;; + --otel_endpoint=*) + OTEL_ENDPOINT="${1#*=}" + shift + ;; *) shift ;; @@ -104,6 +114,7 @@ launch_orchestrator() { --cpu_machine=n2-standard-64 \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=12345 \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_id=orchestrator \ --discovery_port=12345 \ @@ -136,6 +147,7 @@ launch_rollout() { --pathways_proxy_server_image=us-central1-docker.pkg.dev/cloud-tpu-multipod-dev/yangmu/tunix/unsanitized_proxy_server:latest \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=$((10000+i)) \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_addrs=orchestrator:12345 \ --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \ @@ -154,6 +166,7 @@ launch_rollout() { --pathways_proxy_server_image=us-central1-docker.pkg.dev/cloud-tpu-multipod-dev/yangmu/tunix/unsanitized_proxy_server:latest \ --worker_container_image="$TUNIX_IMAGE" \ --worker_container_port=11111 \ + --otel_exporter_otlp_endpoint="$OTEL_ENDPOINT" \ --worker_startup_command="python -m tunix.experimental.distributed.runtime.main \ --discovery_addrs=orchestrator:12345 \ --process_executor=tunix.experimental.distributed.runtime.executor.K8sExecutor \