Retry VLLM [release]

Author: AjayP13

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "DataDreamer"
-version = "0.44.0"
+version = "0.45.0"
 description = "Prompt. Generate Synthetic Data. Train & Align Models."
 license = "MIT"
 authors= [
@@ -66,7 +66,7 @@ warn_unused_ignores = true
 mypy_path = "src/_stubs"
 
 [[tool.mypy.overrides]]
-module = "click,wandb,wandb.*,click.testing,flaky,tensorflow,torch_xla,jax,datasets.features.features,datasets.iterable_dataset,datasets.fingerprint,datasets.builder,datasets.arrow_writer,datasets.splits,datasets.utils,datasets.utils.version,pyarrow.lib,huggingface_hub,huggingface_hub.utils._headers,huggingface_hub.errors,dill,dill.source,transformers,bitsandbytes,sqlitedict,optimum.bettertransformer,optimum.bettertransformer.models,optimum.utils,transformers.utils.quantization_config,sortedcontainers,peft,psutil,ring,ctransformers,petals,petals.client.inference_session,hivemind.p2p.p2p_daemon_bindings.utils,huggingface_hub.utils,tqdm,ctransformers.transformers,vllm,litellm,litellm.llms.palm,litellm.exceptions,sentence_transformers,faiss,huggingface_hub.utils._validators,evaluate,transformers.trainer_callback,transformers.training_args,trl,guidance,sentence_transformers.models.Transformer,trl.trainer.utils,transformers.trainer_utils,setfit,joblib,setfit.modeling,transformers.utils.notebook,mistralai.models,mistralai.models.chat_completion,accelerate.utils,accelerate.utils.constants,accelerate,transformers.trainer,sentence_transformers.util,Pyro5,Pyro5.server,Pyro5.api,Pyro5,datadreamer,huggingface_hub.repocard,transformers.trainer_pt_utils,traitlets.utils.warnings,orjson,Pyro5.errors,sympy,tqdm.auto"
+module = "click,wandb,wandb.*,click.testing,flaky,tensorflow,torch_xla,jax,datasets.features.features,datasets.iterable_dataset,datasets.fingerprint,datasets.builder,datasets.arrow_writer,datasets.splits,datasets.utils,datasets.utils.version,pyarrow.lib,huggingface_hub,huggingface_hub.utils._headers,huggingface_hub.errors,dill,dill.source,transformers,bitsandbytes,sqlitedict,optimum.bettertransformer,optimum.bettertransformer.models,optimum.utils,transformers.utils.quantization_config,sortedcontainers,peft,psutil,ring,ctransformers,petals,petals.client.inference_session,hivemind.p2p.p2p_daemon_bindings.utils,huggingface_hub.utils,tqdm,ctransformers.transformers,vllm,litellm,litellm.llms.palm,litellm.exceptions,sentence_transformers,faiss,huggingface_hub.utils._validators,evaluate,transformers.trainer_callback,transformers.training_args,trl,guidance,sentence_transformers.models.Transformer,trl.trainer.utils,transformers.trainer_utils,setfit,joblib,setfit.modeling,transformers.utils.notebook,mistralai.models,mistralai.models.chat_completion,accelerate.utils,accelerate.utils.constants,accelerate,transformers.trainer,sentence_transformers.util,Pyro5,Pyro5.server,Pyro5.api,Pyro5,datadreamer,huggingface_hub.repocard,transformers.trainer_pt_utils,traitlets.utils.warnings,orjson,Pyro5.errors,sympy,tqdm.auto,requests.exceptions"
 ignore_missing_imports = true
 
 [tool.pyright]

src/llms/vllm.py

@@ -6,6 +6,14 @@
 from typing import Any, Callable, Generator, Iterable
 
 import torch
+from tenacity import (
+    after_log,
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    wait_exponential,
+)
+
 from datasets.fingerprint import Hasher
 
 from .. import DataDreamer
@@ -72,6 +80,34 @@ def __init__(
             self.quantization = "awq"
         self.swap_space = swap_space
 
+    @cached_property
+    def retry_wrapper(self):
+        # Create a retry wrapper function
+        tenacity_logger = self.get_logger(key="retry", verbose=True, log_level=None)
+
+        from requests.exceptions import ConnectionError, Timeout
+
+        @retry(
+            retry=retry_if_exception_type(Timeout),
+            wait=wait_exponential(multiplier=1, min=10, max=60),
+            before_sleep=before_sleep_log(tenacity_logger, logging.INFO),
+            after=after_log(tenacity_logger, logging.INFO),
+            reraise=True,
+        )
+        @retry(
+            retry=retry_if_exception_type(ConnectionError),
+            wait=wait_exponential(multiplier=1, min=10, max=60),
+            before_sleep=before_sleep_log(tenacity_logger, logging.INFO),
+            after=after_log(tenacity_logger, logging.INFO),
+            reraise=True,
+        )
+        def _retry_wrapper(func, *args, **kwargs):
+            return func(*args, **kwargs)
+
+        _retry_wrapper.__wrapped__.__module__ = None  # type: ignore[attr-defined]
+        _retry_wrapper.__wrapped__.__qualname__ = f"{self.__class__.__name__}.run"  # type: ignore[attr-defined]
+        return _retry_wrapper
+
     @cached_property
     def model(self) -> Any:
         env = os.environ.copy()
@@ -119,17 +155,21 @@ def _monkey_patch_init_logger(*args, **kwargs):
                     if datadreamer_logger.level > logging.DEBUG
                     else nullcontext()
                 ):
-                    self_resource.model = LLM(
-                        model=self.model_name,
-                        trust_remote_code=self.trust_remote_code,
-                        dtype=str(self.dtype).replace("torch.", "")
-                        if self.dtype is not None
-                        else "auto",
-                        quantization=self.quantization,
-                        revision=self.revision,
-                        swap_space=self.swap_space,
-                        tensor_parallel_size=tensor_parallel_size,
-                        **kwargs,
+                    self_resource.model = self.retry_wrapper(
+                        func=(
+                            lambda: LLM(
+                                model=self.model_name,
+                                trust_remote_code=self.trust_remote_code,
+                                dtype=str(self.dtype).replace("torch.", "")
+                                if self.dtype is not None
+                                else "auto",
+                                quantization=self.quantization,
+                                revision=self.revision,
+                                swap_space=self.swap_space,
+                                tensor_parallel_size=tensor_parallel_size,
+                                **kwargs,
+                            )
+                        )
                     )
 
                 # Finished loading
@@ -144,7 +184,9 @@ def _monkey_patch_init_logger(*args, **kwargs):
 
             @dill_serializer
             def get_generated_texts_batch(self_resource, *args, **kwargs):
-                outputs = self_resource.model.generate(*args, **kwargs)
+                outputs = self.retry_wrapper(
+                    self_resource.model.generate, *args, **kwargs
+                )
                 generated_texts_batch = [
                     [o.text for o in batch.outputs] for batch in outputs
                 ]
@@ -183,9 +225,9 @@ def _run_batch(  # noqa: C901
         **kwargs,
     ) -> list[str] | list[list[str]]:
         prompts = inputs
-        assert (
-            logit_bias is None
-        ), f"`logit_bias` is not supported for {type(self).__name__}"
+        assert logit_bias is None, (
+            f"`logit_bias` is not supported for {type(self).__name__}"
+        )
         assert seed is None, f"`seed` is not supported for {type(self).__name__}"
 
         SamplingParams = import_module("vllm").SamplingParams
@@ -323,6 +365,7 @@ def __getstate__(self):  # pragma: no cover
         state = super().__getstate__()
 
         # Remove cached model or tokenizer before serializing
+        state.pop("retry_wrapper", None)
         state.pop("model", None)
         state.pop("tokenizer", None)
 

src/trainers/_train_hf_base.py

@@ -7,7 +7,6 @@
 from typing import TYPE_CHECKING, Any, Type, cast
 
 import torch
-
 from datasets.fingerprint import Hasher
 
 from .. import DataDreamer
@@ -490,9 +489,9 @@ def export_to_disk(self, path: str, adapter_only: bool = False) -> PreTrainedMod
         from .train_hf_finetune import TrainHFFineTune
         from .train_setfit_classifier import TrainSetFitClassifier
 
-        assert not adapter_only or self.peft_config, (
-            "`adapter_only` can only be used if a `peft_config` was provided."
-        )
+        assert (
+            not adapter_only or self.peft_config
+        ), "`adapter_only` can only be used if a `peft_config` was provided."
 
         # Clear the directory
         clear_dir(path)
@@ -625,9 +624,9 @@ def publish_to_hf_hub(  # noqa: C901
         from .train_hf_finetune import TrainHFFineTune
         from .train_setfit_classifier import TrainSetFitClassifier
 
-        assert not adapter_only or self.peft_config, (
-            "`adapter_only` can only be used if a `peft_config` was provided."
-        )
+        assert (
+            not adapter_only or self.peft_config
+        ), "`adapter_only` can only be used if a `peft_config` was provided."
 
         # Login
         api = hf_hub_login(token=token)