harmony-client 0.1.0__cp312-cp312-win_amd64.whl

harmony_client/runtime/context.py

@@ -0,0 +1,191 @@
+import inspect
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Self
+
+import pandas as pd
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from rich.console import Console
+from rich.table import Table as RichTable
+
+from harmony_client import EvalSample, HarmonyClient, HarmonyJobNotifier, JobNotifier, StringThread, get_client
+from harmony_client.file_storage import FileStorage, FileStorageConfig
+from harmony_client.internal import _extract_model_key, _save_detailed_eval_table
+
+
+class RecipeConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="ADAPTIVE_", cli_parse_args=True, cli_kebab_case=True)
+
+    harmony_url: str = Field(description="url of harmony service")
+    control_plane_url: str | None = Field(
+        default=None,
+        description="URL of the control plane service (Concorde). Required for fetching grader, dataset, and model configurations. Set this if you need to access centralized configuration; may be omitted for local or test runs where such configurations are not needed."
+    )
+    control_plane_api_token: str | None = Field(
+        default=None,
+        description="JWT token for authenticating with the control plane service (Concorde). Required when accessing protected internal API endpoints; may be omitted for local or test runs."
+    )
+    user_input_file: str | None = None
+    job_id: str = "test"
+    use_case: str | None = None
+    api_key: str | None = None
+    compute_pool: str | None = None
+    storage_url: str | None = None
+    num_gpus: int = 0
+
+
+class RecipeContext:
+    client: HarmonyClient
+    job: JobNotifier
+    file_storage: FileStorage
+    config: RecipeConfig
+    # todo: pass world size
+    world_size: int = 1
+
+    def __init__(self, client: HarmonyClient, config: RecipeConfig):
+        self.client = client
+        self.config = config
+        self.job = HarmonyJobNotifier(client, config.job_id)
+        print(f"{config.storage_url=}")
+        if config.storage_url:
+            self.file_storage = FileStorage.new(FileStorageConfig.from_url(config.storage_url))
+        else:
+            self.file_storage = FileStorage.new(FileStorageConfig.from_url("file:///tmp/recipe_storage"))
+
+    @classmethod
+    async def load(cls) -> Self:
+        config = RecipeConfig()  # type: ignore
+        return await cls.from_config(config)
+
+    @classmethod
+    async def from_config(cls, config: RecipeConfig) -> Self:
+        client = await get_client(
+            config.harmony_url,
+            num_gpus=config.num_gpus,
+            api_key=config.api_key,
+            use_case=config.use_case,
+            compute_pool=config.compute_pool,
+            job_id=config.job_id,
+            control_plane_url=config.control_plane_url,
+            control_plane_api_token=config.control_plane_api_token,
+        )
+        return cls(client, config)
+
+    def load_dataset(self, path: str) -> list[StringThread]:
+        lines = self.file_storage.read(path, use_raw_path=True).decode("utf-8").splitlines()
+        threads = []
+        for line in lines:
+            line_dict = json.loads(line)
+            thread = None
+            if "input" in line_dict or "messages" in line_dict:
+                key = "input" if "input" in line_dict else "messages"
+                thread = StringThread(
+                    [(inner_turn_dict["role"], inner_turn_dict["content"]) for inner_turn_dict in line_dict[key]]
+                )
+                if "completion" in line_dict and line_dict["completion"]:
+                    thread = thread.assistant(line_dict["completion"])
+            else:
+                print("Did not find `input`, or `messages` key in sample, ignoring")
+
+            if thread is not None:
+                thread.metadata = line_dict.get("metadata", {})
+                if "other_completion" in line_dict and "preferred_completion" in line_dict:
+                    thread.metadata["other_completion"] = line_dict["other_completion"]
+                    thread.metadata["preferred_completion"] = line_dict["preferred_completion"]
+
+                threads.append(thread)
+
+        if len(threads) == 0:
+            raise ValueError("Did not find any valid format samples in the dataset")
+
+        return threads
+
+    @staticmethod
+    def log_eval_result(eval_samples: list[EvalSample]) -> None:
+        # Convert to DataFrame for easy aggregation
+        data = []
+        for eval_sample in eval_samples:
+            for grade in eval_sample.grades:
+                data.append(
+                    {
+                        "model": _extract_model_key(eval_sample.interaction.source),
+                        "grader": grade.grader_key,
+                        "score": grade.value,
+                    }
+                )
+        df = pd.DataFrame(data)
+
+        # Create pivot table with models as rows and graders as columns
+        if not df.empty:
+            pivot_df = df.pivot_table(index="model", columns="grader", values="score", aggfunc="mean")
+
+            # Create Rich table from pivot
+            console = Console()
+            table = RichTable(title="GRADER EVALUATION RESULTS")
+            table.add_column("Model", style="cyan", no_wrap=True)
+            # One column per grader
+            for grader_name in pivot_df.columns:
+                table.add_column(grader_name, justify="center")
+
+            # Find the maximum score for each grader (excluding NaN values)
+            max_scores = {}
+            for grader_name in pivot_df.columns:
+                valid_scores = pivot_df[grader_name].dropna()
+                if len(valid_scores) > 0:
+                    max_scores[grader_name] = valid_scores.max()
+
+            # One row per model
+            for model_name in pivot_df.index:
+                row = [model_name]
+                for grader_name in pivot_df.columns:
+                    score = pivot_df.loc[model_name, grader_name]
+                    # Avg will be nan if grader failed on all samples
+                    if pd.isna(score):
+                        row.append("All failed")
+                    else:
+                        score_str = f"{score:.3f}"
+                        # Highlight the highest score in green
+                        if grader_name in max_scores and abs(score - max_scores[grader_name]) < 1e-9:
+                            row.append(f"[green]{score_str}[/green]")
+                        else:
+                            row.append(score_str)
+                table.add_row(*row)
+
+            console.print(table)
+
+            # Get the directory of the original caller (skip this method's frame)
+            caller_frame = inspect.currentframe()
+            original_caller_dir = None
+            if caller_frame is not None and caller_frame.f_back is not None:
+                original_caller_file = caller_frame.f_back.f_globals["__file__"]
+                original_caller_dir = str(Path(original_caller_file).parent)
+
+            # Save compact JSON of grader results
+            original_caller_dir = original_caller_dir or Path.cwd()
+            results_json = {}
+            for model_name in pivot_df.index:
+                results_json[model_name] = {}
+                for grader_name in pivot_df.columns:
+                    score = pivot_df.loc[model_name, grader_name]
+                    if pd.isna(score):
+                        results_json[model_name][grader_name] = None
+                    else:
+                        # Convert pandas scalar to native Python type
+                        score_val = score.item() if hasattr(score, "item") else float(score)  # type: ignore
+                        results_json[model_name][grader_name] = round(score_val, 3)  # type:ignore
+
+            timestamp = datetime.now().replace(microsecond=0).isoformat()
+            eval_dir = Path(original_caller_dir) / "adaptive_eval_samples" / timestamp
+            eval_dir.mkdir(parents=True, exist_ok=True)
+
+            results_path = eval_dir / "aggregate_scores.json"
+            with open(results_path, "w") as f:
+                json.dump(results_json, f, indent=2)
+            print(f"\n📁 Aggregate results saved to: {results_path}")
+
+            # Save detailed evaluation samples as html
+            _save_detailed_eval_table(eval_samples, output_dir=str(eval_dir))
+        else:
+            print("No evaluation data to display")

harmony_client/runtime/data.py

@@ -0,0 +1,76 @@
+import json
+import types
+from typing import Annotated, Literal, Self, Union, get_args, get_origin
+
+from pydantic import BaseModel, Field
+
+
+class InputConfig(BaseModel):
+    def __init_subclass__(cls, **kwargs):
+        # Add union_variant BEFORE calling super (before Pydantic processes it)
+        if not hasattr(cls, "__annotations__"):
+            cls.__annotations__ = {}
+        if "union_variant" not in cls.__annotations__:
+            cls.__annotations__["union_variant"] = Annotated[Literal[cls.__name__], Field(default=cls.__name__)]
+        super().__init_subclass__(**kwargs)
+        if hasattr(cls, "__annotations__"):
+            for field_name, field_type in cls.__annotations__.items():
+                origin = get_origin(field_type)
+                if origin is Union or origin is types.UnionType:
+                    # Add discriminator to the union annotation
+                    cls.__annotations__[field_name] = Annotated[field_type, Field(discriminator="union_variant")]
+
+                    # Add union_variant field to each variant class
+                    variant_types = get_args(field_type)
+                    for variant_type in variant_types:
+                        # Skip None (for Optional types)
+                        if variant_type is type(None):
+                            continue
+
+                        # Add union_variant: Literal["ClassName"] to the variant
+                        # if not already present
+                        if (
+                            hasattr(variant_type, "__annotations__")
+                            and "union_variant" not in variant_type.__annotations__
+                        ):
+                            variant_type.__annotations__["union_variant"] = Literal[variant_type.__name__]
+                            # Set default value
+                            setattr(variant_type, "union_variant", variant_type.__name__)
+
+    @classmethod
+    def load_from_file(cls, json_file) -> Self:
+        with open(json_file) as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
+
+from .dto.AdaptiveDataset import AdaptiveDatasetKind
+from .dto.AdaptiveGrader import (
+    AdaptiveGrader as AdaptiveGrader,
+)
+from .dto.AdaptiveGrader import (
+    Judge as CustomJudge,
+)
+from .dto.AdaptiveGrader import (
+    JudgeExample as CustomJudgeExample,
+)
+from .dto.AdaptiveGrader import (
+    Prebuilt as PrebuiltJudge,
+)
+from .dto.AdaptiveGrader import (
+    PrebuiltConfigKey,
+)
+from .dto.AdaptiveGrader import (
+    Remote as RemoteRewardEndpoint,
+)
+
+__all__ = [
+    "InputConfig",
+    "AdaptiveDatasetKind",
+    "AdaptiveGrader",
+    "CustomJudge",
+    "CustomJudgeExample",
+    "PrebuiltJudge",
+    "PrebuiltConfigKey",
+    "RemoteRewardEndpoint",
+]

harmony_client/runtime/decorators.py

@@ -0,0 +1,19 @@
+from typing import Any, Callable, TypeVar, overload
+
+from harmony_client.runtime.context import RecipeContext
+from harmony_client.runtime.data import InputConfig
+
+IN = TypeVar("IN", bound=InputConfig)
+
+
+@overload
+def recipe_main[IN: InputConfig](func: Callable[[IN, RecipeContext], Any]): ...
+
+
+@overload
+def recipe_main(func: Callable[[RecipeContext], Any]): ...
+
+
+def recipe_main(func):
+    func.is_recipe_main = True
+    return func

harmony_client/runtime/dto/AdaptiveDataset.py

@@ -0,0 +1,23 @@
+# generated by datamodel-codegen:
+#   filename:  AdaptiveDataset.json
+
+from __future__ import annotations
+from enum import Enum
+from typing import Optional
+from .base import DtoBaseModel
+
+
+class AdaptiveDatasetKind(Enum):
+    Prompt = "Prompt"
+    Completion = "Completion"
+    Metric = "Metric"
+    Preference = "Preference"
+    Mixed = "Mixed"
+
+
+class AdaptiveDataset(DtoBaseModel):
+    dataset_key: str
+    feedback_key: Optional[str] = None
+    file: str
+    id: str
+    kind: Optional[AdaptiveDatasetKind] = AdaptiveDatasetKind.Mixed

harmony_client/runtime/dto/AdaptiveGrader.py

@@ -0,0 +1,68 @@
+# generated by datamodel-codegen:
+#   filename:  AdaptiveGrader.json
+
+from __future__ import annotations
+from typing import Annotated, Any, List, Literal, Optional, Union
+from .base import DtoBaseModel
+from uuid import UUID
+from pydantic import Field
+from enum import Enum
+
+
+class ChatMessage(DtoBaseModel):
+    content: str
+    metadata: Optional[Any] = None
+    role: str
+
+
+class Custom(DtoBaseModel):
+    description: Optional[str] = None
+    type: Literal["Custom"] = "Custom"
+
+
+class JudgeExample(DtoBaseModel):
+    id: UUID
+    input: List[ChatMessage]
+    output: str
+    pass_: Annotated[bool, Field(alias="pass")]
+    reasoning: Optional[str] = None
+
+
+class PrebuiltConfigKey(Enum):
+    AnswerRelevancy = "AnswerRelevancy"
+    ContextRelevancy = "ContextRelevancy"
+    Faithfulness = "Faithfulness"
+
+
+class Remote(DtoBaseModel):
+    description: str
+    type: Literal["Remote"] = "Remote"
+    url: str
+    version: str
+
+
+class Judge(DtoBaseModel):
+    criteria: str
+    examples: List[JudgeExample]
+    model_key: str
+    model_uri: str
+    system_template: str
+    type: Literal["Judge"] = "Judge"
+    user_template: str
+
+
+class Prebuilt(DtoBaseModel):
+    model_key: str
+    model_uri: str
+    prebuilt_config_key: PrebuiltConfigKey
+    type: Literal["Prebuilt"] = "Prebuilt"
+
+
+class AdaptiveGrader(DtoBaseModel):
+    config: Annotated[
+        Union[Judge, Prebuilt, Remote, Custom], Field(discriminator="type")
+    ]
+    grader_id: UUID
+    key: str
+    metric_id: UUID
+    name: str

harmony_client/runtime/dto/AdaptiveModel.py

@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  AdaptiveModel.json
+
+from __future__ import annotations
+from typing import Annotated, Optional
+from pydantic import Field
+from .base import DtoBaseModel
+
+
+class ModelParams(DtoBaseModel):
+    kv_cache_len: Annotated[Optional[int], Field(ge=0)] = None
+    max_seq_len: Annotated[Optional[int], Field(ge=0)] = None
+    tp: Annotated[Optional[int], Field(ge=0)] = None
+
+
+class AdaptiveModel(DtoBaseModel):
+    model_key: Optional[str] = None
+    params: Optional[ModelParams] = None
+    path: str

harmony_client/runtime/dto/DatasetSampleFormats.py

@@ -0,0 +1,93 @@
+# generated by datamodel-codegen:
+#   filename:  DatasetSampleFormats.json
+
+from __future__ import annotations
+from typing import Annotated, Any, Dict, List, Optional, Union
+from uuid import UUID
+from .base import DtoBaseModel
+from pydantic import Field, RootModel
+
+
+class SampleMetadata(DtoBaseModel):
+    created_at: int
+    external_data: Optional[Any] = None
+    id: UUID
+    labels: Optional[Dict[str, Any]] = None
+    model_id: Optional[UUID] = None
+
+
+class Type(DtoBaseModel):
+    enum: Optional[List[str]] = None
+    type: Optional[str] = None
+
+
+class Value(DtoBaseModel):
+    type: Optional[str] = None
+
+
+class Properties(DtoBaseModel):
+    type: Optional[Type] = None
+    value: Optional[Value] = None
+
+
+class OneOf(DtoBaseModel):
+    properties: Optional[Properties] = None
+    required: Optional[List[str]] = None
+    type: Optional[str] = None
+
+
+class Value1(DtoBaseModel):
+    format: Optional[str] = None
+    type: Optional[str] = None
+
+
+class Properties1(DtoBaseModel):
+    type: Optional[Type] = None
+    value: Optional[Value1] = None
+
+
+class OneOf1(DtoBaseModel):
+    properties: Optional[Properties1] = None
+    required: Optional[List[str]] = None
+    type: Optional[str] = None
+
+
+class TurnContent(DtoBaseModel):
+    oneOf: Optional[List[Union[OneOf, OneOf1]]] = None
+
+
+class TurnTuple(RootModel[List]):
+    root: Annotated[List, Field(max_length=2, min_length=2)]
+
+
+class DatasetMetricSample(DtoBaseModel):
+    completion: TurnTuple
+    metadata: SampleMetadata
+    metrics: Dict[str, float]
+    prompt: List[TurnTuple]
+
+
+class DatasetPreferenceSample(DtoBaseModel):
+    bad_completion: TurnTuple
+    good_completion: TurnTuple
+    metadata: SampleMetadata
+    metric: Optional[str] = "preference"
+    prompt: List[TurnTuple]
+
+
+class DatasetPromptSample(DtoBaseModel):
+    metadata: SampleMetadata
+    prompt: List[TurnTuple]
+
+
+class DatasetSample(DtoBaseModel):
+    completion: TurnTuple
+    metadata: SampleMetadata
+    prompt: List[TurnTuple]
+
+
+class DatasetSampleFormats(DtoBaseModel):
+    completion: DatasetSample
+    metric: DatasetMetricSample
+    preference: DatasetPreferenceSample
+    prompt: DatasetPromptSample

harmony_client/runtime/dto/__init__.py

@@ -0,0 +1,2 @@
+# generated by datamodel-codegen:
+#   filename:  jsonschema

harmony_client/runtime/dto/base.py

@@ -0,0 +1,7 @@
+# generated file
+
+from pydantic import BaseModel, ConfigDict
+
+
+class DtoBaseModel(BaseModel):
+    model_config = ConfigDict()

harmony_client/runtime/model_artifact_save.py

@@ -0,0 +1,23 @@
+from collections.abc import Awaitable, Callable
+
+from harmony_client import TrainingModel
+from harmony_client.artifacts.model_artifact import ModelArtifact
+from harmony_client.runtime import RecipeContext
+
+
+async def save_with_artifact(
+    model: TrainingModel,
+    model_name: str,
+    inference_only: bool = True,
+    ctx: RecipeContext | None = None,
+    original_save_method: Callable[[TrainingModel, str, bool], Awaitable[str]] | None = None,
+) -> str:
+    if original_save_method is None:
+        raise ValueError("original_save_method must be provided")
+
+    real_model_key = await original_save_method(model, model_name, inference_only)
+
+    if ctx is not None:
+        ModelArtifact(real_model_key, ctx)
+
+    return real_model_key