[diffusion] feat: add --strict-ports option for predictable port assignment (#21320)

Co-authored-by: 阳虎 <yanghu@yanghudeMacBook-Pro.local>
This commit is contained in:
yang1002378395-cmyk
2026-03-27 16:40:50 +08:00
committed by GitHub
co-authored by 阳虎
parent 448b528720
commit f83b1b73a8
2 changed files with 42 additions and 12 deletions
+4 -1
View File
@@ -349,7 +349,10 @@ def _get_config_info(
model_id = matched_model_names[0] model_id = matched_model_names[0]
return _CONFIG_REGISTRY.get(model_id) return _CONFIG_REGISTRY.get(model_id)
else: else:
raise RuntimeError(f"No model info found for model path: {model_path}") raise RuntimeError(
f"No model info found for model path: {model_path}. "
f"Please check the model path or specify the model_id explicitly."
)
# --- Part 3: Main Resolver --- # --- Part 3: Main Resolver ---
@@ -202,6 +202,9 @@ class ServerArgs:
scheduler_port: int = 5555 scheduler_port: int = 5555
# Strict port mode: fail if requested port is unavailable instead of auto-selecting
strict_ports: bool = False
output_path: str | None = "outputs/" output_path: str | None = "outputs/"
input_save_path: str | None = "inputs/uploads" input_save_path: str | None = "inputs/uploads"
@@ -384,6 +387,24 @@ class ServerArgs:
) )
def _adjust_network_ports(self): def _adjust_network_ports(self):
if self.strict_ports:
# Strict mode: fail if port is unavailable
if not is_port_available(self.port):
raise RuntimeError(
f"Port {self.port} is unavailable and --strict-ports is enabled. "
f"Either use a different port or remove --strict-ports to allow auto-selection."
)
if not is_port_available(self.scheduler_port):
raise RuntimeError(
f"Scheduler port {self.scheduler_port} is unavailable and --strict-ports is enabled. "
f"Either use a different port or remove --strict-ports to allow auto-selection."
)
if self.master_port is not None and not is_port_available(self.master_port):
raise RuntimeError(
f"Master port {self.master_port} is unavailable and --strict-ports is enabled. "
f"Either use a different port or remove --strict-ports to allow auto-selection."
)
else:
self.port = self.settle_port(self.port) self.port = self.settle_port(self.port)
initial_scheduler_port = self.scheduler_port + ( initial_scheduler_port = self.scheduler_port + (
random.randint(0, 100) if self.scheduler_port == 5555 else 0 random.randint(0, 100) if self.scheduler_port == 5555 else 0
@@ -778,6 +799,12 @@ class ServerArgs:
default=ServerArgs.port, default=ServerArgs.port,
help="Port for the HTTP API server.", help="Port for the HTTP API server.",
) )
parser.add_argument(
"--strict-ports",
action=StoreBoolean,
default=ServerArgs.strict_ports,
help="If enabled, fail when requested ports are unavailable instead of auto-selecting.",
)
parser.add_argument( parser.add_argument(
"--webui", "--webui",
action=StoreBoolean, action=StoreBoolean,