Skip to content

Commit cf04e12

Browse files
author
maxtext authors
committed
Merge pull request #1682 from AI-Hypercomputer:sujinesh/update_pathways_image_paths
PiperOrigin-RevId: 754149868
2 parents c340436 + 973a2b8 commit cf04e12

8 files changed

Lines changed: 22 additions & 23 deletions

‎benchmarks/Getting_Started_Benchmarking.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ python3 -m benchmarks.benchmark_runner xpk --project $PROJECT --zone $ZONE --clu
1919

2020
```shell
2121
# Pathways with XPK
22-
export RUNNER=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable
23-
export PROXY_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/proxy_server:latest
24-
export SERVER_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/server:latest
22+
export RUNNER=us-docker.pkg.dev/path/to/maxtext_runner
23+
export PROXY_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server
24+
export SERVER_IMAGE=us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server
2525

2626
python3 -m benchmarks.benchmark_runner xpk --project $PROJECT --zone $ZONE --cluster_name $CLUSTER --device_type v6e-256 --base_output_directory gs://maxtext-experiments-tpem/ --num_steps=5 --pathways_server_image="${SERVER_IMAGE}" --pathways_proxy_server_image="${PROXY_IMAGE}" --pathways_runner_image="${RUNNER}"
2727
```

‎benchmarks/benchmark_runner.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,20 +47,19 @@ def add_pathways_arguments(parser: argparse.ArgumentParser):
4747
'--pathways_server_image',
4848
type=str,
4949
default=(
50-
'us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/server:latest'
50+
'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server'
5151
),
5252
help='version of pathways server image to be benchmarked command.',
5353
)
5454
parser.add_argument(
5555
'--pathways_proxy_server_image',
5656
type=str,
57-
default='us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/proxy_server:latest',
57+
default='us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server',
5858
help='version of pathways proxy image to be benchmarked command.',
5959
)
6060
parser.add_argument(
6161
'--pathways_runner_image',
6262
type=str,
63-
default='us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest',
6463
help='version of pathways runner image to be benchmarked command.',
6564
)
6665
parser.add_argument(

‎benchmarks/recipes/pw_elastic_training_recipe.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@
3131
import maxtext_xpk_runner as mxr
3232
from xpk_configs import XpkClusterConfig
3333

34-
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_proxy_server:latest"
35-
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_server:latest"
36-
RUNNER = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest"
34+
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
35+
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
36+
RUNNER = "us-docker.pkg.dev/path/to/maxtext_runner"
3737

3838
# Cluster Params
3939
CLUSTER = "v6e-256-cluster"

‎benchmarks/recipes/pw_long_running_recipe.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,9 @@
2626
import maxtext_xpk_runner as mxr
2727
from xpk_configs import XpkClusterConfig
2828

29-
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_proxy_server:latest"
30-
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_server:latest"
31-
RUNNER = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest"
29+
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
30+
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
31+
RUNNER = "us-docker.pkg.dev/path/to/maxtext_runner"
3232

3333
# Cluster Params
3434
CLUSTER = "v6e-256-cluster"

‎benchmarks/recipes/pw_mcjax_benchmark_recipe.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@
2929
from xpk_configs import XpkClusterConfig
3030

3131

32-
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_proxy_server:latest"
33-
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_server:latest"
34-
RUNNER = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest"
32+
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
33+
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
34+
RUNNER = "us-docker.pkg.dev/path/to/maxtext_runner"
3535

3636
# Cluster Params
3737
CLUSTER = "v6e-256-cluster"

‎benchmarks/recipes/pw_mcjax_checkpoint_benchmark_recipe.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@
3434
import maxtext_xpk_runner as mxr
3535
from xpk_configs import XpkClusterConfig
3636

37-
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_proxy_server:latest"
38-
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_server:latest"
39-
RUNNER = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest"
37+
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
38+
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
39+
RUNNER = "us-docker.pkg.dev/path/to/maxtext_runner"
4040

4141
# Cluster Params
4242
CLUSTER = "v6e-256-cluster"

‎benchmarks/recipes/pw_remote_python_recipe.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,11 @@ def main() -> int:
4848
user = os.environ["USER"]
4949
region = "-".join(cluster_config.zone.split("-")[:-1])
5050
proxy_image = (
51-
f"us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/{user}/"
51+
f"us-docker.pkg.dev/cloud-tpu-v2-images/pathways/gke/{user}/"
5252
"proxy_server:latest"
5353
)
5454
server_image = (
55-
f"us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/{user}/"
55+
f"us-docker.pkg.dev/cloud-tpu-v2-images/pathways/gke/{user}/"
5656
"server:latest"
5757
)
5858
colocated_python_image = f"gcr.io/{cluster_config.project}/{user}/colocated_python_sidecar_latest:latest"

‎benchmarks/recipes/pw_suspend_resume.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,9 @@
3030
import maxtext_xpk_runner as mxr
3131
from xpk_configs import XpkClusterConfig
3232

33-
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_proxy_server:latest"
34-
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/unsanitized_server:latest"
35-
RUNNER = "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest"
33+
PROXY_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
34+
SERVER_IMAGE = "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
35+
RUNNER = "us-docker.pkg.dev/path/to/maxtext_runner"
3636

3737
# Cluster Params
3838
CLUSTER = "v6e-256-cluster"

0 commit comments

Comments
 (0)