Merge branch 'main' into taegyunkim/prof-12434-uwsgi

Author: taegyunkim

benchmarks/django_simple/scenario.py

@@ -69,17 +69,17 @@ def _(loops):
                 from ddtrace.contrib.internal.django import database
 
                 try:
-                    database.get_conn_config.invalidate()
+                    database.get_conn_config.cache_clear()
                 except Exception:
                     pass
 
                 try:
-                    database.get_service_name.invalidate()
+                    database.get_service_name.cache_clear()
                 except Exception:
                     pass
 
                 try:
-                    database.get_conn_service_name.invalidate()
+                    database.get_conn_service_name.cache_clear()
                 except Exception:
                     pass
             except Exception:
@@ -90,12 +90,12 @@ def _(loops):
                 from ddtrace.contrib.internal.django import cache
 
                 try:
-                    cache.get_service_name.invalidate()
+                    cache.get_service_name.cache_clear()
                 except Exception:
                     pass
 
                 try:
-                    cache.func_cache_operation.invalidate()
+                    cache.func_cache_operation.cache_clear()
                 except Exception:
                     pass
             except Exception:

ddtrace/appsec/_utils.py

@@ -367,15 +367,20 @@ def unpatching_popen():
     Context manager to temporarily unpatch `subprocess.Popen` for testing purposes.
     This is useful to ensure that the original `Popen` behavior is restored after the context.
     """
+    import os
     import subprocess  # nosec B404
 
+    from ddtrace.internal._unpatched import unpatched_close
     from ddtrace.internal._unpatched import unpatched_Popen
 
+    original_os_close = os.close
+    os.close = unpatched_close
     original_popen = subprocess.Popen
     subprocess.Popen = unpatched_Popen
     asm_config._bypass_instrumentation_for_waf = True
     try:
         yield
     finally:
         subprocess.Popen = original_popen
+        os.close = original_os_close
         asm_config._bypass_instrumentation_for_waf = False

ddtrace/contrib/internal/django/middleware.py

@@ -1,3 +1,4 @@
+from inspect import iscoroutinefunction
 from inspect import isfunction
 from types import FunctionType
 from typing import Any
@@ -134,17 +135,42 @@ def traced_auth_middleware_process_request(func: FunctionType, args: Tuple[Any],
 def traced_middleware_factory(func: FunctionType, args: Tuple[Any], kwargs: Dict[str, Any]) -> Any:
     middleware = func(*args, **kwargs)
 
-    if isfunction(middleware):
-        if hasattr(func, "__module__") and hasattr(func, "__qualname__"):
-            resource = f"{func.__module__}.{func.__qualname__}"
-        else:
-            resource = func_name(func)
+    if not isfunction(middleware):
+        return middleware
+
+    if hasattr(func, "__module__") and hasattr(func, "__qualname__"):
+        resource = f"{func.__module__}.{func.__qualname__}"
+    else:
+        resource = func_name(func)
+
+    if iscoroutinefunction(middleware):
+        # Handle async middleware - create async wrapper
+        async def traced_async_middleware_func(*args, **kwargs):
+            # The first argument for all middleware is the request object
+            # DEV: Do `optional=true` to avoid raising an error for middleware that don't follow the convention
+            # DEV: This is a function, so no `self` argument, so request is at position 0
+            request = get_argument_value(args, kwargs, 0, "request", optional=True)
+
+            with core.context_with_data(
+                "django.middleware.func",
+                span_name="django.middleware",
+                resource=resource,
+                tags={
+                    COMPONENT: config_django.integration_name,
+                },
+                tracer=config_django._tracer,
+                request=request,
+            ):
+                return await middleware(*args, **kwargs)
 
+        return traced_async_middleware_func
+    else:
+        # Handle sync middleware - use original wrapping approach
         def traced_middleware_func(func: FunctionType, args: Tuple[Any], kwargs: Dict[str, Any]) -> Any:
             # The first argument for all middleware is the request object
             # DEV: Do `optional=true` to avoid raising an error for middleware that don't follow the convention
             # DEV: This is a function, so no `self` argument, so request is at position 0
-            request = get_argument_value(args, kwargs, 0, "request")
+            request = get_argument_value(args, kwargs, 0, "request", optional=True)
 
             with core.context_with_data(
                 "django.middleware.func",

ddtrace/internal/_unpatched.py

@@ -15,6 +15,7 @@
 
 previous_loaded_modules = frozenset(sys.modules.keys())
 from subprocess import Popen as unpatched_Popen  # noqa # nosec B404
+from os import close as unpatched_close  # noqa: F401, E402
 
 loaded_modules = frozenset(sys.modules.keys())
 for module in previous_loaded_modules - loaded_modules:

ddtrace/internal/utils/cache.py

@@ -1,13 +1,13 @@
+from functools import lru_cache
 from functools import wraps
 from inspect import FullArgSpec
 from inspect import getfullargspec
 from inspect import isgeneratorfunction
-from threading import RLock
 from typing import Any  # noqa:F401
 from typing import Callable  # noqa:F401
 from typing import Optional  # noqa:F401
 from typing import Type  # noqa:F401
-from typing import TypeVar  # noqa:F401
+from typing import TypeVar
 
 
 miss = object()
@@ -17,78 +17,14 @@
 M = Callable[[Any, T], Any]
 
 
-class LFUCache(dict):
-    """Simple LFU cache implementation.
+def cached(maxsize: int = 256) -> Callable[[Callable], Callable]:
+    def _(f: Callable) -> Callable:
+        return lru_cache(maxsize)(f)
 
-    This cache is designed for memoizing functions with a single hashable
-    argument. The eviction policy is LFU, i.e. the least frequently used values
-    are evicted when the cache is full. The amortized cost of shrinking the
-    cache when it grows beyond the requested size is O(log(size)).
-    """
-
-    def __init__(self, maxsize=256):
-        # type: (int) -> None
-        self.maxsize = maxsize
-        self.lock = RLock()
-        self.count_lock = RLock()
-
-    def get(self, key, f):  # type: ignore[override]
-        # type: (T, F) -> Any
-        """Get a value from the cache.
-
-        If the value with the given key is not in the cache, the expensive
-        function ``f`` is called on the key to generate it. The return value is
-        then stored in the cache and returned to the caller.
-        """
-
-        _ = super(LFUCache, self).get(key, miss)
-        if _ is not miss:
-            with self.count_lock:
-                value, count = _
-                self[key] = (value, count + 1)
-            return value
-
-        with self.lock:
-            _ = super(LFUCache, self).get(key, miss)
-            if _ is not miss:
-                with self.count_lock:
-                    value, count = _
-                    self[key] = (value, count + 1)
-                return value
-
-            # Cache miss: ensure that we have enough space in the cache
-            # by evicting half of the entries when we go over the threshold
-            while len(self) >= self.maxsize:
-                for h in sorted(self, key=lambda h: self[h][1])[: self.maxsize >> 1]:
-                    del self[h]
-
-            value = f(key)
-
-            self[key] = (value, 1)
-
-            return value
-
-
-def cached(maxsize=256):
-    # type: (int) -> Callable[[F], F]
-    """Decorator for memoizing functions of a single argument (LFU policy)."""
-
-    def cached_wrapper(f):
-        # type: (F) -> F
-        cache = LFUCache(maxsize)
-
-        def cached_f(key):
-            # type: (T) -> Any
-            return cache.get(key, f)
-
-        cached_f.invalidate = cache.clear  # type: ignore[attr-defined]
-
-        return cached_f
-
-    return cached_wrapper
+    return _
 
 
-class CachedMethodDescriptor(object):
+class CachedMethodDescriptor:
     def __init__(self, method, maxsize):
         # type: (M, int) -> None
         self._method = method

ddtrace/llmobs/_experiment.py

@@ -44,6 +44,11 @@
 DatasetRecordInputType = Dict[str, NonNoneJSONType]
 
 
+class Project(TypedDict):
+    name: str
+    _id: str
+
+
 class DatasetRecordRaw(TypedDict):
     input_data: DatasetRecordInputType
     expected_output: JSONType
@@ -106,13 +111,15 @@ class Dataset:
     def __init__(
         self,
         name: str,
+        project: Project,
         dataset_id: str,
         records: List[DatasetRecord],
         description: str,
         version: int,
         _dne_client: "LLMObsExperimentsClient",
     ) -> None:
         self.name = name
+        self.project = project
         self.description = description
         self._id = dataset_id
         self._version = version
@@ -335,8 +342,8 @@ def run(
             )
             return []
 
-        project_id = self._llmobs_instance._dne_client.project_create_or_get(self._project_name)
-        self._project_id = project_id
+        project = self._llmobs_instance._dne_client.project_create_or_get(self._project_name)
+        self._project_id = project.get("_id", "")
 
         experiment_id, experiment_run_name = self._llmobs_instance._dne_client.experiment_create(
             self.name,
@@ -416,6 +423,7 @@ def _run_task(self, jobs: int, raise_errors: bool = False, sample_size: Optional
             subset_name = "[Test subset of {} records] {}".format(sample_size, self._dataset.name)
             subset_dataset = Dataset(
                 name=subset_name,
+                project=self._dataset.project,
                 dataset_id=self._dataset._id,
                 records=subset_records,
                 description=self._dataset.description,

ddtrace/llmobs/_llmobs.py

@@ -92,6 +92,7 @@
 from ddtrace.llmobs._experiment import Experiment
 from ddtrace.llmobs._experiment import ExperimentConfigType
 from ddtrace.llmobs._experiment import JSONType
+from ddtrace.llmobs._experiment import Project
 from ddtrace.llmobs._utils import AnnotationContext
 from ddtrace.llmobs._utils import LinkTracker
 from ddtrace.llmobs._utils import _get_ml_app
@@ -212,6 +213,7 @@ def __init__(
             interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
             timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
             _app_key=self._app_key,
+            _default_project=Project(name=self._project_name, _id=""),
             is_agentless=True,  # agent proxy doesn't seem to work for experiments
         )
 
@@ -645,15 +647,21 @@ def enable(
             )
 
     @classmethod
-    def pull_dataset(cls, name: str) -> Dataset:
-        ds = cls._instance._dne_client.dataset_get_with_records(name)
+    def pull_dataset(cls, dataset_name: str, project_name: Optional[str] = None) -> Dataset:
+        ds = cls._instance._dne_client.dataset_get_with_records(dataset_name, (project_name or cls._project_name))
         return ds
 
     @classmethod
-    def create_dataset(cls, name: str, description: str = "", records: Optional[List[DatasetRecord]] = None) -> Dataset:
+    def create_dataset(
+        cls,
+        dataset_name: str,
+        project_name: Optional[str] = None,
+        description: str = "",
+        records: Optional[List[DatasetRecord]] = None,
+    ) -> Dataset:
         if records is None:
             records = []
-        ds = cls._instance._dne_client.dataset_create(name, description)
+        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
         for r in records:
             ds.append(r)
         if len(records) > 0:
@@ -669,19 +677,20 @@ def create_dataset_from_csv(
         expected_output_columns: Optional[List[str]] = None,
         metadata_columns: Optional[List[str]] = None,
         csv_delimiter: str = ",",
-        description="",
+        description: str = "",
+        project_name: Optional[str] = None,
     ) -> Dataset:
         if expected_output_columns is None:
             expected_output_columns = []
         if metadata_columns is None:
             metadata_columns = []
-        ds = cls._instance._dne_client.dataset_create(dataset_name, description)
 
         # Store the original field size limit to restore it later
         original_field_size_limit = csv.field_size_limit()
 
         csv.field_size_limit(EXPERIMENT_CSV_FIELD_MAX_SIZE)  # 10mb
 
+        records = []
         try:
             with open(csv_path, mode="r") as csvfile:
                 content = csvfile.readline().strip()
@@ -708,7 +717,7 @@ def create_dataset_from_csv(
                     raise ValueError(f"Metadata columns not found in CSV header: {missing_metadata_columns}")
 
                 for row in rows:
-                    ds.append(
+                    records.append(
                         DatasetRecord(
                             input_data={col: row[col] for col in input_data_columns},
                             expected_output={col: row[col] for col in expected_output_columns},
@@ -721,6 +730,9 @@ def create_dataset_from_csv(
             # Always restore the original field size limit
             csv.field_size_limit(original_field_size_limit)
 
+        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
+        for r in records:
+            ds.append(r)
         if len(ds) > 0:
             cls._instance._dne_client.dataset_bulk_upload(ds._id, ds._records)
         return ds