codekingpro/portable-devtools

1import importlib.util2import logging3import pickle4from typing import Any, Callable, List, Mapping, Optional5 6from langchain_core.callbacks import CallbackManagerForLLMRun7from langchain_core.language_models.llms import LLM8from pydantic import ConfigDict9 10from langchain_community.llms.utils import enforce_stop_tokens11 12logger = logging.getLogger(__name__)13 14 15def _generate_text(16    pipeline: Any,17    prompt: str,18    *args: Any,19    stop: Optional[List[str]] = None,20    **kwargs: Any,21) -> str:22    """Inference function to send to the remote hardware.23 24    Accepts a pipeline callable (or, more likely,25    a key pointing to the model on the cluster's object store)26    and returns text predictions for each document27    in the batch.28    """29    text = pipeline(prompt, *args, **kwargs)30    if stop is not None:31        text = enforce_stop_tokens(text, stop)32    return text33 34 35def _send_pipeline_to_device(pipeline: Any, device: int) -> Any:36    """Send a pipeline to a device on the cluster."""37    if isinstance(pipeline, str):38        with open(pipeline, "rb") as f:39            # This code path can only be triggered if the user40            # passed allow_dangerous_deserialization=True41            pipeline = pickle.load(f)  # ignore[pickle]: explicit-opt-in42 43    if importlib.util.find_spec("torch") is not None:44        import torch45 46        cuda_device_count = torch.cuda.device_count()47        if device < -1 or (device >= cuda_device_count):48            raise ValueError(49                f"Got device=={device}, "50                f"device is required to be within [-1, {cuda_device_count})"51            )52        if device < 0 and cuda_device_count > 0:53            logger.warning(54                "Device has %d GPUs available. "55                "Provide device={deviceId} to `from_model_id` to use available"56                "GPUs for execution. deviceId is -1 for CPU and "57                "can be a positive integer associated with CUDA device id.",58                cuda_device_count,59            )60 61        pipeline.device = torch.device(device)62        pipeline.model = pipeline.model.to(pipeline.device)63    return pipeline64 65 66class SelfHostedPipeline(LLM):67    """Model inference on self-hosted remote hardware.68 69    Supported hardware includes auto-launched instances on AWS, GCP, Azure,70    and Lambda, as well as servers specified71    by IP address and SSH credentials (such as on-prem, or another72    cloud like Paperspace, Coreweave, etc.).73 74    To use, you should have the ``runhouse`` python package installed.75 76    Example for custom pipeline and inference functions:77        .. code-block:: python78 79            from langchain_community.llms import SelfHostedPipeline80            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline81            import runhouse as rh82 83            def load_pipeline():84                tokenizer = AutoTokenizer.from_pretrained("gpt2")85                model = AutoModelForCausalLM.from_pretrained("gpt2")86                return pipeline(87                    "text-generation", model=model, tokenizer=tokenizer,88                    max_new_tokens=1089                )90            def inference_fn(pipeline, prompt, stop = None):91                return pipeline(prompt)[0]["generated_text"]92 93            gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")94            llm = SelfHostedPipeline(95                model_load_fn=load_pipeline,96                hardware=gpu,97                model_reqs=model_reqs, inference_fn=inference_fn98            )99    Example for <2GB model (can be serialized and sent directly to the server):100        .. code-block:: python101 102            from langchain_community.llms import SelfHostedPipeline103            import runhouse as rh104            gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")105            my_model = ...106            llm = SelfHostedPipeline.from_pipeline(107                pipeline=my_model,108                hardware=gpu,109                model_reqs=["./", "torch", "transformers"],110            )111    Example passing model path for larger models:112        .. code-block:: python113 114            from langchain_community.llms import SelfHostedPipeline115            import runhouse as rh116            import pickle117            from transformers import pipeline118 119            generator = pipeline(model="gpt2")120            rh.blob(pickle.dumps(generator), path="models/pipeline.pkl"121                ).save().to(gpu, path="models")122            llm = SelfHostedPipeline.from_pipeline(123                pipeline="models/pipeline.pkl",124                hardware=gpu,125                model_reqs=["./", "torch", "transformers"],126            )127    """128 129    pipeline_ref: Any = None  #: :meta private:130    client: Any = None  #: :meta private:131    inference_fn: Callable = _generate_text  #: :meta private:132    """Inference function to send to the remote hardware."""133    hardware: Any = None134    """Remote hardware to send the inference function to."""135    model_load_fn: Callable136    """Function to load the model remotely on the server."""137    load_fn_kwargs: Optional[dict] = None138    """Keyword arguments to pass to the model load function."""139    model_reqs: List[str] = ["./", "torch"]140    """Requirements to install on hardware to inference the model."""141 142    allow_dangerous_deserialization: bool = False143    """Allow deserialization using pickle which can be dangerous if 144    loading compromised data.145    """146 147    model_config = ConfigDict(148        extra="forbid",149    )150 151    def __init__(self, **kwargs: Any):152        """Init the pipeline with an auxiliary function.153 154        The load function must be in global scope to be imported155        and run on the server, i.e. in a module and not a REPL or closure.156        Then, initialize the remote inference function.157        """158        if not kwargs.get("allow_dangerous_deserialization"):159            raise ValueError(160                "SelfHostedPipeline relies on the pickle module. "161                "You will need to set allow_dangerous_deserialization=True "162                "if you want to opt-in to allow deserialization of data using pickle."163                "Data can be compromised by a malicious actor if "164                "not handled properly to include "165                "a malicious payload that when deserialized with "166                "pickle can execute arbitrary code. "167            )168        super().__init__(**kwargs)169        try:170            import runhouse as rh171 172        except ImportError:173            raise ImportError(174                "Could not import runhouse python package. "175                "Please install it with `pip install runhouse`."176            )177 178        remote_load_fn = rh.function(fn=self.model_load_fn).to(179            self.hardware, reqs=self.model_reqs180        )181        _load_fn_kwargs = self.load_fn_kwargs or {}182        self.pipeline_ref = remote_load_fn.remote(**_load_fn_kwargs)183 184        self.client = rh.function(fn=self.inference_fn).to(185            self.hardware, reqs=self.model_reqs186        )187 188    @classmethod189    def from_pipeline(190        cls,191        pipeline: Any,192        hardware: Any,193        model_reqs: Optional[List[str]] = None,194        device: int = 0,195        **kwargs: Any,196    ) -> LLM:197        """Init the SelfHostedPipeline from a pipeline object or string."""198        if not isinstance(pipeline, str):199            logger.warning(200                "Serializing pipeline to send to remote hardware. "201                "Note, it can be quite slow"202                "to serialize and send large models with each execution. "203                "Consider sending the pipeline"204                "to the cluster and passing the path to the pipeline instead."205            )206 207        load_fn_kwargs = {"pipeline": pipeline, "device": device}208        return cls(209            load_fn_kwargs=load_fn_kwargs,210            model_load_fn=_send_pipeline_to_device,211            hardware=hardware,212            model_reqs=["transformers", "torch"] + (model_reqs or []),213            **kwargs,214        )215 216    @property217    def _identifying_params(self) -> Mapping[str, Any]:218        """Get the identifying parameters."""219        return {220            **{"hardware": self.hardware},221        }222 223    @property224    def _llm_type(self) -> str:225        return "self_hosted_llm"226 227    def _call(228        self,229        prompt: str,230        stop: Optional[List[str]] = None,231        run_manager: Optional[CallbackManagerForLLMRun] = None,232        **kwargs: Any,233    ) -> str:234        return self.client(235            pipeline=self.pipeline_ref, prompt=prompt, stop=stop, **kwargs236        )237