minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Release mm features on session close and support multiple /rerun-ut specs (#21501)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-03-26 18:31:29 -07:00
committed by GitHub
parent c2b3e42ad6
commit e1ee68d0fc
5 changed files with 140 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+5
View File
@@ -366,6 +366,11 @@ class MultimodalInputs:
mrope_position_delta: Optional[torch.Tensor] = None
mrope_position_delta_repeated_cache: Optional[torch.Tensor] = None
def release_features(self):
"""Release feature tensors to free GPU memory."""
for item in self.mm_items:
item.feature = None
@staticmethod
def from_dict(obj: dict):
# Check if MM splitting is enabled
python/sglang/srt/managers/scheduler.py
+1 -2
View File
@@ -1674,8 +1674,7 @@ class Scheduler(
if req.session:
continue
# For non-session requests, clear features and mm_inputs
for item in mm_inputs.mm_items:
item.feature = None
mm_inputs.release_features()
req.multimodal_inputs = None
def handle_generate_request(
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+1 -5
View File
@@ -438,11 +438,7 @@ class SchedulerOutputProcessorMixin:
if req.finished():
# delete feature to save memory
if req.multimodal_inputs is not None and req.session is None:
for mm_item in req.multimodal_inputs.mm_items:
pixel_values = mm_item.feature
if isinstance(pixel_values, torch.Tensor):
mm_item.feature = None
del pixel_values
req.multimodal_inputs.release_features()
self.maybe_collect_routed_experts(req)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
python/sglang/srt/managers/session_controller.py
+20 -1
View File
@@ -169,7 +169,14 @@ class Session:
if req.mm_inputs:
for item in req.mm_inputs.get("mm_items", []):
if item.offsets:
item.offsets = [(s - 1, e - 1) for s, e in item.offsets]
if any(s == 0 for s, _ in item.offsets):
logging.warning(
"mm_item offset starts at 0 (BOS position), "
"clamping to 0 after BOS strip"
)
item.offsets = [
(max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
]
input_ids = (
last_req.origin_input_ids
@@ -284,6 +291,18 @@ class SessionController:
req = next(iter(session.req_nodes.values())).req
if not req.finished():
req.session = None
# Release multimodal features held by session requests.
# Session reqs skip the normal mm cleanup path (scheduler and
# output_processor) so features stay alive until the session closes.
seen_mm = set()
for node in session.req_nodes.values():
mm = node.req.multimodal_inputs
if mm is not None and id(mm) not in seen_mm:
seen_mm.add(id(mm))
mm.release_features()
node.req.multimodal_inputs = None
if isinstance(self.tree_cache, SessionAwareCache):
self.tree_cache.release_session(session_id)
del self.sessions[session_id]
scripts/ci/utils/slash_command_handler.py
+99 -73
View File
@@ -515,47 +515,12 @@ def detect_cuda_suite(file_path_from_test):
return suite, runner, use_deepep, None
def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
"""
Handles the /rerun-ut <file>::<TestClass.test_method> command.
Dispatches a lightweight workflow to run a single test on the correct CUDA runner.
Resolve a single test spec and dispatch a workflow run.
Returns a dict with keys: spec, success, test_command, runner_label, run_url, error.
"""
# SECURITY: For fork PRs, only allow /rerun-ut if the commenter has write+ permission.
# This command checks out and executes code from the PR branch on self-hosted GPU
# runners, so we must ensure the commenter is a trusted collaborator.
is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
if is_fork:
commenter = comment.user.login
perm = gh_repo.get_collaborator_permission(commenter)
if perm not in ("admin", "write"):
print(f"Permission denied: /rerun-ut on fork PR by {commenter}.")
comment.create_reaction("confused")
pr.create_issue_comment(
"❌ `/rerun-ut` is not available for fork PRs unless the commenter "
"has write permission on the repo.\n\n"
"Please ask a maintainer to run this command, or use the normal CI flow."
)
return False
print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.")
if not (
user_perms.get("can_rerun_ut", False)
or user_perms.get("can_rerun_stage", False)
):
print("Permission denied: neither can_rerun_ut nor can_rerun_stage is true.")
return False
if not test_spec:
comment.create_reaction("confused")
pr.create_issue_comment(
"❌ Please specify a test: `/rerun-ut <file>::<TestClass.test_method>`\n\n"
"Examples:\n"
"- `/rerun-ut test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n"
"- `/rerun-ut registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n"
"- `/rerun-ut test_srt_endpoint.py`"
)
return False
# Parse spec: split on :: to get file path and optional test selector
if "::" in test_spec:
file_part, test_selector = test_spec.split("::", 1)
@@ -570,16 +535,12 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
# Resolve file path
resolved_path, err = resolve_test_file(file_part)
if err:
comment.create_reaction("confused")
pr.create_issue_comment(f"❌ {err}")
return False
return {"spec": test_spec, "success": False, "error": err}
# Detect suite and runner
suite, runner_label, use_deepep, err = detect_cuda_suite(resolved_path)
if err:
comment.create_reaction("confused")
pr.create_issue_comment(f"❌ {err}")
return False
return {"spec": test_spec, "success": False, "error": err}
# Build test_command: file path (+ optional test selector as unittest arg)
test_command = resolved_path
@@ -601,13 +562,15 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
break
if not target_workflow:
print(f"Error: {workflow_name} workflow not found")
return False
return {
"spec": test_spec,
"success": False,
"error": f"{workflow_name} workflow not found",
}
is_fork = (
pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
)
print(f"PR is from fork: {is_fork}")
pr_head_sha = None
if is_fork:
@@ -641,13 +604,14 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
success = dispatch_resp.status_code in (200, 204)
if not success:
print(f"Dispatch failed: {dispatch_resp.status_code} {dispatch_resp.text}")
return {
"spec": test_spec,
"success": False,
"error": f"Dispatch failed: {dispatch_resp.status_code}",
}
if success:
print(f"Successfully triggered rerun-ut: {test_command}")
comment.create_reaction("+1")
# Include test_command in expected title to distinguish
# concurrent /rerun-ut dispatches (run-name includes test_command)
run_url = find_workflow_run_url(
gh_repo,
target_workflow.id,
@@ -659,32 +623,95 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
max_wait=30,
test_command=test_command,
)
if run_url:
return {
"spec": test_spec,
"success": True,
"test_command": test_command,
"runner_label": runner_label,
"run_url": run_url,
}
except Exception as e:
print(f"Error triggering rerun-ut for {test_spec}: {e}")
return {"spec": test_spec, "success": False, "error": str(e)}
def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_specs, token):
"""
Handles the /rerun-ut command. Accepts a list of test specs and dispatches
a workflow run for each, posting a single consolidated comment.
"""
# SECURITY: For fork PRs, only allow /rerun-ut if the commenter has write+ permission.
# This command checks out and executes code from the PR branch on self-hosted GPU
# runners, so we must ensure the commenter is a trusted collaborator.
is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
if is_fork:
commenter = comment.user.login
perm = gh_repo.get_collaborator_permission(commenter)
if perm not in ("admin", "write"):
print(f"Permission denied: /rerun-ut on fork PR by {commenter}.")
comment.create_reaction("confused")
pr.create_issue_comment(
f"✅ Triggered `/rerun-ut` on `{runner_label}` runner:"
f" [View workflow run]({run_url})\n"
f"```\ncd test/ && python3 {test_command}\n```"
"❌ `/rerun-ut` is not available for fork PRs unless the commenter "
"has write permission on the repo.\n\n"
"Please ask a maintainer to run this command, or use the normal CI flow."
)
return False
print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.")
if not (
user_perms.get("can_rerun_ut", False)
or user_perms.get("can_rerun_stage", False)
):
print("Permission denied: neither can_rerun_ut nor can_rerun_stage is true.")
return False
if not test_specs:
comment.create_reaction("confused")
pr.create_issue_comment(
"❌ Please specify a test: `/rerun-ut <file>::<TestClass.test_method>`\n\n"
"Examples:\n"
"- `/rerun-ut test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n"
"- `/rerun-ut registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n"
"- `/rerun-ut test_srt_endpoint.py`\n"
"- `/rerun-ut test_a.py test_b.py test_c.py` (multiple tests)"
)
return False
results = []
for spec in test_specs:
results.append(_resolve_and_dispatch_ut(gh_repo, pr, spec, token))
# Build consolidated comment
successes = [r for r in results if r["success"]]
failures = [r for r in results if not r["success"]]
lines = []
for r in successes:
if r.get("run_url"):
lines.append(
f"✅ `{r['runner_label']}`: [View workflow run]({r['run_url']})\n"
f"```\ncd test/ && python3 {r['test_command']}\n```"
)
else:
pr.create_issue_comment(
f"✅ Triggered `/rerun-ut` on `{runner_label}` runner:\n"
f"```\ncd test/ && python3 {test_command}\n```\n"
lines.append(
f"✅ `{r['runner_label']}`:\n"
f"```\ncd test/ && python3 {r['test_command']}\n```\n"
f"⚠️ Could not retrieve workflow run URL. "
f"Check the [Actions tab](https://github.com/{gh_repo.full_name}/actions) for progress."
)
return True
else:
print("Failed to trigger workflow_dispatch")
return False
for r in failures:
lines.append(f"❌ `{r['spec']}`: {r['error']}")
except Exception as e:
print(f"Error triggering rerun-ut: {e}")
body = "\n\n".join(lines)
if successes:
comment.create_reaction("+1")
if failures and not successes:
comment.create_reaction("confused")
pr.create_issue_comment(
f"❌ Failed to trigger rerun-ut: {str(e)}\n\n"
f"Please check the logs or contact maintainers."
)
return False
pr.create_issue_comment(body)
return len(successes) > 0
def main():
@@ -769,9 +796,8 @@ def main():
handle_rerun_stage(repo, pr, comment, user_perms, stage_name, token)
elif first_line.startswith("/rerun-ut"):
parts = first_line.split(maxsplit=1)
test_spec = parts[1].strip() if len(parts) > 1 else None
handle_rerun_ut(repo, pr, comment, user_perms, test_spec, token)
test_specs = first_line.split()[1:]
handle_rerun_ut(repo, pr, comment, user_perms, test_specs or None, token)
else:
print(f"Unknown or ignored command: {first_line}")