Update codeowner and common file path (#508)

GoogleCloudPlatform · Dec 18, 2024 · f061011 · f061011
1 parent 2b23fcd
commit f061011
Show file tree

Hide file tree

Showing 81 changed files with 169 additions and 132 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -1,20 +1,22 @@
 # Default owners for everything in the repo, unless a later match takes precedence.
-* @mbzomowski @RissyRan @allenwang28
+* @richardsliu @yixinshi @RissyRan @allenwang28
+
+dags/common @richardsliu @yixinshi @RissyRan @allenwang28 @zpcore @ManfeiBai @gobbleturk @shralex @jiangjy1982 @ortibazar @parambole @vipannalla @crankshaw-google @polydier1
 
 dags/solutions_team/configs/tensorflow @chandrasekhard2 @ZhaoyueCheng @richardsliu
 dags/solutions_team/solutionsteam_tf* @chandrasekhard2 @ZhaoyueCheng @richardsliu
 
-dags/pytorch_xla @JackCaoG @vanbasten23 @zpcore @ManfeiBai
-dags/legacy_test/tests/pytorch @JackCaoG @vanbasten23 @zpcore @ManfeiBai
+dags/pytorch_xla @vanbasten23 @zpcore @ManfeiBai
+dags/legacy_test/tests/pytorch @vanbasten23 @zpcore @ManfeiBai
 
-dags/multipod @jonb377 @tonyjohnchen @raymondzouu @gobbleturk @shralex @RissyRan @jiangjy1982
+dags/multipod @tonyjohnchen @raymondzouu @gobbleturk @shralex @RissyRan @jiangjy1982
 
 dags/mlcompass @ortibazar @sganeshb @brajiang @wlzhg
 
 dags/sparsity_diffusion_devx @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh
 dags/sparsity_diffusion_devx/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @jiya-zhang
 dags/sparsity_diffusion_devx/configs/project_bite* @RissyRan @parambole @jiangjy1982 @aireenmei @michelle-yooh @jiya-zhang
 
-dags/inference @yeandy @vipannalla @morgandu @mailvijayasingh @sixiang-google @joezijunzhou @singh-mitali
+dags/inference @vipannalla @mailvijayasingh @sixiang-google @joezijunzhou @singh-mitali
 
 dags/map_reproducibility @crankshaw-google @polydier1
diff --git a/dags/common/__init__.py b/dags/common/__init__.py
diff --git a/dags/quarantined_tests.py → dags/common/quarantined_tests.py b/dags/quarantined_tests.py → dags/common/quarantined_tests.py
@@ -15,7 +15,7 @@
 """Lists all currently broken tests."""
 
 import dataclasses
-from dags.test_owner import Team as team
+from dags.common.test_owner import Team as team
 
 
 @dataclasses.dataclass

diff --git a/dags/test_owner.py → dags/common/test_owner.py b/dags/test_owner.py → dags/common/test_owner.py
diff --git a/dags/vm_resource.py → dags/common/vm_resource.py b/dags/vm_resource.py → dags/common/vm_resource.py
diff --git a/dags/examples/configs/xpk_example_config.py b/dags/examples/configs/xpk_example_config.py
@@ -17,8 +17,8 @@
 import datetime
 from xlml.apis import gcp_config, metric_config, task, test_config
 from xlml.apis.xpk_cluster_config import XpkClusterConfig
-from dags import test_owner
-from dags.vm_resource import TpuVersion
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion
 
 
 def get_flax_resnet_xpk_config(

diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py
@@ -20,8 +20,8 @@
 
 import datetime
 from airflow import models
-from dags import test_owner
-from dags.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
 from dags.multipod.configs import maxtext_sweep_gke_config
 
 # Set concurrency to number of workers otherwise tasks may time out

diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py
@@ -20,8 +20,8 @@
 
 import datetime
 from airflow import models
-from dags import test_owner
-from dags.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage
 from dags.multipod.configs import maxtext_sweep_gke_config
 
 # Set concurrency to number of workers otherwise tasks may time out

diff --git a/dags/examples/xpk_example_dag.py b/dags/examples/xpk_example_dag.py
@@ -16,9 +16,9 @@
 
 import datetime
 from airflow import models
-from dags.vm_resource import TpuVersion, Project, Zone, XpkClusters, DockerImage
+from dags.common.vm_resource import TpuVersion, Project, Zone, XpkClusters, DockerImage
 from dags.examples.configs import xpk_example_config as config
-from dags import test_owner
+from dags.common import test_owner
 from xlml.utils import name_format
 from airflow.utils.task_group import TaskGroup
 

diff --git a/dags/framework3p/configs/microbenchmarks_config.py b/dags/framework3p/configs/microbenchmarks_config.py
@@ -1,7 +1,8 @@
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import gcs_bucket, test_owner
+from dags import gcs_bucket
 import datetime
-import dags.vm_resource as resource
+import dags.common.vm_resource as resource
 
 
 def get_microbenchmark_config(

diff --git a/dags/framework3p/microbenchmarks_dag.py b/dags/framework3p/microbenchmarks_dag.py
@@ -1,6 +1,7 @@
 import datetime
 from airflow import models
-from dags import composer_env, vm_resource, test_owner
+from dags import composer_env
+from dags.common import test_owner, vm_resource
 from dags.framework3p.configs.microbenchmarks_config import get_microbenchmark_config, get_microbenchmark_xpk_config
 
 

diff --git a/dags/inference/configs/jetstream_benchmark_serving_gce_config.py b/dags/inference/configs/jetstream_benchmark_serving_gce_config.py
@@ -18,9 +18,9 @@
 import json
 from typing import Dict
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value

diff --git a/dags/inference/configs/jetstream_pytorch_gce_config.py b/dags/inference/configs/jetstream_pytorch_gce_config.py
@@ -18,9 +18,9 @@
 import json
 from typing import Dict
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value

diff --git a/dags/inference/configs/maxtext_inference_microbenchmark_gce_config.py b/dags/inference/configs/maxtext_inference_microbenchmark_gce_config.py
@@ -18,9 +18,9 @@
 import json
 from typing import Dict
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value

diff --git a/dags/inference/configs/trt_llm_inference_config.py b/dags/inference/configs/trt_llm_inference_config.py
@@ -15,9 +15,10 @@
 """Utilities to construct configs for TensorRT-LLM inference DAG."""
 
 import datetime
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner, vm_resource
-from dags.vm_resource import Project, RuntimeVersion
+from dags.common import vm_resource
+from dags.common.vm_resource import Project, RuntimeVersion
 
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value
 GCS_SUBFOLDER_PREFIX = test_owner.Team.INFERENCE.value

diff --git a/dags/inference/configs/trt_llm_mlperf_v40_config.py b/dags/inference/configs/trt_llm_mlperf_v40_config.py
@@ -16,9 +16,10 @@
 
 import datetime
 from typing import Dict
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner, vm_resource
-from dags.vm_resource import Project, RuntimeVersion
+from dags.common import vm_resource
+from dags.common.vm_resource import Project, RuntimeVersion
 
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value
 GCS_SUBFOLDER_PREFIX = test_owner.Team.INFERENCE.value

diff --git a/dags/inference/configs/trt_llm_mlperf_v41_config.py b/dags/inference/configs/trt_llm_mlperf_v41_config.py
@@ -16,9 +16,10 @@
 
 import datetime
 from typing import Dict, List
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner, vm_resource
-from dags.vm_resource import GpuVersion, Project, RuntimeVersion
+from dags.common import vm_resource
+from dags.common.vm_resource import GpuVersion, Project, RuntimeVersion
 
 RUNTIME_IMAGE = RuntimeVersion.TPU_UBUNTU2204_BASE.value
 GCS_SUBFOLDER_PREFIX = test_owner.Team.INFERENCE.value

diff --git a/dags/inference/jetstream_inference_e2e.py b/dags/inference/jetstream_inference_e2e.py
@@ -16,7 +16,7 @@
 
 import datetime
 from airflow import models
-from dags.vm_resource import TpuVersion
+from dags.common.vm_resource import TpuVersion
 from dags.inference.maxtext_model_config_generator import generate_model_configs
 
 """A JetStream inference E2E test (JAX nightly, no schedule) DAG.

diff --git a/dags/inference/jetstream_pytorch_inference.py b/dags/inference/jetstream_pytorch_inference.py
@@ -3,8 +3,9 @@
 import datetime
 from airflow import models
 from airflow.models.baseoperator import chain
-from dags import composer_env, test_owner
-from dags.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
+from dags import composer_env
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
 from dags.inference.configs import jetstream_pytorch_gce_config
 from dags.multipod.configs.common import SetupMode, Platform
 import numpy as np

diff --git a/dags/inference/maxtext_inference.py b/dags/inference/maxtext_inference.py
@@ -18,7 +18,7 @@
 import numpy as np
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import TpuVersion
+from dags.common.vm_resource import TpuVersion
 from dags.inference.maxtext_model_config_generator import generate_model_configs
 
 USER_PREFIX = ""

diff --git a/dags/inference/maxtext_inference_microbenchmark.py b/dags/inference/maxtext_inference_microbenchmark.py
@@ -19,7 +19,7 @@
 import itertools
 import numpy
 from airflow import models
-from dags.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
+from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
 from dags.inference.configs import maxtext_inference_microbenchmark_gce_config
 from dags.multipod.configs.common import SetupMode
 

diff --git a/dags/inference/maxtext_inference_offline_benchmark.py b/dags/inference/maxtext_inference_offline_benchmark.py
@@ -17,8 +17,9 @@
 import datetime
 
 from airflow import models
-from dags import test_owner, composer_env
-from dags.vm_resource import TpuVersion, Zone, Project, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
+from dags import composer_env
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion, Zone, Project, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
 from dags.multipod.configs import common
 from dags.multipod.configs.common import SetupMode
 from xlml.apis import gcp_config, metric_config, task, test_config

diff --git a/dags/inference/maxtext_model_config_generator.py b/dags/inference/maxtext_model_config_generator.py
@@ -14,7 +14,7 @@
 
 """A helper to generate maxtext model configs."""
 
-from dags.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
+from dags.common.vm_resource import TpuVersion, Zone, Project, V5_NETWORKS, V5E_SUBNETWORKS, V5P_SUBNETWORKS, RuntimeVersion, V6E_GCE_NETWORK, V6E_GCE_SUBNETWORK
 from dags.inference.configs import jetstream_benchmark_serving_gce_config
 from dags.multipod.configs.common import SetupMode
 

diff --git a/dags/inference/trt_llm_inference.py b/dags/inference/trt_llm_inference.py
@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import H100_INFERENCE_SUBNETWORKS, INFERENCE_NETWORKS, GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project
+from dags.common.vm_resource import H100_INFERENCE_SUBNETWORKS, INFERENCE_NETWORKS, GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project
 from dags.inference.configs import trt_llm_inference_config
 
 # Run once a day at 4 am UTC (8 pm PST)

diff --git a/dags/inference/trt_llm_mlperf_v40_inference.py b/dags/inference/trt_llm_mlperf_v40_inference.py
@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project, INFERENCE_NETWORKS, H100_INFERENCE_SUBNETWORKS
+from dags.common.vm_resource import GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project, INFERENCE_NETWORKS, H100_INFERENCE_SUBNETWORKS
 from dags.inference.configs import trt_llm_mlperf_v40_config
 
 # Run once a day at 4 am UTC (8 pm PST)

diff --git a/dags/inference/trt_llm_mlperf_v41_inference.py b/dags/inference/trt_llm_mlperf_v41_inference.py
@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import A100_INFERENCE_SUBNETWORKS, H100_INFERENCE_SUBNETWORKS, GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project, INFERENCE_NETWORKS, L4_INFERENCE_SUBNETWORKS
+from dags.common.vm_resource import A100_INFERENCE_SUBNETWORKS, H100_INFERENCE_SUBNETWORKS, GpuVersion, Zone, ImageFamily, ImageProject, MachineVersion, Project, INFERENCE_NETWORKS, L4_INFERENCE_SUBNETWORKS
 from dags.inference.configs import trt_llm_mlperf_v41_config
 
 # Run once a day at 1 pm UTC (5 am PST)

diff --git a/dags/infra/clean_up.py b/dags/infra/clean_up.py
@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import Project, Zone
+from dags.common.vm_resource import Project, Zone
 from xlml.utils import tpu
 
 

diff --git a/dags/mlcompass/configs/simple_config.py b/dags/mlcompass/configs/simple_config.py
@@ -16,8 +16,8 @@
 
 import datetime
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner
-from dags.vm_resource import TpuVersion, Zone, Project, RuntimeVersion
+from dags.common import test_owner
+from dags.common.vm_resource import TpuVersion, Zone, Project, RuntimeVersion
 
 
 def get_simple_config():

diff --git a/dags/mlcompass/maxtext_gke.py b/dags/mlcompass/maxtext_gke.py
@@ -28,7 +28,7 @@
 from airflow import models
 from airflow.decorators import task
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
-from dags import test_owner
+from dags.common import test_owner
 from xlml.utils import xpk
 
 

diff --git a/dags/multipod/configs/gke_config.py b/dags/multipod/configs/gke_config.py
@@ -14,10 +14,11 @@
 
 """Utilities to construct configs for maxtext DAG on GKE."""
 
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
 from xlml.apis.xpk_cluster_config import XpkClusterConfig
-from dags import test_owner, gcs_bucket
-from dags.vm_resource import TpuVersion, Project, XpkClusters, GpuVersion, CpuVersion
+from dags import gcs_bucket
+from dags.common.vm_resource import TpuVersion, Project, XpkClusters, GpuVersion, CpuVersion
 from typing import Iterable
 import datetime
 

diff --git a/dags/multipod/configs/jax_tests_gce_config.py b/dags/multipod/configs/jax_tests_gce_config.py
@@ -15,9 +15,9 @@
 """Utilities to construct configs for JAX tests for GCE."""
 
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 import datetime
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value

diff --git a/dags/multipod/configs/jax_tests_gke_config.py b/dags/multipod/configs/jax_tests_gke_config.py
@@ -14,9 +14,9 @@
 
 """Utilities to construct configs for JAX tests for GCE."""
 
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import gke_config
-from dags.vm_resource import XpkClusterConfig
+from dags.common.vm_resource import XpkClusterConfig
 
 
 def get_jax_distributed_initialize_config(

diff --git a/dags/multipod/configs/legacy_unit_test.py b/dags/multipod/configs/legacy_unit_test.py
@@ -20,9 +20,9 @@
 from xlml.apis.xpk_cluster_config import XpkClusterConfig
 from base64 import b64encode
 from collections.abc import Iterable
-from dags import test_owner
+from dags.common import test_owner
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion, XpkClusters
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion, XpkClusters
 
 
 def get_legacy_unit_test_config(

diff --git a/dags/multipod/configs/maxtext_gce_config.py b/dags/multipod/configs/maxtext_gce_config.py
@@ -14,10 +14,11 @@
 
 """Utilities to construct configs for maxtext DAG."""
 
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner, gcs_bucket
+from dags import gcs_bucket
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 import datetime
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value

diff --git a/dags/multipod/configs/maxtext_sweep_gke_config.py b/dags/multipod/configs/maxtext_sweep_gke_config.py
@@ -16,7 +16,7 @@
 
 import datetime
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags.vm_resource import TpuVersion, XpkClusterConfig
+from dags.common.vm_resource import TpuVersion, XpkClusterConfig
 import itertools
 from typing import List, Iterable, Dict, Any
 

diff --git a/dags/multipod/configs/mxla_collective_config.py b/dags/multipod/configs/mxla_collective_config.py
@@ -14,10 +14,11 @@
 
 """Utilities to construct configs for maxtext DAG."""
 
+from dags.common import test_owner
 from xlml.apis import gcp_config, metric_config, task, test_config
-from dags import test_owner, gcs_bucket
+from dags import gcs_bucket
 from dags.multipod.configs import common
-from dags.vm_resource import TpuVersion, Project, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Project, RuntimeVersion
 import datetime
 
 PROJECT_NAME = Project.CLOUD_ML_AUTO_SOLUTIONS.value

diff --git a/dags/multipod/configs/pytorch_config.py b/dags/multipod/configs/pytorch_config.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dags.vm_resource import TpuVersion, Zone, DockerImage
+from dags.common.vm_resource import TpuVersion, Zone, DockerImage
 from dags.multipod.configs import gke_config
 from xlml.apis.xpk_cluster_config import XpkClusterConfig
 from xlml.apis import task

diff --git a/dags/multipod/jax_functional_tests.py b/dags/multipod/jax_functional_tests.py
@@ -17,7 +17,7 @@
 import datetime
 from airflow import models
 from dags import composer_env
-from dags.vm_resource import DockerImage, TpuVersion, Zone, Project, V5_NETWORKS, V5P_SUBNETWORKS, RuntimeVersion, XpkClusters
+from dags.common.vm_resource import DockerImage, TpuVersion, Zone, Project, V5_NETWORKS, V5P_SUBNETWORKS, RuntimeVersion, XpkClusters
 from dags.multipod.configs import jax_tests_gce_config, jax_tests_gke_config
 from dags.multipod.configs.common import SetupMode