Guard against random.seed() when making job names unique

Differential Revision: D88241480

Pull Request resolved: #1170

Author: kiukchung

torchx/schedulers/ids.py

@@ -8,9 +8,9 @@
 # pyre-strict
 
 import os
-import random
 import struct
 
+
 START_CANDIDATES: str = "bcdfghjklmnpqrstvwxz"
 END_CANDIDATES: str = START_CANDIDATES + "012345679"
 
@@ -19,14 +19,19 @@ def make_unique(name: str, string_length: int = 0) -> str:
     """
     Appends a unique 64-bit string to the input argument.
 
+    Note that the unique string pulls entropy from `/dev/urandom` hence is not
+    affected by `random.seed()`
+
+    Args:
+        name: the name string to unique-ify
+        string_length: max length of the unique 64-bit string to append to the ``name``.
+          Default is 0, which returns the length of a randomly generated 64-bit string (typically 11-14 characters long).
+
     Returns:
-        string in format $name-$unique_suffix
+        string in format ``{name}-{unique_suffix}`
     """
-    return (
-        f"{name}-{random_id()}"
-        if string_length == 0
-        else f"{name}-{get_len_random_id(string_length)}"
-    )
+    max_length = None if string_length == 0 else string_length
+    return f"{name}-{random_id(max_length)}"
 
 
 def random_uint64() -> int:
@@ -36,13 +41,24 @@ def random_uint64() -> int:
     return struct.unpack("!Q", os.urandom(8))[0]
 
 
-def random_id() -> str:
+def random_id(max_length: int | None = None) -> str:
     """
     Generates an alphanumeric string ID that matches the requirements from
     https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
+
+    Note that the unique string pulls entropy from `/dev/urandom` hence is not
+    affected by `random.seed()`
+
+    If ``max_length`` is provided, the returned ID will be at most that many characters long.
+
     """
+    # If a max_length is provided and is non-positive, return empty string
+    if max_length is not None and max_length <= 0:
+        return ""
+
     out = ""
     v = random_uint64()
+
     while v > 0:
         if out == "":
             candidates = START_CANDIDATES
@@ -52,21 +68,9 @@ def random_id() -> str:
         char = v % len(candidates)
         v = v // len(candidates)
         out += candidates[char]
-    return out
-
-
-def get_len_random_id(string_length: int) -> str:
-    """
-    Generates an alphanumeric string ID that matches the requirements from
-    https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
-    """
-    out = ""
-    for i in range(string_length):
-        if out == "":
-            candidates = START_CANDIDATES
-        else:
-            candidates = END_CANDIDATES
 
-        out += random.choice(candidates)
+        if max_length is not None and len(out) >= max_length:
+            break
 
+    # NOTE: statistically the length of `out` is typically between 12-14 characters long
     return out

torchx/schedulers/test/ids_test.py

@@ -8,15 +8,27 @@
 # pyre-strict
 
 
+import typing
 import unittest
+from contextlib import contextmanager
 from unittest.mock import MagicMock, patch
 
-from torchx.schedulers.ids import (
-    get_len_random_id,
-    make_unique,
-    random_id,
-    random_uint64,
-)
+from torchx.schedulers.ids import make_unique, random_id, random_uint64
+
+
+@contextmanager
+def scoped_random_seed(seed: int) -> typing.Generator[None, None, None]:
+    """
+    Temporarily set the random module's seed and restore its state afterward.
+    """
+    import random
+
+    state = random.getstate()
+    try:
+        random.seed(seed)
+        yield
+    finally:
+        random.setstate(state)
 
 
 class IdsTest(unittest.TestCase):
@@ -42,15 +54,34 @@ def test_random_id(self) -> None:
         self.assertIn(v[0], ALPHAS)
         self.assertGreater(len(v), 5)
 
-    def test_get_len_random_id(self) -> None:
-        size = 6
-        self.assertNotEqual(get_len_random_id(size), get_len_random_id(size))
-        self.assertEqual(size, len(get_len_random_id(size)))
+    def test_random_id_max_length(self) -> None:
+        for max_length in range(6, 10):
+            with self.subTest(max_length=max_length):
+                self.assertLessEqual(len(random_id(max_length)), max_length)
+                self.assertNotEqual(random_id(max_length), random_id(max_length))
+
+    def test_random_id_zero_max_length(self) -> None:
+        self.assertEqual("", random_id(max_length=0))
 
     @patch("os.urandom", return_value=bytes(range(8)))
     def test_random_id_seed(self, urandom: MagicMock) -> None:
         self.assertEqual(random_id(), "fzfjxlmln9")
+        self.assertEqual(random_id(max_length=6), "fzfjxl")
 
     @patch("os.urandom", return_value=bytes(range(8)))
     def test_make_unique_seed(self, urandom: MagicMock) -> None:
         self.assertEqual(make_unique("test"), "test-fzfjxlmln9")
+
+    def test_make_unique_not_affected_by_random_seed(self) -> None:
+        # Seeding the Python random module should not affect make_unique(),
+        # which relies on os.urandom for entropy.
+        with scoped_random_seed(0):
+            v1 = make_unique("test")
+
+        with scoped_random_seed(0):
+            v2 = make_unique("test")
+
+        # Even with the same random seed, make_unique should produce different values.
+        self.assertNotEqual(v1, v2)
+        self.assertTrue(v1.startswith("test-"))
+        self.assertTrue(v2.startswith("test-"))

torchx/schedulers/test/kubernetes_mcad_scheduler_test.py

@@ -482,20 +482,19 @@ def test_cleanup_str(self) -> None:
         self.assertEqual("abcd1234", cleanup_str("1234abcd1234"))
 
     def test_get_unique_truncated_appid(self) -> None:
-        scheduler = create_scheduler("test")
         app = _test_app()
         app.name = "abcde"
-        self.assertEqual(20, len(get_unique_truncated_appid(app)))
+        self.assertLessEqual(len(get_unique_truncated_appid(app)), 20)
         self.assertIn(app.name, get_unique_truncated_appid(app))
 
         app.name = "abcdefghijklmnopqrstuvwxyz012345678910111213141516"
-        self.assertEqual(56, len(get_unique_truncated_appid(app)))
+        self.assertLessEqual(len(get_unique_truncated_appid(app)), 56)
         self.assertIn(app.name, get_unique_truncated_appid(app))
 
         app.name = (
             "abcdefghijklmnopqrstuvwxyz012345678910111213141516171819202122232425"
         )
-        self.assertEqual(59, len(get_unique_truncated_appid(app)))
+        self.assertLessEqual(len(get_unique_truncated_appid(app)), 59)
         self.assertIn(
             "abcdefghijklmnopqrstuvwxyz01234567891011121314151617181",
             get_unique_truncated_appid(app),