Class Session

tensorrt_llm/runtime/session.py:83–324

Session is a managed TensorRT runtime.

class Session(object):
    ''' Session is a managed TensorRT runtime.  '''

    def __init__(self, **kwargs):
        # use Session.from_serialized_engine to create a session
        pass

    def _init(self, engine_buffer=None):
        '''
        @brief: Setup TensorRT engines and context from a serialized engine file
        @param engine_buffer: a buffer holds the serialized TRT engine
        '''
        self._runtime = trt.Runtime(logger.trt_logger)
        if engine_buffer is not None:
            self._engine = self.runtime.deserialize_cuda_engine(engine_buffer)

        self._context = None
        if not self.engine.streamable_weights_size:
            self.__prepare_execution_contexts()
        return self

    def __prepare_execution_contexts(self):
        self._context = self.engine.create_execution_context()
        assert self._context is not None, "Failed to create an execution context!"
        with _scoped_stream() as stream:
            self._context.set_optimization_profile_async(0, stream)

    @staticmethod
    def from_serialized_engine(engine) -> Session:
        '''
        @brief: Create a session from a serialized engine
        @param engine: a serialized engine
        @return: a Session object
        '''
        session = Session()
        return session._init(engine)

    @staticmethod
    def from_engine(engine) -> Session:
        '''
        @brief: Create a session from an existing ICudaEngine engine
        @param engine: an ICudaEngine
        @return: a Session object
        '''
        session = Session()
        session.engine = engine
        return session._init()

    @property
    def runtime(self) -> trt.Runtime:
        return self._runtime

    @property
    def engine(self) -> trt.ICudaEngine:
        return self._engine

    @engine.setter
    def engine(self, engine: trt.ICudaEngine):