kiln-ai 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

kiln_ai/utils/async_job_runner.py

@@ -0,0 +1,106 @@
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import AsyncGenerator, Awaitable, Callable, List, TypeVar
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class Progress:
+    complete: int
+    total: int
+    errors: int
+
+
+class AsyncJobRunner:
+    def __init__(self, concurrency: int = 1):
+        if concurrency < 1:
+            raise ValueError("concurrency must be ≥ 1")
+        self.concurrency = concurrency
+
+    async def run(
+        self,
+        jobs: List[T],
+        run_job: Callable[[T], Awaitable[bool]],
+    ) -> AsyncGenerator[Progress, None]:
+        """
+        Runs the jobs with parallel workers and yields progress updates.
+        """
+        complete = 0
+        errors = 0
+        total = len(jobs)
+
+        # Send initial status
+        yield Progress(complete=complete, total=total, errors=errors)
+
+        worker_queue: asyncio.Queue[T] = asyncio.Queue()
+        for job in jobs:
+            worker_queue.put_nowait(job)
+
+        # simple status queue to return progress. True=success, False=error
+        status_queue: asyncio.Queue[bool] = asyncio.Queue()
+
+        workers = []
+        for _ in range(self.concurrency):
+            task = asyncio.create_task(
+                self._run_worker(worker_queue, status_queue, run_job),
+            )
+            workers.append(task)
+
+        try:
+            # Send status updates until workers are done, and they are all sent
+            while not status_queue.empty() or not all(
+                worker.done() for worker in workers
+            ):
+                try:
+                    # Use timeout to prevent hanging if all workers complete
+                    # between our while condition check and get()
+                    success = await asyncio.wait_for(status_queue.get(), timeout=0.1)
+                    if success:
+                        complete += 1
+                    else:
+                        errors += 1
+
+                    yield Progress(complete=complete, total=total, errors=errors)
+                except asyncio.TimeoutError:
+                    # Timeout is expected, just continue to recheck worker status
+                    # Don't love this but beats sentinels for reliability
+                    continue
+        finally:
+            # Cancel outstanding workers on early exit or error
+            for w in workers:
+                w.cancel()
+
+            # These are redundant, but keeping them will catch async errors
+            await asyncio.gather(*workers)
+            await worker_queue.join()
+
+    async def _run_worker(
+        self,
+        worker_queue: asyncio.Queue[T],
+        status_queue: asyncio.Queue[bool],
+        run_job: Callable[[T], Awaitable[bool]],
+    ):
+        while True:
+            try:
+                job = worker_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                # worker can end when the queue is empty
+                break
+
+            try:
+                success = await run_job(job)
+            except Exception:
+                logger.error("Job failed to complete", exc_info=True)
+                success = False
+
+            try:
+                await status_queue.put(success)
+            except Exception:
+                logger.error("Failed to enqueue status for job", exc_info=True)
+            finally:
+                # Always mark the dequeued task as done, even on exceptions
+                worker_queue.task_done()

kiln_ai/utils/config.py

@@ -138,6 +138,7 @@ class Config:
                 sensitive_keys=["api_key"],
             ),
         }
+        self._lock = threading.Lock()
         self._settings = self.load_settings()
 
     @classmethod
@@ -180,7 +181,7 @@ class Config:
         return None if value is None else property_config.type(value)
 
     def __setattr__(self, name, value):
-        if name in ("_properties", "_settings"):
+        if name in ("_properties", "_settings", "_lock"):
             super().__setattr__(name, value)
         elif name in self._properties:
             self.update_settings({name: value})
@@ -234,7 +235,7 @@ class Config:
 
     def update_settings(self, new_settings: Dict[str, Any]):
         # Lock to prevent race conditions in multi-threaded scenarios
-        with threading.Lock():
+        with self._lock:
             # Fresh load to avoid clobbering changes from other instances
             current_settings = self.load_settings()
             current_settings.update(new_settings)

kiln_ai/utils/dataset_import.py

@@ -1,11 +1,12 @@
 import csv
 import logging
+import random
 import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, Protocol
 
-from pydantic import BaseModel, Field, ValidationError, field_validator
+from pydantic import BaseModel, Field, ValidationError
 
 from kiln_ai.datamodel import DataSource, DataSourceType, Task, TaskOutput, TaskRun
 
@@ -20,14 +21,36 @@ class DatasetImportFormat(str, Enum):
     CSV = "csv"
 
 
+@dataclass
+class ImportConfig:
+    """Configuration for importing a dataset"""
+
+    dataset_type: DatasetImportFormat
+    dataset_path: str
+    dataset_name: str
+    """
+    A set of splits to assign to the import (as dataset tags).
+    The keys are the names of the splits (tag name), and the values are the proportions of the dataset to include in each split (should sum to 1).
+    """
+    tag_splits: Dict[str, float] | None = None
+
+    def validate_tag_splits(self) -> None:
+        if self.tag_splits:
+            EPSILON = 0.001  # Allow for small floating point errors
+            if abs(sum(self.tag_splits.values()) - 1) > EPSILON:
+                raise ValueError(
+                    "Splits must sum to 1. The following splits do not: "
+                    + ", ".join(f"{k}: {v}" for k, v in self.tag_splits.items())
+                )
+
+
 class Importer(Protocol):
     """Protocol for dataset importers"""
 
     def __call__(
         self,
         task: Task,
-        dataset_path: str,
-        dataset_name: str,
+        config: ImportConfig,
     ) -> int: ...
 
 
@@ -90,6 +113,44 @@ def without_none_values(d: dict) -> dict:
     return {k: v for k, v in d.items() if v is not None}
 
 
+def add_tag_splits(runs: list[TaskRun], tag_splits: Dict[str, float] | None) -> None:
+    """Assign split tags to runs according to configured proportions.
+
+    Args:
+        runs: List of TaskRun objects to assign tags to
+        tag_splits: Dictionary mapping tag names to their desired proportions
+
+    The assignment is random but ensures the proportions match the configured splits
+    as closely as possible given the number of runs.
+    """
+    if not tag_splits:
+        return
+
+    # Calculate exact number of runs for each split
+    total_runs = len(runs)
+    split_counts = {
+        tag: int(proportion * total_runs) for tag, proportion in tag_splits.items()
+    }
+
+    # Handle rounding errors by adjusting the largest split
+    remaining = total_runs - sum(split_counts.values())
+    if remaining != 0:
+        largest_split = max(split_counts.items(), key=lambda x: x[1])
+        split_counts[largest_split[0]] += remaining
+
+    # Create a list of tags with the correct counts
+    tags_to_assign = []
+    for tag, count in split_counts.items():
+        tags_to_assign.extend([tag] * count)
+
+    # Shuffle the tags to randomize assignment
+    random.shuffle(tags_to_assign)
+
+    # Assign tags to runs
+    for run, tag in zip(runs, tags_to_assign):
+        run.tags.append(tag)
+
+
 def create_task_run_from_csv_row(
     task: Task,
     row: dict[str, str],
@@ -143,18 +204,24 @@ def create_task_run_from_csv_row(
     return run
 
 
-def import_csv(task: Task, dataset_path: str, dataset_name: str) -> int:
+def import_csv(
+    task: Task,
+    config: ImportConfig,
+) -> int:
     """Import a CSV dataset.
 
     All rows are validated before any are persisted to files to avoid partial imports."""
 
     session_id = str(int(time.time()))
+    dataset_path = config.dataset_path
+    dataset_name = config.dataset_name
+    tag_splits = config.tag_splits
 
     required_headers = {"input", "output"}  # minimum required headers
     optional_headers = {"reasoning", "tags", "chain_of_thought"}  # optional headers
 
     rows: list[TaskRun] = []
-    with open(dataset_path, "r", newline="") as csvfile:
+    with open(dataset_path, "r", newline="", encoding="utf-8") as csvfile:
         reader = csv.DictReader(csvfile)
 
         # Check if we have headers
@@ -197,6 +264,8 @@ def import_csv(task: Task, dataset_path: str, dataset_name: str) -> int:
                 ) from e
             rows.append(run)
 
+    add_tag_splits(rows, tag_splits)
+
     # now that we know all rows are valid, we can save them
     for run in rows:
         run.save_to_file()
@@ -209,24 +278,17 @@ DATASET_IMPORTERS: Dict[DatasetImportFormat, Importer] = {
 }
 
 
-@dataclass
-class ImportConfig:
-    """Configuration for importing a dataset"""
-
-    dataset_type: DatasetImportFormat
-    dataset_path: str
-    dataset_name: str
-
-
 class DatasetFileImporter:
     """Import a dataset from a file"""
 
     def __init__(self, task: Task, config: ImportConfig):
         self.task = task
-        self.dataset_type = config.dataset_type
-        self.dataset_path = config.dataset_path
-        self.dataset_name = config.dataset_name
+        config.validate_tag_splits()
+        self.config = config
 
     def create_runs_from_file(self) -> int:
-        fn = DATASET_IMPORTERS[self.dataset_type]
-        return fn(self.task, self.dataset_path, self.dataset_name)
+        fn = DATASET_IMPORTERS[self.config.dataset_type]
+        return fn(
+            self.task,
+            self.config,
+        )

kiln_ai/utils/logging.py

@@ -0,0 +1,165 @@
+import datetime
+import json
+import logging
+import logging.handlers
+import os
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+
+from kiln_ai.utils.config import Config
+
+
+def get_default_formatter() -> str:
+    return "%(asctime)s.%(msecs)03d - %(levelname)s - %(name)s - %(message)s"
+
+
+def get_log_file_path(filename: str) -> str:
+    """Get the path to the log file, using environment override if specified.
+
+    Returns:
+        str: The path to the log file
+    """
+    log_path_default = os.path.join(Config.settings_dir(), "logs", filename)
+    log_path = os.getenv("KILN_LOG_FILE", log_path_default)
+
+    # Ensure the log directory exists
+    os.makedirs(os.path.dirname(log_path), exist_ok=True)
+    return log_path
+
+
+class CustomLiteLLMLogger(CustomLogger):
+    def __init__(self, logger: logging.Logger):
+        self.logger = logger
+
+    def log_pre_api_call(self, model, messages, kwargs):
+        api_base = kwargs.get("litellm_params", {}).get("api_base", "")
+        headers = kwargs.get("additional_args", {}).get("headers", {})
+        data = kwargs.get("additional_args", {}).get("complete_input_dict", {})
+
+        try:
+            # Print the curl command for the request
+            logger = Logging(
+                model=model,
+                messages=messages,
+                stream=False,
+                call_type="completion",
+                start_time=datetime.datetime.now(),
+                litellm_call_id="",
+                function_id="na",
+                kwargs=kwargs,
+            )
+            curl_command = logger._get_request_curl_command(
+                api_base=api_base,
+                headers=headers,
+                additional_args=kwargs,
+                data=data,
+            )
+            self.logger.info(f"{curl_command}")
+        except Exception as e:
+            self.logger.info(f"Curl Command: Could not print {e}")
+
+        # Print the formatted input data for the request in API format, pretty print
+        try:
+            self.logger.info(
+                f"Formatted Input Data (API):\n{json.dumps(data, indent=2)}"
+            )
+        except Exception as e:
+            self.logger.info(f"Formatted Input Data (API): Could not print {e}")
+
+        # Print the messages for the request in LiteLLM Message list, pretty print
+        try:
+            json_messages = json.dumps(messages, indent=2)
+            self.logger.info(f"Messages:\n{json_messages}")
+        except Exception as e:
+            self.logger.info(f"Messages: Could not print {e}")
+
+    def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
+        # No op
+        pass
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        litellm_logger = logging.getLogger("LiteLLM")
+        litellm_logger.error(
+            "Used a sync call in Litellm. Kiln should use async calls."
+        )
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        # This logging method is supposed to be called by Litellm in synchronous error cases (Kiln should use async calls instead)
+        # but it appears to also be getting called in async calls that fail early (e.g. UnsupportedParamsError).
+        litellm_logger = logging.getLogger("LiteLLM")
+        litellm_logger.error(
+            "LiteLLM logged a synchronous failure event. This may result from a sync call, or from an async call failing early (e.g. invalid parameters). Make sure you are using async calls.",
+        )
+
+    #### ASYNC #### - for acompletion/aembeddings
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        try:
+            if len(response_obj.choices) == 1:
+                if response_obj.choices[0].message.tool_calls:
+                    for tool_call in response_obj.choices[0].message.tool_calls:
+                        try:
+                            args = tool_call.function.arguments
+                            function_name = tool_call.function.name
+                            self.logger.info(
+                                f"Model Response Tool Call Arguments [{function_name}]:\n{args}"
+                            )
+                        except Exception:
+                            self.logger.info(f"Model Response Tool Call:\n{tool_call}")
+
+                content = response_obj.choices[0].message.content
+                if content:
+                    try:
+                        # JSON format logs if possible
+                        json_content = json.loads(content)
+                        self.logger.info(
+                            f"Model Response Content:\n{json.dumps(json_content, indent=2)}"
+                        )
+                    except Exception:
+                        self.logger.info(f"Model Response Content:\n{content}")
+            elif len(response_obj.choices) > 1:
+                self.logger.info(
+                    f"Model Response (multiple choices):\n{response_obj.choices}"
+                )
+            else:
+                self.logger.info("Model Response: No choices returned")
+
+        except Exception as e:
+            self.logger.info(f"Model Response: Could not print {e}")
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self.logger.info(f"LiteLLM Failure: {response_obj}")
+
+
+def setup_litellm_logging(filename: str = "model_calls.log"):
+    # Check if we already have a custom litellm logger
+    for callback in litellm.callbacks or []:
+        if isinstance(callback, CustomLiteLLMLogger):
+            return  # We already have a custom litellm logger
+
+    # If we don't have a custom litellm logger, create one
+    # Disable the default litellm logger except for errors. It's ugly, hard to use, and we don't want it to mix with kiln logs.
+    litellm_logger = logging.getLogger("LiteLLM")
+    litellm_logger.setLevel(logging.ERROR)
+
+    # Create a logger that logs to files, with a max size of 5MB and 3 backup files
+    handler = logging.handlers.RotatingFileHandler(
+        get_log_file_path(filename),
+        maxBytes=5 * 1024 * 1024,  # 5MB
+        backupCount=3,
+    )
+
+    # Set formatter to match the default formatting
+    formatter = logging.Formatter(get_default_formatter())
+    handler.setFormatter(formatter)
+
+    # Create a new logger for model calls
+    model_calls_logger = logging.getLogger("ModelCalls")
+    model_calls_logger.setLevel(logging.INFO)
+    model_calls_logger.propagate = False  # Only log to file
+    model_calls_logger.addHandler(handler)
+
+    # Tell litellm to use our custom logger
+    litellm.callbacks = [CustomLiteLLMLogger(model_calls_logger)]

kiln_ai/utils/test_async_job_runner.py

@@ -0,0 +1,199 @@
+from typing import List
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress
+
+
+@pytest.mark.parametrize("concurrency", [0, -1, -25])
+def test_invalid_concurrency_raises(concurrency):
+    with pytest.raises(ValueError):
+        AsyncJobRunner(concurrency=concurrency)
+
+
+# Test with and without concurrency
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_status_updates(concurrency):
+    job_count = 50
+    jobs = [{"id": i} for i in range(job_count)]
+
+    runner = AsyncJobRunner(concurrency=concurrency)
+
+    # fake run_job that succeeds
+    mock_run_job_success = AsyncMock(return_value=True)
+
+    # Expect the status updates in order, and 1 for each job
+    expected_completed_count = 0
+    async for progress in runner.run(jobs, mock_run_job_success):
+        assert progress.complete == expected_completed_count
+        expected_completed_count += 1
+        assert progress.errors == 0
+        assert progress.total == job_count
+
+    # Verify last status update was complete
+    assert expected_completed_count == job_count + 1
+
+    # Verify run_job was called for each job
+    assert mock_run_job_success.call_count == job_count
+
+    # Verify run_job was called with the correct arguments
+    for i in range(job_count):
+        mock_run_job_success.assert_any_await(jobs[i])
+
+
+# Test with and without concurrency
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_status_updates_empty_job_list(concurrency):
+    empty_job_list = []
+
+    runner = AsyncJobRunner(concurrency=concurrency)
+
+    # fake run_job that succeeds
+    mock_run_job_success = AsyncMock(return_value=True)
+
+    updates: List[Progress] = []
+    async for progress in runner.run(empty_job_list, mock_run_job_success):
+        updates.append(progress)
+
+    # Verify last status update was complete
+    assert len(updates) == 1
+
+    assert updates[0].complete == 0
+    assert updates[0].errors == 0
+    assert updates[0].total == 0
+
+    # Verify run_job was called for each job
+    assert mock_run_job_success.call_count == 0
+
+
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_all_failures(concurrency):
+    job_count = 50
+    jobs = [{"id": i} for i in range(job_count)]
+
+    runner = AsyncJobRunner(concurrency=concurrency)
+
+    # fake run_job that fails
+    mock_run_job_failure = AsyncMock(return_value=False)
+
+    # Expect the status updates in order, and 1 for each job
+    expected_error_count = 0
+    async for progress in runner.run(jobs, mock_run_job_failure):
+        assert progress.complete == 0
+        assert progress.errors == expected_error_count
+        expected_error_count += 1
+        assert progress.total == job_count
+
+    # Verify last status update was complete
+    assert expected_error_count == job_count + 1
+
+    # Verify run_job was called for each job
+    assert mock_run_job_failure.call_count == job_count
+
+    # Verify run_job was called with the correct arguments
+    for i in range(job_count):
+        mock_run_job_failure.assert_any_await(jobs[i])
+
+
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_partial_failures(concurrency):
+    job_count = 50
+    jobs = [{"id": i} for i in range(job_count)]
+
+    # we want to fail on some jobs and succeed on others
+    jobs_to_fail = set([0, 2, 4, 6, 8, 20, 25])
+
+    runner = AsyncJobRunner(concurrency=concurrency)
+
+    # fake run_job that fails
+    mock_run_job_partial_success = AsyncMock(
+        # return True for jobs that should succeed
+        side_effect=lambda job: job["id"] not in jobs_to_fail
+    )
+
+    # Expect the status updates in order, and 1 for each job
+    async for progress in runner.run(jobs, mock_run_job_partial_success):
+        assert progress.total == job_count
+
+    # Verify last status update was complete
+    expected_error_count = len(jobs_to_fail)
+    expected_success_count = len(jobs) - expected_error_count
+    assert progress.errors == expected_error_count
+    assert progress.complete == expected_success_count
+
+    # Verify run_job was called for each job
+    assert mock_run_job_partial_success.call_count == job_count
+
+    # Verify run_job was called with the correct arguments
+    for i in range(job_count):
+        mock_run_job_partial_success.assert_any_await(jobs[i])
+
+
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_partial_raises(concurrency):
+    job_count = 50
+    jobs = [{"id": i} for i in range(job_count)]
+
+    runner = AsyncJobRunner(concurrency=concurrency)
+
+    ids_to_fail = set([10, 25])
+
+    def failure_fn(job):
+        if job["id"] in ids_to_fail:
+            raise Exception("job failed unexpectedly")
+        return True
+
+    # fake run_job that fails
+    mock_run_job_partial_success = AsyncMock(side_effect=failure_fn)
+
+    # generate all the values we expect to see in progress updates
+    complete_values_expected = set([i for i in range(job_count - len(ids_to_fail) + 1)])
+    errors_values_expected = set([i for i in range(len(ids_to_fail) + 1)])
+
+    # keep track of all the updates we see
+    updates: List[Progress] = []
+
+    # we keep track of the progress values we have actually seen
+    complete_values_actual = set()
+    errors_values_actual = set()
+
+    # Expect the status updates in order, and 1 for each job
+    async for progress in runner.run(jobs, mock_run_job_partial_success):
+        updates.append(progress)
+        complete_values_actual.add(progress.complete)
+        errors_values_actual.add(progress.errors)
+
+        assert progress.total == job_count
+
+    # complete values should be all the jobs, except for the ones that failed
+    assert progress.complete == job_count - len(ids_to_fail)
+
+    # check that the actual updates and expected updates are equivalent sets
+    assert complete_values_actual == complete_values_expected
+    assert errors_values_actual == errors_values_expected
+
+    # we should have seen one update for each job, plus one for the initial status update
+    assert len(updates) == job_count + 1
+
+
+@pytest.mark.parametrize("concurrency", [1, 25])
+@pytest.mark.asyncio
+async def test_async_job_runner_cancelled(concurrency):
+    runner = AsyncJobRunner(concurrency=concurrency)
+    jobs = [{"id": i} for i in range(10)]
+
+    with patch.object(
+        runner,
+        "_run_worker",
+        side_effect=Exception("run_worker raised an exception"),
+    ):
+        # if an exception is raised in the task, we should see it bubble up
+        with pytest.raises(Exception, match="run_worker raised an exception"):
+            async for _ in runner.run(jobs, AsyncMock(return_value=True)):
+                pass

kiln_ai/utils/test_config.py

@@ -1,5 +1,6 @@
 import getpass
 import os
+import threading
 from unittest.mock import patch
 
 import pytest
@@ -299,3 +300,25 @@ def test_yaml_persistence_structured_data(config_with_yaml, mock_yaml_file):
     with open(mock_yaml_file, "r") as f:
         saved_settings = yaml.safe_load(f)
     assert saved_settings["list_of_objects"] == new_settings
+
+
+def test_update_settings_thread_safety(config_with_yaml):
+    config = config_with_yaml
+
+    exceptions = []
+
+    def update(val):
+        try:
+            config.update_settings({"int_property": val})
+        except Exception as e:
+            exceptions.append(e)
+
+    threads = [threading.Thread(target=update, args=(i,)) for i in range(5)]
+
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert not exceptions
+    assert config.int_property in range(5)