Add GRAND_CHALLENGE_COMPONENT_USER

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sagemaker-shim"
-version = "0.2.0"
+version = "0.2.1-alpha"
 description = "Adapts algorithms that implement the Grand Challenge inference API for running in SageMaker"
 authors = ["James Meakin <[email protected]>"]
 license = "Apache-2.0"

sagemaker_shim/models.py

@@ -205,6 +205,10 @@ def entrypoint(self) -> Any:
         logger.debug(f"{entrypoint=}")
         return entrypoint
 
+    @property
+    def user(self) -> str:
+        return os.environ.get("GRAND_CHALLENGE_COMPONENT_USER", "")
+
     @property
     def input_path(self) -> Path:
         """Local path where the subprocess is expected to read its input files"""
@@ -274,6 +278,20 @@ def proc_env(self) -> dict[str, str]:
 
         return env
 
+    @property
+    def proc_user(self) -> dict[str, str | None]:
+        match = re.fullmatch(
+            r"^(?P<user>[0-9a-zA-Z]*):?(?P<group>[0-9a-zA-Z]*)$", self.user
+        )
+
+        if match:
+            return {
+                "user": match.group("user") or None,
+                "group": match.group("group") or None,
+            }
+        else:
+            return {"user": None, "group": None}
+
     async def invoke(self) -> InferenceResult:
         """Run the inference on a single case"""
 
@@ -400,6 +418,8 @@ async def execute(self) -> int:
 
         process = await asyncio.create_subprocess_exec(
             *self.proc_args,
+            user=self.proc_user["user"],
+            group=self.proc_user["group"],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             env=self.proc_env,

tests/test_container_image.py

@@ -93,6 +93,7 @@ def _container(*, base_image="hello-world:latest", host_port=8080, cmd=None):
                 init=False,
                 environment=container_env,
                 links={minio.container.name: "minio"},
+                user=0,
             )
 
             # Wait for startup

tests/test_models.py

@@ -57,3 +57,37 @@ def test_removing_ld_library_path(monkeypatch):
 
     assert "LD_LIBRARY_PATH" not in t.proc_env
     assert env["LD_LIBRARY_PATH"] == "present"
+
+
+@pytest.mark.parametrize(
+    "user,expected_user,expected_group",
+    (
+        ("1000", "1000", None),
+        ("1000:1000", "1000", "1000"),
+        (":1000", None, "1000"),
+        ("", None, None),
+        ("myuser", "myuser", None),
+        ("myuser:myuser", "myuser", "myuser"),
+        (":myuser", None, "myuser"),
+        ("", None, None),
+        ("🙈:🙉", None, None),
+    ),
+)
+def test_proc_user(monkeypatch, user, expected_user, expected_group):
+    monkeypatch.setenv("GRAND_CHALLENGE_COMPONENT_USER", user)
+
+    t = InferenceTask(
+        pk="test", inputs=[], output_bucket_name="test", output_prefix="test"
+    )
+
+    assert t.user == user
+    assert t.proc_user == {"user": expected_user, "group": expected_group}
+
+
+def test_proc_user_unset():
+    t = InferenceTask(
+        pk="test", inputs=[], output_bucket_name="test", output_prefix="test"
+    )
+
+    assert t.user == ""
+    assert t.proc_user == {"user": None, "group": None}

tests/test_patch_image.py

@@ -71,6 +71,7 @@ def test_patch_image(registry):
     assert env_vars == {
         "GRAND_CHALLENGE_COMPONENT_CMD_B64J": "WyJzaCJd",
         "GRAND_CHALLENGE_COMPONENT_ENTRYPOINT_B64J": "bnVsbA==",
+        "GRAND_CHALLENGE_COMPONENT_USER": "0",
     }
     assert new_config["config"]["Entrypoint"] == ["/sagemaker-shim"]
     assert "Cmd" not in new_config["config"]

tests/utils.py

@@ -62,6 +62,7 @@ def get_new_env_vars(*, existing_config: dict[str, Any]) -> dict[str, str]:
         "GRAND_CHALLENGE_COMPONENT_ENTRYPOINT_B64J": encode_b64j(
             val=entrypoint
         ),
+        "GRAND_CHALLENGE_COMPONENT_USER": "0",
     }