Set srun job name to "<account>-CloudAI_install_docker_image.%Y%m%d_%H%M%S" (#544)

Author: TaekyungHeo

src/cloudai/util/docker_image_cache_manager.py

@@ -18,6 +18,7 @@
 import os
 import shutil
 import subprocess
+from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
@@ -185,7 +186,15 @@ def check_docker_image_exists(self, docker_image_url: str, docker_image_filename
     def _import_docker_image(
         self, srun_prefix: str, docker_image_url: str, docker_image_path: Path
     ) -> DockerImageCacheResult:
-        enroot_import_cmd = f"{srun_prefix} enroot import -o {docker_image_path} docker://{docker_image_url}"
+        job_name = "CloudAI_install_docker_image"
+        if self.system.account:
+            job_name = f"{self.system.account}-{job_name}.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        else:
+            job_name = f"{job_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+
+        enroot_import_cmd = (
+            f"{srun_prefix} --job-name={job_name} enroot import -o {docker_image_path} docker://{docker_image_url}"
+        )
         logging.debug(f"Importing Docker image: {enroot_import_cmd}")
         try:
             p = subprocess.run(enroot_import_cmd, shell=True, check=True, capture_output=True, text=True)

tests/test_docker_image_cache_manager.py

@@ -108,16 +108,14 @@ def test_cache_docker_image(
     ]  # Ensure all path checks return True
     mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stderr="")
     result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz")
-    mock_run.assert_called_once_with(
-        (
-            f"srun --export=ALL --partition={slurm_system.default_partition} enroot "
-            f"import -o {slurm_system.install_path}/image.tar.gz docker://docker.io/hello-world"
-        ),
-        shell=True,
-        check=True,
-        capture_output=True,
-        text=True,
-    )
+
+    assert mock_run.call_count == 1
+    actual_command = mock_run.call_args[0][0]
+    assert f"srun --export=ALL --partition={slurm_system.default_partition}" in actual_command
+    assert "--job-name=CloudAI_install_docker_image_" in actual_command
+    assert f"enroot import -o {slurm_system.install_path}/image.tar.gz docker://docker.io/hello-world" in actual_command
+    assert mock_run.call_args[1] == {"shell": True, "check": True, "capture_output": True, "text": True}
+
     assert result.success
     assert result.message == f"Docker image cached successfully at {slurm_system.install_path}/image.tar.gz."
 
@@ -215,17 +213,25 @@ def test_docker_import_with_extra_system_config(slurm_system: SlurmSystem, accou
         res = manager.cache_docker_image("docker.io/hello-world", "docker_image.sqsh")
         assert res.success
 
-    command = f"srun --export=ALL --partition={slurm_system.default_partition}"
-    if slurm_system.account:
-        command += f" --account={slurm_system.account}"
-    if slurm_system.gpus_per_node:
-        command += " --gres=gpu:1"
-    command += f" enroot import -o {slurm_system.install_path}/docker_image.sqsh docker://docker.io/hello-world"
-
-    mock_run.assert_called_once_with(
-        command,
-        shell=True,
-        check=True,
-        capture_output=True,
-        text=True,
+    assert mock_run.call_count == 1
+
+    actual_command = mock_run.call_args[0][0]
+
+    expected_prefix = f"srun --export=ALL --partition={slurm_system.default_partition}"
+    assert expected_prefix in actual_command
+
+    if account:
+        assert f"--account={account}" in actual_command
+        assert f"--job-name={account}-CloudAI_install_docker_image." in actual_command
+    else:
+        assert "--job-name=CloudAI_install_docker_image_" in actual_command
+
+    if gpus_per_node:
+        assert "--gres=gpu:1" in actual_command
+
+    assert (
+        f"enroot import -o {slurm_system.install_path}/docker_image.sqsh docker://docker.io/hello-world"
+        in actual_command
     )
+
+    assert mock_run.call_args[1] == {"shell": True, "check": True, "capture_output": True, "text": True}