o bi@s@UddlZddlZddlZddlZddlZddlmZddlmZddlm Z m Z ddl m Z m Z mZmZmZmZmZmZddlmZddlZddlZddlmZddlmZdd lmZmZm Z m!Z!m"Z"dd l#m$Z$m%Z%dd l&m'Z'dd l(m)Z)dd l*m+Z+ddl,m-Z-ddl.m/Z/m0Z0m1Z1m2Z2ddl3m4Z4e5e6Z7dZ8dZ9dZ:e e;ee)fZd3de?de?de;fddZ@GdddZAGdd d ZBejCdd!Gd"d#d#ZDeEZFejEeGd$<d%d&ZHGd'd(d(ZIeIZJGd)d*d*ejKZLGd+d,d,ZMe-e Gd-d.d.ZNe Gd/d0d0ZOe Gd1d2d2ZPdS)4N) defaultdict)contextmanager) dataclassfields)AnyDictListMappingOptionalSetTupleUnion)uuid4) ActorHandle) BlockList) NODE_UNKNOWN MetricsGroup MetricsType NodeMetricsOpRuntimeMetrics)Topologyget_dataset_metadata_exporter)capfirst) BlockStats) DataContext) DeveloperAPI)CounterGauge HistogramMetric)NodeAffinitySchedulingStrategydatasets_stats_actor_dataset_stats_actorunknownsecondsreturncCsR|dkr tt|ddS|dkrtt|dddStt|ddddS)NsgMbP?msus)strround)r$r.L/home/ubuntu/.local/lib/python3.10/site-packages/ray/data/_internal/stats.pyfmt*s r0lvlspaces_per_indentcCs d||S)zReturns a string of spaces which contains `level` indents, each indent containing `spaces_per_indent` spaces. For example: >>> leveled_indent(2, 3) ' '  r.)r2r3r.r.r/leveled_indent3 r5c@speZdZdZddZedddZdeddfd d Zdefd d Z defd dZ defddZ defddZ dS)Timerz8Helper class for tracking accumulated time (in seconds).cCs d|_td|_d|_d|_dS)Nrinf)_totalfloat_min_max _total_countselfr.r.r/__init__?s  zTimer.__init__r%Nc cs>t}zdVW|t|dS|t|wN)time perf_counteradd)r? time_startr.r.r/timerEs *z Timer.timervaluecCs@|j|7_||jkr||_||jkr||_|jd7_dSNr&)r9r;r<r=)r?rGr.r.r/rDMs   z Timer.addcC|jSrA)r9r>r.r.r/getUz Timer.getcCrIrA)r;r>r.r.r/minXrKz Timer.mincCrIrA)r<r>r.r.r/max[rKz Timer.maxcCs|jr |j|jStdS)Nr8)r=r9r:r>r.r.r/avg^sz Timer.avgr%N) __name__ __module__ __qualname____doc__r@rrFr:rDrJrLrMrNr.r.r.r/r7<s r7c@sNeZdZdZdedddeefddZded dfd d Z d e d dfd dZ dS)_DatasetStatsBuilderzHelper class for building dataset stats. When this class is created, we record the start time. When build() is called with the final blocks of the new dataset, the time delta is saved as part of the stats. operator_nameparent DatasetStatsoverride_start_timecCs||_||_|p t|_dSrA)rUrVrBrC start_time)r?rUrVrXr.r.r/r@isz_DatasetStatsBuilder.__init__metadatar%cCsi}t|D].\}\}}t|}t|dkr1|dkr$|||j|<q|||jdd|<q|||j<qt||j|jd}t |j |_ |S)Nr&rz->)rZrV base_name) enumerateitemsrlenrUsplitrWrVrBrCrY time_total_s)r?rZ op_metadataikvcapped_kstatsr.r.r/build_multioperatorss  z(_DatasetStatsBuilder.build_multioperator final_blockscCs,t|j|i|jd}t|j|_|S)N)rZrV)rWrU get_metadatarVrBrCrYra)r?rirgr.r.r/builds  z_DatasetStatsBuilder.buildN) rPrQrRrSr,r r:r@ StatsDictrhrrkr.r.r.r/rTbs  rT)num_cpusc@sPeZdZdZd,ddZdedeedfdeee ffd d Z deee ffd d Z d dZ ddZ d-dedeeeeeeffdeedeeefdeeeeeeeeffff ddZd.ddZ  d/ddZdededeed ed!ef d"d#Zdedeeeffd$d%Zd-deefd&d'Z  d0ded(eed)eefd*d+ZdS)1 _StatsActoraActor holding stats for blocks created by LazyBlockList. This actor is shared across all datasets created in the same cluster. In order to cap memory usage, we set a max number of stats to keep in the actor. When this limit is exceeded, the stats will be garbage collected in FIFO order. TODO(ekl) we should consider refactoring LazyBlockList so stats can be extracted without using an out-of-band actor.r)cCstt|_i|_i|_||_g|_d|_i|_ i|_ t |_ d}t dd|d|_t dd|d|_t dd |d|_t d d |d|_t d d |d|_t dd|d|_t dd|d|_|jtj|d|_|jtj|d|_|jtj|d|_|jtj|d|_|jtj|d|_ |jtj!|d|_"|#|_$d}t dd|d|_%t dd|d|_&t dd|d|_'d}t dd|d|_(t dd|d|_)t dd d!*d"d#t+Dd$|d|_,d}t d%d&|d|_-t d'd(|d|_.t d)d*|d|_/t d+d,d!*d-d#t+Dd$|d|_0dS).Nrdatasetoperatordata_spilled_byteszBytes spilled by dataset operators. DataContext.enable_get_object_locations_for_metrics must be set to True to report this metric descriptiontag_keysdata_freed_bytesz Bytes freed by dataset operatorsdata_current_bytesz9Bytes currently in memory store used by dataset operatorsdata_cpu_usage_coresz#CPUs allocated to dataset operatorsdata_gpu_usage_coresz#GPUs allocated to dataset operatorsdata_output_bytesz$Bytes outputted by dataset operatorsdata_output_rowsz#Rows outputted by dataset operators) metrics_groupru)rpdata_iter_total_blocked_secondsz0Seconds user thread is blocked by iter_batches()data_iter_user_secondszSeconds spent in user codedata_iter_initialize_secondsz-Seconds spent in iterator initialization coderpjob_idrY#data_dataset_estimated_total_blocksz&Total work units in blocks for dataset!data_dataset_estimated_total_rowsz$Total work units in rows for datasetdata_dataset_statezState of dataset (z, cSg|] }|jd|jqS=rGname.0r(r.r.r/ .z(_StatsActor.__init__..)$data_operator_estimated_total_blocksz'Total work units in blocks for operator"data_operator_estimated_total_rowsz%Total work units in rows for operatordata_operator_queued_blocksz$Number of queued blocks for operatordata_operator_statezState of operator (cSrrrrr.r.r/rDr)1 collectionsrdictrZ last_timerY max_stats fifo_queuenext_dataset_iddatasets_ray_nodes_cacher_metadata_exporterr spilled_bytes freed_bytes current_bytescpu_usage_coresgpu_usage_cores output_bytes output_rows0_create_prometheus_metrics_for_execution_metricsrINPUTSexecution_metrics_inputsOUTPUTSexecution_metrics_outputsTASKSexecution_metrics_tasksOBJECT_STORE_MEMORY"execution_metrics_obj_store_memoryACTORSexecution_metrics_actorsMISCexecution_metrics_misc/_create_prometheus_metrics_for_per_node_metricsper_node_metricsiter_total_blocked_s iter_user_siter_initialize_srrjoin DatasetStaterrrrr)r?r op_tags_keys iter_tag_keys dataset_tags operator_tagsr.r.r/r@s         z_StatsActor.__init__r|ru.r%cCsi}tD]H}|j|ksqd|j}|j}|jtjkr(t|||d||j<q|jtjkr>t|f||d|j ||j<q|jtj krNt |||d||j<q|S)Ndata_rs) r get_metricsr|rrt metrics_typerrr metrics_argsr)r?r|rumetricsmetric metric_namemetric_descriptionr.r.r/rHs:       z<_StatsActor._create_prometheus_metrics_for_execution_metricscCs8i}ttD]}d|jd}t|ddd||j<q|S)Nr _per_node)rpnode_iprs)rrrr)r?rfieldrr.r.r/rfs z;_StatsActor._create_prometheus_metrics_for_per_node_metricscCst|j}|jd7_|SrH)r,r)r? dataset_idr.r.r/get_dataset_idqs z_StatsActor.get_dataset_idcCs,|D]}|j|q|D]}|j|q dSrA)update_execution_metricsupdate_iteration_metrics)r?execution_metricsiteration_metricsrr.r.r/update_metricsvs   z_StatsActor.update_metricsN dataset_tag op_metricsrstaterc CsZ ddtdtttfdtttffdd}t||D]\}}|||} |j | dd| |j | dd| |j | d d| |j | d d| |j | d d| |j | d d| |j | d d| |jD]\} } || | | d| qv|jD]\} } || | | d| q|jD]\} } || | | d| q|jD]\} } || | | d| q|jD]\} } || | | d| q|jD]\} } || | | d| qq|dur%|D]1\} } | |jvr||j | t}|j||d} | D]\}}|j|} || || qq|||dS)N prom_metricrGtagscSsRt|tr |||dSt|tr|||dSt|tr'|||dSdSrA) isinstancersetrincrobserve)rrGrr.r.r/_records   z5_StatsActor.update_execution_metrics.._recordobj_store_mem_spilledrobj_store_mem_freedobj_store_mem_usedbytes_task_outputs_generatedrow_outputs_taken cpu_usage gpu_usage)r node_ip_tagrA)rr intr:rr,zip _create_tagsrrrJrrrrrrrr^rrrrrr_rebuild_ray_nodes_cacherrupdate_dataset)r?rrrrrrrg operator_tagr field_namernode_id node_metricsrr metric_valuer.r.r/r|sX        z$_StatsActor.update_execution_metricscCsHt}|D]}|dd}|dd}|dur!|dur!||j|<qdS)NNodeIDNodeName)raynodesrJr)r? current_nodesnoder node_namer.r.r/rs   z$_StatsActor._rebuild_ray_nodes_cachergrWcCsJ||}|j|j||j|j||j|j|dSrA)rrrrJrr)r?rgrrr.r.r/rs z$_StatsActor.update_iteration_metricsrtopology data_contextc Csnt}|tjjddd|ddd|Dd|j|<|jdur5ddlm}||||||d}|j|dSdS)NrcSs i|] }|tjjddddqS)r)rprogresstotal queued_blocks)rRUNNINGr)rrqr.r.r/ sz0_StatsActor.register_dataset..)rrrr total_rowsrYend_time operators)DatasetMetadata)rrrrYr) rBrrrrr$ray.data._internal.metadata_exporterrexport_dataset_metadata) r?rrrrrrYrdataset_metadatar.r.r/register_datasets.  z_StatsActor.register_datasetc Cs<|j|||j|}|j|dd}tt|j|dd}|||d}|j|dd||j|dd||dtj j }t |}|j |j ||d iD]?\}} ||d } |j| dd| |j| dd| |j| d d| | dtj j }t |}|j|j | q\dS) NrNonerYrrrrrrror)rupdaterJr,rrrrrUNKNOWNr from_stringrrGr^rrrr) r?rrrrYr state_string state_enumrqop_staterr.r.r/rsD        z_StatsActor.update_datasetcs"s|jSfdd|jDS)Ncs"i|] \}}|dkr||qSrr.rrdrerr.r/r+s"z,_StatsActor.get_datasets..)rr^)r?rr.rr/ get_datasets(sz_StatsActor.get_datasetsrrcCs,d|i}|dur ||d<|dur||d<|S)Nrprqrr.)r?rrrrr.r.r/r-s z_StatsActor._create_tags)r)rArO)rgrW)NN)rPrQrRrSr@rr r,rrrrrrrrr rr:rr rrrrrrrr rr.r.r.r/rnsh .       G   ',rn_stats_actor_lockcCspt}|j}tjjjstt dd}t t j t tdd|dWdS1s1wYdS)NF)softTdetached)r namespace get_if_existslifetimescheduling_strategy)r get_currentrrutilclient is_connectedr get_runtime_context get_node_idr rnoptionsSTATS_ACTOR_NAMESTATS_ACTOR_NAMESPACEremote)ctxrr.r.r/_get_or_create_stats_actor@s" $rc @seZdZdZdZdZddZd$deefddZ d d Z d e e dee ee eeeeffffd d Z d%ded e e de edeeefdef ddZdefddZdddefddZdefddZdede ededefdd Zdefd!d"Zd#S)& _StatsManagerafA Class containing util functions that manage remote calls to _StatsActor. This class collects stats from execution and iteration codepaths and keeps track of the latest snapshot. An instance of this class runs a single background thread that periodically forwards the latest execution/iteration stats to the _StatsActor. This thread will terminate itself after being inactive (meaning that there are no active executors or iterators) for STATS_ACTOR_UPDATE_THREAD_INACTIVITY_LIMIT iterations. After terminating, a new thread will start if more calls are made to this class. cCs6d|_d|_i|_i|_t|_d|_t|_dSrA) _stats_actor_handle_stats_actor_cluster_id_last_execution_stats_last_iteration_stats threadingLock _stats_lock_update_thread_update_thread_lockr>r.r.r/r@js z_StatsManager.__init__Tr%cCsztjjjdur tdtjjjj}|jdus|j|kr:|r"t|_nz tj t t d|_Wn t y6YdSw||_|jS)NzGlobal node is not initialized.)rr ) r_privateworker _global_node RuntimeError cluster_idrr r get_actorrr ValueError)r?create_if_not_existscurrent_cluster_idr.r.r/ _stats_actor|s      z_StatsManager._stats_actorcsvj.jdusjs)fdd}tj|dd_jWddSWddS1s4wYdS)Ncsd} js jr=z"jdd}|durWq|jjtjtjdd}Wn!ty<tj dddYdSw|d7}|t j krMt d dSt tjq) NrTF)r/)rrz1Error occurred during remote call to _StatsActor.)exc_infor&z2Terminating StatsManager thread due to inactivity.)r"r!r1rrlistvalues ExceptionloggerdebugrUPDATE_THREAD_INACTIVITY_LIMITrBsleep StatsManager#STATS_ACTOR_UPDATE_INTERVAL_SECONDS)iter_stats_inactivity stats_actorr>r.r/_run_update_loopsF   zD_StatsManager._start_thread_if_not_running.._run_update_loopT)targetdaemon)r'r&is_aliver#Threadstart)r?r>r.r>r/_start_thread_if_not_runnings ) "z*_StatsManager._start_thread_if_not_runningrc CsltjsdStdd}|D]$}|jD]\}}||}ttD]}||jt ||j7<q"qq|S)aR Aggregate per-node metrics from a list of OpRuntimeMetrics objects. If per-node metrics are disabled in the current DataContext, returns None. Otherwise, it sums up all NodeMetrics fields across the provided metrics and returns a nested dictionary mapping each node ID to a dict of field values. NcSsttSrA)rrr.r.r.r/sz;_StatsManager._aggregate_per_node_metrics..) rrenable_per_node_metricsr_per_node_metricsr^rrrgetattr)r?raggregated_by_noderrragg_node_metricsfr.r.r/_aggregate_per_node_metricss   z)_StatsManager._aggregate_per_node_metricsFrrr force_updatec Cszdd|D}||}|||||f}|r|jj|dS|j ||j|<Wdn1s2wY|dS)NcSg|]}|qSr.)as_dict)rrr.r.r/rz:_StatsManager.update_execution_metrics..)rLr1rrr%r!rD) r?rrrrrMop_metrics_dictsrargsr.r.r/rs   z&_StatsManager.update_execution_metricscCN|j||jvr|j|=WddSWddS1s wYdSrA)r%r!r?rr.r.r/clear_last_execution_statss   "z(_StatsManager.clear_last_execution_statsrgrWcCs@|j||f|j|<Wdn1swY|dSrA)r%r"rD)r?rgrr.r.r/rs z&_StatsManager.update_iteration_metricscCrSrA)r%r"rTr.r.r/clear_iteration_metricss   "z%_StatsManager.clear_iteration_metricsrrcCs$|jt||||dS)a1Register a dataset with the stats actor. Args: dataset_tag: Tag for the dataset operator_tags: List of operator tags topology: Optional Topology representing the DAG structure to export data_context: The DataContext attached to the dataset N)r1rrrr get_job_id)r?rrrrr.r.r/register_dataset_to_stats_actors  z-_StatsManager.register_dataset_to_stats_actorcCs2z t|jWStytjYSwrA)rrJr1rrr5rhexr>r.r.r/get_dataset_id_from_stats_actors   z-_StatsManager.get_dataset_id_from_stats_actorN)T)F)rPrQrRrSr;r8r@r rr1rDrrr r,r rr:rLrrboolrrUrrVrrrXrZr.r.r.r/rTsJ5      rc@s4eZdZdZdZdZdZdZddZe dd Z d S) rzDEnum representing the possible states of a dataset during execution.rr&r'r1cCrIrA)rr>r.r.r/__str__2rKzDatasetState.__str__cCs$z||WSty|jYSw)zGet enum by name.)KeyErrorr)clstextr.r.r/r5s    zDatasetState.from_stringN) rPrQrRrSrrFINISHEDFAILEDr\ classmethodrr.r.r.r/r*src@s~eZdZdZdddedeededfdefddZ e d d Z dd ed ee d e fddZdddZd efddZdS)rWzHolds the execution times for a given Dataset. This object contains a reference to the parent Dataset's stats as well, but not the Dataset object itself, to allow its blocks to be dropped from memory.N)r\rZrVr\cCs||_|durt|ts|g}|pg|_|jsdn tdd|jDd|_||_d|_d|_t |_ t |_ t |_ t |_ t |_t |_t |_t |_t |_t |_t |_i|_d|_d|_d|_d|_d|_d|_t |_dS)aCreate dataset stats. Args: metadata: Dict of operators used to create this Dataset from the previous one. Typically one entry, e.g., {"map": [...]}. parent: Reference to parent Dataset's stats, or a list of parents if there are multiple. base_name: The name of the base operation for a multi-operator operation. Nrcs|]}|jVqdSrA)numberrpr.r.r/ [z(DatasetStats.__init__..r& unknown_uuid)rZrr3parentsrMrdr\ dataset_uuidrar7streaming_exec_schedule_s iter_wait_s iter_get_siter_next_batch_siter_format_batch_siter_collate_batch_siter_finalize_batch_srrr iter_total_s extra_metricsiter_blocks_localiter_blocks_remoteiter_unknown_locationglobal_bytes_spilledglobal_bytes_restoreddataset_bytes_spilledstreaming_split_coordinator_s)r?rZrVr\r.r.r/r@Es8   zDatasetStats.__init__cCstSrA)rr>r.r.r/r=szDatasetStats.stats_actorrrXr%cCs t|||S)z>Start recording stats for an op of the given name (e.g., map).)rT)r?rrXr.r.r/ child_builders zDatasetStats.child_builderDatasetStatsSummarycCsg}t|jdk}|jD]\}}|tj|||dqt|j|j|j |j |j |j |j |j|j|j|j|j|j|j}g}|jdurLdd|jD}|jrT|jnd}t||||j|j|j|j|j|j|j|j | S)zGenerate a `DatasetStatsSummary` object from the given `DatasetStats` object, which can be used to generate a summary string.r&)is_sub_operatorNcSrNr.) to_summaryrer.r.r/rrPz+DatasetStats.to_summary..r)!r_rZr^appendOperatorStatsSummaryfrom_block_metadataIterStatsSummaryrmrnrorprqrrrrrrsr{rurvrwrjrlrJr}rdrkrar\rtrxryrz)r?operators_statsr~rrg iter_statsstats_summary_parentsrlr.r.r/rs^  zDatasetStats.to_summarycCs |S)axGenerate a string representing the runtime metrics of a Dataset. This is a high level summary of the time spent in Ray Data code broken down by operator. It also includes the time spent in the scheduler. Times are shown as the total time for each operator and percentages of time are shown as a fraction of the total time for the whole dataset.)rruntime_metricsr>r.r.r/rr6zDatasetStats.runtime_metricsrA)r%r})rPrQrRrSrlr r rr,r@propertyr=r:rTr|rrr.r.r.r/rW>s,  @  5rWc@s2eZdZUeded<ded<eded<eed<eed<eed<eed <eee fed <eed <eed <eed <eed<   d)de e ede defddZ edddedfddZedddeeeffddZdefddZd*defdd Zdefd!d"Zdefd#d$Zdefd%d&Zdefd'd(ZdS)+r}rrrrrjrdrkrar\rtrxryrzrlNTalready_printedinclude_parentr%cCs|durt}d}|jr%|r%|jD]}|j|dd}|r$||7}|d7}qd}t|jdkrY|jd}|j}|j|} |d|j|7}| |vrM|d 7}n\| | |t |7}nPt|jdkrt |j d } | dkrld} |d |j|j | 7}t|jD],\} }|j}|j|} |d7}|d | |7}| |vr|d 7}q|| | |t |7}q|tj} | r|jr|r|jrdnd} || 7}|dt |jd7}|t |j7}t|jdkrS|rSt |jd}t |jd}|s|r|d7}|d|7}|d|7}t |jd}|r|d7}|d|7}|jdj}|r$|dnd}|}|}|rS|rS|rS|d7}|d7}|d||d7}|d||d7}| ra|ra|d|7}|S)aReturn a human-readable summary of this Dataset's stats. Args: already_printed: Set of operator IDs that have already had its stats printed out. include_parent: If true, also include parent stats summary; otherwise, only log stats of the latest operator. add_global_stats: If true, includes global stats to this summary. Returns: String with summary statistics for executing the Dataset. NrF)add_global_stats r&rzOperator {} {}: z[execution cached] r'z Operator {} {}: executed in {}s z Suboperator {} {}: z [execution cached]  z* Extra metrics: .Az Cluster memory: z* Spilled to disk: {}MB z* Restored from disk: {}MB z Dataset memory: r[sumzDataset throughput:  * Ray Data throughput:  rows/s % * Estimated single node throughput: )rrj to_stringr_rrUrkformatrdrDr,r-rar\r]rrverbose_stats_logsrtr~rrxryrzoutput_num_rowsget_total_wall_timeget_total_time_all_blocksr)r?rrroutrf parent_sumoperators_stats_summaryrU operator_uuid rounded_totalnrindent mb_spilled mb_restoreddataset_mb_spilledrtotal_num_out_rows wall_timetotal_time_all_blocksr.r.r/rs                zDatasetStatsSummary.to_stringcurrcCs4g}|jD]}|r|jr|t|q||gSrA)rjextendr} _collect_dataset_stats_summaries)rsummsrfr.r.r/rQs    z4DatasetStatsSummary._collect_dataset_stats_summariessummcCs0tdd|jD}tdd|jD}||fS)NcsrcrA)earliest_start_timeropsr.r.r/rg^rhz:DatasetStatsSummary._find_start_and_end..csrcrA)latest_end_timerr.r.r/rg_rh)rLrrM)rearliest_start latest_endr.r.r/_find_start_and_end\sz'DatasetStatsSummary._find_start_and_endcs|dtdtdtffdd }t|}d}|D]}t|jdkr6t|\}}||}|||j|7}q||d|j 7}||d 7}|S) NrrBr%cs6dkr|nd}d|dt|d|dddS)Nrz* : z (dz.3fz%) r0)rrBfractiontotal_wall_timer.r/fmt_linees"z5DatasetStatsSummary.runtime_metrics..fmt_linezRuntime Metrics: r SchedulingTotal) rr,r:r}rr_rrr\rl)r?r summariesrrrr op_total_timer.rr/rbs z#DatasetStatsSummary.runtime_metricsrcst}dfdd|jD}dfdd|jD}dfdd|jD}|r8d|d|dnd }|rEd|d|dnd }|rRd|d|dnd }d g|d |d |jd|d |jd|d |jd|d|d|d|d|d|j dd|d|j dd|d|j dd|d|j dd|d|d|dS)Nrcg|] }|dqSr'__repr__rsslevelr.r/ryz0DatasetStatsSummary.__repr__..crrr)rpsrr.r/r{rc3s0|]\}}td|d|dVqdS)r'r,N)r5rrr.r/rg|s  z/DatasetStatsSummary.__repr__.., z rzDatasetStatsSummary( z dataset_uuid=z base_name=z number=z extra_metrics={z}, z operators_stats=[z], z iter_stats=r&z global_bytes_spilled=rzMB, z global_bytes_restored=z dataset_bytes_spilled=z parents=[r)r5rrrjrtr^rkr\rdrrrxryrz)r?rrr parent_statsrtr.rr/rvsp     zDatasetStatsSummary.__repr__cCsPddt|D}t|dkrdStdd|D}tdd|D}||S)zCalculate the total wall time for the dataset, this is done by finding the earliest start time and latest end time for any block in any operator. The wall time is the difference of these two times. cSs$g|]}t|jdkrt|qSr)r_rr}rrrr.r.r/rs z;DatasetStatsSummary.get_total_wall_time..rcs|]}|dVqdS)rNr.r start_endr.r.r/rgz:DatasetStatsSummary.get_total_wall_time..csr)r&Nr.rr.r.r/rgr)r}rr_rLrM)r? start_endsrrr.r.r/rs z'DatasetStatsSummary.get_total_wall_timecCst|}tdd|DS)zGCalculate the sum of the wall times across all blocks of all operators.css$|] }tdd|jDVqdS)css(|]}|jr|jddndVqdSrrN)rrJrr.r.r/rgs  zJDatasetStatsSummary.get_total_time_all_blocks...N)rrrr.r.r/rgs z@DatasetStatsSummary.get_total_time_all_blocks..)r}rr)r?rr.r.r/rs z-DatasetStatsSummary.get_total_time_all_blockscCs,tdd|jD}|tdd|jDS)Ncss|]}|VqdSrA)get_total_cpu_timerer.r.r/rgrz9DatasetStatsSummary.get_total_cpu_time..css|] }|jddVqdSr)cpu_timerJrr.r.r/rgs )rrjr)r?rr.r.r/rs  z&DatasetStatsSummary.get_total_cpu_timecCsFdd|jD}|rt|nd}|js|St|gdd|jDRS)NcSrNr.)get_max_heap_memoryrer.r.r/rrPz;DatasetStatsSummary.get_max_heap_memory..rcSsg|] }|jddqS)rMr)memoryrJrr.r.r/rr)rjrMr)r? parent_memory parent_maxr.r.r/rsz'DatasetStatsSummary.get_max_heap_memory)NTTr)rPrQrRr__annotations__rr,r:rrr r r[r staticmethodrr rrrrrrrr.r.r.r/r}sJ    t    r}c @s>eZdZUeed<eed<eed<eed<eed<eed<dZee eefed<dZ ee eefed <dZ ee eefed <dZ ee eefed <dZ ee eefed <dZee eefed <dZee eefed<dZee eefed<ededeededdfddZdefddZddefddZdS)rrUr~rarrblock_execution_summary_strNrrudf_timerroutput_size_bytes node_count task_rows block_statsr%cCsVdd|D}d}d}d\}}|r'tdd|D}tdd|D}||}|r1dt|} n|rGt|d }|dkr>d}d t||} nd } | d 7} tt} |D]} | jd urk| j d urk| | j j | j7<qTd } t| dkrt| t| tt t| t| d} dt| | } d\} }}}|r!tdd|Dtdd|Dt dd|Dtdd|Dd} tdd|Dtdd|Dt dd|Dtdd|Dd}dd|D}t|t|tt |d}tdd|Dtdd|Dt dd|Dtdd|Dd}d }d d|D}|r?t|t|tt |t|d}d }d!d|D}|r]t|t|tt |t|d}d }|rtt}|D] }||j|j qid"d#|D}t| t| tt t| t|d}t|||||| | ||||||| d$S)%aCalculate the stats for a operator from a given list of blocks, and generates a `OperatorStatsSummary` object with the results. Args: block_stats: List of `BlockStats` to calculate stats of operator_name: Name of operator associated with `blocks` is_sub_operator: Whether this set of blocks belongs to a sub operator. Returns: A `OperatorStatsSummary` object initialized with the calculated statistics cSg|] }|jdur|jqSrA) exec_statsrmr.r.r/rz.r)rrcsrcrA) start_time_srr.r.r/rgrhz;OperatorStatsSummary.from_block_metadata..csrcrA) end_time_srr.r.r/rgrhz{} blocks produced r'z{} blocks produced in {}srrN)rLrMmeancountz{} tasks executed, {})NNNNcSg|]}|jqSr. wall_time_srer.r.r/r#cSrr.rrr.r.r/r$rcSrr.rrr.r.r/r%rcSrr.rrr.r.r/r&r)rLrMrrcSrr. cpu_time_srr.r.r/r)rcSrr.rrr.r.r/r*rcSrr.rrr.r.r/r+rcSrr.rrr.r.r/r,rcSs g|] }t|jp dddqS)rir')r- max_uss_bytesrr.r.r/r/s)rLrMrcSrr. udf_time_srr.r.r/r9rcSrr.rrr.r.r/r:rcSrr.rrr.r.r/r;rcSrr.rrr.r.r/r<rcSrrA)num_rowsrr.r.r/r@rcSrrA) size_bytesrr.r.r/rJscSi|] \}}|t|qSr.)r_)rrtasksr.r.r/r[rz.)rUr~rarrrrrrrrrrr)rLrMrr_r-rrrrrtask_idxr4nprr3rrrrDr^r)r^rUrr~rrrarrexec_summary_strrmetatask_rows_statswall_time_stats cpu_stats memory_stats udf_statsmemory_stats_mboutput_num_rows_statsroutput_size_bytes_statsrnode_counts_stats node_tasksr( node_countsr.r.r/rs            z(OperatorStatsSummary.from_block_metadatac CsT|jrdnd}|j}|j}|r-||7}|dt|dt|dt|dt|d7}|j}|rP||7}|dt|dt|dt|dt|d7}|j}|rs||7}|d t|dt|dt|dt|d7}|j}|r||7}|d |d|d|d7}|j}|r||7}|d |d|d|d|d7}|j }|r||7}|d |d|d|d|d7}|j } | r||7}|d | d| d| d| d7}|j } | r||7}|d| d| d| d| d7}|r(|j r(|r(|d} ||7}|d7}||d| |j d7}||d| |dd7}|S)EFor a given (pre-calculated) `OperatorStatsSummary` object (e.g. generated from `OperatorStatsSummary.from_block_metadata()`), returns a human-friendly string that summarizes operator execution statistics. Returns: String with summary statistics for executing the given operator. rrz6* Remote wall time: {} min, {} max, {} mean, {} total rLrMrrz5* Remote cpu time: {} min, {} max, {} mean, {} total z.* UDF time: {} min, {} max, {} mean, {} total z8* Peak heap memory usage (MiB): {} min, {} max, {} mean z?* Output num rows per block: {} min, {} max, {} mean, {} total zA* Output size bytes per block: {} min, {} max, {} mean, {} total z?* Output rows per task: {} min, {} max, {} mean, {} tasks used rz9* Tasks per node: {} min, {} max, {} mean; {} nodes used z* Operator throughput: rrr) r~rrrr0rrrrrrrra) r?rrrrrrrrrnode_count_statsrr.r.r/r\ts                 zOperatorStatsSummary.__str__rc Cst|}||jr tdnd7}dd|jpiD}dd|jp"iD}dd|jp.iD}dd|jp:iD}dd|jpFiD}d d|jpRiD}d g|d |d |j d |d |jd|dt |j d|d|j |d|pdd|d|pdd|d|pdd|d|pdd|d|pdd|d|pdd|d} | S)rr&rcSrr.rrr.r.r/rrz1OperatorStatsSummary.__repr__..cSrr.rrr.r.r/rrcSrr.rrr.r.r/rrcSrr.rrr.r.r/rcSrr.rrr.r.r/rr cSrr.rrr.r.r/rrzOperatorStatsSummary( z operator_name='z', z is_suboperator=rz time_total_s=z block_execution_summary_str=z wall_time=Nz cpu_time=z memory=z output_num_rows=z output_size_bytes=z node_count=r)r5r~rr^rrrrrrrUr0rar) r?rrrrrrrnode_conut_statsrr.r.r/rsr         zOperatorStatsSummary.__repr__r)rPrQrRr,rr[r:rr rrrrrrrrrbrrrr\rr.r.r.r/rs8  zrc@seZdZUeed<eed<eed<eed<eed<eed<eed<eed<eed <eed <eed <eed <eed <eed<defddZdefddZddefddZ dS)r wait_timeget_time next_time format_time collate_timefinalize_batch_time block_time user_timeinitialize_time total_timestreaming_split_coord_timerurvrwr%cCs|SrA)rr>r.r.r/r\3szIterStatsSummary.__str__c Csd}|js&|js&|js&|js&|js&|js&|jrc|d7}|jr;|dt |j7}|j rL|dt |j 7}|jr]|dt |j7}|j rn|dt |j 7}|d7}|jr|dt |j t |j t |jt |j7}|jrd }||t |j t |j t |jt |j7}|jrd }||t |j t |j t |jt |j7}|jr|d t |j t |j t |jt |j7}|jr)d }||t |j t |j t |jt |j7}tjrK|d 7}|d|j7}|d|j7}|d|j7}|jdkrc|d7}|t |jd7}|S)Nrz" Dataset iterator time breakdown: z* Total time overall: {} z> * Total time in Ray Data iterator initialization code: {} zE * Total time user thread is blocked by Ray Data iter_batches: {} z/ * Total execution time for user thread: {} zC* Batch iteration time breakdown (summed across prefetch threads): z5 * In ray.get(): {} min, {} max, {} avg, {} total z: * In batch creation: {} min, {} max, {} avg, {} total z< * In batch formatting: {} min, {} max, {} avg, {} total z6 * In collate_fn: {} min, {} max, {} avg, {} total zA * In host->device transfer: {} min, {} max, {} avg, {} total zBlock locations: z * Num blocks local: {} z * Num blocks remote: {} z& * Num blocks unknown location: {} rz+Streaming split coordinator overhead time: r)rrJrr r rrrrr0rrrLrMrNrr'enable_get_object_locations_for_metricsrurvrwr)r?rbatch_creation_str format_strr.r.r/r6s                               zIterStatsSummary.to_stringrcCs:t|}dgd|dt|jpdd|dt|jp'dd|d|jp4dd|d|jpAdd|d|jpNdd|d t|j p_dd|d t|j ppdd|d t|j pdd|d t|j pdd|d S)NrzIterStatsSummary( z wait_time=rz get_time=z iter_blocks_local=z iter_blocks_remote=z iter_unknown_location=z next_time=z format_time=z user_time=z total_time=r) r5rr0r rJr rurvrwr rrr)r?rrr.r.r/rsR    zIterStatsSummary.__repr__Nr) rPrQrRr7rrr,r\rrr.r.r.r/rs$ Xr)rr1)Qrenumloggingr#rBr contextlibr dataclassesrrtypingrrrr r r r r uuidrnumpyrr ray.actorrray.data._internal.block_listr:ray.data._internal.execution.interfaces.op_runtime_metricsrrrrrrrrray.data._internal.utilrray.data.blockrray.data.contextrray.util.annotationsrray.util.metricsrrrrray.util.scheduling_strategiesr getLoggerrPr6rrrr,rlr:r0rr5r7rTrrnRLockr rrrr:IntEnumrrWr}rrr.r.r.r/sl   (           & -0TxP