feat: use cloudpickle to serialize transformers_modules by value

Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>

vllm/transformers_utils/config.py

@@ -191,6 +191,29 @@ def get_config(
             )
             config.update({key: value})
 
+    if trust_remote_code:
+        # With trust_remote_code, the config is typically an instance of a
+        # custom class imported from the HF modules cache.
+        #
+        # The class will not be importable in Ray workers by default (and won't
+        # exist at all on other nodes), which breaks serialization of the
+        # config. Here we tell the serialization library used by Ray to pass
+        # instances of these generated classes by value instead of by reference
+        # (eg. the class definition is serialized along with its data).
+        #
+        # See: https://github.com/cloudpipe/cloudpickle?tab=readme-ov-file#overriding-pickles-serialization-mechanism-for-importable-constructs
+        try:
+            import cloudpickle
+            import ray
+            import transformers_modules
+            cloudpickle.register_pickle_by_value(transformers_modules)
+            # Ray vendors its own version of cloudpickle
+            ray.cloudpickle.register_pickle_by_value(transformers_modules)
+        # ignore import errors in the case that trust_remote_code is set
+        # unnecessarily and transformers_modules does not exist
+        except ImportError:
+            pass
+
     return config
 
 

vllm/utils.py

@@ -945,6 +945,8 @@ def flatten_2d_lists(lists: List[List[T]]) -> List[T]:
     return [item for sublist in lists for item in sublist]
 
 
+# TODO: This function can be removed if transformer_modules classes are
+# serialized by value when communicating between processes
 def init_cached_hf_modules() -> None:
     """
     Lazy initialization of the Hugging Face modules.