Skip to content

Fix TF-PjRt tests #980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chandrasekhard2 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Fix TF-PjRt tests #980

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
82e8764
Fix pjrt tests
chandrasekhard2 Sep 2, 2023
b17f271
Update tpu health agent
chandrasekhard2 Sep 8, 2023
58ede58
Since PjRt is functional, opening up all the tests
chandrasekhard2 Sep 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tests/tensorflow/experimental.libsonnet
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ local mixins = import 'templates/mixins.libsonnet';
tpuSettings+: {
tpuVmEnvVars+: {
PYTHONPATH: '${PWD}',
WRAPT_DISABLE_EXTENSIONS: 'true',
} + if config.accelerator.replicas > 1 then {
TPU_LOAD_LIBRARY: '0',
} else {},
Expand Down
12 changes: 8 additions & 4 deletions tests/tensorflow/nightly/common.libsonnet
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,10 @@ local volumes = import 'templates/volumes.libsonnet';
tpuVm:: experimental.TensorFlowTpuVmMixin {
local config = self,
tpuSettings+: {
tpuVmEnvVars+: {
TF_PLUGGABLE_DEVICE_LIBRARY_PATH: '/lib/libtpu.so',
NEXT_PLUGGABLE_DEVICE_USE_C_API: 'true',
},
softwareVersion: if config.accelerator.replicas == 1 then
'v2-nightly'
else
Expand Down Expand Up @@ -178,6 +182,9 @@ local volumes = import 'templates/volumes.libsonnet';

softwareVersion=%(softwareVersion)s
gcloud alpha compute tpus tpu-vm ssh ${tpu_name} --zone=${zone} --project=${project} --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "echo 'WRAPT_DISABLE_EXTENSIONS=true' | sudo tee -a /etc/environment"
gcloud alpha compute tpus tpu-vm ssh ${tpu_name} --zone=${zone} --project=${project} --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command 'sudo sed -i "/HEALTH_AGENT_DOCKER_URL/c\HEALTH_AGENT_DOCKER_URL=\"gcr.io/cloud-tpu-v2-images/tpu_agents:cl_562025307\"" /home/tpu-runtime/tpu-env'

gcloud alpha compute tpus tpu-vm ssh ${tpu_name} --zone=${zone} --project=${project} --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo systemctl daemon-reload && sudo systemctl restart healthagent.service"
if [[ ${softwareVersion: -3} == "pod" ]]; then
gcloud alpha compute tpus tpu-vm ssh ${tpu_name} --zone=${zone} --project=${project} --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images-dev\/grpc_tpu_worker:nightly\"/' /etc/systemd/system/tpu-runtime.service"
gcloud alpha compute tpus tpu-vm ssh ${tpu_name} --zone=${zone} --project=${project} --internal-ip --ssh-key-file=/scripts/id_rsa --worker=all --command "sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime"
Expand Down Expand Up @@ -295,10 +302,7 @@ local volumes = import 'templates/volumes.libsonnet';
},
local functional_schedule = '0 9 * * *',
Functional:: mixins.Functional {
schedule: if !(self.accelerator.type == 'tpu') || self.accelerator.name == 'v3-8' || self.accelerator.name == 'v4-8' then
functional_schedule
else
null,
schedule: functional_schedule,
metricConfig+: {
sourceMap+:: {
tensorboard+: {
Expand Down
14 changes: 7 additions & 7 deletions tests/tensorflow/nightly/targets.jsonnet
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,13 @@ local wmt = import 'tf-wmt-wmt14_translate.libsonnet';

// Add new models here
std.flattenArrays([
// dlrm.configs,
dlrm.configs,
keras_api.configs,
// bert.configs,
// wmt.configs,
// maskrcnn.configs,
// retinanet.configs,
bert.configs,
wmt.configs,
maskrcnn.configs,
retinanet.configs,
resnet.configs,
// resnetrs.configs,
// gpt2.configs,
resnetrs.configs,
gpt2.configs,
])