o
    xi                     @  sP  d dl mZ d dlZd dlZd dlZd dlZd dlmZ d dlm	Z	 d dl
mZ d dlmZmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZmZ d dlmZmZ d dlm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z& ddl'm(Z(m)Z) ddl*m+Z, dZ-erd dl.m/Z/ d dl0m1Z1 ddl2m3Z3 e4dZ5d"ddZ6d#ddZ7G d d! d!ej8Z9dS )$    )annotationsN)Iterable)Path)	token_hex)TYPE_CHECKINGAny)termwarn)wandb_internal_pb2)wandb_telemetry_pb2)	json_util)	FilesDict
PolicyName)HandleAbandonedErrorMailboxHandle)WandBJSONEncoderOldget_h5_typenamejson_dumps_saferjson_dumps_safer_historyjson_friendlyjson_friendly_valmaybe_compress_summary   )history_dict_to_jsonval_to_json   )summary_recordi )Artifact)ArtifactManifest)Runwandbpolicyr   returnpb.FilesItem.PolicyType.Vc                 C  sB   | dkrt jjj}|S | dkrt jjj}|S | dkrt jjj}|S Nnowendlivepb	FilesItem
PolicyTypeNOWENDLIVE)r    enum r/   Q/home/ubuntu/.local/lib/python3.10/site-packages/wandb/sdk/interface/interface.pyfile_policy_to_enum+   s   


r1   r.   c                 C  sB   | t jjjkrd}|S | t jjjkrd}|S | t jjjkrd}|S r#   r'   )r.   r    r/   r/   r0   file_enum_to_policy5   s   r2   c                   @  s  e Zd ZU dZded< dlddZejdmddZdlddZ	ejdnddZ
doddZejdpddZ				dqdrd#d$Zdsd(d)Zdtd*d+Zejdud,d-Zdvd0d1Zejdwd4d5Z			dxdyd6d7Zejdzd9d:Zejd{d=d>Zd|dBdCZd}dGdHZd~dKdLZddMdNZejddPdQZddUdVZddWdXZejddZd[Zdld\d]Zejdd`daZddedfZ 	dddkdlZ!ddndoZ"			dxddwdxZ#ejdd{d|Z$e%dddZ&dddZ'dddZ(ejdddZ)					ddddZ*ejdddZ+dddZ,ejdddZ-				ddddZ.ejdddZ/ddddZ0ejdddZ1ejdddZ2dddZ3ejdddZ4			ddddƄZ5ejdddʄZ6		dddd̄Z7ejdddτZ8dlddфZ9ejdddՄZ:dd֜dddڄZ;ejdddބZ<dd֜dddZ=ejdddZ>dlddZ?ejdddZ@dlddZAejdddZBdddZCejdddZDdddZEdddZFejdddZGdlddZHejdddZI		ddddZJejdddZKdlddZLejdddZMdlddZNejdoddZOddd ZPdod!d"ZQejdd%d&ZRejdd'd(ZSdd)d*ZTejdd-d.ZUdd0d1ZVejdd3d4ZWdod5d6ZXejdd8d9ZYdod:d;ZZejdd=d>Z[dod?d@Z\ejddCdDZ]dodEdFZ^ejddIdJZ_dodKdLZ`ejddNdOZaddPdQZbejddRdSZcdodTdUZdejddXdYZedodZd[Zfejdd^d_Zgdod`daZhejddddeZidodfdgZjejddjdkZkdS (  InterfaceBasezMethods for sending run messages (Records) to the service.

    None of the methods may be called from an asyncio context other than
    deliver_async() or those with a `nowait=True` argument.
    bool_dropr!   Nonec                 C  s
   d| _ d S )NF)r5   selfr/   r/   r0   __init__H   s   
zInterfaceBase.__init__record	pb.RecordMailboxHandle[pb.Result]c                   s   dS )a;  Send a record and create a handle to wait for the response.

        The synchronous publish and deliver methods on this class cannot be
        called in the asyncio thread because they block. Instead of having
        an async copy of every method, this is a general method for sending
        any kind of record in the asyncio thread.

        Args:
            record: The record to send. This method takes ownership of the
                record and it must not be used afterward.

        Returns:
            A handle to wait for a response to the record.
        Nr/   )r8   r:   r/   r/   r0   deliver_asyncK       zInterfaceBase.deliver_asyncc                 C     t  }| | d S N)r(   HeaderRecord_publish_headerr8   headerr/   r/   r0   publish_header_      zInterfaceBase.publish_headerrD   pb.HeaderRecordc                 C     t r@   NotImplementedErrorrC   r/   r/   r0   rB   c      zInterfaceBase._publish_headerc                 C  s   |  t S r@   )_deliver_statusr(   StatusRequestr7   r/   r/   r0   deliver_statusg      zInterfaceBase.deliver_statusstatuspb.StatusRequestc                 C  rH   r@   rI   r8   rP   r/   r/   r0   rL   j      zInterfaceBase._deliver_statusNdatadict | Nonekeytuple[str, ...] | str | Noneval
Any | Noneobjpb.ConfigRecord | Nonepb.ConfigRecordc           	      C  s   |pt  }|r"| D ]\}}|j }||_tt|d |_q|rF|j }t	|t
r:|D ]}|j| q0n||_tt|d |_|S )Nr   )r(   ConfigRecorditemsupdateaddrV   r   r   
value_json
isinstancetuple
nested_keyappend)	r8   rT   rV   rX   rZ   configkvr_   r/   r/   r0   _make_configq   s   


zInterfaceBase._make_configrunr   pb.RunRecordc                 C  s  t  }|jjd ur|jj|_|jjd ur|jj|_|jjd ur%|jj|_|jjd ur0|jj|_|jjd ur;|jj|_|jj	d urF|jj	|_
|jjd urQ|jj|_|jjd ur_|j|jj |jd uro|jt|jd  |jd urx|j|_|jjd ur|jj|j_|jjd ur|jj|j_|jjd ur|jj|_|jjr|jj|_|jjr|jj|_|jjr|jj}|j|j _|j!|j _!|j"|j _"|jj#r|jj#}|j|j _|j!|j _!|j"|j _"|j$r|j$|_%|j&d ur|j&' }| j(||j)d |j*r|j+,|j* |j-r|j-|_.|S )Ng    .A)rT   rZ   )/r(   	RunRecord	_settingsentityproject	run_grouprun_job_typejob_typerun_idrun_namedisplay_name	run_notesnotesrun_tagstagsextend_start_time
start_timeFromMicrosecondsint_starting_stepstarting_stepgit_remote_urlgit
remote_url
git_commitcommitsweep_idhostresumed	fork_fromrj   branch_pointmetricvalueresume_from_forkedforked_config_as_dictri   rf   _telemetry_obj	telemetry	MergeFrom_start_runtimeruntime)r8   rj   	proto_run
run_momentconfig_dictr/   r/   r0   	_make_run   sf   



















zInterfaceBase._make_runc                 C     |  |}| | d S r@   )r   _publish_runr8   rj   
run_recordr/   r/   r0   publish_run      
zInterfaceBase.publish_runc                 C  rH   r@   rI   r8   rj   r/   r/   r0   r      rK   zInterfaceBase._publish_runcancel_slotstrc                 C  s   t j|d}| | d S )N)r   )r(   CancelRequest_publish_cancel)r8   r   cancelr/   r/   r0   publish_cancel   s   zInterfaceBase.publish_cancelr   pb.CancelRequestc                 C  rH   r@   rI   )r8   r   r/   r/   r0   r      rK   zInterfaceBase._publish_cancelc                 C  s   | j |||d}| | d S )N)rT   rV   rX   )ri   _publish_config)r8   rT   rV   rX   cfgr/   r/   r0   publish_config   s   zInterfaceBase.publish_configr   c                 C  rH   r@   rI   )r8   r   r/   r/   r0   r      rK   zInterfaceBase._publish_configr   pb.MetricRecordc                 C  rH   r@   rI   )r8   r   r/   r/   r0   _publish_metric   rK   zInterfaceBase._publish_metricsummary_dictdictpb.SummaryRecordc                 C  s:   t  }| D ]\}}|j }||_t||_q|S r@   )	r(   SummaryRecordr^   r_   r`   rV   jsondumpsra   )r8   r   summaryrg   rh   r_   r/   r/   r0   _make_summary_from_dict   s   
z%InterfaceBase._make_summary_from_dictr   r   path_from_rootc           	      C  st   t |tr i }| D ]\}}| j||d | |d||< q|S tt|||dd\}}t|t|\}}|r8	 |S )a  Normalize, compress, and encode sub-objects for backend storage.

        value: Object to encode.
        path_from_root: `str` dot separated string from the top-level summary to the
            current `value`.

        Returns:
            A new tree of dict's with large objects replaced with dictionaries
            with "_type" entries that say which type the original data was.
        .rj   r   )	namespace)rb   r   r^   _summary_encoder   r   r   r   )	r8   r   r   rj   
json_valuerV   friendly_value	converted
compressedr/   r/   r0   r      s$   

zInterfaceBase._summary_encoder   sr.SummaryRecordc           
      C  s   t  }|jD ]A}|j }t|j}|dksJ |dkr%|j|j n|jd |_d|j}| j	|j
||d}t|\}}	tj|td|_q|jD ]$}|j }t|j}|dks^J |dkrj|j|j qL|jd |_qL|S )Nr   r   r   r   )cls)r(   r   r_   r`   lenrV   rd   rz   joinr   r   r   r   r   r   ra   remove)
r8   r   rj   pb_summary_recorditempb_summary_item
key_lengthr   r   _r/   r/   r0   _make_summary  s6   






zInterfaceBase._make_summaryc                 C  s   | j ||d}| | d S Nr   )r   _publish_summary)r8   rj   r   r   r/   r/   r0   publish_summary=  s   zInterfaceBase.publish_summaryr   c                 C  rH   r@   rI   )r8   r   r/   r/   r0   r   E  rK   zInterfaceBase._publish_summary
files_dictr   pb.FilesRecordc                 C  s8   t  }|d D ]\}}|j }||_t||_q|S )Nfiles)r(   FilesRecordr   r`   pathr1   r    )r8   r   r   r   r    fr/   r/   r0   _make_filesI  s   
zInterfaceBase._make_filesc                 C  r   r@   )r   _publish_files)r8   r   r   r/   r/   r0   publish_filesQ  r   zInterfaceBase.publish_filesr   c                 C  rH   r@   rI   )r8   r   r/   r/   r0   r   U  rK   zInterfaceBase._publish_filesc                 C  s4   t  }|D ]}|jj|j|jd q| | d S )N)nameversion)r(   PythonPackagesRequestpackager`   rV   r   _publish_python_packages)r8   working_setpython_packagespkgr/   r/   r0   publish_python_packagesY  s   z%InterfaceBase.publish_python_packagesr   pb.PythonPackagesRequestc                 C  rH   r@   rI   )r8   r   r/   r/   r0   r   _     z&InterfaceBase._publish_python_packagesartifactr   pb.ArtifactRecordc                 C  s   t  }|j|_|j|_|j|_|j|_|j|_|j	r|j	|_	|j
r&|j
|_
|jr2tt|j|_|jr9|j|_| }|rB||_|j|_| j|j|jd |S )N)rZ   )r(   ArtifactRecordtyper   
_client_id	client_id_sequence_client_idsequence_client_iddigestdistributed_iddescriptionmetadatar   r   r   _base_idbase_id_ttl_duration_seconds_to_gqlttl_duration_secondsincrementalincremental_beta1_make_artifact_manifestmanifest)r8   r   proto_artifactttl_duration_inputr/   r/   r0   _make_artifacte  s(   zInterfaceBase._make_artifactartifact_manifestr   pb.ArtifactManifest | Nonepb.ArtifactManifestc                 C  s(  |pt  }| |_|j |_t|tkr!| |}||_|S |j	 
 p+i 
 D ]\}}|j }||_t||_q,t|j dd dD ]G}|j }	|j|	_|j|	_|jr`|j|	_|jrg|j|	_|jrn|j|	_|jru|j|	_|j|	_|j
 D ]\}}|	j }
||
_t||
_q~qJ|S )Nc                 S  s   | j S r@   r   )rg   r/   r/   r0   <lambda>  r>   z7InterfaceBase._make_artifact_manifest.<locals>.<lambda>)rV   )r(   r   r   storage_policyr   r   MANIFEST_FILE_SIZE_THRESHOLD_write_artifact_manifest_filemanifest_file_pathrf   r^   storage_policy_configr`   rV   r   r   ra   sortedentriesvaluescontentsr   r   sizebirth_artifact_idref
local_path
skip_cacheextra)r8   r   rZ   proto_manifestr   rg   rh   r   entryproto_entryproto_extrar/   r/   r0   r   |  s<   




z%InterfaceBase._make_artifact_manifestr   c                 C  s   ddl m} t| d }|jddd t  dtd d}|| }tj|d	d
d }|j	 D ]}|
t|  d q2W d    t|S 1 sOw   Y  t|S )Nr   get_staging_dirartifact_manifestsT)parentsexist_okr      z.manifest_contents.jl.gzwtr   )modecompresslevel
)wandb.sdk.artifacts.stagingr  r   mkdirtimer   gzipopenr  r  writer   r   to_jsonr   )r8   r   r  manifest_dirfilenamer   r   r  r/   r/   r0   r     s   
z+InterfaceBase._write_artifact_manifest_fileportfolio_namealiasesIterable[str]rn   
str | Nonero   organizationc                 C  sf   t  }| r|j|_n	|jr|jnd|_||_|pd|_|p!d|_	|p&d|_
|j| | |S N )r(   LinkArtifactRequestis_draftr   r   id	server_idr#  portfolio_entityportfolio_organizationportfolio_projectportfolio_aliasesrz   _deliver_link_artifact)r8   r   r#  r$  rn   ro   r'  link_artifactr/   r/   r0   deliver_link_artifact  s   	




z#InterfaceBase.deliver_link_artifactr3  pb.LinkArtifactRequestc                 C  rH   r@   rI   )r8   r3  r/   r/   r0   r2    r   z$InterfaceBase._deliver_link_artifactsourcejob_infodict[str, Any]r   c                 C  sF  | dd }|dkrD| di }| dd| j_| jj| dg  | dd| j_| d}|r8|| j_| d	}|rC|| j_nY|d
kr| di  dd| jj	_
| di  dd| jj	_| jj| dg  | dd| j_| d}|r|| j_| d	}|r|| j_n|dkr| dd| j_ntd|  }|S )z9Construct use_artifact.partial.source_info.source as str.source_typer)  r   r6  
entrypointnotebookFbuild_context
dockerfilerepor   remoter   imagedockerzInvalid source type)getstripr   r:  rz   r;  r<  r=  r   git_infor?  r   r@  
ValueErrorSerializeToString)r6  r7  r   r9  info_sourcer<  r=  
source_strr/   r/   r0   _make_partial_source_str  s<   



z&InterfaceBase._make_partial_source_struse_artifactpb.UseArtifactRecordjob_namec                 C  sh   ||j _|dd|j j_|dd|j j_|dd|j j_| j|j jj||d}|j jj	| |S )N_versionr)  r9  r   )r6  r7  r   )
partialrL  rB  source_inforM  r9  r   rI  r6  ParseFromString)r8   rJ  rL  r7  r   src_strr/   r/   r0   _make_proto_use_artifact  s   z&InterfaceBase._make_proto_use_artifactc              
   C  s:  |j d us	J dtj|j |j|jd}d|jv ri }z"|d }t|}t	
|}W d    n1 s7w   Y  W n( tye } ztd| d|  td| d|  W Y d }~d S d }~ww z| j||j||jd}W n" ty } ztd|  td|  W Y d }~d S d }~ww | | d S )	NzArtifact must have an id)r,  r   r   _partialzwandb-job.jsonz2Failed to download partial job info from artifact z, : )rJ  rL  r7  r   z(Failed to construct use artifact proto: )r,  r(   UseArtifactRecordr   r   r   	get_entrydownloadr  r   load	Exceptionloggerwarningr   rR  _publish_use_artifact)r8   r   rJ  r7  r   r   er/   r/   r0   publish_use_artifact
  sL   

	
z"InterfaceBase.publish_use_artifactr   c                 C  rH   r@   rI   r8   r   r/   r/   r0   r[  6  rK   z#InterfaceBase._publish_use_artifactFTry   Iterable[str] | Nonehistory_step
int | Noneis_user_createduse_after_commitfinalizec	                 C  s   ddl m}	 | |}
| |}|
j|_|
j|_|
j|_||_||_||_	|j
|p+g  |j|p3g  t }|j| |d urF||_|	 |_| |}|S )Nr   r  )r  r  r   r   rs   ro   rn   user_createdrc  rd  r$  rz   ry   r(   LogArtifactRequestr   CopyFromr`  staging_dir_deliver_artifact)r8   rj   r   r$  ry   r`  rb  rc  rd  r  r   r   log_artifactrespr/   r/   r0   deliver_artifact:  s$   


zInterfaceBase.deliver_artifactrj  pb.LogArtifactRequestc                 C  rH   r@   rI   )r8   rj  r/   r/   r0   ri  [  rS   zInterfaceBase._deliver_artifactartifact_iddownload_rootallow_missing_referencesr
  path_prefixc                 C  s8   t  }||_||_||_||_|pd|_| |}|S r(  )r(   DownloadArtifactRequestrn  ro  rp  r
  rq  _deliver_download_artifact)r8   rn  ro  rp  r
  rq  download_artifactrk  r/   r/   r0   deliver_download_artifactb  s   

z'InterfaceBase.deliver_download_artifactrt  pb.DownloadArtifactRequestc                 C  rH   r@   rI   )r8   rt  r/   r/   r0   rs  s  r   z(InterfaceBase._deliver_download_artifactc           
      C  sl   |  |}| |}	|j|	_|j|	_|j|	_||	_||	_||	_|	j	|p%g  |	j
	|p-g  | |	 d S r@   )r   r   rs   ro   rn   re  rc  rd  r$  rz   ry   _publish_artifact)
r8   rj   r   r$  ry   rb  rc  rd  r   r   r/   r/   r0   publish_artifacty  s   


zInterfaceBase.publish_artifactc                 C  rH   r@   rI   r^  r/   r/   r0   rw    rK   zInterfaceBase._publish_artifactr)  log_dirsaveroot_logdirc                 C  s(   t  }||_||_||_| | d S r@   )r(   TBRecordry  rz  root_dir_publish_tbdata)r8   ry  rz  r{  tbrecordr/   r/   r0   publish_tbdata  s
   zInterfaceBase.publish_tbdatar  pb.TBRecordc                 C  rH   r@   rI   )r8   r  r/   r/   r0   r~    rK   zInterfaceBase._publish_tbdatatelemtpb.TelemetryRecordc                 C  rH   r@   rI   )r8   r  r/   r/   r0   _publish_telemetry  rK   z InterfaceBase._publish_telemetryenvironmentpb.EnvironmentRecordc                 C  s   |  | d S r@   )_publish_environmentr8   r  r/   r/   r0   publish_environment  rO   z!InterfaceBase.publish_environmentc                 C  rH   r@   rI   r  r/   r/   r0   r    rK   z"InterfaceBase._publish_environment	user_stepr~   stepflushbool | Nonepublish_stepc                 C  s   t |||dd}|dd  d|vrt |d< t }| D ]\}}	|j }
||
_t	|	|
_
q |r<|d ur<||j_|d urD||j_| | d S )NT)r  ignore_copy_err_step
_timestamp)r   popr  r(   PartialHistoryRequestr^   r   r`   rV   r   ra   r  numactionr  _publish_partial_history)r8   rj   rT   r  r  r  r  partial_historyrg   rh   r   r/   r/   r0   publish_partial_history  s   	
z%InterfaceBase.publish_partial_historyhistorypb.PartialHistoryRequestc                 C  rH   r@   rI   r8   r  r/   r/   r0   r    rK   z&InterfaceBase._publish_partial_historyc           	      C  st   t |||d}t }|r|d usJ ||j_|dd  | D ]\}}|j }||_	t
||_q!| | d S )N)r  r  )r   r(   HistoryRecordr  r  r  r^   r   r`   rV   r   ra   _publish_history)	r8   rj   rT   r  r  r  rg   rh   r   r/   r/   r0   publish_history  s   
zInterfaceBase.publish_historypb.HistoryRecordc                 C  rH   r@   rI   r  r/   r/   r0   r    rK   zInterfaceBase._publish_historyc                 C  r?   r@   )r(   RunPreemptingRecord_publish_preemptingr8   preempt_recr/   r/   r0   publish_preempting  rF   z InterfaceBase.publish_preemptingr  pb.RunPreemptingRecordc                 C  rH   r@   rI   r  r/   r/   r0   r    rK   z!InterfaceBase._publish_preemptingnowaitr   r  c                C  Z   |dkr
t jjj}n|dkrt jjj}ntd t j||d}|j  | j||d d S Nstdoutstderrzunknown type)output_typeliner  )	r(   OutputRecord
OutputTypeSTDOUTSTDERRr   	timestampGetCurrentTime_publish_outputr8   r   rT   r  otypeor/   r/   r0   publish_output     
zInterfaceBase.publish_outputoutdatapb.OutputRecordc                C  rH   r@   rI   r8   r  r  r/   r/   r0   r    rK   zInterfaceBase._publish_outputc                C  r  r  )	r(   OutputRawRecordr  r  r  r   r  r  _publish_output_rawr  r/   r/   r0   publish_output_raw  r  z InterfaceBase.publish_output_rawpb.OutputRawRecordc                C  rH   r@   rI   r  r/   r/   r0   r    s   z!InterfaceBase._publish_output_rawc                 C  r?   r@   )r(   PauseRequest_publish_pauser8   pauser/   r/   r0   publish_pause"  rF   zInterfaceBase.publish_pauser  pb.PauseRequestc                 C  rH   r@   rI   r  r/   r/   r0   r  &  rK   zInterfaceBase._publish_pausec                 C  r?   r@   )r(   ResumeRequest_publish_resumer8   resumer/   r/   r0   publish_resume*  rF   zInterfaceBase.publish_resumer  pb.ResumeRequestc                 C  rH   r@   rI   r  r/   r/   r0   r  .  rK   zInterfaceBase._publish_resumetitletextlevelwait_durationc                 C  s.   t  }||_||_||_||_| | d S r@   )r(   AlertRecordr  r  r  r  _publish_alert)r8   r  r  r  r  proto_alertr/   r/   r0   publish_alert2  s   zInterfaceBase.publish_alertalertpb.AlertRecordc                 C  rH   r@   rI   )r8   r  r/   r/   r0   r  <  rK   zInterfaceBase._publish_alert	exit_codepb.RunExitRecordc                 C  s   t  }|d ur||_|S r@   )r(   RunExitRecordr  )r8   r  exitr/   r/   r0   
_make_exit@  s   zInterfaceBase._make_exitc                 C  r   r@   )r  _publish_exitr8   r  	exit_datar/   r/   r0   publish_exitF  r   zInterfaceBase.publish_exitr  c                 C  rH   r@   rI   r8   r  r/   r/   r0   r  J  rK   zInterfaceBase._publish_exitc                 C  r?   r@   )r(   KeepaliveRequest_publish_keepaliver8   	keepaliver/   r/   r0   publish_keepaliveN  rF   zInterfaceBase.publish_keepaliver  pb.KeepaliveRequestc                 C  rH   r@   rI   r  r/   r/   r0   r  R  rK   z InterfaceBase._publish_keepaliveinclude_pathslist[list[str]]exclude_pathsinput_schema
run_config	file_pathc           
      C  s   |r|rt dt }dd |D }dd |D }|j| |j| tjtj d}	|r;|	j	tj  n|	j
	tjj|d |j	|	 |rSt||_| |S )a  Publishes a request to add inputs to the job.

        If run_config is True, the wandb.config will be added as a job input.
        If file_path is provided, the file at file_path will be added as a job
        input.

        The paths provided as arguments are sequences of dictionary keys that
        specify a path within the wandb.config. If a path is included, the
        corresponding field will be treated as a job input. If a path is
        excluded, the corresponding field will not be treated as a job input.

        Args:
            include_paths: paths within config to include as job inputs.
            exclude_paths: paths within config to exclude as job inputs.
            input_schema: A JSON Schema describing which attributes will be
                editable from the Launch drawer.
            run_config: bool indicating whether wandb.config is the input source.
            file_path: path to file to include as a job input.
        z:run_config and file_path are mutually exclusive arguments.c                 S     g | ]}t j|d qS r   r(   JobInputPath.0r   r/   r/   r0   
<listcomp>v      z3InterfaceBase.publish_job_input.<locals>.<listcomp>c                 S  r  r  r  r  r/   r/   r0   r  w  r  )r  r   )rE  r(   JobInputRequestr  rz   r  JobInputSourceRunConfigSourcer  rg  fileConfigFileSourceinput_sourcer   r  _publish_job_input)
r8   r  r  r  r  r  requestinclude_recordsexclude_recordsr6  r/   r/   r0   publish_job_inputV  s*   

zInterfaceBase.publish_job_inputr  pb.JobInputRequestc                 C  rH   r@   rI   )r8   r  r/   r/   r0   r    r   z InterfaceBase._publish_job_inputc                 C     t  }| |S r@   )r(   ProbeSystemInfoRequest_publish_probe_system_infor8   probe_system_infor/   r/   r0   publish_probe_system_info     
z'InterfaceBase.publish_probe_system_infor  pb.ProbeSystemInfoRequestc                 C  rH   r@   rI   r  r/   r/   r0   r    r   z(InterfaceBase._publish_probe_system_infoc                 C  s`   | j rd S |  }z	|jdd W d S  ty!   td Y d S  ty/   td Y d S w )N   )timeoutz timed out communicating shutdownz-handle abandoned while communicating shutdown)r5   _deliver_shutdownwait_orTimeoutErrorrY  rZ  r   )r8   handler/   r/   r0   r     s   zInterfaceBase.joinc                 C  rH   r@   rI   r7   r/   r/   r0   r    rK   zInterfaceBase._deliver_shutdownc                 C     |  |}| |S r@   )r   _deliver_runr   r/   r/   r0   deliver_run     

zInterfaceBase.deliver_runc                 C  r  r@   )r(   SyncFinishRequest_deliver_finish_syncr8   syncr/   r/   r0   deliver_finish_sync  s   
z!InterfaceBase.deliver_finish_syncr  pb.SyncFinishRequestc                 C  rH   r@   rI   r  r/   r/   r0   r    r   z"InterfaceBase._deliver_finish_syncc                 C  rH   r@   rI   r   r/   r/   r0   r    rK   zInterfaceBase._deliver_runc                 C  s   t j| |d}| |S r   )r(   RunStartRequestr   _deliver_run_start)r8   rj   	run_startr/   r/   r0   deliver_run_start  s   
zInterfaceBase.deliver_run_startr  pb.RunStartRequestc                 C  rH   r@   rI   )r8   r  r/   r/   r0   r    r   z InterfaceBase._deliver_run_start	attach_idc                 C  s   t j|d}| |S )N)r  )r(   AttachRequest_deliver_attach)r8   r  attachr/   r/   r0   deliver_attach  s   
zInterfaceBase.deliver_attachpb.AttachRequestc                 C  rH   r@   rI   rR   r/   r/   r0   r     rS   zInterfaceBase._deliver_attachc                 C  r  r@   )r(   StopStatusRequest_deliver_stop_statusrR   r/   r/   r0   deliver_stop_status  r  z!InterfaceBase.deliver_stop_statuspb.StopStatusRequestc                 C  rH   r@   rI   rR   r/   r/   r0   r%    rS   z"InterfaceBase._deliver_stop_statusc                 C  r  r@   )r(   NetworkStatusRequest_deliver_network_statusrR   r/   r/   r0   deliver_network_status  r  z$InterfaceBase.deliver_network_statuspb.NetworkStatusRequestc                 C  rH   r@   rI   rR   r/   r/   r0   r)    rS   z%InterfaceBase._deliver_network_statusc                 C  r  r@   )r(   InternalMessagesRequest_deliver_internal_messagesr8   internal_messager/   r/   r0   deliver_internal_messages  r  z'InterfaceBase.deliver_internal_messagesr/  pb.InternalMessagesRequestc                 C  rH   r@   rI   r.  r/   r/   r0   r-    r   z(InterfaceBase._deliver_internal_messagesc                 C  r  r@   )r(   GetSummaryRequest_deliver_get_summaryr8   get_summaryr/   r/   r0   deliver_get_summary  r  z!InterfaceBase.deliver_get_summaryr5  pb.GetSummaryRequestc                 C  rH   r@   rI   r4  r/   r/   r0   r3    rS   z"InterfaceBase._deliver_get_summaryc                 C  r  r@   )r(   GetSystemMetricsRequest_deliver_get_system_metrics)r8   get_system_metricsr/   r/   r0   deliver_get_system_metrics  r  z(InterfaceBase.deliver_get_system_metricspb.GetSystemMetricsRequestc                 C  rH   r@   rI   r4  r/   r/   r0   r9    r   z)InterfaceBase._deliver_get_system_metricsc                 C  r  r@   )r  _deliver_exitr  r/   r/   r0   deliver_exit  r  zInterfaceBase.deliver_exitc                 C  rH   r@   rI   r  r/   r/   r0   r=    rS   zInterfaceBase._deliver_exitc                 C  r  r@   )r(   PollExitRequest_deliver_poll_exitr8   	poll_exitr/   r/   r0   deliver_poll_exit  r  zInterfaceBase.deliver_poll_exitrB  pb.PollExitRequestc                 C  rH   r@   rI   rA  r/   r/   r0   r@    rS   z InterfaceBase._deliver_poll_exitc                 C  r  r@   )r(   RunFinishWithoutExitRequest_deliver_finish_without_exitr8   run_finish_without_exitr/   r/   r0   deliver_finish_without_exit#  r  z)InterfaceBase.deliver_finish_without_exitrH  pb.RunFinishWithoutExitRequestc                 C  rH   r@   rI   rG  r/   r/   r0   rF  '  r   z*InterfaceBase._deliver_finish_without_exitc                 C  r  r@   )r(   SampledHistoryRequest _deliver_request_sampled_historyr8   sampled_historyr/   r/   r0   deliver_request_sampled_history-  r  z-InterfaceBase.deliver_request_sampled_historyrN  pb.SampledHistoryRequestc                 C  rH   r@   rI   rM  r/   r/   r0   rL  1  r   z.InterfaceBase._deliver_request_sampled_historyc                 C  r  r@   )r(   RunStatusRequest_deliver_request_run_statusr8   
run_statusr/   r/   r0   deliver_request_run_status7  r  z(InterfaceBase.deliver_request_run_statusrT  pb.RunStatusRequestc                 C  rH   r@   rI   rS  r/   r/   r0   rR  ;  r   z)InterfaceBase._deliver_request_run_status)r!   r6   )r:   r;   r!   r<   )rD   rG   r!   r6   )r!   r<   )rP   rQ   r!   r<   )NNNN)
rT   rU   rV   rW   rX   rY   rZ   r[   r!   r\   )rj   r   r!   rk   )rj   r   r!   r6   )rj   rk   r!   r6   )r   r   r!   r6   )r   r   r!   r6   )NNN)rT   rU   rV   rW   rX   rY   r!   r6   )r   r\   r!   r6   )r   r   r!   r6   )r   r   r!   r   )r   r   r   r   rj   r   r!   r   )r   r   rj   r   r!   r   )rj   r   r   r   r!   r6   )r   r   r!   r6   )r   r   r!   r   )r   r   r!   r6   )r   r   r!   r6   )r   r   r!   r6   )r   r   r!   r   r@   )r   r   rZ   r   r!   r   )r   r   r!   r   )r   r   r#  r   r$  r%  rn   r&  ro   r&  r'  r&  r!   r<   )r3  r5  r!   r<   )r6  r   r7  r8  r   r8  r!   r   )
rJ  rK  rL  r   r7  r8  r   r8  r!   rK  )r   r   r!   r6   )r   rK  r!   r6   )NNFFT)rj   r   r   r   r$  r%  ry   r_  r`  ra  rb  r4   rc  r4   rd  r4   r!   r<   )rj  rm  r!   r<   )rn  r   ro  r   rp  r4   r
  r4   rq  r&  r!   r<   )rt  rv  r!   r<   )NFFT)rj   r   r   r   r$  r%  ry   r_  rb  r4   rc  r4   rd  r4   r!   r6   )r   r   r!   r6   )r)  )ry  r   rz  r4   r{  r   r!   r6   )r  r  r!   r6   )r  r  r!   r6   )r  r  r!   r6   )NNT)rj   r   rT   r   r  r~   r  ra  r  r  r  r4   r!   r6   )r  r  r!   r6   )NT)
rj   r   rT   r   r  ra  r  r4   r!   r6   )r  r  r!   r6   )r  r  r!   r6   )r   r   rT   r   r  r4   r!   r6   )r  r  r  r4   r!   r6   )r  r  r  r4   r!   r6   )r  r  r!   r6   )r  r  r!   r6   )
r  r   r  r   r  r   r  r~   r!   r6   )r  r  r!   r6   )r  ra  r!   r  )r  ra  r!   r6   )r  r  r!   r6   )r  r  r!   r6   )Fr)  )
r  r  r  r  r  rU   r  r4   r  r   )r  r   r!   r<   )r  r  r!   r6   )rj   r   r!   r<   )r  r  r!   r<   )rj   rk   r!   r<   )r  r  r!   r<   )r  r   r!   r<   )rP   r#  r!   r<   )rP   r'  r!   r<   )rP   r+  r!   r<   )r/  r1  r!   r<   )r5  r7  r!   r<   )r5  r<  r!   r<   )r  ra  r!   r<   )r  r  r!   r<   )rB  rD  r!   r<   )rH  rJ  r!   r<   )rN  rP  r!   r<   )rT  rV  r!   r<   )l__name__
__module____qualname____doc____annotations__r9   abcabstractmethodr=   rE   rB   rN   rL   ri   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r4  r2  staticmethodrI  rR  r]  r[  rl  ri  ru  rs  rx  rw  r  r~  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r   r  r  r  r  r  r  r  r"  r   r&  r%  r*  r)  r0  r-  r6  r3  r;  r9  r>  r=  rC  r@  rI  rF  rO  rL  rU  rR  r/   r/   r/   r0   r3   ?   sd  
 5
*,+#,!

3r3   )r    r   r!   r"   )r.   r"   r!   r   ):
__future__r   r\  r  loggingr  collections.abcr   pathlibr   secretsr   typingr   r   r   r   wandb.protor	   r(   r
   tpbwandb.sdk.libr   r   wandb.sdk.lib.filesystemr   r   wandb.sdk.mailboxr   r   
wandb.utilr   r   r   r   r   r   r   data_types.utilsr   r   r)  r   srr   wandb.sdk.artifacts.artifactr   %wandb.sdk.artifacts.artifact_manifestr   	wandb_runr   	getLoggerrY  r1   r2   ABCr3   r/   r/   r/   r0   <module>   s6    $





