Feat: OpenAI TITIO support by middleware

.github/workflows/pr-test.yml

@@ -25,6 +25,46 @@ concurrency:
 
 jobs:
 
+  fast:
+    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request)
+    runs-on: self-hosted
+    container:
+      image: radixark/miles:latest
+      options: >
+        --gpus all
+        --ipc=host
+        --shm-size=16g
+        --ulimit memlock=-1
+        --ulimit stack=67108864
+        --memory=0
+        --memory-swap=0
+        -v /mnt/nvme0n1/miles_ci:/data/miles_ci
+        -v /mnt/nvme0n1/miles_ci/models:/root/models
+        -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
+    strategy:
+      fail-fast: false
+      matrix:
+        info: [{"num_gpus": 0, "test_file": "fast"}]
+    defaults:
+      run:
+        working-directory: ${{ github.workspace }}
+    env:
+      GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+      MILES_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }}
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Install
+        shell: bash
+        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages
+
+      - name: Execute
+        shell: bash
+        run: python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- pytest tests/${{ matrix.info.test_file }}
+
   e2e-test-short:
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-short'))
     runs-on: self-hosted
@@ -38,10 +78,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -82,10 +118,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -126,10 +158,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -170,10 +198,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -214,10 +238,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -258,10 +278,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -302,10 +318,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets

.github/workflows/pr-test.yml.j2

@@ -1,4 +1,10 @@
 <% set jobs = {
+    'fast': {
+      'test_executor': 'pytest',
+      'tests': [
+        {'test_file': 'fast', 'num_gpus': 0},
+      ],
+    },
     'e2e-test-short': {
       'label': 'run-ci-short',
       'tests': [
@@ -95,7 +101,7 @@ concurrency:
 jobs:
 <% for job_name, config in jobs.items() %>
   << job_name >>:
-    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, '<< config.label >>'))
+    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request<% if config.label %> && contains(github.event.pull_request.labels.*.name, '<< config.label >>')<% endif %>)
     runs-on: self-hosted
     container:
       image: << config.image if config.image else 'radixark/miles:latest' >>
@@ -107,10 +113,6 @@ jobs:
         --ulimit stack=67108864
         --memory=0
         --memory-swap=0
-        -e http_proxy=$http_proxy
-        -e https_proxy=$https_proxy
-        -e HTTP_PROXY=$HTTP_PROXY
-        -e HTTPS_PROXY=$HTTPS_PROXY
         -v /mnt/nvme0n1/miles_ci:/data/miles_ci
         -v /mnt/nvme0n1/miles_ci/models:/root/models
         -v /mnt/nvme0n1/miles_ci/datasets:/root/datasets
@@ -136,5 +138,5 @@ jobs:
 
       - name: Execute
         shell: bash
-        run: python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python tests/${{ matrix.info.test_file }}
+        run: python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- << config.test_executor | default('python') >> tests/${{ matrix.info.test_file }}
 <% endfor %>

examples/openai_format/__init__.py

@@ -0,0 +1 @@
+"""OpenAI format examples."""

examples/openai_format/dapo_math.py

@@ -0,0 +1,57 @@
+"""
+DAPO math OpenAI format example for token in/out verification.
+"""
+
+import argparse
+from typing import Any
+
+from openai import AsyncOpenAI
+
+from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
+from miles.rollout.generate_utils.openai_endpoint_utils import (
+    OpenAIEndpointTracer,
+    compute_samples_from_openai_records,
+)
+from miles.rollout.generate_utils.sample_utils import merge_samples
+
+_DAPO_MATH_SYSTEM_PROMPT = (
+    "Solve the math problem and return the final answer as \\boxed{integer}. "
+    "Keep the reasoning concise and finish with the boxed answer."
+)
+
+
+async def generate(input: GenerateFnInput) -> GenerateFnOutput:
+    tracer = await OpenAIEndpointTracer.create(input.args)
+    messages = _normalize_prompt(input.sample.prompt)
+    await _run_single_turn_openai(base_url=tracer.base_url, messages=messages)
+
+    records = await tracer.collect_records()
+    samples = compute_samples_from_openai_records(input.sample, records, input.state.tokenizer)
+    if not input.args.generate_multi_samples:
+        samples = merge_samples(samples, input.state.tokenizer)
+    return GenerateFnOutput(samples=samples)
+
+
+def _add_arguments(parser: argparse.ArgumentParser):
+    parser.add_argument("--generate-multi-samples", action="store_true")
+
+
+generate.add_arguments = _add_arguments
+
+
+def build_dapo_math_messages(question: str) -> list[dict[str, str]]:
+    return [
+        {"role": "system", "content": _DAPO_MATH_SYSTEM_PROMPT},
+        {"role": "user", "content": question},
+    ]
+
+
+def _normalize_prompt(prompt: Any) -> list[dict[str, Any]]:
+    if isinstance(prompt, list):
+        return prompt
+    return build_dapo_math_messages(prompt)
+
+
+async def _run_single_turn_openai(base_url: str, messages: list[dict[str, Any]]) -> None:
+    client = AsyncOpenAI(base_url=base_url, api_key="empty")
+    await client.chat.completions.create(model="default", messages=messages)

miles/ray/rollout.py

@@ -13,8 +13,15 @@
 from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
 from miles.backends.sglang_utils.sglang_engine import SGLangEngine
-from miles.rollout.base_types import call_rollout_fn
+from miles.rollout.base_types import (
+    RolloutFnConstructorInput,
+    RolloutFnEvalInput,
+    RolloutFnTrainInput,
+    call_rollout_fn,
+)
+from miles.rollout.inference_rollout.compatibility import call_rollout_function, load_rollout_function
 from miles.utils import tracking_utils
+from miles.utils.environ import enable_experimental_rollout_refactor
 from miles.utils.health_monitor import RolloutHealthMonitor
 from miles.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client
 from miles.utils.iter_utils import group_by
@@ -53,8 +60,14 @@ def __init__(self, args, pg):
         data_source_cls = load_function(self.args.data_source_path)
         self.data_source = data_source_cls(args)
 
-        self.generate_rollout = load_function(self.args.rollout_function_path)
-        self.eval_generate_rollout = load_function(self.args.eval_function_path)
+        self.use_experimental_refactor = enable_experimental_rollout_refactor()
+        if self.use_experimental_refactor:
+            input = RolloutFnConstructorInput(args=args, data_source=self.data_source)
+            self.generate_rollout = load_rollout_function(input, self.args.rollout_function_path)
+            self.eval_generate_rollout = load_rollout_function(input, self.args.eval_function_path)
+        else:
+            self.generate_rollout = load_function(self.args.rollout_function_path)
+            self.eval_generate_rollout = load_function(self.args.eval_function_path)
         self.custom_reward_post_process_func = None
         if self.args.custom_reward_post_process_path is not None:
             self.custom_reward_post_process_func = load_function(self.args.custom_reward_post_process_path)
@@ -142,7 +155,12 @@ def eval(self, rollout_id):
             return
         self.health_monitoring_resume()
 
-        result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True)
+        if self.use_experimental_refactor:
+            result = call_rollout_function(self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id))
+        else:
+            result = call_rollout_fn(
+                self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True
+            )
         data = result.data
         self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True)
         metrics = _log_eval_rollout_data(rollout_id, self.args, data, result.metrics)
@@ -224,7 +242,12 @@ def _get_rollout_data(self, rollout_id):
                 )
             metrics = None
         else:
-            data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False)
+            if self.use_experimental_refactor:
+                data = call_rollout_function(self.generate_rollout, RolloutFnTrainInput(rollout_id=rollout_id))
+            else:
+                data = call_rollout_fn(
+                    self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False
+                )
             metrics = data.metrics
             data = data.samples
             # flatten the data if it is a list of lists

miles/rollout/base_types.py

@@ -1,22 +1,86 @@
+from __future__ import annotations
+
+from argparse import Namespace
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
+from miles.rollout.data_source import DataSource
 from miles.utils.types import Sample
 
+if TYPE_CHECKING:
+    from miles.rollout.inference_rollout.inference_rollout_common import GenerateState
+
+
+@dataclass(frozen=True)
+class RolloutFnConstructorInput:
+    args: Namespace
+    # TODO may refactor DataSource API
+    data_source: DataSource
+
+
+@dataclass(frozen=True)
+class RolloutFnBaseInput:
+    rollout_id: int
+
+    @property
+    def evaluation(self):
+        raise NotImplementedError
+
+
+# subclassing for different data in the future
+@dataclass(frozen=True)
+class RolloutFnTrainInput(RolloutFnBaseInput):
+    @property
+    def evaluation(self):
+        return False
+
 
+@dataclass(frozen=True)
+class RolloutFnEvalInput(RolloutFnBaseInput):
+    @property
+    def evaluation(self):
+        return True
+
+
+# TODO make it frozen
 @dataclass
 class RolloutFnTrainOutput:
     samples: list[list[Sample]]
     metrics: dict[str, Any] = None
 
 
+# TODO make it frozen
 @dataclass
 class RolloutFnEvalOutput:
     data: dict[str, dict[str, Any]]
     metrics: dict[str, Any] = None
 
 
+RolloutFnInput = RolloutFnTrainInput | RolloutFnEvalInput
+RolloutFnOutput = RolloutFnTrainOutput | RolloutFnEvalOutput
+
+
+@dataclass(frozen=True)
+class GenerateFnInput:
+    state: GenerateState
+    sample: Sample
+    sampling_params: dict[str, Any]
+    evaluation: bool
+
+    @property
+    def args(self) -> Namespace:
+        return self.state.args
+
+
+@dataclass(frozen=True)
+class GenerateFnOutput:
+    # One generate may lead to multiple samples, such as multi-agent, tree-like exploration, or
+    # multi-turn with removing thinking tokens.
+    samples: Sample | list[Sample]
+
+
 def call_rollout_fn(fn, *args, evaluation: bool, **kwargs):
+    """Legacy rollout function call interface. Used when MILES_EXPERIMENTAL_ROLLOUT_REFACTOR is disabled."""
     output = fn(*args, **kwargs, evaluation=evaluation)
 
     # compatibility for legacy version