[diffusion] chore: revise process title (#18446)

This commit is contained in:
Mick
2026-02-09 00:14:06 +08:00
committed by GitHub
parent 031a652b93
commit 6601bc24da

View File

@@ -12,11 +12,20 @@ from setproctitle import setproctitle
from sglang.multimodal_gen import envs
from sglang.multimodal_gen.runtime.distributed import (
get_sp_group,
get_tp_rank,
get_tp_world_size,
maybe_init_distributed_environment_and_model_parallel,
model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
get_classifier_free_guidance_world_size,
get_ring_parallel_rank,
get_ring_parallel_world_size,
get_tp_group,
get_ulysses_parallel_rank,
get_ulysses_parallel_world_size,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs
from sglang.multimodal_gen.runtime.pipelines_core import (
@@ -70,7 +79,6 @@ class GPUWorker:
def init_device_and_model(self) -> None:
"""Initialize the device and load the model."""
setproctitle(f"sgl_diffusion::scheduler_TP{self.local_rank}")
torch.get_device_module().set_device(self.local_rank)
# Set environment variables for distributed initialization
os.environ["MASTER_ADDR"] = "localhost"
@@ -78,7 +86,7 @@ class GPUWorker:
os.environ["LOCAL_RANK"] = str(self.local_rank)
os.environ["RANK"] = str(self.rank)
os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus)
# Initialize the distributed environment
# initialize the distributed environment
maybe_init_distributed_environment_and_model_parallel(
tp_size=self.server_args.tp_size,
enable_cfg_parallel=self.server_args.enable_cfg_parallel,
@@ -90,6 +98,25 @@ class GPUWorker:
dist_timeout=self.server_args.dist_timeout,
)
# set proc title
if model_parallel_is_initialized():
suffix = ""
if get_tp_world_size() != 1:
tp_rank = get_tp_rank()
suffix += f"_TP{tp_rank}"
if get_ulysses_parallel_world_size() != 1:
u_rank = get_ulysses_parallel_rank()
suffix += f"_U{u_rank}"
if get_ring_parallel_world_size() != 1:
r_rank = get_ring_parallel_rank()
suffix += f"_R{r_rank}"
if get_classifier_free_guidance_world_size() != 1:
c_rank = get_classifier_free_guidance_rank()
suffix += f"_C{c_rank}"
setproctitle(f"sgl_diffusion::scheduler{suffix}")
else:
setproctitle(f"sgl_diffusion::scheduler_{self.local_rank}")
self.pipeline = build_pipeline(self.server_args)
# apply layerwise offload after lora is applied while building LoRAPipeline