Add OpenTelemetry (OTel) endpoint injection to Tunix deployment pipeline.

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)
 

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

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

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 \

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 \