|
242 | 242 | <div class="pytorch-left-menu-search">
|
243 | 243 |
|
244 | 244 | <div class="version">
|
245 |
| - <a href='https://pytorch.org/docs/versions.html'>main (2.4.0a0+git6b8205d ) ▼</a> |
| 245 | + <a href='https://pytorch.org/docs/versions.html'>main (2.4.0a0+git2b1ba0c ) ▼</a> |
246 | 246 | </div>
|
247 | 247 |
|
248 | 248 |
|
|
482 | 482 | <h1>Source code for torch._utils</h1><div class="highlight"><pre>
|
483 | 483 | <span></span><span class="kn">import</span> <span class="nn">copyreg</span>
|
484 | 484 | <span class="kn">import</span> <span class="nn">functools</span>
|
| 485 | +<span class="kn">import</span> <span class="nn">logging</span> |
485 | 486 | <span class="kn">import</span> <span class="nn">sys</span>
|
486 | 487 | <span class="kn">import</span> <span class="nn">traceback</span>
|
487 | 488 | <span class="kn">import</span> <span class="nn">warnings</span>
|
488 | 489 | <span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">defaultdict</span>
|
489 |
| -<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">DefaultDict</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span> |
| 490 | +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">DefaultDict</span><span class="p">,</span> <span class="n">Generic</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span> |
| 491 | + |
| 492 | +<span class="kn">from</span> <span class="nn">typing_extensions</span> <span class="kn">import</span> <span class="n">ParamSpec</span> |
490 | 493 |
|
491 | 494 | <span class="kn">import</span> <span class="nn">torch</span>
|
492 | 495 |
|
@@ -1417,6 +1420,28 @@ <h1>Source code for torch._utils</h1><div class="highlight"><pre>
|
1417 | 1420 |
|
1418 | 1421 | <span class="k">def</span> <span class="nf">get_calls</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">List</span><span class="p">:</span>
|
1419 | 1422 | <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">call_order</span>
|
| 1423 | + |
| 1424 | + |
| 1425 | +<span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span> |
| 1426 | +<span class="n">P</span> <span class="o">=</span> <span class="n">ParamSpec</span><span class="p">(</span><span class="s2">"P"</span><span class="p">)</span> |
| 1427 | + |
| 1428 | + |
| 1429 | +<span class="k">class</span> <span class="nc">CallbackRegistry</span><span class="p">(</span><span class="n">Generic</span><span class="p">[</span><span class="n">P</span><span class="p">]):</span> |
| 1430 | + <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">name</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span> |
| 1431 | + <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="n">name</span> |
| 1432 | + <span class="bp">self</span><span class="o">.</span><span class="n">callback_list</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Callable</span><span class="p">[</span><span class="n">P</span><span class="p">,</span> <span class="kc">None</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[]</span> |
| 1433 | + |
| 1434 | + <span class="k">def</span> <span class="nf">add_callback</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">cb</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[</span><span class="n">P</span><span class="p">,</span> <span class="kc">None</span><span class="p">])</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span> |
| 1435 | + <span class="bp">self</span><span class="o">.</span><span class="n">callback_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cb</span><span class="p">)</span> |
| 1436 | + |
| 1437 | + <span class="k">def</span> <span class="nf">fire_callbacks</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">P</span><span class="o">.</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">P</span><span class="o">.</span><span class="n">kwargs</span><span class="p">)</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span> |
| 1438 | + <span class="k">for</span> <span class="n">cb</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">callback_list</span><span class="p">:</span> |
| 1439 | + <span class="k">try</span><span class="p">:</span> |
| 1440 | + <span class="n">cb</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span> |
| 1441 | + <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span> |
| 1442 | + <span class="n">logger</span><span class="o">.</span><span class="n">exception</span><span class="p">(</span> |
| 1443 | + <span class="s2">"Exception in callback for </span><span class="si">%s</span><span class="s2"> registered with gpu trace"</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span> |
| 1444 | + <span class="p">)</span> |
1420 | 1445 | </pre></div>
|
1421 | 1446 |
|
1422 | 1447 | </article>
|
|
0 commit comments