Class LoraBaseModel

trainer/src/lora_inference/main.py:85–161

    container_idle_timeout=300,
)
class LoraBaseModel:
    def __init__(self, base_model_id: str):
        model_dir = merged_model_cache_dir(base_model_id)
        logging.info(f"Loading base model {model_dir}")
        self.seen_models = set()

        self.engine = AsyncLLMEngine.from_engine_args(
            AsyncEngineArgs(
                model=base_model_id,
                enable_lora=True,
                max_loras=MAX_INPUTS,
                max_lora_rank=8,
                max_cpu_loras=500,
            )
        )

    @modal.method()
    async def generate(self, request: Input) -> Output:
        logging.info(f"Processing for model {request.lora_model}")
        lora_dir = lora_model_cache_dir(request.lora_model)

        if request.lora_model not in self.seen_models:
            if not os.path.exists(lora_model_cache_dir(request.lora_model)):
                logging.info(f"Couldn't find model, reloading {lora_dir}")
                stub.volume.reload()

            if not os.path.exists(lora_model_cache_dir(request.lora_model)):
                raise Exception(f"Couldn't find model {lora_dir}after reloading!")

            self.seen_models.add(request.lora_model)

        sample_params = SamplingParams(
            n=request.n,
            temperature=request.temperature,
            max_tokens=request.max_tokens,
        )

        lora_request = LoRARequest(
            request.lora_model,
            abs(hash(request.lora_model)),
            lora_dir,
        )

        request_id = random_uuid()

        logging.info(f"Generating for request {request_id}")
        output_generator = self.engine.generate(
            request.prompt,
            sample_params,
            request_id=request_id,
            lora_request=lora_request,
        )

        final_output: Union[RequestOutput, None] = None
        async for request_output in output_generator:
            # TODO: support streaming
            final_output = request_output