nvidia-nat-openpipe-art 1.4.0a20260109__py3-none-any.whl

nat/plugins/openpipe/trainer_adapter.py

@@ -0,0 +1,339 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import art
+import httpx
+
+from nat.data_models.finetuning import EpisodeItem
+from nat.data_models.finetuning import EpisodeItemRole
+from nat.data_models.finetuning import FinetuneConfig
+from nat.data_models.finetuning import TrainingJobRef
+from nat.data_models.finetuning import TrainingJobStatus
+from nat.data_models.finetuning import TrainingStatusEnum
+from nat.data_models.finetuning import Trajectory
+from nat.data_models.finetuning import TrajectoryCollection
+from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
+
+from .config import ARTTrainerAdapterConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ARTTrainerAdapter(TrainerAdapter):
+    """
+    Adapter for the ART Trainer backend.
+    """
+
+    def __init__(self, adapter_config: ARTTrainerAdapterConfig):
+        super().__init__(adapter_config)
+
+        self.adapter_config: ARTTrainerAdapterConfig = adapter_config
+
+        self.remote_backend: art.Backend = art.Backend(
+            base_url=f"http://{adapter_config.backend.ip}:{adapter_config.backend.port}")
+
+        self._model_internal_config: art.dev.InternalModelConfig = art.dev.InternalModelConfig(
+            init_args=self.adapter_config.backend.init_args,
+            engine_args=self.adapter_config.backend.engine_args,
+            torchtune_args=self.adapter_config.backend.torchtune_args,
+            trainer_args=self.adapter_config.training)
+
+        self.model: art.TrainableModel = art.TrainableModel(
+            name=self.adapter_config.backend.name,
+            project=self.adapter_config.backend.project,
+            base_model=self.adapter_config.backend.base_model,
+            _internal_config=self._model_internal_config,
+        )
+
+        self._training_jobs: dict[str, asyncio.Task[None]] = {}
+
+        logger.info(f"Initialized ARTTrainerAdapter with model: {self.model}")
+
+    @property
+    def training_jobs(self) -> dict[str, asyncio.Task[None]]:
+        return self._training_jobs
+
+    async def initialize(self, run_config: FinetuneConfig) -> None:
+
+        await super().initialize(run_config)
+
+        await self.model.register(self.remote_backend, _openai_client_config=self.adapter_config.backend.server_config)
+
+        health = await self.is_healthy()
+
+        if not health:
+            raise ConnectionError("Failed to connect to ART backend.")
+
+        logger.info("Successfully registered with ART backend.")
+
+    async def is_healthy(self) -> bool:
+        try:
+            async with httpx.AsyncClient() as c:
+                await c.get(f"http://{self.adapter_config.backend.ip}:8000/v1/models",
+                            headers={"Authorization": f"Bearer {self.adapter_config.backend.api_key}"})
+            return True
+        except httpx.HTTPError as e:
+            logger.error(f"Health check failed: {e}")
+            return False
+
+    async def _validate_episode_order(self, traj: Trajectory):
+        """
+        Checks all EpisodeItem in traj.episode to validate:
+
+        - Every EpisodeItem.role is EpisodeItemRole.USER, SYSTEM, or ASSISTANT
+        - The first EpisodeItem.role is SYSTEM or USER
+        - The last EpisodeItem.role is ASSISTANT
+        - No two consecutive EpisodeItem.role are the same, except for SYSTEM
+
+        Args:
+            traj: Trajectory to validate
+
+        Raises:
+            ValueError: If any of the above conditions are not met.
+        """
+        if not traj.episode:
+            raise ValueError("Trajectory episode is empty.")
+
+        if traj.episode[0].role not in {EpisodeItemRole.USER, EpisodeItemRole.SYSTEM}:
+            raise ValueError("The first message in the trajectory must be from 'user' or 'system'.")
+
+        # if traj.episode[-1].role != EpisodeItemRole.ASSISTANT:
+        #     raise ValueError("The last message in the trajectory must be from 'assistant'.")
+
+        for i in range(1, len(traj.episode)):
+            if traj.episode[i].role == traj.episode[i - 1].role and traj.episode[i].role == EpisodeItemRole.ASSISTANT:
+                raise ValueError("Consecutive assistant messages from the same role found in trajectory.")
+
+    async def _construct_trajectory_groups(self, trajectory_lists: list[list[Trajectory]]) -> list[art.TrajectoryGroup]:
+        """
+        Convert list of lists of NAT Trajectory to list of ART TrajectoryGroup.
+
+        Args:
+            trajectory_lists: List of lists of NAT Trajectory (each inner list
+                contains trajectories for one example).
+
+        Returns:
+            List of ART TrajectoryGroup.
+
+        Raises:
+            ValueError: If any trajectory is invalid.
+        """
+
+        from openai.types.chat.chat_completion import Choice
+
+        # ---------- helpers ----------
+        def _as_text(obj: Any) -> str:
+            return obj if isinstance(obj, str) else json.dumps(obj, ensure_ascii=False)
+
+        def _to_chat_msg(d: EpisodeItem) -> dict:
+
+            if d.role == EpisodeItemRole.USER:
+                return {
+                    "role": "user",
+                    "content": _as_text(d.content),
+                }
+            elif d.role == EpisodeItemRole.SYSTEM:
+                return {
+                    "role": "system",
+                    "content": _as_text(d.content),
+                }
+            else:
+                return {"role": "assistant", "content": _as_text(d.content)}
+
+        output_trajectory_groups = []
+
+        for trajectory_list in trajectory_lists:
+            art_trajectories = []
+
+            for traj in trajectory_list:
+                episode = traj.episode
+                reward = traj.reward
+
+                # Validate episode order
+                await self._validate_episode_order(traj)
+
+                try:
+                    first_msg = _to_chat_msg(episode[0])
+
+                    t = art.Trajectory(messages_and_choices=[first_msg], reward=reward)
+
+                    for msg in episode[1:]:
+                        if msg.role == EpisodeItemRole.ASSISTANT:
+                            t.messages_and_choices.append(
+                                Choice(index=0, logprobs=msg.logprobs, message=_to_chat_msg(msg), finish_reason="stop"))
+                        else:
+                            t.messages_and_choices.append(_to_chat_msg(msg))
+
+                    # Sanity check: art.Trajectory.model_validate()
+                    t.model_validate(t.model_dump())
+
+                    art_trajectories.append(t)
+
+                except Exception as e:
+                    logger.error(f"Error constructing trajectory: {e}. Skipping.")
+                    continue
+
+            # Create TrajectoryGroup for this list of trajectories
+            if art_trajectories:
+                trajectory_group = art.TrajectoryGroup(trajectories=art_trajectories)
+                output_trajectory_groups.append(trajectory_group)
+
+        return output_trajectory_groups
+
+    async def submit(self, trajectories: TrajectoryCollection) -> TrainingJobRef:
+        """
+        Submit trajectories to ART backend for training.
+
+        Args:
+            trajectories: TrajectoryCollection with list of lists of NAT Trajectory.
+
+        Returns:
+            TrainingJobRef: Reference to the submitted training job.
+        """
+
+        trajectory_groups = await self._construct_trajectory_groups(trajectories.trajectories)
+        if not trajectory_groups:
+            raise ValueError("No valid trajectory groups to submit.")
+
+        assert trajectories.run_id not in self.training_jobs, (f"Training job "
+                                                               f"with run_id {trajectories.run_id} already exists.")
+
+        # Delete old remote checkpoints
+        if self.adapter_config.backend.delete_old_checkpoints:
+            try:
+                logger.info("Deleting old checkpoints on ART backend...")
+                await self.model.delete_checkpoints()
+            except Exception as e:
+                logger.warning(f"Failed to delete old checkpoints: {e}")
+
+        # Submit new trajectories
+        task = asyncio.create_task(
+            self.model.train(trajectory_groups=trajectory_groups,
+                             verbose=False,
+                             config=art.types.TrainConfig(
+                                 beta=getattr(self.adapter_config.training, "beta", 0),
+                                 learning_rate=getattr(self.adapter_config.training, "learning_rate", 5e-5),
+                             )),
+            name=f"art-train:{trajectories.run_id}",
+        )
+
+        # Optional: log + cleanup on completion to avoid leaks
+        def _on_done(t: asyncio.Task, rid: str = trajectories.run_id) -> None:
+            if t.cancelled():
+                logger.info(f"Training {rid} was cancelled.")
+            elif (exc := t.exception()) is not None:
+                logger.exception(f"Training {rid} failed", exc_info=exc)
+            else:
+                logger.info(f"Training {rid} completed successfully.")
+
+        task.add_done_callback(_on_done)
+
+        self.training_jobs[trajectories.run_id] = task
+
+        total_trajectories = sum(len(group.trajectories) for group in trajectory_groups)
+        logger.info(f"Submitted {total_trajectories} trajectories in {len(trajectory_groups)} groups for "
+                    f"training with run_id {trajectories.run_id}.")
+
+        return TrainingJobRef(run_id=trajectories.run_id, backend="openpipe-art")
+
+    async def status(self, ref: TrainingJobRef) -> TrainingJobStatus:
+        task = self.training_jobs.get(ref.run_id)
+        if task is None:
+            raise ValueError(f"No training job found with run_id {ref.run_id}.")
+
+        if task.done():
+            if task.cancelled():
+                status = TrainingStatusEnum.CANCELED
+                progress = None
+                message = "Training was cancelled."
+            else:
+                exc = task.exception()
+                if exc is not None:
+                    status = TrainingStatusEnum.FAILED
+                    progress = None
+                    message = f"Training failed with error: {exc!r}"
+                else:
+                    status = TrainingStatusEnum.COMPLETED
+                    progress = 100.0
+                    message = "Training completed successfully."
+
+            _ = self.training_jobs.pop(ref.run_id, None)  # Clean up completed job
+
+        else:
+            status = TrainingStatusEnum.RUNNING
+            progress = None
+            message = "Training is in progress."
+
+        return TrainingJobStatus(
+            run_id=ref.run_id,
+            backend=ref.backend,
+            status=status,
+            progress=progress,
+            message=message,
+        )
+
+    async def wait_until_complete(self, ref: TrainingJobRef, poll_interval: float = 10.0) -> TrainingJobStatus:
+        task = self.training_jobs.get(ref.run_id)
+        if task is None:
+            raise ValueError(f"No training job found with run_id {ref.run_id}.")
+
+        while not task.done():
+            await asyncio.sleep(poll_interval)
+
+        return await self.status(ref)
+
+    def log_progress(self, ref: TrainingJobRef, metrics: dict[str, Any], output_dir: str | None = None) -> None:
+        """
+        Log training adapter progress.
+
+        Args:
+            ref: Training job reference
+            metrics: Dictionary of metrics to log
+            output_dir: Optional output directory override
+        """
+        # Use default output directory if not provided
+        out_dir = Path(output_dir) if output_dir else Path("./.tmp/nat/finetuning/trainer_adapter")
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create log file for trainer adapter
+        log_file = out_dir / f"trainer_adapter_{ref.run_id}.jsonl"
+
+        # Prepare log entry
+        log_entry = {
+            "timestamp": datetime.now().isoformat(),
+            "run_id": ref.run_id,
+            "backend": ref.backend,
+            "trainer_config": {
+                "base_model": self.adapter_config.backend.base_model,
+                "project": self.adapter_config.backend.project,
+                "name": self.adapter_config.backend.name,
+            },
+            **metrics
+        }
+
+        # Append to log file
+        with open(log_file, 'a', encoding='utf-8') as f:
+            f.write(json.dumps(log_entry) + '\n')
+
+        logger.info("Trainer adapter progress logged for job %s: status=%s",
+                    ref.run_id,
+                    metrics.get("status", "unknown"))

nat/plugins/openpipe/trajectory_builder.py

@@ -0,0 +1,333 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import json
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from nat.data_models.finetuning import EpisodeItem
+from nat.data_models.finetuning import EpisodeItemRole
+from nat.data_models.finetuning import Trajectory
+from nat.data_models.finetuning import TrajectoryCollection
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepCategory
+from nat.eval.config import EvaluationRunOutput
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
+from nat.finetuning.utils.parsers.base_parser import parse_to_openai_messages
+
+from .config import ARTTrajectoryBuilderConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ARTTrajectoryBuilder(TrajectoryBuilder):
+    """
+    Trajectory builder for the ART backend.
+    """
+
+    def __init__(
+        self,
+        trajectory_builder_config: ARTTrajectoryBuilderConfig,
+    ):
+        super().__init__(trajectory_builder_config=trajectory_builder_config)
+        self.evaluation_runs: dict[str, list[asyncio.Task[EvaluationRunOutput]]] = {}
+
+    @property
+    def num_generations(self) -> int:
+        return self.trajectory_builder_config.num_generations
+
+    async def start_run(self, run_id: str, meta: dict | None = None) -> None:
+        """
+        Start multiple evaluation runs to collect trajectories.
+
+        Args:
+            run_id (str): The ID of the run.
+            meta (dict): Metadata for the run.
+        """
+
+        if run_id in self.evaluation_runs:
+            raise ValueError(f"Run {run_id} is already in progress.")
+
+        logger.info("Starting %d evaluation runs for run_id: %s", self.num_generations, run_id)
+        tasks = []
+
+        for gen_idx in range(self.num_generations):
+            task = asyncio.create_task(self.run_eval(), name=f"eval-run-{run_id}-gen-{gen_idx}")
+
+            def _on_done(t: asyncio.Task[EvaluationRunOutput], generation_index: int = gen_idx) -> None:
+                if t.cancelled():
+                    logger.info("Evaluation run for run_id: %s, generation: %d was cancelled.",
+                                run_id,
+                                generation_index)
+                elif exc := t.exception():
+                    logger.error(
+                        "Evaluation run for run_id: %s, generation: %d failed with exception: %s",
+                        run_id,
+                        generation_index,
+                        exc,
+                    )
+                else:
+                    logger.info(
+                        "Evaluation run for run_id: %s, generation: %d completed successfully.",
+                        run_id,
+                        generation_index,
+                    )
+
+            task.add_done_callback(_on_done)
+            tasks.append(task)
+
+        self.evaluation_runs[run_id] = tasks
+
+    async def finalize(self, run_id: str, meta: dict | None = None) -> TrajectoryCollection:
+        """
+        Waits for all evaluation runs to finalize and builds trajectories from
+        the episode items, grouping them by example ID.
+
+        Args:
+            run_id (str): The ID of the run.
+            meta (dict): Metadata for the run.
+
+        Returns:
+            TrajectoryCollection: The collection of built trajectories grouped by example.
+        """
+        from nat.eval.evaluator.evaluator_model import EvalOutputItem
+
+        if run_id not in self.evaluation_runs:
+            raise ValueError(f"No evaluation runs found for run_id: {run_id}")
+
+        # Wait for all evaluation runs to complete
+        tasks = self.evaluation_runs[run_id]
+        eval_results = await asyncio.gather(*tasks)
+
+        # Dictionary to group trajectories by example ID
+        trajectories_by_id: dict[str, list[Trajectory]] = {}
+
+        # Process each evaluation result
+        for gen_idx, eval_result in enumerate(eval_results):
+            reward_results: list[EvalOutputItem] | None = None
+            for metric_name, metric_value in eval_result.evaluation_results:
+                if metric_name == self.run_config.reward_function.name:
+                    reward_results = metric_value.eval_output_items
+                    break
+
+            if not reward_results:
+                logger.warning(f"No reward results found for run_id: {run_id}, generation: {gen_idx}")
+                continue
+
+            logger.info("Building trajectories for run_id: %s, generation: %d", run_id, gen_idx)
+
+            # ---------- helpers ----------
+            def _as_text(obj: Any) -> str:
+                return (obj if isinstance(obj, str) else json.dumps(obj, ensure_ascii=False))
+
+            def _parse_trajectory_from_steps(steps: list[IntermediateStep], ) -> list[EpisodeItem]:
+                """Parse trajectory from intermediate steps using parser."""
+                episode_items = []
+
+                try:
+                    # Use the base parser to convert to OpenAI messages
+                    openai_messages = parse_to_openai_messages(steps)
+
+                    # Convert OpenAI messages to EpisodeItems
+                    for msg in openai_messages:
+                        # Map OpenAI role to EpisodeItemRole
+                        role_mapping = {
+                            "user": EpisodeItemRole.USER,
+                            "assistant": EpisodeItemRole.ASSISTANT,
+                            "system": EpisodeItemRole.SYSTEM,
+                            "tool": EpisodeItemRole.TOOL,
+                            "function": EpisodeItemRole.FUNCTION,
+                            "human": EpisodeItemRole.USER,
+                            "ai": EpisodeItemRole.ASSISTANT,
+                        }
+
+                        role = role_mapping.get(msg.get("role"), EpisodeItemRole.OTHER)
+                        content = msg.get("content", "")
+                        logprobs = msg.get("logprobs")
+
+                        # For assistant messages, skip if no logprobs
+                        if role == EpisodeItemRole.ASSISTANT and not logprobs:
+                            logger.debug("Skipping assistant message without logprobs")
+                            continue
+
+                        # Build metadata from message attributes
+                        metadata = {}
+
+                        # Add tool/function specific metadata
+                        if "tool_call_id" in msg:
+                            metadata["tool_call_id"] = msg["tool_call_id"]
+                        if "tool_calls" in msg:
+                            metadata["tool_calls"] = msg["tool_calls"]
+                        if "function_call" in msg:
+                            metadata["function_call"] = msg["function_call"]
+                        if "name" in msg:
+                            metadata["name"] = msg["name"]
+
+                        episode_items.append(
+                            EpisodeItem(
+                                role=role,
+                                content=content,
+                                logprobs=logprobs,
+                                metadata=metadata if metadata else None,
+                            ))
+
+                except ValueError as e:
+                    logger.warning(
+                        "Failed to parse trajectory using base parser: %s. "
+                        "Falling back to empty episode.", str(e))
+                    # Return empty list on parse failure
+                    return []
+
+                return episode_items
+
+            # Create a mapping of id to input item for quick lookup
+            input_items_map = {item.id: item for item in eval_result.eval_input.eval_input_items}
+
+            for reward_item in reward_results:
+                # Find the corresponding input item
+                input_item: EvalInputItem = input_items_map.get(reward_item.id)
+                if not input_item:
+                    logger.warning(
+                        "No input item found for reward item id: %s",
+                        reward_item.id,
+                    )
+                    continue
+
+                filtered_trajectory = []
+                for item in input_item.trajectory:
+                    if item.function_ancestry.function_name in self.run_config.target_functions:
+                        # If target model is specified, filter by model name
+                        if (self.run_config.target_model and item.event_category == IntermediateStepCategory.LLM
+                                and item.payload.name != self.run_config.target_model):
+                            continue
+                        filtered_trajectory.append(item)
+
+                if not filtered_trajectory:
+                    logger.warning(
+                        "No trajectory steps found for target function '%s' in item id: %s",
+                        self.run_config.target_functions,
+                        reward_item.id,
+                    )
+                    continue
+
+                # Parse episode from intermediate steps
+                episode = _parse_trajectory_from_steps(filtered_trajectory)
+
+                # If no episode was parsed from steps, try to build from
+                # input/output
+                if not episode:
+                    continue
+
+                # Ensure we have at least a user and assistant message
+                if len(episode) < 2:
+                    logger.warning(
+                        "Episode for item %s has less than 2 messages, skipping",
+                        reward_item.id,
+                    )
+                    continue
+
+                # Validate that assistant messages have logprobs
+                # (required for training)
+                has_valid_assistant = False
+                for item in episode:
+                    if item.role == EpisodeItemRole.ASSISTANT and item.logprobs:
+                        has_valid_assistant = True
+                        break
+
+                if not has_valid_assistant:
+                    logger.warning(
+                        "Episode for item %s has no assistant messages with "
+                        "logprobs, skipping as it cannot be used for training",
+                        reward_item.id,
+                    )
+                    continue
+
+                # Create trajectory
+                trajectory = Trajectory(
+                    episode=episode,
+                    reward=(await self.compute_reward(reward_item, meta=meta)),
+                    metadata={
+                        "id": reward_item.id,
+                        "reasoning": reward_item.reasoning,
+                        "run_id": run_id,
+                        "generation": gen_idx,
+                    },
+                )
+
+                # Group by example ID
+                if reward_item.id not in trajectories_by_id:
+                    trajectories_by_id[reward_item.id] = []
+                trajectories_by_id[reward_item.id].append(trajectory)
+
+        # Clean up completed runs
+        self.evaluation_runs.pop(run_id, None)
+
+        # Convert dictionary to list of lists, maintaining order
+        trajectories_list = list(trajectories_by_id.values())
+
+        total_trajectories = sum(len(traj_list) for traj_list in trajectories_list)
+        logger.info("Built %d trajectories across %d examples for run_id: %s",
+                    total_trajectories,
+                    len(trajectories_list),
+                    run_id)
+
+        # Flatten the trajectories list into a 1 d list of trajectories
+        if not trajectories_list:
+            logger.warning("No trajectories were built for run_id: %s", run_id)
+            return TrajectoryCollection(trajectories=[], run_id=run_id)
+
+        if self.num_generations == 1:
+            # If only one generation, return flat list
+            flat_trajectories = [traj for sublist in trajectories_list for traj in sublist]
+            return TrajectoryCollection(trajectories=[flat_trajectories], run_id=run_id)
+
+        return TrajectoryCollection(trajectories=trajectories_list, run_id=run_id)
+
+    def log_progress(self, run_id: str, metrics: dict[str, Any], output_dir: str | None = None) -> None:
+        """
+        Log trajectory building progress.
+
+        Args:
+            run_id: The training run ID
+            metrics: Dictionary of metrics to log
+            output_dir: Optional output directory override
+        """
+        # Use default output directory if not provided
+        out_dir = Path(output_dir) if output_dir else Path("./.tmp/nat/finetuning/trajectory_builder")
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create log file for trajectory builder
+        log_file = out_dir / f"trajectory_builder_{run_id}.jsonl"
+
+        # Prepare log entry
+        log_entry = {
+            "timestamp": datetime.now().isoformat(),
+            "run_id": run_id,
+            "num_generations": self.num_generations,
+            **metrics
+        }
+
+        # Append to log file
+        with open(log_file, 'a', encoding='utf-8') as f:
+            f.write(json.dumps(log_entry) + '\n')
+
+        logger.debug("Trajectory builder progress logged for run %s: %d trajectories",
+                     run_id,
+                     metrics.get("num_trajectories", 0))