Nahs / tests /single_controller /base /test_decorator.py

Add files using upload-large-folder tool

7b9e83e verified about 1 month ago

2.39 kB

# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest

import verl.single_controller.base.decorator as decorator_module
from verl.single_controller.base.decorator import (
    DISPATCH_MODE_FN_REGISTRY,
    Dispatch,
    _check_dispatch_mode,
    get_predefined_dispatch_fn,
    register_dispatch_mode,
    update_dispatch_mode,
)


@pytest.fixture
def reset_dispatch_registry():
    # Store original state
    original_registry = DISPATCH_MODE_FN_REGISTRY.copy()
    yield
    # Reset registry after test
    decorator_module.DISPATCH_MODE_FN_REGISTRY.clear()
    decorator_module.DISPATCH_MODE_FN_REGISTRY.update(original_registry)


def test_register_new_dispatch_mode(reset_dispatch_registry):
    # Test registration
    def dummy_dispatch(worker_group, *args, **kwargs):
        return args, kwargs

    def dummy_collect(worker_group, output):
        return output

    register_dispatch_mode("TEST_MODE", dummy_dispatch, dummy_collect)

    # Verify enum extension
    _check_dispatch_mode(Dispatch.TEST_MODE)

    # Verify registry update
    assert get_predefined_dispatch_fn(Dispatch.TEST_MODE) == {
        "dispatch_fn": dummy_dispatch,
        "collect_fn": dummy_collect,
    }
    # Clean up
    Dispatch.remove("TEST_MODE")


def test_update_existing_dispatch_mode(reset_dispatch_registry):
    # Store original implementation
    original_mode = Dispatch.ONE_TO_ALL

    # New implementations
    def new_dispatch(worker_group, *args, **kwargs):
        return args, kwargs

    def new_collect(worker_group, output):
        return output

    # Test update=
    update_dispatch_mode(original_mode, new_dispatch, new_collect)

    # Verify update
    assert get_predefined_dispatch_fn(original_mode)["dispatch_fn"] == new_dispatch
    assert get_predefined_dispatch_fn(original_mode)["collect_fn"] == new_collect

	# Copyright 2024 Bytedance Ltd. and/or its affiliates
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import pytest

	import verl.single_controller.base.decorator as decorator_module
	from verl.single_controller.base.decorator import (
	DISPATCH_MODE_FN_REGISTRY,
	Dispatch,
	_check_dispatch_mode,
	get_predefined_dispatch_fn,
	register_dispatch_mode,
	update_dispatch_mode,
	)


	@pytest.fixture
	def reset_dispatch_registry():
	# Store original state
	original_registry = DISPATCH_MODE_FN_REGISTRY.copy()
	yield
	# Reset registry after test
	decorator_module.DISPATCH_MODE_FN_REGISTRY.clear()
	decorator_module.DISPATCH_MODE_FN_REGISTRY.update(original_registry)


	def test_register_new_dispatch_mode(reset_dispatch_registry):
	# Test registration
	def dummy_dispatch(worker_group, args, *kwargs):
	return args, kwargs

	def dummy_collect(worker_group, output):
	return output

	register_dispatch_mode("TEST_MODE", dummy_dispatch, dummy_collect)

	# Verify enum extension
	_check_dispatch_mode(Dispatch.TEST_MODE)

	# Verify registry update
	assert get_predefined_dispatch_fn(Dispatch.TEST_MODE) == {
	"dispatch_fn": dummy_dispatch,
	"collect_fn": dummy_collect,
	}
	# Clean up
	Dispatch.remove("TEST_MODE")


	def test_update_existing_dispatch_mode(reset_dispatch_registry):
	# Store original implementation
	original_mode = Dispatch.ONE_TO_ALL

	# New implementations
	def new_dispatch(worker_group, args, *kwargs):
	return args, kwargs

	def new_collect(worker_group, output):
	return output

	# Test update=
	update_dispatch_mode(original_mode, new_dispatch, new_collect)

	# Verify update
	assert get_predefined_dispatch_fn(original_mode)["dispatch_fn"] == new_dispatch
	assert get_predefined_dispatch_fn(original_mode)["collect_fn"] == new_collect