From 1f8f65fe89e37a119f6ea2fd5e3dbec9ba6aca1d Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 12 Oct 2026 20:46:53 +0800 Subject: [PATCH 1/5] fix: keep worker-owned processors local in safe mode --- src/twinkle/infra/__init__.py | 21 ++- .../model/transformers/transformers.py | 4 +- src/twinkle/processor/base.py | 6 + tests/model/test_safe_mode.py | 159 ++++++++++++++++++ 4 files changed, 187 insertions(+), 3 deletions(-) create mode 100644 tests/model/test_safe_mode.py diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index 6d6ae8d07..eaca7d8e0 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -6,6 +6,7 @@ import os import sys import threading +from contextvars import ContextVar from typing import Any, Callable, List, Literal, Optional, TypeVar, Union from twinkle.notifier import Notifier, notify_exception @@ -18,6 +19,24 @@ _mode: Optional[Literal['local', 'ray']] = 'local' +_local_call_targets = ContextVar('twinkle_local_call_targets', default=()) + + +def _is_local_call(instance): + return any(target is instance for target in _local_call_targets.get()) + + +def _call_local(instance, method_name, *args, **kwargs): + """Call a worker-owned object's normal entry point without redispatching it.""" + if hasattr(instance, '_actors'): + raise ValueError('Local calls require a worker-owned object, not a remote handle.') + token = _local_call_targets.set(_local_call_targets.get() + (instance, )) + try: + return getattr(instance, method_name)(*args, **kwargs) + finally: + _local_call_targets.reset(token) + + if os.environ.get('TWINKLE_MODE', 'local') == 'ray': _mode = 'ray' @@ -1025,7 +1044,7 @@ def wrapper(self, *args, **kwargs) -> T1: _ctx = f'{_ctx} <- {_caller}' try: device_mesh = getattr(self, 'device_mesh', None) - if _mode == 'local': + if _mode == 'local' or _is_local_call(self): return func(self, *args, **kwargs) elif _mode == 'ray': check_unsafe(*args, **kwargs) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5540f058f..673991993 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -562,7 +562,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor._process( inputs, sp_strategy=self.sp_strategy, model=self.model, @@ -687,7 +687,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor._process( inputs, sp_strategy=self.sp_strategy, model=self.model, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 131ae81cc..e679aa931 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -90,6 +90,12 @@ def __call__(self, inputs: Union[InputFeature, List[InputFeature]], inputs = pipe(inputs, **kwargs) return inputs + def _process(self, inputs: Union[InputFeature, List[InputFeature]], + **kwargs) -> Union[InputFeature, List[InputFeature]]: + """Preserve custom entry points while processing worker-owned objects locally.""" + from twinkle.infra import _call_local + return _call_local(self, '__call__', inputs, **kwargs) + def prepare_outputs(self, inputs: List[InputFeature], **kwargs) -> Union[List[InputFeature], InputFeature]: if self.framework == 'transformers': return inputs[0] diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py new file mode 100644 index 000000000..1e1dac6a4 --- /dev/null +++ b/tests/model/test_safe_mode.py @@ -0,0 +1,159 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CPU training through production methods, processors and Ray-mode decorators.""" +from types import SimpleNamespace + +import pytest +import torch + +import twinkle.infra as infra +from twinkle import remote_function +from twinkle.loss import CrossEntropyLoss +from twinkle.model.transformers.transformers import OptimizerGroup, TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Template +from twinkle.utils import Platform + + +class _TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embedding = torch.nn.Embedding(8, 8) + + def forward(self, input_ids, **kwargs): + return {'logits': self.embedding(input_ids)} + + +@pytest.fixture +def cpu_model(monkeypatch): + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: torch.device('cpu')) + + def build(processor_cls=InputProcessor, model_cls=TransformersModel): + # Construct worker-owned objects without Ray driver handles or tokenizer + # downloads. Keep the full processor pipeline, loss and backward methods. + with monkeypatch.context() as local: + local.setattr(infra, '_mode', 'local') + processor = processor_cls(padding_free=False) + group = OptimizerGroup( + processor=processor, + template=object.__new__(Template), + loss_instance=CrossEntropyLoss(), + ) + model = object.__new__(model_cls) + torch.nn.Module.__init__(model) + model.model = _TinyModel() + model.optimizer_group = {'': group} + model.device_mesh = None + model.mixed_precision = 'no' + model.sp_strategy = None + model.hf_config = SimpleNamespace() + model._model_wrapped = True + model._router_replay_setup = lambda *args: lambda: None + model.strategy = SimpleNamespace(unwrap_model=lambda obj: obj) + return model, group + + return build + + +def _features(): + return [{'input_ids': [1, 2, 3, 4], 'labels': [2, 3, 4, 5]}] + + +def _assert_gradients(model): + grad = model.model.embedding.weight.grad + assert grad is not None and torch.isfinite(grad).all() + assert grad.norm().item() > 0 + + +@pytest.mark.parametrize('mode', ['local', 'ray']) +@pytest.mark.parametrize('trust', ['0', '1']) +@pytest.mark.parametrize('entry', ['forward_backward', 'micro_batch', 'forward_only']) +def test_training_entries(cpu_model, monkeypatch, mode, trust, entry): + monkeypatch.setattr(infra, '_mode', mode) + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) + model, _ = cpu_model() + if entry == 'forward_only': + model.forward_only(inputs=_features()) + loss = model.calculate_loss() + assert model.model.embedding.weight.grad is None + else: + kwargs = {'micro_batch_size': 1} if entry == 'micro_batch' else {} + inputs = _features() * 2 if kwargs else _features() + loss = model.forward_backward(inputs=inputs, **kwargs)['loss'] + _assert_gradients(model) + assert torch.isfinite(torch.tensor(loss)) + + +@pytest.mark.parametrize('trust', ['0', '1']) +def test_tinker_forward_backward(cpu_model, monkeypatch, trust): + import numpy as np + types = pytest.importorskip('tinker').types + from twinkle.server.model.backends.transformers_model import TwinkleCompatFullTransformersModel + + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) + model, _ = cpu_model(model_cls=TwinkleCompatFullTransformersModel) + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3, 4]), + loss_fn_inputs={ + 'target_tokens': types.TensorData.from_numpy(np.array([2, 3, 4, 5], dtype=np.int64)), + 'weights': types.TensorData.from_numpy(np.ones(4, dtype=np.float32)), + }, + ) + results, loss = model.tinker_forward_backward(inputs=[datum], adapter_name='', loss_fn='cross_entropy') + assert results and torch.isfinite(torch.tensor(loss)) + _assert_gradients(model) + + +@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) +@pytest.mark.parametrize('remote_handle', [False, True]) +def test_public_processor_rejects_callable(cpu_model, value, remote_handle): + _, group = cpu_model() + if remote_handle: + group.processor._actors = [] + with pytest.raises(ValueError, match='Callable or Type'): + group.processor(_features(), external=value) + + +@pytest.mark.parametrize('mode', ['local', 'ray']) +@pytest.mark.parametrize('decorated', [False, True]) +def test_custom_processor_entry_point_is_preserved(cpu_model, monkeypatch, mode, decorated): + class CustomProcessor(InputProcessor): + def __call__(self, inputs, **kwargs): + self.seen_model = kwargs['model'] + return super().__call__(inputs, **kwargs) + + if decorated: + CustomProcessor.__call__ = remote_function()(CustomProcessor.__call__) + monkeypatch.setattr(infra, '_mode', mode) + model, group = cpu_model(processor_cls=CustomProcessor) + model.forward(inputs=_features()) + assert group.processor.seen_model is model.model + if mode == 'ray': + with pytest.raises(ValueError, match='Callable or Type'): + InputProcessor.__call__(group.processor, _features(), external=lambda: None) + + +def test_local_scope_keeps_other_objects_checked_and_resets_on_error(cpu_model): + _, group = cpu_model() + _, other = cpu_model() + + def pipeline(inputs, **kwargs): + other.processor(inputs, external=lambda: None) + + group.processor.process_pipeline = [pipeline] + with pytest.raises(ValueError, match='Callable or Type'): + group.processor._process(_features(), model=lambda: None) + with pytest.raises(ValueError, match='Callable or Type'): + group.processor(_features(), external=lambda: None) + + +def test_local_call_rejects_remote_handle(cpu_model): + _, group = cpu_model() + group.processor._actors = [] + with pytest.raises(ValueError, match='remote handle'): + group.processor._process(_features()) + + +def test_internal_processor_entry_is_not_exported(): + assert not hasattr(InputProcessor._process, '_execute') From 579c0abdb8b706db2473a317394658a31d45ba9d Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 12 Oct 2026 21:27:25 +0800 Subject: [PATCH 2/5] refactor: simplify worker-local input processing --- src/twinkle/infra/__init__.py | 21 +-------------------- src/twinkle/processor/base.py | 13 +++++++------ tests/model/test_safe_mode.py | 28 +++------------------------- 3 files changed, 11 insertions(+), 51 deletions(-) diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index eaca7d8e0..6d6ae8d07 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -6,7 +6,6 @@ import os import sys import threading -from contextvars import ContextVar from typing import Any, Callable, List, Literal, Optional, TypeVar, Union from twinkle.notifier import Notifier, notify_exception @@ -19,24 +18,6 @@ _mode: Optional[Literal['local', 'ray']] = 'local' -_local_call_targets = ContextVar('twinkle_local_call_targets', default=()) - - -def _is_local_call(instance): - return any(target is instance for target in _local_call_targets.get()) - - -def _call_local(instance, method_name, *args, **kwargs): - """Call a worker-owned object's normal entry point without redispatching it.""" - if hasattr(instance, '_actors'): - raise ValueError('Local calls require a worker-owned object, not a remote handle.') - token = _local_call_targets.set(_local_call_targets.get() + (instance, )) - try: - return getattr(instance, method_name)(*args, **kwargs) - finally: - _local_call_targets.reset(token) - - if os.environ.get('TWINKLE_MODE', 'local') == 'ray': _mode = 'ray' @@ -1044,7 +1025,7 @@ def wrapper(self, *args, **kwargs) -> T1: _ctx = f'{_ctx} <- {_caller}' try: device_mesh = getattr(self, 'device_mesh', None) - if _mode == 'local' or _is_local_call(self): + if _mode == 'local': return func(self, *args, **kwargs) elif _mode == 'ray': check_unsafe(*args, **kwargs) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index e679aa931..8bc3b8d99 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -86,15 +86,16 @@ def __init__(self, @remote_function() def __call__(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: - for pipe in self.process_pipeline: - inputs = pipe(inputs, **kwargs) - return inputs + return self._process(inputs, **kwargs) def _process(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: - """Preserve custom entry points while processing worker-owned objects locally.""" - from twinkle.infra import _call_local - return _call_local(self, '__call__', inputs, **kwargs) + """Process worker-owned inputs without remote dispatch.""" + if hasattr(self, '_actors'): + raise ValueError('Local processing requires a worker-owned object, not a remote handle.') + for pipe in self.process_pipeline: + inputs = pipe(inputs, **kwargs) + return inputs def prepare_outputs(self, inputs: List[InputFeature], **kwargs) -> Union[List[InputFeature], InputFeature]: if self.framework == 'transformers': diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py index 1e1dac6a4..5ece362da 100644 --- a/tests/model/test_safe_mode.py +++ b/tests/model/test_safe_mode.py @@ -6,7 +6,6 @@ import torch import twinkle.infra as infra -from twinkle import remote_function from twinkle.loss import CrossEntropyLoss from twinkle.model.transformers.transformers import OptimizerGroup, TransformersModel from twinkle.processor import InputProcessor @@ -29,12 +28,12 @@ def cpu_model(monkeypatch): monkeypatch.setattr(infra, '_mode', 'ray') monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: torch.device('cpu')) - def build(processor_cls=InputProcessor, model_cls=TransformersModel): + def build(model_cls=TransformersModel): # Construct worker-owned objects without Ray driver handles or tokenizer # downloads. Keep the full processor pipeline, loss and backward methods. with monkeypatch.context() as local: local.setattr(infra, '_mode', 'local') - processor = processor_cls(padding_free=False) + processor = InputProcessor(padding_free=False) group = OptimizerGroup( processor=processor, template=object.__new__(Template), @@ -115,26 +114,7 @@ def test_public_processor_rejects_callable(cpu_model, value, remote_handle): group.processor(_features(), external=value) -@pytest.mark.parametrize('mode', ['local', 'ray']) -@pytest.mark.parametrize('decorated', [False, True]) -def test_custom_processor_entry_point_is_preserved(cpu_model, monkeypatch, mode, decorated): - class CustomProcessor(InputProcessor): - def __call__(self, inputs, **kwargs): - self.seen_model = kwargs['model'] - return super().__call__(inputs, **kwargs) - - if decorated: - CustomProcessor.__call__ = remote_function()(CustomProcessor.__call__) - monkeypatch.setattr(infra, '_mode', mode) - model, group = cpu_model(processor_cls=CustomProcessor) - model.forward(inputs=_features()) - assert group.processor.seen_model is model.model - if mode == 'ray': - with pytest.raises(ValueError, match='Callable or Type'): - InputProcessor.__call__(group.processor, _features(), external=lambda: None) - - -def test_local_scope_keeps_other_objects_checked_and_resets_on_error(cpu_model): +def test_internal_processing_keeps_public_safety_checks(cpu_model): _, group = cpu_model() _, other = cpu_model() @@ -144,8 +124,6 @@ def pipeline(inputs, **kwargs): group.processor.process_pipeline = [pipeline] with pytest.raises(ValueError, match='Callable or Type'): group.processor._process(_features(), model=lambda: None) - with pytest.raises(ValueError, match='Callable or Type'): - group.processor(_features(), external=lambda: None) def test_local_call_rejects_remote_handle(cpu_model): From f6edcc8014bbe9c95fd16ea6026251c99995bcf6 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 12 Oct 2026 21:35:02 +0800 Subject: [PATCH 3/5] refactor: make local processor entry explicit --- src/twinkle/model/transformers/transformers.py | 4 ++-- src/twinkle/processor/base.py | 12 ++++++++---- tests/model/test_safe_mode.py | 6 +++--- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 673991993..340b353e6 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -562,7 +562,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor._process( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, @@ -687,7 +687,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) - inputs: Dict[str, Any] = processor._process( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 8bc3b8d99..816c47d5e 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -86,11 +86,15 @@ def __init__(self, @remote_function() def __call__(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: - return self._process(inputs, **kwargs) + return self.process_local(inputs, **kwargs) - def _process(self, inputs: Union[InputFeature, List[InputFeature]], - **kwargs) -> Union[InputFeature, List[InputFeature]]: - """Process worker-owned inputs without remote dispatch.""" + def process_local(self, inputs: Union[InputFeature, List[InputFeature]], + **kwargs) -> Union[InputFeature, List[InputFeature]]: + """Process inputs on this worker without remote dispatch. + + Use this for worker-owned inputs such as the model passed by Transformers. + This method is not remotely exported; use ``__call__`` for remote calls. + """ if hasattr(self, '_actors'): raise ValueError('Local processing requires a worker-owned object, not a remote handle.') for pipe in self.process_pipeline: diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py index 5ece362da..7fcb3ea08 100644 --- a/tests/model/test_safe_mode.py +++ b/tests/model/test_safe_mode.py @@ -123,15 +123,15 @@ def pipeline(inputs, **kwargs): group.processor.process_pipeline = [pipeline] with pytest.raises(ValueError, match='Callable or Type'): - group.processor._process(_features(), model=lambda: None) + group.processor.process_local(_features(), model=lambda: None) def test_local_call_rejects_remote_handle(cpu_model): _, group = cpu_model() group.processor._actors = [] with pytest.raises(ValueError, match='remote handle'): - group.processor._process(_features()) + group.processor.process_local(_features()) def test_internal_processor_entry_is_not_exported(): - assert not hasattr(InputProcessor._process, '_execute') + assert not hasattr(InputProcessor.process_local, '_execute') From 358f768dd8e01f8909fdbac5a0e5f7d88d32bebf Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 13 Oct 2026 14:24:19 +0800 Subject: [PATCH 4/5] fix: check unsafe inputs at the remote dispatch boundary --- src/twinkle/infra/__init__.py | 3 +- .../model/transformers/transformers.py | 4 +- src/twinkle/processor/base.py | 11 ---- tests/infra/test_safe_mode.py | 51 +++++++++++++++++++ tests/model/test_safe_mode.py | 29 ++++++----- 5 files changed, 70 insertions(+), 28 deletions(-) create mode 100644 tests/infra/test_safe_mode.py diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index 6d6ae8d07..3d85d9b5e 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -1028,7 +1028,6 @@ def wrapper(self, *args, **kwargs) -> T1: if _mode == 'local': return func(self, *args, **kwargs) elif _mode == 'ray': - check_unsafe(*args, **kwargs) if not hasattr(self, '_actors'): # This is the worker from ._ray import RayHelper @@ -1044,6 +1043,8 @@ def wrapper(self, *args, **kwargs) -> T1: _, args, kwargs = _workers_and_args[rank] return func(self, *args, **kwargs) else: + # Check inputs before remote dispatch, not worker-local calls. + check_unsafe(*args, **kwargs) # This is the driver from ._ray import RayHelper diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 340b353e6..5540f058f 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -562,7 +562,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor.process_local( + inputs: Dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, model=self.model, @@ -687,7 +687,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) - inputs: Dict[str, Any] = processor.process_local( + inputs: Dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, model=self.model, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 816c47d5e..131ae81cc 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -86,17 +86,6 @@ def __init__(self, @remote_function() def __call__(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: - return self.process_local(inputs, **kwargs) - - def process_local(self, inputs: Union[InputFeature, List[InputFeature]], - **kwargs) -> Union[InputFeature, List[InputFeature]]: - """Process inputs on this worker without remote dispatch. - - Use this for worker-owned inputs such as the model passed by Transformers. - This method is not remotely exported; use ``__call__`` for remote calls. - """ - if hasattr(self, '_actors'): - raise ValueError('Local processing requires a worker-owned object, not a remote handle.') for pipe in self.process_pipeline: inputs = pipe(inputs, **kwargs) return inputs diff --git a/tests/infra/test_safe_mode.py b/tests/infra/test_safe_mode.py new file mode 100644 index 000000000..7d8a7cf37 --- /dev/null +++ b/tests/infra/test_safe_mode.py @@ -0,0 +1,51 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Safe-mode checks at the remote dispatch boundary, using real Ray actors.""" +import pytest +import ray +import os + +import twinkle.infra as infra +from twinkle import remote_function + + +class _Component: + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def run(self, value): + assert infra._mode == 'ray' + assert os.environ['TWINKLE_TRUST_REMOTE_CODE'] == '0' + return self.apply(value, transform=lambda x: x + 1) + + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def apply(self, value, transform): + return transform(value) + + +@pytest.fixture(scope='module') +def actor(): + ray.init(num_cpus=1, include_dashboard=False, logging_level='ERROR') + try: + yield ray.remote(_Component).options( + runtime_env={'env_vars': {'TWINKLE_MODE': 'ray', 'TWINKLE_TRUST_REMOTE_CODE': '0'}}).remote() + finally: + ray.shutdown() + + +@pytest.fixture +def handle(actor, monkeypatch): + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') + instance = _Component() + instance._actors = [actor] + return instance + + +@pytest.mark.parametrize('trust', ['0', '1']) +def test_worker_can_call_another_method_with_local_callable(handle, monkeypatch, trust): + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) + assert handle.run(3) == 4 + + +@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) +def test_remote_handle_rejects_callable_inputs(handle, value): + with pytest.raises(ValueError, match='Callable or Type'): + handle.apply(3, transform=value) diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py index 7fcb3ea08..07eb73a55 100644 --- a/tests/model/test_safe_mode.py +++ b/tests/model/test_safe_mode.py @@ -106,32 +106,33 @@ def test_tinker_forward_backward(cpu_model, monkeypatch, trust): @pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) @pytest.mark.parametrize('remote_handle', [False, True]) -def test_public_processor_rejects_callable(cpu_model, value, remote_handle): +def test_processor_checks_only_remote_dispatch(cpu_model, value, remote_handle): _, group = cpu_model() if remote_handle: group.processor._actors = [] - with pytest.raises(ValueError, match='Callable or Type'): - group.processor(_features(), external=value) + with pytest.raises(ValueError, match='Callable or Type'): + group.processor(_features(), external=value) + else: + result = group.processor(_features(), external=value) + assert torch.equal(result['input_ids'], torch.tensor([[1, 2, 3, 4]])) -def test_internal_processing_keeps_public_safety_checks(cpu_model): +def test_worker_call_to_remote_handle_keeps_safety_checks(cpu_model): _, group = cpu_model() _, other = cpu_model() + other.processor._actors = [] def pipeline(inputs, **kwargs): other.processor(inputs, external=lambda: None) group.processor.process_pipeline = [pipeline] with pytest.raises(ValueError, match='Callable or Type'): - group.processor.process_local(_features(), model=lambda: None) - - -def test_local_call_rejects_remote_handle(cpu_model): - _, group = cpu_model() - group.processor._actors = [] - with pytest.raises(ValueError, match='remote handle'): - group.processor.process_local(_features()) + group.processor(_features(), model=lambda: None) -def test_internal_processor_entry_is_not_exported(): - assert not hasattr(InputProcessor.process_local, '_execute') +@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) +def test_processor_constructor_keeps_safety_checks(monkeypatch, value): + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') + with pytest.raises(ValueError, match='Callable or Type'): + InputProcessor(external=value) From 373bda80a7f447899863283f101fbfb093be0318 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 13 Oct 2026 15:25:48 +0800 Subject: [PATCH 5/5] fix: restore worker-local processor path and receiver safety checks --- src/twinkle/infra/__init__.py | 3 +- .../model/transformers/transformers.py | 4 +- src/twinkle/processor/base.py | 11 ++++ tests/infra/test_safe_mode.py | 51 ------------------- tests/model/test_safe_mode.py | 29 +++++------ 5 files changed, 28 insertions(+), 70 deletions(-) delete mode 100644 tests/infra/test_safe_mode.py diff --git a/src/twinkle/infra/__init__.py b/src/twinkle/infra/__init__.py index 3d85d9b5e..6d6ae8d07 100644 --- a/src/twinkle/infra/__init__.py +++ b/src/twinkle/infra/__init__.py @@ -1028,6 +1028,7 @@ def wrapper(self, *args, **kwargs) -> T1: if _mode == 'local': return func(self, *args, **kwargs) elif _mode == 'ray': + check_unsafe(*args, **kwargs) if not hasattr(self, '_actors'): # This is the worker from ._ray import RayHelper @@ -1043,8 +1044,6 @@ def wrapper(self, *args, **kwargs) -> T1: _, args, kwargs = _workers_and_args[rank] return func(self, *args, **kwargs) else: - # Check inputs before remote dispatch, not worker-local calls. - check_unsafe(*args, **kwargs) # This is the driver from ._ray import RayHelper diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5540f058f..340b353e6 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -562,7 +562,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, @@ -687,7 +687,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 131ae81cc..816c47d5e 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -86,6 +86,17 @@ def __init__(self, @remote_function() def __call__(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: + return self.process_local(inputs, **kwargs) + + def process_local(self, inputs: Union[InputFeature, List[InputFeature]], + **kwargs) -> Union[InputFeature, List[InputFeature]]: + """Process inputs on this worker without remote dispatch. + + Use this for worker-owned inputs such as the model passed by Transformers. + This method is not remotely exported; use ``__call__`` for remote calls. + """ + if hasattr(self, '_actors'): + raise ValueError('Local processing requires a worker-owned object, not a remote handle.') for pipe in self.process_pipeline: inputs = pipe(inputs, **kwargs) return inputs diff --git a/tests/infra/test_safe_mode.py b/tests/infra/test_safe_mode.py deleted file mode 100644 index 7d8a7cf37..000000000 --- a/tests/infra/test_safe_mode.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Safe-mode checks at the remote dispatch boundary, using real Ray actors.""" -import pytest -import ray -import os - -import twinkle.infra as infra -from twinkle import remote_function - - -class _Component: - @remote_function(dispatch='all', collect='first', lazy_collect=False) - def run(self, value): - assert infra._mode == 'ray' - assert os.environ['TWINKLE_TRUST_REMOTE_CODE'] == '0' - return self.apply(value, transform=lambda x: x + 1) - - @remote_function(dispatch='all', collect='first', lazy_collect=False) - def apply(self, value, transform): - return transform(value) - - -@pytest.fixture(scope='module') -def actor(): - ray.init(num_cpus=1, include_dashboard=False, logging_level='ERROR') - try: - yield ray.remote(_Component).options( - runtime_env={'env_vars': {'TWINKLE_MODE': 'ray', 'TWINKLE_TRUST_REMOTE_CODE': '0'}}).remote() - finally: - ray.shutdown() - - -@pytest.fixture -def handle(actor, monkeypatch): - monkeypatch.setattr(infra, '_mode', 'ray') - monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') - instance = _Component() - instance._actors = [actor] - return instance - - -@pytest.mark.parametrize('trust', ['0', '1']) -def test_worker_can_call_another_method_with_local_callable(handle, monkeypatch, trust): - monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) - assert handle.run(3) == 4 - - -@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) -def test_remote_handle_rejects_callable_inputs(handle, value): - with pytest.raises(ValueError, match='Callable or Type'): - handle.apply(3, transform=value) diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py index 07eb73a55..7fcb3ea08 100644 --- a/tests/model/test_safe_mode.py +++ b/tests/model/test_safe_mode.py @@ -106,33 +106,32 @@ def test_tinker_forward_backward(cpu_model, monkeypatch, trust): @pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) @pytest.mark.parametrize('remote_handle', [False, True]) -def test_processor_checks_only_remote_dispatch(cpu_model, value, remote_handle): +def test_public_processor_rejects_callable(cpu_model, value, remote_handle): _, group = cpu_model() if remote_handle: group.processor._actors = [] - with pytest.raises(ValueError, match='Callable or Type'): - group.processor(_features(), external=value) - else: - result = group.processor(_features(), external=value) - assert torch.equal(result['input_ids'], torch.tensor([[1, 2, 3, 4]])) + with pytest.raises(ValueError, match='Callable or Type'): + group.processor(_features(), external=value) -def test_worker_call_to_remote_handle_keeps_safety_checks(cpu_model): +def test_internal_processing_keeps_public_safety_checks(cpu_model): _, group = cpu_model() _, other = cpu_model() - other.processor._actors = [] def pipeline(inputs, **kwargs): other.processor(inputs, external=lambda: None) group.processor.process_pipeline = [pipeline] with pytest.raises(ValueError, match='Callable or Type'): - group.processor(_features(), model=lambda: None) + group.processor.process_local(_features(), model=lambda: None) -@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) -def test_processor_constructor_keeps_safety_checks(monkeypatch, value): - monkeypatch.setattr(infra, '_mode', 'ray') - monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') - with pytest.raises(ValueError, match='Callable or Type'): - InputProcessor(external=value) +def test_local_call_rejects_remote_handle(cpu_model): + _, group = cpu_model() + group.processor._actors = [] + with pytest.raises(ValueError, match='remote handle'): + group.processor.process_local(_features()) + + +def test_internal_processor_entry_is_not_exported(): + assert not hasattr(InputProcessor.process_local, '_execute')