Spaces:
Running on Zero
Running on Zero
File size: 2,202 Bytes
290fe6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import os
import unittest
from unittest.mock import patch
from zerogpu_runtime import configure_zerogpu_cgroup
class ZeroGPUCgroupTests(unittest.TestCase):
def configure(self, content, *, exists=False, enabled="true"):
env = {
"SPACES_ZERO_GPU": enabled,
"ZEROGPU_PROC_SELF_CGROUP_PATH": "/etc/host-cgroup",
}
with patch.dict(os.environ, env, clear=True):
with patch("zerogpu_runtime.Path.exists", return_value=exists):
with patch("zerogpu_runtime.Path.read_text", return_value=content) as read:
configure_zerogpu_cgroup()
return os.environ["ZEROGPU_PROC_SELF_CGROUP_PATH"], read.call_count
def test_non_zerogpu_is_unchanged(self):
self.assertEqual(self.configure("", enabled="false"), ("/etc/host-cgroup", 0))
def test_existing_host_mount_is_preserved(self):
self.assertEqual(self.configure("", exists=True), ("/etc/host-cgroup", 0))
def test_missing_host_mount_with_v1_devices(self):
self.assertEqual(
self.configure("5:memory:/container\n6:devices:/kubepods/pod/container\n"),
("/proc/self/cgroup", 1),
)
def test_missing_host_mount_with_v2(self):
self.assertEqual(
self.configure("0::/kubepods.slice/pod/container\n"),
("/proc/self/cgroup", 1),
)
def test_unusable_procfs_is_rejected(self):
for content in ("0::/\n", "6:devices:/\n", "", "garbage", "0::relative", "0::/../container"):
with self.subTest(content=content):
with self.assertRaisesRegex(RuntimeError, "Restart the Hugging Face Space"):
self.configure(content)
def test_unreadable_procfs_is_actionable(self):
with patch.dict(os.environ, {"SPACES_ZERO_GPU": "true"}, clear=True):
with patch("zerogpu_runtime.Path.exists", return_value=False):
with patch("zerogpu_runtime.Path.read_text", side_effect=FileNotFoundError):
with self.assertRaisesRegex(RuntimeError, "unreadable"):
configure_zerogpu_cgroup()
if __name__ == "__main__":
unittest.main()
|