inspect-ai 0.3.88__py3-none-any.whl → 0.3.90__py3-none-any.whl

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(
@@ -655,18 +642,6 @@ async def task_run_sample(
                                 )
                             )
 
-                        # sample init event (remove file bodies as they have content or absolute paths)
-                        event_sample = sample.model_copy(
-                            update=dict(files={k: "" for k in sample.files.keys()})
-                            if sample.files
-                            else None
-                        )
-                        transcript()._event(
-                            SampleInitEvent(
-                                sample=event_sample, state=state_jsonable(state)
-                            )
-                        )
-
                         # set progress for plan then run it
                         state = await plan(state, generate)
 
@@ -914,7 +889,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,

inspect_ai/_eval/task/task.py

@@ -54,6 +54,7 @@ class Task:
         metrics: list[Metric] | dict[str, list[Metric]] | None = None,
         model: str | Model | None = None,
         config: GenerateConfig = GenerateConfig(),
+        model_roles: dict[str, str | Model] | None = None,
         sandbox: SandboxEnvironmentType | None = None,
         approval: str | list[ApprovalPolicy] | None = None,
         epochs: int | Epochs | None = None,
@@ -79,7 +80,8 @@ class Task:
             scorer: Scorer used to evaluate model output.
             metrics: Alternative metrics (overrides the metrics provided by the specified scorer).
             model: Default model for task (Optional, defaults to eval model).
-            config: Model generation config.
+            config: Model generation config for default model (does not apply to model roles)
+            model_roles: Named roles for use in `get_model()`.
             sandbox: Sandbox environment type (or optionally a str or tuple with a shorthand spec)
             approval: Tool use approval policies.
                 Either a path to an approval policy config file or a list of approval policies. Defaults to no approval policy.
@@ -136,6 +138,7 @@ class Task:
         self.metrics = metrics
         self.model = resolve_model(model)
         self.config = config
+        self.model_roles = resolve_model_roles(model_roles)
         self.sandbox = resolve_sandbox_environment(sandbox)
         self.approval = resolve_approval(approval)
         epochs = resolve_epochs(epochs)
@@ -185,6 +188,7 @@ def task_with(
     metrics: list[Metric] | dict[str, list[Metric]] | None | NotGiven = NOT_GIVEN,
     model: str | Model | NotGiven = NOT_GIVEN,
     config: GenerateConfig | NotGiven = NOT_GIVEN,
+    model_roles: dict[str, str | Model] | NotGiven = NOT_GIVEN,
     sandbox: SandboxEnvironmentType | None | NotGiven = NOT_GIVEN,
     approval: str | list[ApprovalPolicy] | None | NotGiven = NOT_GIVEN,
     epochs: int | Epochs | None | NotGiven = NOT_GIVEN,
@@ -214,7 +218,8 @@ def task_with(
         scorer: Scorer used to evaluate model output.
         metrics: Alternative metrics (overrides the metrics provided by the specified scorer).
         model: Default model for task (Optional, defaults to eval model).
-        config: Model generation config.
+        config: Model generation config for default model (does not apply to model roles)
+        model_roles: Named roles for use in `get_model()`.
         sandbox: Sandbox environment type (or optionally a str or tuple with a shorthand spec)
         approval: Tool use approval policies.
             Either a path to an approval policy config file or a list of approval policies. Defaults to no approval policy.
@@ -257,6 +262,8 @@ def task_with(
         task.model = resolve_model(model)
     if not isinstance(config, NotGiven):
         task.config = config
+    if not isinstance(model_roles, NotGiven):
+        task.model_roles = resolve_model_roles(model_roles)
     if not isinstance(sandbox, NotGiven):
         task.sandbox = resolve_sandbox_environment(sandbox)
     if not isinstance(approval, NotGiven):
@@ -315,6 +322,7 @@ class PreviousTask:
     task: str | Task
     task_args: dict[str, Any]
     model: Model | None
+    model_roles: dict[str, Model] | None
     log: EvalLog
 
 
@@ -365,6 +373,21 @@ def resolve_model(model: str | Model | None) -> Model | None:
         return model
 
 
+def resolve_model_roles(
+    model_roles: dict[str, str | Model] | None,
+) -> dict[str, Model] | None:
+    if model_roles is not None:
+        resolved_model_roles = {
+            k: get_model(v, memoize=False) if isinstance(v, str) else v
+            for k, v in model_roles.items()
+        }
+        for k, v in resolved_model_roles.items():
+            v._set_role(k)
+        return resolved_model_roles
+    else:
+        return None
+
+
 def resolve_scorer(scorer: Scorer | list[Scorer] | None) -> list[Scorer] | None:
     return (
         scorer if isinstance(scorer, list) else [scorer] if scorer is not None else None

inspect_ai/_eval/task/util.py

@@ -25,13 +25,6 @@ def task_run_dir(task: Task) -> str:
     return getattr(task, TASK_RUN_DIR_ATTR, os.getcwd())
 
 
-def task_chdir(task: Task) -> str | None:
-    if task.attribs.get("chdir", False) is True:
-        return task_run_dir(task)
-    else:
-        return None
-
-
 def task_file(task: Task, relative: bool = False) -> str | None:
     file = cast(str | None, getattr(task, TASK_FILE_ATTR, None))
     if file:
@@ -46,7 +39,7 @@ def task_file(task: Task, relative: bool = False) -> str | None:
 def slice_dataset(
     dataset: Dataset,
     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,
 ) -> Dataset:
     def normalise(id: str | int | None) -> str:
         if isinstance(id, str) and id.isdigit():