github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/worker/statesampler_slow.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  # pytype: skip-file
    19  
    20  from typing import Optional
    21  
    22  from apache_beam.runners import common
    23  from apache_beam.utils import counters
    24  
    25  
    26  class StateSampler(object):
    27    def __init__(self, sampling_period_ms):
    28      self._state_stack = [
    29          ScopedState(self, counters.CounterName('unknown'), None)
    30      ]
    31      self.state_transition_count = 0
    32      self.time_since_transition = 0
    33  
    34    def current_state(self):
    35      # type: () -> ScopedState
    36  
    37      """Returns the current execution state.
    38  
    39      This operation is not thread safe, and should only be called from the
    40      execution thread."""
    41      return self._state_stack[-1]
    42  
    43    def _scoped_state(self,
    44                      counter_name,  # type: counters.CounterName
    45                      name_context,  # type: common.NameContext
    46                      output_counter,
    47                      metrics_container=None):
    48      # type: (...) -> ScopedState
    49      assert isinstance(name_context, common.NameContext)
    50      return ScopedState(
    51          self, counter_name, name_context, output_counter, metrics_container)
    52  
    53    def update_metric(self, typed_metric_name, value):
    54      metrics_container = self.current_state().metrics_container
    55      if metrics_container is not None:
    56        metrics_container.get_metric_cell(typed_metric_name).update(value)
    57  
    58    def _enter_state(self, state):
    59      # type: (ScopedState) -> None
    60      self.state_transition_count += 1
    61      self._state_stack.append(state)
    62  
    63    def _exit_state(self):
    64      # type: () -> None
    65      self.state_transition_count += 1
    66      self._state_stack.pop()
    67  
    68    def start(self):
    69      # type: () -> None
    70      # Sampling not yet supported. Only state tracking at the moment.
    71      pass
    72  
    73    def stop(self):
    74      # type: () -> None
    75      pass
    76  
    77    def reset(self):
    78      # type: () -> None
    79      pass
    80  
    81  
    82  class ScopedState(object):
    83  
    84    def __init__(self,
    85                 sampler,  # type: StateSampler
    86                 name,  # type: counters.CounterName
    87                 step_name_context,  # type: Optional[common.NameContext]
    88                 counter=None,
    89                 metrics_container=None):
    90      self.state_sampler = sampler
    91      self.name = name
    92      self.name_context = step_name_context
    93      self.counter = counter
    94      self.nsecs = 0
    95      self.metrics_container = metrics_container
    96  
    97    def sampled_seconds(self):
    98      # type: () -> float
    99      return 1e-9 * self.nsecs
   100  
   101    def sampled_msecs_int(self):
   102      # type: () -> int
   103      return int(1e-6 * self.nsecs)
   104  
   105    def __repr__(self):
   106      return "ScopedState[%s, %s]" % (self.name, self.nsecs)
   107  
   108    def __enter__(self):
   109      self.state_sampler._enter_state(self)
   110  
   111    def __exit__(self, exc_type, exc_value, traceback):
   112      self.state_sampler._exit_state()