chore: import upstream snapshot with attribution

tests/trainer/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/trainer/sample_components.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class SampleBase:
+    pass
+
+
+class SampleComponent(SampleBase):
+    def __init__(self, required: int, optional: int = 7) -> None:
+        self.required = required
+        self.optional = optional
+
+
+class SimpleComponent(SampleBase):
+    def __init__(self) -> None:
+        self.created = True
+
+
+class ComponentWithoutOptional(SampleBase):
+    def __init__(self, value: int) -> None:
+        self.value = value
+
+
+class ComponentWithKwargs(SampleBase):
+    def __init__(self, **kwargs: Any) -> None:
+        self.kwargs = kwargs
+
+
+class Unrelated:
+    pass

tests/trainer/test_init_utils.py

@@ -0,0 +1,286 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+
+from agentlightning.trainer.init_utils import build_component, instantiate_component, instantiate_from_spec, load_class
+
+from .sample_components import (
+    ComponentWithKwargs,
+    ComponentWithoutOptional,
+    SampleBase,
+    SampleComponent,
+    SimpleComponent,
+    Unrelated,
+)
+
+
+def test_load_class_resolves_fully_qualified_name() -> None:
+    resolved = load_class("tests.trainer.sample_components.SampleComponent")
+    assert resolved is SampleComponent
+
+
+def test_instantiate_component_uses_provided_kwargs() -> None:
+    component = instantiate_component(SampleComponent, {"required": 3})
+
+    assert isinstance(component, SampleComponent)
+    assert component.required == 3
+    assert component.optional == 7
+
+
+def test_instantiate_component_applies_optional_defaults_callable() -> None:
+    component = instantiate_component(
+        SampleComponent,
+        {"required": 2},
+        {"optional": lambda: 11},
+    )
+
+    assert component.optional == 11
+
+
+def test_instantiate_component_ignores_unknown_optional_defaults() -> None:
+    component = instantiate_component(
+        ComponentWithoutOptional,
+        {"value": 5},
+        {"missing": lambda: pytest.fail("should not be called")},
+    )
+
+    assert isinstance(component, ComponentWithoutOptional)
+    assert component.value == 5
+
+
+def test_instantiate_from_spec_with_string_path() -> None:
+    instance = instantiate_from_spec("tests.trainer.sample_components.SimpleComponent", spec_name="component")
+
+    assert isinstance(instance, SimpleComponent)
+    assert instance.created
+
+
+def test_instantiate_from_spec_with_dict_type() -> None:
+    instance = instantiate_from_spec(
+        {"type": "tests.trainer.sample_components.SampleComponent", "required": 4},
+        spec_name="component",
+    )
+
+    assert isinstance(instance, SampleComponent)
+    assert instance.required == 4
+
+
+def test_instantiate_from_spec_missing_type_uses_default_cls() -> None:
+    instance = instantiate_from_spec(
+        {"required": 8},
+        spec_name="adapter",
+        dict_requires_type=False,
+        dict_default_cls=SampleComponent,
+    )
+
+    assert isinstance(instance, SampleComponent)
+    assert instance.required == 8
+
+
+def test_instantiate_from_spec_missing_type_raises_when_required() -> None:
+    with pytest.raises(ValueError, match="component dict must have a 'type' key"):
+        instantiate_from_spec({}, spec_name="component")
+
+
+def test_build_component_returns_existing_instance() -> None:
+    existing = SampleComponent(required=3)
+    result = build_component(
+        existing,
+        expected_type=SampleComponent,
+        spec_name="component",
+    )
+
+    assert result is existing
+
+
+def test_build_component_uses_default_factory_for_none() -> None:
+    result = build_component(
+        None,
+        expected_type=SampleComponent,
+        spec_name="component",
+        default_factory=lambda: SampleComponent(required=5),
+    )
+
+    assert isinstance(result, SampleComponent)
+    assert result.required == 5
+
+
+def test_build_component_returns_none_when_allowed() -> None:
+    assert (
+        build_component(
+            None,
+            expected_type=SampleComponent,
+            spec_name="component",
+            allow_none=True,
+        )
+        is None
+    )
+
+
+def test_build_component_from_string_spec() -> None:
+    result = build_component(
+        "tests.trainer.sample_components.SimpleComponent",
+        expected_type=SampleBase,
+        spec_name="component",
+    )
+
+    assert isinstance(result, SimpleComponent)
+
+
+def test_build_component_from_registry_key_string_spec() -> None:
+    result = build_component(
+        "simple",
+        expected_type=SampleBase,
+        spec_name="component",
+        registry={"simple": "tests.trainer.sample_components.SimpleComponent"},
+    )
+
+    assert isinstance(result, SimpleComponent)
+
+
+def test_build_component_from_dict_spec_with_optional_defaults() -> None:
+    result = build_component(
+        {"type": "tests.trainer.sample_components.SampleComponent", "required": 10},
+        expected_type=SampleComponent,
+        spec_name="component",
+        optional_defaults={"optional": lambda: 42},
+    )
+
+    assert result is not None
+    assert result.optional == 42
+
+
+def test_build_component_from_type_spec() -> None:
+    result = build_component(
+        SimpleComponent,
+        expected_type=SampleBase,
+        spec_name="component",
+    )
+
+    assert isinstance(result, SimpleComponent)
+
+
+def test_build_component_from_callable_spec() -> None:
+    result = build_component(
+        lambda: SampleComponent(required=12),
+        expected_type=SampleComponent,
+        spec_name="component",
+    )
+
+    assert isinstance(result, SampleComponent)
+    assert result.required == 12
+
+
+def test_build_component_invalid_type_error() -> None:
+    with pytest.raises(
+        ValueError, match="Invalid component type: <class 'int'>. Expected SampleComponent, str, dict, or None."
+    ):
+        build_component(
+            1,
+            expected_type=SampleComponent,
+            spec_name="component",
+            invalid_spec_error_fmt="Invalid component type: {actual_type}. Expected SampleComponent, str, dict, or None.",
+        )
+
+
+def test_build_component_type_error_fmt() -> None:
+    with pytest.raises(TypeError, match="Custom type error message"):
+        build_component(
+            "tests.trainer.sample_components.SimpleComponent",
+            expected_type=ComponentWithoutOptional,
+            spec_name="component",
+            type_error_fmt="Custom type error message",
+        )
+
+
+def test_instantiate_component_applies_literal_optional_default() -> None:
+    component = instantiate_component(SampleComponent, {"required": 9}, {"optional": 13})
+
+    assert component.optional == 13
+
+
+def test_instantiate_from_spec_invalid_type_error_message() -> None:
+    with pytest.raises(TypeError, match="component spec must be a string or dict"):
+        instantiate_from_spec(42, spec_name="component")  # type: ignore[arg-type]
+
+
+def test_build_component_from_string_type_mismatch() -> None:
+    with pytest.raises(
+        TypeError,
+        match="component factory returned <class 'tests.trainer.sample_components.SimpleComponent'>, which is not a SampleComponent subclass.",
+    ):
+        build_component(
+            "tests.trainer.sample_components.SimpleComponent",
+            expected_type=SampleComponent,
+            spec_name="component",
+            invalid_spec_error_fmt="Invalid component type: {actual_type}. Expected SampleComponent, str, dict, or None.",
+        )
+
+
+def test_build_component_none_without_factory_raises_default_message() -> None:
+    with pytest.raises(ValueError, match="component cannot be None."):
+        build_component(
+            None,
+            expected_type=SampleComponent,
+            spec_name="component",
+        )
+
+
+def test_build_component_with_type_not_subclass_raises() -> None:
+    with pytest.raises(
+        ValueError,
+        match="Unsupported component type <class 'type'> for component",
+    ):
+        build_component(
+            Unrelated,
+            expected_type=SampleComponent,
+            spec_name="component",
+            invalid_spec_error_fmt="Unsupported component type {actual_type} for component",
+        )
+
+
+def test_build_component_callable_returning_wrong_type_raises() -> None:
+    with pytest.raises(
+        TypeError,
+        match="component factory returned <class 'tests.trainer.sample_components.Unrelated'>, which is not a SampleComponent subclass.",
+    ):
+        build_component(
+            lambda: Unrelated(),
+            expected_type=SampleComponent,
+            spec_name="component",
+        )
+
+
+def test_build_component_optional_defaults_skipped_when_not_supported() -> None:
+    component = build_component(
+        {"type": "tests.trainer.sample_components.ComponentWithKwargs"},
+        expected_type=ComponentWithKwargs,
+        spec_name="component",
+        optional_defaults={"missing": lambda: pytest.fail("should not be invoked")},
+    )
+
+    assert isinstance(component, ComponentWithKwargs)
+    assert component.kwargs == {}
+
+
+def test_build_component_dict_spec_uses_registry_type_lookup() -> None:
+    component = build_component(
+        {"type": "simple"},
+        expected_type=SampleBase,
+        spec_name="component",
+        registry={"simple": "tests.trainer.sample_components.SimpleComponent"},
+    )
+
+    assert isinstance(component, SimpleComponent)
+
+
+def test_build_component_dict_spec_supports_name_registry_lookup() -> None:
+    component = build_component(
+        {"name": "simple"},
+        expected_type=SampleBase,
+        spec_name="component",
+        dict_requires_type=False,
+        registry={"simple": "tests.trainer.sample_components.SimpleComponent"},
+    )
+
+    assert isinstance(component, SimpleComponent)

tests/trainer/test_trainer_dev.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for Trainer.dev requirements."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from agentlightning.algorithm import Algorithm, Baseline
+from agentlightning.execution.base import ExecutionStrategy
+from agentlightning.litagent import LitAgent
+from agentlightning.trainer import Trainer
+
+
+class DummyStrategy(ExecutionStrategy):
+    """Execution strategy that only records invocation."""
+
+    def __init__(self) -> None:
+        self.called = False
+
+    def execute(self, algorithm_bundle, runner_bundle, store) -> None:  # type: ignore[override]
+        self.called = True
+
+
+class DummyAgent(LitAgent[Any]):
+    """Minimal agent for exercising Trainer.dev."""
+
+
+class SlowAlgorithm(Algorithm):
+    """Algorithm that does not qualify as FastAlgorithm."""
+
+    def run(self, train_dataset=None, val_dataset=None):  # type: ignore[override]
+        return None
+
+
+def test_dev_requires_fast_algorithm() -> None:
+    trainer = Trainer(strategy=DummyStrategy(), algorithm=SlowAlgorithm())
+    agent = DummyAgent()
+
+    with pytest.raises(TypeError):
+        trainer.dev(agent)
+
+
+def test_dev_allows_fast_algorithm() -> None:
+    strategy = DummyStrategy()
+    trainer = Trainer(strategy=strategy, algorithm=Baseline())
+    agent = DummyAgent()
+
+    trainer.dev(agent)
+
+    assert strategy.called is True

tests/trainer/test_trainer_init.py

@@ -0,0 +1,158 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import pytest
+
+import agentlightning as agl
+
+
+def test_trainer_with_predefined_tracer() -> None:
+    """Test trainer initialization with predefined tracer."""
+    algorithm = agl.Baseline()
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        tracer=agl.OtelTracer(),
+    )
+    # Runner is initialized to be the default runner: LitAgentRunner
+    assert isinstance(trainer.runner, agl.LitAgentRunner)
+    assert isinstance(trainer.runner.tracer, agl.OtelTracer)
+
+
+def test_trainer_with_strategy_alias_shm() -> None:
+    """Test trainer initialization with strategy alias 'shm'."""
+    algorithm = agl.Baseline()
+    # Use strategy alias "shm"
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=1,  # n_runners must be 1 here
+        strategy="shm",
+    )
+    assert isinstance(trainer.strategy, agl.SharedMemoryExecutionStrategy)
+
+
+def test_trainer_with_strategy_dict_main_thread() -> None:
+    """Test trainer initialization with strategy dict allowing n_runners > 1."""
+    algorithm = agl.Baseline()
+    # Use dict. Now n_runners can be >1 because algorithm is on the main thread
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        strategy={"type": "shm", "main_thread": "algorithm", "managed_store": False},
+    )
+    assert isinstance(trainer.strategy, agl.SharedMemoryExecutionStrategy)
+    assert trainer.strategy.main_thread == "algorithm"
+    assert trainer.strategy.managed_store is False
+
+
+def test_trainer_with_initialized_strategy_ignores_n_runners() -> None:
+    """Test that n_runners is ignored when strategy is already initialized."""
+    algorithm = agl.Baseline()
+    # n_runners is ignored in the trainer because strategy has been initialized with n_runners=4
+    strategy = agl.SharedMemoryExecutionStrategy(main_thread="algorithm", n_runners=4)
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        strategy=strategy,
+    )
+    assert trainer.strategy is strategy
+    assert trainer.strategy.n_runners == 4  # type: ignore
+
+
+def test_trainer_with_client_server_strategy_dict() -> None:
+    """Test trainer initialization with client-server strategy dict."""
+    algorithm = agl.Baseline()
+    # By default, strategy is client-server, but you can also use a string alias to specify it again
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        strategy={
+            # This line is optional
+            "type": "cs",
+            "server_port": 9999,
+        },
+    )
+    assert isinstance(trainer.strategy, agl.ClientServerExecutionStrategy)
+    assert trainer.strategy.server_port == 9999
+
+
+def test_trainer_port_forwarded_to_client_server_strategy() -> None:
+    """Test that the top-level port argument configures the client-server strategy."""
+    trainer = agl.Trainer(
+        algorithm=agl.Baseline(),
+        n_runners=4,
+        port=8081,
+    )
+
+    assert isinstance(trainer.strategy, agl.ClientServerExecutionStrategy)
+    assert trainer.strategy.server_port == 8081
+
+
+def test_trainer_port_ignored_for_non_client_server_strategy() -> None:
+    """Test that port has no effect when using a non client-server strategy."""
+    trainer = agl.Trainer(
+        algorithm=agl.Baseline(),
+        n_runners=1,
+        port=8082,
+        strategy="shm",
+    )
+
+    assert isinstance(trainer.strategy, agl.SharedMemoryExecutionStrategy)
+    assert not hasattr(trainer.strategy, "server_port")
+
+
+def test_trainer_port_overrides_existing_client_server_strategy() -> None:
+    """Test that provided port overrides an initialized client-server strategy."""
+    strategy = agl.ClientServerExecutionStrategy(server_port=9000)
+
+    trainer = agl.Trainer(
+        algorithm=agl.Baseline(),
+        n_runners=1,
+        strategy=strategy,
+        port=9100,
+    )
+
+    assert trainer.strategy is strategy
+    assert trainer.strategy.server_port == 9100  # type: ignore
+
+
+def test_trainer_with_env_vars_for_execution_strategy(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that execution strategy supports environment variables to override values."""
+    algorithm = agl.Baseline()
+    # Execution strategy supports using environment variables to override the values
+    monkeypatch.setenv("AGL_SERVER_PORT", "10000")
+    monkeypatch.setenv("AGL_CURRENT_ROLE", "algorithm")
+    monkeypatch.setenv("AGL_MANAGED_STORE", "0")
+
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        # This line is optional
+        strategy="cs",
+    )
+    assert isinstance(trainer.strategy, agl.ClientServerExecutionStrategy)
+    assert trainer.strategy.server_port == 10000
+    assert trainer.strategy.role == "algorithm"
+    assert trainer.strategy.managed_store is False
+
+
+def test_trainer_with_string_adapter() -> None:
+    """Test trainer initialization with adapter specified as string."""
+    algorithm = agl.Baseline()
+    trainer = agl.Trainer(algorithm=algorithm, n_runners=8, adapter="agentlightning.adapter.TraceToMessages")
+    assert isinstance(trainer.adapter, agl.TraceToMessages)
+
+
+def test_trainer_with_adapter_dict_no_type() -> None:
+    """Test trainer initialization with adapter dict without type field."""
+    algorithm = agl.Baseline()
+    # If it's a dict and type is not provided, it will use the default class
+    trainer = agl.Trainer(
+        algorithm=algorithm,
+        n_runners=8,
+        adapter={"agent_match": "plan_agent", "repair_hierarchy": False},
+    )
+    assert isinstance(trainer.adapter, agl.TracerTraceToTriplet)
+    assert trainer.adapter.agent_match == "plan_agent"
+    assert trainer.adapter.repair_hierarchy is False