o -wi@sddlmZddlZddlZddlZddlZddlmZddlm Z ddl m Z ddl mZeeZGdd d eZGd d d ZdS) ) annotationsN)wandb_internal_pb2)wandb_server_pb2) MailboxHandle)MailboxResponseHandlec@seZdZdZdS)MailboxClosedErrorz/The mailbox has been closed and cannot be used.N)__name__ __module__ __qualname____doc__r r V/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/wandb/sdk/mailbox/mailbox.pyrsrc@sBeZdZdZdddZdd d Zdd d ZdddZdddZdS)Mailboxa9Matches service responses to requests. The mailbox can set an address on a server request and create a handle for waiting for a response to that record. Responses are delivered by calling `deliver()`. The `close()` method abandons all handles in case the service process becomes unreachable. returnNonecCsi|_t|_d|_dS)NF)_handles threadingLock _handles_lock_closed)selfr r r__init__s  zMailbox.__init__requestspb.ServerRequest | pb.Record!MailboxHandle[spb.ServerResponse]cCst|tjr|j}rtd|d|}||_n|jj}r)td|d|}||j_|j|j r;t t |}||j |<Wd|S1sOwY|S)aSet a response address on a request. Args: request: The request on which to set a request ID or mailbox slot. This is mutated. An address must not already be set. Returns: A handle for waiting for the response to the request. Raises: MailboxClosedError: If the mailbox has been closed, in which case no new responses are expected to be delivered and new handles cannot be created. z Request already has an address ()N) isinstancespb ServerRequest request_id ValueError _new_addresscontrol mailbox_slotrrrrr)rraddresshandler r rrequire_response#s$     zMailbox.require_responsestrcCs,dd}|}||jvr|}||jvs |S)z[Returns an unused address for a request. Assumes `_handles_lock` is held. cSsdddtdDS)Ncss"|] }ttjtjVqdS)N)secretschoicestringascii_lowercasedigits).0ir r r Rs  z9Mailbox._new_address..generate.. )joinranger r r rgenerateQs z&Mailbox._new_address..generate)r)rr5r%r r rr"Ks   zMailbox._new_addressresponsespb.ServerResponsecCs|j}|s#|d}|dkr|jd}d|}td|dS|j|j|d}Wdn1s8wY|rF||dSdS)zDeliver a response from the service. If the response address is invalid, this does nothing. It is a no-op if the mailbox has been closed. server_response_typeresult_communicate result_typezresult_communicate.z(Received response with no mailbox slot: N) r WhichOneofr9_loggererrorrrpopdeliver)rr6r%kindr:r&r r rr?_s   zMailbox.delivercCsl|j)d|_tdt|jd|jD]}|q|jWddS1s/wYdS)zXIndicate no further responses will be delivered. Abandons all handles. TzClosing mailbox, abandoning z handles.N) rrr<infolenrvaluesabandonclear)rr&r r rcloseys  "z Mailbox.closeN)rr)rrrr)rr()r6r7rr) r r r r rr'r"r?rFr r r rrs   ( r) __future__rloggingr*r,r wandb.protorpbrrmailbox_handlerresponse_handler getLoggerr r< Exceptionrrr r r rs