| |
| |
| <!DOCTYPE html> |
| <!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]--> |
| <!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]--> |
| <head> |
| <meta charset="utf-8"> |
| |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| |
| <title>apache_beam.ml.inference.base — Apache Beam 2.47.0 documentation</title> |
| |
| |
| |
| |
| |
| |
| |
| |
| <script type="text/javascript" src="../../../../_static/js/modernizr.min.js"></script> |
| |
| |
| <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script> |
| <script type="text/javascript" src="../../../../_static/jquery.js"></script> |
| <script type="text/javascript" src="../../../../_static/underscore.js"></script> |
| <script type="text/javascript" src="../../../../_static/doctools.js"></script> |
| <script type="text/javascript" src="../../../../_static/language_data.js"></script> |
| <script async="async" type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script> |
| |
| <script type="text/javascript" src="../../../../_static/js/theme.js"></script> |
| |
| |
| |
| |
| <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" /> |
| <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" /> |
| <link rel="index" title="Index" href="../../../../genindex.html" /> |
| <link rel="search" title="Search" href="../../../../search.html" /> |
| </head> |
| |
| <body class="wy-body-for-nav"> |
| |
| |
| <div class="wy-grid-for-nav"> |
| |
| <nav data-toggle="wy-nav-shift" class="wy-nav-side"> |
| <div class="wy-side-scroll"> |
| <div class="wy-side-nav-search" > |
| |
| |
| |
| <a href="../../../../index.html" class="icon icon-home"> Apache Beam |
| |
| |
| |
| </a> |
| |
| |
| |
| |
| <div class="version"> |
| 2.47.0 |
| </div> |
| |
| |
| |
| |
| <div role="search"> |
| <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get"> |
| <input type="text" name="q" placeholder="Search docs" /> |
| <input type="hidden" name="check_keywords" value="yes" /> |
| <input type="hidden" name="area" value="default" /> |
| </form> |
| </div> |
| |
| |
| </div> |
| |
| <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation"> |
| |
| |
| |
| |
| |
| |
| <ul> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.coders.html">apache_beam.coders package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.dataframe.html">apache_beam.dataframe package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.io.html">apache_beam.io package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.metrics.html">apache_beam.metrics package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.ml.html">apache_beam.ml package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.options.html">apache_beam.options package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.portability.html">apache_beam.portability package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.runners.html">apache_beam.runners package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.testing.html">apache_beam.testing package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.transforms.html">apache_beam.transforms package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.typehints.html">apache_beam.typehints package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.utils.html">apache_beam.utils package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.yaml.html">apache_beam.yaml package</a></li> |
| </ul> |
| <ul> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.error.html">apache_beam.error module</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.pipeline.html">apache_beam.pipeline module</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.pvalue.html">apache_beam.pvalue module</a></li> |
| </ul> |
| |
| |
| |
| </div> |
| </div> |
| </nav> |
| |
| <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"> |
| |
| |
| <nav class="wy-nav-top" aria-label="top navigation"> |
| |
| <i data-toggle="wy-nav-top" class="fa fa-bars"></i> |
| <a href="../../../../index.html">Apache Beam</a> |
| |
| </nav> |
| |
| |
| <div class="wy-nav-content"> |
| |
| <div class="rst-content"> |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| <div role="navigation" aria-label="breadcrumbs navigation"> |
| |
| <ul class="wy-breadcrumbs"> |
| |
| <li><a href="../../../../index.html">Docs</a> »</li> |
| |
| <li><a href="../../../index.html">Module code</a> »</li> |
| |
| <li>apache_beam.ml.inference.base</li> |
| |
| |
| <li class="wy-breadcrumbs-aside"> |
| |
| </li> |
| |
| </ul> |
| |
| |
| <hr/> |
| </div> |
| <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article"> |
| <div itemprop="articleBody"> |
| |
| <h1>Source code for apache_beam.ml.inference.base</h1><div class="highlight"><pre> |
| <span></span><span class="c1">#</span> |
| <span class="c1"># Licensed to the Apache Software Foundation (ASF) under one or more</span> |
| <span class="c1"># contributor license agreements. See the NOTICE file distributed with</span> |
| <span class="c1"># this work for additional information regarding copyright ownership.</span> |
| <span class="c1"># The ASF licenses this file to You under the Apache License, Version 2.0</span> |
| <span class="c1"># (the "License"); you may not use this file except in compliance with</span> |
| <span class="c1"># the License. You may obtain a copy of the License at</span> |
| <span class="c1">#</span> |
| <span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span> |
| <span class="c1">#</span> |
| <span class="c1"># Unless required by applicable law or agreed to in writing, software</span> |
| <span class="c1"># distributed under the License is distributed on an "AS IS" BASIS,</span> |
| <span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span> |
| <span class="c1"># See the License for the specific language governing permissions and</span> |
| <span class="c1"># limitations under the License.</span> |
| <span class="c1">#</span> |
| <span class="c1"># TODO: https://github.com/apache/beam/issues/21822</span> |
| <span class="c1"># mypy: ignore-errors</span> |
| |
| <span class="sd">"""An extensible run inference transform.</span> |
| |
| <span class="sd">Users of this module can extend the ModelHandler class for any machine learning</span> |
| <span class="sd">framework. A ModelHandler implementation is a required parameter of</span> |
| <span class="sd">RunInference.</span> |
| |
| <span class="sd">The transform handles standard inference functionality, like metric</span> |
| <span class="sd">collection, sharing model between threads, and batching elements.</span> |
| <span class="sd">"""</span> |
| |
| <span class="kn">import</span> <span class="nn">logging</span> |
| <span class="kn">import</span> <span class="nn">pickle</span> |
| <span class="kn">import</span> <span class="nn">sys</span> |
| <span class="kn">import</span> <span class="nn">threading</span> |
| <span class="kn">import</span> <span class="nn">time</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Generic</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Iterable</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Mapping</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">NamedTuple</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Sequence</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">TypeVar</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span> |
| |
| <span class="kn">import</span> <span class="nn">apache_beam</span> <span class="k">as</span> <span class="nn">beam</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.utils</span> <span class="kn">import</span> <span class="n">shared</span> |
| |
| <span class="k">try</span><span class="p">:</span> |
| <span class="c1"># pylint: disable=wrong-import-order, wrong-import-position</span> |
| <span class="kn">import</span> <span class="nn">resource</span> |
| <span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span> |
| <span class="n">resource</span> <span class="o">=</span> <span class="kc">None</span> <span class="c1"># type: ignore[assignment]</span> |
| |
| <span class="n">_NANOSECOND_TO_MILLISECOND</span> <span class="o">=</span> <span class="mi">1_000_000</span> |
| <span class="n">_NANOSECOND_TO_MICROSECOND</span> <span class="o">=</span> <span class="mi">1_000</span> |
| |
| <span class="n">ModelT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'ModelT'</span><span class="p">)</span> |
| <span class="n">ExampleT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'ExampleT'</span><span class="p">)</span> |
| <span class="n">PredictionT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'PredictionT'</span><span class="p">)</span> |
| <span class="n">_INPUT_TYPE</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'_INPUT_TYPE'</span><span class="p">)</span> |
| <span class="n">_OUTPUT_TYPE</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'_OUTPUT_TYPE'</span><span class="p">)</span> |
| <span class="n">KeyT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'KeyT'</span><span class="p">)</span> |
| |
| |
| <span class="c1"># We use NamedTuple to define the structure of the PredictionResult,</span> |
| <span class="c1"># however, as support for generic NamedTuples is not available in Python</span> |
| <span class="c1"># versions prior to 3.11, we use the __new__ method to provide default</span> |
| <span class="c1"># values for the fields while maintaining backwards compatibility.</span> |
| <div class="viewcode-block" id="PredictionResult"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.PredictionResult">[docs]</a><span class="k">class</span> <span class="nc">PredictionResult</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">(</span><span class="s1">'PredictionResult'</span><span class="p">,</span> |
| <span class="p">[(</span><span class="s1">'example'</span><span class="p">,</span> <span class="n">_INPUT_TYPE</span><span class="p">),</span> |
| <span class="p">(</span><span class="s1">'inference'</span><span class="p">,</span> <span class="n">_OUTPUT_TYPE</span><span class="p">),</span> |
| <span class="p">(</span><span class="s1">'model_id'</span><span class="p">,</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">])])):</span> |
| <span class="vm">__slots__</span> <span class="o">=</span> <span class="p">()</span> |
| |
| <span class="k">def</span> <span class="fm">__new__</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">example</span><span class="p">,</span> <span class="n">inference</span><span class="p">,</span> <span class="n">model_id</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span> |
| <span class="k">return</span> <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__new__</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">example</span><span class="p">,</span> <span class="n">inference</span><span class="p">,</span> <span class="n">model_id</span><span class="p">)</span></div> |
| |
| |
| <span class="n">PredictionResult</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""A NamedTuple containing both input and output</span> |
| <span class="s2"> from the inference."""</span> |
| <span class="n">PredictionResult</span><span class="o">.</span><span class="n">example</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""The input example."""</span> |
| <span class="n">PredictionResult</span><span class="o">.</span><span class="n">inference</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""Results for the inference on the model</span> |
| <span class="s2"> for the given example."""</span> |
| <span class="n">PredictionResult</span><span class="o">.</span><span class="n">model_id</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""Model ID used to run the prediction."""</span> |
| |
| |
| <div class="viewcode-block" id="ModelMetadata"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelMetadata">[docs]</a><span class="k">class</span> <span class="nc">ModelMetadata</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">):</span> |
| <span class="n">model_id</span><span class="p">:</span> <span class="nb">str</span> |
| <span class="n">model_name</span><span class="p">:</span> <span class="nb">str</span></div> |
| |
| |
| <span class="n">ModelMetadata</span><span class="o">.</span><span class="n">model_id</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""Unique identifier for the model. This can be</span> |
| <span class="s2"> a file path or a URL where the model can be accessed. It is used to load</span> |
| <span class="s2"> the model for inference."""</span> |
| <span class="n">ModelMetadata</span><span class="o">.</span><span class="n">model_name</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="s2">"""Human-readable name for the model. This</span> |
| <span class="s2"> can be used to identify the model in the metrics generated by the</span> |
| <span class="s2"> RunInference transform."""</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_to_milliseconds</span><span class="p">(</span><span class="n">time_ns</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span> |
| <span class="k">return</span> <span class="nb">int</span><span class="p">(</span><span class="n">time_ns</span> <span class="o">/</span> <span class="n">_NANOSECOND_TO_MILLISECOND</span><span class="p">)</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_to_microseconds</span><span class="p">(</span><span class="n">time_ns</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span> |
| <span class="k">return</span> <span class="nb">int</span><span class="p">(</span><span class="n">time_ns</span> <span class="o">/</span> <span class="n">_NANOSECOND_TO_MICROSECOND</span><span class="p">)</span> |
| |
| |
| <div class="viewcode-block" id="ModelHandler"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler">[docs]</a><span class="k">class</span> <span class="nc">ModelHandler</span><span class="p">(</span><span class="n">Generic</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">ModelT</span><span class="p">]):</span> |
| <span class="w"> </span><span class="sd">"""Has the ability to load and apply an ML model."""</span> |
| <div class="viewcode-block" id="ModelHandler.load_model"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.load_model">[docs]</a> <span class="k">def</span> <span class="nf">load_model</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">ModelT</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""Loads and initializes a model for processing."""</span> |
| <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.run_inference"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.run_inference">[docs]</a> <span class="k">def</span> <span class="nf">run_inference</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">],</span> |
| <span class="n">model</span><span class="p">:</span> <span class="n">ModelT</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-></span> <span class="n">Iterable</span><span class="p">[</span><span class="n">PredictionT</span><span class="p">]:</span> |
| <span class="w"> </span><span class="sd">"""Runs inferences on a batch of examples.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> batch: A sequence of examples or features.</span> |
| <span class="sd"> model: The model used to make inferences.</span> |
| <span class="sd"> inference_args: Extra arguments for models whose inference call requires</span> |
| <span class="sd"> extra parameters.</span> |
| |
| <span class="sd"> Returns:</span> |
| <span class="sd"> An Iterable of Predictions.</span> |
| <span class="sd"> """</span> |
| <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.get_num_bytes"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.get_num_bytes">[docs]</a> <span class="k">def</span> <span class="nf">get_num_bytes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">])</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> The number of bytes of data for a batch.</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="n">pickle</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">batch</span><span class="p">))</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.get_metrics_namespace"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.get_metrics_namespace">[docs]</a> <span class="k">def</span> <span class="nf">get_metrics_namespace</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="nb">str</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> A namespace for metrics collected by the RunInference transform.</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="s1">'RunInference'</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.get_resource_hints"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.get_resource_hints">[docs]</a> <span class="k">def</span> <span class="nf">get_resource_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="nb">dict</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> Resource hints for the transform.</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="p">{}</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.batch_elements_kwargs"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.batch_elements_kwargs">[docs]</a> <span class="k">def</span> <span class="nf">batch_elements_kwargs</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> kwargs suitable for beam.BatchElements.</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="p">{}</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.validate_inference_args"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.validate_inference_args">[docs]</a> <span class="k">def</span> <span class="nf">validate_inference_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]):</span> |
| <span class="w"> </span><span class="sd">"""Validates inference_args passed in the inference call.</span> |
| |
| <span class="sd"> Because most frameworks do not need extra arguments in their predict() call,</span> |
| <span class="sd"> the default behavior is to error out if inference_args are present.</span> |
| <span class="sd"> """</span> |
| <span class="k">if</span> <span class="n">inference_args</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span> |
| <span class="s1">'inference_args were provided, but should be None because this '</span> |
| <span class="s1">'framework does not expect extra arguments on inferences.'</span><span class="p">)</span></div> |
| |
| <div class="viewcode-block" id="ModelHandler.update_model_path"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.ModelHandler.update_model_path">[docs]</a> <span class="k">def</span> <span class="nf">update_model_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_path</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""Update the model paths produced by side inputs."""</span> |
| <span class="k">pass</span></div></div> |
| |
| |
| <div class="viewcode-block" id="KeyedModelHandler"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler">[docs]</a><span class="k">class</span> <span class="nc">KeyedModelHandler</span><span class="p">(</span><span class="n">Generic</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">ModelT</span><span class="p">],</span> |
| <span class="n">ModelHandler</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">],</span> |
| <span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">],</span> |
| <span class="n">ModelT</span><span class="p">]):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">unkeyed</span><span class="p">:</span> <span class="n">ModelHandler</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">ModelT</span><span class="p">]):</span> |
| <span class="w"> </span><span class="sd">"""A ModelHandler that takes keyed examples and returns keyed predictions.</span> |
| |
| <span class="sd"> For example, if the original model is used with RunInference to take a</span> |
| <span class="sd"> PCollection[E] to a PCollection[P], this ModelHandler would take a</span> |
| <span class="sd"> PCollection[Tuple[K, E]] to a PCollection[Tuple[K, P]], making it possible</span> |
| <span class="sd"> to use the key to associate the outputs with the inputs.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> unkeyed: An implementation of ModelHandler that does not require keys.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span> <span class="o">=</span> <span class="n">unkeyed</span> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.load_model"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.load_model">[docs]</a> <span class="k">def</span> <span class="nf">load_model</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">ModelT</span><span class="p">:</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">load_model</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.run_inference"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.run_inference">[docs]</a> <span class="k">def</span> <span class="nf">run_inference</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">]],</span> |
| <span class="n">model</span><span class="p">:</span> <span class="n">ModelT</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="p">)</span> <span class="o">-></span> <span class="n">Iterable</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">]]:</span> |
| <span class="n">keys</span><span class="p">,</span> <span class="n">unkeyed_batch</span> <span class="o">=</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">batch</span><span class="p">)</span> |
| <span class="k">return</span> <span class="nb">zip</span><span class="p">(</span> |
| <span class="n">keys</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">run_inference</span><span class="p">(</span><span class="n">unkeyed_batch</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">))</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.get_num_bytes"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.get_num_bytes">[docs]</a> <span class="k">def</span> <span class="nf">get_num_bytes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">]])</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span> |
| <span class="n">keys</span><span class="p">,</span> <span class="n">unkeyed_batch</span> <span class="o">=</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">batch</span><span class="p">)</span> |
| <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="n">pickle</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">keys</span><span class="p">))</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_num_bytes</span><span class="p">(</span><span class="n">unkeyed_batch</span><span class="p">)</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.get_metrics_namespace"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.get_metrics_namespace">[docs]</a> <span class="k">def</span> <span class="nf">get_metrics_namespace</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="nb">str</span><span class="p">:</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_metrics_namespace</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.get_resource_hints"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.get_resource_hints">[docs]</a> <span class="k">def</span> <span class="nf">get_resource_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_resource_hints</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.batch_elements_kwargs"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.batch_elements_kwargs">[docs]</a> <span class="k">def</span> <span class="nf">batch_elements_kwargs</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">batch_elements_kwargs</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.validate_inference_args"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.validate_inference_args">[docs]</a> <span class="k">def</span> <span class="nf">validate_inference_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">validate_inference_args</span><span class="p">(</span><span class="n">inference_args</span><span class="p">)</span></div> |
| |
| <div class="viewcode-block" id="KeyedModelHandler.update_model_path"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.KeyedModelHandler.update_model_path">[docs]</a> <span class="k">def</span> <span class="nf">update_model_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_path</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">update_model_path</span><span class="p">(</span><span class="n">model_path</span><span class="o">=</span><span class="n">model_path</span><span class="p">)</span></div></div> |
| |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler">[docs]</a><span class="k">class</span> <span class="nc">MaybeKeyedModelHandler</span><span class="p">(</span><span class="n">Generic</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">ModelT</span><span class="p">],</span> |
| <span class="n">ModelHandler</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> |
| <span class="n">ExampleT</span><span class="p">]],</span> |
| <span class="n">Union</span><span class="p">[</span><span class="n">PredictionT</span><span class="p">,</span> |
| <span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">]],</span> |
| <span class="n">ModelT</span><span class="p">]):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">unkeyed</span><span class="p">:</span> <span class="n">ModelHandler</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">ModelT</span><span class="p">]):</span> |
| <span class="w"> </span><span class="sd">"""A ModelHandler that takes examples that might have keys and returns</span> |
| <span class="sd"> predictions that might have keys.</span> |
| |
| <span class="sd"> For example, if the original model is used with RunInference to take a</span> |
| <span class="sd"> PCollection[E] to a PCollection[P], this ModelHandler would take either</span> |
| <span class="sd"> PCollection[E] to a PCollection[P] or PCollection[Tuple[K, E]] to a</span> |
| <span class="sd"> PCollection[Tuple[K, P]], depending on the whether the elements are</span> |
| <span class="sd"> tuples. This pattern makes it possible to associate the outputs with the</span> |
| <span class="sd"> inputs based on the key.</span> |
| |
| <span class="sd"> Note that you cannot use this ModelHandler if E is a tuple type.</span> |
| <span class="sd"> In addition, either all examples should be keyed, or none of them.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> unkeyed: An implementation of ModelHandler that does not require keys.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span> <span class="o">=</span> <span class="n">unkeyed</span> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.load_model"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.load_model">[docs]</a> <span class="k">def</span> <span class="nf">load_model</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">ModelT</span><span class="p">:</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">load_model</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.run_inference"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.run_inference">[docs]</a> <span class="k">def</span> <span class="nf">run_inference</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">]]],</span> |
| <span class="n">model</span><span class="p">:</span> <span class="n">ModelT</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="p">)</span> <span class="o">-></span> <span class="n">Union</span><span class="p">[</span><span class="n">Iterable</span><span class="p">[</span><span class="n">PredictionT</span><span class="p">],</span> <span class="n">Iterable</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">]]]:</span> |
| <span class="c1"># Really the input should be</span> |
| <span class="c1"># Union[Sequence[ExampleT], Sequence[Tuple[KeyT, ExampleT]]]</span> |
| <span class="c1"># but there's not a good way to express (or check) that.</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="nb">tuple</span><span class="p">):</span> |
| <span class="n">is_keyed</span> <span class="o">=</span> <span class="kc">True</span> |
| <span class="n">keys</span><span class="p">,</span> <span class="n">unkeyed_batch</span> <span class="o">=</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">batch</span><span class="p">)</span> <span class="c1"># type: ignore[arg-type]</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">is_keyed</span> <span class="o">=</span> <span class="kc">False</span> |
| <span class="n">unkeyed_batch</span> <span class="o">=</span> <span class="n">batch</span> <span class="c1"># type: ignore[assignment]</span> |
| <span class="n">unkeyed_results</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">run_inference</span><span class="p">(</span> |
| <span class="n">unkeyed_batch</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">if</span> <span class="n">is_keyed</span><span class="p">:</span> |
| <span class="k">return</span> <span class="nb">zip</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">unkeyed_results</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">return</span> <span class="n">unkeyed_results</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.get_num_bytes"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.get_num_bytes">[docs]</a> <span class="k">def</span> <span class="nf">get_num_bytes</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">KeyT</span><span class="p">,</span> <span class="n">ExampleT</span><span class="p">]]])</span> <span class="o">-></span> <span class="nb">int</span><span class="p">:</span> |
| <span class="c1"># MyPy can't follow the branching logic.</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="nb">tuple</span><span class="p">):</span> |
| <span class="n">keys</span><span class="p">,</span> <span class="n">unkeyed_batch</span> <span class="o">=</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">batch</span><span class="p">)</span> <span class="c1"># type: ignore[arg-type]</span> |
| <span class="k">return</span> <span class="nb">len</span><span class="p">(</span> |
| <span class="n">pickle</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">keys</span><span class="p">))</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_num_bytes</span><span class="p">(</span><span class="n">unkeyed_batch</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_num_bytes</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="c1"># type: ignore[arg-type]</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.get_metrics_namespace"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.get_metrics_namespace">[docs]</a> <span class="k">def</span> <span class="nf">get_metrics_namespace</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="nb">str</span><span class="p">:</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_metrics_namespace</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.get_resource_hints"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.get_resource_hints">[docs]</a> <span class="k">def</span> <span class="nf">get_resource_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">get_resource_hints</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.batch_elements_kwargs"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.batch_elements_kwargs">[docs]</a> <span class="k">def</span> <span class="nf">batch_elements_kwargs</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">batch_elements_kwargs</span><span class="p">()</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.validate_inference_args"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.validate_inference_args">[docs]</a> <span class="k">def</span> <span class="nf">validate_inference_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">validate_inference_args</span><span class="p">(</span><span class="n">inference_args</span><span class="p">)</span></div> |
| |
| <div class="viewcode-block" id="MaybeKeyedModelHandler.update_model_path"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.MaybeKeyedModelHandler.update_model_path">[docs]</a> <span class="k">def</span> <span class="nf">update_model_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_path</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_unkeyed</span><span class="o">.</span><span class="n">update_model_path</span><span class="p">(</span><span class="n">model_path</span><span class="o">=</span><span class="n">model_path</span><span class="p">)</span></div></div> |
| |
| |
| <div class="viewcode-block" id="RunInference"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.RunInference">[docs]</a><span class="k">class</span> <span class="nc">RunInference</span><span class="p">(</span><span class="n">beam</span><span class="o">.</span><span class="n">PTransform</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">],</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">PredictionT</span><span class="p">]]):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">model_handler</span><span class="p">:</span> <span class="n">ModelHandler</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span> |
| <span class="n">clock</span><span class="o">=</span><span class="n">time</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">metrics_namespace</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="o">*</span><span class="p">,</span> |
| <span class="n">model_metadata_pcoll</span><span class="p">:</span> <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ModelMetadata</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> A transform that takes a PCollection of examples (or features) for use</span> |
| <span class="sd"> on an ML model. The transform then outputs inferences (or predictions) for</span> |
| <span class="sd"> those examples in a PCollection of PredictionResults that contains the input</span> |
| <span class="sd"> examples and the output inferences.</span> |
| |
| <span class="sd"> Models for supported frameworks can be loaded using a URI. Supported</span> |
| <span class="sd"> services can also be used.</span> |
| |
| <span class="sd"> This transform attempts to batch examples using the beam.BatchElements</span> |
| <span class="sd"> transform. Batching can be configured using the ModelHandler.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> model_handler: An implementation of ModelHandler.</span> |
| <span class="sd"> clock: A clock implementing time_ns. *Used for unit testing.*</span> |
| <span class="sd"> inference_args: Extra arguments for models whose inference call requires</span> |
| <span class="sd"> extra parameters.</span> |
| <span class="sd"> metrics_namespace: Namespace of the transform to collect metrics.</span> |
| <span class="sd"> model_metadata_pcoll: PCollection that emits Singleton ModelMetadata</span> |
| <span class="sd"> containing model path and model name, that is used as a side input</span> |
| <span class="sd"> to the _RunInferenceDoFn.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span> <span class="o">=</span> <span class="n">model_handler</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_args</span> <span class="o">=</span> <span class="n">inference_args</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_clock</span> <span class="o">=</span> <span class="n">clock</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_namespace</span> <span class="o">=</span> <span class="n">metrics_namespace</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_metadata_pcoll</span> <span class="o">=</span> <span class="n">model_metadata_pcoll</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_enable_side_input_loading</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_metadata_pcoll</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> |
| |
| <span class="c1"># TODO(BEAM-14046): Add and link to help documentation.</span> |
| <div class="viewcode-block" id="RunInference.from_callable"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.RunInference.from_callable">[docs]</a> <span class="nd">@classmethod</span> |
| <span class="k">def</span> <span class="nf">from_callable</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">model_handler_provider</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""Multi-language friendly constructor.</span> |
| |
| <span class="sd"> Use this constructor with fully_qualified_named_transform to</span> |
| <span class="sd"> initialize the RunInference transform from PythonCallableSource provided</span> |
| <span class="sd"> by foreign SDKs.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> model_handler_provider: A callable object that returns ModelHandler.</span> |
| <span class="sd"> kwargs: Keyword arguments for model_handler_provider.</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="bp">cls</span><span class="p">(</span><span class="n">model_handler_provider</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">))</span></div> |
| |
| <span class="c1"># TODO(https://github.com/apache/beam/issues/21447): Add batch_size back off</span> |
| <span class="c1"># in the case there are functional reasons large batch sizes cannot be</span> |
| <span class="c1"># handled.</span> |
| <div class="viewcode-block" id="RunInference.expand"><a class="viewcode-back" href="../../../../apache_beam.ml.inference.base.html#apache_beam.ml.inference.base.RunInference.expand">[docs]</a> <span class="k">def</span> <span class="nf">expand</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> <span class="n">pcoll</span><span class="p">:</span> <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">])</span> <span class="o">-></span> <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">PredictionT</span><span class="p">]:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">validate_inference_args</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_inference_args</span><span class="p">)</span> |
| <span class="n">resource_hints</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">get_resource_hints</span><span class="p">()</span> |
| <span class="n">batched_elements_pcoll</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">pcoll</span> |
| <span class="c1"># TODO(https://github.com/apache/beam/issues/21440): Hook into the</span> |
| <span class="c1"># batching DoFn APIs.</span> |
| <span class="o">|</span> <span class="n">beam</span><span class="o">.</span><span class="n">BatchElements</span><span class="p">(</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">batch_elements_kwargs</span><span class="p">()))</span> |
| |
| <span class="k">return</span> <span class="p">(</span> |
| <span class="n">batched_elements_pcoll</span> |
| <span class="o">|</span> <span class="s1">'BeamML_RunInference'</span> <span class="o">>></span> <span class="p">(</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">ParDo</span><span class="p">(</span> |
| <span class="n">_RunInferenceDoFn</span><span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="p">,</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_clock</span><span class="p">,</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_namespace</span><span class="p">,</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_enable_side_input_loading</span><span class="p">),</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_args</span><span class="p">,</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">pvalue</span><span class="o">.</span><span class="n">AsSingleton</span><span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_metadata_pcoll</span><span class="p">,</span> |
| <span class="p">)</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_enable_side_input_loading</span> <span class="k">else</span> |
| <span class="kc">None</span><span class="p">)</span><span class="o">.</span><span class="n">with_resource_hints</span><span class="p">(</span><span class="o">**</span><span class="n">resource_hints</span><span class="p">)))</span></div></div> |
| |
| |
| <span class="k">class</span> <span class="nc">_MetricsCollector</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""A metrics collector that tracks ML related performance and memory usage."""</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">namespace</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">prefix</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">''</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> namespace: Namespace for the metrics.</span> |
| <span class="sd"> prefix: Unique identifier for metrics, used when models</span> |
| <span class="sd"> are updated using side input.</span> |
| <span class="sd"> """</span> |
| <span class="c1"># Metrics</span> |
| <span class="k">if</span> <span class="n">prefix</span><span class="p">:</span> |
| <span class="n">prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s1">'</span><span class="si">{</span><span class="n">prefix</span><span class="si">}</span><span class="s1">_'</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_counter</span> <span class="o">=</span> <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">counter</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'num_inferences'</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">failed_batches_counter</span> <span class="o">=</span> <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">counter</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'failed_batches_counter'</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_request_batch_size</span> <span class="o">=</span> <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">distribution</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'inference_request_batch_size'</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_request_batch_byte_size</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">distribution</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'inference_request_batch_byte_size'</span><span class="p">))</span> |
| <span class="c1"># Batch inference latency in microseconds.</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_batch_latency_micro_secs</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">distribution</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'inference_batch_latency_micro_secs'</span><span class="p">))</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size</span> <span class="o">=</span> <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">distribution</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'model_byte_size'</span><span class="p">)</span> |
| <span class="c1"># Model load latency in milliseconds.</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs</span> <span class="o">=</span> <span class="n">beam</span><span class="o">.</span><span class="n">metrics</span><span class="o">.</span><span class="n">Metrics</span><span class="o">.</span><span class="n">distribution</span><span class="p">(</span> |
| <span class="n">namespace</span><span class="p">,</span> <span class="n">prefix</span> <span class="o">+</span> <span class="s1">'load_model_latency_milli_secs'</span><span class="p">)</span> |
| |
| <span class="c1"># Metrics cache</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs_cache</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size_cache</span> <span class="o">=</span> <span class="kc">None</span> |
| |
| <span class="k">def</span> <span class="nf">update_metrics_with_cache</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs_cache</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs</span><span class="o">.</span><span class="n">update</span><span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs_cache</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs_cache</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size_cache</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size_cache</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size_cache</span> <span class="o">=</span> <span class="kc">None</span> |
| |
| <span class="k">def</span> <span class="nf">cache_load_model_metrics</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">load_model_latency_ms</span><span class="p">,</span> <span class="n">model_byte_size</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_load_model_latency_milli_secs_cache</span> <span class="o">=</span> <span class="n">load_model_latency_ms</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_byte_size_cache</span> <span class="o">=</span> <span class="n">model_byte_size</span> |
| |
| <span class="k">def</span> <span class="nf">update</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">examples_count</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> |
| <span class="n">examples_byte_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> |
| <span class="n">latency_micro_secs</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_batch_latency_micro_secs</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">latency_micro_secs</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_counter</span><span class="o">.</span><span class="n">inc</span><span class="p">(</span><span class="n">examples_count</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_request_batch_size</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">examples_count</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_inference_request_batch_byte_size</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">examples_byte_size</span><span class="p">)</span> |
| |
| |
| <span class="k">class</span> <span class="nc">_RunInferenceDoFn</span><span class="p">(</span><span class="n">beam</span><span class="o">.</span><span class="n">DoFn</span><span class="p">,</span> <span class="n">Generic</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">]):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">model_handler</span><span class="p">:</span> <span class="n">ModelHandler</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">PredictionT</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span> |
| <span class="n">clock</span><span class="p">,</span> |
| <span class="n">metrics_namespace</span><span class="p">,</span> |
| <span class="n">enable_side_input_loading</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""A DoFn implementation generic to frameworks.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> model_handler: An implementation of ModelHandler.</span> |
| <span class="sd"> clock: A clock implementing time_ns. *Used for unit testing.*</span> |
| <span class="sd"> metrics_namespace: Namespace of the transform to collect metrics.</span> |
| <span class="sd"> enable_side_input_loading: Bool to indicate if model updates</span> |
| <span class="sd"> with side inputs.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span> <span class="o">=</span> <span class="n">model_handler</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_shared_model_handle</span> <span class="o">=</span> <span class="n">shared</span><span class="o">.</span><span class="n">Shared</span><span class="p">()</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_clock</span> <span class="o">=</span> <span class="n">clock</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_namespace</span> <span class="o">=</span> <span class="n">metrics_namespace</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_enable_side_input_loading</span> <span class="o">=</span> <span class="n">enable_side_input_loading</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_side_input_path</span> <span class="o">=</span> <span class="kc">None</span> |
| |
| <span class="k">def</span> <span class="nf">_load_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">side_input_model_path</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="k">def</span> <span class="nf">load</span><span class="p">():</span> |
| <span class="w"> </span><span class="sd">"""Function for constructing shared LoadedModel."""</span> |
| <span class="n">memory_before</span> <span class="o">=</span> <span class="n">_get_current_process_memory_in_bytes</span><span class="p">()</span> |
| <span class="n">start_time</span> <span class="o">=</span> <span class="n">_to_milliseconds</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_clock</span><span class="o">.</span><span class="n">time_ns</span><span class="p">())</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">update_model_path</span><span class="p">(</span><span class="n">side_input_model_path</span><span class="p">)</span> |
| <span class="n">model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">load_model</span><span class="p">()</span> |
| <span class="n">end_time</span> <span class="o">=</span> <span class="n">_to_milliseconds</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_clock</span><span class="o">.</span><span class="n">time_ns</span><span class="p">())</span> |
| <span class="n">memory_after</span> <span class="o">=</span> <span class="n">_get_current_process_memory_in_bytes</span><span class="p">()</span> |
| <span class="n">load_model_latency_ms</span> <span class="o">=</span> <span class="n">end_time</span> <span class="o">-</span> <span class="n">start_time</span> |
| <span class="n">model_byte_size</span> <span class="o">=</span> <span class="n">memory_after</span> <span class="o">-</span> <span class="n">memory_before</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span><span class="o">.</span><span class="n">cache_load_model_metrics</span><span class="p">(</span> |
| <span class="n">load_model_latency_ms</span><span class="p">,</span> <span class="n">model_byte_size</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">model</span> |
| |
| <span class="c1"># TODO(https://github.com/apache/beam/issues/21443): Investigate releasing</span> |
| <span class="c1"># model.</span> |
| <span class="n">model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_shared_model_handle</span><span class="o">.</span><span class="n">acquire</span><span class="p">(</span><span class="n">load</span><span class="p">,</span> <span class="n">tag</span><span class="o">=</span><span class="n">side_input_model_path</span><span class="p">)</span> |
| <span class="c1"># since shared_model_handle is shared across threads, the model path</span> |
| <span class="c1"># might not get updated in the model handler</span> |
| <span class="c1"># because we directly get cached weak ref model from shared cache, instead</span> |
| <span class="c1"># of calling load(). For sanity check, call update_model_path again.</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">update_model_path</span><span class="p">(</span><span class="n">side_input_model_path</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">model</span> |
| |
| <span class="k">def</span> <span class="nf">get_metrics_collector</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">prefix</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">''</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> prefix: Unique identifier for metrics, used when models</span> |
| <span class="sd"> are updated using side input.</span> |
| <span class="sd"> """</span> |
| <span class="n">metrics_namespace</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_namespace</span><span class="p">)</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_namespace</span> <span class="k">else</span> <span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">get_metrics_namespace</span><span class="p">())</span> |
| <span class="k">return</span> <span class="n">_MetricsCollector</span><span class="p">(</span><span class="n">metrics_namespace</span><span class="p">,</span> <span class="n">prefix</span><span class="o">=</span><span class="n">prefix</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">setup</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_metrics_collector</span><span class="p">()</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">_enable_side_input_loading</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_model</span><span class="p">()</span> |
| |
| <span class="k">def</span> <span class="nf">update_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">side_input_model_path</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_model</span><span class="p">(</span><span class="n">side_input_model_path</span><span class="o">=</span><span class="n">side_input_model_path</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">_run_inference</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">):</span> |
| <span class="n">start_time</span> <span class="o">=</span> <span class="n">_to_microseconds</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_clock</span><span class="o">.</span><span class="n">time_ns</span><span class="p">())</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="n">result_generator</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">run_inference</span><span class="p">(</span> |
| <span class="n">batch</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">except</span> <span class="ne">BaseException</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span><span class="o">.</span><span class="n">failed_batches_counter</span><span class="o">.</span><span class="n">inc</span><span class="p">()</span> |
| <span class="k">raise</span> <span class="n">e</span> |
| <span class="n">predictions</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">result_generator</span><span class="p">)</span> |
| |
| <span class="n">end_time</span> <span class="o">=</span> <span class="n">_to_microseconds</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_clock</span><span class="o">.</span><span class="n">time_ns</span><span class="p">())</span> |
| <span class="n">inference_latency</span> <span class="o">=</span> <span class="n">end_time</span> <span class="o">-</span> <span class="n">start_time</span> |
| <span class="n">num_bytes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_model_handler</span><span class="o">.</span><span class="n">get_num_bytes</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> |
| <span class="n">num_elements</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">num_elements</span><span class="p">,</span> <span class="n">num_bytes</span><span class="p">,</span> <span class="n">inference_latency</span><span class="p">)</span> |
| |
| <span class="k">return</span> <span class="n">predictions</span> |
| |
| <span class="k">def</span> <span class="nf">process</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">,</span> <span class="n">si_model_metadata</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">ModelMetadata</span><span class="p">]):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> When side input is enabled:</span> |
| <span class="sd"> The method checks if the side input model has been updated, and if so,</span> |
| <span class="sd"> updates the model and runs inference on the batch of data. If the</span> |
| <span class="sd"> side input is empty or the model has not been updated, the method</span> |
| <span class="sd"> simply runs inference on the batch of data.</span> |
| <span class="sd"> """</span> |
| <span class="k">if</span> <span class="n">si_model_metadata</span><span class="p">:</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">si_model_metadata</span><span class="p">,</span> <span class="n">beam</span><span class="o">.</span><span class="n">pvalue</span><span class="o">.</span><span class="n">EmptySideInput</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">update_model</span><span class="p">(</span><span class="n">side_input_model_path</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_run_inference</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">_side_input_path</span> <span class="o">!=</span> <span class="n">si_model_metadata</span><span class="o">.</span><span class="n">model_id</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_side_input_path</span> <span class="o">=</span> <span class="n">si_model_metadata</span><span class="o">.</span><span class="n">model_id</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_metrics_collector</span><span class="p">(</span> |
| <span class="n">prefix</span><span class="o">=</span><span class="n">si_model_metadata</span><span class="o">.</span><span class="n">model_name</span><span class="p">)</span> |
| <span class="k">with</span> <span class="n">threading</span><span class="o">.</span><span class="n">Lock</span><span class="p">():</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">update_model</span><span class="p">(</span><span class="n">si_model_metadata</span><span class="o">.</span><span class="n">model_id</span><span class="p">)</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_run_inference</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_run_inference</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">finish_bundle</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="c1"># TODO(https://github.com/apache/beam/issues/21435): Figure out why there</span> |
| <span class="c1"># is a cache.</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_metrics_collector</span><span class="o">.</span><span class="n">update_metrics_with_cache</span><span class="p">()</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_is_darwin</span><span class="p">()</span> <span class="o">-></span> <span class="nb">bool</span><span class="p">:</span> |
| <span class="k">return</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span> <span class="o">==</span> <span class="s1">'darwin'</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_get_current_process_memory_in_bytes</span><span class="p">():</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> memory usage in bytes.</span> |
| <span class="sd"> """</span> |
| |
| <span class="k">if</span> <span class="n">resource</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="n">usage</span> <span class="o">=</span> <span class="n">resource</span><span class="o">.</span><span class="n">getrusage</span><span class="p">(</span><span class="n">resource</span><span class="o">.</span><span class="n">RUSAGE_SELF</span><span class="p">)</span><span class="o">.</span><span class="n">ru_maxrss</span> |
| <span class="k">if</span> <span class="n">_is_darwin</span><span class="p">():</span> |
| <span class="k">return</span> <span class="n">usage</span> |
| <span class="k">return</span> <span class="n">usage</span> <span class="o">*</span> <span class="mi">1024</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">logging</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span> |
| <span class="s1">'Resource module is not available for current platform, '</span> |
| <span class="s1">'memory usage cannot be fetched.'</span><span class="p">)</span> |
| <span class="k">return</span> <span class="mi">0</span> |
| </pre></div> |
| |
| </div> |
| |
| </div> |
| <footer> |
| |
| |
| <hr/> |
| |
| <div role="contentinfo"> |
| <p> |
| © Copyright |
| |
| </p> |
| </div> |
| Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>. |
| |
| </footer> |
| |
| </div> |
| </div> |
| |
| </section> |
| |
| </div> |
| |
| |
| |
| <script type="text/javascript"> |
| jQuery(function () { |
| SphinxRtdTheme.Navigation.enable(true); |
| }); |
| </script> |
| |
| |
| |
| |
| |
| |
| </body> |
| </html> |
