[diffusion] Bound overlong weight lock filenames (#34825)
This commit is contained in:
@@ -91,6 +91,12 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
|
|||||||
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
||||||
# add hash to avoid conflict with old users' lock files
|
# add hash to avoid conflict with old users' lock files
|
||||||
lock_file_name = hash_name + model_name + ".lock"
|
lock_file_name = hash_name + model_name + ".lock"
|
||||||
|
# Linux filesystems commonly cap one filename at 255 bytes. Absolute
|
||||||
|
# snapshot paths can exceed that even though the full path is valid.
|
||||||
|
# The digest is already collision-resistant, so fall back to it alone
|
||||||
|
# while preserving the historical name for ordinary paths.
|
||||||
|
if len(os.fsencode(lock_file_name)) > 255:
|
||||||
|
lock_file_name = hash_name + ".lock"
|
||||||
# mode 0o666 is required for the filelock to be shared across users
|
# mode 0o666 is required for the filelock to be shared across users
|
||||||
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
||||||
return lock
|
return lock
|
||||||
|
|||||||
@@ -1,11 +1,14 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
"""The patched find_local_ranks() must never touch torch.distributed."""
|
"""Unit tests for multimodal weight-loading utilities."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
import unittest
|
import unittest
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
from sglang.multimodal_gen.runtime.loader.weight_utils import (
|
from sglang.multimodal_gen.runtime.loader.weight_utils import (
|
||||||
_disable_runai_streamer_rank_discovery_collective,
|
_disable_runai_streamer_rank_discovery_collective,
|
||||||
|
get_lock,
|
||||||
)
|
)
|
||||||
|
|
||||||
_DIST_STREAMER_MOD = "runai_model_streamer.distributed_streamer.distributed_streamer"
|
_DIST_STREAMER_MOD = "runai_model_streamer.distributed_streamer.distributed_streamer"
|
||||||
@@ -83,5 +86,25 @@ class TestDisableRunaiStreamerRankDiscoveryCollective(unittest.TestCase):
|
|||||||
self.assertFalse(hasattr(_StubParams, "find_local_ranks"))
|
self.assertFalse(hasattr(_StubParams, "find_local_ranks"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestDiffusionWeightLock(unittest.TestCase):
|
||||||
|
def test_long_snapshot_path_uses_bounded_lock_filename(self):
|
||||||
|
component_path = os.path.join(
|
||||||
|
"/scratch",
|
||||||
|
"models--" + "very-long-repository-name-" * 8,
|
||||||
|
"snapshots",
|
||||||
|
"a" * 64,
|
||||||
|
"transformer",
|
||||||
|
"config.json",
|
||||||
|
)
|
||||||
|
|
||||||
|
with tempfile.TemporaryDirectory() as lock_dir:
|
||||||
|
lock = get_lock(component_path, lock_dir)
|
||||||
|
lock_filename = os.path.basename(lock.lock_file)
|
||||||
|
|
||||||
|
self.assertLessEqual(len(os.fsencode(lock_filename)), 255)
|
||||||
|
with lock:
|
||||||
|
self.assertTrue(os.path.exists(lock.lock_file))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user