changes for reporting on costing

mpharrigan · mpharrigan · commit 0338e4db1b4d · 2024-08-29T14:31:35.000-07:00
diff --git a/dev_tools/bloq-method-overrides-report.py b/dev_tools/bloq-method-overrides-report.py
@@ -43,7 +43,7 @@ def _call_graph(bc: Type[Bloq]):
     if annot['ssa'] != 'SympySymbolAllocator':
         print(f"{bc}.build_call_graph `ssa: 'SympySymbolAllocator'`")
     if annot['return'] != Set[ForwardRef('BloqCountT')]:  # type: ignore[misc]
-        print(f"{bc}.build_call_graph -> 'BloqCountT'")
+        print(f"{bc}.build_call_graph -> Set['BloqCountT'], not {annot['return']}")
 
 
 def report_call_graph_methods():
diff --git a/dev_tools/qualtran_dev_tools/bloq_report_card.py b/dev_tools/qualtran_dev_tools/bloq_report_card.py
@@ -19,6 +19,8 @@
 import pandas.io.formats.style
 
 from qualtran import Bloq, BloqExample
+from qualtran.resource_counting import get_cost_value, QubitCount
+from qualtran.simulation.tensor import cbloq_to_quimb
 from qualtran.testing import (
     BloqCheckResult,
     check_bloq_example_decompose,
@@ -134,3 +136,44 @@ def summarize_results(report_card: pd.DataFrame) -> pd.DataFrame:
     )
     summary.columns = [v.name.lower() for v in summary.columns]
     return summary
+
+
+def report_on_tensors(name: str, cls_name: str, bloq: Bloq, cxn) -> None:
+    """Get timing information for tensor functionality.
+
+    This should be used with `ExecuteWithTimeout`. The resultant
+    record dictionary is sent over `cxn`.
+    """
+    record: Dict[str, Any] = {'name': name, 'cls': cls_name}
+
+    try:
+        start = time.perf_counter()
+        flat = bloq.as_composite_bloq().flatten()
+        record['flat_dur'] = time.perf_counter() - start
+
+        start = time.perf_counter()
+        tn = cbloq_to_quimb(flat)
+        record['tn_dur'] = time.perf_counter() - start
+
+        start = time.perf_counter()
+        record['width'] = tn.contraction_width()
+        record['width_dur'] = time.perf_counter() - start
+
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        record['err'] = str(e)
+
+    cxn.send(record)
+
+
+def report_on_cost_timings(name: str, cls_name: str, bloq: Bloq, cxn) -> None:
+    record: Dict[str, Any] = {'name': name, 'cls': cls_name}
+
+    try:
+        start = time.perf_counter()
+        _ = get_cost_value(bloq, QubitCount())
+        record['qubitcount_dur'] = time.perf_counter() - start
+
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        record['err'] = str(e)
+
+    cxn.send(record)
diff --git a/dev_tools/qualtran_dev_tools/execute_with_timeout.py b/dev_tools/qualtran_dev_tools/execute_with_timeout.py
@@ -0,0 +1,123 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import multiprocessing.connection
+import time
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from attrs import define
+
+
+@define
+class _Pending:
+    """Helper dataclass to track currently executing processes in `ExecuteWithTimeout`."""
+
+    p: multiprocessing.Process
+    recv: multiprocessing.connection.Connection
+    start_time: float
+    kwargs: Dict[str, Any]
+
+
+class ExecuteWithTimeout:
+    """Execute tasks in processes where each task will be killed if it exceeds `timeout`.
+
+    Seemingly all the existing "timeout" parameters in the various built-in concurrency
+    primitives in Python won't actually terminate the process. This one does.
+    """
+
+    def __init__(self, timeout: float, max_workers: int):
+        self.timeout = timeout
+        self.max_workers = max_workers
+
+        self.queued: List[Tuple[Callable, Dict[str, Any]]] = []
+        self.pending: List[_Pending] = []
+
+    @property
+    def work_to_be_done(self) -> int:
+        """The number of tasks currently executing or queued."""
+        return len(self.queued) + len(self.pending)
+
+    def submit(self, func: Callable, kwargs: Dict[str, Any]) -> None:
+        """Add a task to the queue.
+
+        `func` must be a callable that can accept `kwargs` in addition to
+        a keyword argument `cxn` which is a multiprocessing `Connection` object that forms
+        the sending-half of a `mp.Pipe`. The callable must call `cxn.send(...)`
+        to return a result.
+        """
+        self.queued.append((func, kwargs))
+
+    def _submit_from_queue(self):
+        # helper method that takes an item from the queue, launches a process,
+        # and records it in the `pending` attribute. This must only be called
+        # if we're allowed to spawn a new process.
+        func, kwargs = self.queued.pop(0)
+        recv, send = multiprocessing.Pipe(duplex=False)
+        kwargs['cxn'] = send
+        p = multiprocessing.Process(target=func, kwargs=kwargs)
+        start_time = time.time()
+        p.start()
+        self.pending.append(_Pending(p=p, recv=recv, start_time=start_time, kwargs=kwargs))
+
+    def _scan_pendings(self) -> Optional[_Pending]:
+        # helper method that goes through the currently pending tasks, terminates the ones
+        # that have been going on too long, and accounts for ones that have finished.
+        # Returns the `_Pending` of the killed or completed job or `None` if each pending
+        # task is still running but none have exceeded the timeout.
+        for i in range(len(self.pending)):
+            pen = self.pending[i]
+
+            if not pen.p.is_alive():
+                self.pending.pop(i)
+                pen.p.join()
+                return pen
+
+            if time.time() - pen.start_time > self.timeout:
+                pen.p.terminate()
+                self.pending.pop(i)
+                return pen
+
+        return None
+
+    def next_result(self) -> Tuple[Dict[str, Any], Optional[Any]]:
+        """Get the next available result.
+
+        This call is blocking, but should never take longer than `self.timeout`. This should
+        be called in a loop to make sure the queue continues to be processed.
+
+        Returns:
+            task kwargs: The keyword arguments used to submit the task.
+            result: If the process finished successfully, this is the object that was
+                sent through the multiprocessing pipe as the result. Otherwise, the result
+                is None.
+        """
+        while len(self.queued) > 0 and len(self.pending) < self.max_workers:
+            self._submit_from_queue()
+
+        while True:
+            finished = self._scan_pendings()
+            if finished is not None:
+                break
+
+        if finished.p.exitcode == 0:
+            result = finished.recv.recv()
+        else:
+            result = None
+
+        finished.recv.close()
+
+        while len(self.queued) > 0 and len(self.pending) < self.max_workers:
+            self._submit_from_queue()
+
+        return (finished.kwargs, result)
diff --git a/dev_tools/qualtran_dev_tools/execute_with_timeout_test.py b/dev_tools/qualtran_dev_tools/execute_with_timeout_test.py
@@ -0,0 +1,38 @@
+#  Copyright 2024 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import time
+
+from .execute_with_timeout import ExecuteWithTimeout
+
+
+def a_long_function(n_seconds: int, cxn) -> None:
+    time.sleep(n_seconds)
+    cxn.send("Done")
+
+
+def test_execute_with_timeout():
+    exec = ExecuteWithTimeout(timeout=1, max_workers=1)
+
+    for ns in [0.1, 100]:
+        exec.submit(a_long_function, {'n_seconds': ns})
+
+    results = []
+    while exec.work_to_be_done:
+        kwargs, result = exec.next_result()
+        if result is None:
+            results.append('Timeout')
+        else:
+            results.append(result)
+
+    assert set(results) == {'Done', 'Timeout'}
diff --git a/dev_tools/qualtran_dev_tools/tensor_report_card.py b/dev_tools/qualtran_dev_tools/tensor_report_card.py
@@ -11,142 +11,8 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-import multiprocessing.connection
 import time
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-from attrs import define
+from typing import Any, Dict
 
 from qualtran import Bloq
 from qualtran.simulation.tensor import cbloq_to_quimb
-
-
-@define
-class _Pending:
-    """Helper dataclass to track currently executing processes in `ExecuteWithTimeout`."""
-
-    p: multiprocessing.Process
-    recv: multiprocessing.connection.Connection
-    start_time: float
-    kwargs: Dict[str, Any]
-
-
-class ExecuteWithTimeout:
-    """Execute tasks in processes where each task will be killed if it exceeds `timeout`.
-
-    Seemingly all the existing "timeout" parameters in the various built-in concurrency
-    primitives in Python won't actually terminate the process. This one does.
-    """
-
-    def __init__(self, timeout: float, max_workers: int):
-        self.timeout = timeout
-        self.max_workers = max_workers
-
-        self.queued: List[Tuple[Callable, Dict[str, Any]]] = []
-        self.pending: List[_Pending] = []
-
-    @property
-    def work_to_be_done(self) -> int:
-        """The number of tasks currently executing or queued."""
-        return len(self.queued) + len(self.pending)
-
-    def submit(self, func: Callable, kwargs: Dict[str, Any]) -> None:
-        """Add a task to the queue.
-
-        `func` must be a callable that can accept `kwargs` in addition to
-        a keyword argument `cxn` which is a multiprocessing `Connection` object that forms
-        the sending-half of a `mp.Pipe`. The callable must call `cxn.send(...)`
-        to return a result.
-        """
-        self.queued.append((func, kwargs))
-
-    def _submit_from_queue(self):
-        # helper method that takes an item from the queue, launches a process,
-        # and records it in the `pending` attribute. This must only be called
-        # if we're allowed to spawn a new process.
-        func, kwargs = self.queued.pop(0)
-        recv, send = multiprocessing.Pipe(duplex=False)
-        kwargs['cxn'] = send
-        p = multiprocessing.Process(target=func, kwargs=kwargs)
-        start_time = time.time()
-        p.start()
-        self.pending.append(_Pending(p=p, recv=recv, start_time=start_time, kwargs=kwargs))
-
-    def _scan_pendings(self) -> Optional[_Pending]:
-        # helper method that goes through the currently pending tasks, terminates the ones
-        # that have been going on too long, and accounts for ones that have finished.
-        # Returns the `_Pending` of the killed or completed job or `None` if each pending
-        # task is still running but none have exceeded the timeout.
-        for i in range(len(self.pending)):
-            pen = self.pending[i]
-
-            if not pen.p.is_alive():
-                self.pending.pop(i)
-                pen.p.join()
-                return pen
-
-            if time.time() - pen.start_time > self.timeout:
-                pen.p.terminate()
-                self.pending.pop(i)
-                return pen
-
-        return None
-
-    def next_result(self) -> Tuple[Dict[str, Any], Optional[Any]]:
-        """Get the next available result.
-
-        This call is blocking, but should never take longer than `self.timeout`. This should
-        be called in a loop to make sure the queue continues to be processed.
-
-        Returns:
-            task kwargs: The keyword arguments used to submit the task.
-            result: If the process finished successfully, this is the object that was
-                sent through the multiprocessing pipe as the result. Otherwise, the result
-                is None.
-        """
-        while len(self.queued) > 0 and len(self.pending) < self.max_workers:
-            self._submit_from_queue()
-
-        while True:
-            finished = self._scan_pendings()
-            if finished is not None:
-                break
-
-        if finished.p.exitcode == 0:
-            result = finished.recv.recv()
-        else:
-            result = None
-
-        finished.recv.close()
-
-        while len(self.queued) > 0 and len(self.pending) < self.max_workers:
-            self._submit_from_queue()
-
-        return (finished.kwargs, result)
-
-
-def report_on_tensors(name: str, cls_name: str, bloq: Bloq, cxn) -> None:
-    """Get timing information for tensor functionality.
-
-    This should be used with `ExecuteWithTimeout`. The resultant
-    record dictionary is sent over `cxn`.
-    """
-    record: Dict[str, Any] = {'name': name, 'cls': cls_name}
-
-    try:
-        start = time.perf_counter()
-        flat = bloq.as_composite_bloq().flatten()
-        record['flat_dur'] = time.perf_counter() - start
-
-        start = time.perf_counter()
-        tn = cbloq_to_quimb(flat)
-        record['tn_dur'] = time.perf_counter() - start
-
-        start = time.perf_counter()
-        record['width'] = tn.contraction_width()
-        record['width_dur'] = time.perf_counter() - start
-
-    except Exception as e:  # pylint: disable=broad-exception-caught
-        record['err'] = str(e)
-
-    cxn.send(record)
