fix(ci): configure image inspection for arbitrary Slurm UIDs

Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com>

Author: functionstackx

scripts/ci/rocm/build_images.py

@@ -43,9 +43,12 @@ def record_image(docker, image_id, results, scratch):
     # Bind node-local inputs and copy evidence back as the submitting user.
     inspection = scratch / "image-inspection"
     inspection.mkdir(mode=0o700)
+    (inspection / "home").mkdir(mode=0o700)
     shutil.copytree(results / "controller", inspection / "controller")
     for name in ("request.json", "image-build.json"):
         shutil.copyfile(results / name, inspection / name)
+    # The Slurm UID need not have an image passwd entry. Libraries still need
+    # a username and writable home/cache directories during import.
     try:
         run(
             docker.command(
@@ -56,6 +59,16 @@ def record_image(docker, image_id, results, scratch):
                 docker.cgroup_parent,
                 "--user",
                 f"{os.getuid()}:{os.getgid()}",
+                "--env",
+                "HOME=/results/home",
+                "--env",
+                "USER=dynamo",
+                "--env",
+                "LOGNAME=dynamo",
+                "--env",
+                "XDG_CACHE_HOME=/results/home/.cache",
+                "--env",
+                "TORCHINDUCTOR_CACHE_DIR=/results/home/.cache/torchinductor",
                 "--mount",
                 f"type=bind,src={inspection},dst=/results",
                 "--mount",

scripts/ci/tests/test_slurm_build_images.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Exercise the official image-build sequence without starting a daemon."""
 
+import getpass
 import json
 import os
 import subprocess
@@ -206,6 +207,34 @@ def fail(command, log, **kwargs):
         )
         self.assertFalse((self.results / "build-manifest.json").exists())
 
+    def test_inspection_supports_uid_without_passwd_entry(self):
+        (self.results / "image-build.json").write_text("{}")
+
+        def inspect(command, log, **kwargs):
+            environment = dict(
+                command[index + 1].split("=", 1)
+                for index, argument in enumerate(command)
+                if argument == "--env"
+            )
+            with (
+                patch.dict(os.environ, environment, clear=True),
+                patch("pwd.getpwuid", side_effect=KeyError("unmapped uid")),
+            ):
+                self.assertEqual(getpass.getuser(), "dynamo")
+                self.assertEqual(Path("~").expanduser(), Path("/results/home"))
+            home = self.scratch / "image-inspection/home"
+            (home / "probe").write_text("writable")
+            self.assertTrue(
+                Path(environment["TORCHINDUCTOR_CACHE_DIR"]).is_relative_to(
+                    "/results/home"
+                )
+            )
+
+        with patch.object(build_images, "run", side_effect=inspect):
+            build_images.record_image(
+                self.daemon, self.test["Id"], self.results, self.scratch
+            )
+
 
 if __name__ == "__main__":
     unittest.main()