o biq@sddlZddlZddlZddlZddlmZmZmZmZm Z m Z er&ddl Z ddl Z ddlZddlmZddlZddlmZddlmZddlmZmZmZmZmZmZmZm Z m!Z!ddl"m#Z#m$Z$ddlm%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z:m;Z;mZ>ddl?m@Z@dd lAmBZBmCZCeDeEZFeGd d ZHGd d d eIZJdedeKfddZLddZMddZNGdddZOdS)N) TYPE_CHECKINGAnyListOptionalTupleUnion) ray_constants)TensorTransportEnum) DynamicObjectRefGeneratorMessagePackSerializedObjectMessagePackSerializerPickle5SerializedObject Pickle5WriterRawSerializedObjectSerializedRayObject split_bufferunpack_pickle5_buffers) ErrorType RayErrorInfo)ActorDiedErrorActorPlacementGroupRemovedActorUnavailableErrorActorUnschedulableErrorLocalRayletDiedError NodeDiedErrorObjectFetchTimedOutErrorObjectFreedErrorObjectLostErrorObjectReconstructionFailedError-ObjectReconstructionFailedLineageEvictedError2ObjectReconstructionFailedMaxAttemptsExceededErrorObjectRefStreamEndOfStreamErrorOutOfDiskErrorOutOfMemoryErrorOwnerDiedErrorPlasmaObjectNotAvailableRayErrorRaySystemError RayTaskErrorReferenceCountingAssertionErrorRuntimeEnvSetupErrorTaskCancelledErrorTaskPlacementGroupRemovedTaskUnschedulableErrorWorkerCrashedError)CompiledDAGRef)inspect_serializabilityserialization_addons.RAY_allow_out_of_band_object_ref_serializationTc@s eZdZdS)DeserializationErrorN)__name__ __module__ __qualname__r7r7N/home/ubuntu/.local/lib/python3.10/site-packages/ray/_private/serialization.pyr3Bsr3obj error_msgc Cstzt|WSttjjfy9}z#t}t||d|d| }t |tr/t||tj|d}~ww)zdWrap cloudpickle.dumps to provide better error message when the object is not serializable. ) print_filez: N) pickledumps TypeErrorray exceptions(OufOfBandObjectRefSerializationExceptionioStringIOr0getvalue isinstance)r9r:esiomsgr7r7r8 pickle_dumpsFs     rIcCs`t|||}|r.tjjj}||}|}|dur#tj}|j | ||||SN) r? ObjectRef_privateworker global_workercheck_connectedget_serialization_contextget_outer_object_refnil core_worker#deserialize_and_register_object_refbinary)rU call_site owner_address object_statusobj_refrMcontextouter_idr7r7r8_object_ref_deserializerVs    r\cCs(tjjj}|}tjj|||SrJ) r?rLrMrNrPrQactor ActorHandle_deserialization_helper)serialized_objweak_refrZr[r7r7r8_actor_handle_deserializerss rbc @sNeZdZdZddZddZddZdd Zd d Zd d Z ddZ ddZ ddZ ddddde deefddZ d:dededeedefd d!Zdejfdeedeefd"d#Zd$d%Zd&d'Zdeefd(d)Zd*eefd+d,Zd-d.Zd/d0Zd1ed2edefd3d4Z d1ede!e"effd5d6Z#d1ede$eed7ffd8d9Z%dS);SerializationContextzInitialize the serialization library. This defines a custom serializer for object refs and also tells ray to serialize several exception classes that we define for error handling. cs||_t_fdd}tjj|dd}t|fdd}tj |dd}t |t dS) Ncs:tjjj|\}}}|sj|ddt||ffS)NT)allow_out_of_band_serialization)r?rLrMrNrO_serialization_helperadd_contained_object_refrb)r9 serializedactor_handle_idraselfr7r8actor_handle_reducers z;SerializationContext.__init__..actor_handle_reducercSstd)Nz1Serialization of CompiledDAGRef is not supported.)r>r9r7r7r8compiled_dag_ref_reducerz?SerializationContext.__init__..compiled_dag_ref_reducercsPtjjj}|j|t|d|j |\}}}t | |||ffS)N)rdrV) r?rLrMrNrOrf*ALLOW_OUT_OF_BAND_OBJECT_REF_SERIALIZATIONrVrSserialize_object_refr\rU)r9rMrWrXrir7r8object_ref_reducers  z9SerializationContext.__init__..object_ref_reducercSs t|jffSrJ)r _refsrlr7r7r8object_ref_generator_reducer zCSerializationContext.__init__..object_ref_generator_reducer) rM threadinglocal _thread_local_register_cloudpickle_reducerr?r]r^r/rKr r1apply)rjrMrkrmrqrsr7rir8__init__s     zSerializationContext.__init__cCs|tjj|<dSrJr< CloudPicklerdispatch)rjclsreducerr7r7r8rxsz2SerializationContext._register_cloudpickle_reducercCstjj|ddSrJ)r<r|r}pop)rjr~r7r7r8_unregister_cloudpickle_reducersz4SerializationContext._unregister_cloudpickle_reducercsfdd}|tjj|<dS)Ncs|ffSrJr7rlcustom_deserializercustom_serializerr7r8_CloudPicklerReducerzSSerializationContext._register_cloudpickle_serializer.._CloudPicklerReducerr{)rjr~rrrr7rr8 _register_cloudpickle_serializersz5SerializationContext._register_cloudpickle_serializercCst|jddS)Nin_bandFgetattrrwrir7r7r8is_in_band_serializationrz-SerializationContext.is_in_band_serializationcC d|j_dS)NTrwrrir7r7r8set_in_band_serializationrtz.SerializationContext.set_in_band_serializationcCr)NFrrir7r7r8set_out_of_band_serializationrtz2SerializationContext.set_out_of_band_serializationcCst|jdg}|r |dSdS)Nobject_ref_stackr)rjstackr7r7r8rQsz)SerializationContext.get_outer_object_refcCs2t|jdst|j_tS|jj}t|j_|S)N object_refs)hasattrrwsetr)rjrr7r7r8#get_and_clear_contained_object_refss   z8SerializationContext.get_and_clear_contained_object_refsN)rV object_refz ray.ObjectRefrdrVcCsj|rt|jdst|j_|jj|dS|s*tjd| d|p&dtj j j j |dS)Nrz-It is not allowed to serialize ray.ObjectRef a&. If you want to allow serialization, set `RAY_allow_out_of_band_object_ref_serialization=1.` If you set the env var, the object is pinned forever in the lifetime of the worker process and can cause Ray object leaks. See the callsite and trace to find where the serialization occurs. Callsite: z-Disabled. Set RAY_record_ref_creation_sites=1)rrrwrraddr?r@rAhexrLrMrNrSadd_object_ref_reference)rjrrdrVr7r7r8rfs    z-SerializationContext.add_contained_object_refdatatensor_transport object_idreturnc Csddlm}|j}|tjk}|rAtjjj j }|j |s0| |s+Jd|d|||j |}|||j |z2zt|\} } t| dkrWtj| | d} nt| } Wn tjjyitwW|rr|g| S|r||gww)a Args: data: The data to deserialize. tensor_transport: The tensor transport to use. If not equal to OBJECT_STORE, it means that any tensors in the object are sent out-of-band instead of through the object store. In this case, we need to retrieve the tensors from the in-actor object store. Then, we deserialize `data` with the retrieved tensors in the serialization context. object_id: The object ID to use as the key for the in-actor object store to retrieve tensors. Returns: Any: The deserialized object. rChannelContextzobj_id=z not found in GPU object store. This error is unexpected. Please report this issue on GitHub: https://github.com/ray-project/ray/issues/new/choose)buffers)ray.experimental.channelr get_currentserialization_contextr OBJECT_STOREr?rLrMrNgpu_object_managergpu_object_storehas_gpu_objectis_managed_gpu_objectfetch_gpu_objectget_gpu_objectreset_out_of_band_tensorsremove_gpu_objectrlenr<loads PicklingErrorr3) rjrrrrctxenable_gpu_objectsrtensorsrrr9r7r7r8_deserialize_pickle5_datas>               z.SerializationContext._deserialize_pickle5_datac sbt|\}}|dtjkr||||ngzfdd}t||}W|Sty0tw)Nrcs|SrJr7)indexpython_objectsr7r8_python_deserializerMrnzLSerializationContext._deserialize_msgpack_data.._python_deserializer)rrOBJECT_METADATA_TYPE_PYTHONrr r Exceptionr3) rjrmetadata_fieldsrr msgpack_data pickle5_datarr9r7rr8_deserialize_msgpack_data9s   z.SerializationContext._deserialize_msgpack_datacCs0|sJ|||}|sJt}|||SrJ)rrParseFromString)rjrrpb_bytesray_error_infor7r7r8_deserialize_error_infoUs   z,SerializationContext._deserialize_error_infocCs\|stS|||}|dsJ|jdrt|jjS|jds'Jt|jjdS)Nactor_died_errorcreation_task_failure_contextactor_died_error_context)cause)rrHasFieldrr&from_ray_exceptionrr)rjrrrr7r7r8_deserialize_actor_died_error^s  z2SerializationContext._deserialize_actor_died_errorcCs|durtj}|r;|d}|dtjtjfvr#|||||S|dtjkr4|dur0dS| S|dtj krU|||}|dd|dddk}}t ||Szt |d} Wnt ynt d|d|w| tdkr|||}t|S| td krtS| td kr|||S| td krtS| td krzd } |r|||} | j} t| dWStjjjy|||}t|YSw| tdkrt|||S| tdkrt|||S| tdkrt |||S| tdkr!|||} t!| jS| tdkr4|||} t"| jS| tdkrHt#|||S| tdkr\t$|||S| tdkrpt%|||S| tdkrt&|||S| tdkrt'|||S| tdkrt(|||S| tdkr|||} d } | )dr| j*j} t+| dS| tdkrt,S| tdkrt-S| tdkr|||} t.| jS| tdkr|||} t/| jS| td krt0S| td!kr3|||} | )d"r+| j1j2} nd} t3| j| St4d#t5| S|rBt6d$t7S)%N,rr1zCan't deserialize object: z , metadata: TASK_EXECUTION_EXCEPTION WORKER_DIED ACTOR_DIEDLOCAL_RAYLET_DIEDTASK_CANCELLED) error_message OBJECT_LOSTOBJECT_FETCH_TIMED_OUTOUT_OF_DISK_ERROR OUT_OF_MEMORY NODE_DIEDOBJECT_DELETED OBJECT_FREED OWNER_DIEDOBJECT_UNRECONSTRUCTABLE.OBJECT_UNRECONSTRUCTABLE_MAX_ATTEMPTS_EXCEEDED(OBJECT_UNRECONSTRUCTABLE_LINEAGE_EVICTEDRUNTIME_ENV_SETUP_FAILEDruntime_env_setup_failed_errorTASK_PLACEMENT_GROUP_REMOVEDACTOR_PLACEMENT_GROUP_REMOVEDTASK_UNSCHEDULABLE_ERRORACTOR_UNSCHEDULABLE_ERROREND_OF_STREAMING_GENERATORACTOR_UNAVAILABLEactor_unavailable_errorzUnrecognized error type z+non-null object should always have metadata)8r rsplitr#OBJECT_METADATA_TYPE_CROSS_LANGUAGErrrOBJECT_METADATA_TYPE_RAW to_pybytes!OBJECT_METADATA_TYPE_ACTOR_HANDLErbintrrValuer& from_bytesr.rrrrr+googleprotobufmessage DecodeErrorrrWrVrr"r#rr)rr$rr rrrr*r,rr-rr!ractor_idrr'str ValueErrorr%)rjrmetadatarrrr9rgra error_typer error_infor:rr7r7r8_deserialize_objectms                          z(SerializationContext._deserialize_objectserialized_ray_objectsc Cst|t|ks Jt|jdsg|j_g}t||D]U\}\}}}z._python_serializer)rEr( issubclassr __class__r+rrrencodeto_bytesr?r]r^rerrrrr r=rr r ) rjrcontained_object_refsrrgrhrarrpickle5_serialized_objectr7rr8_serialize_to_msgpack0s8      z*SerializationContext._serialize_to_msgpackrobj_idcCsN|dusJd||\}}|r%|d}tjjj}|j}|j|||S)aRetrieve GPU data from `value` and store it in the GPU object store. Then, return the serialized value. Args: value: The value to serialize. obj_id: The object ID of the value. `obj_id` is required, and the GPU data (e.g. tensors) in `value` will be stored in the GPU object store with the key `obj_id`. Returns: Serialized value. Nzd`obj_id` is required, and it is the key to retrieve corresponding tensors from the GPU object store.r ) _serialize_and_retrieve_tensorsdecoder?rLrMrNrradd_gpu_object)rjrrserialized_valrrMrr7r7r8serialize_and_store_gpu_objectsgs   z4SerializationContext.serialize_and_store_gpu_objectscCst|tr t|S||S)zSerialize an object. Args: value: The value to serialize. Returns: Serialized value. )rEbytesrr)rjrr7r7r8 serializes  zSerializationContext.serializez torch.Tensorc Csbddlm}|j}|j}|dz ||}W||n||w|g\}}||fS)z Serialize `value` and return the serialized value and any tensors retrieved from `value`. This is only used for GPU objects. rrT)rrrruse_external_transportset_use_external_transportrr)rjrrrprev_use_external_transportrr_r7r7r8rs    z4SerializationContext._serialize_and_retrieve_tensorsrJ)&r4r5r6__doc__rzrxrrrrrrQrboolrrrfrr rrrrrrrrrr rrr rrrrrrr7r7r7r8rc}s7  ' :    7   rc)PrBloggingrurtypingrrrrrrtorchgoogle.protobuf.messagerray._private.utilsr?ray.cloudpickle cloudpickler<ray.exceptions ray._privaterray._private.custom_typesr ray._rayletr r r r rrrrrray.core.generated.common_pb2rrrrrrrrrrrrrr r!r"r#r$r%r&r'r(r)r*r+r,r-r.!ray.experimental.compiled_dag_refr/ray.utilr0r1 getLoggerr4renv_boolrorr3rrIr\rbrcr7r7r7r8s6    , p