chalkruntime 3.36.2__tar.gz → 3.36.4__tar.gz

{chalkruntime-3.36.2 → chalkruntime-3.36.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkruntime
-Version: 3.36.2
+Version: 3.36.4
 Summary: Runtime support library for Chalk AI
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{chalkruntime-3.36.2 → chalkruntime-3.36.4}/chalkruntime/server/entrypoint.py

@@ -79,6 +79,7 @@ from chalkruntime.loader.importer import (
 from chalkruntime.server.config import Config
 from chalkruntime.server.service import ImportGraphResult, InvokerService
 from chalkruntime.utils.datadog import get_datadog_kwargs
+from chalkruntime.utils.proto_bytes_cache import load_proto_bytes_with_cache
 
 if TYPE_CHECKING:
     from azure.core.credentials import TokenCredential
@@ -199,11 +200,20 @@ def _grpc_server_thread(grpc_query_server: CppGrpcServer, finished_event: Future
 
 def _export_from_bucket(filename: str, storage_client: StorageClient) -> export_pb2.Export:
     """Download a pre-generated Export proto from cloud storage."""
-    buf = io.BytesIO()
-    storage_client.download_object(filename, buf)
-    buf.seek(0)
+
+    def _download(uri: str) -> bytes:
+        buf = io.BytesIO()
+        storage_client.download_object(uri, buf)
+        return buf.getvalue()
+
+    proto_export_bytes = load_proto_bytes_with_cache(
+        uri=filename,
+        download=_download,
+        cache_dir=os.getenv("CHALK_PROTO_GRAPH_CACHE_DIR"),
+        proto_label="proto export",
+    )
     proto_export = export_pb2.Export()
-    proto_export.ParseFromString(buf.read())
+    proto_export.ParseFromString(proto_export_bytes)
     return proto_export
 
 

chalkruntime-3.36.4/chalkruntime/utils/proto_bytes_cache.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+
+from chalkruntime.utils.posix_file_lock import posix_file_lock
+
+_logger = logging.getLogger(__name__)
+
+DEFAULT_PROTO_BYTES_CACHE_DIR = "/tmp/chalk_proto_graph_cache"
+
+
+def load_proto_bytes_with_cache(
+    *,
+    uri: str,
+    download: Callable[[str], bytes],
+    cache_dir: str | os.PathLike[str] | None = None,
+    proto_label: str = "proto bytes",
+) -> bytes:
+    """Return downloaded proto bytes for `uri`, using an on-disk cache + file lock."""
+    resolved_cache_dir = Path(cache_dir) if cache_dir else Path(DEFAULT_PROTO_BYTES_CACHE_DIR)
+    cache_key = hashlib.sha256(uri.encode("utf-8")).hexdigest()
+    cache_path = resolved_cache_dir / f"{cache_key}.pb"
+    lock_path = resolved_cache_dir / f"{cache_key}.lock"
+
+    with posix_file_lock(lock_path):
+        if cache_path.exists():
+            try:
+                cached_bytes = cache_path.read_bytes()
+            except OSError:
+                _logger.warning(
+                    "Failed to read %s cache at %s; re-downloading from %s",
+                    proto_label,
+                    cache_path,
+                    uri,
+                    exc_info=True,
+                )
+            else:
+                _logger.info(
+                    "Loaded %s (%d bytes) from cache %s for uri=%s",
+                    proto_label,
+                    len(cached_bytes),
+                    cache_path,
+                    uri,
+                )
+                return cached_bytes
+
+        proto_bytes = download(uri)
+        _logger.info(
+            "Downloaded %s (%d bytes) from %s; writing cache to %s",
+            proto_label,
+            len(proto_bytes),
+            uri,
+            cache_path,
+        )
+        _atomic_write_bytes(cache_path, proto_bytes)
+        return proto_bytes
+
+
+def _atomic_write_bytes(path: Path, data: bytes) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp_name = tempfile.mkstemp(
+        dir=path.parent,
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+    )
+    try:
+        with os.fdopen(fd, "wb") as tmp:
+            tmp.write(data)
+            tmp.flush()
+            os.fsync(tmp.fileno())
+        os.replace(tmp_name, path)
+    except BaseException:
+        try:
+            os.unlink(tmp_name)
+        except FileNotFoundError:
+            pass
+        raise

{chalkruntime-3.36.2 → chalkruntime-3.36.4}/chalkruntime.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkruntime
-Version: 3.36.2
+Version: 3.36.4
 Summary: Runtime support library for Chalk AI
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{chalkruntime-3.36.2 → chalkruntime-3.36.4}/chalkruntime.egg-info/SOURCES.txt

@@ -141,6 +141,7 @@ chalkruntime/utils/contextvars.py
 chalkruntime/utils/datadog.py
 chalkruntime/utils/internal_pl_utils.py
 chalkruntime/utils/posix_file_lock.py
+chalkruntime/utils/proto_bytes_cache.py
 chalkruntime/utils/tracing.py
 chalkruntime/utils/viztracer_profiling.py
 tests/test_branch_model_endpoint_config.py

{chalkruntime-3.36.2 → chalkruntime-3.36.4}/pyproject.toml

@@ -42,7 +42,7 @@ name = "chalkruntime"
 description = "Runtime support library for Chalk AI"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "3.36.2"
+version = "3.36.4"
 
 
 [tool.deptry]