datachain 0.31.2__py3-none-any.whl → 0.31.3__py3-none-any.whl

datachain/catalog/catalog.py

@@ -144,19 +144,26 @@ def shutdown_process(
             return proc.wait()
 
 
-def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
+def process_output(stream: IO[bytes], callback: Callable[[str], None]) -> None:
     buffer = b""
-    while byt := stream.read(1):  # Read one byte at a time
-        buffer += byt
 
-        if byt in (b"\n", b"\r"):  # Check for newline or carriage return
-            line = buffer.decode("utf-8")
-            callback(line)
-            buffer = b""  # Clear buffer for next line
+    try:
+        while byt := stream.read(1):  # Read one byte at a time
+            buffer += byt
 
-    if buffer:  # Handle any remaining data in the buffer
-        line = buffer.decode("utf-8")
-        callback(line)
+            if byt in (b"\n", b"\r"):  # Check for newline or carriage return
+                line = buffer.decode("utf-8", errors="replace")
+                callback(line)
+                buffer = b""  # Clear buffer for the next line
+
+        if buffer:  # Handle any remaining data in the buffer
+            line = buffer.decode("utf-8", errors="replace")
+            callback(line)
+    finally:
+        try:
+            stream.close()  # Ensure output is closed
+        except Exception:  # noqa: BLE001, S110
+            pass
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1760,13 +1767,13 @@ class Catalog:
             recursive=recursive,
         )
 
+    @staticmethod
     def query(
-        self,
         query_script: str,
         env: Optional[Mapping[str, str]] = None,
         python_executable: str = sys.executable,
-        capture_output: bool = False,
-        output_hook: Callable[[str], None] = noop,
+        stdout_callback: Optional[Callable[[str], None]] = None,
+        stderr_callback: Optional[Callable[[str], None]] = None,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
         interrupt_timeout: Optional[int] = None,
@@ -1781,13 +1788,18 @@ class Catalog:
             },
         )
         popen_kwargs: dict[str, Any] = {}
-        if capture_output:
-            popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
+
+        if stdout_callback is not None:
+            popen_kwargs = {"stdout": subprocess.PIPE}
+        if stderr_callback is not None:
+            popen_kwargs["stderr"] = subprocess.PIPE
 
         def raise_termination_signal(sig: int, _: Any) -> NoReturn:
             raise TerminationSignal(sig)
 
-        thread: Optional[Thread] = None
+        stdout_thread: Optional[Thread] = None
+        stderr_thread: Optional[Thread] = None
+
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             logger.info("Starting process %s", proc.pid)
 
@@ -1801,10 +1813,20 @@ class Catalog:
             orig_sigterm_handler = signal.getsignal(signal.SIGTERM)
             signal.signal(signal.SIGTERM, raise_termination_signal)
             try:
-                if capture_output:
-                    args = (proc.stdout, output_hook)
-                    thread = Thread(target=_process_stream, args=args, daemon=True)
-                    thread.start()
+                if stdout_callback is not None:
+                    stdout_thread = Thread(
+                        target=process_output,
+                        args=(proc.stdout, stdout_callback),
+                        daemon=True,
+                    )
+                    stdout_thread.start()
+                if stderr_callback is not None:
+                    stderr_thread = Thread(
+                        target=process_output,
+                        args=(proc.stderr, stderr_callback),
+                        daemon=True,
+                    )
+                    stderr_thread.start()
 
                 proc.wait()
             except TerminationSignal as exc:
@@ -1822,8 +1844,22 @@ class Catalog:
             finally:
                 signal.signal(signal.SIGTERM, orig_sigterm_handler)
                 signal.signal(signal.SIGINT, orig_sigint_handler)
-                if thread:
-                    thread.join()  # wait for the reader thread
+                # wait for the reader thread
+                thread_join_timeout_seconds = 30
+                if stdout_thread is not None:
+                    stdout_thread.join(timeout=thread_join_timeout_seconds)
+                    if stdout_thread.is_alive():
+                        logger.warning(
+                            "stdout thread is still alive after %s seconds",
+                            thread_join_timeout_seconds,
+                        )
+                if stderr_thread is not None:
+                    stderr_thread.join(timeout=thread_join_timeout_seconds)
+                    if stderr_thread.is_alive():
+                        logger.warning(
+                            "stderr thread is still alive after %s seconds",
+                            thread_join_timeout_seconds,
+                        )
 
         logger.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode in (

{datachain-0.31.2.dist-info → datachain-0.31.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.31.2
+Version: 0.31.3
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.31.2.dist-info → datachain-0.31.3.dist-info}/RECORD

@@ -21,7 +21,7 @@ datachain/studio.py,sha256=IS8o4BZnhUo73Bd8m4CJxFc5utdmh2miIs25WswkFBA,15283
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=5ehFeqXau7MFmGUQRsjRyPfDMPoOF1ojpfVciYUo5fE,15659
 datachain/catalog/__init__.py,sha256=9NBaywvAOaXdkyqiHjbBEiXs7JImR1OJsY9r8D5Q16g,403
-datachain/catalog/catalog.py,sha256=a1AN6eDHWWzII1wi46T_1JvTsW1AeMudwR_6sVQ4f7I,67588
+datachain/catalog/catalog.py,sha256=oI4YBuuOJGVx_Fp1cDoFb56lPV7Or27ZquzR8oM1m3Y,69133
 datachain/catalog/datasource.py,sha256=IkGMh0Ttg6Q-9DWfU_H05WUnZepbGa28HYleECi6K7I,1353
 datachain/catalog/loader.py,sha256=53VnuSRkt_CO9RdlHWkzQsPF55qMxcXvEm3ecsZREw8,6150
 datachain/cli/__init__.py,sha256=so3WxEQF03KdGvjav15Sw7a6-lriiE24uDSGbBDBp8o,8298
@@ -161,9 +161,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.31.2.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.31.2.dist-info/METADATA,sha256=ALo4Vp6w2VSanACVy1xv6aHWzbdasSKzD2U8_SybXBU,13898
-datachain-0.31.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.31.2.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.31.2.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.31.2.dist-info/RECORD,,
+datachain-0.31.3.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.31.3.dist-info/METADATA,sha256=dZjBfjFrwEjatAGqlONnD8fIO6H-2Njw1rHyvvZQ1kU,13898
+datachain-0.31.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.31.3.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.31.3.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.31.3.dist-info/RECORD,,