diff --git a/bin/activate b/bin/activate new file mode 100644 index 0000000000000000000000000000000000000000..3a5c2a8da4a1ac5715e3ccd68e94c56a0b2d032b --- /dev/null +++ b/bin/activate @@ -0,0 +1,76 @@ +# This file must be used with "source bin/activate" *from bash* +# You cannot run it directly + +deactivate () { + # reset old environment variables + if [ -n "${_OLD_VIRTUAL_PATH:-}" ] ; then + PATH="${_OLD_VIRTUAL_PATH:-}" + export PATH + unset _OLD_VIRTUAL_PATH + fi + if [ -n "${_OLD_VIRTUAL_PYTHONHOME:-}" ] ; then + PYTHONHOME="${_OLD_VIRTUAL_PYTHONHOME:-}" + export PYTHONHOME + unset _OLD_VIRTUAL_PYTHONHOME + fi + + # Call hash to forget past locations. Without forgetting + # past locations the $PATH changes we made may not be respected. + # See "man bash" for more details. hash is usually a builtin of your shell + hash -r 2> /dev/null + + if [ -n "${_OLD_VIRTUAL_PS1:-}" ] ; then + PS1="${_OLD_VIRTUAL_PS1:-}" + export PS1 + unset _OLD_VIRTUAL_PS1 + fi + + unset VIRTUAL_ENV + unset VIRTUAL_ENV_PROMPT + if [ ! "${1:-}" = "nondestructive" ] ; then + # Self destruct! + unset -f deactivate + fi +} + +# unset irrelevant variables +deactivate nondestructive + +# on Windows, a path can contain colons and backslashes and has to be converted: +case "$(uname)" in + CYGWIN*|MSYS*|MINGW*) + # transform D:\path\to\venv to /d/path/to/venv on MSYS and MINGW + # and to /cygdrive/d/path/to/venv on Cygwin + VIRTUAL_ENV=$(cygpath /storage/openpsi/users/zzy/.threadweaver) + export VIRTUAL_ENV + ;; + *) + # use the path as-is + export VIRTUAL_ENV=/storage/openpsi/users/zzy/.threadweaver + ;; +esac + +_OLD_VIRTUAL_PATH="$PATH" +PATH="$VIRTUAL_ENV/"bin":$PATH" +export PATH + +VIRTUAL_ENV_PROMPT='(.threadweaver) ' +export VIRTUAL_ENV_PROMPT + +# unset PYTHONHOME if set +# this will fail if PYTHONHOME is set to the empty string (which is bad anyway) +# could use `if (set -u; : $PYTHONHOME) ;` in bash +if [ -n "${PYTHONHOME:-}" ] ; then + _OLD_VIRTUAL_PYTHONHOME="${PYTHONHOME:-}" + unset PYTHONHOME +fi + +if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT:-}" ] ; then + _OLD_VIRTUAL_PS1="${PS1:-}" + PS1="("'(.threadweaver) '") ${PS1:-}" + export PS1 +fi + +# Call hash to forget past commands. Without forgetting +# past commands the $PATH changes we made may not be respected +hash -r 2> /dev/null diff --git a/bin/deepspeed b/bin/deepspeed new file mode 100644 index 0000000000000000000000000000000000000000..d545c08e0148d4ff20d759210264a1ebe9c70b4f --- /dev/null +++ b/bin/deepspeed @@ -0,0 +1,6 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 + +from deepspeed.launcher.runner import main + +if __name__ == '__main__': + main() diff --git a/bin/ds_report b/bin/ds_report new file mode 100644 index 0000000000000000000000000000000000000000..db4df4227d91682d0415cf21ac25a8234d23e034 --- /dev/null +++ b/bin/ds_report @@ -0,0 +1,6 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 + +from deepspeed.env_report import cli_main + +if __name__ == '__main__': + cli_main() diff --git a/bin/fastapi b/bin/fastapi new file mode 100644 index 0000000000000000000000000000000000000000..310ed431e5ebe1d834c064c5f3cc17c7bc791b41 --- /dev/null +++ b/bin/fastapi @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from fastapi.cli import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/huggingface-cli b/bin/huggingface-cli new file mode 100644 index 0000000000000000000000000000000000000000..4922d0d70a1d29955735ac83c26bb40f4347aad2 --- /dev/null +++ b/bin/huggingface-cli @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from huggingface_hub.commands.huggingface_cli import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/ipython3 b/bin/ipython3 new file mode 100644 index 0000000000000000000000000000000000000000..e85bce488945faabd123f9f582c30c998a6a9ef3 --- /dev/null +++ b/bin/ipython3 @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from IPython import start_ipython +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(start_ipython()) diff --git a/bin/isympy b/bin/isympy new file mode 100644 index 0000000000000000000000000000000000000000..d120c2bdac4fbe99453c1ad5a74b59f84d31764b --- /dev/null +++ b/bin/isympy @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from isympy import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/markdown_py b/bin/markdown_py new file mode 100644 index 0000000000000000000000000000000000000000..d32eb437fd755583e28005bdb2fd9bc6fc9db9ea --- /dev/null +++ b/bin/markdown_py @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from markdown.__main__ import run +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(run()) diff --git a/bin/pip3 b/bin/pip3 new file mode 100644 index 0000000000000000000000000000000000000000..3d1df378e55c3783a215f57de7b8f82af1cec126 --- /dev/null +++ b/bin/pip3 @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from pip._internal.cli.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/pip3.12 b/bin/pip3.12 new file mode 100644 index 0000000000000000000000000000000000000000..3d1df378e55c3783a215f57de7b8f82af1cec126 --- /dev/null +++ b/bin/pip3.12 @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from pip._internal.cli.main import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/proton b/bin/proton new file mode 100644 index 0000000000000000000000000000000000000000..17a8f9836cafff5935ff391f18de5a6586b3c2d1 --- /dev/null +++ b/bin/proton @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from triton.profiler.proton import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/pyftmerge b/bin/pyftmerge new file mode 100644 index 0000000000000000000000000000000000000000..df25d39e96cba7f46faa76e875a3347836bb0324 --- /dev/null +++ b/bin/pyftmerge @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from fontTools.merge import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/ray b/bin/ray new file mode 100644 index 0000000000000000000000000000000000000000..2a2f6f4f1a4d728d1b60e447f36b4feda6271d33 --- /dev/null +++ b/bin/ray @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from ray.scripts.scripts import main +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(main()) diff --git a/bin/serve b/bin/serve new file mode 100644 index 0000000000000000000000000000000000000000..fbe492020d0e5d816dd56cfa8039e42e59e3f4d5 --- /dev/null +++ b/bin/serve @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from ray.serve.scripts import cli +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(cli()) diff --git a/bin/wandb b/bin/wandb new file mode 100644 index 0000000000000000000000000000000000000000..0e789f319b3e463c5cab449ed26f0f3c7a608232 --- /dev/null +++ b/bin/wandb @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from wandb.cli.cli import cli +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(cli()) diff --git a/bin/wb b/bin/wb new file mode 100644 index 0000000000000000000000000000000000000000..0e789f319b3e463c5cab449ed26f0f3c7a608232 --- /dev/null +++ b/bin/wb @@ -0,0 +1,8 @@ +#!/storage/openpsi/users/zzy/.threadweaver/bin/python3 +# -*- coding: utf-8 -*- +import re +import sys +from wandb.cli.cli import cli +if __name__ == '__main__': + sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0]) + sys.exit(cli()) diff --git a/lib/python3.12/site-packages/decorator.py b/lib/python3.12/site-packages/decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..40a39f2f1b6d07036d98260318108fdaf173de44 --- /dev/null +++ b/lib/python3.12/site-packages/decorator.py @@ -0,0 +1,459 @@ +# ######################### LICENSE ############################ # + +# Copyright (c) 2005-2025, Michele Simionato +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are +# met: + +# Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# Redistributions in bytecode form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in +# the documentation and/or other materials provided with the +# distribution. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +# HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS +# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR +# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE +# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +# DAMAGE. + +""" +Decorator module, see +https://github.com/micheles/decorator/blob/master/docs/documentation.md +for the documentation. +""" +import re +import sys +import inspect +import operator +import itertools +import functools +from contextlib import _GeneratorContextManager +from inspect import getfullargspec, iscoroutinefunction, isgeneratorfunction + +__version__ = '5.2.1' + +DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(') +POS = inspect.Parameter.POSITIONAL_OR_KEYWORD +EMPTY = inspect.Parameter.empty + + +# this is not used anymore in the core, but kept for backward compatibility +class FunctionMaker(object): + """ + An object with the ability to create functions with a given signature. + It has attributes name, doc, module, signature, defaults, dict and + methods update and make. + """ + + # Atomic get-and-increment provided by the GIL + _compile_count = itertools.count() + + # make pylint happy + args = varargs = varkw = defaults = kwonlyargs = kwonlydefaults = () + + def __init__(self, func=None, name=None, signature=None, + defaults=None, doc=None, module=None, funcdict=None): + self.shortsignature = signature + if func: + # func can be a class or a callable, but not an instance method + self.name = func.__name__ + if self.name == '': # small hack for lambda functions + self.name = '_lambda_' + self.doc = func.__doc__ + self.module = func.__module__ + if inspect.isroutine(func) or isinstance(func, functools.partial): + argspec = getfullargspec(func) + self.annotations = getattr(func, '__annotations__', {}) + for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs', + 'kwonlydefaults'): + setattr(self, a, getattr(argspec, a)) + for i, arg in enumerate(self.args): + setattr(self, 'arg%d' % i, arg) + allargs = list(self.args) + allshortargs = list(self.args) + if self.varargs: + allargs.append('*' + self.varargs) + allshortargs.append('*' + self.varargs) + elif self.kwonlyargs: + allargs.append('*') # single star syntax + for a in self.kwonlyargs: + allargs.append('%s=None' % a) + allshortargs.append('%s=%s' % (a, a)) + if self.varkw: + allargs.append('**' + self.varkw) + allshortargs.append('**' + self.varkw) + self.signature = ', '.join(allargs) + self.shortsignature = ', '.join(allshortargs) + self.dict = func.__dict__.copy() + # func=None happens when decorating a caller + if name: + self.name = name + if signature is not None: + self.signature = signature + if defaults: + self.defaults = defaults + if doc: + self.doc = doc + if module: + self.module = module + if funcdict: + self.dict = funcdict + # check existence required attributes + assert hasattr(self, 'name') + if not hasattr(self, 'signature'): + raise TypeError('You are decorating a non function: %s' % func) + + def update(self, func, **kw): + """ + Update the signature of func with the data in self + """ + func.__name__ = self.name + func.__doc__ = getattr(self, 'doc', None) + func.__dict__ = getattr(self, 'dict', {}) + func.__defaults__ = self.defaults + func.__kwdefaults__ = self.kwonlydefaults or None + func.__annotations__ = getattr(self, 'annotations', None) + try: + frame = sys._getframe(3) + except AttributeError: # for IronPython and similar implementations + callermodule = '?' + else: + callermodule = frame.f_globals.get('__name__', '?') + func.__module__ = getattr(self, 'module', callermodule) + func.__dict__.update(kw) + + def make(self, src_templ, evaldict=None, addsource=False, **attrs): + """ + Make a new function from a given template and update the signature + """ + src = src_templ % vars(self) # expand name and signature + evaldict = evaldict or {} + mo = DEF.search(src) + if mo is None: + raise SyntaxError('not a valid function template\n%s' % src) + name = mo.group(1) # extract the function name + names = set([name] + [arg.strip(' *') for arg in + self.shortsignature.split(',')]) + for n in names: + if n in ('_func_', '_call_'): + raise NameError('%s is overridden in\n%s' % (n, src)) + + if not src.endswith('\n'): # add a newline for old Pythons + src += '\n' + + # Ensure each generated function has a unique filename for profilers + # (such as cProfile) that depend on the tuple of (, + # , ) being unique. + filename = '' % next(self._compile_count) + try: + code = compile(src, filename, 'single') + exec(code, evaldict) + except Exception: + print('Error in generated code:', file=sys.stderr) + print(src, file=sys.stderr) + raise + func = evaldict[name] + if addsource: + attrs['__source__'] = src + self.update(func, **attrs) + return func + + @classmethod + def create(cls, obj, body, evaldict, defaults=None, + doc=None, module=None, addsource=True, **attrs): + """ + Create a function from the strings name, signature and body. + evaldict is the evaluation dictionary. If addsource is true an + attribute __source__ is added to the result. The attributes attrs + are added, if any. + """ + if isinstance(obj, str): # "name(signature)" + name, rest = obj.strip().split('(', 1) + signature = rest[:-1] # strip a right parens + func = None + else: # a function + name = None + signature = None + func = obj + self = cls(func, name, signature, defaults, doc, module) + ibody = '\n'.join(' ' + line for line in body.splitlines()) + caller = evaldict.get('_call_') # when called from `decorate` + if caller and iscoroutinefunction(caller): + body = ('async def %(name)s(%(signature)s):\n' + ibody).replace( + 'return', 'return await') + else: + body = 'def %(name)s(%(signature)s):\n' + ibody + return self.make(body, evaldict, addsource, **attrs) + + +def fix(args, kwargs, sig): + """ + Fix args and kwargs to be consistent with the signature + """ + ba = sig.bind(*args, **kwargs) + ba.apply_defaults() # needed for test_dan_schult + return ba.args, ba.kwargs + + +def decorate(func, caller, extras=(), kwsyntax=False): + """ + Decorates a function/generator/coroutine using a caller. + If kwsyntax is True calling the decorated functions with keyword + syntax will pass the named arguments inside the ``kw`` dictionary, + even if such argument are positional, similarly to what functools.wraps + does. By default kwsyntax is False and the the arguments are untouched. + """ + sig = inspect.signature(func) + if isinstance(func, functools.partial): + func = functools.update_wrapper(func, func.func) + if iscoroutinefunction(caller): + async def fun(*args, **kw): + if not kwsyntax: + args, kw = fix(args, kw, sig) + return await caller(func, *(extras + args), **kw) + elif isgeneratorfunction(caller): + def fun(*args, **kw): + if not kwsyntax: + args, kw = fix(args, kw, sig) + for res in caller(func, *(extras + args), **kw): + yield res + else: + def fun(*args, **kw): + if not kwsyntax: + args, kw = fix(args, kw, sig) + return caller(func, *(extras + args), **kw) + + fun.__name__ = func.__name__ + fun.__doc__ = func.__doc__ + fun.__wrapped__ = func + fun.__signature__ = sig + fun.__qualname__ = func.__qualname__ + # builtin functions like defaultdict.__setitem__ lack many attributes + try: + fun.__defaults__ = func.__defaults__ + except AttributeError: + pass + try: + fun.__kwdefaults__ = func.__kwdefaults__ + except AttributeError: + pass + try: + fun.__annotations__ = func.__annotations__ + except AttributeError: + pass + try: + fun.__module__ = func.__module__ + except AttributeError: + pass + try: + fun.__name__ = func.__name__ + except AttributeError: # happens with old versions of numpy.vectorize + func.__name__ == 'noname' + try: + fun.__dict__.update(func.__dict__) + except AttributeError: + pass + return fun + + +def decoratorx(caller): + """ + A version of "decorator" implemented via "exec" and not via the + Signature object. Use this if you are want to preserve the `.__code__` + object properties (https://github.com/micheles/decorator/issues/129). + """ + def dec(func): + return FunctionMaker.create( + func, + "return _call_(_func_, %(shortsignature)s)", + dict(_call_=caller, _func_=func), + __wrapped__=func, __qualname__=func.__qualname__) + return dec + + +def decorator(caller, _func=None, kwsyntax=False): + """ + decorator(caller) converts a caller function into a decorator + """ + if _func is not None: # return a decorated function + # this is obsolete behavior; you should use decorate instead + return decorate(_func, caller, (), kwsyntax) + # else return a decorator function + sig = inspect.signature(caller) + dec_params = [p for p in sig.parameters.values() if p.kind is POS] + + def dec(func=None, *args, **kw): + na = len(args) + 1 + extras = args + tuple(kw.get(p.name, p.default) + for p in dec_params[na:] + if p.default is not EMPTY) + if func is None: + return lambda func: decorate(func, caller, extras, kwsyntax) + else: + return decorate(func, caller, extras, kwsyntax) + dec.__signature__ = sig.replace(parameters=dec_params) + dec.__name__ = caller.__name__ + dec.__doc__ = caller.__doc__ + dec.__wrapped__ = caller + dec.__qualname__ = caller.__qualname__ + dec.__kwdefaults__ = getattr(caller, '__kwdefaults__', None) + dec.__dict__.update(caller.__dict__) + return dec + + +# ####################### contextmanager ####################### # + + +class ContextManager(_GeneratorContextManager): + def __init__(self, g, *a, **k): + _GeneratorContextManager.__init__(self, g, a, k) + + def __call__(self, func): + def caller(f, *a, **k): + with self.__class__(self.func, *self.args, **self.kwds): + return f(*a, **k) + return decorate(func, caller) + + +_contextmanager = decorator(ContextManager) + + +def contextmanager(func): + # Enable Pylint config: contextmanager-decorators=decorator.contextmanager + return _contextmanager(func) + + +# ############################ dispatch_on ############################ # + +def append(a, vancestors): + """ + Append ``a`` to the list of the virtual ancestors, unless it is already + included. + """ + add = True + for j, va in enumerate(vancestors): + if issubclass(va, a): + add = False + break + if issubclass(a, va): + vancestors[j] = a + add = False + if add: + vancestors.append(a) + + +# inspired from simplegeneric by P.J. Eby and functools.singledispatch +def dispatch_on(*dispatch_args): + """ + Factory of decorators turning a function into a generic function + dispatching on the given arguments. + """ + assert dispatch_args, 'No dispatch args passed' + dispatch_str = '(%s,)' % ', '.join(dispatch_args) + + def check(arguments, wrong=operator.ne, msg=''): + """Make sure one passes the expected number of arguments""" + if wrong(len(arguments), len(dispatch_args)): + raise TypeError('Expected %d arguments, got %d%s' % + (len(dispatch_args), len(arguments), msg)) + + def gen_func_dec(func): + """Decorator turning a function into a generic function""" + + # first check the dispatch arguments + argset = set(getfullargspec(func).args) + if not set(dispatch_args) <= argset: + raise NameError('Unknown dispatch arguments %s' % dispatch_str) + + typemap = {} + + def vancestors(*types): + """ + Get a list of sets of virtual ancestors for the given types + """ + check(types) + ras = [[] for _ in range(len(dispatch_args))] + for types_ in typemap: + for t, type_, ra in zip(types, types_, ras): + if issubclass(t, type_) and type_ not in t.mro(): + append(type_, ra) + return [set(ra) for ra in ras] + + def ancestors(*types): + """ + Get a list of virtual MROs, one for each type + """ + check(types) + lists = [] + for t, vas in zip(types, vancestors(*types)): + n_vas = len(vas) + if n_vas > 1: + raise RuntimeError( + 'Ambiguous dispatch for %s: %s' % (t, vas)) + elif n_vas == 1: + va, = vas + mro = type('t', (t, va), {}).mro()[1:] + else: + mro = t.mro() + lists.append(mro[:-1]) # discard t and object + return lists + + def register(*types): + """ + Decorator to register an implementation for the given types + """ + check(types) + + def dec(f): + check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__) + typemap[types] = f + return f + return dec + + def dispatch_info(*types): + """ + An utility to introspect the dispatch algorithm + """ + check(types) + lst = [] + for ancs in itertools.product(*ancestors(*types)): + lst.append(tuple(a.__name__ for a in ancs)) + return lst + + def _dispatch(dispatch_args, *args, **kw): + types = tuple(type(arg) for arg in dispatch_args) + try: # fast path + f = typemap[types] + except KeyError: + pass + else: + return f(*args, **kw) + combinations = itertools.product(*ancestors(*types)) + next(combinations) # the first one has been already tried + for types_ in combinations: + f = typemap.get(types_) + if f is not None: + return f(*args, **kw) + + # else call the default implementation + return func(*args, **kw) + + return FunctionMaker.create( + func, 'return _f_(%s, %%(shortsignature)s)' % dispatch_str, + dict(_f_=_dispatch), register=register, default=func, + typemap=typemap, vancestors=vancestors, ancestors=ancestors, + dispatch_info=dispatch_info, __wrapped__=func) + + gen_func_dec.__name__ = 'dispatch_on' + dispatch_str + return gen_func_dec diff --git a/lib/python3.12/site-packages/ipython_pygments_lexers.py b/lib/python3.12/site-packages/ipython_pygments_lexers.py new file mode 100644 index 0000000000000000000000000000000000000000..c655e233654b8ddf8d9a18638451a1ad6937e1d8 --- /dev/null +++ b/lib/python3.12/site-packages/ipython_pygments_lexers.py @@ -0,0 +1,582 @@ +# -*- coding: utf-8 -*- +""" +Defines a variety of Pygments lexers for highlighting IPython code. + +This includes: + + IPythonLexer, IPython3Lexer + Lexers for pure IPython (python + magic/shell commands) + + IPythonPartialTracebackLexer, IPythonTracebackLexer + Supports 2.x and 3.x via keyword `python3`. The partial traceback + lexer reads everything but the Python code appearing in a traceback. + The full lexer combines the partial lexer with an IPython lexer. + + IPythonConsoleLexer + A lexer for IPython console sessions, with support for tracebacks. + + IPyLexer + A friendly lexer which examines the first line of text and from it, + decides whether to use an IPython lexer or an IPython console lexer. + This is probably the only lexer that needs to be explicitly added + to Pygments. + +""" +# ----------------------------------------------------------------------------- +# Copyright (c) 2013, the IPython Development Team. +# +# Distributed under the terms of the Modified BSD License. +# +# The full license is in the file COPYING.txt, distributed with this software. +# ----------------------------------------------------------------------------- + +__version__ = "1.1.1" + +# Standard library +import re + +# Third party +from pygments.lexers import ( + BashLexer, + HtmlLexer, + JavascriptLexer, + RubyLexer, + PerlLexer, + Python2Lexer, + Python3Lexer, + TexLexer, +) +from pygments.lexer import ( + Lexer, + DelegatingLexer, + RegexLexer, + do_insertions, + bygroups, + using, +) +from pygments.token import ( + Generic, + Keyword, + Literal, + Name, + Operator, + Other, + Text, + Error, +) + + +line_re = re.compile(".*?\n") + +__all__ = [ + "IPython3Lexer", + "IPythonLexer", + "IPythonPartialTracebackLexer", + "IPythonTracebackLexer", + "IPythonConsoleLexer", + "IPyLexer", +] + + +ipython_tokens = [ + ( + r"(?s)(\s*)(%%capture)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%debug)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?is)(\s*)(%%html)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(HtmlLexer)), + ), + ( + r"(?s)(\s*)(%%javascript)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(JavascriptLexer)), + ), + ( + r"(?s)(\s*)(%%js)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(JavascriptLexer)), + ), + ( + r"(?s)(\s*)(%%latex)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(TexLexer)), + ), + ( + r"(?s)(\s*)(%%perl)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(PerlLexer)), + ), + ( + r"(?s)(\s*)(%%prun)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%pypy)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%python2)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python2Lexer)), + ), + ( + r"(?s)(\s*)(%%python3)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%python)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%ruby)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(RubyLexer)), + ), + ( + r"(?s)(\s*)(%%timeit)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%time)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%writefile)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + ( + r"(?s)(\s*)(%%file)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(Python3Lexer)), + ), + (r"(?s)(\s*)(%%)(\w+)(.*)", bygroups(Text, Operator, Keyword, Text)), + ( + r"(?s)(^\s*)(%%!)([^\n]*\n)(.*)", + bygroups(Text, Operator, Text, using(BashLexer)), + ), + (r"(%%?)(\w+)(\?\??)$", bygroups(Operator, Keyword, Operator)), + (r"\b(\?\??)(\s*)$", bygroups(Operator, Text)), + (r"(%)(sx|sc|system)(.*)(\n)", bygroups(Operator, Keyword, using(BashLexer), Text)), + (r"(%)(\w+)(.*\n)", bygroups(Operator, Keyword, Text)), + (r"^(!!)(.+)(\n)", bygroups(Operator, using(BashLexer), Text)), + (r"(!)(?!=)(.+)(\n)", bygroups(Operator, using(BashLexer), Text)), + (r"^(\s*)(\?\??)(\s*%{0,2}[\w\.\*]*)", bygroups(Text, Operator, Text)), + (r"(\s*%{0,2}[\w\.\*]*)(\?\??)(\s*)$", bygroups(Text, Operator, Text)), +] + + +class IPython3Lexer(Python3Lexer): + """IPython code lexer (based on Python 3)""" + + name = "IPython" + aliases = ["ipython", "ipython3"] + + tokens = Python3Lexer.tokens.copy() + tokens["root"] = ipython_tokens + tokens["root"] + + +IPythonLexer = IPython3Lexer + + +class IPythonPartialTracebackLexer(RegexLexer): + """ + Partial lexer for IPython tracebacks. + + Handles all the non-python output. + + """ + + name = "IPython Partial Traceback" + + tokens = { + "root": [ + # Tracebacks for syntax errors have a different style. + # For both types of tracebacks, we mark the first line with + # Generic.Traceback. For syntax errors, we mark the filename + # as we mark the filenames for non-syntax tracebacks. + # + # These two regexps define how IPythonConsoleLexer finds a + # traceback. + # + ## Non-syntax traceback + (r"^(\^C)?(-+\n)", bygroups(Error, Generic.Traceback)), + ## Syntax traceback + ( + r"^( File)(.*)(, line )(\d+\n)", + bygroups( + Generic.Traceback, + Name.Namespace, + Generic.Traceback, + Literal.Number.Integer, + ), + ), + # (Exception Identifier)(Whitespace)(Traceback Message) + ( + r"(?u)(^[^\d\W]\w*)(\s*)(Traceback.*?\n)", + bygroups(Name.Exception, Generic.Whitespace, Text), + ), + # (Module/Filename)(Text)(Callee)(Function Signature) + # Better options for callee and function signature? + ( + r"(.*)( in )(.*)(\(.*\)\n)", + bygroups(Name.Namespace, Text, Name.Entity, Name.Tag), + ), + # Regular line: (Whitespace)(Line Number)(Python Code) + ( + r"(\s*?)(\d+)(.*?\n)", + bygroups(Generic.Whitespace, Literal.Number.Integer, Other), + ), + # Emphasized line: (Arrow)(Line Number)(Python Code) + # Using Exception token so arrow color matches the Exception. + ( + r"(-*>?\s?)(\d+)(.*?\n)", + bygroups(Name.Exception, Literal.Number.Integer, Other), + ), + # (Exception Identifier)(Message) + (r"(?u)(^[^\d\W]\w*)(:.*?\n)", bygroups(Name.Exception, Text)), + # Tag everything else as Other, will be handled later. + (r".*\n", Other), + ], + } + + +class IPythonTracebackLexer(DelegatingLexer): + """ + IPython traceback lexer. + + For doctests, the tracebacks can be snipped as much as desired with the + exception to the lines that designate a traceback. For non-syntax error + tracebacks, this is the line of hyphens. For syntax error tracebacks, + this is the line which lists the File and line number. + + """ + + # The lexer inherits from DelegatingLexer. The "root" lexer is an + # appropriate IPython lexer, which depends on the value of the boolean + # `python3`. First, we parse with the partial IPython traceback lexer. + # Then, any code marked with the "Other" token is delegated to the root + # lexer. + # + name = "IPython Traceback" + aliases = ["ipythontb", "ipython3tb"] + + def __init__(self, **options): + """ + A subclass of `DelegatingLexer` which delegates to the appropriate to either IPyLexer, + IPythonPartialTracebackLexer. + """ + # note we need a __init__ doc, as otherwise it inherits the doc from the super class + # which will fail the documentation build as it references section of the pygments docs that + # do not exists when building IPython's docs. + DelegatingLexer.__init__( + self, IPython3Lexer, IPythonPartialTracebackLexer, **options + ) + + +class IPythonConsoleLexer(Lexer): + """ + An IPython console lexer for IPython code-blocks and doctests, such as: + + .. code-block:: rst + + .. code-block:: ipythonconsole + + In [1]: a = 'foo' + + In [2]: a + Out[2]: 'foo' + + In [3]: print(a) + foo + + + Support is also provided for IPython exceptions: + + .. code-block:: rst + + .. code-block:: ipythonconsole + + In [1]: raise Exception + Traceback (most recent call last): + ... + Exception + + """ + + name = "IPython console session" + aliases = ["ipythonconsole", "ipython3console"] + mimetypes = ["text/x-ipython-console"] + + # The regexps used to determine what is input and what is output. + # The default prompts for IPython are: + # + # in = 'In [#]: ' + # continuation = ' .D.: ' + # template = 'Out[#]: ' + # + # Where '#' is the 'prompt number' or 'execution count' and 'D' + # D is a number of dots matching the width of the execution count + # + in1_regex = r"In \[[0-9]+\]: " + in2_regex = r" \.\.+\.: " + out_regex = r"Out\[[0-9]+\]: " + + #: The regex to determine when a traceback starts. + ipytb_start = re.compile(r"^(\^C)?(-+\n)|^( File)(.*)(, line )(\d+\n)") + + def __init__(self, **options): + """Initialize the IPython console lexer. + + Parameters + ---------- + in1_regex : RegexObject + The compiled regular expression used to detect the start + of inputs. Although the IPython configuration setting may have a + trailing whitespace, do not include it in the regex. If `None`, + then the default input prompt is assumed. + in2_regex : RegexObject + The compiled regular expression used to detect the continuation + of inputs. Although the IPython configuration setting may have a + trailing whitespace, do not include it in the regex. If `None`, + then the default input prompt is assumed. + out_regex : RegexObject + The compiled regular expression used to detect outputs. If `None`, + then the default output prompt is assumed. + + """ + in1_regex = options.get("in1_regex", self.in1_regex) + in2_regex = options.get("in2_regex", self.in2_regex) + out_regex = options.get("out_regex", self.out_regex) + + # So that we can work with input and output prompts which have been + # rstrip'd (possibly by editors) we also need rstrip'd variants. If + # we do not do this, then such prompts will be tagged as 'output'. + # The reason can't just use the rstrip'd variants instead is because + # we want any whitespace associated with the prompt to be inserted + # with the token. This allows formatted code to be modified so as hide + # the appearance of prompts, with the whitespace included. One example + # use of this is in copybutton.js from the standard lib Python docs. + in1_regex_rstrip = in1_regex.rstrip() + "\n" + in2_regex_rstrip = in2_regex.rstrip() + "\n" + out_regex_rstrip = out_regex.rstrip() + "\n" + + # Compile and save them all. + attrs = [ + "in1_regex", + "in2_regex", + "out_regex", + "in1_regex_rstrip", + "in2_regex_rstrip", + "out_regex_rstrip", + ] + for attr in attrs: + self.__setattr__(attr, re.compile(locals()[attr])) + + Lexer.__init__(self, **options) + + self.pylexer = IPython3Lexer(**options) + self.tblexer = IPythonTracebackLexer(**options) + + self.reset() + + def reset(self): + self.mode = "output" + self.index = 0 + self.buffer = "" + self.insertions = [] + + def buffered_tokens(self): + """ + Generator of unprocessed tokens after doing insertions and before + changing to a new state. + + """ + if self.mode == "output": + tokens = [(0, Generic.Output, self.buffer)] + elif self.mode == "input": + tokens = self.pylexer.get_tokens_unprocessed(self.buffer) + else: # traceback + tokens = self.tblexer.get_tokens_unprocessed(self.buffer) + + for i, t, v in do_insertions(self.insertions, tokens): + # All token indexes are relative to the buffer. + yield self.index + i, t, v + + # Clear it all + self.index += len(self.buffer) + self.buffer = "" + self.insertions = [] + + def get_mci(self, line): + """ + Parses the line and returns a 3-tuple: (mode, code, insertion). + + `mode` is the next mode (or state) of the lexer, and is always equal + to 'input', 'output', or 'tb'. + + `code` is a portion of the line that should be added to the buffer + corresponding to the next mode and eventually lexed by another lexer. + For example, `code` could be Python code if `mode` were 'input'. + + `insertion` is a 3-tuple (index, token, text) representing an + unprocessed "token" that will be inserted into the stream of tokens + that are created from the buffer once we change modes. This is usually + the input or output prompt. + + In general, the next mode depends on current mode and on the contents + of `line`. + + """ + # To reduce the number of regex match checks, we have multiple + # 'if' blocks instead of 'if-elif' blocks. + + # Check for possible end of input + in2_match = self.in2_regex.match(line) + in2_match_rstrip = self.in2_regex_rstrip.match(line) + if ( + in2_match and in2_match.group().rstrip() == line.rstrip() + ) or in2_match_rstrip: + end_input = True + else: + end_input = False + if end_input and self.mode != "tb": + # Only look for an end of input when not in tb mode. + # An ellipsis could appear within the traceback. + mode = "output" + code = "" + insertion = (0, Generic.Prompt, line) + return mode, code, insertion + + # Check for output prompt + out_match = self.out_regex.match(line) + out_match_rstrip = self.out_regex_rstrip.match(line) + if out_match or out_match_rstrip: + mode = "output" + if out_match: + idx = out_match.end() + else: + idx = out_match_rstrip.end() + code = line[idx:] + # Use the 'heading' token for output. We cannot use Generic.Error + # since it would conflict with exceptions. + insertion = (0, Generic.Heading, line[:idx]) + return mode, code, insertion + + # Check for input or continuation prompt (non stripped version) + in1_match = self.in1_regex.match(line) + if in1_match or (in2_match and self.mode != "tb"): + # New input or when not in tb, continued input. + # We do not check for continued input when in tb since it is + # allowable to replace a long stack with an ellipsis. + mode = "input" + if in1_match: + idx = in1_match.end() + else: # in2_match + idx = in2_match.end() + code = line[idx:] + insertion = (0, Generic.Prompt, line[:idx]) + return mode, code, insertion + + # Check for input or continuation prompt (stripped version) + in1_match_rstrip = self.in1_regex_rstrip.match(line) + if in1_match_rstrip or (in2_match_rstrip and self.mode != "tb"): + # New input or when not in tb, continued input. + # We do not check for continued input when in tb since it is + # allowable to replace a long stack with an ellipsis. + mode = "input" + if in1_match_rstrip: + idx = in1_match_rstrip.end() + else: # in2_match + idx = in2_match_rstrip.end() + code = line[idx:] + insertion = (0, Generic.Prompt, line[:idx]) + return mode, code, insertion + + # Check for traceback + if self.ipytb_start.match(line): + mode = "tb" + code = line + insertion = None + return mode, code, insertion + + # All other stuff... + if self.mode in ("input", "output"): + # We assume all other text is output. Multiline input that + # does not use the continuation marker cannot be detected. + # For example, the 3 in the following is clearly output: + # + # In [1]: print(3) + # 3 + # + # But the following second line is part of the input: + # + # In [2]: while True: + # print(True) + # + # In both cases, the 2nd line will be 'output'. + # + mode = "output" + else: + mode = "tb" + + code = line + insertion = None + + return mode, code, insertion + + def get_tokens_unprocessed(self, text): + self.reset() + for match in line_re.finditer(text): + line = match.group() + mode, code, insertion = self.get_mci(line) + + if mode != self.mode: + # Yield buffered tokens before transitioning to new mode. + for token in self.buffered_tokens(): + yield token + self.mode = mode + + if insertion: + self.insertions.append((len(self.buffer), [insertion])) + self.buffer += code + + for token in self.buffered_tokens(): + yield token + + +class IPyLexer(Lexer): + r""" + Primary lexer for all IPython-like code. + + This is a simple helper lexer. If the first line of the text begins with + "In \[[0-9]+\]:", then the entire text is parsed with an IPython console + lexer. If not, then the entire text is parsed with an IPython lexer. + + The goal is to reduce the number of lexers that are registered + with Pygments. + + """ + + name = "IPy session" + aliases = ["ipy", "ipy3"] + + def __init__(self, **options): + """ + Create a new IPyLexer instance which dispatch to either an + IPythonCOnsoleLexer (if In prompts are present) or and IPythonLexer (if + In prompts are not present). + """ + # init docstring is necessary for docs not to fail to build do to parent + # docs referenceing a section in pygments docs. + Lexer.__init__(self, **options) + + self.IPythonLexer = IPythonLexer(**options) + self.IPythonConsoleLexer = IPythonConsoleLexer(**options) + + def get_tokens_unprocessed(self, text): + # Search for the input prompt anywhere...this allows code blocks to + # begin with comments as well. + if re.match(r".*(In \[[0-9]+\]:)", text.strip(), re.DOTALL): + lex = self.IPythonConsoleLexer + else: + lex = self.IPythonLexer + for token in lex.get_tokens_unprocessed(text): + yield token diff --git a/lib/python3.12/site-packages/isympy.py b/lib/python3.12/site-packages/isympy.py new file mode 100644 index 0000000000000000000000000000000000000000..50e9bc78d08904b8c177105ee90d984ea4b01d20 --- /dev/null +++ b/lib/python3.12/site-packages/isympy.py @@ -0,0 +1,342 @@ +""" +Python shell for SymPy. + +This is just a normal Python shell (IPython shell if you have the +IPython package installed), that executes the following commands for +the user: + + >>> from __future__ import division + >>> from sympy import * + >>> x, y, z, t = symbols('x y z t') + >>> k, m, n = symbols('k m n', integer=True) + >>> f, g, h = symbols('f g h', cls=Function) + >>> init_printing() + +So starting 'isympy' is equivalent to starting Python (or IPython) and +executing the above commands by hand. It is intended for easy and quick +experimentation with SymPy. isympy is a good way to use SymPy as an +interactive calculator. If you have IPython and Matplotlib installed, then +interactive plotting is enabled by default. + +COMMAND LINE OPTIONS +-------------------- + +-c CONSOLE, --console=CONSOLE + + Use the specified shell (Python or IPython) shell as the console + backend instead of the default one (IPython if present, Python + otherwise), e.g.: + + $isympy -c python + + CONSOLE must be one of 'ipython' or 'python' + +-p PRETTY, --pretty PRETTY + + Setup pretty-printing in SymPy. When pretty-printing is enabled, + expressions can be printed with Unicode or ASCII. The default is + to use pretty-printing (with Unicode if the terminal supports it). + When this option is 'no', expressions will not be pretty-printed + and ASCII will be used: + + $isympy -p no + + PRETTY must be one of 'unicode', 'ascii', or 'no' + +-t TYPES, --types=TYPES + + Setup the ground types for the polys. By default, gmpy ground types + are used if gmpy2 or gmpy is installed, otherwise it falls back to python + ground types, which are a little bit slower. You can manually + choose python ground types even if gmpy is installed (e.g., for + testing purposes): + + $isympy -t python + + TYPES must be one of 'gmpy', 'gmpy1' or 'python' + + Note that the ground type gmpy1 is primarily intended for testing; it + forces the use of gmpy version 1 even if gmpy2 is available. + + This is the same as setting the environment variable + SYMPY_GROUND_TYPES to the given ground type (e.g., + SYMPY_GROUND_TYPES='gmpy') + + The ground types can be determined interactively from the variable + sympy.polys.domains.GROUND_TYPES. + +-o ORDER, --order ORDER + + Setup the ordering of terms for printing. The default is lex, which + orders terms lexicographically (e.g., x**2 + x + 1). You can choose + other orderings, such as rev-lex, which will use reverse + lexicographic ordering (e.g., 1 + x + x**2): + + $isympy -o rev-lex + + ORDER must be one of 'lex', 'rev-lex', 'grlex', 'rev-grlex', + 'grevlex', 'rev-grevlex', 'old', or 'none'. + + Note that for very large expressions, ORDER='none' may speed up + printing considerably but the terms will have no canonical order. + +-q, --quiet + + Print only Python's and SymPy's versions to stdout at startup. + +-d, --doctest + + Use the same format that should be used for doctests. This is + equivalent to -c python -p no. + +-C, --no-cache + + Disable the caching mechanism. Disabling the cache may slow certain + operations down considerably. This is useful for testing the cache, + or for benchmarking, as the cache can result in deceptive timings. + + This is equivalent to setting the environment variable + SYMPY_USE_CACHE to 'no'. + +-a, --auto-symbols (requires at least IPython 0.11) + + Automatically create missing symbols. Normally, typing a name of a + Symbol that has not been instantiated first would raise NameError, + but with this option enabled, any undefined name will be + automatically created as a Symbol. + + Note that this is intended only for interactive, calculator style + usage. In a script that uses SymPy, Symbols should be instantiated + at the top, so that it's clear what they are. + + This will not override any names that are already defined, which + includes the single character letters represented by the mnemonic + QCOSINE (see the "Gotchas and Pitfalls" document in the + documentation). You can delete existing names by executing "del + name". If a name is defined, typing "'name' in dir()" will return True. + + The Symbols that are created using this have default assumptions. + If you want to place assumptions on symbols, you should create them + using symbols() or var(). + + Finally, this only works in the top level namespace. So, for + example, if you define a function in isympy with an undefined + Symbol, it will not work. + + See also the -i and -I options. + +-i, --int-to-Integer (requires at least IPython 0.11) + + Automatically wrap int literals with Integer. This makes it so that + things like 1/2 will come out as Rational(1, 2), rather than 0.5. This + works by preprocessing the source and wrapping all int literals with + Integer. Note that this will not change the behavior of int literals + assigned to variables, and it also won't change the behavior of functions + that return int literals. + + If you want an int, you can wrap the literal in int(), e.g. int(3)/int(2) + gives 1.5 (with division imported from __future__). + +-I, --interactive (requires at least IPython 0.11) + + This is equivalent to --auto-symbols --int-to-Integer. Future options + designed for ease of interactive use may be added to this. + +-D, --debug + + Enable debugging output. This is the same as setting the + environment variable SYMPY_DEBUG to 'True'. The debug status is set + in the variable SYMPY_DEBUG within isympy. + +-- IPython options + + Additionally you can pass command line options directly to the IPython + interpreter (the standard Python shell is not supported). However you + need to add the '--' separator between two types of options, e.g the + startup banner option and the colors option. You need to enter the + options as required by the version of IPython that you are using, too: + + in IPython 0.11, + + $isympy -q -- --colors=NoColor + + or older versions of IPython, + + $isympy -q -- -colors NoColor + +See also isympy --help. +""" + +import os +import sys + +# DO NOT IMPORT SYMPY HERE! Or the setting of the sympy environment variables +# by the command line will break. + +def main() -> None: + from argparse import ArgumentParser, RawDescriptionHelpFormatter + + VERSION = None + if '--version' in sys.argv: + # We cannot import sympy before this is run, because flags like -C and + # -t set environment variables that must be set before SymPy is + # imported. The only thing we need to import it for is to get the + # version, which only matters with the --version flag. + import sympy + VERSION = sympy.__version__ + + usage = 'isympy [options] -- [ipython options]' + parser = ArgumentParser( + usage=usage, + description=__doc__, + formatter_class=RawDescriptionHelpFormatter, + ) + + parser.add_argument('--version', action='version', version=VERSION) + + parser.add_argument( + '-c', '--console', + dest='console', + action='store', + default=None, + choices=['ipython', 'python'], + metavar='CONSOLE', + help='select type of interactive session: ipython | python; defaults ' + 'to ipython if IPython is installed, otherwise python') + + parser.add_argument( + '-p', '--pretty', + dest='pretty', + action='store', + default=None, + metavar='PRETTY', + choices=['unicode', 'ascii', 'no'], + help='setup pretty printing: unicode | ascii | no; defaults to ' + 'unicode printing if the terminal supports it, otherwise ascii') + + parser.add_argument( + '-t', '--types', + dest='types', + action='store', + default=None, + metavar='TYPES', + choices=['gmpy', 'gmpy1', 'python'], + help='setup ground types: gmpy | gmpy1 | python; defaults to gmpy if gmpy2 ' + 'or gmpy is installed, otherwise python') + + parser.add_argument( + '-o', '--order', + dest='order', + action='store', + default=None, + metavar='ORDER', + choices=['lex', 'grlex', 'grevlex', 'rev-lex', 'rev-grlex', 'rev-grevlex', 'old', 'none'], + help='setup ordering of terms: [rev-]lex | [rev-]grlex | [rev-]grevlex | old | none; defaults to lex') + + parser.add_argument( + '-q', '--quiet', + dest='quiet', + action='store_true', + default=False, + help='print only version information at startup') + + parser.add_argument( + '-d', '--doctest', + dest='doctest', + action='store_true', + default=False, + help='use the doctest format for output (you can just copy and paste it)') + + parser.add_argument( + '-C', '--no-cache', + dest='cache', + action='store_false', + default=True, + help='disable caching mechanism') + + parser.add_argument( + '-a', '--auto-symbols', + dest='auto_symbols', + action='store_true', + default=False, + help='automatically construct missing symbols') + + parser.add_argument( + '-i', '--int-to-Integer', + dest='auto_int_to_Integer', + action='store_true', + default=False, + help="automatically wrap int literals with Integer") + + parser.add_argument( + '-I', '--interactive', + dest='interactive', + action='store_true', + default=False, + help="equivalent to -a -i") + + parser.add_argument( + '-D', '--debug', + dest='debug', + action='store_true', + default=False, + help='enable debugging output') + + (options, ipy_args) = parser.parse_known_args() + if '--' in ipy_args: + ipy_args.remove('--') + + if not options.cache: + os.environ['SYMPY_USE_CACHE'] = 'no' + + if options.types: + os.environ['SYMPY_GROUND_TYPES'] = options.types + + if options.debug: + os.environ['SYMPY_DEBUG'] = str(options.debug) + + if options.doctest: + options.pretty = 'no' + options.console = 'python' + + session = options.console + + if session is not None: + ipython = session == 'ipython' + else: + try: + import IPython + ipython = True + except ImportError: + if not options.quiet: + from sympy.interactive.session import no_ipython + print(no_ipython) + ipython = False + + args = { + 'pretty_print': True, + 'use_unicode': None, + 'use_latex': None, + 'order': None, + 'argv': ipy_args, + } + + if options.pretty == 'unicode': + args['use_unicode'] = True + elif options.pretty == 'ascii': + args['use_unicode'] = False + elif options.pretty == 'no': + args['pretty_print'] = False + + if options.order is not None: + args['order'] = options.order + + args['quiet'] = options.quiet + args['auto_symbols'] = options.auto_symbols or options.interactive + args['auto_int_to_Integer'] = options.auto_int_to_Integer or options.interactive + + from sympy.interactive import init_session + init_session(ipython, **args) + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/package_readme.md b/lib/python3.12/site-packages/package_readme.md new file mode 100644 index 0000000000000000000000000000000000000000..2e935d793b15fa6cb2043cbd326a44e3fb539adb --- /dev/null +++ b/lib/python3.12/site-packages/package_readme.md @@ -0,0 +1,97 @@ +
+

+
+ +# Weights and Biases [![PyPI](https://img.shields.io/pypi/v/wandb)](https://pypi.python.org/pypi/wandb) [![Conda (channel only)](https://img.shields.io/conda/vn/conda-forge/wandb)](https://anaconda.org/conda-forge/wandb) [![CircleCI](https://img.shields.io/circleci/build/github/wandb/wandb/main)](https://circleci.com/gh/wandb/wandb) [![Codecov](https://img.shields.io/codecov/c/gh/wandb/wandb)](https://codecov.io/gh/wandb/wandb) + +Use W&B to build better models faster. Track and visualize all the pieces of your machine learning pipeline, from datasets to production machine learning models. Get started with W&B today, [sign up for an account!](https://wandb.com?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=readme) + + + +See the [W&B Developer Guide](https://docs.wandb.ai/?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=documentation) and [API Reference Guide](https://docs.wandb.ai/ref?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=documentation) for a full technical description of the W&B platform. + +  + +# Quickstart + +Get started with W&B in four steps: + +1. First, sign up for a [W&B account](https://wandb.ai/login?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=quickstart). + +2. Second, install the W&B SDK with [pip](https://pip.pypa.io/en/stable/). Navigate to your terminal and type the following command: + +```shell +pip install wandb +``` + +3. Third, log into W&B: + +```python +wandb.login() +``` + +4. Use the example code snippet below as a template to integrate W&B to your Python script: + +```python +import wandb + +# Start a W&B Run with wandb.init +run = wandb.init(project="my_first_project") + +# Save model inputs and hyperparameters in a wandb.config object +config = run.config +config.learning_rate = 0.01 + +# Model training code here ... + +# Log metrics over time to visualize performance with wandb.log +for i in range(10): + run.log({"loss": ...}) + +# Mark the run as finished, and finish uploading all data +run.finish() +``` + +For example, if the preceding code was stored in a script called train.py: + +```shell +python train.py +``` + +You will see a URL in your terminal logs when your script starts and finishes. Data is staged locally in a directory named _wandb_ relative to your script. Navigate to the W&B App to view a dashboard of your first W&B Experiment. Use the W&B App to compare multiple experiments in a unified place, dive into the results of a single run, and much more! + +  + +# Integrations + +Use your favorite framework with W&B. W&B integrations make it fast and easy to set up experiment tracking and data versioning inside existing projects. For more information on how to integrate W&B with the framework of your choice, see [W&B Integrations](https://docs.wandb.ai/guides/integrations) in the W&B Developer Guide. + +  + +# Python Version Support + +We are committed to supporting our minimum required Python version for *at least* six months after its official end-of-life (EOL) date, as defined by the Python Software Foundation. You can find a list of Python EOL dates [here](https://devguide.python.org/versions/). + +When we discontinue support for a Python version, we will increment the library’s minor version number to reflect this change. + +  + +# Contribution guidelines +Weights & Biases ❤️ open source, and we welcome contributions from the community! See the [Contribution guide](https://github.com/wandb/wandb/blob/main/CONTRIBUTING.md) for more information on the development workflow and the internals of the wandb library. For wandb bugs and feature requests, visit [GitHub Issues](https://github.com/wandb/wandb/issues) or contact support@wandb.com. + +  + +# Academic Researchers +Reach out to W&B Support at support@wandb.com to get a [free academic license](https://www.wandb.com/academic) for you and your research group. + +  + +# W&B Community + +Be a part of the growing W&B Community and interact with the W&B team in our [Discord](https://wandb.me/discord). Stay connected with the latest ML updates and tutorials with [W&B Fully Connected](https://wandb.ai/fully-connected). + +  + +# License + +[MIT License](https://github.com/wandb/wandb/blob/main/LICENSE) diff --git a/lib/python3.12/site-packages/pylab.py b/lib/python3.12/site-packages/pylab.py new file mode 100644 index 0000000000000000000000000000000000000000..64ece55b1cb2c312d8a9c2ac1fd264fa11727270 --- /dev/null +++ b/lib/python3.12/site-packages/pylab.py @@ -0,0 +1,3 @@ +from matplotlib.pylab import * # noqa: F401, F403 +import matplotlib.pylab +__doc__ = matplotlib.pylab.__doc__ diff --git a/lib/python3.12/site-packages/sglang/__init__.py b/lib/python3.12/site-packages/sglang/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..91dc58e53feec780ce22c658851c36cce953f0f7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/__init__.py @@ -0,0 +1,76 @@ +# SGLang public APIs + +# Frontend Language APIs +from sglang.api import ( + Engine, + Runtime, + assistant, + assistant_begin, + assistant_end, + flush_cache, + function, + gen, + gen_int, + gen_string, + get_server_info, + image, + select, + set_default_backend, + system, + system_begin, + system_end, + user, + user_begin, + user_end, + video, +) +from sglang.global_config import global_config +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint +from sglang.lang.choices import ( + greedy_token_selection, + token_length_normalized, + unconditional_likelihood_normalized, +) +from sglang.utils import LazyImport +from sglang.version import __version__ + +ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs") +Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic") +LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM") +OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI") +VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI") + +__all__ = [ + "Engine", + "Runtime", + "assistant", + "assistant_begin", + "assistant_end", + "flush_cache", + "function", + "gen", + "gen_int", + "gen_string", + "get_server_info", + "image", + "select", + "set_default_backend", + "system", + "system_begin", + "system_end", + "user", + "user_begin", + "user_end", + "video", + "RuntimeEndpoint", + "greedy_token_selection", + "token_length_normalized", + "unconditional_likelihood_normalized", + "ServerArgs", + "Anthropic", + "LiteLLM", + "OpenAI", + "VertexAI", + "global_config", + "__version__", +] diff --git a/lib/python3.12/site-packages/sglang/api.py b/lib/python3.12/site-packages/sglang/api.py new file mode 100644 index 0000000000000000000000000000000000000000..50319c45178e37d094429704601218c3a4472a7f --- /dev/null +++ b/lib/python3.12/site-packages/sglang/api.py @@ -0,0 +1,279 @@ +"""Public APIs of the language.""" + +import re +from typing import Callable, List, Optional, Union + +from sglang.global_config import global_config +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized +from sglang.lang.ir import ( + SglExpr, + SglExprList, + SglFunction, + SglGen, + SglImage, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglVideo, +) + + +def function( + func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None +): + if func: + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + def decorator(func): + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + return decorator + + +def Runtime(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.lang.backend.runtime_endpoint import Runtime + + return Runtime(*args, **kwargs) + + +def Engine(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.srt.entrypoints.engine import Engine + + return Engine(*args, **kwargs) + + +def set_default_backend(backend: BaseBackend): + global_config.default_backend = backend + + +def flush_cache(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return False + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.flush_cache() + + +def get_server_info(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return None + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.get_server_info() + + +def gen( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + min_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + dtype: Optional[Union[type, str]] = None, + choices: Optional[List[str]] = None, + choices_method: Optional[ChoicesSamplingMethod] = None, + regex: Optional[str] = None, + json_schema: Optional[str] = None, +): + """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" + + if choices: + return SglSelect( + name, + choices, + 0.0 if temperature is None else temperature, + token_length_normalized if choices_method is None else choices_method, + ) + + # check regex is valid + if regex is not None: + try: + re.compile(regex) + except re.error as e: + raise e + + return SglGen( + name, + max_tokens, + min_tokens, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + dtype, + regex, + json_schema, + ) + + +def gen_int( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + int, + None, + ) + + +def gen_string( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + str, + None, + ) + + +def image(expr: SglExpr): + return SglImage(expr) + + +def video(path: str, num_frames: int): + return SglVideo(path, num_frames) + + +def select( + name: Optional[str] = None, + choices: Optional[List[str]] = None, + temperature: float = 0.0, + choices_method: ChoicesSamplingMethod = token_length_normalized, +): + assert choices is not None + return SglSelect(name, choices, temperature, choices_method) + + +def _role_common(name: str, expr: Optional[SglExpr] = None): + if expr is None: + return SglExprList([SglRoleBegin(name), SglRoleEnd(name)]) + else: + return SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) + + +def system(expr: Optional[SglExpr] = None): + return _role_common("system", expr) + + +def user(expr: Optional[SglExpr] = None): + return _role_common("user", expr) + + +def assistant(expr: Optional[SglExpr] = None): + return _role_common("assistant", expr) + + +def system_begin(): + return SglRoleBegin("system") + + +def system_end(): + return SglRoleEnd("system") + + +def user_begin(): + return SglRoleBegin("user") + + +def user_end(): + return SglRoleEnd("user") + + +def assistant_begin(): + return SglRoleBegin("assistant") + + +def assistant_end(): + return SglRoleEnd("assistant") diff --git a/lib/python3.12/site-packages/sglang/bench_offline_throughput.py b/lib/python3.12/site-packages/sglang/bench_offline_throughput.py new file mode 100644 index 0000000000000000000000000000000000000000..69bbc3e4de01066bdccf1d156c37eed59f777360 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/bench_offline_throughput.py @@ -0,0 +1,424 @@ +""" +Benchmark the throughput in the offline mode. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as bench_serving.py). + +# Usage +## Sharegpt dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + +## Random dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 +""" + +import argparse +import dataclasses +import json +import logging +import os +import random +import time +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from sglang.bench_serving import ( + get_dataset, + get_tokenizer, + sample_random_requests, + set_ulimit, +) +from sglang.lang.backend.runtime_endpoint import Runtime +from sglang.srt.entrypoints.engine import Engine +from sglang.srt.server_args import ServerArgs + + +@dataclasses.dataclass +class BenchArgs: + backend: str = "engine" + result_filename: str = "" + dataset_name: str = "sharegpt" + dataset_path: str = "" + num_prompts: int = 1000 + sharegpt_output_len: Optional[int] = None + sharegpt_context_len: Optional[int] = None + random_input_len: int = 1024 + random_output_len: int = 1024 + random_range_ratio: float = 0.0 + gsp_num_groups: int = 64 + gsp_prompts_per_group: int = 16 + gsp_system_prompt_len: int = 2048 + gsp_question_len: int = 128 + gsp_output_len: int = 256 + seed: int = 1 + disable_ignore_eos: bool = False + extra_request_body: Optional[str] = None + apply_chat_template: bool = False + profile: bool = False + skip_warmup: bool = False + do_not_exit: bool = False + prompt_suffix: str = "" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--backend", type=str, default=BenchArgs.backend) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["sharegpt", "random", "generated-shared-prefix"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--num-prompts", + type=int, + default=BenchArgs.num_prompts, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=BenchArgs.sharegpt_output_len, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=BenchArgs.sharegpt_context_len, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=BenchArgs.random_input_len, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + type=int, + default=BenchArgs.random_output_len, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=BenchArgs.random_range_ratio, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--gsp-num-groups", + type=int, + default=BenchArgs.gsp_num_groups, + help="Number of groups with shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-prompts-per-group", + type=int, + default=BenchArgs.gsp_prompts_per_group, + help="Number of prompts per group of shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-system-prompt-len", + type=int, + default=BenchArgs.gsp_system_prompt_len, + help="System prompt length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-question-len", + type=int, + default=BenchArgs.gsp_question_len, + help="Question length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-output-len", + type=int, + default=BenchArgs.gsp_output_len, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignore EOS token", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + default=BenchArgs.extra_request_body, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--skip-warmup", + action="store_true", + help="Skip the warmup batches.", + ) + parser.add_argument( + "--do-not-exit", + action="store_true", + help="Do not exit the program. This is useful for nsys profile with --duration and --delay.", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls(**{attr: getattr(args, attr) for attr in attrs}) + + +def throughput_test_once( + backend_name: str, + backend, + reqs: List[Tuple[str, int, int]], + ignore_eos: bool, + extra_request_body: Dict, + profile: bool, +): + measurement_results = { + "backend": backend_name, + "successful_requests": len(reqs), + "total_latency": -1, + "total_input_tokens": sum(r[1] for r in reqs), + "total_output_tokens": -1, + "request_throughput": -1, + "input_throughput": -1, + "output_throughput": -1, + "total_throughput": -1, + } + + prompt = [r[0] for r in reqs] + sampling_params = [ + { + "temperature": 0, + "max_new_tokens": r[2], + "ignore_eos": ignore_eos, + **extra_request_body, + } + for r in reqs + ] + + if profile: + assert ( + "SGLANG_TORCH_PROFILER_DIR" in os.environ + ), "Please set SGLANG_TORCH_PROFILER_DIR." + os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) + backend.start_profile() + + st = time.perf_counter() + gen_out = backend.generate(prompt=prompt, sampling_params=sampling_params) + latency = time.perf_counter() - st + + if profile: + backend.stop_profile() + monitor_trace_file(os.getenv("SGLANG_TORCH_PROFILER_DIR")) + + if backend_name == "runtime": + gen_out = json.loads(gen_out) + + server_info = backend.get_server_info() + + measurement_results["total_latency"] = latency + measurement_results["total_output_tokens"] = sum( + o["meta_info"]["completion_tokens"] for o in gen_out + ) + measurement_results["request_throughput"] = ( + measurement_results["successful_requests"] / latency + ) + measurement_results["input_throughput"] = ( + measurement_results["total_input_tokens"] / latency + ) + measurement_results["output_throughput"] = ( + measurement_results["total_output_tokens"] / latency + ) + measurement_results["total_throughput"] = ( + measurement_results["total_input_tokens"] + + measurement_results["total_output_tokens"] + ) / latency + measurement_results["last_gen_throughput"] = server_info["last_gen_throughput"] + + return measurement_results + + +def monitor_trace_file(directory, interval=1): + + print(f"Monitoring {directory} for new trace files...") + + known_files = set(os.listdir(directory)) + + while True: + flag = False + time.sleep(interval) + current_files = set(os.listdir(directory)) + + new_files = current_files - known_files + for new_file in new_files: + new_file_path = os.path.join(directory, new_file) + print(f"New file detected: {new_file}") + + previous_size = 0 + while True: + try: + current_size = os.path.getsize(new_file_path) + except FileNotFoundError: + print(f"File {new_file} is no longer accessible.") + break + + if current_size > previous_size: + previous_size = current_size + else: + flag = True + break + + time.sleep(interval) + if flag: + break + + +def throughput_test( + server_args: ServerArgs, + bench_args: BenchArgs, +): + if bench_args.backend == "engine": + backend = Engine(**dataclasses.asdict(server_args)) + if not backend: + raise ValueError("Please provide valid engine arguments") + elif bench_args.backend == "runtime": + backend = Runtime(**dataclasses.asdict(server_args)) + else: + raise ValueError('Please set backend to either "engine" or "runtime"') + + tokenizer_id = server_args.tokenizer_path or server_args.model_path + tokenizer = get_tokenizer(tokenizer_id) + + # Set global environmnets + set_ulimit() + random.seed(bench_args.seed) + np.random.seed(bench_args.seed) + + # Parse args + extra_request_body = {} + if bench_args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Read dataset + input_requests = get_dataset(bench_args, tokenizer) + + warmup_requests = sample_random_requests( + input_len=256, + output_len=16, + num_prompts=min(bench_args.num_prompts, 16), + range_ratio=1.0, + tokenizer=tokenizer, + dataset_path=bench_args.dataset_path, + ) + + # Warm up + if not bench_args.skip_warmup: + logging.info("\nWarmup...") + throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=warmup_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=False, + ) + time.sleep(0.5) + + logging.info("\nBenchmark...") + result = throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=input_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=bench_args.profile, + ) + backend.shutdown() + + if bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + fout.write(json.dumps(result) + "\n") + + print( + "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=") + ) + print("{:<40} {:<10}".format("Backend:", result["backend"])) + print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"])) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"])) + print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"])) + print( + "{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"]) + ) + print( + "{:<40} {:<10.2f}".format( + "Last generation throughput (tok/s):", result["last_gen_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", result["request_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", result["input_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", result["output_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", result["total_throughput"] + ) + ) + print("=" * 50) + + return result + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + throughput_test(server_args, bench_args) + + while bench_args.do_not_exit: + pass diff --git a/lib/python3.12/site-packages/sglang/bench_one_batch.py b/lib/python3.12/site-packages/sglang/bench_one_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..cdf8e9ea3ccbcef287347ea0ca639938b041ec80 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/bench_one_batch.py @@ -0,0 +1,568 @@ +""" +Benchmark the latency of running a single static batch without a server. + +This script does not launch a server and uses the low-level APIs. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +# Usage (latency test) +## with dummy weights: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy +## sweep through multiple data points and store (append) the results in a jsonl file: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --run-name test_run +## run with profiling: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --profile +# Usage (correctness test): +python -m sglang.bench_one_batch --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct + +## Reference output (of the correctness test above, can be gpu dependent): +input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]] + +prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]], + device='cuda:0') + +prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141], + [-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781], + [-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]], + device='cuda:0') + +========== Prompt 0 ========== + The capital of France is Paris. +The capital of the United States is Washington, D.C. + + +========== Prompt 1 ========== + The capital of the United Kindom is London. +The capital of the United Kingdom is London. +The capital of the + +========== Prompt 2 ========== + Today is a sunny day and I like to go for a walk in the park. +I'm going to the park +""" + +import argparse +import dataclasses +import itertools +import json +import logging +import multiprocessing +import os +import time +from typing import Tuple + +import numpy as np +import torch +import torch.distributed as dist + +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.distributed.parallel_state import destroy_distributed_environment +from sglang.srt.entrypoints.engine import _set_envs_and_config +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils import ( + configure_logger, + get_bool_env_var, + kill_process_tree, + set_gpu_proc_affinity, + suppress_other_loggers, +) + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + result_filename: str = "result.jsonl" + correctness_test: bool = False + # This is only used for correctness test + cut_len: int = 4 + log_decode_step: int = 0 + profile: bool = False + profile_filename_prefix: str = "profile" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--correctness-test", action="store_true") + parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len) + parser.add_argument( + "--log-decode-step", + type=int, + default=BenchArgs.log_decode_step, + help="Log decode latency by step, default is set to zero to disable.", + ) + parser.add_argument( + "--profile", action="store_true", help="Use Torch Profiler." + ) + parser.add_argument( + "--profile-filename-prefix", + type=str, + default=BenchArgs.profile_filename_prefix, + help="Prefix of the profiling file names. The full profiling result file(s) be " + '"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"', + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def load_model(server_args, port_args, tp_rank): + suppress_other_loggers() + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + model_config = ModelConfig( + server_args.model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + context_length=server_args.context_length, + model_override_args=server_args.json_model_override_args, + is_embedding=server_args.is_embedding, + enable_multimodal=server_args.enable_multimodal, + dtype=server_args.dtype, + quantization=server_args.quantization, + ) + model_runner = ModelRunner( + model_config=model_config, + mem_fraction_static=server_args.mem_fraction_static, + gpu_id=tp_rank, + tp_rank=tp_rank, + tp_size=server_args.tp_size, + nccl_port=port_args.nccl_port, + server_args=server_args, + ) + rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}") + tokenizer = get_tokenizer( + server_args.tokenizer_path, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + ) + if server_args.tp_size > 1: + dist.barrier() + return model_runner, tokenizer + + +def prepare_inputs_for_correctness_test(bench_args, tokenizer): + prompts = [ + "The capital of France is", + "The capital of the United Kindom is", + "Today is a sunny day and I like", + ] + input_ids = [tokenizer.encode(p) for p in prompts] + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(prompts)): + assert len(input_ids[i]) > bench_args.cut_len + + tmp_input_ids = input_ids[i][: bench_args.cut_len] + req = Req( + rid=i, + origin_input_text=prompts[i], + origin_input_ids=tmp_input_ids, + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + + return input_ids, reqs + + +def prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner +): + for i in range(len(reqs)): + req = reqs[i] + req.fill_ids += input_ids[i][bench_args.cut_len :] + req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ + i, : bench_args.cut_len + ] + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + return reqs + + +def prepare_synthetic_inputs_for_latency_test(batch_size, input_len): + input_ids = np.random.randint(0, 10000, (batch_size, input_len), dtype=np.int32) + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(input_ids)): + req = Req( + rid=i, + origin_input_text="", + origin_input_ids=list(input_ids[i]), + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + + return reqs + + +@torch.no_grad +def extend(reqs, model_runner): + batch = ScheduleBatch.init_new( + reqs=reqs, + req_to_token_pool=model_runner.req_to_token_pool, + token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator, + tree_cache=None, + model_config=model_runner.model_config, + enable_overlap=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + enable_custom_logit_processor=False, + ) + batch.prepare_for_extend() + _maybe_prepare_dp_attn_batch(batch, model_runner) + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits, batch + + +@torch.no_grad +def decode(input_token_ids, batch, model_runner): + batch.output_ids = input_token_ids + batch.prepare_for_decode() + _maybe_prepare_dp_attn_batch(batch, model_runner) + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits + + +def _maybe_prepare_dp_attn_batch(batch: ScheduleBatch, model_runner): + if model_runner.server_args.enable_dp_attention: + Scheduler.prepare_dp_attn_batch_raw( + batch, + dp_size=model_runner.server_args.dp_size, + attn_tp_size=1, + tp_cpu_group=model_runner.tp_group.cpu_group, + get_idle_batch=None, + disable_cuda_graph=model_runner.server_args.disable_cuda_graph, + spec_algorithm=SpeculativeAlgorithm.NONE, + speculative_num_draft_tokens=None, + ) + + +def correctness_test( + server_args, + port_args, + bench_args, + tp_rank, +): + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs + input_ids, reqs = prepare_inputs_for_correctness_test(bench_args, tokenizer) + rank_print(f"\n{input_ids=}\n") + + if bench_args.cut_len > 0: + # Prefill + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (first half): {next_token_logits} \n") + + # Prepare extend inputs + reqs = prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner + ) + + # Extend (prefill w/ KV cache) + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (final): {next_token_logits} \n") + + # Decode + output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))] + for _ in range(bench_args.output_len[0] - 1): + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + next_token_ids_list = next_token_ids.tolist() + for i in range(len(reqs)): + output_ids[i].append(next_token_ids_list[i]) + + # Print output texts + for i in range(len(reqs)): + rank_print(f"========== Prompt {i} ==========") + rank_print(tokenizer.decode(output_ids[i]), "\n") + + +def synchronize(device): + torch.get_device_module(device).synchronize() + + +def latency_test_run_once( + run_name, + model_runner, + rank_print, + reqs, + batch_size, + input_len, + output_len, + device, + log_decode_step, + profile, + profile_filename_prefix, +): + max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len) + if batch_size > max_batch_size: + rank_print( + f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit" + ) + return + + # Clear the pools. + model_runner.req_to_token_pool.clear() + model_runner.token_to_kv_pool_allocator.clear() + + measurement_results = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + } + + tot_latency = 0 + + profiler = None + if profile: + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + ) + profiler.start() + + # Prefill + synchronize(device) + tic = time.time() + next_token_ids, _, batch = extend(reqs, model_runner) + synchronize(device) + prefill_latency = time.time() - tic + tot_latency += prefill_latency + throughput = input_len * batch_size / prefill_latency + rank_print( + f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["prefill_latency"] = prefill_latency + measurement_results["prefill_throughput"] = throughput + + # Decode + decode_latencies = [] + for i in range(output_len - 1): + synchronize(device) + tic = time.time() + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + synchronize(device) + latency = time.time() - tic + tot_latency += latency + throughput = batch_size / latency + decode_latencies.append(latency) + if i < 5 or (log_decode_step > 0 and i % log_decode_step == 0): + rank_print( + f"Decode {i}. Batch size: {batch_size}, latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + + if profile: + profiler.stop() + profile_filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}.trace.json.gz" + parent_dir = os.path.dirname(os.path.abspath(profile_filename)) + os.makedirs(parent_dir, exist_ok=True) + profiler.export_chrome_trace(profile_filename) + rank_print(f"torch profiler chrome trace saved to {profile_filename}") + + # Record decode timing from 2nd output + if output_len > 1: + med_decode_latency = np.median(decode_latencies) + med_decode_throughput = batch_size / med_decode_latency + rank_print( + f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s" + ) + measurement_results["median_decode_latency"] = med_decode_latency + measurement_results["median_decode_throughput"] = med_decode_throughput + + throughput = (input_len + output_len) * batch_size / tot_latency + rank_print( + f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["total_latency"] = tot_latency + measurement_results["overall_throughput"] = throughput + return measurement_results + + +def latency_test( + server_args, + port_args, + bench_args, + tp_rank, +): + # Set CPU affinity + if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): + set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, tp_rank) + + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs for warm up + reqs = prepare_synthetic_inputs_for_latency_test( + bench_args.batch_size[0], bench_args.input_len[0] + ) + + # Warm up + rank_print("Warmup ...") + latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bench_args.batch_size[0], + bench_args.input_len[0], + min(32, bench_args.output_len[0]), # shorter decoding to speed up the warmup + server_args.device, + log_decode_step=0, + profile=False, + profile_filename_prefix="", # not used + ) + + rank_print("Benchmark ...") + + # Run the sweep + result_list = [] + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + reqs = prepare_synthetic_inputs_for_latency_test(bs, il) + ret = latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bs, + il, + ol, + server_args.device, + bench_args.log_decode_step, + bench_args.profile if tp_rank == 0 else None, + bench_args.profile_filename_prefix, + ) + if ret is not None: + result_list.append(ret) + + # Write results in jsonlines format on rank 0. + if tp_rank == 0 and bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + for result in result_list: + fout.write(json.dumps(result) + "\n") + + if server_args.tp_size > 1: + destroy_distributed_environment() + + +def main(server_args, bench_args): + server_args.cuda_graph_max_bs = max(bench_args.batch_size) + + _set_envs_and_config(server_args) + + if server_args.model_path: + if bench_args.correctness_test: + work_func = correctness_test + else: + work_func = latency_test + else: + raise ValueError( + "Provide --model-path for running the tests or " + "provide --result-filename for plotting the results" + ) + + port_args = PortArgs.init_new(server_args) + + if server_args.tp_size == 1: + work_func(server_args, port_args, bench_args, 0) + else: + workers = [] + for tp_rank in range(server_args.tp_size): + proc = multiprocessing.Process( + target=work_func, + args=( + server_args, + port_args, + bench_args, + tp_rank, + ), + ) + proc.start() + workers.append(proc) + + for proc in workers: + proc.join() + + proc.terminate() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + try: + main(server_args, bench_args) + finally: + if server_args.tp_size != 1: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/lib/python3.12/site-packages/sglang/bench_one_batch_server.py b/lib/python3.12/site-packages/sglang/bench_one_batch_server.py new file mode 100644 index 0000000000000000000000000000000000000000..1fc4ff58db0dd185bf6495e663aa555bc855b3be --- /dev/null +++ b/lib/python3.12/site-packages/sglang/bench_one_batch_server.py @@ -0,0 +1,192 @@ +""" +Benchmark the latency of running a single batch with a server. + +This script launches a server and uses the HTTP interface. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +Usage: +python3 -m sglang.bench_one_batch_server --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8 + +python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 +""" + +import argparse +import dataclasses +import itertools +import json +import multiprocessing +import os +import time +from typing import Tuple + +import numpy as np +import requests + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + result_filename: str = "result.jsonl" + base_url: str = "" + skip_warmup: bool = False + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--base-url", type=str, default=BenchArgs.base_url) + parser.add_argument("--skip-warmup", action="store_true") + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def launch_server_internal(server_args): + try: + launch_server(server_args) + except Exception as e: + raise e + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +def launch_server_process(server_args: ServerArgs): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = 600 + + start_time = time.time() + while time.time() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + response = requests.get(f"{base_url}/v1/models", headers=headers) + if response.status_code == 200: + return proc, base_url + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError("Server failed to start within the timeout period.") + + +def run_one_case( + url: str, + batch_size: int, + input_len: int, + output_len: int, + run_name: str, + result_filename: str, +): + input_ids = [ + [int(x) for x in np.random.randint(0, high=16384, size=(input_len,))] + for _ in range(batch_size) + ] + + tic = time.time() + response = requests.post( + url + "/generate", + json={ + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": output_len, + "ignore_eos": True, + }, + }, + ) + latency = time.time() - tic + + _ = response.json() + output_throughput = batch_size * output_len / latency + overall_throughput = batch_size * (input_len + output_len) / latency + + print(f"batch size: {batch_size}") + print(f"latency: {latency:.2f} s") + print(f"output throughput: {output_throughput:.2f} token/s") + print(f"(input + output) throughput: {overall_throughput:.2f} token/s") + + if result_filename: + with open(result_filename, "a") as fout: + res = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "latency": round(latency, 4), + "output_throughput": round(output_throughput, 2), + "overall_throughput": round(overall_throughput, 2), + } + fout.write(json.dumps(res) + "\n") + + +def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): + if bench_args.base_url: + proc, base_url = None, bench_args.base_url + else: + proc, base_url = launch_server_process(server_args) + + # warmup + if not bench_args.skip_warmup: + run_one_case( + base_url, + batch_size=16, + input_len=1024, + output_len=16, + run_name="", + result_filename="", + ) + + # benchmark + try: + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + run_one_case( + base_url, + bs, + il, + ol, + bench_args.run_name, + bench_args.result_filename, + ) + finally: + if proc: + kill_process_tree(proc.pid) + + print(f"\nResults are saved to {bench_args.result_filename}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + run_benchmark(server_args, bench_args) diff --git a/lib/python3.12/site-packages/sglang/bench_serving.py b/lib/python3.12/site-packages/sglang/bench_serving.py new file mode 100644 index 0000000000000000000000000000000000000000..9b272212681578fb39331e46c24d7eac259854b7 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/bench_serving.py @@ -0,0 +1,1623 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py + +""" +Benchmark online serving with dynamic requests. + +Usage: +python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 +""" + +import argparse +import asyncio +import json +import os +import pickle +import random +import resource +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm +from transformers import ( + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) +ASSISTANT_SUFFIX = "Assistant:" + +global args + + +# don't want to import sglang package here +def _get_bool_env_var(name: str, default: str = "false") -> bool: + value = os.getenv(name, default) + return value.lower() in ("true", "1") + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + lora_name: str + extra_request_body: Dict[str, Any] + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field(default_factory=list) # List of inter-token latencies + prompt_len: int = 0 + error: str = "" + output_len: int = 0 + + +def remove_prefix(text: str, prefix: str) -> str: + return text[len(prefix) :] if text.startswith(prefix) else text + + +def remove_suffix(text: str, suffix: str) -> str: + return text[: -len(suffix)] if text.endswith(suffix) else text + + +def get_auth_headers() -> Dict[str, str]: + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + return {"Authorization": f"Bearer {api_key}"} + else: + return {} + + +# trt llm does not support ignore_eos +# https://github.com/triton-inference-server/tensorrtllm_backend/issues/505 +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.000001, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + "min_length": request_func_input.output_len, + "end_id": 1048576, + **request_func_input.extra_request_body, + } + if args.disable_ignore_eos: + del payload["min_length"] + del payload["end_id"] + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.output_len = request_func_input.output_len + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "completions" + ), "OpenAI Completions API URL must end with 'completions'." + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["text"] + output_len = (data.get("usage") or {}).get( + "completion_tokens", output_len + ) + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_truss( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = request_func_input.output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_sglang_generate( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": request_func_input.output_len, + "ignore_eos": not args.disable_ignore_eos, + }, + "stream": not args.disable_stream, + "lora_path": request_func_input.lora_name, + "return_logprob": args.return_logprob, + "logprob_start_len": -1, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + last_output_len = 0 + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + # print(chunk_bytes) + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["text"]: + timestamp = time.perf_counter() + generated_text = data["text"] + output_len = data["meta_info"]["completion_tokens"] + + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + num_new_tokens = output_len - last_output_len + if num_new_tokens == 0: + continue + adjust_itl = ( + timestamp - most_recent_timestamp + ) / num_new_tokens + output.itl.extend([adjust_itl] * num_new_tokens) + + most_recent_timestamp = timestamp + last_output_len = output_len + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + print(f"{output.error=}") + + if pbar: + pbar.update(1) + return output + + +async def async_request_gserver( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + raise NotImplementedError() + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": + import huggingface_hub.constants + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], + ) + + return model_path + return pretrained_model_name_or_path + + +def get_tokenizer( + pretrained_model_name_or_path: str, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + if pretrained_model_name_or_path.endswith( + ".json" + ) or pretrained_model_name_or_path.endswith(".model"): + from sglang.srt.hf_transformers_utils import get_tokenizer + + return get_tokenizer(pretrained_model_name_or_path) + + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path + ): + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) + return AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=True + ) + + +def get_dataset(args, tokenizer): + if args.dataset_name == "sharegpt": + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + context_len=args.sharegpt_context_len, + prompt_suffix=args.prompt_suffix, + apply_chat_template=args.apply_chat_template, + ) + elif args.dataset_name.startswith("random"): + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + random_sample=args.dataset_name == "random", + ) + elif args.dataset_name == "generated-shared-prefix": + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gsp_num_groups, + prompts_per_group=args.gsp_prompts_per_group, + system_prompt_len=args.gsp_system_prompt_len, + question_len=args.gsp_question_len, + output_len=args.gsp_output_len, + tokenizer=tokenizer, + args=args, + ) + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_sglang_generate, + "sglang-native": async_request_sglang_generate, + "sglang-oai": async_request_openai_completions, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, + "trt": async_request_trt_llm, + "gserver": async_request_gserver, + "truss": async_request_truss, +} + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p95_itl_ms: float + p99_itl_ms: float + max_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" + + +def download_and_cache_file(url: str, filename: Optional[str] = None): + """Read and cache a file from a url.""" + if filename is None: + filename = os.path.join("/tmp", url.split("/")[-1]) + + # Check if the cache file already exists + if os.path.exists(filename): + return filename + + print(f"Downloading from {url} to {filename}") + + # Stream the response to show the progress bar + response = requests.get(url, stream=True) + response.raise_for_status() # Check for request errors + + # Total size of the file in bytes + total_size = int(response.headers.get("content-length", 0)) + chunk_size = 1024 # Download in chunks of 1KB + + # Use tqdm to display the progress bar + with open(filename, "wb") as f, tqdm( + desc=filename, + total=total_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for chunk in response.iter_content(chunk_size=chunk_size): + f.write(chunk) + bar.update(len(chunk)) + + return filename + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + context_len: Optional[int] = None, + prompt_suffix: Optional[str] = "", + apply_chat_template=False, +) -> List[Tuple[str, int, int]]: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path) and dataset_path == "": + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: List[Tuple[str, int, int]] = [] + for i in range(len(dataset)): + if len(filtered_dataset) == num_requests: + break + + # Tokenize the prompts and completions. + prompt = dataset[i][0] + if prompt_suffix: + prompt = ( + remove_suffix(prompt, ASSISTANT_SUFFIX) + + prompt_suffix + + ASSISTANT_SUFFIX + ) + + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + prompt = prompt.replace(tokenizer.bos_token, "") + + prompt_token_ids = tokenizer.encode(prompt) + completion = dataset[i][1] + completion_token_ids = tokenizer.encode(completion) + prompt_len = len(prompt_token_ids) + output_len = ( + len(completion_token_ids) if fixed_output_len is None else fixed_output_len + ) + + if prompt_len < 2 or output_len < 2: + # Prune too short sequences. + continue + + if context_len and prompt_len + output_len > context_len: + # Prune too long sequences. + continue + + filtered_dataset.append((prompt, prompt_len, output_len)) + + print(f"#Input tokens: {np.sum([x[1] for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x[2] for x in filtered_dataset])}") + return filtered_dataset + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + random_sample: bool = True, +) -> List[Tuple[str, int, int]]: + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), + input_len + 1, + size=num_prompts, + ) + output_lens = np.random.randint( + int(output_len * range_ratio), + output_len + 1, + size=num_prompts, + ) + + if random_sample: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: List[Tuple[str, int, int]] = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + prompt = tokenizer.decode(input_ids) + input_requests.append((prompt, int(input_lens[i]), int(output_lens[i]))) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + prompt = tokenizer.decode( + [ + (offsets[i] + i + j) % tokenizer.vocab_size + for j in range(input_lens[i]) + ] + ) + input_requests.append((prompt, int(input_lens[i]), int(output_lens[i]))) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_shared_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, + args: argparse.Namespace, +) -> List[Tuple[str, int, int]]: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, tokenizer) + + # Try to load from cache first + if cache_path.exists(): + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print("\nGenerating new input data...") + + # Generate system prompts for each group + system_prompts = [] + for _ in range(num_groups): + system_prompt = gen_prompt(tokenizer, system_prompt_len) + system_prompts.append(system_prompt) + + # Generate questions + questions = [] + for _ in range(num_groups * prompts_per_group): + question = gen_prompt(tokenizer, question_len) + questions.append(question) + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + question = questions[group_idx * prompts_per_group + prompt_idx] + full_prompt = f"{system_prompt}\n\n{question}" + prompt_len = len(tokenizer.encode(full_prompt)) + + input_requests.append((full_prompt, prompt_len, output_len)) + total_input_tokens += prompt_len + total_output_tokens += output_len + + # Shuffle questions + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Total prompts: {len(input_requests)}") + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests + + +async def get_request( + input_requests: List[Tuple[str, int, int]], + request_rate: float, +) -> AsyncGenerator[Tuple[str, int, int], None]: + input_requests = iter(input_requests) + for request in input_requests: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the exponential distribution. + interval = np.random.exponential(1.0 / request_rate) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[Tuple[str, int, int]], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + backend: str, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_len + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode(outputs[i].generated_text, add_special_tokens=False) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += input_requests[i][1] + if output_len > 1: + tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1)) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + + e2e_latencies.append(outputs[i].latency) + + completed += 1 + else: + output_lens.append(0) + retokenized_output_lens.append(0) + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2, + ) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(output_lens), + total_output_retokenized=sum(retokenized_output_lens), + request_throughput=completed / dur_s, + input_throughput=total_input / dur_s, + output_throughput=sum(output_lens) / dur_s, + output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, + total_throughput=(total_input + sum(output_lens)) / dur_s, + total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) + / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) + * 1000, # ttfts is empty if streaming is not supported by backend + median_ttft_ms=np.median(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, + mean_tpot_ms=np.mean(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, + mean_itl_ms=np.mean(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + p95_itl_ms=np.percentile(itls or 0, 95) * 1000, + p99_itl_ms=np.percentile(itls or 0, 99) * 1000, + max_itl_ms=np.max(itls or 0) * 1000, + mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, + median_e2e_latency_ms=np.median(e2e_latencies) * 1000, + std_e2e_latency_ms=np.std(e2e_latencies) * 1000, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + + return metrics, output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + request_rate: float, + max_concurrency: Optional[int], + disable_tqdm: bool, + lora_names: List[str], + extra_request_body: Dict[str, Any], + profile: bool, + pd_seperated: bool = False, + flush_cache: bool = False, + warmup_requests: int = 1, +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + # Limit concurrency + # From https://github.com/vllm-project/vllm/pull/9390 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, pbar=pbar) + + # Warmup + print(f"Starting warmup with {warmup_requests} sequences...") + + # Use the first request for all warmup iterations + test_prompt, test_prompt_len, test_output_len = input_requests[0] + if lora_names is not None and len(lora_names) != 0: + lora_name = lora_names[0] + else: + lora_name = None + + # Create the test input once + test_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=min(test_output_len, 32), + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + + # Run warmup requests + warmup_tasks = [] + for _ in range(warmup_requests): + warmup_tasks.append( + asyncio.create_task(request_func(request_func_input=test_input)) + ) + + warmup_outputs = await asyncio.gather(*warmup_tasks) + + # Check if at least one warmup request succeeded + if warmup_requests > 0 and not any(output.success for output in warmup_outputs): + raise ValueError( + "Warmup failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {warmup_outputs[0].error}" + ) + else: + print( + f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..." + ) + + # Flush cache + if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache: + requests.post(base_url + "/flush_cache", headers=get_auth_headers()) + + time.sleep(1.0) + + # Start profiler + if profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=base_url + "/start_profile" + ) + if profile_output.success: + print("Profiler started") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + # Run all requests + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate): + prompt, prompt_len, output_len = request + if lora_names is not None and len(lora_names) != 0: + idx = random.randint(0, len(lora_names) - 1) + lora_name = lora_names[idx] + else: + lora_name = None + + request_func_input = RequestFuncInput( + model=model_id, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=base_url + "/stop_profile") + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + if "sglang" in backend: + server_info = requests.get(base_url + "/get_server_info") + if pd_seperated: + accept_length = server_info.json()["decode"][0].get( + "avg_spec_accept_length", None + ) + else: + accept_length = server_info.json().get("avg_spec_accept_length", None) + else: + accept_length = None + + # Compute metrics and print results + benchmark_duration = time.perf_counter() - benchmark_start_time + metrics, output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print( + "{:<40} {:<10}".format( + "Max reqeuest concurrency:", + max_concurrency if max_concurrency else "not set", + ) + ) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + if accept_length: + print("{:<40} {:<10.2f}".format("Accept length:", accept_length)) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + # Arguments + "backend": args.backend, + "dataset_name": args.dataset_name, + "request_rate": request_rate, + "max_concurrency": max_concurrency, + "sharegpt_output_len": args.sharegpt_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + # Results + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "std_e2e_latency_ms": metrics.std_e2e_latency_ms, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p95_itl_ms": metrics.p95_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "concurrency": metrics.concurrency, + "accept_length": accept_length, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + if args.dataset_name.startswith("random"): + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + else: + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_sharegpt.jsonl" + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result.update( + { + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + ) + return result + + +def check_chat_template(model_path): + try: + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + return "chat_template" in tokenizer.init_kwargs + except Exception as e: + print(f"Fail to load tokenizer config with error={e}") + return False + + +def set_global_args(args_: argparse.Namespace): + """Set the global args.""" + global args + args = args_ + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set default value for max_concurrency if not present + if not hasattr(args, "max_concurrency"): + args.max_concurrency = None + + # Set default value for warmup_requests if not present + if not hasattr(args, "warmup_requests"): + args.warmup_requests = 1 + + print(f"benchmark_args={args}") + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + extra_request_body = {} + if args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + "sglang-native": 30000, + "sglang-oai": 30000, + "lmdeploy": 23333, + "vllm": 8000, + "trt": 8000, + "gserver": 9988, + "truss": 8080, + }.get(args.backend, 30000) + + model_url = ( + f"{args.base_url}/v1/models" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models" + ) + + if args.backend in ["sglang", "sglang-native"]: + api_url = ( + f"{args.base_url}/generate" + if args.base_url + else f"http://{args.host}:{args.port}/generate" + ) + elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]: + api_url = ( + f"{args.base_url}/v1/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/completions" + ) + elif args.backend == "trt": + api_url = ( + f"{args.base_url}/v2/models/ensemble/generate_stream" + if args.base_url + else f"http://{args.host}:{args.port}/v2/models/ensemble/generate_stream" + ) + if args.model is None: + print("Please provide a model using `--model` when using `trt` backend.") + sys.exit(1) + elif args.backend == "gserver": + api_url = args.base_url if args.base_url else f"{args.host}:{args.port}" + args.model = args.model or "default" + elif args.backend == "truss": + api_url = ( + f"{args.base_url}/v1/models/model:predict" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models/model:predict" + ) + base_url = ( + f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url + ) + + # Get model name + if args.model is None: + if args.backend == "truss": + print( + "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" + ) + sys.exit(1) + try: + response = requests.get(model_url, headers=get_auth_headers()) + model_list = response.json().get("data", []) + args.model = model_list[0]["id"] if model_list else None + except Exception as e: + print(f"Failed to fetch model from {model_url}. Error: {e}") + print( + "Please specify the correct host and port using `--host` and `--port`." + ) + sys.exit(1) + + if args.model is None: + print("No model specified or found. Please provide a model using `--model`.") + sys.exit(1) + + if not check_chat_template(args.model): + print( + "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n" + "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n" + ) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer = get_tokenizer(tokenizer_id) + input_requests = get_dataset(args, tokenizer) + + # compatible with SimpleNamespace + if not hasattr(args, "flush_cache"): + args.flush_cache = False + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + max_concurrency=args.max_concurrency, + disable_tqdm=args.disable_tqdm, + lora_names=args.lora_name, + extra_request_body=extra_request_body, + profile=args.profile, + pd_seperated=args.pd_seperated, + flush_cache=args.flush_cache, + ) + ) + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}") + + +class LoRAPathAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, []) + for lora_name in values: + getattr(namespace, self.dest).append(lora_name) + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["sharegpt", "random", "random-ids", "generated-shared-prefix"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--model", + type=str, + help="Name or path of the model. If not set, the default model will request /v1/models for conf.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help="Name or path of the tokenizer. If not set, using the model conf.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=None, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + default=1024, + type=int, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--disable-stream", + action="store_true", + help="Disable streaming mode.", + ) + parser.add_argument( + "--return-logprob", + action="store_true", + help="Return logprob.", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--lora-name", + type=str, + nargs="*", + default=None, + action=LoRAPathAction, + help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + parser.add_argument( + "--pd-seperated", + action="store_true", + help="Benchmark PD disaggregation server", + ) + parser.add_argument( + "--flush-cache", + action="store_true", + help="Flush the cache before running the benchmark", + ) + parser.add_argument( + "--warmup-requests", + type=int, + default=1, + help="Number of warmup requests to run before the benchmark", + ) + + group = parser.add_argument_group("generated-shared-prefix dataset arguments") + group.add_argument( + "--gsp-num-groups", + type=int, + default=64, + help="Number of system prompt groups for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-prompts-per-group", + type=int, + default=16, + help="Number of prompts per system prompt group for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-system-prompt-len", + type=int, + default=2048, + help="Target length in tokens for system prompts in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-question-len", + type=int, + default=128, + help="Target length in tokens for questions in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-output-len", + type=int, + default=256, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + args = parser.parse_args() + run_benchmark(args) diff --git a/lib/python3.12/site-packages/sglang/check_env.py b/lib/python3.12/site-packages/sglang/check_env.py new file mode 100644 index 0000000000000000000000000000000000000000..d525da9292f73e951136284322864826671b965b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/check_env.py @@ -0,0 +1,305 @@ +"""Check environment configurations and dependency versions.""" + +import importlib.metadata +import os +import resource +import subprocess +import sys +from collections import OrderedDict, defaultdict + +import torch + +from sglang.srt.utils import is_hip + + +def is_cuda_v2(): + return torch.version.cuda is not None + + +# List of packages to check versions +PACKAGE_LIST = [ + "sglang", + "sgl_kernel", + "flashinfer", + "triton", + "transformers", + "torchao", + "numpy", + "aiohttp", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "modelscope", + "orjson", + "outlines", + "packaging", + "psutil", + "pydantic", + "multipart", + "zmq", + "torchao", + "uvicorn", + "uvloop", + "vllm", + "xgrammar", + "openai", + "tiktoken", + "anthropic", + "litellm", + "decord", +] + + +def get_package_versions(packages): + """ + Get versions of specified packages. + """ + versions = {} + for package in packages: + package_name = package.split("==")[0].split(">=")[0].split("<=")[0] + try: + version = importlib.metadata.version(package_name) + versions[package_name] = version + except ModuleNotFoundError: + versions[package_name] = "Module Not Found" + return versions + + +def get_cuda_info(): + """ + Get CUDA-related information if available. + """ + if is_cuda_v2(): + cuda_info = {"CUDA available": torch.cuda.is_available()} + + if cuda_info["CUDA available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + elif is_hip(): + cuda_info = {"ROCM available": torch.cuda.is_available()} + + if cuda_info["ROCM available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + + +def _get_gpu_info(): + """ + Get information about available GPUs. + """ + devices = defaultdict(list) + capabilities = defaultdict(list) + for k in range(torch.cuda.device_count()): + devices[torch.cuda.get_device_name(k)].append(str(k)) + capability = torch.cuda.get_device_capability(k) + capabilities[f"{capability[0]}.{capability[1]}"].append(str(k)) + + gpu_info = {} + for name, device_ids in devices.items(): + gpu_info[f"GPU {','.join(device_ids)}"] = name + + if len(capabilities) == 1: + # All GPUs have the same compute capability + cap, gpu_ids = list(capabilities.items())[0] + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + else: + # GPUs have different compute capabilities + for cap, gpu_ids in capabilities.items(): + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + + return gpu_info + + +def _get_cuda_version_info(): + """ + Get CUDA version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + cuda_info = {"CUDA_HOME": CUDA_HOME} + + if CUDA_HOME and os.path.isdir(CUDA_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME + + cuda_info = {"ROCM_HOME": ROCM_HOME} + + if ROCM_HOME and os.path.isdir(ROCM_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + else: + cuda_info = {"CUDA_HOME": ""} + return cuda_info + + +def _get_nvcc_info(): + """ + Get NVCC version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + try: + nvcc = os.path.join(CUDA_HOME, "bin/nvcc") + nvcc_output = ( + subprocess.check_output(f'"{nvcc}" -V', shell=True) + .decode("utf-8") + .strip() + ) + return { + "NVCC": nvcc_output[ + nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind( + "Build" + ) + ].strip() + } + except subprocess.SubprocessError: + return {"NVCC": "Not Available"} + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME + + try: + hipcc = os.path.join(ROCM_HOME, "bin/hipcc") + hipcc_output = ( + subprocess.check_output(f'"{hipcc}" --version', shell=True) + .decode("utf-8") + .strip() + ) + return { + "HIPCC": hipcc_output[ + hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang") + ].strip() + } + except subprocess.SubprocessError: + return {"HIPCC": "Not Available"} + else: + return {"NVCC": "Not Available"} + + +def _get_cuda_driver_version(): + """ + Get CUDA driver version. + """ + versions = set() + if is_cuda_v2(): + try: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ] + ) + versions = set(output.decode().strip().split("\n")) + if len(versions) == 1: + return {"CUDA Driver Version": versions.pop()} + else: + return {"CUDA Driver Versions": ", ".join(sorted(versions))} + except subprocess.SubprocessError: + return {"CUDA Driver Version": "Not Available"} + elif is_hip(): + try: + output = subprocess.check_output( + [ + "rocm-smi", + "--showdriverversion", + "--csv", + ] + ) + versions = set(output.decode().strip().split("\n")) + versions.discard("name, value") + ver = versions.pop() + ver = ver.replace('"Driver version", ', "").replace('"', "") + + return {"ROCM Driver Version": ver} + except subprocess.SubprocessError: + return {"ROCM Driver Version": "Not Available"} + else: + return {"CUDA Driver Version": "Not Available"} + + +def get_gpu_topology(): + """ + Get GPU topology information. + """ + if is_cuda_v2(): + try: + result = subprocess.run( + ["nvidia-smi", "topo", "-m"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + elif is_hip(): + try: + result = subprocess.run( + ["rocm-smi", "--showtopotype"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + else: + return None + + +def get_hypervisor_vendor(): + try: + output = subprocess.check_output(["lscpu"], text=True) + for line in output.split("\n"): + if "Hypervisor vendor:" in line: + return line.split(":")[1].strip() + return None + except: + return None + + +def check_env(): + """ + Check and print environment information. + """ + env_info = OrderedDict() + env_info["Python"] = sys.version.replace("\n", "") + env_info.update(get_cuda_info()) + env_info["PyTorch"] = torch.__version__ + env_info.update(get_package_versions(PACKAGE_LIST)) + + gpu_topo = get_gpu_topology() + if gpu_topo: + if is_cuda_v2(): + env_info["NVIDIA Topology"] = gpu_topo + elif is_hip(): + env_info["AMD Topology"] = gpu_topo + + hypervisor_vendor = get_hypervisor_vendor() + if hypervisor_vendor: + env_info["Hypervisor vendor"] = hypervisor_vendor + + ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + env_info["ulimit soft"] = ulimit_soft + + for k, v in env_info.items(): + print(f"{k}: {v}") + + +if __name__ == "__main__": + check_env() diff --git a/lib/python3.12/site-packages/sglang/compile_deep_gemm.py b/lib/python3.12/site-packages/sglang/compile_deep_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..e787140236dbec522ecf6b08acf36234e31c1cf2 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/compile_deep_gemm.py @@ -0,0 +1,177 @@ +""" +Compile DeepGEMM Kernels for a model with specify server arguments + +This script launches a server for capturing DeepGEMM calls and then compiles the kernels. +It accepts server arguments (the same as launch_server.py). + +Usage: +python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code + +""" + +import argparse +import dataclasses +import multiprocessing +import os +import time + +import requests + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree +from sglang.srt.warmup import warmup + +multiprocessing.set_start_method("spawn", force=True) + +# Reduce warning +os.environ["SGL_IN_DEEPGEMM_PRECOMPILE_STAGE"] = "1" +# Force enable deep gemm +os.environ["SGL_ENABLE_JIT_DEEPGEMM"] = "1" +# Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case +os.environ["SGL_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "0" + + +@dataclasses.dataclass +class CompileArgs: + timeout: int = 3600 + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--timeout", type=int, default=CompileArgs.timeout) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +@warmup("compile-deep-gemm") +async def warm_up_compile(tokenizer_manager: TokenizerManager): + print("\nGenerate warm up request for compiling DeepGEMM...\n") + generate_req_input = GenerateReqInput( + input_ids=[0, 1, 2, 3], + sampling_params={ + "temperature": 0.0, + "max_new_tokens": 8, + "ignore_eos": True, + }, + ) + await tokenizer_manager.generate_request(generate_req_input, None).__anext__() + + +def launch_server_internal(server_args): + try: + launch_server(server_args) + except Exception as e: + raise e + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +def launch_server_process_and_send_one_request( + server_args: ServerArgs, compile_args: CompileArgs +): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = compile_args.timeout + + start_time = time.time() + while time.time() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + if server_args.node_rank == 0: + response = requests.get(f"{base_url}/v1/models", headers=headers) + else: + # This http api is created by launch_dummy_health_check_server for none-rank0 node. + response = requests.get(f"{base_url}/health", headers=headers) + if response.status_code == 200: + # Rank-0 node send a request to sync with other node and then return. + if server_args.node_rank == 0: + response = requests.post( + f"{base_url}/generate", + json={ + "input_ids": [0, 1, 2, 3], + "sampling_params": { + "max_new_tokens": 8, + "temperature": 0, + }, + }, + timeout=600, + ) + if response.status_code != 200: + error = response.json() + raise RuntimeError(f"Sync request failed: {error}") + # Other nodes should wait for the exit signal from Rank-0 node. + else: + start_time_waiting = time.time() + while proc.is_alive(): + if time.time() - start_time_waiting < timeout: + time.sleep(10) + else: + raise TimeoutError("Waiting for main node timeout!") + return proc + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError( + "DeepGEMM Kernels compilation timeout." + "\n\nFeel free and please restart the command." + ) + + +def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs): + # Disbale cuda graph and torch compile to save time + server_args.disable_cuda_graph = True + server_args.enable_torch_compile = False + print(f"Disable CUDA Graph and Torch Compile to save time...") + + # Set watchdog timeout to compile_args.timeout because compilation will take a long time + server_args.watchdog_timeout = compile_args.timeout + server_args.warmups = "compile-deep-gemm" + + +def run_compile(server_args: ServerArgs, compile_args: CompileArgs): + print( + "Begin DeepGEMM Kernels compilation...\n" + "It may take a long time and timeout maybe raised " + "while the compilation is still in progress.\n" + "Just feel free to restart the command " + "until the compilation is fully finished.\n" + ) + + proc = launch_server_process_and_send_one_request(server_args, compile_args) + + print("\nDeepGEMM Kernels compilation finished successfully.") + + # Sleep for safety + time.sleep(10) + if proc.is_alive(): + # This is the rank0 node. + kill_process_tree(proc.pid) + else: + try: + kill_process_tree(proc.pid) + except Exception: + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + CompileArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + compile_args = CompileArgs.from_cli_args(args) + + refine_server_args(server_args, compile_args) + + run_compile(server_args, compile_args) diff --git a/lib/python3.12/site-packages/sglang/global_config.py b/lib/python3.12/site-packages/sglang/global_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c331e64db61df851beeaf151518dbb462b84b5c6 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/global_config.py @@ -0,0 +1,49 @@ +"""Global configurations""" + +import os + + +class GlobalConfig: + """ + Store some global constants. + + See also python/sglang/srt/managers/schedule_batch.py::global_server_args_dict, which stores + many global runtime arguments as well. + """ + + def __init__(self): + # Verbosity level + # 0: do not output anything + # 2: output final text after every run + self.verbosity = 0 + + # Default backend of the language + self.default_backend = None + + # Runtime constants: New generation token ratio estimation + self.default_init_new_token_ratio = float( + os.environ.get("SGLANG_INIT_NEW_TOKEN_RATIO", 0.7) + ) + self.default_min_new_token_ratio_factor = float( + os.environ.get("SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR", 0.14) + ) + self.default_new_token_ratio_decay_steps = float( + os.environ.get("SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS", 600) + ) + + # Runtime constants: others + self.retract_decode_steps = 20 + self.flashinfer_workspace_size = os.environ.get( + "FLASHINFER_WORKSPACE_SIZE", 384 * 1024 * 1024 + ) + + # Output tokenization configs + self.skip_special_tokens_in_output = True + self.spaces_between_special_tokens_in_out = True + + # Language frontend interpreter optimization configs + self.enable_precache_with_tracing = True + self.enable_parallel_encoding = True + + +global_config = GlobalConfig() diff --git a/lib/python3.12/site-packages/sglang/launch_server.py b/lib/python3.12/site-packages/sglang/launch_server.py new file mode 100644 index 0000000000000000000000000000000000000000..caae7b0f6cc7615af8af1820ef52980ed0074c06 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/launch_server.py @@ -0,0 +1,16 @@ +"""Launch the inference server.""" + +import os +import sys + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import prepare_server_args +from sglang.srt.utils import kill_process_tree + +if __name__ == "__main__": + server_args = prepare_server_args(sys.argv[1:]) + + try: + launch_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/lib/python3.12/site-packages/sglang/llama3_eval.py b/lib/python3.12/site-packages/sglang/llama3_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..35bd4a7e4d4cef40f30f86ce5a2a7e1240e46217 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/llama3_eval.py @@ -0,0 +1,316 @@ +# Adapt from https://github.com/fw-ai/llm_eval_meta + +import argparse +import asyncio +import os +import pickle +import re +import shutil +from collections import defaultdict +from dataclasses import dataclass + +import httpx +import numpy as np +import openai +import transformers +from datasets import load_dataset +from openai import AsyncOpenAI +from tqdm import tqdm + +# Mapping providers to their clients and models +provider_to_models = { + "b10": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, + "oai": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, + "sgl": { + "8b": "meta-llama/Llama-3.1-8B-Instruct", + "70b": "meta-llama/Llama-3.1-70B-Instruct", + "405b": "meta-llama/Llama-3.1-405B-Instruct", + }, +} + + +async def fetch_responses( + client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens +): + output_file = os.path.join(output_dir, f"response_{index}.pkl") + if os.path.exists(output_file): + print(f"File {output_file} already exists, skipping.") + return + + async with semaphore: + response = await client.completions.create( + model=provider_to_models[provider][model_size], + prompt=prompt, + temperature=0.0, + max_tokens=max_tokens, + ) + if isinstance(response, openai.BadRequestError): + with open(output_file, "wb") as f: + pickle.dump("bad_response", f) + assert isinstance(response, openai.types.completion.Completion) + # Save response to a file + with open(output_file, "wb") as f: + pickle.dump(response, f) + + +TASK_TO_MAX_TOKENS = { + "evals__mmlu__details": 1, + "evals__mmlu__0_shot__cot__details": 1024, + # Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing + "evals__mmlu_pro__details": 2048, + "evals__gsm8k__details": 1024, +} + +TASK_TO_EVAL_SET = { + "mmlu": "evals__mmlu__details", + "mmlu_cot": "evals__mmlu__0_shot__cot__details", + "mmlu_pro": "evals__mmlu_pro__details", + "gsm8k": "evals__gsm8k__details", +} + + +class CustomAsyncHTTPXClient(httpx.AsyncClient): + async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response: + request.url = httpx.URL( + f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict" + ) + return await super().send(request, *args, **kwargs) + + +def get_client(provider): + if provider not in "b10": + if os.getenv("OPENAI_API_KEY") == None: + os.environ["OPENAI_API_KEY"] = "EMPTY" + return { + "oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"), + "b10": AsyncOpenAI( + api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}", + base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict", + http_client=CustomAsyncHTTPXClient(), + ), + "sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"), + }[provider] + + +# Define the benchmark function +async def benchmark(args): + ds = load_dataset( + "meta-llama/Llama-3.1-405B-Instruct-evals", + f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}", + ) + semaphore = asyncio.Semaphore(args.concurrency) # Limit to 16 concurrent tasks + + if args.num_examples is None: + args.num_examples = len(ds["latest"]["input_final_prompts"]) + prompts = ds["latest"]["input_final_prompts"][: args.num_examples] + + # Create the output directory if it does not exist + os.makedirs(args.output_dir, exist_ok=True) + + tasks = [] + # Create the tasks with tqdm progress bar + max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]] + client = get_client(args.provider) + for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")): + tasks.append( + asyncio.create_task( + fetch_responses( + client, + f"<|begin_of_text|>{prompt[0]}", + semaphore, + idx, + args.provider, + args.model_size, + args.output_dir, + max_tokens=max_tokens, + ) + ) + ) + + # Run the tasks with tqdm progress bar + for future in tqdm( + asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks" + ): + await future + + +def get_mmlu_answer(response): + if response is not None: + return response.choices[0].text.lstrip().rstrip().upper().replace(".", "") + return None + + +def get_mmlu_cot_answer(response): + pattern = r"The best answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "").replace("*", "") + + pattern = r"the best answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + pattern = r"The correct answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + pattern = r"the correct answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + return match.group(1).replace(".", "") + + +def get_answer_gsm8k(response): + pattern = r"The final answer is (.+)\.?" + match = re.search(pattern, response.choices[0].text) + if match: + s = match.group(1) + for ok_symbol in ["%", "$"]: + s = s.replace(ok_symbol, "") + return s + + +TASK_TO_ANSWER_EXTRACTOR = { + "evals__mmlu__details": get_mmlu_answer, + "evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer, + "evals__gsm8k__details": get_answer_gsm8k, + "evals__mmlu_pro__details": get_mmlu_cot_answer, +} + + +def get_dataset_from_task(task, response_path, model_size): + ds_405b = load_dataset( + f"meta-llama/Llama-3.1-405B-Instruct-evals", + f"Llama-3.1-405B-Instruct-{task}", + ) + ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]] + + if "70b" in model_size or "8b" in model_size: + if "70" in model_size: + ref_model_ds = load_dataset( + f"meta-llama/Llama-3.1-70B-Instruct-evals", + f"Llama-3.1-70B-Instruct-{task}", + ) + else: + ref_model_ds = load_dataset( + f"meta-llama/Llama-3.1-8B-Instruct-evals", + f"Llama-3.1-8B-Instruct-{task}", + ) + + hash_to_row = {} + for row in ref_model_ds["latest"]: + hash_to_row[row["input_final_prompts_hash"][0]] = row + reordered_rows = [] + for prompt_hash in ds_405b_hash_order: + reordered_rows.append(hash_to_row[prompt_hash]) + ref_model_ds["latest"] = reordered_rows + return ref_model_ds + + return ds_405b + + +def analyze(task, response_path, model_size): + ds = get_dataset_from_task(task, response_path, model_size) + + responses = [] + total = len(ds["latest"]) + + for i in range(0, total): + response = pickle.load( + open(os.path.join(response_path, f"response_{i}.pkl"), "rb") + ) + responses.append(response) + + @dataclass + class Stats: + correct: int = 0 + total: int = 0 + meta_correct: int = 0 + + average: float = None + + subtask_name_to_stats = defaultdict(lambda: Stats()) + + for response, ds_row in zip(responses, ds["latest"]): + model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response) + + subtask = ds_row["subtask_name"] + + is_eval_correct = model_answer in ds_row["input_correct_responses"] + if is_eval_correct: + subtask_name_to_stats[subtask].correct += 1 + + if ds_row["is_correct"]: + subtask_name_to_stats[subtask].meta_correct += 1 + + subtask_name_to_stats[subtask].total += 1 + + micro_stats = Stats() + for subtask, stats in subtask_name_to_stats.items(): + stats.average = stats.correct / stats.total + stats.meta_average = stats.meta_correct / stats.total + + micro_stats.correct += stats.correct + micro_stats.total += stats.total + micro_stats.meta_correct += stats.meta_correct + + micro_stats.average = micro_stats.correct / micro_stats.total + micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total + + print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()])) + print( + "Meta Macro average", + np.mean([x.meta_average for x in subtask_name_to_stats.values()]), + ) + print("Micro average", micro_stats.average) + print("Meta Micro average", micro_stats.meta_average) + + +# Entry point for the script +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Script to run model with specified parameters." + ) + parser.add_argument( + "--model-size", + type=str, + default="8b", + help="Size of the model (e.g., 8b or 70b)", + ) + parser.add_argument( + "--provider", + type=str, + default="sgl", + help="Provider name (e.g., sgl, oai, b10)", + ) + parser.add_argument( + "--task", + type=str, + required=True, + help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)", + ) + parser.add_argument( + "--num-examples", type=int, default=None, help="Number of examples to process" + ) + parser.add_argument("--concurrency", type=int, default=16) + parser.add_argument( + "--output-dir", + type=str, + default="tmp-output-dir", + help="Directory to save responses", + ) + + args = parser.parse_args() + asyncio.run(benchmark(args)) + analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size) + shutil.rmtree("tmp-output-dir", ignore_errors=True) diff --git a/lib/python3.12/site-packages/sglang/srt/_custom_ops.py b/lib/python3.12/site-packages/sglang/srt/_custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..07c087bf6c423fe36b9af47f6c791e176734483d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/_custom_ops.py @@ -0,0 +1,115 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py +import logging +from typing import List, Tuple + +import torch + +from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu + +logger = logging.getLogger(__name__) +use_vllm_custom_allreduce = get_bool_env_var( + "USE_VLLM_CUSTOM_ALLREDUCE", default="false" +) + +if not is_hpu(): + # ROCm does not use vllm custom allreduce + if use_vllm_custom_allreduce and not is_hip(): + try: + import vllm._C + except ImportError as e: + logger.warning("Failed to import from vllm._C with %r", e) + else: + try: + import sgl_kernel + except ImportError as e: + logger.warning("Failed to import from custom_ar with %r", e) + + +if not is_hip(): + if use_vllm_custom_allreduce: + custom_op = torch.ops._C_custom_ar + else: + custom_op = sgl_kernel.allreduce + + # custom allreduce + def init_custom_ar( + ipc_tensors: List[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + full_nvlink: bool, + ) -> int: + return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink) + + def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, + ) -> None: + custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + def dispose(fa: int) -> None: + custom_op.dispose(fa) + + def meta_size() -> int: + return custom_op.meta_size() + + def register_buffer(fa: int, ipc_tensors: List[int]) -> None: + return custom_op.register_buffer(fa, ipc_tensors) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]: + return custom_op.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[List[int]], offsets: List[List[int]] + ) -> None: + custom_op.register_graph_buffers(fa, handles, offsets) + +else: + # ROCM custom allreduce + + def init_custom_ar( + meta: torch.Tensor, + rank_data: torch.Tensor, + handles: List[str], + offsets: List[int], + rank: int, + full_nvlink: bool, + ) -> int: + return sgl_kernel.allreduce.init_custom_ar( + meta, rank_data, handles, offsets, rank, full_nvlink + ) + + def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + sgl_kernel.allreduce.all_reduce_reg(fa, inp, out) + + def all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out) + + def dispose(fa: int) -> None: + sgl_kernel.allreduce.dispose(fa) + + def meta_size() -> int: + return sgl_kernel.allreduce.meta_size() + + def register_buffer( + fa: int, t: torch.Tensor, handles: List[str], offsets: List[int] + ) -> None: + return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]: + return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[str], offsets: List[List[int]] + ) -> None: + sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets) + + def allocate_meta_buffer(size: int) -> torch.Tensor: + return sgl_kernel.allreduce.allocate_meta_buffer(size) + + def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor: + return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp) diff --git a/lib/python3.12/site-packages/sglang/srt/aio_rwlock.py b/lib/python3.12/site-packages/sglang/srt/aio_rwlock.py new file mode 100644 index 0000000000000000000000000000000000000000..deda1fe79030947d46af50c2d8d2c777415ee237 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/aio_rwlock.py @@ -0,0 +1,100 @@ +import asyncio + + +class RWLock: + def __init__(self): + # Protects internal state + self._lock = asyncio.Lock() + + # Condition variable used to wait for state changes + self._cond = asyncio.Condition(self._lock) + + # Number of readers currently holding the lock + self._readers = 0 + + # Whether a writer is currently holding the lock + self._writer_active = False + + # How many writers are queued waiting for a turn + self._waiting_writers = 0 + + @property + def reader_lock(self): + """ + A context manager for acquiring a shared (reader) lock. + + Example: + async with rwlock.reader_lock: + # read-only access + """ + return _ReaderLock(self) + + @property + def writer_lock(self): + """ + A context manager for acquiring an exclusive (writer) lock. + + Example: + async with rwlock.writer_lock: + # exclusive access + """ + return _WriterLock(self) + + async def acquire_reader(self): + async with self._lock: + # Wait until there is no active writer or waiting writer + # to ensure fairness. + while self._writer_active or self._waiting_writers > 0: + await self._cond.wait() + self._readers += 1 + + async def release_reader(self): + async with self._lock: + self._readers -= 1 + # If this was the last reader, wake up anyone waiting + # (potentially a writer or new readers). + if self._readers == 0: + self._cond.notify_all() + + async def acquire_writer(self): + async with self._lock: + # Increment the count of writers waiting + self._waiting_writers += 1 + try: + # Wait while either a writer is active or readers are present + while self._writer_active or self._readers > 0: + await self._cond.wait() + self._writer_active = True + finally: + # Decrement waiting writers only after we've acquired the writer lock + self._waiting_writers -= 1 + + async def release_writer(self): + async with self._lock: + self._writer_active = False + # Wake up anyone waiting (readers or writers) + self._cond.notify_all() + + +class _ReaderLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_reader() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_reader() + + +class _WriterLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_writer() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_writer() diff --git a/lib/python3.12/site-packages/sglang/srt/code_completion_parser.py b/lib/python3.12/site-packages/sglang/srt/code_completion_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..81cdb4a3618d8c08f4056bf8edadafbb961d919d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/code_completion_parser.py @@ -0,0 +1,174 @@ +# Copyright 2023-2024 SGLang Team +# 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. +# ============================================================================== +"""Completion templates.""" + + +import dataclasses +import json +import logging +import os +from enum import auto + +from sglang.srt.openai_api.protocol import ChatCompletionRequest + +logger = logging.getLogger(__name__) +completion_template_name = None + + +class FimPosition: + """Postion of fim middle token.""" + + MIDDLE = auto() + END = auto() + + +@dataclasses.dataclass +class CompletionTemplate: + """A class that manages completion prompt templates. only for code completion currently.""" + + # The name of this template + name: str + + # the fim begin token + fim_begin_token: str + + # The fim middle token + fim_middle_token: str + + # The fim end token + fim_end_token: str + + # The position of the fim middle token + fim_position: FimPosition + + +# A global registry for all completion templates +completion_templates: dict[str, CompletionTemplate] = {} + + +def load_completion_template_for_openai_api(completion_template_arg): + global completion_template_name + + logger.info( + f"Use completion template for the OpenAI-compatible API server: {completion_template_arg}" + ) + + if not completion_template_exists(completion_template_arg): + if not os.path.exists(completion_template_arg): + raise RuntimeError( + f"Completion template {completion_template_arg} is not a built-in template name " + "or a valid completion template file path." + ) + + assert completion_template_arg.endswith( + ".json" + ), "unrecognized format of completion template file" + with open(completion_template_arg, "r") as filep: + template = json.load(filep) + try: + fim_position = FimPosition[template["fim_position"]] + except KeyError: + raise ValueError( + f"Unknown fim position: {template['fim_position']}" + ) from None + register_completion_template( + CompletionTemplate( + name=template["name"], + fim_begin_token=template["fim_begin_token"], + fim_middle_token=template["fim_middle_token"], + fim_end_token=template["fim_end_token"], + fim_position=fim_position, + ), + override=True, + ) + completion_template_name = template["name"] + else: + completion_template_name = completion_template_arg + + +def register_completion_template(template: CompletionTemplate, override: bool = False): + """Register a new completion template.""" + if not override: + assert ( + template.name not in completion_templates + ), f"{template.name} has been registered." + + completion_templates[template.name] = template + + +def completion_template_exists(template_name: str) -> bool: + return template_name in completion_templates + + +def is_completion_template_defined() -> bool: + global completion_template_name + return completion_template_name is not None + + +def generate_completion_prompt_from_request(request: ChatCompletionRequest) -> str: + global completion_template_name + if request.suffix == "": + return request.prompt + + return generate_completion_prompt( + request.prompt, request.suffix, completion_template_name + ) + + +def generate_completion_prompt(prompt: str, suffix: str, template_name: str) -> str: + + completion_template = completion_templates[template_name] + fim_begin_token = completion_template.fim_begin_token + fim_middle_token = completion_template.fim_middle_token + fim_end_token = completion_template.fim_end_token + fim_position = completion_template.fim_position + + if fim_position == FimPosition.MIDDLE: + prompt = f"{fim_begin_token}{prompt}{fim_middle_token}{suffix}{fim_end_token}" + elif fim_position == FimPosition.END: + prompt = f"{fim_begin_token}{prompt}{fim_end_token}{suffix}{fim_middle_token}" + + return prompt + + +register_completion_template( + CompletionTemplate( + name="deepseek_coder", + fim_begin_token="<｜fim▁begin｜>", + fim_middle_token="<｜fim▁hole｜>", + fim_end_token="<｜fim▁end｜>", + fim_position=FimPosition.MIDDLE, + ) +) + + +register_completion_template( + CompletionTemplate( + name="star_coder", + fim_begin_token="", + fim_middle_token="", + fim_end_token="", + fim_position=FimPosition.END, + ) +) + +register_completion_template( + CompletionTemplate( + name="qwen_coder", + fim_begin_token="<|fim_prefix|>", + fim_middle_token="<|fim_middle|>", + fim_end_token="<|fim_suffix|>", + fim_position=FimPosition.END, + ) +) diff --git a/lib/python3.12/site-packages/sglang/srt/conversation.py b/lib/python3.12/site-packages/sglang/srt/conversation.py new file mode 100644 index 0000000000000000000000000000000000000000..3fd45b467240794c6f57bda6bed7b6426b586346 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/conversation.py @@ -0,0 +1,794 @@ +# Copyright 2023-2024 SGLang Team +# 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. +# ============================================================================== +"""Conversation chat templates.""" + +# Adapted from +# https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py +import dataclasses +from enum import IntEnum, auto +from typing import Dict, List, Optional, Tuple, Union + +from sglang.srt.openai_api.protocol import ChatCompletionRequest + + +class SeparatorStyle(IntEnum): + """Separator styles.""" + + ADD_COLON_SINGLE = auto() + ADD_COLON_TWO = auto() + ADD_COLON_SPACE_SINGLE = auto() + NO_COLON_SINGLE = auto() + NO_COLON_TWO = auto() + ADD_NEW_LINE_SINGLE = auto() + LLAMA2 = auto() + LLAMA3 = auto() + LLAMA4 = auto() + CHATGLM = auto() + CHATML = auto() + CHATINTERN = auto() + DOLLY = auto() + RWKV = auto() + PHOENIX = auto() + ROBIN = auto() + FALCON_CHAT = auto() + CHATGLM3 = auto() + DEEPSEEK_CHAT = auto() + METAMATH = auto() + DeepSeekVL2 = auto() + QWEN2_VL_EMBED = auto() + GEMMA3 = auto() + + +@dataclasses.dataclass +class Conversation: + """A class that manages prompt templates and keeps all conversation history.""" + + # The name of this template + name: str + # The template of the system prompt + system_template: str = "{system_message}" + # The system message + system_message: str = "" + # The names of two roles + roles: Tuple[str] = ("USER", "ASSISTANT") + # All messages. Each item is (role, message). + messages: List[List[str]] = () + # The number of few shot examples + offset: int = 0 + # The separator style and configurations + sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE + sep: str = "\n" + sep2: str = None + # Stop criteria (the default one is EOS token) + stop_str: Union[str, List[str]] = None + # The string that represents an image token in the prompt + image_token: str = "" + audio_token: str = "