S-Dreamer/CodeCraftLab

1"""2Training configuration schemas — Pydantic v2.3 4All training jobs are validated against these models before execution.5No raw dicts escape into the pipeline; everything is typed and constrained.6"""7 8from __future__ import annotations9 10from enum import StrEnum11from typing import Annotated12 13from pydantic import BaseModel, Field, HttpUrl, model_validator14from pydantic import PositiveFloat, PositiveInt15 16 17# ---------------------------------------------------------------------------18# Enums19# ---------------------------------------------------------------------------20class EvalStrategy(StrEnum):21    NO = "no"22    STEPS = "steps"23    EPOCH = "epoch"24 25 26class Precision(StrEnum):27    FP32 = "fp32"28    FP16 = "fp16"29    BF16 = "bf16"30    INT8 = "int8"31 32 33class OptimizerType(StrEnum):34    ADAMW = "adamw_torch"35    ADAMW_8BIT = "adamw_8bit"36    PAGED_ADAMW_8BIT = "paged_adamw_8bit"37    SGD = "sgd"38 39 40class EvalMetric(StrEnum):41    PASS_AT_1 = "pass_at_1"42    PASS_AT_10 = "pass_at_10"43    BLEU = "bleu"44    EXECUTION_ACCURACY = "execution_accuracy"45    EXACT_MATCH = "exact_match"46 47 48# ---------------------------------------------------------------------------49# Sub-configs50# ---------------------------------------------------------------------------51class LoRAConfig(BaseModel):52    """LoRA adapter configuration. Omit to disable LoRA (full fine-tune)."""53 54    enabled: bool = True55    r: Annotated[int, Field(ge=1, le=256)] = 1656    alpha: Annotated[int, Field(ge=1)] = 3257    dropout: Annotated[float, Field(ge=0.0, lt=1.0)] = 0.0558    target_modules: list[str] = Field(59        default_factory=lambda: ["q_proj", "v_proj"],60        min_length=1,61    )62    bias: str = "none"63 64    @model_validator(mode="after")65    def alpha_geq_r(self) -> "LoRAConfig":66        if self.alpha < self.r:67            raise ValueError(f"lora.alpha ({self.alpha}) should be >= lora.r ({self.r})")68        return self69 70 71class TrainingHyperparams(BaseModel):72    num_epochs: Annotated[int, Field(ge=1, le=100)] = 373    batch_size: Annotated[int, Field(ge=1, le=256)] = 874    gradient_accumulation_steps: Annotated[int, Field(ge=1, le=128)] = 475    learning_rate: Annotated[float, Field(gt=0.0, lt=1.0)] = 2e-576    weight_decay: Annotated[float, Field(ge=0.0, lt=1.0)] = 0.0177    warmup_ratio: Annotated[float, Field(ge=0.0, lt=1.0)] = 0.178    max_seq_length: Annotated[int, Field(ge=64, le=32768)] = 102479    max_grad_norm: Annotated[float, Field(gt=0.0)] = 1.080    optimizer: OptimizerType = OptimizerType.ADAMW81    precision: Precision = Precision.BF1682    lr_scheduler: str = "cosine"83    seed: int = 4284    dataloader_num_workers: Annotated[int, Field(ge=0, le=32)] = 485 86    @property87    def effective_batch_size(self) -> int:88        return self.batch_size * self.gradient_accumulation_steps89 90 91class EvaluationConfig(BaseModel):92    enabled: bool = True93    strategy: EvalStrategy = EvalStrategy.EPOCH94    eval_steps: PositiveInt | None = None  # required when strategy=STEPS95    metrics: list[EvalMetric] = Field(96        default_factory=lambda: [EvalMetric.PASS_AT_1, EvalMetric.BLEU]97    )98    num_samples_per_problem: Annotated[int, Field(ge=1, le=200)] = 1099    timeout_seconds: Annotated[int, Field(ge=1, le=60)] = 10100    load_best_model_at_end: bool = True101    metric_for_best_model: EvalMetric = EvalMetric.PASS_AT_1102    greater_is_better: bool = True103 104    @model_validator(mode="after")105    def eval_steps_required_for_steps_strategy(self) -> "EvaluationConfig":106        if self.strategy == EvalStrategy.STEPS and self.eval_steps is None:107            raise ValueError("evaluation.eval_steps is required when strategy='steps'")108        return self109 110 111class CheckpointConfig(BaseModel):112    save_strategy: EvalStrategy = EvalStrategy.EPOCH113    save_steps: PositiveInt | None = None114    save_total_limit: Annotated[int, Field(ge=1, le=20)] = 3115    output_dir: str = "./checkpoints"116    resume_from_checkpoint: str | None = None117 118    @model_validator(mode="after")119    def save_steps_required_for_steps_strategy(self) -> "CheckpointConfig":120        if self.save_strategy == EvalStrategy.STEPS and self.save_steps is None:121            raise ValueError("checkpoint.save_steps required when save_strategy='steps'")122        return self123 124 125class HubConfig(BaseModel):126    push_to_hub: bool = False127    repo_id: str | None = None128    private: bool = True129    commit_message: str = "Training checkpoint"130 131    @model_validator(mode="after")132    def repo_id_required_if_pushing(self) -> "HubConfig":133        if self.push_to_hub and not self.repo_id:134            raise ValueError("hub.repo_id is required when hub.push_to_hub=true")135        return self136 137 138class DatasetConfig(BaseModel):139    dataset_id: str  # internal UUID or HF Hub dataset path140    split_ratio: Annotated[float, Field(gt=0.0, lt=1.0)] = 0.9  # train split141    max_samples: PositiveInt | None = None  # None = use all142    text_column: str = "content"143    shuffle: bool = True144    shuffle_seed: int = 42145 146 147# ---------------------------------------------------------------------------148# Root job config149# ---------------------------------------------------------------------------150class TrainingJobConfig(BaseModel):151    """152    Complete training job specification.153 154    Validated at job submission time. If validation passes, the job is155    guaranteed to reach the pipeline with a coherent configuration.156    """157 158    job_name: Annotated[str, Field(min_length=1, max_length=128, pattern=r"^[\w\-]+$")]159    base_model: str = Field(160        description="HuggingFace model ID or local path",161        examples=["Salesforce/codegen-350M-mono", "deepseek-ai/deepseek-coder-1.3b-base"],162    )163    dataset: DatasetConfig164    training: TrainingHyperparams = Field(default_factory=TrainingHyperparams)165    lora: LoRAConfig | None = Field(default_factory=LoRAConfig)166    evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig)167    checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig)168    hub: HubConfig = Field(default_factory=HubConfig)169    tags: list[str] = Field(default_factory=list, max_length=20)170    notes: str | None = None171 172    model_config = {173        "json_schema_extra": {174            "examples": [175                {176                    "job_name": "codegen-finetune-v1",177                    "base_model": "Salesforce/codegen-350M-mono",178                    "dataset": {"dataset_id": "ds_abc123"},179                    "training": {180                        "num_epochs": 3,181                        "batch_size": 8,182                        "learning_rate": 2e-5,183                    },184                    "hub": {185                        "push_to_hub": True,186                        "repo_id": "your-org/codegen-finetune-v1",187                    },188                }189            ]190        }191    }192 193 194# ---------------------------------------------------------------------------195# Inference config (served separately but validated here for consistency)196# ---------------------------------------------------------------------------197class InferenceConfig(BaseModel):198    model_id: str199    max_new_tokens: Annotated[int, Field(ge=1, le=4096)] = 256200    temperature: Annotated[float, Field(ge=0.0, le=2.0)] = 0.2201    top_p: Annotated[float, Field(ge=0.0, le=1.0)] = 0.95202    top_k: Annotated[int, Field(ge=0, le=1000)] = 50203    do_sample: bool = True204    num_return_sequences: Annotated[int, Field(ge=1, le=200)] = 1205    stop_sequences: list[str] = Field(default_factory=list)206    precision: Precision = Precision.BF16207