Skip to content

Commit 2c6642c

Browse files
authored
Add environment variables for Pathways metrics collection (#404)
Metrics collection is not enabled yet. When we enable the feature, these extra env variables will be used in collecting metrics for Pathways jobs. Testing / Documentation [ y ] Training sucecssful on real TPUs.
1 parent 3e8748e commit 2c6642c

File tree

1 file changed

+49
-0
lines changed

1 file changed

+49
-0
lines changed

src/xpk/commands/workload.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,22 @@
300300
name: shared-tmp
301301
{storage_volume_mounts}
302302
env:
303+
- name: PROJECT_ID
304+
value: {args.project}
305+
- name: LOCATION
306+
value: {args.zone}
307+
- name: CLUSTER_NAME
308+
value: {args.cluster}
309+
- name: POD_NAME
310+
valueFrom:
311+
fieldRef:
312+
fieldPath: metadata.name
313+
- name: CONTAINER_NAME
314+
value: "pathways-worker"
315+
- name: NAMESPACE
316+
valueFrom:
317+
fieldRef:
318+
fieldPath: metadata.namespace
303319
# Workaround for v6e
304320
- name: MEGASCALE_GRPC_ENABLE_XOR_TRACER
305321
value: "false"
@@ -351,6 +367,22 @@
351367
- args:
352368
{pathways_rm_args}
353369
env:
370+
- name: PROJECT_ID
371+
value: {args.project}
372+
- name: LOCATION
373+
value: {args.zone}
374+
- name: CLUSTER_NAME
375+
value: {args.cluster}
376+
- name: POD_NAME
377+
valueFrom:
378+
fieldRef:
379+
fieldPath: metadata.name
380+
- name: CONTAINER_NAME
381+
value: "pathways-rm"
382+
- name: NAMESPACE
383+
valueFrom:
384+
fieldRef:
385+
fieldPath: metadata.namespace
354386
- name: REPLICATED_JOB_NAME
355387
valueFrom:
356388
fieldRef:
@@ -397,6 +429,23 @@
397429
containers:
398430
- args:
399431
{pathways_proxy_args}
432+
env:
433+
- name: PROJECT_ID
434+
value: {args.project}
435+
- name: LOCATION
436+
value: {args.zone}
437+
- name: CLUSTER_NAME
438+
value: {args.cluster}
439+
- name: POD_NAME
440+
valueFrom:
441+
fieldRef:
442+
fieldPath: metadata.name
443+
- name: CONTAINER_NAME
444+
value: "pathways-proxy"
445+
- name: NAMESPACE
446+
valueFrom:
447+
fieldRef:
448+
fieldPath: metadata.namespace
400449
image: {args.proxy_server_image}
401450
imagePullPolicy: Always
402451
name: pathways-proxy

0 commit comments

Comments
 (0)