Skip to content

fix: Run pw_mcjax_benchmark_recipe on a smaller cluster #1176

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

Merged
alfredyu-cienet merged 1 commit into from
Feb 9, 2026
Merged

fix: Run pw_mcjax_benchmark_recipe on a smaller cluster #1176

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
31 changes: 23 additions & 8 deletions dags/maxtext_pathways/configs/parameters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""Configuration parameters for MaxText Pathways DAGs."""

from airflow.models.param import Param
from dags.maxtext_pathways.configs import model_configs as model_cfg
from dags.common.vm_resource import TpuVersion
Expand All @@ -35,10 +37,13 @@
"root",
type="string",
title="User",
description="User name is used to confirm the first three characters of the pod in the cluster.",
description=(
"User name is used to confirm the first three characters "
"of the pod in the cluster."
),
),
"cluster_name": Param(
"pw-v6e-64",
"pw-v6e-8",
type="string",
title="Cluster Name",
description="GCP cluster name for training model.",
Expand All @@ -63,7 +68,7 @@
enum=DEVICE_VERSION,
),
"core_count": Param(
64,
8,
type="integer",
title="Core Count",
description='Device core count for the cluster. ex: v6e-"64"',
Expand Down Expand Up @@ -118,14 +123,23 @@
"default_basic_1", # Available model for device type 'v6e'
type="string",
title="Model Name",
description='Select a model name to run. Only when "customized_model_name" is selected for "Model Name", the input value of "customized_model_name" parameter will take effect.',
description=(
"Select a model name to run. Only when 'customized_model_name' "
"is selected for 'Model Name', the input value of "
"'customized_model_name' parameter will take effect."
),
enum=["customized_model_name"] + MODEL_NAME,
),
"customized_model_name": Param(
None,
type=["null", "string"],
title="Customized Model Name",
description='Select a customized model name to run. Only when "customized_model_name" is selected for "Model Name", the input value of "customized_model_name" parameter will take effect.',
description=(
"Select a customized model name to run. Only when "
"'customized_model_name' is selected for 'Model Name', "
"the input value of 'customized_model_name' parameter "
"will take effect."
),
),
"priority": Param(
"medium",
Expand Down Expand Up @@ -164,9 +178,10 @@
type=["null", "integer"],
title="Override Timeout In Minutes",
description=(
"Timeout in minutes for the workload task. "
"Adjust it when your meet (airflow.exceptions.AirflowException: Timed out after ...) issue. "
"The default value `None` means automatic calculation of timeout."
"Timeout in minutes for the workload task. Adjust it when you "
"meet (airflow.exceptions.AirflowException: Timed out after ...) "
"issue. The default value `None` means automatic calculation "
"of timeout."
),
),
}
Loading