rft create uses sdk now

Author: xzrderek

eval_protocol/cli.py

@@ -402,7 +402,8 @@ def _configure_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParse
         help="Extra flags to pass to 'docker run' when validating evaluator (quoted string, e.g. \"--env-file .env --memory=8g\")",
     )
 
-    # Everything below has to manually be maintained, can't be auto-generated
+    # The flags below are Eval Protocol CLI workflow controls (not part of the Fireworks SDK `create()` signature),
+    # so they can’t be auto-generated via signature introspection and must be maintained here.
     rft_parser.add_argument(
         "--source-job",
         metavar="",
@@ -419,11 +420,9 @@ def _configure_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParse
             "extra_query",
             "extra_body",
             "timeout",
-            "node_count",
             "display_name",
             "account_id",
         },
-        "loss_config": {"kl_beta", "method"},
         "training_config": {"region", "jinja_template"},
         "wandb_config": {"run_id"},
     }
@@ -433,11 +432,15 @@ def _configure_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParse
         "wandb_config.entity": ["--wandb-entity"],
         "wandb_config.enabled": ["--wandb"],
         "reinforcement_fine_tuning_job_id": ["--job-id"],
+        "loss_config.kl_beta": ["--rl-kl-beta"],
+        "loss_config.method": ["--rl-loss-method"],
+        "node_count": ["--nodes"],
     }
     help_overrides = {
         "training_config.gradient_accumulation_steps": "The number of batches to accumulate gradients before updating the model parameters. The effective batch size will be batch-size multiplied by this value.",
         "training_config.learning_rate_warmup_steps": "The number of learning rate warmup steps for the reinforcement fine-tuning job.",
         "mcp_server": "The MCP server resource name to use for the reinforcement fine-tuning job. (Optional)",
+        "loss_config.method": "RL loss method for underlying trainers. One of {grpo,dapo}.",
     }
 
     create_rft_job_fn = Fireworks().reinforcement_fine_tuning_jobs.create

eval_protocol/cli_commands/create_rft.py

@@ -1,10 +1,12 @@
 import argparse
+from fireworks._client import Fireworks
+from fireworks.types.reinforcement_fine_tuning_job import ReinforcementFineTuningJob
 import json
 import os
 import sys
 import time
 from typing import Any, Dict, Optional
-
+import inspect
 import requests
 from pydantic import ValidationError
 
@@ -13,7 +15,6 @@
 from ..fireworks_rft import (
     build_default_output_model,
     create_dataset_from_jsonl,
-    create_reinforcement_fine_tuning_job,
     detect_dataset_builder,
     materialize_dataset_via_builder,
 )
@@ -33,6 +34,8 @@
 )
 from .local_test import run_evaluator_test
 
+from fireworks import Fireworks
+
 
 def _extract_jsonl_from_dataloader(test_file_path: str, test_func_name: str) -> Optional[str]:
     """Import the test module and extract a JSONL path from data_loaders param if present.
@@ -619,126 +622,48 @@ def _create_rft_job(
     args: argparse.Namespace,
     dry_run: bool,
 ) -> int:
-    """Build and submit the RFT job request."""
-    # Build training config/body
-    # Exactly one of base-model or warm-start-from must be provided
-    base_model_raw = getattr(args, "base_model", None)
-    warm_start_from_raw = getattr(args, "warm_start_from", None)
-    # Treat empty/whitespace strings as not provided
-    base_model = base_model_raw.strip() if isinstance(base_model_raw, str) else base_model_raw
-    warm_start_from = warm_start_from_raw.strip() if isinstance(warm_start_from_raw, str) else warm_start_from_raw
-    has_base_model = bool(base_model)
-    has_warm_start = bool(warm_start_from)
-    if (not has_base_model and not has_warm_start) or (has_base_model and has_warm_start):
-        print("Error: exactly one of --base-model or --warm-start-from must be specified.")
-        return 1
+    """Build and submit the RFT job request (via Fireworks SDK)."""
 
-    training_config: Dict[str, Any] = {}
-    if has_base_model:
-        training_config["baseModel"] = base_model
-    if has_warm_start:
-        training_config["warmStartFrom"] = warm_start_from
-
-    # Optional hyperparameters
-    for key, arg_name in [
-        ("epochs", "epochs"),
-        ("batchSize", "batch_size"),
-        ("learningRate", "learning_rate"),
-        ("maxContextLength", "max_context_length"),
-        ("loraRank", "lora_rank"),
-        ("gradientAccumulationSteps", "gradient_accumulation_steps"),
-        ("learningRateWarmupSteps", "learning_rate_warmup_steps"),
-        ("acceleratorCount", "accelerator_count"),
-        ("region", "region"),
-    ]:
-        val = getattr(args, arg_name, None)
-        if val is not None:
-            training_config[key] = val
-
-    inference_params: Dict[str, Any] = {}
-    for key, arg_name in [
-        ("temperature", "temperature"),
-        ("topP", "top_p"),
-        ("topK", "top_k"),
-        ("maxOutputTokens", "max_output_tokens"),
-        ("responseCandidatesCount", "response_candidates_count"),
-    ]:
-        val = getattr(args, arg_name, None)
-        if val is not None:
-            inference_params[key] = val
-    if getattr(args, "extra_body", None):
-        extra = getattr(args, "extra_body")
-        if isinstance(extra, (dict, list)):
-            try:
-                inference_params["extraBody"] = json.dumps(extra, ensure_ascii=False)
-            except (TypeError, ValueError) as e:
-                print(f"Error: --extra-body dict/list must be JSON-serializable: {e}")
-                return 1
-        elif isinstance(extra, str):
-            inference_params["extraBody"] = extra
-        else:
-            print("Error: --extra-body must be a JSON string or a JSON-serializable dict/list.")
-            return 1
+    signature = inspect.signature(Fireworks().reinforcement_fine_tuning_jobs.create)
 
-    wandb_config: Optional[Dict[str, Any]] = None
-    if getattr(args, "enabled", False):
-        wandb_config = {
-            "enabled": True,
-            "apiKey": getattr(args, "api_key", None),
-            "project": getattr(args, "project", None),
-            "entity": getattr(args, "entity", None),
-            "runId": getattr(args, "run_id", None),
-        }
-
-    body: Dict[str, Any] = {
-        "displayName": getattr(args, "display_name", None),
-        "dataset": dataset_resource,
+    # Build top-level SDK kwargs
+    sdk_kwargs: Dict[str, Any] = {
         "evaluator": evaluator_resource_name,
-        "evalAutoCarveout": bool(getattr(args, "eval_auto_carveout", True)),
-        "trainingConfig": training_config,
-        "inferenceParameters": inference_params or None,
-        "wandbConfig": wandb_config,
-        "chunkSize": getattr(args, "chunk_size", None),
-        "outputStats": None,
-        "outputMetrics": None,
-        "mcpServer": getattr(args, "mcp_server", None),
-        "jobId": getattr(args, "reinforcement_fine_tuning_job_id", None),
-        "sourceJob": getattr(args, "source_job", None),
-        "quiet": getattr(args, "quiet", False),
+        "dataset": dataset_resource,
     }
-    # Debug: print minimal summary
-    print(f"Prepared RFT job for evaluator '{evaluator_id}' using dataset '{dataset_id}'")
-    if getattr(args, "evaluation_dataset", None):
-        body["evaluationDataset"] = args.evaluation_dataset
 
-    output_model_arg = getattr(args, "output_model", None)
-    if output_model_arg:
-        if len(output_model_arg) > 63:
-            print(f"Error: Output model name '{output_model_arg}' exceeds 63 characters.")
-            return 1
-        body.setdefault("trainingConfig", {})["outputModel"] = f"accounts/{account_id}/models/{output_model_arg}"
-    else:
-        # Auto-generate output model name if not provided
-        auto_output_model = build_default_output_model(evaluator_id)
-        body.setdefault("trainingConfig", {})["outputModel"] = f"accounts/{account_id}/models/{auto_output_model}"
+    args_dict = vars(args)
+    for name in signature.parameters:
+        prefix = name + "_"
+
+        # Collect "flattened" argparse fields back into the nested dict expected by the SDK.
+        # Example: training_config_epochs=3 becomes sdk_kwargs["training_config"]["epochs"] = 3.
+        nested = {}
+        for k, v in args_dict.items():
+            if v is None:
+                continue
+            if not k.startswith(prefix):
+                continue
+            nested[k[len(prefix) :]] = v
+
+        if nested:
+            sdk_kwargs[name] = nested
+        elif args_dict.get(name) is not None:
+            sdk_kwargs[name] = args_dict[name]
 
-    # Clean None fields to avoid noisy payloads
-    body = {k: v for k, v in body.items() if v is not None}
+    print(f"Prepared RFT job for evaluator '{evaluator_id}' using dataset '{dataset_id}'")
 
     if dry_run:
-        print("--dry-run: would create RFT job with body:")
-        print(json.dumps(body, indent=2))
+        print("--dry-run: would call Fireworks().reinforcement_fine_tuning_jobs.create with kwargs:")
+        print(json.dumps(sdk_kwargs, indent=2))
         _print_links(evaluator_id, dataset_id, None)
         return 0
 
     try:
-        result = create_reinforcement_fine_tuning_job(
-            account_id=account_id, api_key=api_key, api_base=api_base, body=body
-        )
-        job_name = result.get("name") if isinstance(result, dict) else None
-        print("\n✅ Created Reinforcement Fine-tuning Job")
-        if job_name:
-            print(f"   name: {job_name}")
+        fw: Fireworks = Fireworks(api_key=api_key, base_url=api_base)
+        job: ReinforcementFineTuningJob = fw.reinforcement_fine_tuning_jobs.create(account_id=account_id, **sdk_kwargs)
+        job_name = job.name
+        print(f"\n✅ Created Reinforcement Fine-tuning Job: {job_name}")
         _print_links(evaluator_id, dataset_id, job_name)
         return 0
     except Exception as e:

eval_protocol/cli_commands/utils.py

@@ -660,7 +660,6 @@ def _add_flag(
 def add_args_from_callable_signature(
     parser: argparse.ArgumentParser,
     fn: Callable[..., Any],
-    *,
     overrides: dict[str, str] | None = None,
     skip_fields: dict[str, set[str]] | None = None,
     aliases: dict[str, list[str]] | None = None,
@@ -676,7 +675,7 @@ def add_args_from_callable_signature(
     help = _parse_args_section_from_doc(inspect.getdoc(fn) or "")
     hints = typing.get_type_hints(fn, include_extras=True)
 
-    for name, param in sig.parameters.items():
+    for name in sig.parameters.keys():
         resolved_type = unwrap_union(hints.get(name))
 
         # Allow one nested layer of TypeDicts
@@ -688,8 +687,10 @@ def add_args_from_callable_signature(
             for field_name, field_type in resolved_type.__annotations__.items():
                 if field_name in field_skip:
                     continue
-                flag_name = "--" + field_name.replace("_", "-")
-                flags = [flag_name] + aliases.get(f"{name}.{field_name}", [])
+                prefix = name.replace("_", "-")
+                field_kebab = field_name.replace("_", "-")
+                flag_name = f"--{prefix}-{field_kebab}"
+                flags = [flag_name] + aliases.get(f"{name}.{field_name}", []) + [f"--{field_kebab}"]
                 help_text = help_overrides.get(f"{name}.{field_name}", field_help.get(field_name))
 
                 _add_flag(parser, flags, field_hints.get(field_name, field_type), help_text)

eval_protocol/fireworks_rft.py

@@ -181,33 +181,6 @@ def create_dataset_from_jsonl(
     return dataset_id, ds
 
 
-def create_reinforcement_fine_tuning_job(
-    account_id: str,
-    api_key: str,
-    api_base: str,
-    body: Dict[str, Any],
-) -> Dict[str, Any]:
-    url = f"{api_base.rstrip('/')}/v1/accounts/{account_id}/reinforcementFineTuningJobs"
-    # Move optional jobId from body to query parameter if provided
-    job_id = body.get("jobId")
-    if isinstance(job_id, str):
-        job_id = job_id.strip()
-    if job_id:
-        # Remove from body and append as query param
-        body.pop("jobId", None)
-        url = f"{url}?{urlencode({'reinforcementFineTuningJobId': job_id})}"
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json",
-        "Accept": "application/json",
-        "User-Agent": get_user_agent(),
-    }
-    resp = requests.post(url, json=body, headers=headers, timeout=60)
-    if resp.status_code not in (200, 201):
-        raise RuntimeError(f"RFT job creation failed: {resp.status_code} {resp.text}")
-    return resp.json()
-
-
 def build_default_dataset_id(evaluator_id: str) -> str:
     ts = time.strftime("%Y%m%d%H%M%S")
     base = evaluator_id.lower().replace("_", "-")
@@ -217,22 +190,22 @@ def build_default_dataset_id(evaluator_id: str) -> str:
 def build_default_output_model(evaluator_id: str) -> str:
     base = evaluator_id.lower().replace("_", "-")
     uuid_suffix = str(uuid.uuid4())[:4]
-    
+
     # suffix is "-rft-{4chars}" -> 9 chars
     suffix_len = 9
     max_len = 63
-    
+
     # Check if we need to truncate
     if len(base) + suffix_len > max_len:
         # Calculate hash of the full base to preserve uniqueness
         hash_digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:6]
         # New structure: {truncated_base}-{hash}-{uuid_suffix}
         # Space needed for "-{hash}" is 1 + 6 = 7
         hash_part_len = 7
-        
+
         allowed_base_len = max_len - suffix_len - hash_part_len
         truncated_base = base[:allowed_base_len].strip("-")
-        
+
         return f"{truncated_base}-{hash_digest}-rft-{uuid_suffix}"
 
     return f"{base}-rft-{uuid_suffix}"
@@ -242,7 +215,6 @@ def build_default_output_model(evaluator_id: str) -> str:
     "detect_dataset_builder",
     "materialize_dataset_via_builder",
     "create_dataset_from_jsonl",
-    "create_reinforcement_fine_tuning_job",
     "build_default_dataset_id",
     "build_default_output_model",
     "_map_api_host_to_app_host",