1# /// script
2# dependencies = [
3#   "openai>=1.0",
4#   "requests",
5#   "azure-identity",
6#   "azure-ai-projects",
7# ]
8# ///
9"""
10submit_training.py — Submit SFT, DPO, or RFT training jobs on Azure AI Foundry.
11 
12Handles both SDK and REST API submission (REST fallback for OSS models).
13Supports /v1/ project endpoint (preferred) and Azure endpoint (fallback).
14 
15Usage:
16  python submit_training.py --base-url https://<resource>.services.ai.azure.com/api/projects/<project>/openai/v1/ \
17      --api-key KEY --training-file training.jsonl --validation-file validation.jsonl \
18      --model gpt-4.1-mini --type sft --epochs 2 --lr 1.0
19 
20  python submit_training.py --endpoint https://<resource>.openai.azure.com --api-key KEY \
21      --training-file-id file-abc123 --validation-file-id file-def456 \
22      --model gpt-oss-20b --type sft --epochs 2 --lr 0.5 --use-rest
23 
24  python submit_training.py --base-url <url> --api-key KEY \
25      --training-file-id file-abc123 --validation-file-id file-def456 \
26      --model o4-mini-2025-04-16 --type rft --grader-file grader.py
27"""
28 
29import json
30import os
31import sys
32 
33 
34try:
35    sys.stdout.reconfigure(encoding="utf-8")
36    sys.stderr.reconfigure(encoding="utf-8")
37except (AttributeError, OSError):
38    pass  # Stream not reconfigurable (older Python or non-tty); default encoding is fine
39sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
40from common import HelpOnErrorParser, get_clients, upload_file
41 
42import requests
43 
44 
45def submit_sft_sdk(client, model, train_id, val_id, epochs=2, lr=1.0, batch_size=None, suffix=None, training_type="globalStandard"):
46    """Submit SFT job using the Python SDK."""
47    hp = {"n_epochs": epochs, "learning_rate_multiplier": lr}
48    if batch_size:
49        hp["batch_size"] = batch_size
50 
51    kwargs = dict(
52        model=model,
53        training_file=train_id,
54        validation_file=val_id,
55        method={"type": "supervised"},
56        hyperparameters=hp,
57        # Azure-specific: passed via extra_body since the OpenAI SDK has no
58        # top-level trainingType kwarg.
59        extra_body={"trainingType": training_type},
60    )
61    if suffix:
62        kwargs["suffix"] = suffix
63 
64    job = client.fine_tuning.jobs.create(**kwargs)
65    return {"id": job.id, "status": job.status, "model": model, "method": "sdk"}
66 
67 
68def submit_sft_rest(endpoint, api_key, model, train_id, val_id, epochs=2, lr=1.0, batch_size=None, suffix=None, training_type="globalStandard"):
69    """Submit SFT job via REST API (fallback for models like gpt-oss-20b)."""
70    url = f"{endpoint}/openai/fine_tuning/jobs?api-version=2025-04-01-preview"
71    body = {
72        "model": model,
73        "training_file": train_id,
74        "validation_file": val_id,
75        "method": {"type": "supervised"},
76        "hyperparameters": {"n_epochs": epochs, "learning_rate_multiplier": lr},
77        "trainingType": training_type,
78    }
79    if batch_size:
80        body["hyperparameters"]["batch_size"] = batch_size
81    if suffix:
82        body["suffix"] = suffix
83 
84    resp = requests.post(url, headers={
85        "Content-Type": "application/json",
86        "api-key": api_key,
87    }, json=body, timeout=(10, 60))
88 
89    if resp.status_code in (200, 201):
90        try:
91            data = resp.json()
92        except ValueError:
93            raise RuntimeError(
94                f"REST submission returned {resp.status_code} but body was not JSON: {resp.text[:200]}"
95            )
96        if "id" not in data or "status" not in data:
97            raise RuntimeError(f"REST response missing 'id' or 'status' fields: {data}")
98        return {"id": data["id"], "status": data["status"], "model": model, "method": "rest"}
99    else:
100        try:
101            err_msg = resp.json().get('error', {}).get('message', 'Unknown error')
102        except (ValueError, KeyError):
103            err_msg = resp.text[:200] if resp.text else "Unknown error"
104        raise RuntimeError(
105            f"REST submission failed ({resp.status_code}): {err_msg}"
106        )
107 
108 
109def submit_rft(client, model, train_id, val_id, grader_source):
110    """Submit RFT job."""
111    job = client.fine_tuning.jobs.create(
112        model=model,
113        training_file=train_id,
114        validation_file=val_id,
115        method={
116            "type": "reinforcement",
117            "reinforcement": {
118                "grader": {
119                    "type": "python",
120                    "name": "custom_grader",
121                    "source": grader_source,
122                },
123            },
124        },
125    )
126    return {"id": job.id, "status": job.status, "model": model, "method": "sdk-rft"}
127 
128 
129def submit_dpo(client, model, train_id, val_id, epochs=2, lr=1.0, beta=0.1, suffix=None):
130    """Submit DPO job."""
131    job = client.fine_tuning.jobs.create(
132        model=model,
133        training_file=train_id,
134        validation_file=val_id,
135        suffix=suffix or None,
136        method={
137            "type": "dpo",
138            "dpo": {
139                "hyperparameters": {
140                    "n_epochs": epochs,
141                    "beta": beta,
142                    "learning_rate_multiplier": lr,
143                },
144            },
145        },
146    )
147    return {"id": job.id, "status": job.status, "model": model, "method": "sdk-dpo"}
148 
149 
150def main():
151    parser = HelpOnErrorParser(description="Submit fine-tuning jobs on Azure AI Foundry")
152    parser.add_argument("--base-url", default=os.environ.get("OPENAI_BASE_URL"),
153                        help="Project /v1/ URL (preferred)")
154    parser.add_argument("--endpoint", default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
155                        help="Azure OpenAI endpoint (fallback)")
156    parser.add_argument("--project-endpoint", default=os.environ.get("AZURE_AI_PROJECT_ENDPOINT"),
157                        help="Azure AI project endpoint (Foundry SDK)")
158    parser.add_argument("--api-key", default=os.environ.get("AZURE_OPENAI_API_KEY"),
159                        help="API key")
160    parser.add_argument("--model", required=True, help="Base model name (e.g., gpt-4.1-mini)")
161    parser.add_argument("--type", choices=["sft", "dpo", "rft"], default="sft",
162                        help="Training type: sft, dpo, or rft")
163 
164    # Data files — either paths (will upload) or IDs (already uploaded)
165    parser.add_argument("--training-file", help="Path to training JSONL file (will upload)")
166    parser.add_argument("--validation-file", help="Path to validation JSONL file (will upload)")
167    parser.add_argument("--training-file-id", help="Already-uploaded training file ID")
168    parser.add_argument("--validation-file-id", help="Already-uploaded validation file ID")
169 
170    # Hyperparameters
171    parser.add_argument("--epochs", type=int, default=2)
172    parser.add_argument("--lr", type=float, default=1.0, help="Learning rate multiplier")
173    parser.add_argument("--batch-size", type=int, default=None)
174    parser.add_argument("--suffix", help="Model suffix for identification")
175 
176    # DPO-specific
177    parser.add_argument("--beta", type=float, default=0.1, help="DPO beta (alignment strength)")
178 
179    # RFT-specific
180    parser.add_argument("--grader-file", help="Path to Python grader file (for RFT)")
181 
182    # REST fallback
183    parser.add_argument("--use-rest", action="store_true",
184                        help="Force REST API (needed for gpt-oss-20b and other OSS models)")
185    parser.add_argument("--training-type", choices=["globalStandard", "developerTier", "standard"],
186                        default="globalStandard",
187                        help="Azure training tier (default: globalStandard). developerTier is ~50%% off "
188                             "globalStandard with lower quotas. OSS models (gpt-oss-20b, Ministral, "
189                             "Llama, Qwen) only support globalStandard.")
190 
191    args = parser.parse_args()
192 
193    client, method = get_clients(
194        base_url=args.base_url, azure_endpoint=args.endpoint,
195        project_endpoint=args.project_endpoint, api_key=args.api_key
196    )
197 
198    # Resolve file IDs
199    train_id = args.training_file_id
200    val_id = args.validation_file_id
201    if args.training_file:
202        train_id = upload_file(client, args.training_file)
203    if args.validation_file:
204        val_id = upload_file(client, args.validation_file)
205 
206    if not train_id or not val_id:
207        print("Error: Provide training and validation file paths or IDs")
208        sys.exit(1)
209 
210    # Submit
211    if args.type == "rft":
212        if not args.grader_file:
213            print("Error: --grader-file required for RFT")
214            sys.exit(1)
215        with open(args.grader_file, encoding="utf-8") as f:
216            grader_source = f.read()
217        result = submit_rft(client, args.model, train_id, val_id, grader_source)
218    elif args.type == "dpo":
219        result = submit_dpo(client, args.model, train_id, val_id,
220                            args.epochs, args.lr, args.beta, args.suffix)
221    elif args.use_rest:
222        if not args.endpoint or not args.api_key:
223            print("Error: --use-rest requires --endpoint and --api-key (REST does not support DefaultAzureCredential)")
224            sys.exit(1)
225        result = submit_sft_rest(args.endpoint, args.api_key, args.model,
226                                 train_id, val_id, args.epochs, args.lr, args.batch_size, args.suffix,
227                                 args.training_type)
228    else:
229        # SFT via SDK with REST fallback for OSS models
230        try:
231            result = submit_sft_sdk(client, args.model, train_id, val_id,
232                                    args.epochs, args.lr, args.batch_size, args.suffix,
233                                    args.training_type)
234        except Exception as e:
235            err_str = str(e).lower()
236            # Match a wider set of "use REST instead" signals than the original
237            # exact-string comparison: Azure changes error text periodically.
238            if ("trainingtype" in err_str
239                    or "globalstandard" in err_str
240                    or "global_standard" in err_str
241                    or "does not support fine-tuning" in err_str):
242                if not args.endpoint or not args.api_key:
243                    print(f"SDK failed for {args.model}. REST fallback requires --endpoint and --api-key.")
244                    sys.exit(1)
245                print(f"SDK failed for {args.model}, falling back to REST API...")
246                result = submit_sft_rest(args.endpoint, args.api_key, args.model,
247                                         train_id, val_id, args.epochs, args.lr, args.batch_size, args.suffix,
248                                         args.training_type)
249            else:
250                raise
251 
252    print(f"\nJob submitted successfully:")
253    print(json.dumps(result, indent=2))
254 
255    # Save job info
256    outfile = f"ft_job_{result['id']}.json"
257    with open(outfile, "w", encoding="utf-8") as f:
258        json.dump({**result, "epochs": args.epochs, "lr": args.lr,
259                    "batch_size": args.batch_size, "train_file": train_id,
260                    "val_file": val_id}, f, indent=2)
261    print(f"Job info saved to {outfile}")
262 
263 
264if __name__ == "__main__":
265    main()
266