[Improve] Free routed experts ray obj ref to avoid memory leak

[RangiLyu]

No description provided.

[YanhuiDua]

这个是修改是因为ray的api有更新吗？

[RangiLyu]

调用外部库的私有方法不安全，统一改了

[YanhuiDua]

ray文档说在释放显式序列化的object.ref时，需要显式调用ray._private.internal_api.free(obj_ref)来释放 https://docs.ray.io/en/latest/ray-core/objects/serialization.html#serializing-objectrefs

[RangiLyu]

https://github.com/ray-project/ray/blob/cb9c80fee6a700efe61ea97987248ce82e3fa2e2/python/ray/internal/__init__.py#L1

这两者等价的

[YanhuiDua]

@claude review

[claude]

Claude: Warning: The del then re-assign pattern here is functionally a no-op — del dict[key] followed by dict[key] = value is equivalent to just dict[key] = value. This works, but the del adds no value since you're immediately re-assigning to the same key.

More importantly, ray.internal.free is called on the ObjectRef before the key is deleted or reassigned. At this point routed_experts (the local variable) already holds the resolved tensor from ray.get(), so the ObjectRef being freed is the right one. The logic is correct, but the del/reassign could be simplified to just:

Suggested change

	routed_experts = ray.get(data_item.env.rollout.extra_info["routed_experts"])
	ray.internal.free(data_item.env.rollout.extra_info["routed_experts"], local_only=False)
	del data_item.env.rollout.extra_info["routed_experts"]
	data_item.env.rollout.extra_info["routed_experts"] = routed_experts
	routed_experts = ray.get(data_item.env.rollout.extra_info["routed_experts"])
	ray.internal.free(data_item.env.rollout.extra_info["routed_experts"], local_only=False)
	data_item.env.rollout.extra_info["routed_experts"] = routed_experts

[claude]

Claude: Critical: Good fix — the old code mutated self.sample_params and self.extra_params in place on every rollout() call, which meant evaluator's disable_routed_experts=True would permanently pollute the shared state for subsequent training rollouts. Using shallow copies here is the correct approach.

One note: since sample_params.dict() returns a new dict, the {**self.sample_params, ...} merge is correct. But if self.sample_params or self.extra_params contain mutable values (e.g., nested dicts/lists), downstream code could still mutate those shared objects through the shallow copy. If that's a concern, copy.deepcopy would be safer, though the shallow copy is likely fine for the current usage.

[claude]

Claude: Nit: The docstring update for extra_params mentions "disable_routed_experts=True" as a dict key, which leaks an implementation detail into the environment interface. Consider describing it more generically, e.g., "Extra parameters for generation, passed through to the rollout controller."

[claude]

Claude: Warning: This changes the behavior of _save_trajectories — previously it would raise a KeyError if "ground_truth" was missing from reward_model. Now it silently saves an empty string. If the missing key indicates a data integrity issue, this could mask bugs. Consider logging a warning when the key is missing, or documenting why this fallback is expected.

[claude]

Claude: Nit: Good use of .pop() here — it removes the key from meta_info in one step, which is cleaner than the previous access-then-leave pattern. Consistent with the memory leak fix theme.

[claude]

Claude: ## Summary

This PR makes two logical changes: (1) fixes memory leaks by properly freeing Ray ObjectRefs for routed experts, and (2) disables routed expert tracking during evaluation to reduce overhead.

Issues

Critical

• [controller.py:371-373] The fix to use local copies instead of mutating self.sample_params/self.extra_params is the most important change in this PR. Without it, the evaluator's disable_routed_experts=True would permanently pollute shared state for training rollouts. Good catch.

Warning

• [rl_trainer.py:896] The reward_model.get("ground_truth", "") change silently handles missing keys that previously would raise KeyError. If missing ground_truth indicates a data integrity bug, this could mask issues. Consider adding a warning log.
• [replay_buffer.py:465-468] The del + re-assign pattern (del dict[key] then dict[key] = value) is functionally equivalent to just dict[key] = value. The del is unnecessary.

Nit

• [single_turn_env.py:87] Docstring leaks implementation detail (disable_routed_experts=True). Consider a more generic description.
• PR bundles two distinct changes (memory leak fix + evaluation optimization). These could have been separate PRs per project guidelines, but the changes are cohesive enough that this is minor.

Verdict

APPROVE — the core changes are correct and address real issues (memory leak via unfree'd ObjectRefs, state pollution via in-place dict mutation). The migration from ray._private.internal_api.free to ray.internal.free uses the proper public API. Minor suggestions above are non-blocking.