chore: import upstream snapshot with attribution

tests/utils/test_assemble_teacher_topk_logprobs.py

@@ -0,0 +1,131 @@
+"""Unit tests for ``assemble_teacher_topk_logprobs``.
+
+Covers padding_free (packed) and non-packed modes.
+"""
+import pytest
+import torch
+
+from swift.rlhf_trainers.utils import assemble_teacher_topk_logprobs
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_parsed(seq_len: int, topk: int):
+    """Create ``parsed`` data mimicking ``parse_prompt_logprobs`` output.
+
+    ``parse_prompt_logprobs`` skips position 0, so ``len(lps) == seq_len - 1``.
+    Each position gets ``topk`` values; we use distinct floats so we can verify
+    the mapping.
+    """
+    lps = []
+    ixs = []
+    for pos in range(seq_len - 1):  # skip position 0
+        lps.append([float(pos * 100 + k) for k in range(topk)])
+        ixs.append([pos * 10 + k for k in range(topk)])
+    return lps, ixs
+
+
+# ---------------------------------------------------------------------------
+# 1. padding_free=True (packed mode)
+# ---------------------------------------------------------------------------
+
+
+class TestPacked:
+
+    def test_single_sample(self):
+        """One sample, sequential mapping."""
+        topk = 2
+        seq_len = 6
+        parsed = [_make_parsed(seq_len, topk)]
+        cu_seqlens = [0, seq_len]
+
+        out_lp, out_ix = assemble_teacher_topk_logprobs(
+            parsed, batch_size=1, seq_len=seq_len, cu_seqlens=cu_seqlens, topk=topk, device=torch.device('cpu'))
+
+        assert out_lp.shape == (1, seq_len, topk)
+        # Positions 0..4 filled, position 5 = -inf
+        for i in range(seq_len - 1):
+            for k in range(topk):
+                assert out_lp[0, i, k].item() == pytest.approx(float(i * 100 + k))
+                assert out_ix[0, i, k].item() == i * 10 + k
+        assert torch.isinf(out_lp[0, seq_len - 1, 0])
+
+    def test_two_samples(self):
+        """Two samples packed together."""
+        topk = 1
+        s1, s2 = 4, 5
+        parsed = [_make_parsed(s1, topk), _make_parsed(s2, topk)]
+        cu_seqlens = [0, s1, s1 + s2]
+
+        out_lp, out_ix = assemble_teacher_topk_logprobs(
+            parsed, batch_size=1, seq_len=s1 + s2, cu_seqlens=cu_seqlens, topk=topk, device=torch.device('cpu'))
+
+        assert out_lp.shape == (1, s1 + s2, topk)
+        # Sample 1: positions 0..2 filled, position 3 = -inf
+        assert out_lp[0, 0, 0].item() == pytest.approx(0.0)
+        assert out_lp[0, 2, 0].item() == pytest.approx(200.0)
+        assert torch.isinf(out_lp[0, 3, 0])
+        # Sample 2: positions 4..7 filled, position 8 = -inf
+        assert out_lp[0, 4, 0].item() == pytest.approx(0.0)
+        assert out_lp[0, 7, 0].item() == pytest.approx(300.0)
+        assert torch.isinf(out_lp[0, 8, 0])
+
+
+# ---------------------------------------------------------------------------
+# 2. padding_free=False (non-packed mode)
+# ---------------------------------------------------------------------------
+
+
+class TestNonPacked:
+
+    def test_no_offset(self):
+        """Batch of 2, no left padding (offsets=0)."""
+        topk = 2
+        seq_len = 5
+        batch_size = 2
+        parsed = [_make_parsed(seq_len, topk), _make_parsed(seq_len, topk)]
+
+        out_lp, out_ix = assemble_teacher_topk_logprobs(
+            parsed, batch_size=batch_size, seq_len=seq_len, cu_seqlens=None, topk=topk, device=torch.device('cpu'))
+
+        assert out_lp.shape == (batch_size, seq_len, topk)
+        for b in range(batch_size):
+            lps = parsed[b][0]
+            for i in range(seq_len - 1):
+                for k in range(topk):
+                    assert out_lp[b, i, k].item() == pytest.approx(lps[i][k])
+            assert torch.isinf(out_lp[b, seq_len - 1, 0])
+
+    def test_with_offsets(self):
+        """Batch of 2 with left padding (offsets=[2, 0])."""
+        topk = 1
+        seq_len = 6
+        batch_size = 2
+        parsed = [_make_parsed(4, topk), _make_parsed(6, topk)]
+        offsets = [2, 0]
+
+        out_lp, out_ix = assemble_teacher_topk_logprobs(
+            parsed,
+            batch_size=batch_size,
+            seq_len=seq_len,
+            cu_seqlens=None,
+            topk=topk,
+            device=torch.device('cpu'),
+            offsets=offsets)
+
+        assert out_lp.shape == (batch_size, seq_len, topk)
+        # Sample 0: starts at offset 2, has 3 logprobs (4 tokens - 1)
+        lps0 = parsed[0][0]
+        assert out_lp[0, 2, 0].item() == pytest.approx(lps0[0][0])
+        assert out_lp[0, 4, 0].item() == pytest.approx(lps0[2][0])
+        assert torch.isinf(out_lp[0, 5, 0])  # last position for sample 0
+        assert torch.isinf(out_lp[0, 0, 0])  # left padding
+        assert torch.isinf(out_lp[0, 1, 0])  # left padding
+
+        # Sample 1: starts at offset 0, has 5 logprobs (6 tokens - 1)
+        lps1 = parsed[1][0]
+        assert out_lp[1, 0, 0].item() == pytest.approx(lps1[0][0])
+        assert out_lp[1, 4, 0].item() == pytest.approx(lps1[4][0])
+        assert torch.isinf(out_lp[1, 5, 0])

tests/utils/test_async_rewards.py

@@ -0,0 +1,222 @@
+"""Unit tests for async reward functions.
+
+Tests the async reward function implementation including:
+- AsyncORM class functionality
+- Event loop daemon thread management
+- Performance comparison between sync and async reward functions
+"""
+import asyncio
+import time
+import unittest
+from typing import List
+
+
+class TestAsyncRewardFunctions(unittest.TestCase):
+    """Test async reward function utilities and base class."""
+
+    def test_start_and_shutdown_event_loop_in_daemon(self):
+        """Test that we can start and shutdown an event loop in a daemon thread."""
+        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon
+
+        # Start the event loop
+        thread, loop, ready_event = start_event_loop_in_daemon(name='TestLoop')
+
+        # Wait for the loop to be ready
+        self.assertTrue(ready_event.wait(timeout=5))
+
+        # Verify the loop is running
+        self.assertTrue(loop.is_running())
+        self.assertTrue(thread.is_alive())
+
+        # Shutdown the event loop
+        shutdown_event_loop_in_daemon(thread, loop)
+
+        # Give the thread time to finish
+        thread.join(timeout=2)
+        self.assertFalse(thread.is_alive())
+
+    def test_run_async_function_in_daemon_loop(self):
+        """Test running an async function in the daemon event loop."""
+        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon
+
+        thread, loop, ready_event = start_event_loop_in_daemon(name='TestLoop')
+        ready_event.wait(timeout=5)
+
+        async def async_add(a, b):
+            await asyncio.sleep(0.01)  # Simulate async work
+            return a + b
+
+        # Run the async function in the daemon loop
+        future = asyncio.run_coroutine_threadsafe(async_add(2, 3), loop)
+        result = future.result(timeout=5)
+
+        self.assertEqual(result, 5)
+
+        shutdown_event_loop_in_daemon(thread, loop)
+
+    def test_async_orm_base_class(self):
+        """Test that AsyncORM can be subclassed and used correctly."""
+        from swift.rewards import AsyncORM
+
+        class TestAsyncReward(AsyncORM):
+
+            async def __call__(self, completions: List[str], **kwargs) -> List[float]:
+                await asyncio.sleep(0.01)
+                return [float(len(c)) for c in completions]
+
+        reward_func = TestAsyncReward()
+
+        # Check that it's detected as async
+        self.assertTrue(asyncio.iscoroutinefunction(reward_func.__call__))
+
+        # Run in an event loop to verify it works
+        async def run_test():
+            result = await reward_func(['hello', 'world!'])
+            return result
+
+        result = asyncio.get_event_loop().run_until_complete(run_test())
+        self.assertEqual(result, [5.0, 6.0])
+
+    def test_async_reward_is_detected(self):
+        """Test that async reward functions are correctly detected."""
+        from swift.rewards import AsyncORM
+
+        class SyncReward:
+
+            def __call__(self, completions, **kwargs):
+                return [1.0] * len(completions)
+
+        class AsyncReward(AsyncORM):
+
+            async def __call__(self, completions, **kwargs):
+                return [1.0] * len(completions)
+
+        sync_func = SyncReward()
+        async_func = AsyncReward()
+
+        # Check detection
+        self.assertFalse(asyncio.iscoroutinefunction(sync_func))
+        self.assertFalse(asyncio.iscoroutinefunction(sync_func.__call__))
+        self.assertTrue(asyncio.iscoroutinefunction(async_func.__call__))
+
+
+class TestAsyncRewardPerformance(unittest.TestCase):
+    """Test performance benefits of async reward functions."""
+
+    def test_parallel_async_execution(self):
+        """Test that multiple async reward functions execute in parallel."""
+        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon
+
+        thread, loop, ready_event = start_event_loop_in_daemon(name='PerfTestLoop')
+        ready_event.wait(timeout=5)
+
+        sleep_time = 0.1  # 100ms per call
+        num_calls = 5
+
+        async def slow_async_func(idx):
+            await asyncio.sleep(sleep_time)
+            return idx
+
+        # Test sequential execution time
+        start_seq = time.time()
+        for i in range(num_calls):
+            future = asyncio.run_coroutine_threadsafe(slow_async_func(i), loop)
+            future.result(timeout=5)
+        sequential_time = time.time() - start_seq
+
+        # Test parallel execution time using asyncio.gather
+        async def run_parallel():
+            tasks = [slow_async_func(i) for i in range(num_calls)]
+            return await asyncio.gather(*tasks)
+
+        start_par = time.time()
+        future = asyncio.run_coroutine_threadsafe(run_parallel(), loop)
+        results = future.result(timeout=5)
+        parallel_time = time.time() - start_par
+
+        shutdown_event_loop_in_daemon(thread, loop)
+
+        # Verify results
+        self.assertEqual(results, list(range(num_calls)))
+
+        # Parallel should be significantly faster than sequential
+        # Sequential: ~num_calls * sleep_time = 0.5s
+        # Parallel: ~sleep_time = 0.1s
+        # Allow some margin for overhead
+        self.assertLess(parallel_time, sequential_time * 0.5)
+
+        print('\nPerformance test results:')
+        print(f'  Sequential time: {sequential_time:.3f}s (expected ~{sleep_time * num_calls:.1f}s)')
+        print(f'  Parallel time:   {parallel_time:.3f}s (expected ~{sleep_time:.1f}s)')
+        print(f'  Speedup:         {sequential_time / parallel_time:.1f}x')
+
+    def test_async_reward_function_batch_performance(self):
+        """Test performance of async reward function with batch processing."""
+        from swift.rewards import AsyncORM
+        from swift.utils import shutdown_event_loop_in_daemon, start_event_loop_in_daemon
+
+        sleep_per_item = 0.05  # 50ms per item
+        batch_size = 8
+
+        class SlowSyncReward:
+
+            def __call__(self, completions, **kwargs):
+                rewards = []
+                for c in completions:
+                    time.sleep(sleep_per_item)  # Blocking sleep
+                    rewards.append(float(len(c)))
+                return rewards
+
+        class FastAsyncReward(AsyncORM):
+
+            async def __call__(self, completions, **kwargs):
+
+                async def score_single(text):
+                    await asyncio.sleep(sleep_per_item)  # Non-blocking sleep
+                    return float(len(text))
+
+                # Process all in parallel
+                tasks = [score_single(c) for c in completions]
+                return await asyncio.gather(*tasks)
+
+        completions = [f'text_{i}' for i in range(batch_size)]
+
+        # Test sync reward function
+        sync_reward = SlowSyncReward()
+        start_sync = time.time()
+        sync_results = sync_reward(completions)
+        sync_time = time.time() - start_sync
+
+        # Test async reward function
+        thread, loop, ready_event = start_event_loop_in_daemon(name='BatchPerfLoop')
+        ready_event.wait(timeout=5)
+
+        async_reward = FastAsyncReward()
+
+        async def run_async():
+            return await async_reward(completions)
+
+        start_async = time.time()
+        future = asyncio.run_coroutine_threadsafe(run_async(), loop)
+        async_results = future.result(timeout=10)
+        async_time = time.time() - start_async
+
+        shutdown_event_loop_in_daemon(thread, loop)
+
+        # Verify results are the same
+        self.assertEqual(len(sync_results), len(async_results))
+        self.assertEqual(sync_results, list(async_results))
+
+        # Async should be significantly faster
+        # Sync: ~batch_size * sleep_per_item = 0.4s
+        # Async: ~sleep_per_item = 0.05s
+        self.assertLess(async_time, sync_time * 0.5)
+
+        print('\nBatch processing performance:')
+        print(f'  Sync time:  {sync_time:.3f}s (expected ~{sleep_per_item * batch_size:.2f}s)')
+        print(f'  Async time: {async_time:.3f}s (expected ~{sleep_per_item:.2f}s)')
+        print(f'  Speedup:    {sync_time / async_time:.1f}x')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils/test_file_utils.py

@@ -0,0 +1,32 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+from swift.utils import copy_files_by_pattern
+
+
+class TestFileUtils(unittest.TestCase):
+
+    def setUp(self):
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = self._tmp_dir.name
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def test_copy_files(self):
+        os.makedirs(os.path.join(self.tmp_dir, 'source'))
+        os.makedirs(os.path.join(self.tmp_dir, 'source', 'subfolder'))
+        with open(os.path.join(self.tmp_dir, 'source', '1.txt'), 'w') as f:
+            f.write('')
+        with open(os.path.join(self.tmp_dir, 'source', 'subfolder', '2.txt'), 'w') as f:
+            f.write('')
+        copy_files_by_pattern(
+            os.path.join(self.tmp_dir, 'source'), os.path.join(self.tmp_dir, 'target'), ['*.txt', 'subfolder/*.txt'])
+        self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'target', '1.txt')))
+        self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'target', 'subfolder', '2.txt')))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils/test_io_utils.py

@@ -0,0 +1,42 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+from swift.utils import append_to_jsonl, get_logger, read_from_jsonl, write_to_jsonl
+
+logger = get_logger()
+
+
+class TestIOUtils(unittest.TestCase):
+
+    def setUp(self):
+        self._tmp_dir = tempfile.TemporaryDirectory()
+        self.tmp_dir = self._tmp_dir.name
+        # self.tmp_dir = 'test'
+        logger.info(f'self.tmp_dir: {self.tmp_dir}')
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def test_jsonl(self):
+        fpath = os.path.join(self.tmp_dir, '1.jsonl')
+        obj_list = [{'aaa': 'bbb'}, 111, [1.1]]
+        write_to_jsonl(fpath, obj_list)
+        new_obj = {'bbb': 'aaa'}
+        obj_list.append(new_obj)
+        append_to_jsonl(fpath, new_obj)
+        new_obj_list = read_from_jsonl(fpath)
+        self.assertTrue(new_obj_list == obj_list)
+
+    def test_jsonl2(self):
+        fpath = os.path.join(self.tmp_dir, '1.jsonl')
+        obj_list = [{'aaa': 'bbb'}, 111, [1.1]]
+        for obj in obj_list:
+            append_to_jsonl(fpath, obj)
+        new_obj_list = read_from_jsonl(fpath)
+        self.assertTrue(new_obj_list == obj_list)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils/test_multi_teacher.py

@@ -0,0 +1,405 @@
+"""Unit tests for multi-teacher model server support.
+
+Tests cover:
+- parse_teacher_model_server: single URL / JSON parsing + non-empty & non-overlapping tag validation
+- route_samples_to_teachers: single-teacher (all samples) and multi-teacher tag routing + fail-fast
+- fetch_teacher_parsed_by_routing: teacher-count-agnostic fetch + scatter back to sample order
+- expand_advantage_to_per_token: scalar teacher KL coefficient (all teachers share --teacher_kl_coef)
+- streaming multi-dataset load: auto-injected ``dataset`` column + routing parity with non-streaming
+"""
+import json
+import os
+import tempfile
+import torch
+import unittest
+from typing import List, Optional
+
+from swift.dataset import load_dataset
+from swift.rl_core.advantage import expand_advantage_to_per_token
+from swift.rl_core.data import OnPolicySample
+from swift.rlhf_trainers.gkd_helpers import (TeacherServerConfig, fetch_teacher_parsed_by_routing,
+                                             parse_teacher_model_server, route_samples_to_teachers)
+
+
+def _make_sample(tag: Optional[str] = None) -> OnPolicySample:
+    extra = {}
+    if tag is not None:
+        extra['dataset'] = tag
+    return OnPolicySample(messages=[], extra=extra)
+
+
+class TestParseTeacherModelServer(unittest.TestCase):
+
+    def test_parse_none(self):
+        self.assertIsNone(parse_teacher_model_server(None))
+
+    def test_parse_single_url(self):
+        result = parse_teacher_model_server('http://localhost:8000')
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].url, 'http://localhost:8000')
+        self.assertEqual(result[0].tags, [])
+
+    def test_parse_multi_json(self):
+        config = json.dumps([
+            {
+                'url': 'http://localhost:8000',
+                'tags': ['math']
+            },
+            {
+                'url': 'http://localhost:8001',
+                'tags': ['code']
+            },
+        ])
+        result = parse_teacher_model_server(config)
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0].tags, ['math'])
+        self.assertEqual(result[1].tags, ['code'])
+
+    def test_parse_empty_json_list(self):
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server('[]')
+
+    def test_parse_missing_url(self):
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server(json.dumps([{'tags': ['math']}]))
+
+    def test_parse_invalid_tags(self):
+        config = json.dumps([{'url': 'http://localhost:8000', 'tags': 'math'}])
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server(config)
+
+    def test_parse_multi_empty_tags_rejected(self):
+        """With multiple teachers, empty tags are rejected (each sample needs exactly one teacher)."""
+        config = json.dumps([
+            {
+                'url': 'http://localhost:8000',
+                'tags': ['math']
+            },
+            {
+                'url': 'http://localhost:8001',
+                'tags': []
+            },
+        ])
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server(config)
+
+    def test_parse_overlapping_tags_rejected(self):
+        """A tag may not appear in more than one teacher."""
+        config = json.dumps([
+            {
+                'url': 'http://localhost:8000',
+                'tags': ['math', 'shared']
+            },
+            {
+                'url': 'http://localhost:8001',
+                'tags': ['shared', 'code']
+            },
+        ])
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server(config)
+
+    def test_parse_non_dict_entry_rejected(self):
+        """A non-dict list entry is rejected instead of raising AttributeError."""
+        with self.assertRaises(ValueError):
+            parse_teacher_model_server(json.dumps(['http://localhost:8000']))
+
+    def test_parse_tags_coerced_to_str(self):
+        """Non-string tags are normalized to str so they match get_tag's str output."""
+        config = json.dumps([
+            {
+                'url': 'http://localhost:8000',
+                'tags': [1]
+            },
+            {
+                'url': 'http://localhost:8001',
+                'tags': [2]
+            },
+        ])
+        result = parse_teacher_model_server(config)
+        self.assertEqual(result[0].tags, ['1'])
+        self.assertEqual(result[1].tags, ['2'])
+
+
+class TestRouteSamplesToTeachers(unittest.TestCase):
+
+    def test_route_single_teacher_all_samples(self):
+        """Single teacher (empty tags) handles all samples, tags ignored."""
+        samples = [_make_sample('math'), _make_sample(), _make_sample('code')]
+        configs = [TeacherServerConfig(url='http://t0', tags=[])]
+        routing = route_samples_to_teachers(samples, configs)
+        self.assertEqual(routing[0], [0, 1, 2])
+
+    def test_route_one_to_one(self):
+        samples = [_make_sample('math'), _make_sample('code'), _make_sample('math')]
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        routing = route_samples_to_teachers(samples, configs)
+        self.assertEqual(routing[0], [0, 2])
+        self.assertEqual(routing[1], [1])
+
+    def test_route_unmatched_fails_fast(self):
+        samples = [_make_sample('unknown')]
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        with self.assertRaises(ValueError):
+            route_samples_to_teachers(samples, configs)
+
+    def test_get_tag_reads_tag_key(self):
+        # OnPolicySample.get_tag keys off exactly tag_key (default 'dataset'); no column fallback.
+        self.assertEqual(_make_sample('math').get_tag(), 'math')
+        sample = OnPolicySample(messages=[], extra={'domain': 'code'})
+        self.assertEqual(sample.get_tag(tag_key='domain'), 'code')
+        self.assertIsNone(sample.get_tag())  # default 'dataset' key absent -> None
+
+    def test_get_tag_no_tag_returns_none(self):
+        self.assertIsNone(_make_sample().get_tag())
+
+
+class TestFetchTeacherParsedByRouting(unittest.TestCase):
+
+    def test_multi_teacher_scatter_back(self):
+        samples = [_make_sample('math'), _make_sample('code'), _make_sample('math')]
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        requests = ['r0', 'r1', 'r2']
+        seen = {}
+
+        def fetch_fn(subset_reqs, client):
+            seen[client.url] = list(subset_reqs)
+            return [f'{client.url}:{r}' for r in subset_reqs]
+
+        parsed = fetch_teacher_parsed_by_routing(samples, requests, configs, configs, fetch_fn=fetch_fn)
+        self.assertEqual(parsed, ['http://t0:r0', 'http://t1:r1', 'http://t0:r2'])
+        self.assertEqual(seen['http://t0'], ['r0', 'r2'])
+        self.assertEqual(seen['http://t1'], ['r1'])
+
+    def test_single_teacher_one_fetch(self):
+        """Single teacher: one fetch over all requests in original order."""
+        samples = [_make_sample(), _make_sample(), _make_sample()]
+        configs = [TeacherServerConfig(url='http://t0', tags=[])]
+        requests = ['r0', 'r1', 'r2']
+        calls = []
+
+        def fetch_fn(subset_reqs, client):
+            calls.append(list(subset_reqs))
+            return [f'p:{r}' for r in subset_reqs]
+
+        parsed = fetch_teacher_parsed_by_routing(samples, requests, configs, configs, fetch_fn=fetch_fn)
+        self.assertEqual(parsed, ['p:r0', 'p:r1', 'p:r2'])
+        self.assertEqual(calls, [['r0', 'r1', 'r2']])  # exactly one fetch
+
+    def test_empty_subset_teacher_still_visited(self):
+        """Every teacher index is visited even with an empty subset (collective ordering / no
+        per-rank skip that would desync the DP gather/broadcast)."""
+        samples = [_make_sample('math'), _make_sample('math')]  # nothing routes to 'code'
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        requests = ['r0', 'r1']
+        visited = []
+
+        def fetch_fn(subset_reqs, client):
+            visited.append((client.url, list(subset_reqs)))
+            return [f'{client.url}:{r}' for r in subset_reqs]
+
+        parsed = fetch_teacher_parsed_by_routing(samples, requests, configs, configs, fetch_fn=fetch_fn)
+        self.assertEqual(parsed, ['http://t0:r0', 'http://t0:r1'])
+        # both teachers visited, in order, including the empty one
+        self.assertEqual(visited, [('http://t0', ['r0', 'r1']), ('http://t1', [])])
+
+    def test_phase_split_concurrent_scatter(self):
+        """3-phase (gather/infer/scatter) form scatters back in original sample order, and only
+        the main process runs infer (concurrently across teachers)."""
+        samples = [_make_sample('math'), _make_sample('code'), _make_sample('math')]
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        requests = ['r0', 'r1', 'r2']
+
+        def gather_fn(subset_reqs):
+            return list(subset_reqs)
+
+        def infer_fn(handle, client):
+            return [f'{client.url}:{r}' for r in handle]
+
+        def scatter_fn(handle, parsed_global):
+            return parsed_global
+
+        parsed = fetch_teacher_parsed_by_routing(
+            samples,
+            requests,
+            configs,
+            configs,
+            gather_fn=gather_fn,
+            infer_fn=infer_fn,
+            scatter_fn=scatter_fn,
+            is_main_process=True)
+        self.assertEqual(parsed, ['http://t0:r0', 'http://t1:r1', 'http://t0:r2'])
+
+    def test_phase_split_non_main_no_infer(self):
+        """On non-main ranks infer is skipped; scatter still runs for every teacher (receives the
+        broadcast in real distributed runs). Here scatter returns the per-teacher subset unchanged."""
+        samples = [_make_sample('math'), _make_sample('code')]
+        configs = [
+            TeacherServerConfig(url='http://t0', tags=['math']),
+            TeacherServerConfig(url='http://t1', tags=['code']),
+        ]
+        requests = ['r0', 'r1']
+        infer_called = []
+        scatter_calls = []
+
+        def gather_fn(subset_reqs):
+            return list(subset_reqs)
+
+        def infer_fn(handle, client):
+            infer_called.append(client)
+            return [f'x:{r}' for r in handle]
+
+        def scatter_fn(handle, parsed_global):
+            scatter_calls.append(parsed_global)
+            return [f's:{r}' for r in handle]  # mimic broadcast-then-slice back to local subset
+
+        parsed = fetch_teacher_parsed_by_routing(
+            samples,
+            requests,
+            configs,
+            configs,
+            gather_fn=gather_fn,
+            infer_fn=infer_fn,
+            scatter_fn=scatter_fn,
+            is_main_process=False)
+        self.assertEqual(infer_called, [])  # infer never runs off the main process
+        self.assertEqual(scatter_calls, [None, None])  # parsed_global is None on non-main
+        self.assertEqual(parsed, ['s:r0', 's:r1'])
+
+
+class TestExpandAdvantageScalarCoef(unittest.TestCase):
+
+    def test_scalar_coef(self):
+        B, T = 2, 4
+        result = expand_advantage_to_per_token(
+            torch.tensor([1.0, 1.0]),
+            torch.ones(B, T),
+            teacher_per_token_logps=torch.tensor([[2.0] * T, [3.0] * T]),
+            policy_per_token_logps=torch.tensor([[1.0] * T, [1.0] * T]),
+            teacher_kl_coef=0.5,
+        )
+        # base(1) + 0.5 * (2-1) = 1.5; base(1) + 0.5 * (3-1) = 2.0
+        torch.testing.assert_close(result[0], torch.ones(T) * 1.5)
+        torch.testing.assert_close(result[1], torch.ones(T) * 2.0)
+
+    def test_no_teacher(self):
+        B, T = 2, 4
+        result = expand_advantage_to_per_token(torch.tensor([1.0, 2.0]), torch.ones(B, T))
+        torch.testing.assert_close(result[0], torch.ones(T) * 1.0)
+        torch.testing.assert_close(result[1], torch.ones(T) * 2.0)
+
+    def test_zero_coef_no_injection(self):
+        B, T = 1, 4
+        result = expand_advantage_to_per_token(
+            torch.tensor([1.0]),
+            torch.ones(B, T),
+            teacher_per_token_logps=torch.tensor([[2.0] * T]),
+            policy_per_token_logps=torch.tensor([[1.0] * T]),
+            teacher_kl_coef=0.0,
+        )
+        torch.testing.assert_close(result[0], torch.ones(T) * 1.0)
+
+
+class TestStreamingDatasetMultiTeacherRouting(unittest.TestCase):
+    """End-to-end: streaming load_dataset injects ``dataset`` tags that multi-teacher routing consumes."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.tmpdir = tempfile.mkdtemp()
+        cls.paths: List[str] = []
+        for name, n_rows in (('math.jsonl', 3), ('code.jsonl', 2)):
+            path = os.path.join(cls.tmpdir, name)
+            with open(path, 'w', encoding='utf-8') as f:
+                for i in range(n_rows):
+                    row = {'messages': [{'role': 'user', 'content': f'{name}-{i}'}]}
+                    f.write(json.dumps(row, ensure_ascii=False) + '\n')
+            cls.paths.append(path)
+
+    def _teacher_configs(self) -> List[TeacherServerConfig]:
+        return [
+            TeacherServerConfig(url='http://t0', tags=[self.paths[0]]),
+            TeacherServerConfig(url='http://t1', tags=[self.paths[1]]),
+        ]
+
+    def _load_samples(self, *, streaming: bool, interleave_prob: Optional[List[float]] = None) -> List[OnPolicySample]:
+        kwargs = dict(datasets=self.paths, streaming=streaming, split_dataset_ratio=0.)
+        if interleave_prob is not None:
+            kwargs['interleave_prob'] = interleave_prob
+            kwargs['stopping_strategy'] = 'all_exhausted'
+        train, _ = load_dataset(**kwargs)
+        return [OnPolicySample.from_row(row) for row in train]
+
+    def test_streaming_injected_tags_route_correctly(self):
+        samples = self._load_samples(streaming=True)
+        self.assertEqual(len(samples), 5)
+        for sample in samples:
+            self.assertIn(sample.extra.get('dataset'), self.paths)
+
+        routing = route_samples_to_teachers(samples, self._teacher_configs())
+        math_indices = [i for i, s in enumerate(samples) if s.extra['dataset'] == self.paths[0]]
+        code_indices = [i for i, s in enumerate(samples) if s.extra['dataset'] == self.paths[1]]
+        self.assertEqual(routing[0], math_indices)
+        self.assertEqual(routing[1], code_indices)
+        self.assertEqual(len(math_indices), 3)
+        self.assertEqual(len(code_indices), 2)
+
+    def test_streaming_fetch_by_routing(self):
+        samples = self._load_samples(streaming=True)
+        configs = self._teacher_configs()
+        requests = [f'r{i}' for i in range(len(samples))]
+        seen = {}
+
+        def fetch_fn(subset_reqs, client):
+            seen[client.url] = list(subset_reqs)
+            return [f'{client.url}:{r}' for r in subset_reqs]
+
+        parsed = fetch_teacher_parsed_by_routing(samples, requests, configs, configs, fetch_fn=fetch_fn)
+        for i, sample in enumerate(samples):
+            tag = sample.extra['dataset']
+            expected_url = 'http://t0' if tag == self.paths[0] else 'http://t1'
+            self.assertEqual(parsed[i], f'{expected_url}:r{i}')
+        self.assertEqual(len(seen['http://t0']), 3)
+        self.assertEqual(len(seen['http://t1']), 2)
+
+    def test_streaming_routing_matches_non_streaming(self):
+        stream_samples = self._load_samples(streaming=True)
+        static_samples = self._load_samples(streaming=False)
+        configs = self._teacher_configs()
+        self.assertEqual(
+            route_samples_to_teachers(stream_samples, configs),
+            route_samples_to_teachers(static_samples, configs),
+        )
+        stream_tags = [s.extra['dataset'] for s in stream_samples]
+        static_tags = [s.extra['dataset'] for s in static_samples]
+        self.assertEqual(stream_tags, static_tags)
+
+    def test_streaming_interleave_preserves_per_dataset_tags(self):
+        samples = self._load_samples(streaming=True, interleave_prob=[0.5, 0.5])
+        self.assertGreater(len(samples), 0)
+        for sample in samples:
+            self.assertIn(sample.extra.get('dataset'), self.paths)
+
+        routing = route_samples_to_teachers(samples, self._teacher_configs())
+        self.assertEqual(sorted(routing[0] + routing[1]), list(range(len(samples))))
+        math_count = sum(1 for s in samples if s.extra['dataset'] == self.paths[0])
+        code_count = sum(1 for s in samples if s.extra['dataset'] == self.paths[1])
+        self.assertEqual(len(routing[0]), math_count)
+        self.assertEqual(len(routing[1]), code_count)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils/test_rewards.py

@@ -0,0 +1,260 @@
+import unittest
+
+
+class TestMathAccuracy(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        try:
+            from swift.rewards.orm import MathAccuracy
+            cls.math_accuracy = MathAccuracy()
+            cls.available = True
+        except (ImportError, AssertionError) as e:
+            print(f'Warning: MathAccuracy not available: {e}')
+            cls.available = False
+
+    def setUp(self):
+        if not self.available:
+            self.skipTest('MathAccuracy not available (math_verify not installed)')
+
+    def test_pure_latex_format(self):
+        completions = ['The answer is \\boxed{42}']
+        solutions = ['\\boxed{42}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_latex_in_long_text(self):
+        completions = ['After careful calculation, the final answer is \\boxed{100}']
+        solutions = ['\\boxed{100}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_multiple_steps_with_boxed(self):
+        completions = [
+            'Let me solve step by step:\n'
+            '1. First we have x = 2\n'
+            '2. Then y = 3x = 6\n'
+            '3. Finally z = x + y = 8\n'
+            '\nFinal answer: \\boxed{8}'
+        ]
+        solutions = ['\\boxed{8}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_wrong_answer_no_tag(self):
+        completions = ['The answer is \\boxed{42}']
+        solutions = ['\\boxed{100}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 0.0)
+
+    def test_batch_processing_no_tag(self):
+        completions = ['\\boxed{42}', '\\boxed{100}', '\\boxed{8}']
+        solutions = ['\\boxed{42}', '\\boxed{100}', '\\boxed{8}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 3)
+        self.assertEqual(rewards[0], 1.0)
+        self.assertEqual(rewards[1], 1.0)
+        self.assertEqual(rewards[2], 1.0)
+
+    def test_answer_tag_with_plain_number(self):
+        completions = ['<answer>84</answer>']
+        solutions = ['\\boxed{84}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_answer_tag_with_latex(self):
+        completions = ['<answer>\\boxed{100}</answer>']
+        solutions = ['\\boxed{100}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_long_text_with_answer_tag(self):
+        completions = [
+            'Let me solve:\n'
+            'Step 1: Calculate x = 10\n'
+            'Step 2: Calculate y = 20\n'
+            'Step 3: Sum = 30\n'
+            '\n<answer>54</answer>'
+        ]
+        solutions = ['\\boxed{54}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_answer_tag_with_complex_expression(self):
+        completions = ['<answer>\\frac{1}{2}</answer>']
+        solutions = ['\\boxed{\\frac{1}{2}}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_solution_with_answer_tag(self):
+        completions = ['<answer>84</answer>']
+        solutions = ['<answer>\\boxed{84}</answer>']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_answer_tag_wrong_answer(self):
+        completions = ['<answer>42</answer>']
+        solutions = ['\\boxed{100}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 0.0)
+
+    def test_mixed_batch_with_and_without_tags(self):
+        completions = [
+            '\\boxed{42}',
+            '<answer>100</answer>',
+            'The answer is \\boxed{8}',
+        ]
+        solutions = [
+            '\\boxed{42}',
+            '\\boxed{100}',
+            '\\boxed{8}',
+        ]
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 3)
+        self.assertEqual(rewards[0], 1.0)
+        self.assertEqual(rewards[1], 1.0)
+        self.assertEqual(rewards[2], 1.0)
+
+    def test_empty_solution(self):
+        completions = ['<answer>42</answer>']
+        solutions = ['']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 0.0)
+
+    def test_malformed_latex(self):
+        completions = ['\\boxed{42']
+        solutions = ['\\boxed{42}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 0.0)
+
+    def test_answer_tag_with_extra_whitespace(self):
+        completions = ['<answer>  84  </answer>']
+        solutions = ['\\boxed{84}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_multiple_answer_tags(self):
+        completions = ['<answer>42</answer> Some text <answer>100</answer>']
+        solutions = ['\\boxed{42}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_real_world_example_from_user(self):
+        completions = [
+            'We are given a geometric sequence $\\{a_n\\}$ with:\n\n'
+            '- $a_3 = 2$\n- $a_5 = 6$\n\n'
+            'We are to find $a_9$.\n\n---\n\n'
+            '### Step 1: Recall the formula\n\n'
+            '$$a_n = a_1 \\cdot r^{n-1}$$\n\n---\n\n'
+            '### Step 2: Use the given terms\n\n'
+            '$$a_3 = a_1 \\cdot r^2 = 2$$\n'
+            '$$a_5 = a_1 \\cdot r^4 = 6$$\n\n'
+            'Divide equation (2) by equation (1):\n'
+            '$$r^2 = 3$$\n\n---\n\n'
+            '### Step 3: Find $a_9$\n\n'
+            '$$a_9 = a_1 \\cdot r^8 = \\frac{2}{3} \\cdot 81 = 54$$\n\n'
+            '### ✅ Final Answer:\n\n'
+            '<answer>54</answer>'
+        ]
+        solutions = ['\\boxed{54}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_equivalent_fractions(self):
+        completions = ['<answer>0.5</answer>']
+        solutions = ['\\boxed{\\frac{1}{2}}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_different_forms_same_answer(self):
+        completions = ['<answer>2</answer>']
+        solutions = ['\\boxed{\\sqrt{4}}']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_latex_inline_math_delimiters(self):
+        completions = ['<answer>84</answer>', '<answer>3</answer>']
+        solutions = ['\n\n\\[\n\\boxed{84}\n\\]', 'Therefore, the value of \\(a^2 - a + 2\\) is \\(\\boxed{3}\\).']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 2)
+        self.assertEqual(rewards[0], 1.0)
+        self.assertEqual(rewards[1], 1.0)
+
+    def test_latex_display_math_delimiters(self):
+        completions = ['<answer>100</answer>']
+        solutions = ['\\[\\boxed{100}\\]']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+    def test_mixed_latex_delimiters(self):
+        completions = ['<answer>\\(x = 42\\)</answer>']
+        solutions = ['\\[\\boxed{x = 42}\\]']
+
+        rewards = self.math_accuracy(completions, solutions)
+
+        self.assertEqual(len(rewards), 1)
+        self.assertEqual(rewards[0], 1.0)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils/test_split_str_parts_by.py

@@ -0,0 +1,13 @@
+from swift.template import split_str_parts_by
+
+
+def test_split_str_parts_by():
+    print(split_str_parts_by('aaaAction:bb\nbAction Inputs:\nabbb', ['Action:', 'Action Inputs:'], regex_mode=False))
+    print(split_str_parts_by('aaaAction:bb\nbAction Inputs:\nabbb', ['Action:', 'Action Inputs:'], regex_mode=True))
+    print(split_str_parts_by('aaa<tool_call>bbb</tool_call>ccc', ['<tool_call>.+?</tool_call>'], regex_mode=True))
+    print(split_str_parts_by('aaa<image>\nbb\nb<audio>\nabbb', ['<image>', '<audio>', '<video>'], regex_mode=False))
+    print(split_str_parts_by('aaa<image>\nbb\nb<audio>\nabbb', ['<image>', '<audio>', '<video>'], regex_mode=True))
+
+
+if __name__ == '__main__':
+    test_split_str_parts_by()

tests/utils/test_torch_utils.py

@@ -0,0 +1,15 @@
+import unittest
+from modelscope import Model
+
+from swift.utils import find_sub_module
+
+
+class TestTorchUtils(unittest.TestCase):
+
+    def test_find_sub_module(self):
+        model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
+        self.assertTrue(find_sub_module(model, 'query') is not None)
+
+
+if __name__ == '__main__':
+    unittest.main()