Apply changes to MaxText

GoogleCloudPlatform · Jan 12, 2024 · 4c06ea7 · 4c06ea7
1 parent 24967e5
commit 4c06ea7
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 8 deletions.
diff --git a/dags/multipod/configs/__init__.py b/dags/multipod/configs/__init__.py
diff --git a/configs/maxtext/common.py → dags/multipod/configs/common.py b/configs/maxtext/common.py → dags/multipod/configs/common.py
diff --git a/configs/maxtext/maxtext_gce_config.py → dags/multipod/configs/maxtext_gce_config.py b/configs/maxtext/maxtext_gce_config.py → dags/multipod/configs/maxtext_gce_config.py
@@ -14,12 +14,10 @@
 
 """Utilities to construct configs for maxtext DAG."""
 
-from typing import Tuple
-import uuid
-from apis import gcp_config, metric_config, task, test_config
-from configs import gcs_bucket, test_owner
-from configs.maxtext import common
-from configs.vm_resource import TpuVersion, Project, RuntimeVersion
+from xlml.apis import gcp_config, metric_config, task, test_config
+from dags import test_owner
+from dags.multipod.configs import common
+from dags.vm_resource import TpuVersion, Project, RuntimeVersion
 import datetime
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value

diff --git a/dags/maxtext/maxtext_nightly.py → dags/multipod/maxtext_nightly.py b/dags/maxtext/maxtext_nightly.py → dags/multipod/maxtext_nightly.py
@@ -17,8 +17,8 @@
 import datetime
 from airflow import models
 from configs import composer_env
-from configs.vm_resource import Project, TpuVersion, Zone, RuntimeVersion, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS
-from configs.maxtext import maxtext_gce_config
+from dags.vm_resource import TpuVersion, Zone
+from dags.multipod.configs import maxtext_gce_config
 
 
 # Run once a day at 2 am UTC (6 pm PST)