Skip to content

Name pending rollout tasks in abort logs #75

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
nightlessbaron wants to merge 1 commit into
base: prod
Choose a base branch
from
Draft

Name pending rollout tasks in abort logs #75

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
34 changes: 21 additions & 13 deletions miles/rollout/inference_rollout/inference_rollout_train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,9 +87,13 @@ async def abort(state: GenerateState, pendings: set, rollout_id: int) -> list[li
assert not state.aborted
state.aborted = True

# How many rollout tasks are still in flight when the abort fires. The
# specific instances harbor cancels are reported by /abort_all below.
logger.info(f"[abort] rollout_id={rollout_id} draining {len(pendings)} in-flight rollout tasks")
# How many rollout tasks are still in flight when the abort fires, and which
# groups they are. The specific instances harbor cancels are reported by
# /abort_all below.
cancelled_names = sorted(task.get_name() for task in pendings)
logger.info(
f"[abort] rollout_id={rollout_id} draining {len(pendings)} in-flight rollout tasks: {cancelled_names}"
)

is_agentic = bool(getattr(args, "use_session_server", False) and getattr(args, "custom_agent_function_path", None))
if is_agentic:
Expand Down Expand Up @@ -145,18 +149,22 @@ def stamp_rollout_id(samples: list[list[Sample]], rollout_id: int) -> None:


def submit_generate_tasks(state: GenerateState, samples: list[list[Sample]]):
return [
asyncio.create_task(
# submit a group of samples as a single task.
generate_and_rm_group(
state,
group,
sampling_params=state.sampling_params.copy(),
evaluation=False,
tasks = []
for group in samples:
first = group[0][0] if isinstance(group[0], list) else group[0]
tasks.append(
asyncio.create_task(
# submit a group of samples as a single task.
generate_and_rm_group(
state,
group,
sampling_params=state.sampling_params.copy(),
evaluation=False,
),
name=f"group-{first.index}",
)
)
for group in samples
]
return tasks


async def generate_rollout_async(
Expand Down
Loading