feat(rollout): add TransferQueue trajectory ingestion for closed-loop VIME training by aoshen02 · Pull Request #320 · vllm-project/vime

aoshen02 · 2026-07-07T03:58:46Z

Summary

Add TransferQueue trajectory ingestion for closed-loop VIME training.
Thread artifact-consumer plumbing through the rollout path.
Remove the unit-test files from this branch per request.

Duplicate-work check

Checked open PRs for aoshen02:feat/tq-artifact-consumer: none found.
Checked open PRs for huangyibo:feat/tq-artifact-consumer: none found in this repo.
Checked open PRs for tq artifact consumer: no matching duplicate found.

Tests

Not run locally after the test-file removals.

AI assistance

AI assistance was used to set up the local worktree, trim the tests, and publish this draft PR.

… VIME training - Add TransferQueue trajectory source support in RolloutManager - Keep normal VIME rollout-engine ownership with --tq-manage-rollout-servers - Decode vLLM trajectory artifacts into native VIME Samples - Add rollout dispatch and TransferQueue service helper scripts - Document the closed-loop artifact-transfer flow - Add focused tests for TQ args, consumers, trajectory conversion, and dispatch

gemini-code-assist

Code Review

This pull request integrates TransferQueue as a new rollout trajectory source for VIME, adding service scripts, a rollout dispatcher, and consumer-side trajectory parsers. The review feedback highlights several important issues: a tensor unwrapping bug in trajectory_artifact.py that causes errors for 1D tensors of length 1, potential resource leaks in rollout_dispatcher.py due to unhandled task cancellation in asyncio.gather, incorrect usage of kv_batch_get in transfer_queue_consumer.py, unhandled StopIteration exceptions in trajectory_source.py, and hardcoded user-specific absolute paths in start_services.sh and arguments.py that hinder portability.

gemini-code-assist · 2026-07-07T04:00:45Z

+        if value.ndim > 0 and value.shape[0] == 1:
+            return value[0]


The condition value.ndim > 0 incorrectly unwraps 1D tensors of shape (1,) into 0D scalar tensors. When _vector is subsequently called on this 0D tensor, it raises a ValueError because tensor.ndim != 1. Changing the condition to value.ndim > 1 ensures that 1D sequence tensors of length 1 are preserved as 1D vectors, while only batch dimensions of size 1 (e.g., shape (1, N)) are unwrapped.

Suggested change

if value.ndim > 0 and value.shape[0] == 1:

return value[0]

if value.ndim > 1 and value.shape[0] == 1:

return value[0]

gemini-code-assist · 2026-07-07T04:00:45Z

+        tasks = [
+            self._dispatch_one(client, semaphore, prompt, sample_index)
+            for prompt in prompts
+            for sample_index in range(self.config.n_samples_per_prompt)
+        ]
+        try:
+            return await asyncio.gather(*tasks)
+        finally:
+            if self._owns_client:
+                await client.aclose()
+                self._client = None


When asyncio.gather is called with raw coroutines, any exception raised by one task will propagate immediately, but the remaining coroutines will continue running in the background. When the finally block executes and closes the client, these background tasks will fail with closed-client errors or leak resources. Wrapping the coroutines in asyncio.create_task and explicitly cancelling any pending tasks in an except block prevents this.

Suggested change

tasks = [

self._dispatch_one(client, semaphore, prompt, sample_index)

for prompt in prompts

for sample_index in range(self.config.n_samples_per_prompt)

]

try:

return await asyncio.gather(*tasks)

finally:

if self._owns_client:

await client.aclose()

self._client = None

tasks = [

asyncio.create_task(

self._dispatch_one(client, semaphore, prompt, sample_index)

)

for prompt in prompts

for sample_index in range(self.config.n_samples_per_prompt)

]

try:

return await asyncio.gather(*tasks)

except Exception:

for task in tasks:

if not task.done():

task.cancel()

raise

finally:

if self._owns_client:

await client.aclose()

self._client = None

gemini-code-assist · 2026-07-07T04:00:45Z

+                fields = self._tq.kv_batch_get(
+                    keys=key,
+                    partition_id=partition,
+                    select_fields=self.config.data_fields,
+                )
+                tags = self._tq.kv_list(partition)
+                tag = tags[partition][key]
+                return TrajectoryArtifactV1Alpha1.from_transfer_queue(
+                    fields,
+                    tag,
+                )


Since kv_batch_get is a batch operation, the keys parameter likely expects an iterable of keys (e.g., [key]) rather than a single string key. Passing a string directly can cause Python to iterate over the characters of the string. Furthermore, if kv_batch_get returns a dictionary mapping keys to their respective fields, you should extract the fields for the specific key using fields[key].

Suggested change

fields = self._tq.kv_batch_get(

keys=key,

partition_id=partition,

select_fields=self.config.data_fields,

)

tags = self._tq.kv_list(partition)

tag = tags[partition][key]

return TrajectoryArtifactV1Alpha1.from_transfer_queue(

fields,

tag,

)

fields_dict = self._tq.kv_batch_get(

keys=[key],

partition_id=partition,

select_fields=self.config.data_fields,

)

fields = fields_dict[key]

tags = self._tq.kv_list(partition)

tag = tags[partition][key]

return TrajectoryArtifactV1Alpha1.from_transfer_queue(

fields,

tag,

)

gemini-code-assist · 2026-07-07T04:00:45Z

+        return [
+            trajectory_to_sample(
+                next(self._iterator),
+                require_reward=self.require_reward,
+            )
+            for _ in range(count)
+        ]


Calling next(self._iterator) inside a list comprehension can propagate a raw StopIteration (or a generic RuntimeError under PEP 479) if the stream runs out of elements. It is much safer and more informative to catch StopIteration explicitly and raise a descriptive RuntimeError indicating that the stream was exhausted prematurely.

Suggested change

return [

trajectory_to_sample(

next(self._iterator),

require_reward=self.require_reward,

)

for _ in range(count)

]

samples = []

for _ in range(count):

try:

artifact = next(self._iterator)

except StopIteration:

raise RuntimeError(

f"TransferQueue trajectory stream exhausted prematurely; "

f"expected {count} samples, but only got {len(samples)}."

) from None

samples.append(

trajectory_to_sample(

artifact,

require_reward=self.require_reward,

)

)

return samples

gemini-code-assist · 2026-07-07T04:00:45Z

+VENV="${VIME_VENV:-/mnt/data1/yibo/vime-workspace/.venv}"
+SERVICE_DIR="${TQ_SERVICE_DIR:-/mnt/data1/yibo/vime-workspace/services/transfer_queue}"
+RAY_NODE_IP="${RAY_NODE_IP:-172.16.1.248}"
+RAY_GCS_PORT="${RAY_GCS_PORT:-6380}"
+RAY_CLIENT_PORT="${RAY_CLIENT_PORT:-20001}"
+RAY_ADDRESS="${RAY_NODE_IP}:${RAY_GCS_PORT}"
+PID_FILE="${SERVICE_DIR}/transfer_queue.pid"
+LOG_FILE="${SERVICE_DIR}/transfer_queue.log"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SERVICE_SCRIPT="${SCRIPT_DIR}/tq_service.py"
+EXPORT_SCRIPT="${SCRIPT_DIR}/export_config.py"
+CLIENT_CONFIG="${SERVICE_DIR}/client_config.pkl"


The script hardcodes absolute paths to a specific user's home/data directory (/mnt/data1/yibo/vime-workspace/...). This makes the script completely non-portable and fragile across different environments or users. Defining SCRIPT_DIR and REPO_ROOT first allows resolving the defaults dynamically.

Suggested change

VENV="${VIME_VENV:-/mnt/data1/yibo/vime-workspace/.venv}"

SERVICE_DIR="${TQ_SERVICE_DIR:-/mnt/data1/yibo/vime-workspace/services/transfer_queue}"

RAY_NODE_IP="${RAY_NODE_IP:-172.16.1.248}"

RAY_GCS_PORT="${RAY_GCS_PORT:-6380}"

RAY_CLIENT_PORT="${RAY_CLIENT_PORT:-20001}"

RAY_ADDRESS="${RAY_NODE_IP}:${RAY_GCS_PORT}"

PID_FILE="${SERVICE_DIR}/transfer_queue.pid"

LOG_FILE="${SERVICE_DIR}/transfer_queue.log"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

SERVICE_SCRIPT="${SCRIPT_DIR}/tq_service.py"

EXPORT_SCRIPT="${SCRIPT_DIR}/export_config.py"

CLIENT_CONFIG="${SERVICE_DIR}/client_config.pkl"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

VENV="${VIME_VENV:-${REPO_ROOT}/.venv}"

SERVICE_DIR="${TQ_SERVICE_DIR:-${REPO_ROOT}/services/transfer_queue}"

RAY_NODE_IP="${RAY_NODE_IP:-172.16.1.248}"

RAY_GCS_PORT="${RAY_GCS_PORT:-6380}"

RAY_CLIENT_PORT="${RAY_CLIENT_PORT:-20001}"

RAY_ADDRESS="${RAY_NODE_IP}:${RAY_GCS_PORT}"

PID_FILE="${SERVICE_DIR}/transfer_queue.pid"

LOG_FILE="${SERVICE_DIR}/transfer_queue.log"

SERVICE_SCRIPT="${SCRIPT_DIR}/tq_service.py"

EXPORT_SCRIPT="${SCRIPT_DIR}/export_config.py"

CLIENT_CONFIG="${SERVICE_DIR}/client_config.pkl"

gemini-code-assist · 2026-07-07T04:00:45Z

+            parser.add_argument(
+                "--tq-service-config-path",
+                type=str,
+                default=(
+                    "/mnt/data1/yibo/vime-workspace/services/"
+                    "transfer_queue/client_config.pkl"
+                ),
+            )


The default path for --tq-service-config-path is hardcoded to a specific user's directory (/mnt/data1/yibo/...). To make the configuration portable, resolve the default path dynamically relative to the repository root.

parser.add_argument( "--tq-service-config-path", type=str, default=os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "services", "transfer_queue", "client_config.pkl", ), )

read-the-docs-community · 2026-07-07T04:00:48Z

Documentation build overview

📚 vime | 🛠️ Build #33471737 | 📁 Comparing 9237e03 against latest (e70b319)

🔍 Preview build

42 files changed · ± 26 modified · - 16 deleted

± Modified

- Deleted

huangyibo added 2 commits July 7, 2026 03:14

test: remove unit tests

9237e03

aoshen02 closed this Jul 7, 2026

aoshen02 reopened this Jul 7, 2026

gemini-code-assist Bot reviewed Jul 7, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(rollout): add TransferQueue trajectory ingestion for closed-loop VIME training#320

feat(rollout): add TransferQueue trajectory ingestion for closed-loop VIME training#320
aoshen02 wants to merge 2 commits into
vllm-project:mainfrom
aoshen02:feat/tq-artifact-consumer

aoshen02 commented Jul 7, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Uh oh!

read-the-docs-community Bot commented Jul 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

-        return [
-            trajectory_to_sample(
-                next(self._iterator),
-                require_reward=self.require_reward,
-            )
-            for _ in range(count)
-        ]
+        samples = []
+        for _ in range(count):
+            try:
+                artifact = next(self._iterator)
+            except StopIteration:
+                raise RuntimeError(
+                    f"TransferQueue trajectory stream exhausted prematurely; "
+                    f"expected {count} samples, but only got {len(samples)}."
+                ) from None
+            samples.append(
+                trajectory_to_sample(
+                    artifact,
+                    require_reward=self.require_reward,
+                )
+            )
+        return samples

Uh oh!

Conversation

aoshen02 commented Jul 7, 2026

Summary

Duplicate-work check

Tests

AI assistance

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

read-the-docs-community Bot commented Jul 7, 2026

Documentation build overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants