Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion python/sglang/srt/entrypoints/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@
MultiTokenizerRouter,
get_main_process_id,
get_tokenizer_worker_class,
multi_tokenizer_args_shm_name,
read_from_shared_memory,
run_multi_detokenizer_router_process,
write_data_for_multi_tokenizer,
Expand Down Expand Up @@ -1347,7 +1348,7 @@ def attach_tokenizer_worker(cls, parent_pid: Optional[int] = None) -> Engine:
"""
parent_pid = get_main_process_id() if parent_pid is None else parent_pid
port_args, server_args, scheduler_info = read_from_shared_memory(
f"multi_tokenizer_args_{parent_pid}"
multi_tokenizer_args_shm_name(parent_pid)
)
publish(server_args, role="tokenizer")
self = cls.__new__(cls)
Expand Down
3 changes: 2 additions & 1 deletion python/sglang/srt/entrypoints/http_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@
TokenizerWorker,
get_main_process_id,
get_tokenizer_worker_class,
multi_tokenizer_args_shm_name,
read_from_shared_memory,
write_data_for_multi_tokenizer,
)
Expand Down Expand Up @@ -231,7 +232,7 @@ async def init_multi_tokenizer() -> ServerArgs:
# Read configuration from shared memory
main_pid = get_main_process_id()
port_args, server_args, scheduler_info = read_from_shared_memory(
f"multi_tokenizer_args_{main_pid}"
multi_tokenizer_args_shm_name(main_pid)
)
server_args: ServerArgs
port_args: PortArgs
Expand Down
25 changes: 24 additions & 1 deletion python/sglang/srt/managers/multi_tokenizer_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -769,6 +769,29 @@ def get_main_process_id() -> int:
return multiprocessing.current_process()._parent_pid


def _pid_namespace_id() -> Optional[int]:
"""Inode of this process's PID namespace, or None without /proc (e.g. macOS)."""
try:
return os.stat("/proc/self/ns/pid").st_ino
except OSError:
return None


def multi_tokenizer_args_shm_name(main_pid: int) -> str:
"""Name of the shared memory holding the multi-tokenizer args of `main_pid`.

A pid is unique only within its PID namespace. Containers that share /dev/shm
(e.g. docker --ipc=host) each run their server as pid 1, so a pid-only name
lets one server read, or overwrite, another server's args. Readers (tokenizer
workers, Engine.attach_tokenizer_worker) address `main_pid` from the same PID
namespace, so they derive the same name.
"""
pid_ns = _pid_namespace_id()
if pid_ns is None:
return f"multi_tokenizer_args_{main_pid}"
return f"multi_tokenizer_args_{pid_ns}_{main_pid}"


def write_to_shared_memory(obj, name: str) -> shared_memory.SharedMemory:
"""Write data to shared memory"""
serialized = pickle.dumps(obj)
Expand Down Expand Up @@ -811,7 +834,7 @@ def write_data_for_multi_tokenizer(
current_pid = os.getpid()
logger.info(f"main process ID: {main_pid}, current process ID: {current_pid}")
args = (port_args, server_args, scheduler_info)
args_shm = write_to_shared_memory(args, f"multi_tokenizer_args_{current_pid}")
args_shm = write_to_shared_memory(args, multi_tokenizer_args_shm_name(current_pid))
args_shm.close()

return args_shm
40 changes: 40 additions & 0 deletions test/registered/unit/managers/test_multi_tokenizer_mixin.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import unittest
from unittest.mock import patch

import numpy as np

Expand All @@ -9,11 +10,15 @@

maybe_stub_sgl_kernel()

from sglang.srt.managers import multi_tokenizer_mixin
from sglang.srt.managers.io_struct import BatchStrOutput
from sglang.srt.managers.multi_tokenizer_mixin import (
TokenizerWorker,
_handle_output_by_index,
get_tokenizer_worker_class,
multi_tokenizer_args_shm_name,
read_from_shared_memory,
write_to_shared_memory,
)

register_cpu_ci(est_time=12, suite="base-a-test-cpu")
Expand Down Expand Up @@ -162,5 +167,40 @@ def test_get_tokenizer_worker_class_rejects_non_worker(self):
get_tokenizer_worker_class(InvalidServerArgs())


class TestMultiTokenizerArgsShmName(unittest.TestCase):
def test_same_pid_in_two_pid_namespaces_keeps_args_apart(self):
# Two containers sharing /dev/shm (docker --ipc=host) both run their server as pid 1.
names = []
for pid_ns in (4026532001, 4026532002):
with patch.object(
multi_tokenizer_mixin, "_pid_namespace_id", return_value=pid_ns
):
names.append(multi_tokenizer_args_shm_name(1))
self.assertNotEqual(names[0], names[1])
shms = [
write_to_shared_memory(("port-a", "server-a", {}), names[0]),
write_to_shared_memory(("port-b", "server-b", {}), names[1]),
]
try:
self.assertEqual(
read_from_shared_memory(names[0]), ("port-a", "server-a", {})
)
self.assertEqual(
read_from_shared_memory(names[1]), ("port-b", "server-b", {})
)
finally:
for shm in shms:
shm.close()
shm.unlink()

def test_without_pid_namespace_falls_back_to_pid(self):
with patch.object(
multi_tokenizer_mixin, "_pid_namespace_id", return_value=None
):
self.assertEqual(
multi_tokenizer_args_shm_name(1234), "multi_tokenizer_args_1234"
)


if __name__ == "__main__":
unittest.main()
2 changes: 2 additions & 0 deletions test/registered/utils/test_stale_shm_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@ def test_embeds_pid_and_is_unique(self):
def test_creator_pid_parsing(self):
self.assertEqual(_creator_pid("sgl_shm_mq_1234_abcd1234"), 1234)
self.assertEqual(_creator_pid("multi_tokenizer_args_5678"), 5678)
# Qualified by the PID namespace (multi_tokenizer_args_shm_name).
self.assertEqual(_creator_pid("multi_tokenizer_args_4026531836_5678"), 5678)
self.assertIsNone(_creator_pid("psm_deadbeef"))
self.assertIsNone(_creator_pid("sgl_shm_garbage"))
self.assertIsNone(_creator_pid("multi_tokenizer_args_notanint"))
Expand Down
Loading