fix to the expected json format (GoogleCloudPlatform#1060)

Author: severus-ho

xlml/apis/task.py

@@ -979,6 +979,7 @@ def run(self) -> DAGNode:
 
       gke_run = gke.run_job.override(group_id="run_model")(
           job_body,
+          self.task_test_config,
           self.task_gcp_config,
           self.cluster_name,
           self.job_create_timeout,

xlml/utils/gke.py

@@ -12,7 +12,7 @@
 from google.cloud import container_v1
 import kubernetes
 
-from xlml.apis import gcp_config
+from xlml.apis import gcp_config, test_config
 from xlml.utils import composer
 
 """Utilities for GKE."""
@@ -55,6 +55,7 @@ def get_authenticated_client(
 def run_job(
     body: Dict[str, Any],
     gcp: gcp_config.GCPConfig,
+    gke_test_config: test_config.GpuGkeTest,
     cluster_name: str,
     job_create_timeout: datetime.timedelta,
     gcs_location: str = '',
@@ -64,6 +65,7 @@ def run_job(
   Args:
     body: Dict that defines a Kubernetes `Job`.
     gcp: GCP config with the project name and zone of the GKE cluster.
+    gke_test_config: Test config with the accelerator information of the GKE cluster
     cluster_name: Name of the GCP cluster.
     job_create_timeout: Amount of time to wait for all pods to become active.
   """
@@ -72,12 +74,13 @@ def run_job(
   def deploy_job(gcs_location):
     # Log required info for XLML PLX Dashboard
     composer.log_metadata_for_xlml_dashboard({
-        'project_name': gcp.project_name,
+        'cluster_project': gcp.project_name,
         'zone': gcp.zone,
         'dataset_name': gcp.dataset_name.value,
         'composer_project': gcp.composer_project,
         'dataset_project': gcp.dataset_project,
         'cluster_name': cluster_name,
+        'accelerator_type': gke_test_config.accelerator.machine_type,
     })
 
     body['spec']['template']['spec']['containers'][0]['env'].append(

xlml/utils/gpu.py

@@ -319,7 +319,7 @@ def create_resource_request(
     # Log required info for XLML PLX Dashboard
     composer.log_metadata_for_xlml_dashboard({
         "instance_name": instance_name,
-        "project_name": gcp.project_name,
+        "cluster_project": gcp.project_name,
         "zone": gcp.zone,
         "dataset_name": gcp.dataset_name.value,
         "composer_project": gcp.composer_project,
@@ -331,6 +331,7 @@ def create_resource_request(
             "machine_type": accelerator.machine_type,
             "image_family": accelerator.image_family,
         },
+        "accelerator_type": accelerator.machine_type,
     })
 
     machine_type = accelerator.machine_type

xlml/utils/tpu.py

@@ -86,7 +86,7 @@ def create_queued_resource_request(
     # Log required info for XLML PLX Dashboard
     composer.log_metadata_for_xlml_dashboard({
         'instance_name': tpu_name,
-        'project_name': gcp.project_name,
+        'cluster_project': gcp.project_name,
         'zone': gcp.zone,
         'dataset_name': gcp.dataset_name.value,
         'composer_project': gcp.composer_project,
@@ -97,6 +97,7 @@ def create_queued_resource_request(
             'runtime_version': task_test_config.accelerator.runtime_version,
             'version': task_test_config.accelerator.version.value,
         },
+        'accelerator_type': task_test_config.accelerator.name,
     })
 
     creds, _ = google.auth.default()