Merge branch 'main' into jax-config-common-func

Author: samos123

axlearn/cloud/common/bastion.py

@@ -1037,15 +1037,30 @@ def _update_single_job(self, job: Job) -> Job:
             serialized_jobspec = io.StringIO()
             serialize_jobspec(job.spec, serialized_jobspec)
             env_vars |= {_BASTION_SERIALIZED_JOBSPEC_ENV_VAR: serialized_jobspec.getvalue()}
-            _start_command(
-                job,
-                remote_log_dir=self._log_dir,
-                env_vars=env_vars,
-            )
-            assert job.command_proc is not None
 
-            # If command is completed, move to CLEANING. Otherwise, it's still RUNNING.
-            if _is_proc_complete(job.command_proc):
+            try:
+                _start_command(
+                    job,
+                    remote_log_dir=self._log_dir,
+                    env_vars=env_vars,
+                )
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                job.command_proc = None
+                logging.warning(
+                    "Failed to start command for the job %s: %s with error: %s",
+                    job.spec.name,
+                    job.spec.command,
+                    e,
+                )
+
+            if job.command_proc is None:
+                # If failed to start command, moving the job to CLEANING directly.
+                self._append_to_job_history(
+                    job, msg="Failed to start job command", state=JobLifecycleState.CLEANING
+                )
+                job.state.status = JobStatus.CLEANING
+            elif _is_proc_complete(job.command_proc):
+                # If command is completed, move to CLEANING. Otherwise, it's still RUNNING.
                 self._append_to_job_history(
                     job, msg="CLEANING: process finished", state=JobLifecycleState.CLEANING
                 )

axlearn/cloud/common/bastion_test.py

@@ -1112,6 +1112,46 @@ def mock_tfio_exists(f):
                 updated_job.state, JobState(status=JobStatus.CLEANING, metadata={"tier": 1})
             )
 
+    @mock.patch("subprocess.Popen")
+    def test_active_command_malformed(self, mock_popen):
+        """If the command is malformed, bastion should gracefully move the job to CLEANING state."""
+        job = Job(
+            spec=new_jobspec(
+                name="test_job",
+                command="command",
+                cleanup_command="cleanup",
+                metadata=JobMetadata(
+                    user_id="test_user",
+                    project_id="test_job",
+                    creation_time=datetime.now(),
+                    resources={"v4": 8},
+                ),
+            ),
+            state=JobState(status=JobStatus.ACTIVE, metadata={"tier": 1}),
+            command_proc=None,  # Initially, command is None.
+            cleanup_proc=None,
+        )
+        patch_fns = mock.patch.multiple(
+            bastion.__name__,
+            _upload_job_state=mock.DEFAULT,
+        )
+        mock_popen.side_effect = ValueError("Command malformed")
+
+        with patch_fns, self._patch_bastion(None) as mock_bastion:
+            # Initially, job should have no command.
+            self.assertIsNone(job.command_proc)
+
+            # Run single update step to start the job.
+            updated_job = mock_bastion._update_single_job(job)
+
+            # Command failed to be started.
+            self.assertIsNone(updated_job.command_proc)
+
+            # Job state should be CLEANING.
+            self.assertEqual(
+                updated_job.state, JobState(status=JobStatus.CLEANING, metadata={"tier": 1})
+            )
+
     # pylint: disable-next=too-many-branches
     def test_update_jobs(self):
         """Tests the global update step."""

axlearn/cloud/common/git_summary.py

@@ -104,17 +104,20 @@ class GitSummaryMembers(enum.Enum):
     commit = GitSummaryMember(
         label="git-commit", file=".git.commit", cmd=("git", "rev-parse", "HEAD")
     )
+    # For retrieving a git branch, use rev-parse instead of ``git branch --show_current``
+    # Since the latter is not supported on older git versions (was introduced in git version 2.22.0)
     branch = GitSummaryMember(
         label="git-branch",
         file=".git.branch",
-        cmd=("git", "branch", "--show-current"),
+        cmd=("git", "rev-parse", "--abbrev-ref", "HEAD"),
         indicator_label=False,
     )
     remote = GitSummaryMember(
         label="git-remote",
         file=".git.remote",
         # get remote url where branch is tracked (if any)
-        cmd="git remote get-url $(git config --get branch.$(git branch --show-current).remote) "
+        cmd="git remote get-url $(git config --get"
+        " branch.$(git rev-parse --abbrev-ref HEAD).remote) "
         "|| git remote get-url origin "  # fallback to origin
         "|| git config --get-regexp 'remote.*.url' | head -n 1 | cut -d= -f2",  # or 1st remote
         required=False,  # remote is not always available

axlearn/cloud/common/git_summary_test.py

@@ -70,9 +70,9 @@ def test_valid_summary(self, remote):
                 + (
                     (
                         f"&& git remote add origin {remote} "
-                        "&& git config branch.$(git branch --show-current).remote origin "
-                        "&& git config branch.$(git branch --show-current).merge refs/heads/"
-                        "$(git branch --show-current)"
+                        "&& git config branch.$(git rev-parse --abbrev-ref HEAD).remote origin "
+                        "&& git config branch.$(git rev-parse --abbrev-ref HEAD).merge refs/heads/"
+                        "$(git rev-parse --abbrev-ref HEAD)"
                     )
                     if remote
                     else ""

axlearn/cloud/gcp/jobs/cpu_runner.py

@@ -117,7 +117,7 @@ def define_flags(cls, fv: flags.FlagValues):
             "https://cloud.google.com/compute/docs/general-purpose-machines",
             **common_kwargs,
         )
-        flags.DEFINE_integer("disk_size", 64, "Disk size of the VM in GB.", **common_kwargs)
+        flags.DEFINE_integer("disk_size", 128, "Disk size of the VM in GB.", **common_kwargs)
         flags.DEFINE_bool(
             "retain_vm",
             False,

axlearn/cloud/gcp/jobset_utils.py

@@ -579,7 +579,7 @@ def _build_pod(self) -> Nested[Any]:
                             bucketName=parsed.netloc,
                             # pylint: disable=line-too-long
                             mountOptions=f"only-dir={parsed.path.lstrip('/')},implicit-dirs,metadata-cache:ttl-secs:-1,metadata-cache:stat-cache-max-size-mb:-1,metadata-cache:type-cache-max-size-mb:-1,kernel-list-cache-ttl-secs=-1,gcs-connection:http-client-timeout:{cfg.gcsfuse_mount.http_client_timeout}",
-                            gcsfuseMetadataPrefetchOnMount="true",  # Improves first-time read.
+                            gcsfuseMetadataPrefetchOnMount="false",  # Improves first-time read.
                             disableMetrics="false",  # Enables GCSFuse metrics by default.
                         ),
                     ),

axlearn/cloud/gcp/vm.py

@@ -365,19 +365,28 @@ def list_disk_images(creds: Credentials) -> list[str]:
     """
     resource = _compute_resource(creds)
     image_project = gcp_settings("image_project")
-    images = (
-        resource.images()
-        .list(
-            project=image_project,
-            orderBy="creationTimestamp desc",
-            maxResults=50,
-        )
-        .execute()
-    )
     image_names = []
-    for el in images["items"]:
-        if "ubuntu-2004" in el["name"] and "arm" not in el["name"]:
-            image_names.append(el["name"])
+    next_page_token = None
+
+    while True:
+        images = (
+            resource.images()
+            .list(
+                project=image_project,
+                orderBy="creationTimestamp desc",
+                maxResults=100,
+                pageToken=next_page_token,
+            )
+            .execute()
+        )
+        for el in images["items"]:
+            if "ubuntu-2204" in el["name"] and "arm" not in el["name"]:
+                image_names.append(el["name"])
+
+        next_page_token = images.get("nextPageToken")
+        if next_page_token is None:
+            break
+
     return [f"projects/{image_project}/global/images/{name}" for name in image_names]