chalk-remote-call-python 1.5.1__tar.gz → 1.6.1__tar.gz

{chalk_remote_call_python-1.5.1 → chalk_remote_call_python-1.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.5.1
+Version: 1.6.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalk_remote_call_python-1.5.1 → chalk_remote_call_python-1.6.1}/chalk-remote-call-rs/chalk-remote-call-server/src/python_bridge.rs

@@ -15,11 +15,12 @@ pub struct PythonHandler {
     arg_names: Option<Vec<String>>,
 }
 
-/// One caller's input: request bytes + gRPC metadata + peer address.
+/// One caller's input: request bytes + function name + gRPC metadata + peer address.
 /// Used identically in both the single-request and coalesced paths; the
 /// coalesced path just has N of them.
 pub struct CallerInput {
     pub ipc_bytes: Vec<u8>,
+    pub function_name: String,
     pub metadata: HashMap<String, String>,
     pub peer: String,
 }
@@ -40,6 +41,7 @@ impl PythonHandler {
     pub async fn call(
         &self,
         ipc_bytes: Vec<u8>,
+        function_name: String,
         metadata: HashMap<String, String>,
         peer: String,
     ) -> Result<Vec<Vec<u8>>, PythonError> {
@@ -50,7 +52,9 @@ impl PythonHandler {
         tokio::task::spawn_blocking(move || {
             Python::attach(|py| -> Result<Vec<Vec<u8>>, PythonError> {
                 let py_bytes = PyBytes::new(py, &ipc_bytes).into_any().unbind();
-                let py_ctx = build_context(py, &metadata, &peer)?.into_any().unbind();
+                let py_ctx = build_context(py, &function_name, &metadata, &peer)?
+                    .into_any()
+                    .unbind();
                 let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
                 invoke_process_fn(py, &process_fn, &handler, py_bytes, py_arg_names, py_ctx)
             })
@@ -79,7 +83,7 @@ impl PythonHandler {
                         .append(PyBytes::new(py, &c.ipc_bytes))
                         .map_err(|e| into_python_error(py, e))?;
                     py_contexts
-                        .append(build_context(py, &c.metadata, &c.peer)?)
+                        .append(build_context(py, &c.function_name, &c.metadata, &c.peer)?)
                         .map_err(|e| into_python_error(py, e))?;
                 }
                 let py_arg_names = build_arg_names(py, arg_names.as_deref())?;
@@ -100,9 +104,10 @@ impl PythonHandler {
 
 // ---- Shared helpers -------------------------------------------------------
 
-/// Build a `{"peer": str, "metadata": dict}` context dict for one caller.
+/// Build a `{"function_name": str, "peer": str, "metadata": dict}` context dict for one caller.
 fn build_context<'py>(
     py: Python<'py>,
+    function_name: &str,
     metadata: &HashMap<String, String>,
     peer: &str,
 ) -> Result<Bound<'py, PyDict>, PythonError> {
@@ -113,6 +118,8 @@ fn build_context<'py>(
             .map_err(|e| into_python_error(py, e))?;
     }
     let ctx = PyDict::new(py);
+    ctx.set_item("function_name", function_name)
+        .map_err(|e| into_python_error(py, e))?;
     ctx.set_item("peer", peer)
         .map_err(|e| into_python_error(py, e))?;
     ctx.set_item("metadata", meta_dict)

{chalk_remote_call_python-1.5.1 → chalk_remote_call_python-1.6.1}/chalk-remote-call-rs/chalk-remote-call-server/src/service.rs

@@ -40,7 +40,11 @@ impl RemoteCallService for RemoteCallServiceImpl {
 
         // Accumulate all feather_stream bytes from the request stream
         let mut all_bytes: Vec<u8> = Vec::new();
+        let mut function_name = String::new();
         while let Some(req) = stream.message().await? {
+            if function_name.is_empty() && !req.name.is_empty() {
+                function_name = req.name.clone();
+            }
             all_bytes.extend_from_slice(&req.feather_stream);
             if all_bytes.len() > MAX_REQUEST_SIZE {
                 return Err(Status::resource_exhausted(format!(
@@ -74,6 +78,7 @@ impl RemoteCallService for RemoteCallServiceImpl {
                 .submit(BufferedCall {
                     input: CallerInput {
                         ipc_bytes: all_bytes,
+                        function_name,
                         metadata,
                         peer,
                     },
@@ -86,7 +91,7 @@ impl RemoteCallService for RemoteCallServiceImpl {
         // Single-request path.
         let handler = self.python_handler.clone();
         tokio::spawn(async move {
-            match handler.call(all_bytes, metadata, peer).await {
+            match handler.call(all_bytes, function_name, metadata, peer).await {
                 Ok(response_chunks) => {
                     for chunk in response_chunks {
                         let msg = CallFunctionResponse {

chalk_remote_call_python-1.6.1/chalk_remote_call/_version.py

@@ -0,0 +1 @@
+__version__ = "1.6.1"

chalk_remote_call_python-1.6.1/chalk_remote_call/tracing.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import atexit
+import contextlib
+import importlib
+import os
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, cast
+
+_TRACER_NAME = "chalk_remote_call"
+_TRACE_CONTEXT_METADATA_KEYS = ("traceparent", "tracestate", "baggage")
+_REMOTE_FUNCTION_TRACE_POLICY_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_POLICY"
+_REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR = "CHALK_REMOTE_FUNCTION_TRACE_SAMPLE_RATE"
+_TRACE_POLICY_PARENT_BASED_ALWAYS_OFF = "parentbased_always_off"
+_TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO = "parentbased_traceidratio"
+_TRACE_POLICY_ALWAYS_OFF = "always_off"
+_TRACE_POLICIES = frozenset(
+    {
+        _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF,
+        _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO,
+        _TRACE_POLICY_ALWAYS_OFF,
+    }
+)
+_missing_otel_modules = object()
+_missing_always_off_tracer = object()
+_otel_modules: tuple[Any, Any, Any] | object | None = None
+_always_off_tracer: Any | object | None = None
+_runtime_tracing_configured = False
+_runtime_tracer_provider: Any | None = None
+
+
+def _raw_metadata(context_metadata: Any) -> Mapping[str, Any] | None:
+    if not isinstance(context_metadata, Mapping):
+        return None
+
+    raw = context_metadata.get("metadata", context_metadata)
+    if isinstance(raw, Mapping):
+        return raw
+    return None
+
+
+def _collect_trace_metadata(
+    context_metadata: Any,
+) -> tuple[Mapping[str, Any] | None, Mapping[str, Any] | None]:
+    if isinstance(context_metadata, Mapping):
+        first = _raw_metadata(context_metadata)
+        if first is not None and first.get("traceparent"):
+            return first, first
+        return first, None
+
+    first: Mapping[str, Any] | None = None
+    if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
+        for item in context_metadata:
+            metadata = _raw_metadata(item)
+            if metadata is None:
+                continue
+
+            if first is None:
+                first = metadata
+            if metadata.get("traceparent"):
+                return first, metadata
+
+    return first, None
+
+
+def _trace_policy() -> str | None:
+    raw_policy = os.environ.get(_REMOTE_FUNCTION_TRACE_POLICY_ENV_VAR)
+    if raw_policy is None:
+        return None
+
+    policy = raw_policy.strip().lower()
+    if policy in _TRACE_POLICIES:
+        return policy
+    return None
+
+
+def _trace_sample_rate() -> float | None:
+    raw_rate = os.environ.get(_REMOTE_FUNCTION_TRACE_SAMPLE_RATE_ENV_VAR)
+    if raw_rate is None:
+        return None
+
+    try:
+        rate = float(raw_rate)
+    except ValueError:
+        return None
+
+    if rate < 0.0 or rate > 1.0:
+        return None
+    return rate
+
+
+def _otel_exporter_configured() -> bool:
+    return bool(os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") or os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"))
+
+
+def _configure_runtime_tracing(sample_rate: float) -> None:
+    global _runtime_tracing_configured, _runtime_tracer_provider
+    if _runtime_tracing_configured:
+        return
+
+    try:
+        resources = importlib.import_module("opentelemetry.sdk.resources")
+        sdk_trace = importlib.import_module("opentelemetry.sdk.trace")
+        sdk_sampling = importlib.import_module("opentelemetry.sdk.trace.sampling")
+        trace_api = importlib.import_module("opentelemetry.trace")
+    except Exception:
+        _runtime_tracing_configured = True
+        return
+
+    try:
+        provider = sdk_trace.TracerProvider(
+            resource=resources.Resource.create({"service.name": os.environ.get("CHALK_SERVICE") or _TRACER_NAME}),
+            sampler=sdk_sampling.ParentBased(sdk_sampling.TraceIdRatioBased(sample_rate)),
+        )
+
+        if _otel_exporter_configured():
+            try:
+                otlp_exporter = importlib.import_module("opentelemetry.exporter.otlp.proto.grpc.trace_exporter")
+                sdk_export = importlib.import_module("opentelemetry.sdk.trace.export")
+                provider.add_span_processor(sdk_export.BatchSpanProcessor(otlp_exporter.OTLPSpanExporter()))
+            except Exception:
+                pass
+
+        trace_api.set_tracer_provider(provider)
+        _runtime_tracer_provider = provider
+        atexit.register(provider.shutdown)
+    except Exception:
+        pass
+    _runtime_tracing_configured = True
+
+
+def _sample_rate_for_inherited_parent(policy: str) -> float | None:
+    if policy == _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF:
+        return 0.0
+    if policy == _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO:
+        return _trace_sample_rate()
+    return None
+
+
+def _span_metadata(
+    first_metadata: Mapping[str, Any] | None,
+    traceparent_metadata: Mapping[str, Any] | None,
+    policy: str,
+) -> tuple[Mapping[str, Any] | None, Any | None, Any]:
+    if traceparent_metadata is not None:
+        carrier = _trace_carrier(traceparent_metadata)
+        if not carrier:
+            return None, None, None
+
+        otel_modules = _get_otel_modules()
+        if otel_modules is None:
+            return None, None, None
+        _, propagate, trace = otel_modules
+
+        parent_context = propagate.extract(carrier)
+        parent_span_context = trace.get_current_span(parent_context).get_span_context()
+        if not parent_span_context.is_valid:
+            return None, None, None
+
+        sample_rate = _sample_rate_for_inherited_parent(policy)
+        if sample_rate is None:
+            return None, None, None
+        _configure_runtime_tracing(sample_rate)
+        return traceparent_metadata, parent_context, otel_modules
+
+    if policy != _TRACE_POLICY_PARENT_BASED_TRACE_ID_RATIO:
+        return None, None, None
+
+    sample_rate = _trace_sample_rate()
+    if sample_rate is None:
+        return None, None, None
+
+    _configure_runtime_tracing(sample_rate)
+    otel_modules = _get_otel_modules()
+    if otel_modules is None:
+        return None, None, None
+    return first_metadata, None, otel_modules
+
+
+def _always_off_span_metadata(
+    first_metadata: Mapping[str, Any] | None,
+    traceparent_metadata: Mapping[str, Any] | None,
+) -> tuple[Mapping[str, Any] | None, Any | None, Any | None, Any | None]:
+    otel_modules = _get_otel_modules()
+    if otel_modules is None:
+        return None, None, None, None
+    _, propagate, trace = otel_modules
+
+    parent_context = None
+    if traceparent_metadata is not None:
+        carrier = _trace_carrier(traceparent_metadata)
+        if carrier:
+            extracted_context = propagate.extract(carrier)
+            parent_span_context = trace.get_current_span(extracted_context).get_span_context()
+            if parent_span_context.is_valid:
+                parent_context = extracted_context
+
+    tracer = _get_always_off_tracer()
+    if tracer is None:
+        return None, None, None, None
+    return traceparent_metadata or first_metadata, parent_context, otel_modules, tracer
+
+
+def _trace_carrier(metadata: Mapping[str, Any]) -> dict[str, str]:
+    return {
+        str(key).lower(): str(value)
+        for key, value in metadata.items()
+        if str(key).lower() in _TRACE_CONTEXT_METADATA_KEYS and value
+    }
+
+
+def _function_name(context_metadata: Any) -> str:
+    if isinstance(context_metadata, Mapping):
+        return str(context_metadata.get("function_name") or "") or "unknown"
+
+    if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
+        for item in context_metadata:
+            if isinstance(item, Mapping):
+                function_name = str(item.get("function_name") or "")
+                if function_name:
+                    return function_name
+
+    return "unknown"
+
+
+def _get_tracer(trace_api: Any) -> Any:
+    return trace_api.get_tracer(_TRACER_NAME)
+
+
+def _get_otel_modules() -> tuple[Any, Any, Any] | None:
+    global _otel_modules
+    if _otel_modules is None:
+        try:
+            _otel_modules = (
+                importlib.import_module("opentelemetry.context"),
+                importlib.import_module("opentelemetry.propagate"),
+                importlib.import_module("opentelemetry.trace"),
+            )
+        except ImportError:
+            _otel_modules = _missing_otel_modules
+
+    if _otel_modules is _missing_otel_modules:
+        return None
+    return cast(tuple[Any, Any, Any], _otel_modules)
+
+
+def _get_always_off_tracer() -> Any | None:
+    global _always_off_tracer
+    if _always_off_tracer is None:
+        try:
+            sdk_trace = importlib.import_module("opentelemetry.sdk.trace")
+            sdk_sampling = importlib.import_module("opentelemetry.sdk.trace.sampling")
+        except ImportError:
+            _always_off_tracer = _missing_always_off_tracer
+        else:
+            # Keep always-off local to this invocation instead of replacing the
+            # process-wide tracer provider.
+            provider = sdk_trace.TracerProvider(sampler=sdk_sampling.ALWAYS_OFF)
+            _always_off_tracer = provider.get_tracer(_TRACER_NAME)
+
+    if _always_off_tracer is _missing_always_off_tracer:
+        return None
+    return _always_off_tracer
+
+
+@contextlib.contextmanager
+def remote_function_invocation_span(
+    context_metadata: Any,
+    coalesced_count: int | None = None,
+) -> Iterator[None]:
+    policy = _trace_policy()
+    if policy is None:
+        yield
+        return
+
+    first_metadata, traceparent_metadata = _collect_trace_metadata(context_metadata)
+    if policy == _TRACE_POLICY_PARENT_BASED_ALWAYS_OFF and traceparent_metadata is None:
+        yield
+        return
+
+    if policy == _TRACE_POLICY_ALWAYS_OFF:
+        _metadata, parent_context, otel_modules, tracer = _always_off_span_metadata(
+            first_metadata,
+            traceparent_metadata,
+        )
+    else:
+        _metadata, parent_context, otel_modules = _span_metadata(
+            first_metadata,
+            traceparent_metadata,
+            policy,
+        )
+        tracer = _get_tracer(otel_modules[2]) if otel_modules is not None else None
+
+    if otel_modules is None:
+        yield
+        return
+    if tracer is None:
+        yield
+        return
+    otel_context, _, trace = otel_modules
+
+    SpanKind = trace.SpanKind
+    Status = trace.Status
+    StatusCode = trace.StatusCode
+    function_name = _function_name(context_metadata)
+    span_kwargs: dict[str, Any] = {"kind": SpanKind.SERVER}
+    if parent_context is not None:
+        span_kwargs["context"] = parent_context
+
+    with tracer.start_as_current_span("chalkcompute.remote_function.invoke", **span_kwargs) as span:
+        span.set_attribute("chalk.remote_function.name", function_name)
+        if coalesced_count is not None:
+            span.set_attribute(
+                "chalk.remote_function.coalesced_count",
+                coalesced_count,
+            )
+        if parent_context is not None and not span.get_span_context().is_valid:
+            token = otel_context.attach(parent_context)
+            try:
+                yield
+            finally:
+                otel_context.detach(token)
+            return
+        try:
+            yield
+        except Exception as exc:
+            span.record_exception(exc)
+            span.set_status(Status(StatusCode.ERROR, str(exc)))
+            raise

{chalk_remote_call_python-1.5.1 → chalk_remote_call_python-1.6.1}/chalk_remote_call_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.5.1
+Version: 1.6.1
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

chalk_remote_call_python-1.5.1/chalk_remote_call/_version.py

@@ -1 +0,0 @@
-__version__ = "1.5.1"

chalk_remote_call_python-1.5.1/chalk_remote_call/tracing.py

@@ -1,150 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import importlib
-import os
-from collections.abc import Callable, Iterator, Mapping, Sequence
-from typing import Any, cast
-
-_TRACER_NAME = "chalk_remote_call"
-_TRACE_CONTEXT_METADATA_KEYS = ("traceparent", "tracestate", "baggage")
-_REMOTE_FUNCTION_NAME_METADATA = "x-chalk-function-name"
-_REMOTE_FUNCTION_TRACING_ENV_VAR = "CHALK_ENABLE_REMOTE_FUNCTION_TRACING"
-_missing_sdk_tracer = object()
-_missing_otel_modules = object()
-_sdk_get_tracer: Callable[[], Any] | object | None = None
-_otel_modules: tuple[Any, Any, Any] | object | None = None
-
-
-def _raw_metadata(context_metadata: Any) -> Mapping[str, Any] | None:
-    if not isinstance(context_metadata, Mapping):
-        return None
-
-    raw = context_metadata.get("metadata", context_metadata)
-    if isinstance(raw, Mapping):
-        return raw
-    return None
-
-
-def _first_context_with_traceparent(context_metadata: Any) -> Mapping[str, Any] | None:
-    if isinstance(context_metadata, Sequence) and not isinstance(context_metadata, str | bytes):
-        for item in context_metadata:
-            metadata = _raw_metadata(item)
-            if metadata is not None and metadata.get("traceparent"):
-                return metadata
-        return None
-
-    metadata = _raw_metadata(context_metadata)
-    if metadata is not None and metadata.get("traceparent"):
-        return metadata
-    return None
-
-
-def _trace_carrier(metadata: Mapping[str, Any]) -> dict[str, str]:
-    return {
-        str(key).lower(): str(value)
-        for key, value in metadata.items()
-        if str(key).lower() in _TRACE_CONTEXT_METADATA_KEYS and value
-    }
-
-
-def _function_name(metadata: Mapping[str, Any]) -> str:
-    return str(metadata.get(_REMOTE_FUNCTION_NAME_METADATA) or "") or "unknown"
-
-
-def _get_tracer(trace_api: Any) -> Any:
-    global _sdk_get_tracer
-    if _sdk_get_tracer is None:
-        try:
-            sdk_tracing = importlib.import_module("chalkcompute._tracing")
-        except ImportError:
-            _sdk_get_tracer = _missing_sdk_tracer
-        else:
-            _sdk_get_tracer = sdk_tracing.get_tracer
-
-    sdk_get_tracer = _sdk_get_tracer
-    if callable(sdk_get_tracer):
-        try:
-            return sdk_get_tracer()
-        except ImportError:
-            return trace_api.get_tracer(_TRACER_NAME)
-    return trace_api.get_tracer(_TRACER_NAME)
-
-
-def _get_otel_modules() -> tuple[Any, Any, Any] | None:
-    global _otel_modules
-    if _otel_modules is None:
-        try:
-            _otel_modules = (
-                importlib.import_module("opentelemetry.context"),
-                importlib.import_module("opentelemetry.propagate"),
-                importlib.import_module("opentelemetry.trace"),
-            )
-        except ImportError:
-            _otel_modules = _missing_otel_modules
-
-    if _otel_modules is _missing_otel_modules:
-        return None
-    return cast(tuple[Any, Any, Any], _otel_modules)
-
-
-@contextlib.contextmanager
-def remote_function_invocation_span(
-    context_metadata: Any,
-    coalesced_count: int | None = None,
-) -> Iterator[None]:
-    if os.environ.get(_REMOTE_FUNCTION_TRACING_ENV_VAR, "") != "1":
-        yield
-        return
-
-    metadata = _first_context_with_traceparent(context_metadata)
-    if metadata is None:
-        yield
-        return
-
-    carrier = _trace_carrier(metadata)
-    if not carrier:
-        yield
-        return
-
-    otel_modules = _get_otel_modules()
-    if otel_modules is None:
-        yield
-        return
-    otel_context, propagate, trace = otel_modules
-
-    parent_context = propagate.extract(carrier)
-    parent_span_context = trace.get_current_span(parent_context).get_span_context()
-    if not parent_span_context.is_valid:
-        yield
-        return
-
-    SpanKind = trace.SpanKind
-    Status = trace.Status
-    StatusCode = trace.StatusCode
-    function_name = _function_name(metadata)
-    tracer = _get_tracer(trace)
-    with tracer.start_as_current_span(
-        "chalkcompute.remote_function.invoke",
-        context=parent_context,
-        kind=SpanKind.SERVER,
-    ) as span:
-        span.set_attribute("chalk.remote_function.name", function_name)
-        if coalesced_count is not None:
-            span.set_attribute(
-                "chalk.remote_function.coalesced_count",
-                coalesced_count,
-            )
-        if not span.get_span_context().is_valid:
-            token = otel_context.attach(parent_context)
-            try:
-                yield
-            finally:
-                otel_context.detach(token)
-            return
-        try:
-            yield
-        except Exception as exc:
-            span.record_exception(exc)
-            span.set_status(Status(StatusCode.ERROR, str(exc)))
-            raise