Merge branch 'main' into dev/cyx/rec_gspo_variant

yanxi-chen · yanxi-chen · commit c389ba6012af · 2025-12-29T18:26:05.000+08:00
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
@@ -21,7 +21,7 @@ env:
 jobs:
   build-and-push-image:
     runs-on: self-hosted
-    timeout-minutes: 240  # wait up to 4 hours
+    timeout-minutes: 480  # wait up to 8 hours
     # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.
     permissions:
       contents: read
@@ -61,8 +61,8 @@ jobs:
         with:
           context: trinity-${{ github.run_id }}
           push: true
-          file: trinity-${{ github.run_id }}/scripts/docker/Dockerfile
-          shm-size: 64g
+          file: trinity-${{ github.run_id }}/scripts/docker/Dockerfile.uv
+          shm-size: 128g
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
 
diff --git a/.github/workflows/docker/docker-compose.yaml b/.github/workflows/docker/docker-compose.yaml
@@ -1,8 +1,8 @@
 services:
   trinity-node-1:
-    image: trinity-rft-unittest:20251030
+    image: trinity-rft-unittest:20251225
     pull_policy: never
-    command: sh -c "pip install -e .[dev] && ray start --head --dashboard-host 0.0.0.0 --include-dashboard true --block"
+    command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --head --dashboard-host 0.0.0.0 --include-dashboard true --block"
     environment:
       - HF_ENDPOINT=https://hf-mirror.com
       - RAY_ADDRESS=auto
@@ -13,6 +13,7 @@ services:
       - TRINITY_MODEL_PATH=/mnt/models/Qwen3-0.6B
       - TRINITY_API_MODEL_PATH=/mnt/models/Qwen3-1.7B
       - TRINITY_VLM_MODEL_PATH=/mnt/models/Qwen2.5-VL-3B
+      - VIRTUAL_ENV=/opt/venv
     working_dir: /workspace
     networks:
       - trinity-network
@@ -29,14 +30,15 @@ services:
             capabilities: [gpu]
 
   trinity-node-2:
-    image: trinity-rft-unittest:20251030
+    image: trinity-rft-unittest:20251225
     pull_policy: never
-    command: sh -c "pip install -e .[dev] && ray start --address=trinity-node-1:6379 --block"
+    command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --address=trinity-node-1:6379 --block"
     environment:
       - HF_ENDPOINT=https://hf-mirror.com
       - TRINITY_CHECKPOINT_ROOT_DIR=/mnt/checkpoints
       - TRINITY_TASKSET_PATH=/mnt/data
       - TRINITY_MODEL_PATH=/mnt/models/Qwen3-1.7B
+      - VIRTUAL_ENV=/opt/venv
     working_dir: /workspace
     volumes:
       - trinity-volume:/mnt
diff --git a/.github/workflows/pre-commit.yaml b/.github/workflows/pre-commit.yaml
@@ -9,5 +9,5 @@ jobs:
     - uses: actions/checkout@v4
     - uses: actions/setup-python@v4
       with:
-        python-version: '3.10'
+        python-version: '3.12'
     - uses: pre-commit/action@v3.0.0
diff --git a/.github/workflows/sphinx-doc.yaml b/.github/workflows/sphinx-doc.yaml
@@ -16,10 +16,10 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ['3.10']
+        python-version: ['3.12']
     env:
       OS: ${{ matrix.os }}
-      PYTHON: '3.10'
+      PYTHON: '3.12'
     steps:
     - name: Free up disk space
       run: |
diff --git a/.github/workflows/unittest.yaml b/.github/workflows/unittest.yaml
@@ -34,7 +34,10 @@ jobs:
         MAX_RETRIES=20
         RETRY_INTERVAL=5
         for i in $(seq 1 $MAX_RETRIES); do
-          docker compose exec trinity-node-1 ray status && docker compose exec trinity-node-2 ray status && break
+          if docker compose exec trinity-node-1 bash -c "source /opt/venv/bin/activate && ray status" \
+            && docker compose exec trinity-node-2 bash -c "source /opt/venv/bin/activate && ray status"; then
+            break
+          fi
           echo "Waiting for ray cluster to be ready... ($i/$MAX_RETRIES)"
           sleep $RETRY_INTERVAL
           if [ "$i" -eq "$MAX_RETRIES" ]; then
@@ -76,12 +79,12 @@ jobs:
         TYPE="${{ steps.test_type.outputs.type }}"
         if [ "$TYPE" = "all" ]; then
           echo "tests_run=true" >> $GITHUB_ENV
-          docker compose exec trinity-node-1 pytest tests -v -s --ctrf report.json
+          docker compose exec trinity-node-1 bash -c "source /opt/venv/bin/activate && pytest tests -v -s --ctrf report.json"
         elif [ "$TYPE" = "diff" ]; then
           if [ -s ../../../test_dirs.txt ]; then
             echo "tests_run=true" >> $GITHUB_ENV
             TEST_DIRS=$(cat ../../../test_dirs.txt | xargs)
-            docker compose exec trinity-node-1 pytest $TEST_DIRS -v -s --ctrf report.json
+            docker compose exec trinity-node-1 bash -c "source /opt/venv/bin/activate && pytest $TEST_DIRS -v -s --ctrf report.json"
           else
             echo "No changed modules detected, skipping tests."
             echo "tests_run=false" >> $GITHUB_ENV
@@ -90,7 +93,7 @@ jobs:
           MODULE="${{ steps.test_type.outputs.module }}"
           if [ -n "$MODULE" ]; then
             echo "tests_run=true" >> $GITHUB_ENV
-            docker compose exec trinity-node-1 pytest tests/$MODULE -v -s --ctrf report.json
+            docker compose exec trinity-node-1 bash -c "source /opt/venv/bin/activate && pytest tests/$MODULE -v -s --ctrf report.json"
           else
             echo "No module specified, skipping tests."
             echo "tests_run=false" >> $GITHUB_ENV
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -15,7 +15,7 @@ repos:
     rev: 23.7.0
     hooks:
     - id: black
-      language_version: python3.10
+      language_version: python3.12
       args: [--line-length=100]
 
   - repo: https://github.com/pycqa/isort
diff --git a/README.md b/README.md
@@ -183,7 +183,7 @@ Choose one of the following options:
 ###### Using Conda
 
 ```bash
-conda create -n trinity python=3.10
+conda create -n trinity python=3.12
 conda activate trinity
 
 pip install -e ".[dev]"
diff --git a/README_zh.md b/README_zh.md
@@ -182,7 +182,7 @@ cd Trinity-RFT
 #### 使用 Conda
 
 ```bash
-conda create -n trinity python=3.10
+conda create -n trinity python=3.12
 conda activate trinity
 
 pip install -e ".[dev]"
diff --git a/docs/sphinx_doc/source/tutorial/trinity_installation.md b/docs/sphinx_doc/source/tutorial/trinity_installation.md
@@ -29,7 +29,7 @@ Choose one of the following options:
 #### Using Conda
 
 ```bash
-conda create -n trinity python=3.10
+conda create -n trinity python=3.12
 conda activate trinity
 
 pip install -e ".[dev]"
diff --git a/docs/sphinx_doc/source_zh/tutorial/trinity_installation.md b/docs/sphinx_doc/source_zh/tutorial/trinity_installation.md
@@ -29,7 +29,7 @@ cd Trinity-RFT
 #### 使用 Conda
 
 ```bash
-conda create -n trinity python=3.10
+conda create -n trinity python=3.12
 conda activate trinity
 
 pip install -e ".[dev]"
diff --git a/environments/data.yaml b/environments/data.yaml
@@ -2,7 +2,7 @@ name: trinity_data
 channels:
   - defaults
 dependencies:
-  - python=3.10
+  - python=3.12
   - pip:
     - py-data-juicer
     - agentscope
diff --git a/environments/training.yaml b/environments/training.yaml
@@ -2,6 +2,6 @@ name: trinity
 channels:
   - defaults
 dependencies:
-  - python=3.10
+  - python=3.12
   - pip:
     - "-e ..[dev]"
diff --git a/examples/agentscope_react/gsm8k.yaml b/examples/agentscope_react/gsm8k.yaml
@@ -29,7 +29,7 @@ buffer:
         response_key: 'answer'
       rollout_args:
         temperature: 1.0
-      default_workflow_type: 'as_react_workflow'
+      default_workflow_type: 'agentscope_react_workflow'
     eval_tasksets: []
   trainer_input:
     experience_buffer:
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name =  "trinity-rft"
-version = "0.3.3"
+version = "0.4.0"
 authors = [
     {name="Trinity-RFT Team", email="trinity-rft@outlook.com"},
 ]
@@ -43,6 +43,7 @@ dependencies = [
     "sortedcontainers",
     "word2number",
     "transformers",
+    "tinker",
 ]
 
 [project.scripts]
@@ -73,6 +74,8 @@ dev = [
 ]
 megatron = [
     "megatron-core[mlm]==0.13.1",
+    # if you found "undefined symbol" error in transformer engine
+    # reinstall it with --no-build-isolation and `--no-cache-dir` flag
     "transformer_engine[pytorch]==2.8.0",
     "mbridge>=0.13.0",
 ]
diff --git a/scripts/docker/Dockerfile.uv b/scripts/docker/Dockerfile.uv
@@ -39,8 +39,9 @@ RUN . /opt/venv/bin/activate && \
 
 # Install flash_attn and Megatron
 RUN . /opt/venv/bin/activate && \
-    uv pip install flash_attn==2.8.1 --no-cache-dir && \
+    uv pip install flash_attn==2.8.1 --no-build-isolation && \
     uv pip install -e .[megatron] && \
+    uv pip install --reinstall transformer_engine[pytorch]==2.8.0 --no-build-isolation --no-cache-dir && \
     NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 \
     uv pip install -v --no-build-isolation \
     --config-settings="--build-option=--cpp_ext" \
diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py
@@ -1219,3 +1219,91 @@ async def test_generate(self):
             response.prompt_length, 40960
         )  # If not long enough, please add more files to prompt
         self.assertGreater(response.logprobs.shape[0], 1000)
+
+
+class TestTinkerAPI(RayUnittestBaseAysnc):
+    """Test the Tinker API integration with the vLLM engine."""
+
+    def setUp(self):
+        self.config = get_template_config()
+        self.config.mode = "explore"
+        self.config.model.model_path = get_model_path()
+        self.config.explorer.rollout_model.engine_type = "vllm"
+        self.config.explorer.rollout_model.engine_num = 1
+        self.config.explorer.rollout_model.tensor_parallel_size = 1
+        self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE
+        self.config.explorer.rollout_model.enable_openai_api = True
+
+        self.config.check_and_update()
+        self.engines, self.auxiliary_engines = create_inference_models(self.config)
+        self.model_wrapper = ModelWrapper(self.engines[0], engine_type="vllm", enable_history=True)
+
+    async def test_tinker_api(self):
+        from tinker import types
+        from transformers import AutoTokenizer
+
+        engine = self.engines[0]
+        tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path)
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What is your name?"},
+        ]
+        result_dict = tokenizer.apply_chat_template(
+            messages,
+            chat_template=CHAT_TEMPLATE,
+            add_generation_prompt=False,
+            padding=False,
+            truncation=True,
+            return_tensors="pt",
+            add_special_tokens=False,
+            return_assistant_tokens_mask=True,
+            return_dict=True,
+        )
+        prompt = types.ModelInput.from_ints(
+            result_dict["input_ids"][0].tolist(),
+        )
+        # sample api without prompt logprobs
+        num_samples = 4
+        response = await engine.sample.remote(
+            prompt=prompt,
+            num_samples=num_samples,
+            sampling_params=types.SamplingParams(temperature=0.7),  # no limit on length
+        )
+        self.assertEqual(len(response.sequences), num_samples)
+        for sequence in response.sequences:
+            self.assertEqual(len(sequence.tokens), len(sequence.logprobs))
+            self.assertEqual(sequence.stop_reason, "stop")
+        self.assertIsNone(response.prompt_logprobs)
+        self.assertIsNone(response.topk_prompt_logprobs)
+        # sample api with prompt logprobs
+        num_samples = 2
+        topk_prompt_logprobs = 3
+        response = await engine.sample.remote(
+            prompt=prompt,
+            num_samples=num_samples,
+            sampling_params=types.SamplingParams(temperature=0.7, max_tokens=8),
+            include_prompt_logprobs=True,
+            topk_prompt_logprobs=topk_prompt_logprobs,
+        )
+        self.assertEqual(len(response.sequences), num_samples)
+        for sequence in response.sequences:
+            self.assertEqual(len(sequence.tokens), len(sequence.logprobs))
+            self.assertEqual(sequence.stop_reason, "length")
+        self.assertEqual(len(response.prompt_logprobs), len(prompt.to_ints()))
+        self.assertIsNone(response.prompt_logprobs[0])
+        self.assertEqual(len(response.topk_prompt_logprobs), len(prompt.to_ints()))
+        self.assertIsNone(response.topk_prompt_logprobs[0])
+        for topk_logprobs in response.topk_prompt_logprobs[1:]:
+            self.assertIsNotNone(topk_logprobs)
+            self.assertEqual(len(topk_logprobs), topk_prompt_logprobs)
+        # compute_logprob api
+        response = await engine.sample.remote(
+            prompt=prompt,
+            num_samples=1,
+            sampling_params=types.SamplingParams(max_tokens=1),
+            include_prompt_logprobs=True,
+        )
+        self.assertEqual(len(response.sequences), 1)
+        self.assertEqual(response.sequences[0].stop_reason, "length")
+        self.assertEqual(len(prompt.to_ints()), len(response.prompt_logprobs))
+        self.assertIsNone(response.topk_prompt_logprobs)
diff --git a/tests/utils/swanlab_test.py b/tests/utils/swanlab_test.py
@@ -5,6 +5,9 @@
 class TestSwanlabMonitor(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
+        cls._original_env = {
+            "SWANLAB_API_KEY": os.environ.get("SWANLAB_API_KEY"),
+        }
         os.environ["SWANLAB_API_KEY"] = "xxxxxxxxxxxxxxxxxxxxx"
 
     @classmethod
@@ -31,7 +34,3 @@ def test_swanlab_monitor_smoke(self):
         # Log a minimal metric to verify basic flow
         mon.log({"smoke/metric": 1.0}, step=1)
         mon.close()
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/trinity/__init__.py b/trinity/__init__.py
@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
 """Trinity-RFT (Reinforcement Fine-Tuning)"""
 
-__version__ = "0.3.3"
+__version__ = "0.4.0"
diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py
diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py
diff --git a/trinity/common/workflows/agentscope_workflow.py b/trinity/common/workflows/agentscope_workflow.py
diff --git a/trinity/utils/dlc_utils.py b/trinity/utils/dlc_utils.py
