From a3935df73f1218a20eb5261ef7b3007279dfb85d Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 08:44:03 +0000 Subject: [PATCH 1/2] Use shallow copy in _merge_sample_pair; deepcopy of token lists made merging quadratic --- miles/rollout/generate_utils/sample_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 7845731bed..4d8bee2f52 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -1,4 +1,4 @@ -from copy import deepcopy +from copy import copy from dataclasses import fields from miles.utils.types import Sample @@ -30,7 +30,7 @@ def merge_samples(samples: list[Sample], tokenizer) -> Sample: def _merge_sample_pair(a: Sample, b: Sample, tokenizer) -> Sample: """Merge two samples generated from sibling inference engine calls.""" - a, b = deepcopy(a), deepcopy(b) + a, b = copy(a), copy(b) def _merge_equal_value(field): x = getattr(a, field) From 2a6b456878500e79fc3c29ba72f6c04d37a59b70 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Fri, 3 Jul 2026 10:09:45 +0000 Subject: [PATCH 2/2] Shallow-copy input sample per record; only mutated fields get fresh copies --- miles/rollout/generate_utils/openai_endpoint_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index e03197d6a5..5551ca2f41 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -6,7 +6,7 @@ import logging import random from argparse import Namespace -from copy import deepcopy +from copy import copy, deepcopy from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response from miles.rollout.session.session_types import GetSessionResponse, SessionRecord @@ -184,7 +184,10 @@ def _compute_sample_from_openai_record( output_token_ids = [item[1] for item in choice["meta_info"]["output_token_logprobs"]] output_log_probs = [item[0] for item in choice["meta_info"]["output_token_logprobs"]] - sample = deepcopy(input_sample) + sample = copy(input_sample) + sample.metadata = dict(input_sample.metadata) + sample.weight_versions = list(input_sample.weight_versions) + sample.prefix_cache_info = deepcopy(input_sample.prefix_cache_info) request_input_ids = record.request.get("input_ids") if request_input_ids is not None: assert (