inspect-ai 0.3.87__py3-none-any.whl → 0.3.89__py3-none-any.whl

inspect_ai/_eval/score.py

@@ -6,6 +6,7 @@ from typing import Any, Callable, Literal, cast
 import anyio
 
 from inspect_ai._display import display
+from inspect_ai._eval.context import init_task_context
 from inspect_ai._eval.loader import scorer_from_spec
 from inspect_ai._util._async import configured_async_backend, run_coroutine, tg_collect
 from inspect_ai._util.platform import platform_init, running_in_notebook
@@ -14,7 +15,9 @@ from inspect_ai.log import (
     EvalLog,
 )
 from inspect_ai.log._log import EvalMetricDefinition
+from inspect_ai.log._model import model_roles_config_to_model_roles
 from inspect_ai.model import ModelName
+from inspect_ai.model._model import get_model
 from inspect_ai.scorer import Metric, Scorer, Target
 from inspect_ai.scorer._metric import SampleScore
 from inspect_ai.scorer._reducer import (
@@ -122,7 +125,7 @@ async def score_async(
         scores: list[dict[str, SampleScore]] = await tg_collect(
             [
                 functools.partial(
-                    run_score_task, state, Target(sample.target), scorers, progress
+                    run_score_task, log, state, Target(sample.target), scorers, progress
                 )
                 for (sample, state) in zip(log.samples, states)
             ]
@@ -218,11 +221,25 @@ async def task_score(
 
 
 async def run_score_task(
+    log: EvalLog,
     state: TaskState,
     target: Target,
     scorers: list[Scorer],
     progress: Callable[..., None],
 ) -> dict[str, SampleScore]:
+    # get the model then initialize the async context
+    model = get_model(
+        model=log.eval.model,
+        config=log.plan.config.merge(log.eval.model_generate_config),
+        **log.eval.model_args,
+    )
+
+    # get the model roles
+    model_roles = model_roles_config_to_model_roles(log.eval.model_roles)
+
+    # initialize active model
+    init_task_context(model, model_roles)
+
     results: dict[str, SampleScore] = {}
     for scorer in scorers:
         result = await scorer(state, target)

inspect_ai/_eval/task/log.py

@@ -1,6 +1,5 @@
 from importlib import metadata as importlib_metadata
-from inspect import isgenerator
-from typing import Any, Iterator, Literal, cast
+from typing import Any, Literal, cast
 
 from shortuuid import uuid
 
@@ -34,6 +33,7 @@ from inspect_ai.log._log import (
     EvalScorer,
     eval_config_defaults,
 )
+from inspect_ai.log._model import model_args_for_log, model_roles_to_model_roles_config
 from inspect_ai.log._recorders import Recorder
 from inspect_ai.log._recorders.buffer import SampleBufferDatabase
 from inspect_ai.log._recorders.types import SampleEvent, SampleSummary
@@ -63,6 +63,7 @@ class TaskLogger:
         solver: SolverSpec | None,
         tags: list[str] | None,
         model: Model,
+        model_roles: dict[str, Model] | None,
         dataset: Dataset,
         scorer: list[ScorerSpec] | None,
         metrics: list[MetricSpec] | dict[str, list[MetricSpec]] | None,
@@ -84,17 +85,7 @@ class TaskLogger:
         packages = {PKG_NAME: importlib_metadata.version(PKG_NAME)}
 
         # redact authentication oriented model_args
-        model_args = model_args.copy()
-        if "api_key" in model_args:
-            del model_args["api_key"]
-        model_args = {k: v for k, v in model_args.items() if not k.startswith("aws_")}
-
-        # don't try to serialise generators
-        model_args = {
-            k: v
-            for k, v in model_args.items()
-            if not isgenerator(v) and not isinstance(v, Iterator)
-        }
+        model_args = model_args_for_log(model_args)
 
         # cwd_relative_path for sandbox config
         if sandbox and isinstance(sandbox.config, str):
@@ -141,6 +132,7 @@ class TaskLogger:
             model=str(ModelName(model)),
             model_generate_config=model.config,
             model_base_url=model.api.base_url,
+            model_roles=model_roles_to_model_roles_config(model_roles),
             dataset=EvalDataset(
                 name=dataset.name,
                 location=cwd_relative_path(dataset.location),

inspect_ai/_eval/task/resolved.py

@@ -13,6 +13,7 @@ class ResolvedTask:
     task_args: dict[str, Any]
     task_file: str | None
     model: Model
+    model_roles: dict[str, Model] | None
     sandbox: SandboxEnvironmentSpec | None
     sequence: int
     id: str | None = field(default=None)

inspect_ai/_eval/task/run.py

@@ -101,7 +101,6 @@ from .images import (
 )
 from .log import TaskLogger, collect_eval_data, log_start
 from .results import eval_results
-from .rundir import set_task_chdir
 from .sandbox import sandboxenv_context
 from .util import sample_messages, slice_dataset
 
@@ -121,6 +120,7 @@ SAMPLE_TOTAL_PROGRESS_UNITS = 1
 class TaskRunOptions:
     task: Task
     model: Model
+    model_roles: dict[str, Model] | None
     sandbox: SandboxEnvironmentSpec | None
     logger: TaskLogger
     eval_wd: str
@@ -137,6 +137,7 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
     # destructure options
     task = options.task
     model = options.model
+    model_roles = options.model_roles
     sandbox = options.sandbox
     logger = options.logger
     eval_wd = options.eval_wd
@@ -151,156 +152,136 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
     generate_config = task.config.merge(GenerateConfigArgs(**kwargs))
 
     # init task context
-    init_task_context(model, options.task.approval, generate_config)
-
-    # establish chdir for duration of execution (if a task has chdir=True)
-    with set_task_chdir(task):
-        # track stats and error
-        results: EvalResults | None = None
-        reductions: list[EvalSampleReductions] | None = None
-        stats = EvalStats(started_at=iso_now())
-
-        # handle sample errors (raise as required)
-        sample_error_handler = SampleErrorHandler(
-            config.fail_on_error, len(task.dataset)
-        )
-
-        # resolve some config
-        model_name = ModelName(model)
-        epochs = config.epochs if config.epochs else DEFAULT_EPOCHS
-        sandbox_cleanup = config.sandbox_cleanup is not False
-        log_images = config.log_images is not False
-        log_samples = config.log_samples is not False
-
-        # resolve dataset
-        _, samples, states = await resolve_dataset(
-            dataset=task.dataset,
-            model_name=model_name,
-            limit=config.limit,
-            sample_id=config.sample_id,
-            epochs=epochs,
-            log_images=log_images,
-            message_limit=config.message_limit,
-            token_limit=config.token_limit,
-        )
-
-        # resolve the plan (unroll chains)
-        solver = solver or task.solver
-        if isinstance(solver, Plan):
-            plan = solver
-        elif isinstance(solver, Chain):
-            plan = Plan(list(solver), cleanup=task.cleanup, internal=True)
-        else:
-            plan = Plan(unroll(solver), cleanup=task.cleanup, internal=True)
-
-        # add setup solver(s) if specified
-        if task.setup:
-            plan.steps = unroll(task.setup) + plan.steps
-
-        # resolve the scorer
-        score = score and task.scorer is not None
-        scorers: list[Scorer] | None = task.scorer if (score and task.scorer) else None
-        scorer_profiles = (
-            [
-                registry_log_name(scorer)
-                for scorer in scorers
-                if is_registry_object(scorer)
-            ]
-            if scorers is not None
-            else ["(none)"]
-        )
-
-        # compute an eval directory relative log location if we can
-        if PurePath(logger.location).is_relative_to(PurePath(eval_wd)):
-            log_location = PurePath(logger.location).relative_to(eval_wd).as_posix()
-        else:
-            log_location = logger.location
-
-        # create task profile for display
-        profile = TaskProfile(
-            name=task.name,
-            file=logger.eval.task_file,
-            model=model_name,
-            dataset=task.dataset.name or "(samples)",
-            scorer=", ".join(scorer_profiles),
-            samples=len(samples),
-            steps=len(samples) * SAMPLE_TOTAL_PROGRESS_UNITS,
-            eval_config=config,
-            task_args=logger.eval.task_args,
-            generate_config=generate_config,
-            tags=tags,
-            log_location=log_location,
-        )
+    init_task_context(model, model_roles, options.task.approval, generate_config)
+
+    # track stats and error
+    results: EvalResults | None = None
+    reductions: list[EvalSampleReductions] | None = None
+    stats = EvalStats(started_at=iso_now())
+
+    # handle sample errors (raise as required)
+    sample_error_handler = SampleErrorHandler(config.fail_on_error, len(task.dataset))
+
+    # resolve some config
+    model_name = ModelName(model)
+    epochs = config.epochs if config.epochs else DEFAULT_EPOCHS
+    sandbox_cleanup = config.sandbox_cleanup is not False
+    log_images = config.log_images is not False
+    log_samples = config.log_samples is not False
+
+    # resolve dataset
+    _, samples, states = await resolve_dataset(
+        dataset=task.dataset,
+        model_name=model_name,
+        limit=config.limit,
+        sample_id=config.sample_id,
+        epochs=epochs,
+        log_images=log_images,
+        message_limit=config.message_limit,
+        token_limit=config.token_limit,
+    )
 
-        with display().task(
-            profile,
-        ) as td:
-            try:
-                # start the log
-                await log_start(logger, plan, generate_config)
-
-                with td.progress() as p:
-                    # forward progress
-                    def progress(number: int) -> None:
-                        p.update(number)
-
-                    # provide solvers a function that they can use to generate output
-                    async def generate(
-                        state: TaskState,
-                        tool_calls: Literal["loop", "single", "none"] = "loop",
-                        cache: bool | CachePolicy = False,
-                        **kwargs: Unpack[GenerateConfigArgs],
-                    ) -> TaskState:
-                        return await task_generate(
-                            model=model,
-                            state=state,
-                            tool_calls=tool_calls,
-                            cache=cache,
-                            config=generate_config.merge(kwargs),
-                        )
+    # resolve the plan (unroll chains)
+    solver = solver or task.solver
+    if isinstance(solver, Plan):
+        plan = solver
+    elif isinstance(solver, Chain):
+        plan = Plan(list(solver), cleanup=task.cleanup, internal=True)
+    else:
+        plan = Plan(unroll(solver), cleanup=task.cleanup, internal=True)
+
+    # add setup solver(s) if specified
+    if task.setup:
+        plan.steps = unroll(task.setup) + plan.steps
+
+    # resolve the scorer
+    score = score and task.scorer is not None
+    scorers: list[Scorer] | None = task.scorer if (score and task.scorer) else None
+    scorer_profiles = (
+        [registry_log_name(scorer) for scorer in scorers if is_registry_object(scorer)]
+        if scorers is not None
+        else ["(none)"]
+    )
 
-                    # set generate for fork module
-                    set_task_generate(generate)
+    # compute an eval directory relative log location if we can
+    if PurePath(logger.location).is_relative_to(PurePath(eval_wd)):
+        log_location = PurePath(logger.location).relative_to(eval_wd).as_posix()
+    else:
+        log_location = logger.location
+
+    # create task profile for display
+    profile = TaskProfile(
+        name=task.name,
+        file=logger.eval.task_file,
+        model=model_name,
+        dataset=task.dataset.name or "(samples)",
+        scorer=", ".join(scorer_profiles),
+        samples=len(samples),
+        steps=len(samples) * SAMPLE_TOTAL_PROGRESS_UNITS,
+        eval_config=config,
+        task_args=logger.eval.task_args,
+        generate_config=generate_config,
+        tags=tags,
+        log_location=log_location,
+    )
 
-                    # semaphore to limit concurrency
-                    sample_semaphore = create_sample_semaphore(
-                        config, generate_config, model.api
+    with display().task(
+        profile,
+    ) as td:
+        try:
+            # start the log
+            await log_start(logger, plan, generate_config)
+
+            with td.progress() as p:
+                # forward progress
+                def progress(number: int) -> None:
+                    p.update(number)
+
+                # provide solvers a function that they can use to generate output
+                async def generate(
+                    state: TaskState,
+                    tool_calls: Literal["loop", "single", "none"] = "loop",
+                    cache: bool | CachePolicy = False,
+                    **kwargs: Unpack[GenerateConfigArgs],
+                ) -> TaskState:
+                    return await task_generate(
+                        model=model,
+                        state=state,
+                        tool_calls=tool_calls,
+                        cache=cache,
+                        config=generate_config.merge(kwargs),
                     )
 
-                    # track when samples complete and update progress as we go
-                    progress_results: list[dict[str, SampleScore]] = []
+                # set generate for fork module
+                set_task_generate(generate)
 
-                    def update_metrics(metrics: list[TaskDisplayMetric]) -> None:
-                        td.update_metrics(metrics)
-                        logger.update_metrics(metrics)
+                # semaphore to limit concurrency
+                sample_semaphore = create_sample_semaphore(
+                    config, generate_config, model.api
+                )
 
-                    update_metrics_display = update_metrics_display_fn(
-                        update_metrics,
-                        display_metrics=profile.eval_config.score_display is not False,
-                    )
+                # track when samples complete and update progress as we go
+                progress_results: list[dict[str, SampleScore]] = []
 
-                    def sample_complete(sample_score: dict[str, SampleScore]) -> None:
-                        # Capture the result
-                        progress_results.append(sample_score)
+                def update_metrics(metrics: list[TaskDisplayMetric]) -> None:
+                    td.update_metrics(metrics)
+                    logger.update_metrics(metrics)
 
-                        # Increment the segment progress
-                        td.sample_complete(
-                            complete=len(progress_results), total=len(samples)
-                        )
+                update_metrics_display = update_metrics_display_fn(
+                    update_metrics,
+                    display_metrics=profile.eval_config.score_display is not False,
+                )
 
-                        # Update metrics
-                        update_metrics_display(
-                            len(progress_results),
-                            progress_results,
-                            scorers,
-                            task.epochs_reducer,
-                            task.metrics,
-                        )
+                def sample_complete(sample_score: dict[str, SampleScore]) -> None:
+                    # Capture the result
+                    progress_results.append(sample_score)
 
-                    # initial progress
-                    td.sample_complete(complete=0, total=len(samples))
+                    # Increment the segment progress
+                    td.sample_complete(
+                        complete=len(progress_results), total=len(samples)
+                    )
 
-                    # Update metrics to empty state
+                    # Update metrics
                     update_metrics_display(
                         len(progress_results),
                         progress_results,
@@ -309,127 +290,133 @@ async def task_run(options: TaskRunOptions) -> EvalLog:
                         task.metrics,
                     )
 
-                    sample_results = await tg_collect(
-                        [
-                            functools.partial(
-                                task_run_sample,
-                                task_name=task.name,
-                                sample=sample,
-                                state=state,
-                                sandbox=sandbox,
-                                max_sandboxes=config.max_sandboxes,
-                                sandbox_cleanup=sandbox_cleanup,
-                                plan=plan,
-                                scorers=scorers,
-                                generate=generate,
-                                progress=progress,
-                                logger=logger if log_samples else None,
-                                log_images=log_images,
-                                sample_source=sample_source,
-                                sample_error=sample_error_handler,
-                                sample_complete=sample_complete,
-                                fails_on_error=(
-                                    config.fail_on_error is None
-                                    or config.fail_on_error is True
-                                ),
-                                time_limit=config.time_limit,
-                                working_limit=config.working_limit,
-                                semaphore=sample_semaphore,
-                            )
-                            for (sample, state) in zip(samples, states)
-                        ]
-                    )
+                # initial progress
+                td.sample_complete(complete=0, total=len(samples))
 
-                # compute and record metrics if we have scores
-                completed_scores = [
-                    score_dict
-                    for score_dict in sample_results
-                    if isinstance(score_dict, dict)
-                ]
-
-                if len(completed_scores) > 0:
-                    results, reductions = eval_results(
-                        samples=profile.samples,
-                        scores=completed_scores,
-                        reducers=task.epochs_reducer,
-                        scorers=scorers,
-                        metrics=task.metrics,
-                    )
+                # Update metrics to empty state
+                update_metrics_display(
+                    len(progress_results),
+                    progress_results,
+                    scorers,
+                    task.epochs_reducer,
+                    task.metrics,
+                )
 
-                # collect eval data
-                collect_eval_data(stats)
+                sample_results = await tg_collect(
+                    [
+                        functools.partial(
+                            task_run_sample,
+                            task_name=task.name,
+                            sample=sample,
+                            state=state,
+                            sandbox=sandbox,
+                            max_sandboxes=config.max_sandboxes,
+                            sandbox_cleanup=sandbox_cleanup,
+                            plan=plan,
+                            scorers=scorers,
+                            generate=generate,
+                            progress=progress,
+                            logger=logger if log_samples else None,
+                            log_images=log_images,
+                            sample_source=sample_source,
+                            sample_error=sample_error_handler,
+                            sample_complete=sample_complete,
+                            fails_on_error=(
+                                config.fail_on_error is None
+                                or config.fail_on_error is True
+                            ),
+                            time_limit=config.time_limit,
+                            working_limit=config.working_limit,
+                            semaphore=sample_semaphore,
+                        )
+                        for (sample, state) in zip(samples, states)
+                    ]
+                )
 
-                # finish w/ success status
-                eval_log = await logger.log_finish(
-                    "success", stats, results, reductions
+            # compute and record metrics if we have scores
+            completed_scores = [
+                score_dict
+                for score_dict in sample_results
+                if isinstance(score_dict, dict)
+            ]
+
+            if len(completed_scores) > 0:
+                results, reductions = eval_results(
+                    samples=profile.samples,
+                    scores=completed_scores,
+                    reducers=task.epochs_reducer,
+                    scorers=scorers,
+                    metrics=task.metrics,
                 )
 
-                # display task summary
-                td.complete(
-                    TaskSuccess(
-                        samples_completed=logger.samples_completed,
-                        stats=stats,
-                        results=results or EvalResults(),
-                    )
+            # collect eval data
+            collect_eval_data(stats)
+
+            # finish w/ success status
+            eval_log = await logger.log_finish("success", stats, results, reductions)
+
+            # display task summary
+            td.complete(
+                TaskSuccess(
+                    samples_completed=logger.samples_completed,
+                    stats=stats,
+                    results=results or EvalResults(),
                 )
+            )
 
-            except anyio.get_cancelled_exc_class():
-                with anyio.CancelScope(shield=True):
-                    # collect eval data
-                    collect_eval_data(stats)
+        except anyio.get_cancelled_exc_class():
+            with anyio.CancelScope(shield=True):
+                # collect eval data
+                collect_eval_data(stats)
 
-                    # finish w/ cancelled status
-                    eval_log = await logger.log_finish(
-                        "cancelled", stats, results, reductions
-                    )
+                # finish w/ cancelled status
+                eval_log = await logger.log_finish(
+                    "cancelled", stats, results, reductions
+                )
 
-                    # display task cancelled
-                    td.complete(TaskCancelled(logger.samples_completed, stats))
+                # display task cancelled
+                td.complete(TaskCancelled(logger.samples_completed, stats))
 
-            except BaseException as ex:
-                if options.debug_errors:
-                    raise
-                else:
-                    # get exception info
-                    type, value, traceback = sys.exc_info()
-                    type = type if type else BaseException
-                    value = value if value else ex
+        except BaseException as ex:
+            if options.debug_errors:
+                raise
+            else:
+                # get exception info
+                type, value, traceback = sys.exc_info()
+                type = type if type else BaseException
+                value = value if value else ex
 
-                    # build eval error
-                    error = eval_error(ex, type, value, traceback)
+                # build eval error
+                error = eval_error(ex, type, value, traceback)
 
-                    # collect eval data
-                    collect_eval_data(stats)
+                # collect eval data
+                collect_eval_data(stats)
 
-                    # finish with error status
-                    eval_log = await logger.log_finish(
-                        "error", stats, results, reductions, error
-                    )
+                # finish with error status
+                eval_log = await logger.log_finish(
+                    "error", stats, results, reductions, error
+                )
 
-                    # display it
-                    td.complete(
-                        TaskError(logger.samples_completed, type, value, traceback)
-                    )
+                # display it
+                td.complete(TaskError(logger.samples_completed, type, value, traceback))
 
-        # notify the view module that an eval just completed
-        # (in case we have a view polling for new evals)
-        view_notify_eval(logger.location)
+    # notify the view module that an eval just completed
+    # (in case we have a view polling for new evals)
+    view_notify_eval(logger.location)
 
-        try:
-            if (
-                await send_telemetry("eval_log_location", eval_log.location)
-                == "not_handled"
-            ):
-                # Converting the eval log to JSON is expensive. Only do so if
-                # eval_log_location was not handled.
-                await send_telemetry("eval_log", eval_log_json_str(eval_log))
-        except Exception as ex:
-            py_logger.warning(
-                f"Error occurred sending telemetry: {exception_message(ex)}"
-            )
+    try:
+        if (
+            await send_telemetry("eval_log_location", eval_log.location)
+            == "not_handled"
+        ):
+            # Converting the eval log to JSON is expensive. Only do so if
+            # eval_log_location was not handled.
+            await send_telemetry("eval_log", eval_log_json_str(eval_log))
+    except Exception as ex:
+        py_logger.warning(f"Error occurred sending telemetry: {exception_message(ex)}")
 
-        # return eval log
-        return eval_log
+    # return eval log
+    return eval_log
 
 
 def update_metrics_display_fn(
@@ -914,7 +901,7 @@ async def resolve_dataset(
     dataset: Dataset,
     model_name: ModelName,
     limit: int | tuple[int, int] | None,
-    sample_id: str | int | list[str | int] | None,
+    sample_id: str | int | list[str] | list[int] | list[str | int] | None,
     epochs: int,
     log_images: bool,
     message_limit: int | None,