SubstanceSHIFT/SeedVR2-3B

1# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates2# //3# // Licensed under the Apache License, Version 2.0 (the "License");4# // you may not use this file except in compliance with the License.5# // You may obtain a copy of the License at6# //7# //     http://www.apache.org/licenses/LICENSE-2.08# //9# // Unless required by applicable law or agreed to in writing, software10# // distributed under the License is distributed on an "AS IS" BASIS,11# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# // See the License for the specific language governing permissions and13# // limitations under the License.14 15"""16Decorators.17"""18 19import functools20import threading21import time22from typing import Callable23import torch24 25from common.distributed import barrier_if_distributed, get_global_rank, get_local_rank26from common.logger import get_logger27 28logger = get_logger(__name__)29 30 31def log_on_entry(func: Callable) -> Callable:32    """33    Functions with this decorator will log the function name at entry.34    When using multiple decorators, this must be applied innermost to properly capture the name.35    """36 37    def log_on_entry_wrapper(*args, **kwargs):38        logger.info(f"Entering {func.__name__}")39        return func(*args, **kwargs)40 41    return log_on_entry_wrapper42 43 44def barrier_on_entry(func: Callable) -> Callable:45    """46    Functions with this decorator will start executing when all ranks are ready to enter.47    """48 49    def barrier_on_entry_wrapper(*args, **kwargs):50        barrier_if_distributed()51        return func(*args, **kwargs)52 53    return barrier_on_entry_wrapper54 55 56def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable:57    """58    Helper function for local_rank_zero_only and global_rank_zero_only.59    """60 61    def conditional_execute_wrapper(*args, **kwargs):62        # Only execute if needed.63        result = func(*args, **kwargs) if execute else None64        # All GPUs must wait.65        barrier_if_distributed()66        # Return results.67        return result68 69    return conditional_execute_wrapper70 71 72def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable:73    """74    Helper function for some functions with special constraints,75    especially functions called by other global_rank_zero_only / local_rank_zero_only ones,76    in case they are wrongly invoked in other scenarios.77    """78 79    def asserted_execute_wrapper(*args, **kwargs):80        assert condition, err_msg81        result = func(*args, **kwargs)82        return result83 84    return asserted_execute_wrapper85 86 87def local_rank_zero_only(func: Callable) -> Callable:88    """89    Functions with this decorator will only execute on local rank zero.90    """91    return _conditional_execute_wrapper_factory(get_local_rank() == 0, func)92 93 94def global_rank_zero_only(func: Callable) -> Callable:95    """96    Functions with this decorator will only execute on global rank zero.97    """98    return _conditional_execute_wrapper_factory(get_global_rank() == 0, func)99 100 101def assert_only_global_rank_zero(func: Callable) -> Callable:102    """103    Functions with this decorator are only accessible to processes with global rank zero.104    """105    return _asserted_wrapper_factory(106        get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0"107    )108 109 110def assert_only_local_rank_zero(func: Callable) -> Callable:111    """112    Functions with this decorator are only accessible to processes with local rank zero.113    """114    return _asserted_wrapper_factory(115        get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0"116    )117 118 119def new_thread(func: Callable) -> Callable:120    """121    Functions with this decorator will run in a new thread.122    The function will return the thread, which can be joined to wait for completion.123    """124 125    def new_thread_wrapper(*args, **kwargs):126        thread = threading.Thread(target=func, args=args, kwargs=kwargs)127        thread.start()128        return thread129 130    return new_thread_wrapper131 132 133def log_runtime(func: Callable) -> Callable:134    """135    Functions with this decorator will logging the runtime.136    """137 138    @functools.wraps(func)139    def wrapped(*args, **kwargs):140        torch.distributed.barrier()141        start = time.perf_counter()142        result = func(*args, **kwargs)143        torch.distributed.barrier()144        logger.info(f"Completed {func.__name__} in {time.perf_counter() - start:.3f} seconds.")145        return result146 147    return wrapped148