Upgrade ty to 0.0.59

pyproject.toml

@@ -254,6 +254,7 @@ python-version = "3.12"
 # optional deps are installed. The ty:ignore comments are needed in CI (with all deps)
 # but become unused locally (without all deps).
 unused-ignore-comment = "ignore"
+unused-type-ignore-comment = "ignore"
 
 [tool.ty.analysis]
 # Allow unresolved imports for optional dependencies that may not be installed locally.
@@ -295,9 +296,11 @@ allowed-unresolved-imports = [
     "seaborn.**",
     # megatron deps
     "causal_conv1d.**",
+    "einops.**",
     "fla.**",
     "megatron.**",
     "quack.**",
+    "safetensors.**",
     "transformer_engine.**",
     "triton.**",
 ]
@@ -312,7 +315,7 @@ dev = [
     "pytest>=8.4.1",
     "nbval>=0.11.0",
     "pytest-xdist>=3.8.0",
-    "ty==0.0.14",
+    "ty==0.0.59",
     "pytest-asyncio>=1.1.0",
     "duckdb>=1.0.0",
     "pyarrow>=15.0.0",

src/art/backend.py

@@ -1,7 +1,16 @@
-from typing import TYPE_CHECKING, Any, AsyncIterator, Iterable, Protocol, TypeAlias
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Callable,
+    Coroutine,
+    Iterable,
+    Protocol,
+    TypeAlias,
+)
 
 from . import dev
-from .trajectories import Trajectory, TrajectoryGroup
+from .trajectories import Trajectory
 from .types import TrainResult, TrainSFTConfig
 
 if TYPE_CHECKING:
@@ -35,12 +44,11 @@ async def _prepare_backend_for_training(
         config: dev.OpenAIServerConfig | None,
     ) -> tuple[str, str]: ...
 
-    async def train(
-        self,
-        model: AnyTrainableModel,
-        trajectory_groups: Iterable[TrajectoryGroup],
-        **kwargs: Any,
-    ) -> TrainResult: ...
+    # Backends intentionally expose backend-specific optional training arguments.
+    # Callable[..., ...] preserves that extensibility without falsely requiring
+    # every implementation to accept every other backend's keyword arguments.
+    @property
+    def train(self) -> Callable[..., Coroutine[Any, Any, TrainResult]]: ...
 
     def _train_sft(
         self,

src/art/gather.py

@@ -3,7 +3,7 @@
 import contextlib
 import contextvars
 from dataclasses import dataclass, field
-from typing import Awaitable, Callable, Iterable, Iterator, Literal, overload
+from typing import Awaitable, Callable, Iterable, Iterator, Literal, Sequence, overload
 
 from openai.types.chat.chat_completion import Choice
 from tqdm import auto as tqdm
@@ -53,7 +53,7 @@ async def group_forward(g: Awaitable[TrajectoryGroup]):
         context.pbar.close()
 
     # Filter out any None results that may have been returned due to handled exceptions
-    processed_groups = []
+    processed_groups: list[TrajectoryGroup] = []
     for g in result_groups:
         if g is None:
             continue
@@ -113,12 +113,7 @@ async def gather_trajectories(
     pbar_desc: str | None = "gather",
     pbar_total_completion_tokens: bool = False,
     max_exceptions: int | float = 0,
-) -> (
-    list[Trajectory]
-    | list[Trajectory | BaseException]
-    | list[list[Trajectory]]
-    | list[list[Trajectory] | BaseException]
-):
+) -> Sequence[Trajectory | list[Trajectory] | BaseException]:
     if pbar_total_completion_tokens:
         print(
             "pbar_total_completion_tokens is deprecated and will be removed in a future version."

src/art/guided_completion.py

@@ -13,7 +13,9 @@
 from pydantic import create_model
 
 
-def freeze_tool_schema(tool: dict, fixed_args: dict) -> ChatCompletionToolParam:
+def freeze_tool_schema(
+    tool: ChatCompletionToolParam, fixed_args: dict[str, object]
+) -> ChatCompletionToolParam:
     """
     Return a clone of *tool* whose parameters schema permits *only* `fixed_args`.
     Each field is cast to typing.Literal[value] so Pydantic emits an
@@ -27,16 +29,16 @@ def freeze_tool_schema(tool: dict, fixed_args: dict) -> ChatCompletionToolParam:
 
     locked = deepcopy(tool)
     locked["function"]["parameters"] = FrozenModel.model_json_schema()
-    return locked  # type: ignore
+    return locked
 
 
 def get_guided_completion_params(
     completion: ChatCompletion,
     base_tools: Iterable[ChatCompletionToolParam] | None = None,
 ) -> Tuple[
-    List[str] | None,
+    List[str | None] | None,
     ChatCompletionToolChoiceOptionParam | None,
-    ChatCompletionToolParam | None,
+    List[ChatCompletionToolParam] | None,
 ]:
     """
     Given a completion from a teacher model, returns chat completion params that can be used to guide a student model's response.
@@ -48,7 +50,9 @@ def get_guided_completion_params(
 
     Returns a tuple of (guided_choice, tool_choice, tool_params).
     """
-    guided_choice, tool_choice, tool_params = None, None, None
+    guided_choice: List[str | None] | None = None
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = None
+    tool_params: List[ChatCompletionToolParam] | None = None
 
     if (
         completion.choices[0].message.tool_calls
@@ -74,4 +78,4 @@ def get_guided_completion_params(
     else:
         content = completion.choices[0].message.content
         guided_choice = [content]
-    return (guided_choice, tool_choice, tool_params)  # type: ignore
+    return (guided_choice, tool_choice, tool_params)

src/art/langgraph/message_utils.py

@@ -1,5 +1,5 @@
 import json
-from typing import List, Union
+from typing import List, Union, cast
 
 from langchain_core.messages import (
     AIMessage,
@@ -88,7 +88,7 @@ def langchain_msg_to_openai(msg: BaseMessage):
 
 
 def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
-    converted = []
+    converted: MessagesAndChoices = []
 
     for msg in messages:
         response_metadata = getattr(msg, "response_metadata")
@@ -116,7 +116,7 @@ def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
         elif isinstance(msg, BaseMessage):
             converted.append(langchain_msg_to_openai(msg))
         elif isinstance(msg, dict):
-            converted.append(msg)
+            converted.append(cast(Message, msg))
         else:
             raise TypeError(f"Unsupported message type: {type(msg)}")
 

src/art/local/backend.py

@@ -54,7 +54,7 @@
     aggregate_rl_training_metrics,
     build_rl_train_configs,
 )
-from ..backend import AnyTrainableModel, Backend
+from ..backend import AnyTrainableModel
 from ..dev.sequence_lengths import max_seq_length_from_model_config
 from ..metrics_taxonomy import (
     TRAIN_GRADIENT_STEPS_KEY,
@@ -75,6 +75,7 @@
 )
 from ..trajectories import Trajectory, TrajectoryGroup
 from ..types import (
+    Choice,
     LocalTrainResult,
     Message,
     TrainConfig,
@@ -189,7 +190,14 @@ def _tokenizer_cache_key(
     return (base_model, hashlib.sha256(chat_template.encode("utf-8")).hexdigest())
 
 
-class LocalBackend(Backend):
+def _load_training_tokenizer(base_model: str) -> PreTrainedTokenizerBase:
+    return cast(
+        PreTrainedTokenizerBase,
+        AutoTokenizer.from_pretrained(base_model),
+    )
+
+
+class LocalBackend:
     def __init__(
         self,
         *,
@@ -578,7 +586,7 @@ def _get_packed_tensors(
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
             tokenizer = self._configure_training_tokenizer(
-                AutoTokenizer.from_pretrained(model.base_model),
+                _load_training_tokenizer(model.base_model),
                 model=model,
                 internal_config=internal_config,
             )
@@ -777,10 +785,10 @@ def _trajectory_log(self, trajectory: Trajectory) -> str:
         header = f"reward: {trajectory.reward} {' '.join(f'{k}: {v}' for k, v in trajectory.metrics.items())}\n\n"
         formatted_messages = []
         for message_or_choice in trajectory.messages_and_choices:
-            if isinstance(message_or_choice, dict):
-                message = message_or_choice
+            if isinstance(message_or_choice, Choice):
+                message = cast(Message, message_or_choice.message.model_dump())
             else:
-                message = cast(Message, message_or_choice.message.model_dump())  # ty:ignore[possibly-missing-attribute]
+                message = message_or_choice
             formatted_messages.append(format_message(message))
         return header + "\n".join(formatted_messages)
 
@@ -1068,10 +1076,11 @@ async def _train_model(
                 try:
                     # Register the copied checkpoint as a new LoRA adapter
                     # so it's available for inference at the new step
-                    if hasattr(service, "register_lora_for_step"):
-                        await service.register_lora_for_step(  # type: ignore[attr-defined]
-                            next_step, next_checkpoint_dir
-                        )
+                    register_lora_for_step = getattr(
+                        service, "register_lora_for_step", None
+                    )
+                    if callable(register_lora_for_step):
+                        await register_lora_for_step(next_step, next_checkpoint_dir)
                     logger.info(
                         f"[BACKEND] _train_model SKIP: register_lora_for_step "
                         f"completed for step {next_step}"
@@ -1194,7 +1203,7 @@ async def _train_sft(
         tokenizer_key = _tokenizer_cache_key(model.base_model, internal_config)
         if tokenizer_key not in self._tokenizers:
             tokenizer = self._configure_training_tokenizer(
-                AutoTokenizer.from_pretrained(model.base_model),
+                _load_training_tokenizer(model.base_model),
                 model=model,
                 internal_config=internal_config,
             )
@@ -1364,7 +1373,8 @@ async def _experimental_pull_model_checkpoint(
                     f"No checkpoints found for {model.project}/{model.name} in local storage or S3"
                 )
             elif local_latest_step is None:
-                resolved_step = s3_latest_step  # type: ignore[assignment]
+                assert s3_latest_step is not None
+                resolved_step = s3_latest_step
                 if verbose:
                     print(f"Using latest checkpoint from S3: step {resolved_step}")
             elif s3_latest_step is None:

src/art/loss.py

@@ -73,7 +73,7 @@ def align_inputs(self) -> AlignedLossInputs:
             weights=shift_tensor(inputs["weights"], 0.0),
             group_ids=shift_tensor(inputs["group_ids"], 0),
             original_logprobs=(
-                shift_tensor(inputs["original_logprobs"], 0.0)  # ty: ignore[invalid-key]
+                shift_tensor(inputs["original_logprobs"], 0.0)
                 if "original_logprobs" in inputs
                 else None
             ),

src/art/mcp/generate_scenarios.py

@@ -5,6 +5,9 @@
 from typing import Any, Dict, List, Optional
 
 import openai
+from openai.types.shared_params.response_format_json_schema import (
+    ResponseFormatJSONSchema,
+)
 
 from art.mcp.types import GeneratedScenarioCollection, MCPResource, MCPTool
 from art.utils.logging import _C, dim, err, info, ok, step
@@ -137,7 +140,7 @@ async def generate_scenarios(
     if custom_instructions:
         prompt += f"\n\nPay close attention to the following instructions when generating scenarios:\n\n{custom_instructions}"
 
-    response_schema = {
+    response_schema: dict[str, object] = {
         "type": "object",
         "properties": {
             "scenarios": {
@@ -166,14 +169,15 @@ async def generate_scenarios(
     )
 
     t1 = time.perf_counter()
+    response_format: ResponseFormatJSONSchema = {
+        "type": "json_schema",
+        "json_schema": {"name": "scenario_list", "schema": response_schema},
+    }
     response = client_openai.chat.completions.create(
         model=generator_model,
         messages=[{"role": "user", "content": prompt}],
         max_completion_tokens=8000,
-        response_format={
-            "type": "json_schema",
-            "json_schema": {"name": "scenario_list", "schema": response_schema},
-        },
+        response_format=response_format,
     )
     dt = time.perf_counter() - t1
     ok(f"Model responded in {dt:.2f}s.")

src/art/megatron/context_parallel/executor.py

@@ -2235,31 +2235,31 @@ def _run_context_parallel_backward(
             grad_outputs=tuple(stage_output_grads),
             allow_unused=True,
         )
-        grad_map = {
+        grad_map: dict[str, torch.Tensor | None] = {
             name: grad for name, grad in zip(input_names, input_grads, strict=True)
         }
         for grad_name in ("q_input", "k_input", "v_input"):
             grad_map[grad_name] = _sanitize_nested_stage_input_grad(
-                cast(torch.Tensor | None, grad_map.get(grad_name)),
+                grad_map.get(grad_name),
             )
         _scatter_stage_grad(
             target=dq_flat,
-            grad=cast(torch.Tensor | None, grad_map.get("q_input")),
+            grad=grad_map.get("q_input"),
             ranges=stage_plan.owner_local_q_ranges,
             state=state,
             head_major=True,
         )
         if stage_plan.is_local_stage:
             _scatter_stage_grad(
                 target=dk_flat,
-                grad=cast(torch.Tensor | None, grad_map.get("k_input")),
+                grad=grad_map.get("k_input"),
                 ranges=stage_plan.owner_local_k_ranges,
                 state=state,
                 head_major=True,
             )
             _scatter_stage_grad(
                 target=dv_flat,
-                grad=cast(torch.Tensor | None, grad_map.get("v_input")),
+                grad=grad_map.get("v_input"),
                 ranges=stage_plan.owner_local_k_ranges,
                 state=state,
                 head_major=True,
@@ -2269,8 +2269,8 @@ def _run_context_parallel_backward(
             stage_record.clear()
             continue
         if not stage_plan.is_local_stage:
-            dk_remote = cast(torch.Tensor | None, grad_map.get("k_input"))
-            dv_remote = cast(torch.Tensor | None, grad_map.get("v_input"))
+            dk_remote = grad_map.get("k_input")
+            dv_remote = grad_map.get("v_input")
             if dk_remote is None:
                 dk_remote = k_flat.new_empty((k_flat.shape[0], 0, k_flat.shape[2]))
             if dv_remote is None:

src/art/megatron/dsv4/bridge.py

@@ -669,7 +669,7 @@ def __init__(
 
         @property
         def group_key(self) -> str:
-            return cast(dict[str, str], self.export_hf_param)["gate"]
+            return self.export_hf_param["gate"]
 
         def hf_to_megatron(
             self,
@@ -718,7 +718,7 @@ def megatron_to_hf(self, megatron_weights: Any, megatron_module: Any):
             hf_param = cast(dict[str, str], self.hf_param)
             gate_suffix = hf_param["gate"].rpartition(".experts.")[2].split(".", 1)[1]
             remapped: dict[str, torch.Tensor] = {}
-            export = cast(dict[str, str], self.export_hf_param)
+            export = self.export_hf_param
             for gate_key, gate in converted.items():
                 if not gate_key.endswith(gate_suffix):
                     continue
@@ -1269,5 +1269,5 @@ def ensure_dsv4_bridge_registered() -> None:
         target=GPTModel,
         provider=MLAModelProvider,
         model_type="deepseek_v4",
-    )(ArtDeepSeekV4Bridge)
+    )(ArtDeepSeekV4Bridge)  # ty: ignore[invalid-argument-type]
     _DSV4_BRIDGE_REGISTERED = True