addressed comments

Author: Edwinhr716

axlearn/cloud/gcp/job.py

@@ -353,7 +353,7 @@ def _execute(self):
             kind="LeaderWorkerSet",
             **self._build_leaderworkerset(),
         )
-        print(custom_object)
+        logging.info("submitting LeaderWorkerSet: %s", custom_object)
         return k8s.client.CustomObjectsApi().create_namespaced_custom_object(
             namespace=cfg.namespace,
             body=custom_object,

axlearn/cloud/gcp/lws_utils.py

@@ -170,10 +170,10 @@ def _build_pod(self) -> dict:
         cfg: TPULeaderWorkerTemplate.Config = self.config
         system = USER_FACING_NAME_TO_SYSTEM_CHARACTERISTICS[self._tpu_type]
         annotations, labels, selector = {}, {}, {}
-        if cfg.reservation is not None:
+        if cfg.reservation:
             logging.info("Using reservation=%s", cfg.reservation)
             selector.update({"cloud.google.com/reservation-name": cfg.reservation})
-        if cfg.reservation_project is not None:
+        if cfg.reservation_project:
             selector.update({"cloud.google.com/reservation-project": cfg.reservation_project})
 
         spec = dict(

axlearn/cloud/gcp/runners/gke.py

@@ -737,7 +737,6 @@ def _get_status(self) -> Status:
             RuntimeError: When the job fails, and LWS runner will retry it.
         """
         cfg: LWSRunnerJob.Config = self.config
-        logging.info("get status: ")
         try:
             resp = k8s.client.CustomObjectsApi().get_namespaced_custom_object_status(
                 name=cfg.name,
@@ -754,8 +753,6 @@ def _get_status(self) -> Status:
             condition_progressive = None
             condition_update_in_progress = None
 
-            logging.info(conditions)
-
             for condition in conditions:
                 if condition.get("type") == "Progressing":
                     condition_progressive = condition.get("status")
@@ -764,10 +761,6 @@ def _get_status(self) -> Status:
                 if condition.get("type") == "UpdateInProgress":
                     condition_update_in_progress = condition.get("status")
 
-            logging.info(condition_available)
-            logging.info(condition_progressive)
-            logging.info(condition_update_in_progress)
-
             if condition_update_in_progress:
                 return LWSRunnerJob.Status.UPDATING
 
@@ -807,10 +800,7 @@ def _execute(self):
         # Keep track of last status to prevent duplicate events.
         last_job_status = None
         while True:
-            logging.info("In Execute loop: ")
             status = self._get_status()
-            logging.info("Just after getting status: ")
-            logging.info(status)
 
             # Don't retry if FAILED, since we ask GKE to handle retries.
             # Note that LeaderWorkerSet remains ACTIVE until all retries are exhausted.