kiln-ai 0.14.0__py3-none-any.whl → 0.16.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/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,12 +204,18 @@ 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
@@ -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/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