From 489caac9806e79a4419ac55f0a0c613fc9d206e5 Mon Sep 17 00:00:00 2001 From: AB Date: Thu, 23 Jul 2026 13:19:42 +0300 Subject: [PATCH 1/3] feat(examples/finetune) --- crates/alien-aws-clients/src/aws/bedrock.rs | 661 +++++++++++++++++ crates/alien-aws-clients/src/aws/mod.rs | 1 + crates/alien-aws-clients/src/lib.rs | 1 + crates/alien-azure-clients/src/azure/mod.rs | 1 + .../src/azure/openai_finetuning.rs | 434 +++++++++++ crates/alien-azure-clients/src/lib.rs | 3 + .../src/emitters/aws/ai.rs | 5 +- crates/alien-core/src/bin/schema_exporter.rs | 2 + crates/alien-core/src/bindings/ai.rs | 101 +++ crates/alien-core/src/resources/ai.rs | 242 +++++- crates/alien-gateway/src/config.rs | 12 +- crates/alien-gateway/src/lib.rs | 15 + crates/alien-gateway/src/router.rs | 110 ++- crates/alien-gateway/tests/integration.rs | 2 + crates/alien-gateway/tests/live_bedrock.rs | 2 + .../tests/live_foundry_claude.rs | 1 + .../alien-gateway/tests/live_vertex_claude.rs | 1 + .../alien-gcp-clients/src/gcp/aiplatform.rs | 406 +++++++++++ crates/alien-gcp-clients/src/gcp/mod.rs | 1 + crates/alien-gcp-clients/src/lib.rs | 1 + crates/alien-infra/src/ai/aws.rs | 548 +++++++++++++- crates/alien-infra/src/ai/aws_import.rs | 3 + crates/alien-infra/src/ai/azure.rs | 686 +++++++++++++++++- crates/alien-infra/src/ai/azure_import.rs | 3 + crates/alien-infra/src/ai/gcp.rs | 621 +++++++++++++++- crates/alien-infra/src/ai/gcp_import.rs | 4 + .../alien-infra/src/core/service_provider.rs | 51 ++ .../permission-sets/ai/finetune.jsonc | 90 +++ .../alien-permissions/tests/registry_tests.rs | 55 +- crates/alien-terraform/src/emitters/aws/ai.rs | 7 +- examples/README.md | 2 + examples/ai-finetune-inference-ts/README.md | 84 +++ examples/ai-finetune-inference-ts/alien.ts | 57 ++ .../ai-finetune-inference-ts/package.json | 22 + .../sample-training.jsonl | 10 + .../ai-finetune-inference-ts/src/index.ts | 66 ++ .../ai-finetune-inference-ts/template.toml | 3 + .../tests/ai-finetune-inference.test.ts | 68 ++ .../ai-finetune-inference-ts/tsconfig.json | 14 + .../ai-finetune-inference-ts/vitest.config.ts | 10 + examples/pnpm-lock.yaml | 25 + examples/pnpm-workspace.yaml | 1 + packages/core/src/__tests__/ai.test.ts | 42 ++ packages/core/src/ai.ts | 56 +- packages/core/src/generated/index.ts | 4 + packages/core/src/generated/schemas/ai.json | 2 +- .../src/generated/schemas/finetuneMethod.json | 1 + .../src/generated/schemas/finetuneSpec.json | 1 + packages/core/src/generated/zod/ai-schema.ts | 10 +- .../generated/zod/finetune-method-schema.ts | 13 + .../src/generated/zod/finetune-spec-schema.ts | 22 + packages/core/src/generated/zod/index.ts | 4 + 52 files changed, 4536 insertions(+), 51 deletions(-) create mode 100644 crates/alien-aws-clients/src/aws/bedrock.rs create mode 100644 crates/alien-azure-clients/src/azure/openai_finetuning.rs create mode 100644 crates/alien-gcp-clients/src/gcp/aiplatform.rs create mode 100644 crates/alien-permissions/permission-sets/ai/finetune.jsonc create mode 100644 examples/ai-finetune-inference-ts/README.md create mode 100644 examples/ai-finetune-inference-ts/alien.ts create mode 100644 examples/ai-finetune-inference-ts/package.json create mode 100644 examples/ai-finetune-inference-ts/sample-training.jsonl create mode 100644 examples/ai-finetune-inference-ts/src/index.ts create mode 100644 examples/ai-finetune-inference-ts/template.toml create mode 100644 examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts create mode 100644 examples/ai-finetune-inference-ts/tsconfig.json create mode 100644 examples/ai-finetune-inference-ts/vitest.config.ts create mode 100644 packages/core/src/generated/schemas/finetuneMethod.json create mode 100644 packages/core/src/generated/schemas/finetuneSpec.json create mode 100644 packages/core/src/generated/zod/finetune-method-schema.ts create mode 100644 packages/core/src/generated/zod/finetune-spec-schema.ts diff --git a/crates/alien-aws-clients/src/aws/bedrock.rs b/crates/alien-aws-clients/src/aws/bedrock.rs new file mode 100644 index 000000000..326f7dc49 --- /dev/null +++ b/crates/alien-aws-clients/src/aws/bedrock.rs @@ -0,0 +1,661 @@ +//! AWS Bedrock control-plane client (model customization / fine-tuning). +//! +//! Hand-rolled reqwest + SigV4 client for the two model-customization-job REST +//! operations the fine-tuning controller needs. There is deliberately no +//! `aws-sdk-bedrock` dependency; this mirrors the other reqwest-based clients in +//! this crate (see `s3.rs` / `apigatewayv2.rs`) and keeps the dependency surface +//! small. +//! +//! Bedrock control plane host: `https://bedrock.{region}.amazonaws.com` +//! (SigV4 service signing name `bedrock`). +//! - CreateModelCustomizationJob: `POST /model-customization-jobs` +//! - GetModelCustomizationJob: `GET /model-customization-jobs/{jobIdentifier}` +//! +//! Request/response JSON field names follow the AWS Bedrock API reference +//! (CreateModelCustomizationJob / GetModelCustomizationJob, service version +//! 2023-04-20). Only the fields the controller uses are modeled; everything else +//! is ignored on deserialize. + +use crate::aws::aws_request_utils::{AwsRequestBuilderExt, AwsSignConfig}; +use crate::aws::credential_provider::AwsCredentialProvider; +use alien_client_core::{ErrorData, Result}; +use alien_error::{Context, ContextError, IntoAlienError}; +use async_trait::async_trait; +use bon::Builder; +use form_urlencoded; +use reqwest::{Client, Method, StatusCode}; +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +#[cfg(feature = "test-utils")] +use mockall::automock; + +// --------------------------------------------------------------------------- +// Bedrock API trait +// --------------------------------------------------------------------------- + +/// Minimal Bedrock control-plane surface for model customization (fine-tuning). +#[cfg_attr(feature = "test-utils", automock)] +#[cfg_attr(target_arch = "wasm32", async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait)] +pub trait BedrockApi: Send + Sync + std::fmt::Debug { + /// Submit a model customization (fine-tuning) job. Returns the created job's ARN, + /// which is a valid `jobIdentifier` for [`get_model_customization_job`]. + async fn create_model_customization_job( + &self, + request: &CreateModelCustomizationJobRequest, + ) -> Result; + + /// Fetch a model customization job's current status and (once complete) the + /// resulting custom-model ARN. `job_identifier` is the job ARN or job name. + async fn get_model_customization_job( + &self, + job_identifier: &str, + ) -> Result; +} + +// --------------------------------------------------------------------------- +// Client +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub struct BedrockClient { + client: Client, + credentials: AwsCredentialProvider, +} + +impl BedrockClient { + pub fn new(client: Client, credentials: AwsCredentialProvider) -> Self { + Self { + client, + credentials, + } + } + + fn sign_config(&self) -> AwsSignConfig { + AwsSignConfig { + service_name: "bedrock".into(), + region: self.credentials.region().to_string(), + credentials: self.credentials.get_credentials(), + signing_region: None, + } + } + + /// Host header value (always the real Bedrock host so the signature matches + /// what AWS expects, even when the base URL is overridden for tests). + fn host_header(&self) -> String { + format!("bedrock.{}.amazonaws.com", self.credentials.region()) + } + + fn get_base_url(&self) -> String { + if let Some(override_url) = self.credentials.get_service_endpoint_option("bedrock") { + override_url.trim_end_matches('/').to_string() + } else { + format!("https://bedrock.{}.amazonaws.com", self.credentials.region()) + } + } + + async fn send_json( + &self, + method: Method, + path: &str, + body: Option, + operation: &str, + resource: &str, + ) -> Result { + self.credentials.ensure_fresh().await?; + let url = format!("{}{}", self.get_base_url(), path); + + let mut builder = self + .client + .request(method, &url) + .host(&self.host_header()) + .content_type_json(); + + if let Some(body) = body { + builder = builder.content_sha256(&body).body(body.clone()); + let result = + crate::aws::aws_request_utils::sign_send_json(builder, &self.sign_config()).await; + return Self::map_result(result, operation, resource, Some(&body)); + } + + builder = builder.content_sha256(""); + let result = + crate::aws::aws_request_utils::sign_send_json(builder, &self.sign_config()).await; + Self::map_result(result, operation, resource, None) + } + + fn map_result( + result: Result, + operation: &str, + resource: &str, + request_body: Option<&str>, + ) -> Result { + match result { + Ok(v) => Ok(v), + Err(e) => { + if let Some(ErrorData::HttpResponseError { + http_status, + http_response_text: Some(ref text), + .. + }) = &e.error + { + let status = StatusCode::from_u16(*http_status) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + if let Some(mapped) = + Self::map_bedrock_error(status, text, operation, resource, request_body) + { + Err(e.context(mapped)) + } else { + Err(e) + } + } else { + Err(e) + } + } + } + } + + /// Map a Bedrock error JSON body (`{"message": "...", "__type": "..."}`) to a + /// structured client error. Bedrock control-plane errors use the standard AWS + /// JSON error shape. + fn map_bedrock_error( + status: StatusCode, + body: &str, + _operation: &str, + resource: &str, + request_body: Option<&str>, + ) -> Option { + let parsed: std::result::Result = serde_json::from_str(body); + let (code, message) = match parsed { + Ok(e) => { + let code = e + .type_field + .or(e.code) + .unwrap_or_else(|| "UnknownError".into()); + // `__type` is often `com.amazonaws...#ValidationException`; keep the tail. + let code = code.rsplit(['#', '.']).next().unwrap_or(&code).to_string(); + let message = e.message.unwrap_or_else(|| "Unknown error".into()); + (code, message) + } + Err(_) => return None, + }; + + Some(match code.as_str() { + "AccessDeniedException" => ErrorData::RemoteAccessDenied { + resource_type: "BedrockModelCustomizationJob".into(), + resource_name: resource.into(), + }, + "ResourceNotFoundException" => ErrorData::RemoteResourceNotFound { + resource_type: "BedrockModelCustomizationJob".into(), + resource_name: resource.into(), + }, + "ConflictException" => ErrorData::RemoteResourceConflict { + message, + resource_type: "BedrockModelCustomizationJob".into(), + resource_name: resource.into(), + }, + "ThrottlingException" => ErrorData::RateLimitExceeded { message }, + "ServiceQuotaExceededException" => ErrorData::QuotaExceeded { message }, + "ValidationException" => ErrorData::InvalidInput { + message, + field_name: None, + }, + _ => match status { + StatusCode::NOT_FOUND => ErrorData::RemoteResourceNotFound { + resource_type: "BedrockModelCustomizationJob".into(), + resource_name: resource.into(), + }, + StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => ErrorData::RemoteAccessDenied { + resource_type: "BedrockModelCustomizationJob".into(), + resource_name: resource.into(), + }, + StatusCode::TOO_MANY_REQUESTS => ErrorData::RateLimitExceeded { message }, + StatusCode::SERVICE_UNAVAILABLE + | StatusCode::BAD_GATEWAY + | StatusCode::GATEWAY_TIMEOUT => ErrorData::RemoteServiceUnavailable { message }, + _ => ErrorData::HttpResponseError { + message: format!("Bedrock operation failed: {}", message), + url: "bedrock.amazonaws.com".to_string(), + http_status: status.as_u16(), + http_response_text: Some(body.into()), + http_request_text: request_body.map(|s| s.to_string()), + }, + }, + }) + } +} + +#[cfg_attr(target_arch = "wasm32", async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait)] +impl BedrockApi for BedrockClient { + async fn create_model_customization_job( + &self, + request: &CreateModelCustomizationJobRequest, + ) -> Result { + let body = serde_json::to_string(request).into_alien_error().context( + ErrorData::SerializationError { + message: "Failed to serialize CreateModelCustomizationJobRequest".to_string(), + }, + )?; + self.send_json( + Method::POST, + "/model-customization-jobs", + Some(body), + "CreateModelCustomizationJob", + &request.job_name, + ) + .await + } + + async fn get_model_customization_job( + &self, + job_identifier: &str, + ) -> Result { + // The identifier is a job ARN or name; it is a path segment and must be + // percent-encoded (ARNs contain ':' and '/'). + let encoded = form_urlencoded::byte_serialize(job_identifier.as_bytes()).collect::(); + let path = format!("/model-customization-jobs/{}", encoded); + self.send_json( + Method::GET, + &path, + None, + "GetModelCustomizationJob", + job_identifier, + ) + .await + } +} + +// --------------------------------------------------------------------------- +// Error struct +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct BedrockErrorResponse { + pub message: Option, + pub code: Option, + #[serde(rename = "__type")] + pub type_field: Option, +} + +// --------------------------------------------------------------------------- +// Request / response types (subset of the Bedrock API) +// --------------------------------------------------------------------------- + +/// S3 location for a job's input or output data. +/// Matches Bedrock's `TrainingDataConfig` / `OutputDataConfig` (`s3Uri`). +#[derive(Debug, Clone, Serialize, Deserialize, Builder)] +#[serde(rename_all = "camelCase")] +pub struct S3DataConfig { + /// Fully-qualified S3 URI, e.g. `s3://bucket/key`. + pub s3_uri: String, +} + +/// Request body for CreateModelCustomizationJob. +/// Field names and casing follow the AWS Bedrock API reference; only the fields +/// the fine-tuning controller sets are modeled, and optional/None fields are +/// omitted so the wire shape stays minimal. +#[derive(Debug, Clone, Serialize, Builder)] +#[serde(rename_all = "camelCase")] +pub struct CreateModelCustomizationJobRequest { + /// A name for the fine-tuning job. + pub job_name: String, + /// A name for the resulting custom model. + pub custom_model_name: String, + /// ARN of the IAM role Bedrock assumes to read training data / write output. + pub role_arn: String, + /// Base foundation-model identifier to tune. + pub base_model_identifier: String, + /// The customization type, e.g. `FINE_TUNING`. + pub customization_type: String, + /// Where the job reads its training dataset from. + pub training_data_config: S3DataConfig, + /// Where the job writes output/metrics. + pub output_data_config: S3DataConfig, + /// Optional tuning hyperparameters (string→string map). + #[serde(skip_serializing_if = "Option::is_none")] + pub hyper_parameters: Option>, +} + +/// Response body for CreateModelCustomizationJob. +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateModelCustomizationJobResponse { + /// ARN of the fine-tuning job (usable as `jobIdentifier` for polling). + pub job_arn: String, +} + +/// Terminal/non-terminal status of a model customization job, mirroring the +/// Bedrock `status` enum (`InProgress | Completed | Failed | Stopping | Stopped`). +/// An unrecognized value maps to [`Unknown`](ModelCustomizationJobStatus::Unknown) +/// so a future/unexpected status is surfaced rather than silently treated as terminal. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ModelCustomizationJobStatus { + InProgress, + Completed, + Failed, + Stopping, + Stopped, + /// Any status string Bedrock returns that we don't recognize. + Unknown(String), +} + +impl ModelCustomizationJobStatus { + fn from_wire(raw: &str) -> Self { + match raw { + "InProgress" => Self::InProgress, + "Completed" => Self::Completed, + "Failed" => Self::Failed, + "Stopping" => Self::Stopping, + "Stopped" => Self::Stopped, + other => Self::Unknown(other.to_string()), + } + } + + /// True for a status the job will not move out of (success or failure). + pub fn is_terminal(&self) -> bool { + matches!(self, Self::Completed | Self::Failed | Self::Stopped) + } +} + +/// Response body for GetModelCustomizationJob (subset). +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct GetModelCustomizationJobResponse { + /// Raw job status string from Bedrock. Prefer [`status`](Self::status) for the + /// typed enum. + pub status: String, + /// ARN of the resulting custom model, present once the job completes. This is + /// the artifact the OpenAI chat endpoint accepts for a custom model. + #[serde(default)] + pub output_model_arn: Option, + /// Name of the resulting custom model, present once the job completes. + #[serde(default)] + pub output_model_name: Option, + /// Human-readable reason the job failed, present when `status == Failed`. + #[serde(default)] + pub failure_message: Option, + /// ARN of the customization job itself. + #[serde(default)] + pub job_arn: Option, +} + +impl GetModelCustomizationJobResponse { + /// The job's status as a typed enum. + pub fn status(&self) -> ModelCustomizationJobStatus { + ModelCustomizationJobStatus::from_wire(&self.status) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use alien_core::{AwsClientConfig, AwsCredentials, AwsServiceOverrides}; + use std::{ + collections::HashMap, + io::{Read, Write}, + net::TcpListener, + sync::{Arc, Mutex}, + }; + + #[test] + fn create_request_serializes_expected_json_shape() { + let request = CreateModelCustomizationJobRequest::builder() + .job_name("my-ai-job".to_string()) + .custom_model_name("my-ai-model".to_string()) + .role_arn("arn:aws:iam::123456789012:role/test-my-ai".to_string()) + .base_model_identifier("amazon.nova-lite-v1:0".to_string()) + .customization_type("FINE_TUNING".to_string()) + .training_data_config( + S3DataConfig::builder() + .s3_uri("s3://test-bucket/training.jsonl".to_string()) + .build(), + ) + .output_data_config( + S3DataConfig::builder() + .s3_uri("s3://test-bucket/output/my-ai/".to_string()) + .build(), + ) + .build(); + + let value = serde_json::to_value(&request).unwrap(); + assert_eq!(value["jobName"], "my-ai-job"); + assert_eq!(value["customModelName"], "my-ai-model"); + assert_eq!(value["roleArn"], "arn:aws:iam::123456789012:role/test-my-ai"); + assert_eq!(value["baseModelIdentifier"], "amazon.nova-lite-v1:0"); + assert_eq!(value["customizationType"], "FINE_TUNING"); + assert_eq!( + value["trainingDataConfig"]["s3Uri"], + "s3://test-bucket/training.jsonl" + ); + assert_eq!( + value["outputDataConfig"]["s3Uri"], + "s3://test-bucket/output/my-ai/" + ); + // Optional hyperParameters omitted when not set. + assert!(value.get("hyperParameters").is_none()); + } + + #[test] + fn get_response_maps_status_and_output_arn() { + let body = r#"{ + "status": "Completed", + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345", + "outputModelArn": "arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345", + "outputModelName": "my-ai-model" + }"#; + let parsed: GetModelCustomizationJobResponse = serde_json::from_str(body).unwrap(); + assert_eq!(parsed.status(), ModelCustomizationJobStatus::Completed); + assert!(parsed.status().is_terminal()); + assert_eq!( + parsed.output_model_arn.as_deref(), + Some("arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345") + ); + + let in_progress: GetModelCustomizationJobResponse = + serde_json::from_str(r#"{"status":"InProgress"}"#).unwrap(); + assert_eq!( + in_progress.status(), + ModelCustomizationJobStatus::InProgress + ); + assert!(!in_progress.status().is_terminal()); + assert!(in_progress.output_model_arn.is_none()); + } + + /// Read one full HTTP request (headers + body) off a socket, using the + /// Content-Length header to know when the body is complete. + fn read_http_request(stream: &mut std::net::TcpStream) -> (String, String) { + let mut bytes = Vec::new(); + let mut buffer = [0_u8; 4096]; + loop { + let count = stream.read(&mut buffer).expect("read request"); + if count == 0 { + break; + } + bytes.extend_from_slice(&buffer[..count]); + if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") { + let header_end = end + 4; + let headers = String::from_utf8_lossy(&bytes[..header_end]).to_string(); + let length: usize = headers + .lines() + .find_map(|line| { + line.to_ascii_lowercase() + .strip_prefix("content-length: ") + .and_then(|v| v.trim().parse().ok()) + }) + .unwrap_or(0); + if bytes.len() >= header_end + length { + let body = + String::from_utf8_lossy(&bytes[header_end..header_end + length]).to_string(); + return (headers, body); + } + } + } + (String::from_utf8_lossy(&bytes).to_string(), String::new()) + } + + fn test_config(endpoint: String) -> AwsClientConfig { + AwsClientConfig { + account_id: "123456789012".to_string(), + region: "us-east-1".to_string(), + credentials: AwsCredentials::AccessKeys { + access_key_id: "test-access".to_string(), + secret_access_key: "test-secret".to_string(), + session_token: None, + }, + service_overrides: Some(AwsServiceOverrides { + endpoints: HashMap::from([("bedrock".to_string(), endpoint)]), + }), + } + } + + #[tokio::test] + async fn create_model_customization_job_signs_and_sends_expected_request() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind test server"); + let endpoint = format!("http://{}", listener.local_addr().expect("local address")); + let captured: Arc> = + Arc::new(Mutex::new((String::new(), String::new()))); + let sink = captured.clone(); + + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept request"); + let (headers, body) = read_http_request(&mut stream); + *sink.lock().expect("lock") = (headers, body); + let resp_body = r#"{"jobArn":"arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345"}"#; + write!( + stream, + "HTTP/1.1 201 Created\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + resp_body.len(), + resp_body + ) + .expect("write response"); + }); + + let client = + BedrockClient::new(Client::new(), AwsCredentialProvider::from_config_sync(test_config(endpoint))); + let request = CreateModelCustomizationJobRequest::builder() + .job_name("my-ai-job".to_string()) + .custom_model_name("my-ai-model".to_string()) + .role_arn("arn:aws:iam::123456789012:role/test-my-ai".to_string()) + .base_model_identifier("amazon.nova-lite-v1:0".to_string()) + .customization_type("FINE_TUNING".to_string()) + .training_data_config( + S3DataConfig::builder() + .s3_uri("s3://test-bucket/training.jsonl".to_string()) + .build(), + ) + .output_data_config( + S3DataConfig::builder() + .s3_uri("s3://test-bucket/output/my-ai/".to_string()) + .build(), + ) + .build(); + + let response = client + .create_model_customization_job(&request) + .await + .expect("request should succeed"); + server.join().expect("server should finish"); + + assert_eq!( + response.job_arn, + "arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345" + ); + + let (headers, body) = captured.lock().expect("lock").clone(); + let request_line = headers.lines().next().unwrap_or_default(); + // Method + path. + assert!( + request_line.starts_with("POST /model-customization-jobs "), + "unexpected request line: {request_line}" + ); + // SigV4 signing headers must be present and name the bedrock service. + let lower = headers.to_ascii_lowercase(); + assert!( + lower.contains("authorization: aws4-hmac-sha256"), + "missing SigV4 Authorization header:\n{headers}" + ); + assert!( + lower.contains("credential=test-access/") + && lower.contains("/us-east-1/bedrock/aws4_request"), + "Authorization must scope the signature to us-east-1/bedrock:\n{headers}" + ); + assert!(lower.contains("x-amz-date:"), "missing x-amz-date header"); + assert!( + lower.contains("x-amz-content-sha256:"), + "missing x-amz-content-sha256 header" + ); + assert!( + lower.contains("host: bedrock.us-east-1.amazonaws.com"), + "host header must be the real bedrock host for a valid signature:\n{headers}" + ); + + // Body must carry the exact JSON the API expects. + let json: serde_json::Value = serde_json::from_str(&body).expect("body is json"); + assert_eq!(json["jobName"], "my-ai-job"); + assert_eq!(json["customModelName"], "my-ai-model"); + assert_eq!(json["baseModelIdentifier"], "amazon.nova-lite-v1:0"); + assert_eq!(json["customizationType"], "FINE_TUNING"); + assert_eq!( + json["trainingDataConfig"]["s3Uri"], + "s3://test-bucket/training.jsonl" + ); + assert_eq!(json["roleArn"], "arn:aws:iam::123456789012:role/test-my-ai"); + } + + #[tokio::test] + async fn get_model_customization_job_uses_get_and_encoded_path() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind test server"); + let endpoint = format!("http://{}", listener.local_addr().expect("local address")); + let captured: Arc> = Arc::new(Mutex::new(String::new())); + let sink = captured.clone(); + + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept request"); + let (headers, _body) = read_http_request(&mut stream); + *sink.lock().expect("lock") = headers; + let resp_body = r#"{"status":"Completed","outputModelArn":"arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345"}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + resp_body.len(), + resp_body + ) + .expect("write response"); + }); + + let client = + BedrockClient::new(Client::new(), AwsCredentialProvider::from_config_sync(test_config(endpoint))); + let job_arn = + "arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345"; + let response = client + .get_model_customization_job(job_arn) + .await + .expect("request should succeed"); + server.join().expect("server should finish"); + + assert_eq!(response.status(), ModelCustomizationJobStatus::Completed); + assert_eq!( + response.output_model_arn.as_deref(), + Some("arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345") + ); + + let headers = captured.lock().expect("lock").clone(); + let request_line = headers.lines().next().unwrap_or_default(); + assert!( + request_line.starts_with("GET /model-customization-jobs/"), + "unexpected request line: {request_line}" + ); + // ARN separators must be percent-encoded in the path segment. + assert!( + request_line.contains("%3A") && request_line.contains("%2F"), + "job identifier must be percent-encoded in the path: {request_line}" + ); + } +} diff --git a/crates/alien-aws-clients/src/aws/mod.rs b/crates/alien-aws-clients/src/aws/mod.rs index 48ce254aa..423081cfd 100644 --- a/crates/alien-aws-clients/src/aws/mod.rs +++ b/crates/alien-aws-clients/src/aws/mod.rs @@ -48,6 +48,7 @@ pub mod acm; pub mod apigatewayv2; pub mod autoscaling; pub mod aws_request_utils; +pub mod bedrock; pub mod cloudformation; pub mod cloudwatch; pub mod codebuild; diff --git a/crates/alien-aws-clients/src/lib.rs b/crates/alien-aws-clients/src/lib.rs index 953b21287..946c3449d 100644 --- a/crates/alien-aws-clients/src/lib.rs +++ b/crates/alien-aws-clients/src/lib.rs @@ -8,6 +8,7 @@ pub use aws::{AwsClientConfig, AwsClientConfigExt, AwsImpersonationConfig}; // Re-export all client APIs pub use aws::acm::{AcmApi, AcmClient}; pub use aws::apigatewayv2::{ApiGatewayV2Api, ApiGatewayV2Client}; +pub use aws::bedrock::{BedrockApi, BedrockClient}; pub use aws::cloudformation::{CloudFormationApi, CloudFormationClient}; pub use aws::cloudwatch::{CloudWatchApi, CloudWatchClient}; pub use aws::codebuild::{CodeBuildApi, CodeBuildClient}; diff --git a/crates/alien-azure-clients/src/azure/mod.rs b/crates/alien-azure-clients/src/azure/mod.rs index b252da790..8116be840 100644 --- a/crates/alien-azure-clients/src/azure/mod.rs +++ b/crates/alien-azure-clients/src/azure/mod.rs @@ -28,6 +28,7 @@ pub mod managed_identity; pub mod models; pub mod monitor; pub mod network; +pub mod openai_finetuning; pub mod private_networking; pub mod resource_graph; pub mod resource_skus; diff --git a/crates/alien-azure-clients/src/azure/openai_finetuning.rs b/crates/alien-azure-clients/src/azure/openai_finetuning.rs new file mode 100644 index 000000000..7bff1e2a2 --- /dev/null +++ b/crates/alien-azure-clients/src/azure/openai_finetuning.rs @@ -0,0 +1,434 @@ +//! Azure AI Foundry (Azure OpenAI) fine-tuning — a **data-plane** client. +//! +//! Unlike [`cognitive_services`](crate::azure::cognitive_services), which drives +//! the ARM/control-plane to provision the account and its deployments, this +//! client talks to the *account's own OpenAI data-plane endpoint* +//! (`https://{account}.openai.azure.com` / `.cognitiveservices.azure.com`). +//! +//! Two differences from the control-plane client: +//! - The base URL is the **account endpoint**, supplied per call by the caller +//! (the controller already learns it when the account is provisioned). There +//! is no subscription/resource-group path. +//! - The bearer token is minted for the **cognitive-services data-plane scope** +//! `https://cognitiveservices.azure.com/.default`, not +//! `https://management.azure.com/.default`. Data-plane RBAC (e.g. the +//! `Cognitive Services OpenAI Contributor` role the controller applies) is +//! distinct from ARM Contributor. +//! +//! REST surface (Azure OpenAI fine-tuning, API version `2024-10-21`): +//! - `POST {endpoint}/openai/fine_tuning/jobs?api-version=...` +//! - `GET {endpoint}/openai/fine_tuning/jobs/{job_id}?api-version=...` + +use crate::azure::common::{AzureClientBase, AzureRequestBuilder}; +use crate::azure::token_cache::AzureTokenCache; +use alien_client_core::{ErrorData, Result}; +use alien_error::{Context, IntoAlienError}; +use reqwest::{Client, Method}; +use serde::{Deserialize, Serialize}; + +#[cfg(feature = "test-utils")] +use mockall::automock; + +/// Fine-tuning REST API version. Azure pins the whole `fine_tuning.jobs` +/// surface to a dated version supplied as an `api-version` query parameter. +pub const OPENAI_FINE_TUNING_API_VERSION: &str = "2024-10-21"; + +/// OAuth scope for the Azure Cognitive Services **data plane**. +/// +/// Distinct from the ARM management scope (`https://management.azure.com/.default`) +/// used by the control-plane client: fine-tuning is a data action authorized by +/// the data-plane RBAC role assigned on the account. +pub const COGNITIVE_SERVICES_DATA_PLANE_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; + +// ------------------------------------------------------------------------- +// Data-plane models +// ------------------------------------------------------------------------- + +/// Request body for `POST /openai/fine_tuning/jobs`. +/// +/// `training_file` is the identifier Azure OpenAI accepts for the dataset. For +/// a Foundry import from customer Blob storage this is the blob reference; see +/// [`FoundryFineTuningApi::create_fine_tuning_job`] for the network gotcha. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub struct FineTuningJobCreateRequest { + /// Provider-native base model to fine-tune (e.g. `gpt-4o-mini`). + pub model: String, + /// The training dataset reference the job reads from. + pub training_file: String, +} + +/// A fine-tuning job as returned by `POST` (creation) and `GET` (poll). +/// +/// Field names are Azure OpenAI's native `snake_case`, so this type does **not** +/// use the workspace's usual `camelCase` rename — matching the wire shape is the +/// whole point. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FineTuningJob { + /// The job id (e.g. `ftjob-abc123`). Stored by the controller to poll. + pub id: String, + /// Lifecycle status: `pending`, `validating_files`, `queued`, `running`, + /// `succeeded`, `failed`, or `cancelled`. + pub status: String, + /// Populated only once `status == "succeeded"`: the tuned model name the + /// OpenAI chat endpoint accepts as a deployment target. + #[serde(default)] + pub fine_tuned_model: Option, +} + +impl FineTuningJob { + /// Whether the job reached the successful terminal state. + pub fn is_succeeded(&self) -> bool { + self.status.eq_ignore_ascii_case("succeeded") + } + + /// Whether the job reached a terminal *failure* state (`failed`/`cancelled`). + /// The controller fails fast on these rather than polling forever. + pub fn is_terminal_failure(&self) -> bool { + self.status.eq_ignore_ascii_case("failed") + || self.status.eq_ignore_ascii_case("cancelled") + || self.status.eq_ignore_ascii_case("canceled") + } +} + +// ------------------------------------------------------------------------- +// Foundry fine-tuning API trait +// ------------------------------------------------------------------------- + +/// Data-plane fine-tuning operations against an Azure AI Foundry account. +/// +/// The `endpoint` argument is the account's OpenAI data-plane base URL (e.g. +/// `https://my-account.openai.azure.com`), learned by the controller when the +/// account is provisioned. +#[cfg_attr(feature = "test-utils", automock)] +#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] +pub trait FoundryFineTuningApi: Send + Sync + std::fmt::Debug { + /// Submit a fine-tuning job (`POST /openai/fine_tuning/jobs`). + /// + /// `training_file` is the dataset reference the job reads. Foundry supports + /// importing directly from customer Blob storage, but **that import path + /// requires the account to allow public network access** — a private-endpoint + /// account will reject the blob URL at submit time. This client passes the + /// caller's reference through unchanged; enforcing network posture is the + /// controller/account concern. + async fn create_fine_tuning_job( + &self, + endpoint: &str, + model: &str, + training_file: &str, + ) -> Result; + + /// Poll a fine-tuning job's status (`GET /openai/fine_tuning/jobs/{job_id}`). + /// On success the returned job carries `fine_tuned_model`. + async fn get_fine_tuning_job(&self, endpoint: &str, job_id: &str) -> Result; +} + +// ------------------------------------------------------------------------- +// Foundry fine-tuning client struct +// ------------------------------------------------------------------------- + +#[derive(Debug)] +pub struct AzureFoundryFineTuningClient { + pub base: AzureClientBase, + pub token_cache: AzureTokenCache, +} + +impl AzureFoundryFineTuningClient { + pub fn new(client: Client, token_cache: AzureTokenCache) -> Self { + // The base endpoint is supplied per call (the account endpoint), so the + // base's own endpoint field is unused for URL building here. Seed it with + // the management endpoint for parity with the other clients. + let endpoint = token_cache.management_endpoint().to_string(); + Self { + base: AzureClientBase::with_client_config( + client, + endpoint, + token_cache.config().clone(), + ), + token_cache, + } + } + + /// Builds `{endpoint}/openai/fine_tuning/jobs[/{job_id}]?api-version=...`. + fn jobs_url(&self, endpoint: &str, job_id: Option<&str>) -> String { + let base = endpoint.trim_end_matches('/'); + let path = match job_id { + Some(id) => format!("{base}/openai/fine_tuning/jobs/{id}"), + None => format!("{base}/openai/fine_tuning/jobs"), + }; + format!("{path}?api-version={OPENAI_FINE_TUNING_API_VERSION}") + } + + /// Reads and JSON-parses a `FineTuningJob` from a successful response. + async fn parse_job(resp: reqwest::Response, url: &str, op: &str) -> Result { + let status = resp.status().as_u16(); + let body = resp + .text() + .await + .into_alien_error() + .context(ErrorData::HttpResponseError { + message: format!("Azure {op}: failed to read fine-tuning job response body"), + url: url.to_string(), + http_status: status, + http_request_text: None, + http_response_text: None, + })?; + + serde_json::from_str(&body) + .into_alien_error() + .context(ErrorData::HttpResponseError { + message: format!("Azure {op}: JSON parse error for fine-tuning job"), + url: url.to_string(), + http_status: status, + http_request_text: None, + http_response_text: Some(body), + }) + } +} + +#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] +impl FoundryFineTuningApi for AzureFoundryFineTuningClient { + async fn create_fine_tuning_job( + &self, + endpoint: &str, + model: &str, + training_file: &str, + ) -> Result { + let bearer_token = self + .token_cache + .get_bearer_token_with_scope(COGNITIVE_SERVICES_DATA_PLANE_SCOPE) + .await?; + + let url = self.jobs_url(endpoint, None); + + let request = FineTuningJobCreateRequest { + model: model.to_string(), + training_file: training_file.to_string(), + }; + let body = serde_json::to_string(&request) + .into_alien_error() + .context(ErrorData::SerializationError { + message: format!( + "Failed to serialize fine-tuning job create request for model '{model}'" + ), + })?; + + let builder = AzureRequestBuilder::new(Method::POST, url.clone()) + .content_type_json() + .content_length(&body) + .body(body); + + let req = builder.build()?; + let signed = self.base.sign_request(req, &bearer_token).await?; + let resp = self + .base + .execute_request(signed, "CreateFineTuningJob", model) + .await?; + + Self::parse_job(resp, &url, "CreateFineTuningJob").await + } + + async fn get_fine_tuning_job(&self, endpoint: &str, job_id: &str) -> Result { + let bearer_token = self + .token_cache + .get_bearer_token_with_scope(COGNITIVE_SERVICES_DATA_PLANE_SCOPE) + .await?; + + let url = self.jobs_url(endpoint, Some(job_id)); + + let builder = AzureRequestBuilder::new(Method::GET, url.clone()).content_length(""); + let req = builder.build()?; + let signed = self.base.sign_request(req, &bearer_token).await?; + let resp = self + .base + .execute_request(signed, "GetFineTuningJob", job_id) + .await?; + + Self::parse_job(resp, &url, "GetFineTuningJob").await + } +} + +// ------------------------------------------------------------------------- +// Tests +// ------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn uses_pinned_api_version_and_data_plane_scope() { + assert_eq!(OPENAI_FINE_TUNING_API_VERSION, "2024-10-21"); + // The data-plane scope must NOT be the ARM management scope — a token for + // the wrong audience is rejected by Azure. + assert_eq!( + COGNITIVE_SERVICES_DATA_PLANE_SCOPE, + "https://cognitiveservices.azure.com/.default" + ); + assert_ne!( + COGNITIVE_SERVICES_DATA_PLANE_SCOPE, + "https://management.azure.com/.default" + ); + } + + #[test] + fn create_request_serializes_snake_case_wire_shape() { + let req = FineTuningJobCreateRequest { + model: "gpt-4o-mini".to_string(), + training_file: "https://acct.blob.core.windows.net/data/training.jsonl".to_string(), + }; + let value: serde_json::Value = + serde_json::from_str(&serde_json::to_string(&req).unwrap()).unwrap(); + assert_eq!(value["model"], "gpt-4o-mini"); + assert_eq!( + value["training_file"], + "https://acct.blob.core.windows.net/data/training.jsonl" + ); + } + + #[test] + fn job_status_helpers_classify_terminal_states() { + let succeeded = FineTuningJob { + id: "ftjob-1".to_string(), + status: "succeeded".to_string(), + fine_tuned_model: Some("gpt-4o-mini.ft-1".to_string()), + }; + assert!(succeeded.is_succeeded()); + assert!(!succeeded.is_terminal_failure()); + + for failed in ["failed", "cancelled", "canceled", "FAILED"] { + let job = FineTuningJob { + id: "ftjob-2".to_string(), + status: failed.to_string(), + fine_tuned_model: None, + }; + assert!(job.is_terminal_failure(), "'{failed}' must be terminal failure"); + assert!(!job.is_succeeded()); + } + + for pending in ["pending", "running", "queued", "validating_files"] { + let job = FineTuningJob { + id: "ftjob-3".to_string(), + status: pending.to_string(), + fine_tuned_model: None, + }; + assert!(!job.is_succeeded(), "'{pending}' is not succeeded"); + assert!(!job.is_terminal_failure(), "'{pending}' is not terminal failure"); + } + } + + #[test] + fn deserializes_get_response_with_fine_tuned_model() { + // Azure's GET response once the job succeeds. + let json = r#"{ + "id": "ftjob-abc123", + "status": "succeeded", + "model": "gpt-4o-mini", + "fine_tuned_model": "gpt-4o-mini.ft-abc123", + "object": "fine_tuning.job" + }"#; + let job: FineTuningJob = serde_json::from_str(json).expect("should deserialize"); + assert_eq!(job.id, "ftjob-abc123"); + assert!(job.is_succeeded()); + assert_eq!(job.fine_tuned_model.as_deref(), Some("gpt-4o-mini.ft-abc123")); + } + + #[test] + fn deserializes_pending_response_without_fine_tuned_model() { + // A freshly-created job has no fine_tuned_model yet. + let json = r#"{ "id": "ftjob-abc123", "status": "pending", "object": "fine_tuning.job" }"#; + let job: FineTuningJob = serde_json::from_str(json).expect("should deserialize"); + assert_eq!(job.status, "pending"); + assert!(job.fine_tuned_model.is_none()); + } +} + +#[cfg(all(test, feature = "test-utils"))] +mod http_tests { + use super::*; + use crate::azure::{AzureClientConfig, AzureClientConfigExt}; + use httpmock::{Method::GET, Method::POST, MockServer}; + use serde_json::json; + + fn test_client() -> AzureFoundryFineTuningClient { + // The account endpoint is passed per call, so the client itself needs no + // endpoint override — the MockServer URL is the `endpoint` argument. + AzureFoundryFineTuningClient::new( + Client::new(), + AzureTokenCache::new(AzureClientConfig::mock()), + ) + } + + /// The POST must hit the data-plane fine-tuning path with the pinned + /// api-version, carry a bearer token, and send the snake_case body. This is + /// the request-shape contract the gateway/controller relies on. + #[tokio::test] + async fn create_job_builds_correct_request() { + let server = MockServer::start_async().await; + + let request_mock = server + .mock_async(|when, then| { + when.method(POST) + .path("/openai/fine_tuning/jobs") + .query_param("api-version", OPENAI_FINE_TUNING_API_VERSION) + .header_exists("authorization") + .json_body(json!({ + "model": "gpt-4o-mini", + "training_file": "https://acct.blob.core.windows.net/data/training.jsonl" + })); + then.status(201).json_body(json!({ + "id": "ftjob-xyz", + "status": "pending", + "object": "fine_tuning.job" + })); + }) + .await; + + let job = test_client() + .create_fine_tuning_job( + &server.base_url(), + "gpt-4o-mini", + "https://acct.blob.core.windows.net/data/training.jsonl", + ) + .await + .expect("create fine-tuning job should succeed"); + + request_mock.assert_async().await; + assert_eq!(job.id, "ftjob-xyz"); + assert_eq!(job.status, "pending"); + assert!(job.fine_tuned_model.is_none()); + } + + /// The GET must hit `/openai/fine_tuning/jobs/{job_id}` and parse the tuned + /// model name out of a succeeded response. + #[tokio::test] + async fn get_job_builds_correct_request_and_parses_success() { + let server = MockServer::start_async().await; + + let request_mock = server + .mock_async(|when, then| { + when.method(GET) + .path("/openai/fine_tuning/jobs/ftjob-xyz") + .query_param("api-version", OPENAI_FINE_TUNING_API_VERSION) + .header_exists("authorization"); + then.status(200).json_body(json!({ + "id": "ftjob-xyz", + "status": "succeeded", + "fine_tuned_model": "gpt-4o-mini.ft-xyz", + "object": "fine_tuning.job" + })); + }) + .await; + + let job = test_client() + .get_fine_tuning_job(&server.base_url(), "ftjob-xyz") + .await + .expect("get fine-tuning job should succeed"); + + request_mock.assert_async().await; + assert!(job.is_succeeded()); + assert_eq!(job.fine_tuned_model.as_deref(), Some("gpt-4o-mini.ft-xyz")); + } +} diff --git a/crates/alien-azure-clients/src/lib.rs b/crates/alien-azure-clients/src/lib.rs index 8b6fc5373..feff96b8a 100644 --- a/crates/alien-azure-clients/src/lib.rs +++ b/crates/alien-azure-clients/src/lib.rs @@ -33,6 +33,9 @@ pub use azure::managed_identity::{ }; pub use azure::monitor::{AzureMonitorClient, MonitorApi}; pub use azure::network::{AzureNetworkClient, NetworkApi}; +pub use azure::openai_finetuning::{ + AzureFoundryFineTuningClient, FineTuningJob, FineTuningJobCreateRequest, FoundryFineTuningApi, +}; pub use azure::resource_graph::{AzureResourceGraphClient, ResourceGraphApi}; pub use azure::resource_skus::{ AzureResourceSkusClient, ResourceSku, ResourceSkuLocationInfo, ResourceSkuRestriction, diff --git a/crates/alien-cloudformation/src/emitters/aws/ai.rs b/crates/alien-cloudformation/src/emitters/aws/ai.rs index a9c33a2f6..a92428fb7 100644 --- a/crates/alien-cloudformation/src/emitters/aws/ai.rs +++ b/crates/alien-cloudformation/src/emitters/aws/ai.rs @@ -3,8 +3,9 @@ //! AWS Bedrock is a regional, account-scoped service with no per-stack //! cloud resource to provision. The emitter returns zero CFN resources for //! the AI resource itself and emits `AWS::IAM::Policy` resources for every -//! permission profile that references `ai/invoke` on this resource, attaching -//! them to the corresponding service-account role. +//! permission profile that references an `ai/*` set (e.g. `ai/invoke`, or +//! `ai/finetune` when the resource declares a fine-tuning job) on this +//! resource, attaching them to the corresponding service-account role. //! //! The import ref carries the region so the controller can reconstruct the //! Bedrock endpoint without a cloud round-trip. diff --git a/crates/alien-core/src/bin/schema_exporter.rs b/crates/alien-core/src/bin/schema_exporter.rs index 67c45eccf..d46e8fb76 100644 --- a/crates/alien-core/src/bin/schema_exporter.rs +++ b/crates/alien-core/src/bin/schema_exporter.rs @@ -93,6 +93,8 @@ use utoipa::OpenApi; PostgresOutputs, Ai, AiOutputs, + FinetuneSpec, + FinetuneMethod, Queue, QueueOutputs, Email, diff --git a/crates/alien-core/src/bindings/ai.rs b/crates/alien-core/src/bindings/ai.rs index f2552cf1d..f4bf118a4 100644 --- a/crates/alien-core/src/bindings/ai.rs +++ b/crates/alien-core/src/bindings/ai.rs @@ -23,6 +23,22 @@ pub enum AiBinding { External(ExternalAiBinding), } +/// A tuned model the gateway can route to, produced by a completed +/// fine-tuning job. Maps the public `served_id` an app requests to the +/// provider-native artifact the gateway forwards to (a Bedrock custom-model +/// ARN, a Vertex tuned endpoint id, or a Foundry deployment name). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +pub struct TunedModel { + /// The public model id apps send in the `model` field. + pub served_id: String, + /// The provider-native upstream artifact (custom-model ARN / tuned endpoint + /// id / deployment name) the gateway forwards to. + pub upstream_id: String, +} + /// AWS Bedrock AI binding configuration #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] @@ -31,6 +47,10 @@ pub enum AiBinding { pub struct BedrockAiBinding { /// The AWS region where Bedrock is accessed pub region: String, + /// A tuned model served alongside the base catalog, if the resource + /// declared a completed fine-tuning job. Absent for a pure inference gateway. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tuned_model: Option, } /// GCP Vertex AI binding configuration @@ -43,6 +63,10 @@ pub struct VertexAiBinding { pub project: String, /// The Vertex AI region (e.g., "us-central1") pub location: String, + /// A tuned model served alongside the base catalog, if the resource + /// declared a completed fine-tuning job. Absent for a pure inference gateway. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tuned_model: Option, } /// Azure AI Foundry binding configuration @@ -55,6 +79,10 @@ pub struct FoundryAiBinding { pub endpoint: String, /// The Azure account or subscription identifier pub account: String, + /// A tuned model served alongside the base catalog, if the resource + /// declared a completed fine-tuning job. Absent for a pure inference gateway. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tuned_model: Option, } /// External AI provider binding configuration (BYO-key). @@ -98,6 +126,7 @@ impl AiBinding { pub fn bedrock(region: impl Into) -> Self { Self::Bedrock(BedrockAiBinding { region: region.into(), + tuned_model: None, }) } @@ -105,6 +134,7 @@ impl AiBinding { Self::Vertex(VertexAiBinding { project: project.into(), location: location.into(), + tuned_model: None, }) } @@ -112,9 +142,45 @@ impl AiBinding { Self::Foundry(FoundryAiBinding { endpoint: endpoint.into(), account: account.into(), + tuned_model: None, }) } + /// Attach a tuned model to a managed binding, so the gateway routes + /// `served_id` to `upstream_id` alongside the base catalog. A no-op on the + /// `External` (BYO-key) variant, which the gateway does not serve. + pub fn with_tuned_model(self, served_id: impl Into, upstream_id: impl Into) -> Self { + let tuned = TunedModel { + served_id: served_id.into(), + upstream_id: upstream_id.into(), + }; + match self { + Self::Bedrock(b) => Self::Bedrock(BedrockAiBinding { + tuned_model: Some(tuned), + ..b + }), + Self::Vertex(b) => Self::Vertex(VertexAiBinding { + tuned_model: Some(tuned), + ..b + }), + Self::Foundry(b) => Self::Foundry(FoundryAiBinding { + tuned_model: Some(tuned), + ..b + }), + Self::External(b) => Self::External(b), + } + } + + /// The tuned model attached to this binding, if any. + pub fn tuned_model(&self) -> Option<&TunedModel> { + match self { + Self::Bedrock(b) => b.tuned_model.as_ref(), + Self::Vertex(b) => b.tuned_model.as_ref(), + Self::Foundry(b) => b.tuned_model.as_ref(), + Self::External(_) => None, + } + } + pub fn external( provider: impl Into, api_key: impl Into>, @@ -141,6 +207,41 @@ mod tests { assert_eq!(binding, deserialized); } + #[test] + fn test_bedrock_binding_without_tuned_model_omits_field() { + let binding = AiBinding::bedrock("us-east-1"); + let json = serde_json::to_value(&binding).unwrap(); + assert!( + json.get("tunedModel").is_none(), + "an untuned binding must omit tunedModel so the inference-only wire shape is unchanged" + ); + } + + #[test] + fn test_bedrock_binding_with_tuned_model_roundtrip() { + let binding = AiBinding::bedrock("us-east-1") + .with_tuned_model("finance-model", "arn:aws:bedrock:us-east-1:123:custom-model/abc"); + + let json = serde_json::to_value(&binding).unwrap(); + assert_eq!(json["service"], "bedrock"); + assert_eq!(json["tunedModel"]["servedId"], "finance-model"); + assert_eq!( + json["tunedModel"]["upstreamId"], + "arn:aws:bedrock:us-east-1:123:custom-model/abc" + ); + + let back: AiBinding = serde_json::from_value(json).unwrap(); + assert_eq!(binding, back); + assert_eq!(back.tuned_model().unwrap().served_id, "finance-model"); + } + + #[test] + fn test_external_binding_ignores_tuned_model() { + // The gateway does not serve BYO-key providers, so a tuned model is a no-op there. + let binding = AiBinding::external("openai", "sk-x").with_tuned_model("x", "y"); + assert!(binding.tuned_model().is_none()); + } + #[test] fn test_vertex_binding_roundtrip() { let binding = AiBinding::vertex("my-project", "us-central1"); diff --git a/crates/alien-core/src/resources/ai.rs b/crates/alien-core/src/resources/ai.rs index a2180c94c..526d5a56f 100644 --- a/crates/alien-core/src/resources/ai.rs +++ b/crates/alien-core/src/resources/ai.rs @@ -1,12 +1,102 @@ use crate::error::{ErrorData, Result}; use crate::resource::{ResourceDefinition, ResourceOutputsDefinition, ResourceRef}; -use crate::ResourceType; +use crate::{ResourceType, Storage}; use alien_error::AlienError; use bon::Builder; use serde::{Deserialize, Serialize}; use std::any::Any; use std::fmt::Debug; +/// The fine-tuning method applied to the base model. +/// +/// The gateway-side controllers map each variant onto the provider's native +/// technique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is +/// direct preference optimization (Bedrock/Foundry), and `Lora` requests a +/// parameter-efficient adapter where the provider exposes it. Providers that +/// only implement a subset reject unsupported methods at job-submit time. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))] +#[serde(rename_all = "lowercase")] +pub enum FinetuneMethod { + /// Supervised fine-tuning on labelled prompt/response pairs (default). + Sft, + /// Direct preference optimization on chosen/rejected pairs. + Dpo, + /// Low-rank adaptation (parameter-efficient) where the provider exposes it. + Lora, +} + +impl Default for FinetuneMethod { + fn default() -> Self { + Self::Sft + } +} + +/// Declares that an [`Ai`] resource should fine-tune a base model in the +/// customer's cloud before serving it. +/// +/// The training data lives in a customer-owned [`Storage`](crate::Storage) +/// bucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the +/// cloud controller submits the provider's tuning job (Bedrock +/// `CreateModelCustomizationJob`, Vertex tuning job, or Foundry +/// `fine_tuning.jobs`), polls it to completion via the heartbeat loop, and +/// records the resulting artifact so the gateway can route `served_model_id` +/// to it. Base-model inference is unaffected — an `Ai` without a `finetune` +/// spec behaves exactly as before. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct FinetuneSpec { + /// Provider-native base-model identifier to tune (e.g. an Amazon Nova model + /// id on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on + /// Foundry). Validated against the target cloud when the job is submitted. + pub base_model: String, + + /// The storage resource holding the JSONL training dataset. The controller + /// reads it from the customer bucket the storage resolves to; the data + /// never leaves the customer's cloud. + pub training_data: String, + + /// Object key of the training file within `training_data`. + /// Defaults to `training.jsonl`. + #[serde(default = "default_training_key", skip_serializing_if = "is_default_training_key")] + pub training_key: String, + + /// The public model id apps use to invoke the tuned model through the + /// gateway (the `model` field in an OpenAI-compatible request). Defaults to + /// `-tuned`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub served_model_id: Option, + + /// The fine-tuning method. Defaults to supervised fine-tuning. + #[serde(default, skip_serializing_if = "is_default_method")] + pub method: FinetuneMethod, +} + +fn default_training_key() -> String { + "training.jsonl".to_string() +} + +fn is_default_training_key(key: &str) -> bool { + key == "training.jsonl" +} + +fn is_default_method(method: &FinetuneMethod) -> bool { + *method == FinetuneMethod::default() +} + +impl FinetuneSpec { + /// The public model id the gateway serves the tuned model under, falling + /// back to `-tuned` when the spec doesn't set one explicitly. + pub fn served_model_id_or_default(&self, ai_id: &str) -> String { + self.served_model_id + .clone() + .unwrap_or_else(|| format!("{ai_id}-tuned")) + } +} + /// Represents an AI Gateway resource that provides a unified interface to /// managed AI inference services across cloud providers. /// @@ -14,6 +104,10 @@ use std::fmt::Debug; /// other BYO infrastructure (e.g. external Redis for `kv`), an external AI /// provider is supplied at deploy time as an `ExternalBinding::Ai` in the /// stack's external-bindings map; the executor then skips the cloud controller. +/// +/// When `finetune` is set, the resource also tunes a base model in the +/// customer's cloud and serves the result through the same gateway (see +/// [`FinetuneSpec`]). #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Builder)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -23,6 +117,12 @@ pub struct Ai { /// Maximum 64 characters. #[builder(start_fn)] pub id: String, + + /// Optional fine-tuning declaration. When present, the resource tunes + /// `finetune.base_model` on the customer's cloud and serves the result + /// alongside the base models. Absent for a pure inference gateway. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finetune: Option, } impl Ai { @@ -82,7 +182,16 @@ impl ResourceDefinition for Ai { } fn get_dependencies(&self) -> Vec { - Vec::new() + // The training dataset lives in a customer Storage bucket that must be + // provisioned before the tuning job can read it, so a finetune spec adds + // that storage as a dependency. A pure inference gateway has none. + match &self.finetune { + Some(spec) => vec![ResourceRef::new( + Storage::RESOURCE_TYPE, + spec.training_data.clone(), + )], + None => Vec::new(), + } } fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> { @@ -100,6 +209,20 @@ impl ResourceDefinition for Ai { reason: "the 'id' field is immutable".to_string(), })); } + + // The tuned artifact is derived from the training data + base model, so + // repointing them would silently serve a different model under the same + // served id. Require a new resource id (hence a fresh tuning job) instead. + if let (Some(old), Some(new)) = (&self.finetune, &new_ai.finetune) { + if old.base_model != new.base_model || old.training_data != new.training_data { + return Err(AlienError::new(ErrorData::InvalidResourceUpdate { + resource_id: self.id.clone(), + reason: "finetune 'baseModel' and 'trainingData' are immutable; \ + create a new AI resource to retrain" + .to_string(), + })); + } + } Ok(()) } @@ -157,6 +280,121 @@ mod tests { assert!(original.validate_update(&invalid_update).is_err()); } + #[test] + fn test_ai_finetune_dependency() { + let base = Ai::new("llm".to_string()).build(); + assert!( + base.get_dependencies().is_empty(), + "a pure inference gateway has no dependencies" + ); + + let tuned = Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "training-set".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build(); + assert_eq!( + tuned.get_dependencies(), + vec![ResourceRef::new(Storage::RESOURCE_TYPE, "training-set")], + "a finetune spec depends on its training-data storage" + ); + } + + #[test] + fn test_ai_served_model_id_default() { + let spec = FinetuneSpec { + base_model: "b".to_string(), + training_data: "d".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }; + assert_eq!(spec.served_model_id_or_default("llm"), "llm-tuned"); + + let explicit = FinetuneSpec { + served_model_id: Some("finance-model".to_string()), + ..spec + }; + assert_eq!(explicit.served_model_id_or_default("llm"), "finance-model"); + } + + #[test] + fn test_ai_finetune_immutable_fields_rejected() { + let original = Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "set-a".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build(); + + // Same base + data, changed method: allowed. + let ok = Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "set-a".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Dpo, + }) + .build(); + assert!(original.validate_update(&ok).is_ok()); + + // Changed training data: rejected. + let repoint = Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "set-b".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build(); + let err = original + .validate_update(&repoint) + .expect_err("repointing training data must be rejected"); + assert_eq!(err.code, "INVALID_RESOURCE_UPDATE"); + } + + #[test] + fn test_ai_finetune_roundtrip_camel_case() { + let ai = Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "training-set".to_string(), + training_key: "data.jsonl".to_string(), + served_model_id: Some("finance-model".to_string()), + method: FinetuneMethod::Lora, + }) + .build(); + + let json = serde_json::to_value(&ai).unwrap(); + assert_eq!(json["finetune"]["baseModel"], "amazon.nova-lite-v1:0"); + assert_eq!(json["finetune"]["trainingData"], "training-set"); + assert_eq!(json["finetune"]["trainingKey"], "data.jsonl"); + assert_eq!(json["finetune"]["servedModelId"], "finance-model"); + assert_eq!(json["finetune"]["method"], "lora"); + + let back: Ai = serde_json::from_value(json).unwrap(); + assert_eq!(ai, back); + } + + #[test] + fn test_ai_without_finetune_omits_field() { + let ai = Ai::new("llm".to_string()).build(); + let json = serde_json::to_value(&ai).unwrap(); + assert!( + json.get("finetune").is_none(), + "finetune must be omitted when unset so the inference-only shape is unchanged" + ); + } + #[test] fn test_ai_outputs_serialization() { let outputs = AiOutputs { diff --git a/crates/alien-gateway/src/config.rs b/crates/alien-gateway/src/config.rs index 79baa0459..f3734e776 100644 --- a/crates/alien-gateway/src/config.rs +++ b/crates/alien-gateway/src/config.rs @@ -20,7 +20,7 @@ use alien_error::{AlienError, Context, IntoAlienError}; use crate::creds::{AmbientCred, AwsSigV4Cred, BearerTokenCred}; use crate::error::{ErrorData, Result}; -use crate::{GatewayBinding, GatewayRoute}; +use crate::{GatewayBinding, GatewayRoute, TunedRoute}; /// The shared workload credential resolver used for the mint-gated (runtime-less) path. pub type Managed = Arc; @@ -157,6 +157,12 @@ fn bindings_from_pairs( /// does not serve. The binding name is the canonical path segment (lowercased, as the /// env-var key encodes it). fn gateway_binding(name: &str, binding: AiBinding) -> Option { + // A managed binding may carry a tuned model the gateway serves alongside the + // static catalog. Read it before consuming the binding into its variant. + let tuned = binding.tuned_model().map(|t| TunedRoute { + served_id: t.served_id.clone(), + upstream_id: t.upstream_id.clone(), + }); match binding { AiBinding::Bedrock(b) => Some(GatewayBinding { name: name.to_string(), @@ -164,6 +170,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { region: Some(b.region), project: None, azure_endpoint: None, + tuned, }), AiBinding::Vertex(b) => Some(GatewayBinding { name: name.to_string(), @@ -171,6 +178,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { region: Some(b.location), project: Some(b.project), azure_endpoint: None, + tuned, }), AiBinding::Foundry(b) => Some(GatewayBinding { name: name.to_string(), @@ -178,6 +186,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { region: None, project: None, azure_endpoint: Some(b.endpoint), + tuned, }), // External is a BYO-key provider, not an ambient-managed cloud — not served here. AiBinding::External(_) => None, @@ -239,6 +248,7 @@ pub async fn resolve_route(binding: GatewayBinding, managed: Option<&Managed>) - azure_endpoint: binding.azure_endpoint, cred, upstream_base_override: None, + tuned: binding.tuned, }) } diff --git a/crates/alien-gateway/src/lib.rs b/crates/alien-gateway/src/lib.rs index 67838a190..ccd15fb16 100644 --- a/crates/alien-gateway/src/lib.rs +++ b/crates/alien-gateway/src/lib.rs @@ -44,6 +44,21 @@ pub struct GatewayBinding { pub project: Option, /// Azure account endpoint, e.g. `https://acct.openai.azure.com/`. pub azure_endpoint: Option, + /// A tuned model this binding serves alongside the static catalog, produced + /// by a completed fine-tuning job. `None` for a pure inference gateway. + pub tuned: Option, +} + +/// A tuned model the gateway routes to: the public id an app requests mapped to +/// the provider-native upstream artifact (a Bedrock custom-model ARN, a Vertex +/// tuned endpoint id, or a Foundry deployment name). Populated from the +/// binding's `tunedModel`; consumed by the router before the static catalog. +#[derive(Debug, Clone)] +pub struct TunedRoute { + /// The public model id apps send in the `model` field. + pub served_id: String, + /// The provider-native upstream artifact the gateway forwards to. + pub upstream_id: String, } /// A running gateway: its loopback base URL and the server task that keeps it alive diff --git a/crates/alien-gateway/src/router.rs b/crates/alien-gateway/src/router.rs index e59ec8d74..15d111836 100644 --- a/crates/alien-gateway/src/router.rs +++ b/crates/alien-gateway/src/router.rs @@ -45,6 +45,10 @@ pub struct GatewayRoute { /// host (the per-protocol path is still appended). Lets tests aim a binding at a /// mock upstream. pub upstream_base_override: Option, + /// A tuned model this route serves alongside the static catalog. Checked + /// before the catalog so a completed fine-tuning job's artifact is reachable + /// under its public `served_id`. `None` for a pure inference gateway. + pub tuned: Option, } struct AppState { @@ -161,7 +165,20 @@ async fn proxy( }) })?; - let (mut payload, model) = parse_model_request(&body)?; + let (payload, model) = parse_model_request(&body)?; + + // A tuned model produced by a fine-tuning job is served under its public + // `served_id` and is not in the static catalog (its upstream artifact is + // created at runtime). Check it before the catalog. All three clouds serve a + // tuned model over the OpenAI-compatible chat path: on Bedrock the + // `upstream_id` is a custom-model id, on Vertex a tuned-endpoint model id, on + // Foundry a deployment name — in every case it becomes the request body's + // `model` against the cloud's OpenAI endpoint. + if let Some(tuned) = &route.tuned { + if model == tuned.served_id { + return proxy_openai(&state, route, tuned.upstream_id.clone(), payload).await; + } + } // Cloud-scoped resolution: Claude ids appear once per cloud, so a first-match // resolve would always land on another cloud's entry and fail the cloud filter. @@ -193,14 +210,26 @@ async fn proxy( .await; } - payload["model"] = Value::String(cm.upstream_id.to_string()); + proxy_openai(&state, route, cm.upstream_id.to_string(), payload).await +} + +/// Forward an OpenAI-protocol chat request: rewrite the body `model` to +/// `upstream_id`, build the cloud's OpenAI chat endpoint, sign, and stream the +/// reply back untranslated. Shared by base-catalog models and tuned models. +async fn proxy_openai( + state: &AppState, + route: &GatewayRoute, + upstream_id: String, + mut payload: Value, +) -> Result { + payload["model"] = Value::String(upstream_id); let upstream_body = serde_json::to_vec(&payload) .into_alien_error() .context(ErrorData::Other { message: "could not re-serialize the rewritten request body".to_string(), })?; - let (url, aws_service) = upstream_target(route, cm.protocol)?; + let (url, aws_service) = upstream_target(route, Protocol::OpenAi)?; let upstream = sign_and_execute(&state.client, &route.cred, &url, aws_service, upstream_body, &[]).await?; @@ -971,6 +1000,7 @@ mod tests { azure_endpoint: None, cred: test_aws_cred(), upstream_base_override: Some(upstream.to_string()), + tuned: None, } } @@ -983,6 +1013,7 @@ mod tests { azure_endpoint: None, cred: AmbientCred::Bearer(BearerTokenCred::static_token("t")), upstream_base_override: None, + tuned: None, } } @@ -1157,6 +1188,7 @@ mod tests { azure_endpoint: Some(endpoint.to_string()), cred: AmbientCred::Bearer(BearerTokenCred::static_token("t")), upstream_base_override: None, + tuned: None, } } @@ -1269,6 +1301,78 @@ mod tests { mock.assert_async().await; } + #[tokio::test] + async fn tuned_model_routes_to_upstream_artifact() { + // A request for the tuned model's public served id must be rewritten to the + // tuned upstream artifact (here a Bedrock custom-model id) and hit the same + // OpenAI chat path — proving the tuned override fires before catalog lookup. + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(POST) + .path("/openai/v1/chat/completions") + .body_contains("custom-model/finance-abc123") + .header_exists("authorization"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"id":"cmpl-1","choices":[{"message":{"content":"pong"}}]}"#); + }) + .await; + + let mut route = aws_route(&server.base_url()); + route.tuned = Some(crate::TunedRoute { + served_id: "finance-model".to_string(), + upstream_id: "custom-model/finance-abc123".to_string(), + }); + + let url = serve(build_router(vec![route])).await; + let resp = reqwest::Client::new() + .post(format!("{url}/llm/v1/chat/completions")) + .json(&json!({"model":"finance-model","messages":[{"role":"user","content":"hi"}]})) + .send() + .await + .expect("proxy request"); + + assert_eq!(resp.status(), 200); + let text = resp.text().await.unwrap(); + assert!(text.contains("\"pong\""), "upstream body must pass through: {text}"); + mock.assert_async().await; + } + + #[tokio::test] + async fn tuned_model_does_not_shadow_base_catalog() { + // With a tuned model configured, a base catalog id still resolves normally — + // the override only triggers on an exact served-id match. + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(POST) + .path("/openai/v1/chat/completions") + .body_contains("openai.gpt-oss-20b-1:0"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"id":"cmpl-1","choices":[{"message":{"content":"pong"}}]}"#); + }) + .await; + + let mut route = aws_route(&server.base_url()); + route.tuned = Some(crate::TunedRoute { + served_id: "finance-model".to_string(), + upstream_id: "custom-model/finance-abc123".to_string(), + }); + + let url = serve(build_router(vec![route])).await; + let resp = reqwest::Client::new() + .post(format!("{url}/llm/v1/chat/completions")) + .json(&json!({"model":"gpt-oss-20b","messages":[{"role":"user","content":"hi"}]})) + .send() + .await + .expect("proxy request"); + + assert_eq!(resp.status(), 200); + mock.assert_async().await; + } + #[tokio::test] async fn streams_sse_through_unchanged() { let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"po\"}}]}\n\n\ diff --git a/crates/alien-gateway/tests/integration.rs b/crates/alien-gateway/tests/integration.rs index d6b697071..6da65209e 100644 --- a/crates/alien-gateway/tests/integration.rs +++ b/crates/alien-gateway/tests/integration.rs @@ -80,6 +80,7 @@ async fn routes_two_clouds_with_rewrite_auth_and_passthrough() { azure_endpoint: None, cred: aws_cred(), upstream_base_override: Some(aws_upstream.base_url()), + tuned: None, }, GatewayRoute { name: "azllm".to_string(), @@ -89,6 +90,7 @@ async fn routes_two_clouds_with_rewrite_auth_and_passthrough() { azure_endpoint: Some(azure_upstream.base_url()), cred: AmbientCred::Bearer(BearerTokenCred::static_token("test-azure-token")), upstream_base_override: Some(azure_upstream.base_url()), + tuned: None, }, ]; diff --git a/crates/alien-gateway/tests/live_bedrock.rs b/crates/alien-gateway/tests/live_bedrock.rs index 1c7c1a93f..94c215f77 100644 --- a/crates/alien-gateway/tests/live_bedrock.rs +++ b/crates/alien-gateway/tests/live_bedrock.rs @@ -43,6 +43,7 @@ async fn live_bedrock_openai_chat() { azure_endpoint: None, cred, upstream_base_override: None, + tuned: None, }; let base = serve(build_router(vec![route])).await; @@ -98,6 +99,7 @@ async fn live_bedrock_claude_streaming() { azure_endpoint: None, cred, upstream_base_override: None, + tuned: None, }; let base = serve(build_router(vec![route])).await; diff --git a/crates/alien-gateway/tests/live_foundry_claude.rs b/crates/alien-gateway/tests/live_foundry_claude.rs index b9f0c9e5c..74e467eea 100644 --- a/crates/alien-gateway/tests/live_foundry_claude.rs +++ b/crates/alien-gateway/tests/live_foundry_claude.rs @@ -48,6 +48,7 @@ fn foundry_route() -> GatewayRoute { azure_endpoint: Some(endpoint), cred: AmbientCred::Bearer(BearerTokenCred::static_token(token)), upstream_base_override: None, + tuned: None, } } diff --git a/crates/alien-gateway/tests/live_vertex_claude.rs b/crates/alien-gateway/tests/live_vertex_claude.rs index 6b3722d68..c5486db5a 100644 --- a/crates/alien-gateway/tests/live_vertex_claude.rs +++ b/crates/alien-gateway/tests/live_vertex_claude.rs @@ -40,6 +40,7 @@ fn vertex_route() -> GatewayRoute { azure_endpoint: None, cred: AmbientCred::Bearer(BearerTokenCred::static_token(token)), upstream_base_override: None, + tuned: None, } } diff --git a/crates/alien-gcp-clients/src/gcp/aiplatform.rs b/crates/alien-gcp-clients/src/gcp/aiplatform.rs new file mode 100644 index 000000000..509a53ba2 --- /dev/null +++ b/crates/alien-gcp-clients/src/gcp/aiplatform.rs @@ -0,0 +1,406 @@ +//! Vertex AI Platform client covering the tuning-job surface used for +//! fine-tuning. +//! +//! Only the two calls the fine-tuning controller needs are exposed: +//! `create_tuning_job` (POST `.../tuningJobs`) and `get_tuning_job` +//! (GET `.../tuningJobs/{id}`). See: +//! +//! +//! Unlike the other GCP services, Vertex AI is *regional*: the host is +//! `{location}-aiplatform.googleapis.com`. Because [`GcpServiceConfig::base_url`] +//! must return a `&'static str`, the location-specific base URL is built in the +//! client constructor and injected as a service override, so [`GcpClientBase`]'s +//! standard override lookup resolves it. + +use crate::gcp::api_client::{GcpClientBase, GcpServiceConfig}; +use crate::gcp::longrunning::Status; +use crate::gcp::GcpClientConfig; +use crate::gcp::GcpClientConfigExt; +use crate::gcp::ServiceOverrides; +use alien_client_core::Result; +use bon::Builder; +use reqwest::{Client, Method}; +use serde::{Deserialize, Serialize}; +use std::fmt::Debug; + +#[cfg(feature = "test-utils")] +use mockall::automock; + +/// Service key used to inject the regional base-URL override for Vertex AI. +const AIPLATFORM_SERVICE_KEY: &str = "aiplatform"; + +/// Vertex AI (aiplatform) service configuration. +/// +/// The `base_url` returned here is a placeholder for the global host; the real, +/// region-scoped host is always supplied via a service override installed by +/// [`AiPlatformClient::new`], so this default is never used in practice. +#[derive(Debug)] +pub struct AiPlatformServiceConfig; + +impl GcpServiceConfig for AiPlatformServiceConfig { + fn base_url(&self) -> &'static str { + // Overridden per-region in the constructor; kept valid as a fallback. + "https://aiplatform.googleapis.com/v1" + } + + fn default_audience(&self) -> &'static str { + "https://aiplatform.googleapis.com/" + } + + fn service_name(&self) -> &'static str { + "Vertex AI" + } + + fn service_key(&self) -> &'static str { + AIPLATFORM_SERVICE_KEY + } +} + +#[cfg_attr(feature = "test-utils", automock)] +#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] +pub trait AiPlatformApi: Send + Sync + Debug { + /// Submits a supervised tuning job and returns the created [`TuningJob`], + /// whose `name` (`projects/{p}/locations/{l}/tuningJobs/{id}`) is the handle + /// to poll. + async fn create_tuning_job(&self, request: CreateTuningJobRequest) -> Result; + + /// Fetches the current state of a tuning job by its full resource name + /// (`projects/{p}/locations/{l}/tuningJobs/{id}`) or bare id. + async fn get_tuning_job(&self, name: String) -> Result; +} + +/// Vertex AI tuning client. +#[derive(Debug)] +pub struct AiPlatformClient { + base: GcpClientBase, + project_id: String, + location: String, +} + +impl AiPlatformClient { + /// Builds a client pinned to the config's region. Installs the regional + /// `aiplatform` base-URL override (unless the caller already set one, e.g. + /// a test endpoint) so all requests hit `{location}-aiplatform...`. + pub fn new(client: Client, config: GcpClientConfig) -> Self { + let project_id = config.project_id.clone(); + let location = config.region.clone(); + + // Only synthesize the regional host when no override is present; this + // lets tests point the client at a mock endpoint via service overrides. + let mut config = config; + if config.get_service_endpoint_option(AIPLATFORM_SERVICE_KEY).is_none() { + let regional = format!("https://{location}-aiplatform.googleapis.com/v1"); + let mut overrides = config.service_overrides.unwrap_or(ServiceOverrides { + endpoints: std::collections::HashMap::new(), + }); + overrides + .endpoints + .insert(AIPLATFORM_SERVICE_KEY.to_string(), regional); + config.service_overrides = Some(overrides); + } + + Self { + base: GcpClientBase::new(client, config, Box::new(AiPlatformServiceConfig)), + project_id, + location, + } + } + + /// The collection path `projects/{p}/locations/{l}/tuningJobs`. + fn tuning_jobs_path(&self) -> String { + format!( + "projects/{}/locations/{}/tuningJobs", + self.project_id, self.location + ) + } +} + +#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] +impl AiPlatformApi for AiPlatformClient { + async fn create_tuning_job(&self, request: CreateTuningJobRequest) -> Result { + let path = self.tuning_jobs_path(); + self.base + .execute_request( + Method::POST, + &path, + None, + Some(request.clone()), + &request.base_model, + ) + .await + } + + async fn get_tuning_job(&self, name: String) -> Result { + // Vertex returns the full resource name; accept either that or a bare id + // and build the region-scoped collection path either way. + let path = if name.contains("/tuningJobs/") { + name.clone() + } else { + format!("{}/{}", self.tuning_jobs_path(), name) + }; + + self.base + .execute_request(Method::GET, &path, None, Option::<()>::None, &name) + .await + } +} + +// --- Data Structures --- + +/// Request body for `tuningJobs.create`. +/// +/// Only the supervised-tuning surface is modelled: `baseModel`, the training +/// dataset URI (a `gs://` path to JSONL), and an optional display name. See +/// . +#[derive(Debug, Serialize, Deserialize, Clone, Builder)] +#[serde(rename_all = "camelCase")] +pub struct CreateTuningJobRequest { + /// Provider-native base model to tune (e.g. a Gemini model id). + pub base_model: String, + + /// Supervised tuning parameters, including the training dataset URI. + pub supervised_tuning_spec: SupervisedTuningSpec, + + /// Optional human-readable name for the resulting tuned model. + #[serde(skip_serializing_if = "Option::is_none")] + pub tuned_model_display_name: Option, +} + +/// Supervised tuning parameters. +#[derive(Debug, Serialize, Deserialize, Clone, Builder)] +#[serde(rename_all = "camelCase")] +pub struct SupervisedTuningSpec { + /// `gs://` URI of the JSONL training dataset in the customer bucket. + pub training_dataset_uri: String, + + /// Optional `gs://` URI of a JSONL validation dataset. + #[serde(skip_serializing_if = "Option::is_none")] + pub validation_dataset_uri: Option, + + /// Optional supervised-tuning hyperparameters. + #[serde(skip_serializing_if = "Option::is_none")] + pub hyper_parameters: Option, +} + +/// Supervised tuning hyperparameters (all optional; Vertex picks defaults). +#[derive(Debug, Serialize, Deserialize, Clone, Default, Builder)] +#[serde(rename_all = "camelCase")] +pub struct SupervisedHyperParameters { + /// Number of complete passes over the training dataset. + #[serde(skip_serializing_if = "Option::is_none")] + pub epoch_count: Option, + + /// Multiplier applied to the recommended learning rate. + #[serde(skip_serializing_if = "Option::is_none")] + pub learning_rate_multiplier: Option, +} + +/// A Vertex AI tuning job. +/// +/// Mirrors the fields the controller needs from +/// . +#[derive(Debug, Serialize, Deserialize, Clone, Default, Builder)] +#[serde(rename_all = "camelCase")] +pub struct TuningJob { + /// Server-assigned resource name: + /// `projects/{p}/locations/{l}/tuningJobs/{id}`. + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + + /// Current lifecycle state of the job. + #[serde(skip_serializing_if = "Option::is_none")] + pub state: Option, + + /// The tuned model produced once the job reaches `JOB_STATE_SUCCEEDED`. + #[serde(skip_serializing_if = "Option::is_none")] + pub tuned_model: Option, + + /// Populated when the job fails; carries the gRPC-style status. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// Reference to the artifacts a completed tuning job produced. +#[derive(Debug, Serialize, Deserialize, Clone, Default, Builder)] +#[serde(rename_all = "camelCase")] +pub struct TunedModelRef { + /// Resource name of the tuned Model: + /// `projects/{p}/locations/{l}/models/{model}@{version}`. + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// Resource name of the Endpoint serving the tuned model: + /// `projects/{p}/locations/{l}/endpoints/{endpoint}`. This is the id the + /// Vertex OpenAI-compat chat path routes to. + #[serde(skip_serializing_if = "Option::is_none")] + pub endpoint: Option, +} + +/// Lifecycle state of a Vertex AI job (subset of `google.cloud.aiplatform.v1.JobState`). +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum JobState { + /// The job state is unspecified. + JobStateUnspecified, + /// The job has been created and is awaiting resources. + JobStateQueued, + /// The job is pending; not yet running. + JobStatePending, + /// The job is currently running. + JobStateRunning, + /// The job completed successfully. + JobStateSucceeded, + /// The job failed. + JobStateFailed, + /// The job is being cancelled. + JobStateCancelling, + /// The job was cancelled. + JobStateCancelled, + /// The job was paused. + JobStatePaused, + /// The job expired. + JobStateExpired, + /// The job is being updated. + JobStateUpdating, + /// The job partially failed (some outputs missing). + JobStatePartiallySucceeded, +} + +impl JobState { + /// Whether the job is still making progress and should be polled again. + pub fn is_in_progress(self) -> bool { + matches!( + self, + JobState::JobStateUnspecified + | JobState::JobStateQueued + | JobState::JobStatePending + | JobState::JobStateRunning + | JobState::JobStateCancelling + | JobState::JobStatePaused + | JobState::JobStateUpdating + ) + } + + /// Whether the job reached a terminal *failure* state (never produces a + /// usable tuned model). + pub fn is_terminal_failure(self) -> bool { + matches!( + self, + JobState::JobStateFailed + | JobState::JobStateCancelled + | JobState::JobStateExpired + | JobState::JobStatePartiallySucceeded + ) + } +} + +impl TunedModelRef { + /// The id the Vertex OpenAI-compat chat endpoint accepts for a tuned model: + /// the serving `endpoint` resource name, falling back to the `model` + /// resource name if the endpoint is absent. + pub fn upstream_id(&self) -> Option<&str> { + self.endpoint + .as_deref() + .or(self.model.as_deref()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn create_request_serializes_supervised_shape() { + let request = CreateTuningJobRequest::builder() + .base_model("gemini-2.0-flash-001".to_string()) + .supervised_tuning_spec( + SupervisedTuningSpec::builder() + .training_dataset_uri("gs://my-bucket/training.jsonl".to_string()) + .build(), + ) + .tuned_model_display_name("my-ai-tuned".to_string()) + .build(); + + let json = serde_json::to_value(&request).expect("request should serialize"); + + assert_eq!(json["baseModel"], "gemini-2.0-flash-001"); + assert_eq!( + json["supervisedTuningSpec"]["trainingDatasetUri"], + "gs://my-bucket/training.jsonl" + ); + assert_eq!(json["tunedModelDisplayName"], "my-ai-tuned"); + // Optional fields must be omitted, not sent as null (Vertex rejects nulls). + assert!( + json["supervisedTuningSpec"].get("validationDatasetUri").is_none(), + "unset validationDatasetUri must be omitted, got {json:?}" + ); + assert!( + json["supervisedTuningSpec"].get("hyperParameters").is_none(), + "unset hyperParameters must be omitted, got {json:?}" + ); + } + + #[test] + fn tuning_job_deserializes_succeeded_with_endpoint() { + let body = r#"{ + "name": "projects/p/locations/us-central1/tuningJobs/123", + "state": "JOB_STATE_SUCCEEDED", + "tunedModel": { + "model": "projects/p/locations/us-central1/models/456@1", + "endpoint": "projects/p/locations/us-central1/endpoints/789" + } + }"#; + + let job: TuningJob = serde_json::from_str(body).expect("job should deserialize"); + + assert_eq!(job.state, Some(JobState::JobStateSucceeded)); + assert!(job.state.unwrap().is_terminal_failure() == false); + let tuned = job.tuned_model.expect("tuned model present on success"); + // The OpenAI-compat chat path routes to the serving endpoint. + assert_eq!( + tuned.upstream_id(), + Some("projects/p/locations/us-central1/endpoints/789") + ); + } + + #[test] + fn tuning_job_upstream_falls_back_to_model_without_endpoint() { + let tuned = TunedModelRef::builder() + .model("projects/p/locations/us-central1/models/456@1".to_string()) + .build(); + assert_eq!( + tuned.upstream_id(), + Some("projects/p/locations/us-central1/models/456@1") + ); + } + + #[test] + fn tuning_job_deserializes_failed_with_error() { + let body = r#"{ + "name": "projects/p/locations/us-central1/tuningJobs/123", + "state": "JOB_STATE_FAILED", + "error": { "code": 9, "message": "training data invalid" } + }"#; + + let job: TuningJob = serde_json::from_str(body).expect("job should deserialize"); + + let state = job.state.expect("state present"); + assert!(state.is_terminal_failure(), "FAILED must be terminal failure"); + assert!(!state.is_in_progress()); + assert_eq!(job.error.expect("error present").message, "training data invalid"); + } + + #[test] + fn in_progress_states_are_polled_not_terminal() { + for state in [ + JobState::JobStatePending, + JobState::JobStateRunning, + JobState::JobStateQueued, + ] { + assert!(state.is_in_progress(), "{state:?} should be in-progress"); + assert!(!state.is_terminal_failure()); + } + } +} diff --git a/crates/alien-gcp-clients/src/gcp/mod.rs b/crates/alien-gcp-clients/src/gcp/mod.rs index 096e65978..f9cc69d03 100644 --- a/crates/alien-gcp-clients/src/gcp/mod.rs +++ b/crates/alien-gcp-clients/src/gcp/mod.rs @@ -1,3 +1,4 @@ +pub mod aiplatform; pub mod api_client; pub mod artifactregistry; pub mod cloud_sql; diff --git a/crates/alien-gcp-clients/src/lib.rs b/crates/alien-gcp-clients/src/lib.rs index 51e1eccbe..7b43324be 100644 --- a/crates/alien-gcp-clients/src/lib.rs +++ b/crates/alien-gcp-clients/src/lib.rs @@ -12,6 +12,7 @@ pub mod platform { } // Re-export all client APIs +pub use gcp::aiplatform::{AiPlatformApi, AiPlatformClient}; pub use gcp::artifactregistry::{ArtifactRegistryApi, ArtifactRegistryClient}; pub use gcp::cloud_sql::{CloudSqlApi, CloudSqlClient}; pub use gcp::cloudasset::{CloudAssetApi, CloudAssetClient}; diff --git a/crates/alien-infra/src/ai/aws.rs b/crates/alien-infra/src/ai/aws.rs index d3953c830..755d9ad07 100644 --- a/crates/alien-infra/src/ai/aws.rs +++ b/crates/alien-infra/src/ai/aws.rs @@ -1,22 +1,50 @@ use std::time::Duration; -use tracing::info; +use tracing::{info, warn}; use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; +use alien_aws_clients::bedrock::{ + CreateModelCustomizationJobRequest, ModelCustomizationJobStatus, S3DataConfig, +}; use alien_core::{ bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, - AwsBedrockAiHeartbeatData, HeartbeatBackend, Platform, ResourceHeartbeat, - ResourceHeartbeatData, ResourceOutputs, ResourceStatus, + AwsBedrockAiHeartbeatData, FinetuneMethod, HeartbeatBackend, Platform, ResourceHeartbeat, + ResourceHeartbeatData, ResourceOutputs, ResourceRef, ResourceStatus, Storage, }; use alien_error::{AlienError, Context, IntoAlienError}; use alien_macros::controller; use chrono::Utc; +/// Poll interval while a Bedrock model-customization job runs. Matches the +/// heartbeat cadence used elsewhere in this controller. +const TUNING_POLL_INTERVAL: Duration = Duration::from_secs(30); + +/// Bedrock `customizationType` for a supervised / preference / LoRA fine-tune. +/// Bedrock's public model-customization API exposes `FINE_TUNING`; the specific +/// technique (SFT vs DPO vs LoRA) is selected per base model via hyperparameters, +/// not a distinct customizationType, so all `FinetuneMethod` variants map here. +fn bedrock_customization_type(_method: FinetuneMethod) -> &'static str { + "FINE_TUNING" +} + #[controller] pub struct AwsAiController { /// AWS region where Bedrock is accessed. None until create_start runs. pub(crate) region: Option, + /// ARN of the submitted Bedrock model-customization job, set once + /// `SubmittingTuningJob` succeeds. Used as the poll identifier. Only ever set + /// for a finetune-enabled resource. + pub(crate) tuning_job_arn: Option, + /// The completed custom-model ARN Bedrock returns. This is the `upstream_id` + /// the OpenAI chat endpoint accepts for the tuned model, and is attached to + /// the binding via `with_tuned_model`. None until the job completes. + pub(crate) tuned_model_id: Option, + /// The gateway-facing served model id the tuned model is exposed under + /// (`spec.served_model_id_or_default(&config.id)`). Captured when the job + /// completes so `get_binding_params` (which has no `ctx`) can pair it with + /// `tuned_model_id`. None for a pure inference gateway. + pub(crate) served_id: Option, } #[controller] @@ -65,12 +93,197 @@ impl AwsAiController { info!(ai=%config.id, "Successfully applied resource-scoped permissions"); + // A pure inference gateway is Ready as soon as permissions are applied — + // exactly as before. A finetune-enabled resource additionally submits and + // polls a Bedrock model-customization job before serving. + let next_state = if config.finetune.is_some() { + info!(ai=%config.id, "Finetune requested; submitting Bedrock model-customization job"); + SubmittingTuningJob + } else { + Ready + }; + Ok(HandlerAction::Continue { - state: Ready, + state: next_state, suggested_delay: None, }) } + // ─────────────── FINE-TUNING FLOW ────────────────────────── + // Only entered when the resource declares a `finetune` spec. + + #[handler( + state = SubmittingTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn submitting_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let config = ctx.desired_resource_config::()?; + let spec = config.finetune.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "SubmittingTuningJob reached without a finetune spec".to_string(), + }) + })?; + let aws_config = ctx.get_aws_config()?; + + // Resolve the training bucket from the dependency's real state rather than + // re-deriving the name. The Ai resource declares its training-data Storage + // as a dependency (see `Ai::get_dependencies`), and the AWS storage + // controller stores the actual (prefixed) bucket name in `bucket_name`. + // Reading it here keeps this controller correct even if the storage naming + // scheme changes. + let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); + let storage_state = ctx + .require_dependency::(&training_ref)?; + let bucket = storage_state.bucket_name.clone().ok_or_else(|| { + AlienError::new(ErrorData::DependencyNotReady { + resource_id: config.id.clone(), + dependency_id: spec.training_data.clone(), + }) + })?; + + // Bedrock assumes an IAM role to read the training data and write output. + // The Ai resource carries no dedicated service account, so derive the + // execution role ARN deterministically from the deployment account id and + // the stack's resource-role naming convention (`{prefix}-{id}`, matching + // `service_account::aws::get_aws_role_name`). The `ai/finetune` permission + // set grants this role the Bedrock job + S3 read actions it needs. + let role_arn = format!( + "arn:aws:iam::{}:role/{}-{}", + aws_config.account_id, ctx.resource_prefix, config.id + ); + + let training_key = &spec.training_key; + let training_uri = format!("s3://{}/{}", bucket, training_key); + let output_uri = format!("s3://{}/alien-finetune-output/{}/", bucket, config.id); + + // Bedrock job + custom-model names must match `([0-9a-zA-Z][_-]?){1,63}`; + // the resource id is already constrained to `[A-Za-z0-9-_]{1,64}`. + let job_name = format!("{}-{}", ctx.resource_prefix, config.id); + let custom_model_name = format!("{}-{}", ctx.resource_prefix, config.id); + + let request = CreateModelCustomizationJobRequest::builder() + .job_name(job_name) + .custom_model_name(custom_model_name) + .role_arn(role_arn) + .base_model_identifier(spec.base_model.clone()) + .customization_type(bedrock_customization_type(spec.method).to_string()) + .training_data_config(S3DataConfig::builder().s3_uri(training_uri).build()) + .output_data_config(S3DataConfig::builder().s3_uri(output_uri).build()) + .build(); + + let client = ctx.service_provider.get_aws_bedrock_client(aws_config).await?; + let response = client + .create_model_customization_job(&request) + .await + .context(ErrorData::CloudPlatformError { + message: format!( + "Failed to submit Bedrock model-customization job for AI '{}'", + config.id + ), + resource_id: Some(config.id.clone()), + })?; + + info!(ai=%config.id, job_arn=%response.job_arn, "Submitted Bedrock model-customization job"); + self.tuning_job_arn = Some(response.job_arn); + + Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: Some(TUNING_POLL_INTERVAL), + }) + } + + #[handler( + state = WaitingForTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn waiting_for_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let config = ctx.desired_resource_config::()?; + let aws_config = ctx.get_aws_config()?; + + let job_arn = self.tuning_job_arn.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "Tuning job ARN not set in state".to_string(), + }) + })?; + + let client = ctx.service_provider.get_aws_bedrock_client(aws_config).await?; + let job = client + .get_model_customization_job(job_arn) + .await + .context(ErrorData::CloudPlatformError { + message: format!( + "Failed to poll Bedrock model-customization job '{}' for AI '{}'", + job_arn, config.id + ), + resource_id: Some(config.id.clone()), + })?; + + match job.status() { + ModelCustomizationJobStatus::Completed => { + // The custom-model ARN is the artifact the gateway forwards to. + let output_model_arn = job.output_model_arn.clone().ok_or_else(|| { + AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Bedrock job '{}' completed but returned no outputModelArn", + job_arn + ), + resource_id: Some(config.id.clone()), + }) + })?; + + let spec = config.finetune.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "WaitingForTuningJob completed without a finetune spec".to_string(), + }) + })?; + + info!(ai=%config.id, custom_model=%output_model_arn, "Bedrock model-customization job completed"); + self.tuned_model_id = Some(output_model_arn); + self.served_id = Some(spec.served_model_id_or_default(&config.id)); + + Ok(HandlerAction::Continue { + state: Ready, + suggested_delay: None, + }) + } + // Fail loud on any terminal failure, mirroring the Azure controller's + // WaitingForDeployments fail-fast style (don't poll a dead job forever). + status @ (ModelCustomizationJobStatus::Failed + | ModelCustomizationJobStatus::Stopped) => { + let reason = job + .failure_message + .clone() + .unwrap_or_else(|| "no failure message reported".to_string()); + Err(AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Bedrock model-customization job '{}' entered terminal state '{:?}': {}", + job_arn, status, reason + ), + resource_id: Some(config.id.clone()), + })) + } + // InProgress / Stopping / any unrecognized status: keep polling. + other => { + info!(ai=%config.id, ?other, "Bedrock model-customization job not yet complete, re-polling"); + Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: Some(TUNING_POLL_INTERVAL), + }) + } + } + } + // ─────────────── READY STATE ──────────────────────────────── // Loops as a heartbeat tick; Bedrock has no per-stack resource to poll. @@ -110,7 +323,9 @@ impl AwsAiController { } // ─────────────── UPDATE FLOW ────────────────────────────── - // Ai has no mutable fields -- update is a no-op that also recovers RefreshFailed. + // Ai has no mutable inference fields, and finetune base/training are immutable + // (enforced in `Ai::validate_update`), so update is a no-op that also recovers + // RefreshFailed. The tuned model id persists in controller state across updates. #[flow_entry(Update, from = [Ready, RefreshFailed])] #[handler( @@ -177,13 +392,320 @@ impl AwsAiController { Some(r) => r, None => return Ok(None), }; - Ok(Some( - serde_json::to_value(AiBinding::bedrock(region)) - .into_alien_error() - .context(ErrorData::ResourceStateSerializationFailed { - resource_id: "binding".to_string(), - message: "Failed to serialize AI binding parameters".to_string(), - })?, - )) + + let mut binding = AiBinding::bedrock(region); + + // Attach the tuned model only once the job has completed and produced both + // a custom-model id and its served id (both are set together in + // WaitingForTuningJob). Until then (and for a pure inference gateway) the + // base binding is emitted unchanged. + match (&self.tuned_model_id, &self.served_id) { + (Some(upstream_id), Some(served_id)) => { + binding = binding.with_tuned_model(served_id.clone(), upstream_id.clone()); + } + (Some(_), None) => { + // A tuned artifact without a served id is a controller-state + // inconsistency (they are set together). Emit the untuned binding + // rather than serve under an unknown id, and warn loudly. + warn!("Tuned model id present but served id missing; emitting untuned binding"); + } + _ => {} + } + + Ok(Some(serde_json::to_value(binding).into_alien_error().context( + ErrorData::ResourceStateSerializationFailed { + resource_id: "binding".to_string(), + message: "Failed to serialize AI binding parameters".to_string(), + }, + )?)) + } +} + +#[cfg(all(test, feature = "test-utils"))] +mod tests { + use super::*; + use crate::core::controller_test::SingleControllerExecutor; + use crate::core::ResourceController; + use crate::storage::AwsStorageController; + use crate::MockPlatformServiceProvider; + use alien_aws_clients::bedrock::{ + CreateModelCustomizationJobResponse, GetModelCustomizationJobResponse, MockBedrockApi, + }; + use alien_core::bindings::AiBinding; + use alien_core::{Ai, FinetuneMethod, FinetuneSpec, Platform, ResourceStatus, Storage}; + use std::sync::Arc; + + // The training-data storage the finetune resource depends on. Its AWS storage + // controller mock is wired with a real bucket name so the AI controller can + // read it via `require_dependency`. + const TRAINING_STORAGE_ID: &str = "training-set"; + + fn training_storage() -> Storage { + Storage::new(TRAINING_STORAGE_ID.to_string()).build() + } + + fn tuned_ai() -> Ai { + Ai::new("my-ai".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: TRAINING_STORAGE_ID.to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build() + } + + fn untuned_ai() -> Ai { + Ai::new("my-ai".to_string()).build() + } + + /// The completed custom-model ARN the gateway forwards tuned requests to. + const CUSTOM_MODEL_ARN: &str = + "arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345"; + const JOB_ARN: &str = + "arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345"; + + fn completed_job() -> GetModelCustomizationJobResponse { + serde_json::from_value(serde_json::json!({ + "status": "Completed", + "jobArn": JOB_ARN, + "outputModelArn": CUSTOM_MODEL_ARN, + "outputModelName": "test-my-ai", + })) + .expect("valid completed job json") + } + + fn in_progress_job() -> GetModelCustomizationJobResponse { + serde_json::from_value(serde_json::json!({ "status": "InProgress" })) + .expect("valid in-progress job json") + } + + fn failed_job() -> GetModelCustomizationJobResponse { + serde_json::from_value(serde_json::json!({ + "status": "Failed", + "failureMessage": "training data validation failed", + })) + .expect("valid failed job json") + } + + /// Build a mock Bedrock client that submits the job then returns the given + /// sequence of poll responses (one per `get_model_customization_job` call). + fn mock_bedrock(polls: Vec) -> Arc { + let mut mock = MockBedrockApi::new(); + mock.expect_create_model_customization_job().returning(|_| { + Ok(CreateModelCustomizationJobResponse { + job_arn: JOB_ARN.to_string(), + }) + }); + let responses = std::sync::Mutex::new(polls.into_iter()); + mock.expect_get_model_customization_job() + .returning(move |_| { + responses + .lock() + .unwrap() + .next() + .map(Ok) + .unwrap_or_else(|| Ok(completed_job())) + }); + Arc::new(mock) + } + + fn provider_with_bedrock(mock: Arc) -> Arc { + let mut provider = MockPlatformServiceProvider::new(); + provider + .expect_get_aws_bedrock_client() + .returning(move |_| Ok(mock.clone())); + Arc::new(provider) + } + + /// The real (prefixed) bucket name the storage dependency mock resolves to. + /// `AwsStorageController::mock_ready` prefixes with "test-stack". + fn training_bucket_name() -> String { + format!("test-stack-{}", TRAINING_STORAGE_ID) + } + + /// Read the controller's current binding via the ResourceController trait. + fn current_binding(executor: &SingleControllerExecutor) -> AiBinding { + let controller = executor + .internal_state::() + .expect("controller is AwsAiController"); + let value = controller + .get_binding_params() + .expect("binding params resolve") + .expect("binding present once region is set"); + serde_json::from_value(value).expect("binding deserializes") + } + + #[tokio::test] + async fn test_finetune_submit_poll_succeeds_and_binds_tuned_model() { + // Submit -> InProgress -> Completed -> Ready, then the binding must carry + // the tuned model with the right served + upstream ids. + let mock = mock_bedrock(vec![in_progress_job(), completed_job()]); + let provider = provider_with_bedrock(mock); + + let mut executor = SingleControllerExecutor::builder() + .resource(tuned_ai()) + .controller(AwsAiController::default()) + .platform(Platform::Aws) + .service_provider(provider) + .with_dependency( + training_storage(), + AwsStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor builds"); + + // Ready is not terminal (heartbeat loop), so step until Running. + for _ in 0..8 { + if executor.status() == ResourceStatus::Running { + break; + } + executor.step().await.expect("step should not error"); + } + assert_eq!( + executor.status(), + ResourceStatus::Running, + "finetune resource must reach Ready after the job completes" + ); + + let binding = current_binding(&executor); + let tuned = binding + .tuned_model() + .expect("completed finetune must attach a tuned model"); + assert_eq!( + tuned.served_id, "my-ai-tuned", + "served id must default to -tuned" + ); + assert_eq!( + tuned.upstream_id, CUSTOM_MODEL_ARN, + "upstream id must be the completed custom-model ARN" + ); + } + + #[tokio::test] + async fn test_finetune_job_failed_fails_fast_to_create_failed() { + // A terminal Failed status must route to CreateFailed, not poll forever. + let mock = mock_bedrock(vec![failed_job()]); + let provider = provider_with_bedrock(mock); + + let mut executor = SingleControllerExecutor::builder() + .resource(tuned_ai()) + .controller(AwsAiController::default()) + .platform(Platform::Aws) + .service_provider(provider) + .with_dependency( + training_storage(), + AwsStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor builds"); + + // A terminal Failed job must surface as a handler error (which the real + // executor routes to CreateFailed via `on_failure`), NOT an endless poll. + // The test harness's `step()` returns that error rather than applying the + // failure transition, so assert the error surfaces. Bounded so a + // poll-forever regression fails the test instead of hanging. + let mut surfaced_error = false; + for _ in 0..10 { + if executor.step().await.is_err() { + surfaced_error = true; + break; + } + } + assert!( + surfaced_error, + "a terminal Failed job must surface as an error, not a silent retry" + ); + } + + #[tokio::test] + async fn test_no_finetune_reaches_ready_with_untuned_binding() { + // Regression: an inference-only resource must never touch Bedrock tuning + // and must emit an untuned binding. + let mut provider = MockPlatformServiceProvider::new(); + provider.expect_get_aws_bedrock_client().never(); + let provider = Arc::new(provider); + + let mut executor = SingleControllerExecutor::builder() + .resource(untuned_ai()) + .controller(AwsAiController::default()) + .platform(Platform::Aws) + .service_provider(provider) + .with_test_dependencies() + .build() + .await + .expect("executor builds"); + + for _ in 0..6 { + if executor.status() == ResourceStatus::Running { + break; + } + executor.step().await.expect("step should not error"); + } + assert_eq!( + executor.status(), + ResourceStatus::Running, + "inference-only resource must reach Ready" + ); + + let binding = current_binding(&executor); + assert!( + binding.tuned_model().is_none(), + "inference-only binding must omit tuned_model" + ); + } + + #[tokio::test] + async fn test_finetune_uses_training_bucket_and_customization_type() { + // Assert the submitted job carries the dependency's real bucket in both the + // training and output S3 URIs, and FINE_TUNING as the customization type. + let bucket = training_bucket_name(); + let expected_training_uri = format!("s3://{}/training.jsonl", bucket); + let expected_output_prefix = format!("s3://{}/alien-finetune-output/my-ai/", bucket); + + let mut mock = MockBedrockApi::new(); + mock.expect_create_model_customization_job() + .withf(move |req| { + req.customization_type == "FINE_TUNING" + && req.base_model_identifier == "amazon.nova-lite-v1:0" + && req.training_data_config.s3_uri == expected_training_uri + && req.output_data_config.s3_uri == expected_output_prefix + && req.role_arn == "arn:aws:iam::123456789012:role/test-my-ai" + }) + .returning(|_| { + Ok(CreateModelCustomizationJobResponse { + job_arn: JOB_ARN.to_string(), + }) + }); + mock.expect_get_model_customization_job() + .returning(|_| Ok(completed_job())); + + let provider = provider_with_bedrock(Arc::new(mock)); + + let mut executor = SingleControllerExecutor::builder() + .resource(tuned_ai()) + .controller(AwsAiController::default()) + .platform(Platform::Aws) + .service_provider(provider) + .with_dependency( + training_storage(), + AwsStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor builds"); + + // Reaching Running proves the withf predicate matched (a mismatch panics + // the mock, which surfaces as a step error and never reaches Running). + for _ in 0..8 { + if executor.status() == ResourceStatus::Running { + break; + } + executor.step().await.expect("step should not error"); + } + assert_eq!(executor.status(), ResourceStatus::Running); } } diff --git a/crates/alien-infra/src/ai/aws_import.rs b/crates/alien-infra/src/ai/aws_import.rs index ccc8e2269..e35b3cdf2 100644 --- a/crates/alien-infra/src/ai/aws_import.rs +++ b/crates/alien-infra/src/ai/aws_import.rs @@ -24,6 +24,9 @@ impl ResourceImporter for AwsAiImporter { let controller = AwsAiController { state: AwsAiState::Ready, region: Some(data.region), + tuning_job_arn: None, + tuned_model_id: None, + served_id: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-infra/src/ai/azure.rs b/crates/alien-infra/src/ai/azure.rs index 2859cb79b..7f3667a4c 100644 --- a/crates/alien-infra/src/ai/azure.rs +++ b/crates/alien-infra/src/ai/azure.rs @@ -3,12 +3,14 @@ use tracing::info; use crate::azure_utils::get_resource_group_name; use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; +use crate::infra_requirements::azure_utils::get_storage_account_name; use alien_azure_clients::azure::cognitive_services::{ CognitiveServicesAccountCreateParameters, CognitiveServicesAccountCreateProperties, CognitiveServicesDeploymentCreateParameters, CognitiveServicesDeploymentCreateProperties, CognitiveServicesDeploymentModel, CognitiveServicesDeploymentSku, CognitiveServicesSku, }; use alien_azure_clients::long_running_operation::OperationResult; +use alien_core::FinetuneSpec; use alien_client_core::ErrorData as CloudClientErrorData; use alien_core::{ bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, @@ -52,6 +54,15 @@ pub struct AzureAiController { pub(crate) resource_group: Option, /// The Azure region where the account is created. pub(crate) location: Option, + /// The Foundry fine-tuning job id, set once the job is submitted. Only + /// populated when the resource declares a `finetune` spec. + #[serde(default)] + pub(crate) tuning_job_id: Option, + /// The deployment name serving the tuned model, set once the tuned model is + /// deployed. This is the `upstream_id` the gateway forwards to over the + /// OpenAI chat path. Absent for a pure inference gateway. + #[serde(default)] + pub(crate) tuned_deployment_name: Option, } #[controller] @@ -457,12 +468,265 @@ impl AzureAiController { } info!(account_name = %account_name, "All Azure model deployments provisioned"); + + // A pure inference gateway is Ready here. When the resource declares a + // finetune spec, tune the base model and serve the result before Ready. + if config.finetune.is_some() { + return Ok(HandlerAction::Continue { + state: SubmittingTuningJob, + suggested_delay: None, + }); + } + Ok(HandlerAction::Continue { state: Ready, suggested_delay: None, }) } + // ─────────────── FINE-TUNING FLOW ─────────────────────────── + + #[handler( + state = SubmittingTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn submitting_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let azure_config = ctx.get_azure_config()?; + let config = ctx.desired_resource_config::()?; + let spec = require_finetune_spec(&config)?; + + let endpoint = self.require_endpoint(&config.id)?.to_string(); + + // Foundry imports the training dataset directly from the customer's Blob + // container. We derive the blob URL deterministically the same way the + // Azure Storage controller names the container + // (`{prefix}-{training_data}`, lowercased, `_`->`-`) under the shared + // default storage account, rather than requiring a public storage state + // type. See `crates/alien-infra/src/storage/azure.rs`. + // + // GOTCHA: Foundry's Blob import path requires the AIServices account to + // allow *public network access*. A private-endpoint-only account rejects + // the blob URL at job-submit time; the job then surfaces as a terminal + // failure in WaitingForTuningJob (fail-fast), not a silent hang. + let storage_account = get_storage_account_name(ctx.state)?; + let training_file = training_blob_url( + &storage_account, + ctx.resource_prefix, + &spec.training_data, + &spec.training_key, + ); + + info!( + id = %config.id, + endpoint = %endpoint, + base_model = %spec.base_model, + training_file = %training_file, + "Submitting Azure Foundry fine-tuning job" + ); + + let finetuning_client = ctx + .service_provider + .get_azure_foundry_finetuning_client(azure_config)?; + + let job = finetuning_client + .create_fine_tuning_job(&endpoint, &spec.base_model, &training_file) + .await + .context(ErrorData::CloudPlatformError { + message: format!( + "Failed to submit Azure Foundry fine-tuning job for base model '{}'", + spec.base_model + ), + resource_id: Some(config.id.clone()), + })?; + + info!(id = %config.id, job_id = %job.id, status = %job.status, "Fine-tuning job submitted"); + self.tuning_job_id = Some(job.id); + + Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: Some(std::time::Duration::from_secs(30)), + }) + } + + #[handler( + state = WaitingForTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn waiting_for_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let azure_config = ctx.get_azure_config()?; + let config = ctx.desired_resource_config::()?; + + let endpoint = self.require_endpoint(&config.id)?.to_string(); + let job_id = self.tuning_job_id.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Fine-tuning job id not set in state".to_string(), + resource_id: Some(config.id.clone()), + }) + })?; + + info!(id = %config.id, job_id = %job_id, "Polling Azure Foundry fine-tuning job"); + + let finetuning_client = ctx + .service_provider + .get_azure_foundry_finetuning_client(azure_config)?; + + let job = finetuning_client + .get_fine_tuning_job(&endpoint, job_id) + .await + .context(ErrorData::CloudPlatformError { + message: format!("Failed to poll Azure Foundry fine-tuning job '{}'", job_id), + resource_id: Some(config.id.clone()), + })?; + + // Fail fast on a terminal failure rather than polling a dead job forever, + // mirroring the WaitingForDeployments style. + if job.is_terminal_failure() { + return Err(AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Azure Foundry fine-tuning job '{}' entered terminal state '{}'", + job.id, job.status + ), + resource_id: Some(config.id.clone()), + })); + } + + if !job.is_succeeded() { + info!(id = %config.id, job_id = %job.id, status = %job.status, "Fine-tuning job not yet complete, retrying"); + return Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: Some(std::time::Duration::from_secs(30)), + }); + } + + // Succeeded: capture the tuned model name to deploy. Fail fast if the + // provider reports success without one — we cannot serve a missing model. + let fine_tuned_model = job.fine_tuned_model.ok_or_else(|| { + AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Azure Foundry fine-tuning job '{}' succeeded but returned no fine_tuned_model", + job.id + ), + resource_id: Some(config.id.clone()), + }) + })?; + + info!(id = %config.id, fine_tuned_model = %fine_tuned_model, "Fine-tuning job succeeded"); + + // Deploy the tuned model under the served id, so the OpenAI chat path + // accepts it as a deployment target. + let spec = require_finetune_spec(&config)?; + let served_id = spec.served_model_id_or_default(&config.id); + self.deploy_tuned_model(ctx, &config.id, &served_id, &fine_tuned_model) + .await?; + self.tuned_deployment_name = Some(served_id); + + Ok(HandlerAction::Continue { + state: WaitingForTunedDeployment, + suggested_delay: Some(std::time::Duration::from_secs(10)), + }) + } + + #[handler( + state = WaitingForTunedDeployment, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn waiting_for_tuned_deployment( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let azure_config = ctx.get_azure_config()?; + let config = ctx.desired_resource_config::()?; + + let account_name = self.require_account_name(&config.id)?.to_string(); + let resource_group_name = self.require_resource_group(&config.id)?.to_string(); + let deployment_name = self.tuned_deployment_name.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Tuned deployment name not set in state".to_string(), + resource_id: Some(config.id.clone()), + }) + })?; + + let cognitive_client = ctx + .service_provider + .get_azure_cognitive_services_client(azure_config)?; + + match cognitive_client + .get_deployment(&resource_group_name, &account_name, deployment_name) + .await + { + Ok(deployment) => { + let provisioning_state = deployment + .properties + .as_ref() + .and_then(|p| p.provisioning_state.as_deref()) + .unwrap_or(""); + + if provisioning_state.eq_ignore_ascii_case("Failed") + || provisioning_state.eq_ignore_ascii_case("Canceled") + { + return Err(AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Azure tuned-model deployment '{}' entered terminal state '{}'", + deployment_name, provisioning_state + ), + resource_id: Some(config.id.clone()), + })); + } + + if !provisioning_state.eq_ignore_ascii_case("Succeeded") { + info!( + account_name = %account_name, + deployment = %deployment_name, + provisioning_state = %provisioning_state, + "Tuned-model deployment not yet ready, retrying" + ); + return Ok(HandlerAction::Continue { + state: WaitingForTunedDeployment, + suggested_delay: Some(std::time::Duration::from_secs(10)), + }); + } + + info!(account_name = %account_name, deployment = %deployment_name, "Tuned-model deployment provisioned"); + Ok(HandlerAction::Continue { + state: Ready, + suggested_delay: None, + }) + } + Err(e) + if matches!( + &e.error, + Some(CloudClientErrorData::RemoteResourceNotFound { .. }) + ) => + { + info!( + account_name = %account_name, + deployment = %deployment_name, + "Tuned-model deployment not yet visible, retrying" + ); + Ok(HandlerAction::Continue { + state: WaitingForTunedDeployment, + suggested_delay: Some(std::time::Duration::from_secs(10)), + }) + } + Err(e) => Err(e.context(ErrorData::CloudPlatformError { + message: format!( + "Failed to poll Azure tuned-model deployment '{}'", + deployment_name + ), + resource_id: Some(config.id.clone()), + })), + } + } + // ─────────────── READY STATE ──────────────────────────────── #[handler( @@ -705,14 +969,24 @@ impl AzureAiController { Some(a) => a.clone(), None => return Ok(None), }; - Ok(Some( - serde_json::to_value(AiBinding::foundry(endpoint, account_name)) - .into_alien_error() - .context(ErrorData::ResourceStateSerializationFailed { - resource_id: "binding".to_string(), - message: "Failed to serialize Azure AI binding parameters".to_string(), - })?, - )) + + let mut binding = AiBinding::foundry(endpoint, account_name); + + // Attach the tuned model only once its deployment exists, so a pure + // inference gateway keeps the unchanged untuned wire shape. We deploy the + // tuned model under a deployment named `served_id`, and Foundry's OpenAI + // chat path takes that same deployment name as the request-body `model`. + // So served_id and upstream_id are the one stored deployment name. + if let Some(deployment_name) = &self.tuned_deployment_name { + binding = binding.with_tuned_model(deployment_name.clone(), deployment_name.clone()); + } + + Ok(Some(serde_json::to_value(binding).into_alien_error().context( + ErrorData::ResourceStateSerializationFailed { + resource_id: "binding".to_string(), + message: "Failed to serialize Azure AI binding parameters".to_string(), + }, + )?)) } } @@ -722,6 +996,95 @@ impl AzureAiController { self.endpoint = None; self.resource_group = None; self.location = None; + self.tuning_job_id = None; + self.tuned_deployment_name = None; + } + + /// Returns the provisioned account endpoint, or a config error if unset. + fn require_endpoint(&self, resource_id: &str) -> Result<&str> { + self.endpoint.as_deref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Endpoint not set in state".to_string(), + resource_id: Some(resource_id.to_string()), + }) + }) + } + + /// Returns the AIServices account name, or a config error if unset. + fn require_account_name(&self, resource_id: &str) -> Result<&str> { + self.account_name.as_deref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Account name not set in state".to_string(), + resource_id: Some(resource_id.to_string()), + }) + }) + } + + /// Returns the resource group name, or a config error if unset. + fn require_resource_group(&self, resource_id: &str) -> Result<&str> { + self.resource_group.as_deref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Resource group not set in state".to_string(), + resource_id: Some(resource_id.to_string()), + }) + }) + } + + /// Deploys the tuned model under `deployment_name`, reusing the same + /// control-plane `create_deployment` path as the base catalog so the OpenAI + /// chat endpoint accepts the deployment name as a `model`. + async fn deploy_tuned_model( + &self, + ctx: &ResourceControllerContext<'_>, + resource_id: &str, + deployment_name: &str, + fine_tuned_model: &str, + ) -> Result<()> { + let azure_config = ctx.get_azure_config()?; + let account_name = self.require_account_name(resource_id)?.to_string(); + let resource_group_name = self.require_resource_group(resource_id)?.to_string(); + + info!( + account_name = %account_name, + deployment = %deployment_name, + fine_tuned_model = %fine_tuned_model, + "Deploying tuned model" + ); + + let cognitive_client = ctx + .service_provider + .get_azure_cognitive_services_client(azure_config)?; + + // The tuned model's format is OpenAI and its "version" the fine-tuned + // model id Foundry returned. Capacity mirrors the base deployments. + let parameters = CognitiveServicesDeploymentCreateParameters { + sku: CognitiveServicesDeploymentSku { + name: "GlobalStandard".to_string(), + capacity: DEFAULT_DEPLOYMENT_CAPACITY, + }, + properties: CognitiveServicesDeploymentCreateProperties { + model: CognitiveServicesDeploymentModel { + format: "OpenAI".to_string(), + name: fine_tuned_model.to_string(), + version: "1".to_string(), + }, + }, + }; + + cognitive_client + .create_deployment( + &resource_group_name, + &account_name, + deployment_name, + ¶meters, + ) + .await + .context(ErrorData::CloudPlatformError { + message: format!("Failed to deploy tuned model as '{}'", deployment_name), + resource_id: Some(resource_id.to_string()), + })?; + + Ok(()) } /// Creates a controller in the ready state for testing purposes. @@ -733,6 +1096,8 @@ impl AzureAiController { endpoint: Some(endpoint.to_string()), resource_group: Some("mock-rg".to_string()), location: Some("eastus".to_string()), + tuning_job_id: None, + tuned_deployment_name: None, _internal_stay_count: None, } } @@ -747,23 +1112,79 @@ impl AzureAiController { endpoint: Some(endpoint.to_string()), resource_group: Some("mock-rg".to_string()), location: Some("eastus".to_string()), + tuning_job_id: None, + tuned_deployment_name: None, _internal_stay_count: None, } } + + /// Creates a controller poised to submit a fine-tuning job (account and base + /// deployments already provisioned), for testing the tuning states. + #[cfg(feature = "test-utils")] + pub fn mock_submitting_tuning(account_name: &str, endpoint: &str) -> Self { + Self { + state: AzureAiState::SubmittingTuningJob, + account_name: Some(account_name.to_string()), + endpoint: Some(endpoint.to_string()), + resource_group: Some("mock-rg".to_string()), + location: Some("eastus".to_string()), + tuning_job_id: None, + tuned_deployment_name: None, + _internal_stay_count: None, + } + } +} + +/// Extracts the finetune spec from an `Ai` config, erroring if absent. Called +/// only from tuning states, which are unreachable without a spec. +fn require_finetune_spec(config: &Ai) -> Result<&FinetuneSpec> { + config.finetune.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + message: "Reached a fine-tuning state without a finetune spec".to_string(), + resource_id: Some(config.id.clone()), + }) + }) +} + +/// Builds the Blob URL Foundry imports the training dataset from. +/// +/// Derived deterministically to match the Azure Storage controller's container +/// naming (`{prefix}-{name}` lowercased with `_`->`-`) under the shared default +/// storage account, so we don't need a public storage state type. Kept as a +/// pure function so it is unit-testable. +fn training_blob_url( + storage_account: &str, + resource_prefix: &str, + training_data: &str, + training_key: &str, +) -> String { + let container = format!("{}-{}", resource_prefix, training_data) + .to_lowercase() + .replace('_', "-"); + format!( + "https://{}.blob.core.windows.net/{}/{}", + storage_account, container, training_key + ) } #[cfg(all(test, feature = "test-utils"))] mod tests { use super::*; use crate::core::controller_test::SingleControllerExecutor; + use crate::core::ResourceController; use crate::MockPlatformServiceProvider; use alien_azure_clients::azure::cognitive_services::{ CognitiveServicesDeployment, CognitiveServicesDeploymentProperties, MockCognitiveServicesAccountsApi, }; + use alien_azure_clients::azure::openai_finetuning::{FineTuningJob, MockFoundryFineTuningApi}; use alien_azure_clients::long_running_operation::OperationResult; use alien_client_core::ErrorData as CloudClientErrorData; - use alien_core::{Ai, AiOutputs, Platform, ResourceStatus}; + use crate::storage::AzureStorageController; + use alien_core::bindings::AiBinding; + use alien_core::{ + Ai, AiOutputs, FinetuneMethod, FinetuneSpec, Platform, ResourceStatus, Storage, + }; use alien_error::AlienError; use std::sync::Arc; @@ -771,6 +1192,62 @@ mod tests { Ai::new("my-ai".to_string()).build() } + /// The training-data storage id the tuned resource depends on. Matches the + /// Azure test storage account dependency wired by `with_test_dependencies`. + const TRAINING_STORAGE_ID: &str = "training-set"; + + fn tuned_ai() -> Ai { + Ai::new("my-ai".to_string()) + .finetune(FinetuneSpec { + base_model: "gpt-4o-mini".to_string(), + training_data: TRAINING_STORAGE_ID.to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build() + } + + /// The training-data Storage the tuned resource declares as a dependency + /// (via `Ai::get_dependencies`). Must be wired Ready or the executor won't + /// run the AI controller. + fn training_storage() -> Storage { + Storage::new(TRAINING_STORAGE_ID.to_string()).build() + } + + /// A fine-tuning job in the given status, carrying `fine_tuned_model` only + /// when succeeded. + fn job(status: &str, fine_tuned_model: Option<&str>) -> FineTuningJob { + FineTuningJob { + id: "ftjob-abc".to_string(), + status: status.to_string(), + fine_tuned_model: fine_tuned_model.map(|s| s.to_string()), + } + } + + /// A cognitive-services mock that succeeds base + tuned deployments (create + + /// get both report Succeeded). + fn mock_cognitive_all_succeed() -> MockCognitiveServicesAccountsApi { + let mut mock = MockCognitiveServicesAccountsApi::new(); + mock.expect_create_deployment() + .returning(|_, _, _, _| Ok(OperationResult::Completed(succeeded_deployment()))); + mock.expect_get_deployment() + .returning(|_, _, _| Ok(succeeded_deployment())); + mock + } + + /// Read the controller's current binding via the ResourceController trait. + fn current_binding(executor: &SingleControllerExecutor) -> AiBinding { + let controller = executor + .internal_state::() + .expect("controller is AzureAiController"); + let value = controller + .get_binding_params() + .expect("binding params resolve") + .expect("binding present once endpoint + account are set"); + serde_json::from_value(value).expect("binding deserializes") + } + /// A deployment whose model + provisioning state are both reported succeeded. fn succeeded_deployment() -> CognitiveServicesDeployment { CognitiveServicesDeployment { @@ -995,4 +1472,195 @@ mod tests { assert_eq!(executor.status(), ResourceStatus::Deleted); } + + // ─────────────── FINE-TUNING TESTS ────────────────────────── + + #[test] + fn training_blob_url_matches_storage_container_naming() { + // Must match the Azure Storage controller's container naming so Foundry + // imports from the container the storage controller actually created. + let url = training_blob_url("myacct", "test", "training_set", "training.jsonl"); + assert_eq!( + url, + "https://myacct.blob.core.windows.net/test-training-set/training.jsonl", + "underscores must become hyphens and the prefix must be applied, lowercased" + ); + } + + #[tokio::test] + async fn test_finetune_submit_poll_succeeds_and_binds_tuned_model() { + // Full flow from the base deployments (mock_deploying) through submit -> + // pending -> succeeded -> deploy tuned -> Ready, then assert the tuned + // binding. Drives the real WaitingForDeployments -> SubmittingTuningJob + // branch (config.finetune is Some). + // One shared finetuning mock (cloned per getter call) so its poll + // sequence persists across steps: first `running`, then `succeeded`. + let mut finetuning_mock = MockFoundryFineTuningApi::new(); + finetuning_mock + .expect_create_fine_tuning_job() + .returning(|_, _, _| Ok(job("pending", None))); + let polls = std::sync::Mutex::new( + vec![ + job("running", None), + job("succeeded", Some("gpt-4o-mini.ft-abc")), + ] + .into_iter(), + ); + finetuning_mock + .expect_get_fine_tuning_job() + .returning(move |_, _| { + Ok(polls + .lock() + .unwrap() + .next() + .unwrap_or_else(|| job("succeeded", Some("gpt-4o-mini.ft-abc")))) + }); + let finetuning_mock = Arc::new(finetuning_mock); + + let mut mock_provider = MockPlatformServiceProvider::new(); + mock_provider + .expect_get_azure_cognitive_services_client() + .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); + mock_provider + .expect_get_azure_foundry_finetuning_client() + .returning(move |_| Ok(finetuning_mock.clone())); + + let mut executor = SingleControllerExecutor::builder() + .resource(tuned_ai()) + .controller(AzureAiController::mock_deploying( + "default-storage-account", + "https://my-ai.cognitiveservices.azure.com/", + )) + .platform(Platform::Azure) + .service_provider(Arc::new(mock_provider)) + .with_test_dependencies() + .with_dependency( + training_storage(), + AzureStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor builds"); + + // Ready is a heartbeat loop (not terminal), so step until Running. + for _ in 0..12 { + if executor.status() == ResourceStatus::Running { + break; + } + executor.step().await.expect("step should not error"); + } + assert_eq!( + executor.status(), + ResourceStatus::Running, + "finetune resource must reach Ready after the job completes and the tuned model deploys" + ); + + let binding = current_binding(&executor); + let tuned = binding + .tuned_model() + .expect("completed finetune must attach a tuned model"); + assert_eq!( + tuned.served_id, "my-ai-tuned", + "served id must default to -tuned" + ); + assert_eq!( + tuned.upstream_id, "my-ai-tuned", + "upstream id is the tuned deployment name, which equals served_id" + ); + } + + #[tokio::test] + async fn test_finetune_job_failed_fails_fast_to_create_failed() { + // A terminal Failed job status must route to CreateFailed, not poll forever. + let mut mock_provider = MockPlatformServiceProvider::new(); + mock_provider + .expect_get_azure_cognitive_services_client() + .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); + mock_provider + .expect_get_azure_foundry_finetuning_client() + .returning(|_| { + let mut mock = MockFoundryFineTuningApi::new(); + mock.expect_create_fine_tuning_job() + .returning(|_, _, _| Ok(job("pending", None))); + mock.expect_get_fine_tuning_job() + .returning(|_, _| Ok(job("failed", None))); + Ok(Arc::new(mock)) + }); + + let mut executor = SingleControllerExecutor::builder() + .resource(tuned_ai()) + .controller(AzureAiController::mock_submitting_tuning( + "default-storage-account", + "https://my-ai.cognitiveservices.azure.com/", + )) + .platform(Platform::Azure) + .service_provider(Arc::new(mock_provider)) + .with_test_dependencies() + .with_dependency( + training_storage(), + AzureStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor builds"); + + // A terminal Failed job must surface as a handler error (which the + // executor's on_failure routing turns into CreateFailed), not an endless + // poll. Bounded so a poll-forever regression fails instead of hanging. + let mut surfaced_error = false; + for _ in 0..8 { + if executor.step().await.is_err() { + surfaced_error = true; + break; + } + } + assert!( + surfaced_error, + "a terminal Failed job must surface as an error, not a silent retry" + ); + } + + #[tokio::test] + async fn test_no_finetune_reaches_ready_with_untuned_binding() { + // Regression: an inference-only resource must never call the fine-tuning + // client and must emit an untuned binding. + let mut mock_provider = MockPlatformServiceProvider::new(); + mock_provider + .expect_get_azure_cognitive_services_client() + .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); + mock_provider + .expect_get_azure_foundry_finetuning_client() + .never(); + + let mut executor = SingleControllerExecutor::builder() + .resource(basic_ai()) + .controller(AzureAiController::mock_deploying( + "my-ai-account", + "https://my-ai.cognitiveservices.azure.com/", + )) + .platform(Platform::Azure) + .service_provider(Arc::new(mock_provider)) + .with_test_dependencies() + .build() + .await + .expect("executor builds"); + + for _ in 0..6 { + if executor.status() == ResourceStatus::Running { + break; + } + executor.step().await.expect("step should not error"); + } + assert_eq!( + executor.status(), + ResourceStatus::Running, + "inference-only resource must reach Ready" + ); + + let binding = current_binding(&executor); + assert!( + binding.tuned_model().is_none(), + "inference-only binding must omit tuned_model" + ); + } } diff --git a/crates/alien-infra/src/ai/azure_import.rs b/crates/alien-infra/src/ai/azure_import.rs index 1c625dc5d..f9356287e 100644 --- a/crates/alien-infra/src/ai/azure_import.rs +++ b/crates/alien-infra/src/ai/azure_import.rs @@ -28,6 +28,9 @@ impl ResourceImporter for AzureAiImporter { endpoint: Some(data.endpoint), resource_group: Some(data.resource_group), location: Some(data.location), + // An imported resource has no in-flight tuning job or tuned deployment. + tuning_job_id: None, + tuned_deployment_name: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-infra/src/ai/gcp.rs b/crates/alien-infra/src/ai/gcp.rs index 278e451ea..775dc8bb9 100644 --- a/crates/alien-infra/src/ai/gcp.rs +++ b/crates/alien-infra/src/ai/gcp.rs @@ -4,23 +4,45 @@ use tracing::info; use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; +use crate::storage::GcpStorageController; use alien_core::{ - bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, + bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, FinetuneMethod, GcpVertexAiHeartbeatData, HeartbeatBackend, Platform, ResourceHeartbeat, ResourceHeartbeatData, - ResourceOutputs, ResourceStatus, + ResourceOutputs, ResourceRef, ResourceStatus, Storage, }; use alien_error::{AlienError, Context, IntoAlienError}; +use alien_gcp_clients::aiplatform::{ + CreateTuningJobRequest, JobState, SupervisedTuningSpec, +}; use alien_gcp_clients::iam::IamPolicy; use alien_gcp_clients::resource_manager::GetPolicyOptions; use alien_macros::controller; use chrono::Utc; +/// How often to re-poll a still-running Vertex tuning job. +const TUNING_POLL_INTERVAL: Duration = Duration::from_secs(30); + #[controller] pub struct GcpAiController { /// GCP project ID. None until create_start runs. pub(crate) project: Option, /// GCP region (location) for the Vertex AI endpoint. None until create_start runs. pub(crate) location: Option, + /// Full resource name of the submitted Vertex tuning job + /// (`projects/{p}/locations/{l}/tuningJobs/{id}`). Set once + /// `SubmittingTuningJob` runs; used by `WaitingForTuningJob` to poll. + /// `None` for a pure-inference gateway (no `finetune` spec). + pub(crate) tuning_job_name: Option, + /// The tuned model's upstream artifact id (the serving endpoint / model + /// resource name) the Vertex OpenAI-compat chat path routes to. Set only + /// once the tuning job reaches `JOB_STATE_SUCCEEDED`. Attached to the + /// binding via `.with_tuned_model(..)` when present. + pub(crate) tuned_model_upstream_id: Option, + /// The public model id the gateway serves the tuned model under + /// (`spec.served_model_id_or_default(&config.id)`). Captured alongside the + /// upstream id at success so `get_binding_params` (which has no ctx) can + /// build the tuned binding. + pub(crate) tuned_model_served_id: Option, } #[controller] @@ -187,12 +209,218 @@ impl GcpAiController { ) .await?; + // A pure-inference gateway is ready as soon as permissions are applied. + // A finetune spec first submits and awaits a Vertex tuning job. + if config.finetune.is_none() { + return Ok(HandlerAction::Continue { + state: Ready, + suggested_delay: None, + }); + } + + info!(id = %config.id, "Finetune requested; submitting Vertex tuning job"); Ok(HandlerAction::Continue { - state: Ready, + state: SubmittingTuningJob, suggested_delay: None, }) } + // ─────────────── FINETUNE FLOW ───────────────────────────── + // Only reached when the Ai declares a `finetune` spec. Submits a Vertex + // supervised tuning job reading the customer's training data from GCS, then + // polls it to completion before serving the tuned model through the gateway. + + #[handler( + state = SubmittingTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn submitting_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let gcp_config = ctx.get_gcp_config()?; + let config = ctx.desired_resource_config::()?; + let spec = config.finetune.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "SubmittingTuningJob reached without a finetune spec".to_string(), + }) + })?; + + // Vertex exposes only supervised tuning for Gemini; there is no + // user-selectable LoRA/QLoRA or DPO method. Map Sft -> supervised tuning + // and reject the others loudly rather than silently mis-tuning. + match spec.method { + FinetuneMethod::Sft => {} + other => { + return Err(AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: format!( + "Vertex AI supports only supervised fine-tuning (sft); method {other:?} is not available on Vertex" + ), + })); + } + } + + // Resolve the training data's real GCS bucket from the dependency's + // controller state, rather than re-deriving the prefixed name here. The + // Ai resource declares the training Storage as a dependency + // (`Ai::get_dependencies`), and the GCS controller records the exact, + // prefixed bucket name it created in `bucket_name` — reading it keeps + // this in lockstep with the storage controller's naming. + let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); + let storage_state = ctx.require_dependency::(&training_ref)?; + let bucket_name = storage_state.bucket_name.ok_or_else(|| { + AlienError::new(ErrorData::DependencyNotReady { + resource_id: config.id.clone(), + dependency_id: spec.training_data.clone(), + }) + })?; + + let training_uri = format!("gs://{bucket_name}/{}", spec.training_key); + + let request = CreateTuningJobRequest::builder() + .base_model(spec.base_model.clone()) + .supervised_tuning_spec( + SupervisedTuningSpec::builder() + .training_dataset_uri(training_uri.clone()) + .build(), + ) + .tuned_model_display_name(spec.served_model_id_or_default(&config.id)) + .build(); + + info!( + id = %config.id, + base_model = %spec.base_model, + training_uri = %training_uri, + "Submitting Vertex supervised tuning job" + ); + + let client = ctx.service_provider.get_gcp_aiplatform_client(gcp_config)?; + let job = client + .create_tuning_job(request) + .await + .context(ErrorData::CloudPlatformError { + message: "Failed to submit Vertex tuning job".to_string(), + resource_id: Some(config.id.clone()), + })?; + + let job_name = job.name.ok_or_else(|| { + AlienError::new(ErrorData::CloudPlatformError { + message: "Vertex tuning job creation returned no resource name".to_string(), + resource_id: Some(config.id.clone()), + }) + })?; + + info!(id = %config.id, job = %job_name, "Vertex tuning job submitted"); + self.tuning_job_name = Some(job_name); + + // Poll once immediately; only *in-progress* re-polls wait a full interval. + Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: None, + }) + } + + #[handler( + state = WaitingForTuningJob, + on_failure = CreateFailed, + status = ResourceStatus::Provisioning, + )] + async fn waiting_for_tuning_job( + &mut self, + ctx: &ResourceControllerContext<'_>, + ) -> Result { + let gcp_config = ctx.get_gcp_config()?; + let config = ctx.desired_resource_config::()?; + let spec = config.finetune.as_ref().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "WaitingForTuningJob reached without a finetune spec".to_string(), + }) + })?; + let spec_served_id = spec.served_model_id_or_default(&config.id); + let job_name = self.tuning_job_name.clone().ok_or_else(|| { + AlienError::new(ErrorData::ResourceConfigInvalid { + resource_id: Some(config.id.clone()), + message: "WaitingForTuningJob reached without a submitted job name".to_string(), + }) + })?; + + let client = ctx.service_provider.get_gcp_aiplatform_client(gcp_config)?; + let job = client + .get_tuning_job(job_name.clone()) + .await + .context(ErrorData::CloudPlatformError { + message: format!("Failed to poll Vertex tuning job '{job_name}'"), + resource_id: Some(config.id.clone()), + })?; + + let state = job.state.unwrap_or(JobState::JobStateUnspecified); + + if state.is_in_progress() { + info!(id = %config.id, job = %job_name, ?state, "Vertex tuning job still running; re-polling"); + return Ok(HandlerAction::Continue { + state: WaitingForTuningJob, + suggested_delay: Some(TUNING_POLL_INTERVAL), + }); + } + + // Terminal failure states fail loud — a fine-tuning resource whose job + // failed must not silently degrade to serving the untuned base model. + if state.is_terminal_failure() { + let detail = job + .error + .map(|e| e.message) + .unwrap_or_else(|| "no error detail returned".to_string()); + return Err(AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Vertex tuning job '{job_name}' reached terminal state {state:?}: {detail}" + ), + resource_id: Some(config.id.clone()), + })); + } + + // Success: record the upstream artifact the gateway routes to. + if state == JobState::JobStateSucceeded { + let upstream_id = job + .tuned_model + .as_ref() + .and_then(|m| m.upstream_id()) + .ok_or_else(|| { + AlienError::new(ErrorData::CloudPlatformError { + message: format!( + "Vertex tuning job '{job_name}' succeeded but returned no tuned model endpoint/model" + ), + resource_id: Some(config.id.clone()), + }) + })? + .to_string(); + + info!( + id = %config.id, + job = %job_name, + upstream_id = %upstream_id, + "Vertex tuning job succeeded; tuned model ready" + ); + self.tuned_model_upstream_id = Some(upstream_id); + self.tuned_model_served_id = Some(spec_served_id); + + return Ok(HandlerAction::Continue { + state: Ready, + suggested_delay: None, + }); + } + + // Any other (unspecified/unknown) terminal state is unexpected — fail + // rather than loop forever or serve an untuned model. + Err(AlienError::new(ErrorData::CloudPlatformError { + message: format!("Vertex tuning job '{job_name}' in unexpected state {state:?}"), + resource_id: Some(config.id.clone()), + })) + } + // ─────────────── READY STATE ──────────────────────────────── // Loops as a heartbeat tick; Vertex AI has no per-stack resource to poll. @@ -323,13 +551,384 @@ impl GcpAiController { else { return Ok(None); }; - Ok(Some( - serde_json::to_value(AiBinding::vertex(project, location)) - .into_alien_error() - .context(ErrorData::ResourceStateSerializationFailed { - resource_id: "binding".to_string(), - message: "Failed to serialize AI binding parameters".to_string(), - })?, - )) + + let mut binding = AiBinding::vertex(project, location); + + // Attach the tuned model only once the tuning job has completed and we + // recorded both ids; otherwise the base (untuned) binding is returned + // unchanged, matching a pure-inference gateway. + if let (Some(served_id), Some(upstream_id)) = ( + self.tuned_model_served_id.as_ref(), + self.tuned_model_upstream_id.as_ref(), + ) { + binding = binding.with_tuned_model(served_id, upstream_id); + } + + Ok(Some(serde_json::to_value(binding).into_alien_error().context( + ErrorData::ResourceStateSerializationFailed { + resource_id: "binding".to_string(), + message: "Failed to serialize AI binding parameters".to_string(), + }, + )?)) + } +} + +#[cfg(test)] +mod tests { + //! GCP Vertex AI controller tests. + //! + //! These drive the finetune state machine end-to-end against a mocked + //! `AiPlatformApi`: submit -> pending -> succeeded -> Ready, asserting the + //! resulting binding carries the tuned model. They also cover the fail-fast + //! path (job FAILED -> ProvisionFailed) and the pure-inference regression + //! (no finetune -> Ready with an untuned binding). + + use std::sync::{Arc, Mutex}; + + use super::GcpAiController; + use alien_core::bindings::AiBinding; + use alien_core::{Ai, FinetuneMethod, FinetuneSpec, Platform, ResourceStatus, Storage}; + use alien_gcp_clients::aiplatform::{ + JobState, MockAiPlatformApi, TunedModelRef, TuningJob, + }; + use alien_gcp_clients::iam::IamPolicy; + use alien_gcp_clients::longrunning::Operation; + use alien_gcp_clients::resource_manager::MockResourceManagerApi; + use alien_gcp_clients::service_usage::MockServiceUsageApi; + + use crate::core::controller_test::SingleControllerExecutor; + use crate::core::{MockPlatformServiceProvider, PlatformServiceProvider, ResourceController}; + use crate::storage::GcpStorageController; + + const TRAINING_STORAGE_ID: &str = "training-set"; + const TUNED_ENDPOINT: &str = "projects/test-project-123/locations/us-central1/endpoints/9988"; + + // ─────────────── FIXTURES ────────────────────────────────── + + /// A pure-inference gateway (no finetune). + fn base_ai() -> Ai { + Ai::new("llm".to_string()).build() + } + + /// An Ai that fine-tunes a Gemini base model from the training storage. + fn finetune_ai() -> Ai { + Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "gemini-2.0-flash-001".to_string(), + training_data: TRAINING_STORAGE_ID.to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build() + } + + /// An Ai requesting a method Vertex does not expose (LoRA). + fn lora_ai() -> Ai { + Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "gemini-2.0-flash-001".to_string(), + training_data: TRAINING_STORAGE_ID.to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Lora, + }) + .build() + } + + fn training_storage() -> Storage { + Storage::new(TRAINING_STORAGE_ID.to_string()).build() + } + + // ─────────────── MOCK HELPERS ────────────────────────────── + + /// Resource-manager mock that satisfies the read-modify-write IAM step in + /// `applying_resource_permissions` (the AI resource grants no bindings, but + /// the closure always gets then sets the project policy). + fn iam_mock() -> Arc { + let mut mock = MockResourceManagerApi::new(); + mock.expect_get_project_iam_policy() + .returning(|_, _| Ok(IamPolicy::default())); + mock.expect_set_project_iam_policy() + .returning(|_, policy, _| Ok(policy)); + Arc::new(mock) + } + + fn service_usage_mock() -> Arc { + let mut mock = MockServiceUsageApi::new(); + mock.expect_enable_service() + .returning(|_| Ok(Operation::default())); + Arc::new(mock) + } + + /// Wires a service provider with the given aiplatform mock plus the + /// service-usage and resource-manager mocks the create flow needs. + fn provider_with(aiplatform: MockAiPlatformApi) -> Arc { + let aiplatform = Arc::new(aiplatform); + let mut provider = MockPlatformServiceProvider::new(); + provider + .expect_get_gcp_service_usage_client() + .returning({ + let m = service_usage_mock(); + move |_| Ok(m.clone()) + }); + provider + .expect_get_gcp_resource_manager_client() + .returning({ + let m = iam_mock(); + move |_| Ok(m.clone()) + }); + provider + .expect_get_gcp_aiplatform_client() + .returning(move |_| Ok(aiplatform.clone())); + Arc::new(provider) + } + + fn succeeded_job() -> TuningJob { + TuningJob::builder() + .name("projects/test-project-123/locations/us-central1/tuningJobs/42".to_string()) + .state(JobState::JobStateSucceeded) + .tuned_model( + TunedModelRef::builder() + .endpoint(TUNED_ENDPOINT.to_string()) + .build(), + ) + .build() + } + + // ─────────────── TESTS ───────────────────────────────────── + + /// Submit -> pending -> running -> succeeded -> Ready, and the binding + /// carries the tuned model with the right served/upstream ids. + #[tokio::test] + async fn finetune_reaches_ready_with_tuned_binding() { + let mut aiplatform = MockAiPlatformApi::new(); + + // create returns a job with a name to poll. + aiplatform + .expect_create_tuning_job() + .withf(|req| { + req.base_model == "gemini-2.0-flash-001" + // The gs:// URI is built from the dependency bucket + // (test-stack-, from GcpStorageController::mock_ready) + training_key. + && req.supervised_tuning_spec.training_dataset_uri + == "gs://test-stack-training-set/training.jsonl" + && req.tuned_model_display_name.as_deref() == Some("llm-tuned") + }) + .times(1) + .returning(|_| { + Ok(TuningJob::builder() + .name( + "projects/test-project-123/locations/us-central1/tuningJobs/42" + .to_string(), + ) + .state(JobState::JobStatePending) + .build()) + }); + + // Poll: still-running once (proving the re-poll loop), then succeeded. + let poll_count = Arc::new(Mutex::new(0u32)); + aiplatform.expect_get_tuning_job().returning(move |name| { + assert_eq!( + name, + "projects/test-project-123/locations/us-central1/tuningJobs/42" + ); + let mut n = poll_count.lock().unwrap(); + *n += 1; + match *n { + 1 => Ok(TuningJob::builder().state(JobState::JobStateRunning).build()), + _ => Ok(succeeded_job()), + } + }); + + let provider = provider_with(aiplatform); + + let mut executor = SingleControllerExecutor::builder() + .resource(finetune_ai()) + .controller(GcpAiController::default()) + .platform(Platform::Gcp) + .service_provider(provider) + .with_dependency( + training_storage(), + GcpStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor should build"); + + executor + .run_until_terminal() + .await + .expect("finetune create flow should complete"); + + assert_eq!( + executor.status(), + ResourceStatus::Running, + "a completed tuning job leaves the gateway Running" + ); + + // The controller recorded the tuned artifact. + let controller = executor + .internal_state::() + .expect("controller downcast"); + assert_eq!( + controller.tuned_model_upstream_id.as_deref(), + Some(TUNED_ENDPOINT) + ); + assert_eq!(controller.tuned_model_served_id.as_deref(), Some("llm-tuned")); + + // The binding the gateway consumes carries the tuned model. + let params = controller + .get_binding_params() + .expect("binding params ok") + .expect("binding present once project/location are set"); + let binding: AiBinding = + serde_json::from_value(params).expect("binding deserializes"); + let tuned = binding + .tuned_model() + .expect("binding must carry the tuned model"); + assert_eq!(tuned.served_id, "llm-tuned"); + assert_eq!(tuned.upstream_id, TUNED_ENDPOINT); + } + + /// A terminal FAILED job fails the resource loud (no silent fall-back to base). + #[tokio::test] + async fn finetune_failed_job_reaches_provision_failed() { + let mut aiplatform = MockAiPlatformApi::new(); + aiplatform.expect_create_tuning_job().returning(|_| { + Ok(TuningJob::builder() + .name("projects/test-project-123/locations/us-central1/tuningJobs/7".to_string()) + .state(JobState::JobStatePending) + .build()) + }); + aiplatform.expect_get_tuning_job().returning(|_| { + Ok(TuningJob::builder() + .state(JobState::JobStateFailed) + .error( + alien_gcp_clients::longrunning::Status::builder() + .code(3) + .message("training data malformed".to_string()) + .build(), + ) + .build()) + }); + + let provider = provider_with(aiplatform); + + let mut executor = SingleControllerExecutor::builder() + .resource(finetune_ai()) + .controller(GcpAiController::default()) + .platform(Platform::Gcp) + .service_provider(provider) + .with_dependency( + training_storage(), + GcpStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor should build"); + + // Fail-fast: the handler surfaces the terminal-failure as an error rather + // than silently reaching Ready on the untuned base model. In production the + // executor catches this and applies `on_failure = CreateFailed` + // (ProvisionFailed); the test harness surfaces the raw error, which we + // assert on directly. + let err = executor + .run_until_terminal() + .await + .expect_err("a failed tuning job must error out, not reach Ready"); + assert_eq!(err.code, "CLOUD_PLATFORM_ERROR"); + assert!( + err.to_string().contains("JobStateFailed") + && err.to_string().contains("training data malformed"), + "error must name the terminal state and carry the provider detail: {err}" + ); + + // The resource never reached Running and recorded no tuned model, so the + // gateway would not serve an untuned model as if it were tuned. + assert_ne!(executor.status(), ResourceStatus::Running); + let controller = executor + .internal_state::() + .expect("controller downcast"); + assert!(controller.tuned_model_upstream_id.is_none()); + } + + /// Regression: an Ai without a finetune spec reaches Ready with an untuned + /// binding and never touches the aiplatform client. + #[tokio::test] + async fn no_finetune_reaches_ready_with_untuned_binding() { + // A create-tuning-job expectation would fail if the controller called it. + let aiplatform = MockAiPlatformApi::new(); + let provider = provider_with(aiplatform); + + let mut executor = SingleControllerExecutor::builder() + .resource(base_ai()) + .controller(GcpAiController::default()) + .platform(Platform::Gcp) + .service_provider(provider) + .build() + .await + .expect("executor should build"); + + executor + .run_until_terminal() + .await + .expect("inference create flow should complete"); + + assert_eq!(executor.status(), ResourceStatus::Running); + + let controller = executor + .internal_state::() + .expect("controller downcast"); + assert!( + controller.tuning_job_name.is_none(), + "a pure-inference gateway never submits a tuning job" + ); + + let params = controller + .get_binding_params() + .expect("binding params ok") + .expect("binding present"); + let binding: AiBinding = + serde_json::from_value(params).expect("binding deserializes"); + assert!( + binding.tuned_model().is_none(), + "an untuned gateway must not carry a tuned model" + ); + } + + /// Vertex does not expose LoRA for Gemini supervised tuning, so submit fails + /// loud rather than silently mapping it to something else. + #[tokio::test] + async fn unsupported_method_reaches_provision_failed() { + // No create call should happen — the method check rejects before submit. + let aiplatform = MockAiPlatformApi::new(); + let provider = provider_with(aiplatform); + + let mut executor = SingleControllerExecutor::builder() + .resource(lora_ai()) + .controller(GcpAiController::default()) + .platform(Platform::Gcp) + .service_provider(provider) + .with_dependency( + training_storage(), + GcpStorageController::mock_ready(TRAINING_STORAGE_ID), + ) + .build() + .await + .expect("executor should build"); + + // The method check rejects before any client call; the handler errors out + // (mapped to `on_failure = CreateFailed` by the executor in production). + let err = executor + .run_until_terminal() + .await + .expect_err("an unsupported tuning method must fail rather than mis-tune"); + assert_eq!(err.code, "RESOURCE_CONFIG_INVALID"); + assert!( + err.to_string().to_lowercase().contains("supervised"), + "error should explain Vertex only supports supervised tuning: {err}" + ); + assert_ne!(executor.status(), ResourceStatus::Running); } } diff --git a/crates/alien-infra/src/ai/gcp_import.rs b/crates/alien-infra/src/ai/gcp_import.rs index 0de6725b7..fa4cebf3c 100644 --- a/crates/alien-infra/src/ai/gcp_import.rs +++ b/crates/alien-infra/src/ai/gcp_import.rs @@ -29,6 +29,10 @@ impl ResourceImporter for GcpAiImporter { state: GcpAiState::Ready, project: Some(data.project_id), location: Some(data.location), + // Import targets an existing base gateway; no tuning job is involved. + tuning_job_name: None, + tuned_model_upstream_id: None, + tuned_model_served_id: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-infra/src/core/service_provider.rs b/crates/alien-infra/src/core/service_provider.rs index 2219bd2b9..d2c30162e 100644 --- a/crates/alien-infra/src/core/service_provider.rs +++ b/crates/alien-infra/src/core/service_provider.rs @@ -3,6 +3,7 @@ use alien_aws_clients::{ acm::{AcmApi, AcmClient}, apigatewayv2::{ApiGatewayV2Api, ApiGatewayV2Client}, autoscaling::{AutoScalingApi, AutoScalingClient}, + bedrock::{BedrockApi, BedrockClient}, cloudformation::{CloudFormationApi, CloudFormationClient}, codebuild::{CodeBuildApi, CodeBuildClient}, dynamodb::{DynamoDbApi, DynamoDbClient}, @@ -39,6 +40,7 @@ use alien_azure_clients::{ managed_clusters::{AzureManagedClustersClient, ManagedClustersApi}, managed_identity::{AzureManagedIdentityClient, ManagedIdentityApi}, network::{AzureNetworkClient, NetworkApi as AzureNetworkApi}, + openai_finetuning::{AzureFoundryFineTuningClient, FoundryFineTuningApi}, private_networking::{AzurePrivateNetworkingClient, PrivateNetworkingApi}, resource_skus::{AzureResourceSkusClient, ResourceSkusApi}, resources::{AzureResourcesClient, ResourcesApi}, @@ -53,6 +55,7 @@ use alien_azure_clients::{ }; use alien_error::Context; use alien_gcp_clients::{ + aiplatform::{AiPlatformApi, AiPlatformClient}, artifactregistry::{ArtifactRegistryApi, ArtifactRegistryClient}, cloud_sql::{CloudSqlApi, CloudSqlClient}, cloudbuild::{CloudBuildApi, CloudBuildClient}, @@ -129,6 +132,10 @@ pub trait PlatformServiceProvider: Send + Sync { &self, config: &AwsClientConfig, ) -> Result>; + async fn get_aws_bedrock_client( + &self, + config: &AwsClientConfig, + ) -> Result>; // GCP clients fn get_gcp_iam_client(&self, config: &GcpClientConfig) -> Result>; @@ -164,6 +171,10 @@ pub trait PlatformServiceProvider: Send + Sync { &self, config: &GcpClientConfig, ) -> Result>; + fn get_gcp_aiplatform_client( + &self, + config: &GcpClientConfig, + ) -> Result>; // Azure clients fn get_azure_application_gateway_client( @@ -230,6 +241,10 @@ pub trait PlatformServiceProvider: Send + Sync { &self, config: &AzureClientConfig, ) -> Result>; + fn get_azure_foundry_finetuning_client( + &self, + config: &AzureClientConfig, + ) -> Result>; fn get_azure_key_vault_management_client( &self, config: &AzureClientConfig, @@ -682,6 +697,22 @@ impl PlatformServiceProvider for DefaultPlatformServiceProvider { ))) } + async fn get_aws_bedrock_client( + &self, + config: &AwsClientConfig, + ) -> Result> { + let credentials = AwsCredentialProvider::from_config(config.clone()) + .await + .context(crate::error::ErrorData::CloudPlatformError { + message: "Failed to create AWS credential provider".to_string(), + resource_id: None, + })?; + Ok(Arc::new(BedrockClient::new( + reqwest::Client::new(), + credentials, + ))) + } + // GCP implementations fn get_gcp_iam_client(&self, config: &GcpClientConfig) -> Result> { Ok(Arc::new(GcpIamClient::new( @@ -802,6 +833,16 @@ impl PlatformServiceProvider for DefaultPlatformServiceProvider { ))) } + fn get_gcp_aiplatform_client( + &self, + config: &GcpClientConfig, + ) -> Result> { + Ok(Arc::new(AiPlatformClient::new( + reqwest::Client::new(), + config.clone(), + ))) + } + // Azure implementations fn get_azure_authorization_client( &self, @@ -963,6 +1004,16 @@ impl PlatformServiceProvider for DefaultPlatformServiceProvider { ))) } + fn get_azure_foundry_finetuning_client( + &self, + config: &AzureClientConfig, + ) -> Result> { + Ok(Arc::new(AzureFoundryFineTuningClient::new( + reqwest::Client::new(), + AzureTokenCache::new(config.clone()), + ))) + } + fn get_azure_key_vault_management_client( &self, config: &AzureClientConfig, diff --git a/crates/alien-permissions/permission-sets/ai/finetune.jsonc b/crates/alien-permissions/permission-sets/ai/finetune.jsonc new file mode 100644 index 000000000..0909629e4 --- /dev/null +++ b/crates/alien-permissions/permission-sets/ai/finetune.jsonc @@ -0,0 +1,90 @@ +{ + "id": "ai/finetune", + "description": "Allows submitting and monitoring model fine-tuning jobs and reading the training dataset from the customer's object storage", + "platforms": { + "aws": [ + { + // Bedrock model customization is a control-plane job the workload submits, + // polls, and (for on-demand serving) invokes the resulting custom model. + // The job reads its dataset from and writes metrics to a customer S3 + // bucket; those S3 grants are scoped by the storage resource's own + // permission set, so this set only needs the Bedrock job + custom-model + // actions. The custom-model ARN is not known until the job completes, so + // the resource ARN is the account-wide custom-model / job namespace. + "grant": { + "actions": [ + "bedrock:CreateModelCustomizationJob", + "bedrock:GetModelCustomizationJob", + "bedrock:StopModelCustomizationJob", + "bedrock:GetCustomModel", + "bedrock:ListCustomModels", + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream" + ] + }, + "binding": { + "stack": { + "resources": [ + "arn:aws:bedrock:*::foundation-model/*", + "arn:aws:bedrock:*:${awsAccountId}:custom-model/*", + "arn:aws:bedrock:*:${awsAccountId}:model-customization-job/*" + ] + }, + "resource": { + "resources": [ + "arn:aws:bedrock:*::foundation-model/*", + "arn:aws:bedrock:*:${awsAccountId}:custom-model/*", + "arn:aws:bedrock:*:${awsAccountId}:model-customization-job/*" + ] + } + } + } + ], + "gcp": [ + { + // Vertex tuning jobs run under a custom role. roles/aiplatform.user carries + // far more than tuning needs; a custom role with the tuning-job lifecycle + // and the predict permissions for serving the tuned endpoint is the + // least-privilege set. Reading the dataset from GCS is granted by the + // storage resource's own set, not here. + "grant": { + "permissions": [ + "aiplatform.tuningJobs.create", + "aiplatform.tuningJobs.get", + "aiplatform.tuningJobs.list", + "aiplatform.tuningJobs.cancel", + "aiplatform.models.get", + "aiplatform.endpoints.predict" + ] + }, + "binding": { + "stack": { "scope": "projects/${projectName}" }, + "resource": { "scope": "projects/${projectName}" } + } + } + ], + "azure": [ + { + // Foundry fine-tuning + deploying the tuned model are account-scoped + // control-plane operations. "Cognitive Services Contributor" is the + // least-privilege built-in role that can create fine-tuning jobs and + // model deployments on the account; "Cognitive Services OpenAI User" + // is added so the same workload can invoke the tuned deployment. + "grant": { + "predefinedRoles": [ + "Cognitive Services Contributor", + "Cognitive Services OpenAI User" + ] + }, + "binding": { + "stack": { + "scope": "/subscriptions/${subscriptionId}/resourceGroups/${resourceGroup}" + }, + "resource": { + "scope": "/subscriptions/${subscriptionId}/resourceGroups/${resourceGroup}/providers/Microsoft.CognitiveServices/accounts/${resourceName}" + } + } + } + ] + } +} diff --git a/crates/alien-permissions/tests/registry_tests.rs b/crates/alien-permissions/tests/registry_tests.rs index 43e31b21e..4874a588b 100644 --- a/crates/alien-permissions/tests/registry_tests.rs +++ b/crates/alien-permissions/tests/registry_tests.rs @@ -102,6 +102,59 @@ fn test_ai_invoke_is_inference_only() { } } +#[test] +fn test_ai_finetune_grants_job_lifecycle() { + let finetune = get_permission_set("ai/finetune").expect("ai/finetune must resolve"); + + // AWS: the Bedrock customization-job lifecycle plus custom-model invoke. + let aws = finetune.platforms.aws.as_ref().expect("aws platform"); + let aws_actions: Vec<&String> = aws + .iter() + .filter_map(|e| e.grant.actions.as_ref()) + .flatten() + .collect(); + for required in [ + "bedrock:CreateModelCustomizationJob", + "bedrock:GetModelCustomizationJob", + "bedrock:GetCustomModel", + ] { + assert!( + aws_actions.iter().any(|a| a.as_str() == required), + "ai/finetune AWS must grant {required}" + ); + } + + // GCP: the Vertex tuning-job lifecycle as a least-privilege permissions list, + // never a predefined role (same invariant ai/invoke holds). + let gcp = finetune.platforms.gcp.as_ref().expect("gcp platform"); + for (i, entry) in gcp.iter().enumerate() { + assert!( + entry.grant.predefined_roles.is_none(), + "ai/finetune GCP entry {i} must use a permissions list, not predefinedRoles" + ); + let perms = entry.grant.permissions.as_ref().expect("permissions list"); + assert!( + perms.iter().any(|p| p == "aiplatform.tuningJobs.create"), + "ai/finetune GCP must grant aiplatform.tuningJobs.create" + ); + } + + // Azure: fine-tuning + deployment need a management-class role, which is + // exactly what ai/invoke forbids — assert finetune is allowed to carry it. + let azure = finetune.platforms.azure.as_ref().expect("azure platform"); + let has_contributor = azure.iter().any(|e| { + e.grant + .predefined_roles + .as_ref() + .map(|roles| roles.iter().any(|r| r == "Cognitive Services Contributor")) + .unwrap_or(false) + }); + assert!( + has_contributor, + "ai/finetune Azure must grant Cognitive Services Contributor to create tuning jobs and deployments" + ); +} + #[test] fn test_ai_provision_has_deployment_writes() { // The predefined model set is deployed at provision time, so deployments/{write,read} @@ -188,7 +241,7 @@ fn test_openai_user_role_id_resolves() { #[test] fn test_ai_permission_sets_have_all_platforms() { - for id in ["ai/provision", "ai/management", "ai/heartbeat", "ai/invoke"] { + for id in ["ai/provision", "ai/management", "ai/heartbeat", "ai/invoke", "ai/finetune"] { let perm_set = get_permission_set(id).unwrap_or_else(|| panic!("{id} must resolve")); assert_eq!(perm_set.id, id); assert!(!perm_set.description.is_empty(), "{id} must have a description"); diff --git a/crates/alien-terraform/src/emitters/aws/ai.rs b/crates/alien-terraform/src/emitters/aws/ai.rs index 78888111c..97d72caaa 100644 --- a/crates/alien-terraform/src/emitters/aws/ai.rs +++ b/crates/alien-terraform/src/emitters/aws/ai.rs @@ -3,9 +3,10 @@ //! AWS Bedrock is a regional, account-scoped service with no per-stack //! cloud resource to provision. The emitter returns an empty fragment for //! the resource itself and emits any resource-scoped `aws_iam_role_policy` -//! blocks for permission profiles that reference `ai/invoke` on this -//! resource. The region is carried in the import ref so the controller can -//! reconstruct the Bedrock endpoint without a cloud round-trip. +//! blocks for permission profiles that reference an `ai/*` set (e.g. +//! `ai/invoke`, or `ai/finetune` when the resource declares a fine-tuning +//! job) on this resource. The region is carried in the import ref so the +//! controller can reconstruct the Bedrock endpoint without a cloud round-trip. //! //! Stack-level permissions flow through `AwsServiceAccountEmitter` via //! `stack_permission_sets`; resource-scoped grants are emitted here. diff --git a/examples/README.md b/examples/README.md index 76572caf3..902ae4847 100644 --- a/examples/README.md +++ b/examples/README.md @@ -4,6 +4,8 @@ Each example is a self-contained template you can initialize with `alien init`. | Template | Description | Language | |----------|-------------|----------| +| [ai-quickstart-ts](./ai-quickstart-ts) | The smallest AI setup: one worker asking cloud LLMs questions through the embedded AI gateway (no API keys). | TypeScript | +| [ai-finetune-inference-ts](./ai-finetune-inference-ts) | Fine-tune a base model in the customer's cloud and serve it for inference — data and weights never leave the account. | TypeScript | | [remote-worker-ts](./remote-worker-ts) | Execute tool calls in your customer's cloud. The AI worker pattern. | TypeScript | | [basic-worker-ts](./basic-worker-ts) | The simplest Alien worker, in TypeScript. | TypeScript | | [basic-worker-rs](./basic-worker-rs) | The simplest Alien worker, in Rust. | Rust | diff --git a/examples/ai-finetune-inference-ts/README.md b/examples/ai-finetune-inference-ts/README.md new file mode 100644 index 000000000..5b6165ffb --- /dev/null +++ b/examples/ai-finetune-inference-ts/README.md @@ -0,0 +1,84 @@ +# AI Fine-tune + Inference + +Fine-tune a base model **inside the customer's cloud**, then serve it for inference — the training data and the tuned weights never leave the customer's account. Extends the AI gateway (`alien.AI`) with a `.finetune()` declaration. + +## What it shows + +- One `alien.AI` resource that both **tunes** a base model and **serves** it (plus the base foundation models) through the same OpenAI-compatible gateway. +- Training data read from the customer's own object storage (S3 / GCS / Blob) under the workload's ambient identity — no keys, no data egress. +- The same app deploys unchanged to **AWS Bedrock**, **GCP Vertex AI**, and **Azure AI Foundry**; the resource resolves to the deploy-target's managed tuning + inference service. + +## How it works + +``` +alien.Storage("dataset") # S3 / GCS / Blob in the customer's account + │ training.jsonl uploaded by the worker + ▼ +alien.AI("llm").finetune({...}) # on deploy, the controller submits the + │ # provider's tuning job reading `dataset`, + │ # polls to completion, records the artifact + ▼ +gateway serves "support-tuned" # tuned model routed alongside the base catalog +``` + +At deploy time the AI resource's cloud controller submits the tuning job (Bedrock `CreateModelCustomizationJob`, Vertex `tuningJobs`, or Foundry `fine_tuning.jobs`), polls it to completion via its heartbeat loop, and records the tuned artifact. The gateway then routes the public id `support-tuned` to that artifact — so app code calls it exactly like a base model, only the `model` string differs. + +## API + +| Route | Method | Purpose | +|-------|--------|---------| +| `/dataset` | POST | Upload JSONL training data into the customer's bucket (body = JSONL) | +| `/finetune/status` | GET | `ready` once the tuning job has completed, else `pending` | +| `/chat` | POST | Inference against a base foundation model (`{ "message": "..." }`) | +| `/chat-tuned` | POST | Inference against the fine-tuned model — same call, different model id | + +## Run locally + +```bash +npm install +alien dev +``` + +Upload the sample dataset and query the tuned model: + +```bash +curl -X POST --data-binary @sample-training.jsonl http://localhost:8080/dataset +curl http://localhost:8080/finetune/status +curl -X POST http://localhost:8080/chat-tuned -d '{"message":"How do I reset my password?"}' +``` + +On the **local** platform the AI resource is a BYO-key provider (set `OPENAI_API_KEY`); fine-tuning is a managed-cloud capability, so `/finetune/status` reports `pending` locally and the tuned route falls back to the base model. Deploy to a cloud to exercise the real tuning flow. + +## Picking `baseModel` per cloud + +`baseModel` is a **provider-native** id — set it to match the cloud you deploy to (see `alien.ts`): + +| Cloud | Service | Example `baseModel` | Tuning method | +|-------|---------|--------------------|---------------| +| AWS | Bedrock | `amazon.nova-lite-v1:0` | SFT (`sft`) — also RFT on Nova | +| GCP | Vertex AI | a Gemini model id | Supervised (`sft`); Vertex does **not** expose LoRA/QLoRA as a user knob for Gemini | +| Azure | AI Foundry | a `gpt-4o` / `gpt-4.1` family id | `sft`, plus `dpo` on some models (LoRA underneath) | + +## Data residency — read before you ship + +"In the customer's cloud" is not automatic on every tier. What the verified provider docs say: + +- **AWS Bedrock** — training data is S3-in / S3-out in the customer's buckets; the job runs under a customer IAM role; the custom model serves **on-demand** (Provisioned Throughput is *not* required — a common misconception). Private connectivity via PrivateLink. Region-confined by default. +- **Azure AI Foundry** — training data and the tuned model are stored at rest in the customer's Foundry resource, in-tenant, same geography (AES-256, optional CMK). **But**: the *Global Standard* and *Developer* deployment/training tiers may move weights outside the resource's region for cost — pin **Standard** if you need strict residency. Two gotchas: training JSONL must be UTF-8 **with a BOM**, and importing from Blob requires the storage account to allow **public** network access. +- **GCP Vertex AI** — supervised-tunes Gemini and auto-provisions a managed tuned-model endpoint; per-epoch checkpoints are auto-deployed. Residency follows the chosen region, but auto-deployed checkpoints can relax strict regional confinement — verify for your region. + +For a hard "training data **and** tuned weights never leave the chosen region" guarantee, pin the region and the residency-preserving tier on each provider (Bedrock in-region + Standard-equivalent, Azure **Standard**, Vertex single-region), and confirm against current provider docs — these tiers and defaults change. + +## Training data format + +JSONL, one example per line, in the OpenAI chat/conversational shape: + +```json +{"messages":[{"role":"user","content":"How do I reset my password?"},{"role":"assistant","content":"Go to Settings → Security → Reset password, then follow the emailed link."}]} +``` + +See `sample-training.jsonl`. Real fine-tuning needs enough examples to matter (Azure requires ≥ 10; providers recommend hundreds) — the sample is illustrative. + +## License + +ISC diff --git a/examples/ai-finetune-inference-ts/alien.ts b/examples/ai-finetune-inference-ts/alien.ts new file mode 100644 index 000000000..c13b7105e --- /dev/null +++ b/examples/ai-finetune-inference-ts/alien.ts @@ -0,0 +1,57 @@ +// EXAMPLE: Fine-tune a base model in the customer's cloud, then serve it for +// inference — without the training data or the tuned weights ever leaving the +// customer's account. + +import * as alien from "@alienplatform/core" + +// The training dataset lives in the customer's object storage: +// S3 on AWS, Cloud Storage on GCP, Blob Storage on Azure. The worker uploads +// JSONL examples here; the tuning job reads them in-account. +const dataset = new alien.Storage("dataset").build() + +// A model-less AI gateway that also fine-tunes. On deploy, the resource's cloud +// controller submits the provider's tuning job (Bedrock CreateModelCustomizationJob, +// Vertex tuningJobs, or Foundry fine_tuning.jobs) reading `dataset` in the +// customer's account, polls it to completion, and serves the tuned model through +// the same gateway under `servedModelId`. Base models remain callable too. +// +// `baseModel` is a provider-native id; pick the one that matches the cloud you +// deploy to (see the README's per-provider table). The default here targets +// AWS Bedrock (Amazon Nova). +const llm = new alien.AI("llm") + .finetune({ + baseModel: "amazon.nova-lite-v1:0", + trainingData: dataset, + trainingKey: "training.jsonl", + servedModelId: "support-tuned", + method: "sft", + }) + .build() + +const api = new alien.Worker("api") + .code({ type: "source", src: "./", toolchain: { type: "typescript" } }) + // GCP Cloud Run gen2 requires >= 512 MiB; the default 256 MiB fails its preflight. + .memoryMb(512) + .publicEndpoint("api") + // Linking injects the dataset binding and the AI gateway (ALIEN_LLM_BINDING). + .link(dataset) + .link(llm) + .permissions("execution") + .build() + +export default new alien.Stack("ai-finetune-inference") + .platforms(["aws", "gcp", "azure"]) + .add(dataset, "live") + .add(llm, "live") + .add(api, "live") + .permissions({ + profiles: { + execution: { + // Read/write the dataset bucket, and both submit the tuning job and + // invoke models (base + tuned) through the gateway. + dataset: ["storage/data-read", "storage/data-write"], + "*": ["ai/invoke", "ai/finetune"], + }, + }, + }) + .build() diff --git a/examples/ai-finetune-inference-ts/package.json b/examples/ai-finetune-inference-ts/package.json new file mode 100644 index 000000000..a0e540782 --- /dev/null +++ b/examples/ai-finetune-inference-ts/package.json @@ -0,0 +1,22 @@ +{ + "name": "ai-finetune-inference-ts", + "version": "1.0.0", + "type": "module", + "private": true, + "scripts": { + "dev": "alien dev", + "build": "alien build --platform local", + "test": "vitest run" + }, + "dependencies": { + "@alienplatform/sdk": "^1.8.0", + "@alienplatform/core": "^1.8.0", + "hono": "^4.0.0" + }, + "devDependencies": { + "@alienplatform/testing": "^1.8.0", + "@types/node": "^24.0.15", + "typescript": "^5.8.3", + "vitest": "^3.2.7" + } +} diff --git a/examples/ai-finetune-inference-ts/sample-training.jsonl b/examples/ai-finetune-inference-ts/sample-training.jsonl new file mode 100644 index 000000000..33545a4d1 --- /dev/null +++ b/examples/ai-finetune-inference-ts/sample-training.jsonl @@ -0,0 +1,10 @@ +{"messages":[{"role":"user","content":"How do I reset my password?"},{"role":"assistant","content":"Go to Settings → Security → Reset password, then follow the emailed link. The link expires in 30 minutes."}]} +{"messages":[{"role":"user","content":"My invoice looks wrong this month."},{"role":"assistant","content":"I can help. Open Billing → Invoices, select the invoice, and click 'Report an issue'. Include the line item you're disputing and our team responds within one business day."}]} +{"messages":[{"role":"user","content":"How do I invite a teammate?"},{"role":"assistant","content":"Open Team → Members → Invite, enter their work email, pick a role, and send. They'll get an email to join your workspace."}]} +{"messages":[{"role":"user","content":"Can I export my data?"},{"role":"assistant","content":"Yes. Go to Settings → Data → Export, choose CSV or JSON, and we'll email a download link when it's ready — usually within a few minutes."}]} +{"messages":[{"role":"user","content":"The app is slow today."},{"role":"assistant","content":"Sorry about that. First check status.example.com for incidents. If all green, try clearing your cache and reloading; if it persists, send us your account id and region and we'll investigate."}]} +{"messages":[{"role":"user","content":"How do I cancel my subscription?"},{"role":"assistant","content":"Open Billing → Plan → Cancel subscription. You keep access until the end of the current billing period, and no further charges are made."}]} +{"messages":[{"role":"user","content":"Do you support single sign-on?"},{"role":"assistant","content":"Yes, on Business and Enterprise plans. Go to Settings → Security → SSO to configure SAML or OIDC with your identity provider."}]} +{"messages":[{"role":"user","content":"I didn't get my verification email."},{"role":"assistant","content":"Check your spam folder first. If it's not there, request a new one from the login screen — 'Resend verification'. Allow a few minutes for delivery."}]} +{"messages":[{"role":"user","content":"How do I change my billing email?"},{"role":"assistant","content":"Open Billing → Settings → Billing contact, update the email, and save. Future invoices go to the new address."}]} +{"messages":[{"role":"user","content":"Is my data encrypted?"},{"role":"assistant","content":"Yes — data is encrypted in transit (TLS) and at rest (AES-256). Enterprise plans can bring their own encryption key."}]} diff --git a/examples/ai-finetune-inference-ts/src/index.ts b/examples/ai-finetune-inference-ts/src/index.ts new file mode 100644 index 000000000..e92b36ee4 --- /dev/null +++ b/examples/ai-finetune-inference-ts/src/index.ts @@ -0,0 +1,66 @@ +import { ai, storage } from "@alienplatform/sdk" +import { Hono } from "hono" + +// The public id the tuned model is served under (matches `servedModelId` in alien.ts). +const TUNED_MODEL = "support-tuned" +const TRAINING_KEY = "training.jsonl" + +const app = new Hono() + +// 1. Upload the JSONL training set into the customer's bucket. The tuning job +// (submitted by the AI resource's controller at deploy time) reads it from +// there — the data never leaves the customer's cloud. In a real app you'd +// seed this before deploy; exposed here so the flow is runnable end-to-end. +app.post("/dataset", async c => { + const body = await c.req.text() + if (!body.trim()) { + return c.json({ error: "POST JSONL training data as the request body" }, 400) + } + await storage("dataset").put(TRAINING_KEY, new TextEncoder().encode(body)) + const lines = body.split("\n").filter(l => l.trim()).length + return c.json({ uploaded: TRAINING_KEY, examples: lines }) +}) + +// 2. Fine-tune status. The tuned model shows up in the gateway's model list only +// once its job has completed, so its presence is a simple readiness signal. +app.get("/finetune/status", async c => { + const models = await ai("llm").getAvailableModels() + const ready = models.some(m => m.id === TUNED_MODEL) + return c.json({ + tunedModel: TUNED_MODEL, + status: ready ? "ready" : "pending", + availableModels: models.map(m => m.id), + }) +}) + +// 3. Inference against a base foundation model (the per-cloud catalog). +app.post("/chat", async c => { + return chat(c, await defaultBaseModel()) +}) + +// 4. Inference against the fine-tuned model — same OpenAI-compatible call, just a +// different `model` id. The gateway routes it to the tuned artifact in-account. +app.post("/chat-tuned", async c => { + return chat(c, TUNED_MODEL) +}) + +async function defaultBaseModel(): Promise { + const models = await ai("llm").getAvailableModels() + const base = models.find(m => m.id !== TUNED_MODEL)?.id + if (!base) throw new Error("no base models available for this cloud") + return base +} + +async function chat(c: Parameters[1]>[0], model: string) { + const { message } = await c.req.json<{ message?: string }>() + if (!message) { + return c.json({ error: "send { \"message\": \"...\" }" }, 400) + } + const completion = (await ai("llm").chat.completions.create({ + model, + messages: [{ role: "user", content: message }], + })) as { choices?: Array<{ message?: { content?: string } }> } + return c.json({ model, answer: completion.choices?.[0]?.message?.content ?? "" }) +} + +export default app diff --git a/examples/ai-finetune-inference-ts/template.toml b/examples/ai-finetune-inference-ts/template.toml new file mode 100644 index 000000000..8573c7917 --- /dev/null +++ b/examples/ai-finetune-inference-ts/template.toml @@ -0,0 +1,3 @@ +name = "ai-finetune-inference-ts" +description = "Fine-tune a base model in the customer's cloud and serve it for inference through the embedded Alien AI gateway — training data and tuned weights never leave the account." +language = "TypeScript" diff --git a/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts b/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts new file mode 100644 index 000000000..d3696c4f1 --- /dev/null +++ b/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts @@ -0,0 +1,68 @@ +import { readFileSync } from "node:fs" +import { fileURLToPath } from "node:url" +import { type Deployment, deploy } from "@alienplatform/testing" +import { afterAll, beforeAll, describe, expect, it } from "vitest" + +// The local platform serves the AI gateway as a BYO-key provider; fine-tuning is +// a managed-cloud capability, so locally we verify the deployable surface — +// dataset upload, status shape, and (when a key is present) base inference — not +// a real tuning job. The cloud tuning flow is exercised by deploying to a cloud. +const OPENAI_KEY = process.env.OPENAI_API_KEY + +describe("ai-finetune-inference-ts", () => { + let deployment: Deployment + + beforeAll(async () => { + deployment = await deploy({ app: ".", platform: "local" }) + }, 300_000) + + afterAll(async () => { + await deployment?.destroy() + }) + + it("uploads training data into the customer bucket", async () => { + const jsonl = readFileSync( + fileURLToPath(new URL("../sample-training.jsonl", import.meta.url)), + "utf8", + ) + const response = await fetch(`${deployment.url}/dataset`, { + method: "POST", + body: jsonl, + }) + expect(response.status).toBe(200) + const body = (await response.json()) as { uploaded: string; examples: number } + expect(body.uploaded).toBe("training.jsonl") + expect(body.examples).toBe(10) + }) + + it("rejects an empty dataset upload", async () => { + const response = await fetch(`${deployment.url}/dataset`, { method: "POST", body: "" }) + expect(response.status).toBe(400) + }) + + it("reports a well-formed fine-tune status", async () => { + const response = await fetch(`${deployment.url}/finetune/status`) + expect(response.status).toBe(200) + const body = (await response.json()) as { + tunedModel: string + status: string + availableModels: string[] + } + expect(body.tunedModel).toBe("support-tuned") + // Locally the managed tuning job never runs, so the tuned model is not ready. + expect(body.status).toBe("pending") + expect(Array.isArray(body.availableModels)).toBe(true) + }) + + it.skipIf(!OPENAI_KEY)("answers a base-model chat when a provider key is set", async () => { + const response = await fetch(`${deployment.url}/chat`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ message: "Say the single word: pong" }), + }) + expect(response.status).toBe(200) + const body = (await response.json()) as { model: string; answer: string } + expect(body.model).toBeTruthy() + expect(body.answer.length).toBeGreaterThan(0) + }) +}) diff --git a/examples/ai-finetune-inference-ts/tsconfig.json b/examples/ai-finetune-inference-ts/tsconfig.json new file mode 100644 index 000000000..c6051cab0 --- /dev/null +++ b/examples/ai-finetune-inference-ts/tsconfig.json @@ -0,0 +1,14 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ESNext", + "lib": ["ES2022"], + "moduleResolution": "bundler", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "noEmit": true + }, + "include": ["src/**/*", "alien.ts", "tests/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/examples/ai-finetune-inference-ts/vitest.config.ts b/examples/ai-finetune-inference-ts/vitest.config.ts new file mode 100644 index 000000000..5dd3c9e50 --- /dev/null +++ b/examples/ai-finetune-inference-ts/vitest.config.ts @@ -0,0 +1,10 @@ +import { defineConfig } from "vitest/config" + +export default defineConfig({ + test: { + include: ["tests/**/*.test.ts"], + testTimeout: 300_000, // 5 min timeout for tests involving deployment + hookTimeout: 300_000, // 5 min for beforeAll/afterAll + pool: "forks", // Use forks to ensure clean process state + }, +}) diff --git a/examples/pnpm-lock.yaml b/examples/pnpm-lock.yaml index 4688b3b49..d0efa4199 100644 --- a/examples/pnpm-lock.yaml +++ b/examples/pnpm-lock.yaml @@ -17,6 +17,31 @@ importers: .: {} + ai-finetune-inference-ts: + dependencies: + '@alienplatform/core': + specifier: file:../../packages/core + version: file:../packages/core(@types/json-schema@7.0.15)(openapi-types@12.1.3) + '@alienplatform/sdk': + specifier: file:../../packages/sdk + version: file:../packages/sdk(@types/json-schema@7.0.15)(openapi-types@12.1.3) + hono: + specifier: ^4.0.0 + version: 4.12.5 + devDependencies: + '@alienplatform/testing': + specifier: file:../../packages/testing + version: file:../packages/testing(@aws-sdk/client-ssm@3.1004.0)(@azure/identity@4.13.0)(@azure/keyvault-secrets@4.10.0(@azure/core-client@1.10.1))(@google-cloud/secret-manager@5.6.0)(@types/json-schema@7.0.15)(openapi-types@12.1.3) + '@types/node': + specifier: ^24.0.15 + version: 24.12.0 + typescript: + specifier: ^5.8.3 + version: 5.8.3 + vitest: + specifier: ^3.2.7 + version: 3.2.7(@types/debug@4.1.13)(@types/node@24.12.0)(jiti@2.7.0)(lightningcss@1.32.0) + ai-quickstart-ts: dependencies: '@alienplatform/core': diff --git a/examples/pnpm-workspace.yaml b/examples/pnpm-workspace.yaml index 70571183b..45eacc56c 100644 --- a/examples/pnpm-workspace.yaml +++ b/examples/pnpm-workspace.yaml @@ -1,5 +1,6 @@ packages: - ai-quickstart-ts + - ai-finetune-inference-ts - basic-worker-ts - basic-worker-rs - remote-worker-ts diff --git a/packages/core/src/__tests__/ai.test.ts b/packages/core/src/__tests__/ai.test.ts index acc67c39f..683b517b1 100644 --- a/packages/core/src/__tests__/ai.test.ts +++ b/packages/core/src/__tests__/ai.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest" import { AI } from "../ai.js" +import { Storage } from "../storage.js" describe("AI", () => { it("builds with just an id", () => { @@ -26,4 +27,45 @@ describe("AI", () => { expect(r.config).not.toHaveProperty("external") expect(r.config.id).toBe("llm") }) + + it("omits finetune by default", () => { + const r = new AI("llm").build() + expect(r.config).not.toHaveProperty("finetune") + }) + + it("accepts a Storage resource for trainingData and resolves it to its id", () => { + const dataset = new Storage("training-set").build() + const r = new AI("llm") + .finetune({ baseModel: "amazon.nova-lite-v1:0", trainingData: dataset }) + .build() + const finetune = (r.config as { finetune?: Record }).finetune + expect(finetune).toBeDefined() + expect(finetune?.baseModel).toBe("amazon.nova-lite-v1:0") + expect(finetune?.trainingData).toBe("training-set") + }) + + it("accepts a string id for trainingData and carries all fields", () => { + const r = new AI("llm") + .finetune({ + baseModel: "amazon.nova-lite-v1:0", + trainingData: "training-set", + trainingKey: "data.jsonl", + servedModelId: "finance-model", + method: "lora", + }) + .build() + const finetune = (r.config as { finetune?: Record }).finetune + expect(finetune).toEqual({ + baseModel: "amazon.nova-lite-v1:0", + trainingData: "training-set", + trainingKey: "data.jsonl", + servedModelId: "finance-model", + method: "lora", + }) + }) + + it("is chainable and returns the same builder", () => { + const ai = new AI("llm") + expect(ai.finetune({ baseModel: "b", trainingData: "d" })).toBe(ai) + }) }) diff --git a/packages/core/src/ai.ts b/packages/core/src/ai.ts index c805ac380..49e468709 100644 --- a/packages/core/src/ai.ts +++ b/packages/core/src/ai.ts @@ -1,9 +1,34 @@ -import { type Ai as AiConfig, AiSchema, type ResourceType } from "./generated/index.js" +import { + type Ai as AiConfig, + AiSchema, + type FinetuneMethod, + type ResourceType, +} from "./generated/index.js" import { Resource } from "./resource.js" -export type { AiOutputs, Ai as AiConfig } from "./generated/index.js" +export type { AiOutputs, Ai as AiConfig, FinetuneSpec, FinetuneMethod } from "./generated/index.js" export { AiSchema as AiConfigSchema } from "./generated/index.js" +/** + * Options for {@link AI.finetune}. `trainingData` accepts either a built + * {@link Resource} (a Storage bucket) or its id string. + */ +export interface FinetuneOptions { + /** + * Provider-native base-model id to tune (an Amazon Nova id on Bedrock, a + * Gemini model on Vertex, or a gpt-4o family model on Foundry). + */ + baseModel: string + /** The Storage resource (or its id) holding the JSONL training dataset. */ + trainingData: Resource | string + /** Object key of the training file within `trainingData`. Defaults to `training.jsonl`. */ + trainingKey?: string + /** Public model id apps use to invoke the tuned model. Defaults to `-tuned`. */ + servedModelId?: string + /** The fine-tuning method. Defaults to `"sft"` (supervised fine-tuning). */ + method?: FinetuneMethod +} + /** * Represents an AI Gateway resource that provides a unified interface to * managed AI inference services across cloud providers. @@ -19,6 +44,33 @@ export class AI { this._config.id = id } + /** + * Declares that this resource should fine-tune a base model in the customer's + * cloud before serving it. The tuning job reads the JSONL dataset from the + * given Storage bucket (S3 / GCS / Blob) under the workload's ambient + * identity — the data never leaves the customer's cloud — and the tuned model + * is served through the same gateway under `servedModelId` (default + * `-tuned`). Omit this call for a pure inference gateway. + * + * @param options Fine-tuning configuration. + * @returns This builder, for chaining. + */ + public finetune(options: FinetuneOptions): this { + const trainingData = + typeof options.trainingData === "string" + ? options.trainingData + : options.trainingData.config.id + + this._config.finetune = { + baseModel: options.baseModel, + trainingData, + ...(options.trainingKey !== undefined ? { trainingKey: options.trainingKey } : {}), + ...(options.servedModelId !== undefined ? { servedModelId: options.servedModelId } : {}), + ...(options.method !== undefined ? { method: options.method } : {}), + } + return this + } + /** * Returns a ResourceType representing any AI resource. * Used for creating permission targets that apply to all AI resources. diff --git a/packages/core/src/generated/index.ts b/packages/core/src/generated/index.ts index ad0ab8b5d..d74b50a28 100644 --- a/packages/core/src/generated/index.ts +++ b/packages/core/src/generated/index.ts @@ -163,6 +163,8 @@ export type { EventState } from "./zod/event-state-schema.js"; export type { ExposeProtocol } from "./zod/expose-protocol-schema.js"; export type { ExternalAiHeartbeatData } from "./zod/external-ai-heartbeat-data-schema.js"; export type { FailureDomainSelection } from "./zod/failure-domain-selection-schema.js"; +export type { FinetuneMethod } from "./zod/finetune-method-schema.js"; +export type { FinetuneSpec } from "./zod/finetune-spec-schema.js"; export type { GcpArtifactRegistryHeartbeatData } from "./zod/gcp-artifact-registry-heartbeat-data-schema.js"; export type { GcpArtifactRegistryImportData } from "./zod/gcp-artifact-registry-import-data-schema.js"; export type { GcpBuildImportData } from "./zod/gcp-build-import-data-schema.js"; @@ -543,6 +545,8 @@ export { EventStateSchema } from "./zod/event-state-schema.js"; export { ExposeProtocolSchema } from "./zod/expose-protocol-schema.js"; export { ExternalAiHeartbeatDataSchema } from "./zod/external-ai-heartbeat-data-schema.js"; export { FailureDomainSelectionSchema } from "./zod/failure-domain-selection-schema.js"; +export { FinetuneMethodSchema } from "./zod/finetune-method-schema.js"; +export { FinetuneSpecSchema } from "./zod/finetune-spec-schema.js"; export { GcpArtifactRegistryHeartbeatDataSchema } from "./zod/gcp-artifact-registry-heartbeat-data-schema.js"; export { GcpArtifactRegistryImportDataSchema } from "./zod/gcp-artifact-registry-import-data-schema.js"; export { GcpBuildImportDataSchema } from "./zod/gcp-build-import-data-schema.js"; diff --git a/packages/core/src/generated/schemas/ai.json b/packages/core/src/generated/schemas/ai.json index 713b9a9ff..a0f8fc14a 100644 --- a/packages/core/src/generated/schemas/ai.json +++ b/packages/core/src/generated/schemas/ai.json @@ -1 +1 @@ -{"type":"object","description":"Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.","required":["id"],"properties":{"id":{"type":"string","description":"Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters."}},"additionalProperties":false,"x-readme-ref-name":"Ai"} \ No newline at end of file +{"type":"object","description":"Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.\n\nWhen `finetune` is set, the resource also tunes a base model in the\ncustomer's cloud and serves the result through the same gateway (see\n[`FinetuneSpec`]).","required":["id"],"properties":{"finetune":{"oneOf":[{"type":"null"},{"description":"Optional fine-tuning declaration. When present, the resource tunes\n`finetune.base_model` on the customer's cloud and serves the result\nalongside the base models. Absent for a pure inference gateway.","type":"object","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"}]},"id":{"type":"string","description":"Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters."}},"additionalProperties":false,"x-readme-ref-name":"Ai"} \ No newline at end of file diff --git a/packages/core/src/generated/schemas/finetuneMethod.json b/packages/core/src/generated/schemas/finetuneMethod.json new file mode 100644 index 000000000..afaa49566 --- /dev/null +++ b/packages/core/src/generated/schemas/finetuneMethod.json @@ -0,0 +1 @@ +{"type":"string","description":"The fine-tuning method applied to the base model.\n\nThe gateway-side controllers map each variant onto the provider's native\ntechnique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is\ndirect preference optimization (Bedrock/Foundry), and `Lora` requests a\nparameter-efficient adapter where the provider exposes it. Providers that\nonly implement a subset reject unsupported methods at job-submit time.","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"} \ No newline at end of file diff --git a/packages/core/src/generated/schemas/finetuneSpec.json b/packages/core/src/generated/schemas/finetuneSpec.json new file mode 100644 index 000000000..347bf463f --- /dev/null +++ b/packages/core/src/generated/schemas/finetuneSpec.json @@ -0,0 +1 @@ +{"type":"object","description":"Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider's tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before.","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"} \ No newline at end of file diff --git a/packages/core/src/generated/zod/ai-schema.ts b/packages/core/src/generated/zod/ai-schema.ts index c16df54de..8ea89a00b 100644 --- a/packages/core/src/generated/zod/ai-schema.ts +++ b/packages/core/src/generated/zod/ai-schema.ts @@ -4,12 +4,16 @@ */ import * as z from "zod"; +import { FinetuneSpecSchema } from "./finetune-spec-schema.js"; /** - * @description Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack\'s external-bindings map; the executor then skips the cloud controller. + * @description Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack\'s external-bindings map; the executor then skips the cloud controller.\n\nWhen `finetune` is set, the resource also tunes a base model in the\ncustomer\'s cloud and serves the result through the same gateway (see\n[`FinetuneSpec`]). */ export const AiSchema = z.object({ - "id": z.string().describe("Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters.") - }).describe("Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.") + get "finetune"(){ + return z.union([FinetuneSpecSchema, z.null()]).optional() + }, +"id": z.string().describe("Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters.") + }).describe("Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.\n\nWhen `finetune` is set, the resource also tunes a base model in the\ncustomer's cloud and serves the result through the same gateway (see\n[`FinetuneSpec`]).") export type Ai = z.infer \ No newline at end of file diff --git a/packages/core/src/generated/zod/finetune-method-schema.ts b/packages/core/src/generated/zod/finetune-method-schema.ts new file mode 100644 index 000000000..152c63309 --- /dev/null +++ b/packages/core/src/generated/zod/finetune-method-schema.ts @@ -0,0 +1,13 @@ +/** +* Generated by Kubb (https://kubb.dev/). +* Do not edit manually. +*/ + +import * as z from "zod"; + +/** + * @description The fine-tuning method applied to the base model.\n\nThe gateway-side controllers map each variant onto the provider\'s native\ntechnique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is\ndirect preference optimization (Bedrock/Foundry), and `Lora` requests a\nparameter-efficient adapter where the provider exposes it. Providers that\nonly implement a subset reject unsupported methods at job-submit time. + */ +export const FinetuneMethodSchema = z.enum(["sft", "dpo", "lora"]).describe("The fine-tuning method applied to the base model.\n\nThe gateway-side controllers map each variant onto the provider's native\ntechnique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is\ndirect preference optimization (Bedrock/Foundry), and `Lora` requests a\nparameter-efficient adapter where the provider exposes it. Providers that\nonly implement a subset reject unsupported methods at job-submit time.") + +export type FinetuneMethod = z.infer \ No newline at end of file diff --git a/packages/core/src/generated/zod/finetune-spec-schema.ts b/packages/core/src/generated/zod/finetune-spec-schema.ts new file mode 100644 index 000000000..68b608447 --- /dev/null +++ b/packages/core/src/generated/zod/finetune-spec-schema.ts @@ -0,0 +1,22 @@ +/** +* Generated by Kubb (https://kubb.dev/). +* Do not edit manually. +*/ + +import * as z from "zod"; +import { FinetuneMethodSchema } from "./finetune-method-schema.js"; + +/** + * @description Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer\'s cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider\'s tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before. + */ +export const FinetuneSpecSchema = z.object({ + "baseModel": z.string().describe("Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."), +get "method"(){ + return FinetuneMethodSchema.describe("The fine-tuning method applied to the base model.\n\nThe gateway-side controllers map each variant onto the provider's native\ntechnique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is\ndirect preference optimization (Bedrock/Foundry), and `Lora` requests a\nparameter-efficient adapter where the provider exposes it. Providers that\nonly implement a subset reject unsupported methods at job-submit time.").optional() + }, +"servedModelId": z.string().describe("The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`.").nullish(), +"trainingData": z.string().describe("The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."), +"trainingKey": z.optional(z.string().describe("Object key of the training file within `training_data`.\nDefaults to `training.jsonl`.")) + }).describe("Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider's tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before.") + +export type FinetuneSpec = z.infer \ No newline at end of file diff --git a/packages/core/src/generated/zod/index.ts b/packages/core/src/generated/zod/index.ts index abfaf822c..b8369ca6a 100644 --- a/packages/core/src/generated/zod/index.ts +++ b/packages/core/src/generated/zod/index.ts @@ -163,6 +163,8 @@ export type { EventState } from "./event-state-schema.js"; export type { ExposeProtocol } from "./expose-protocol-schema.js"; export type { ExternalAiHeartbeatData } from "./external-ai-heartbeat-data-schema.js"; export type { FailureDomainSelection } from "./failure-domain-selection-schema.js"; +export type { FinetuneMethod } from "./finetune-method-schema.js"; +export type { FinetuneSpec } from "./finetune-spec-schema.js"; export type { GcpArtifactRegistryHeartbeatData } from "./gcp-artifact-registry-heartbeat-data-schema.js"; export type { GcpArtifactRegistryImportData } from "./gcp-artifact-registry-import-data-schema.js"; export type { GcpBuildImportData } from "./gcp-build-import-data-schema.js"; @@ -543,6 +545,8 @@ export { EventStateSchema } from "./event-state-schema.js"; export { ExposeProtocolSchema } from "./expose-protocol-schema.js"; export { ExternalAiHeartbeatDataSchema } from "./external-ai-heartbeat-data-schema.js"; export { FailureDomainSelectionSchema } from "./failure-domain-selection-schema.js"; +export { FinetuneMethodSchema } from "./finetune-method-schema.js"; +export { FinetuneSpecSchema } from "./finetune-spec-schema.js"; export { GcpArtifactRegistryHeartbeatDataSchema } from "./gcp-artifact-registry-heartbeat-data-schema.js"; export { GcpArtifactRegistryImportDataSchema } from "./gcp-artifact-registry-import-data-schema.js"; export { GcpBuildImportDataSchema } from "./gcp-build-import-data-schema.js"; From 3e1392c4b3ab879881a4f10061ad07029daca741 Mon Sep 17 00:00:00 2001 From: AB Date: Thu, 23 Jul 2026 14:26:40 +0300 Subject: [PATCH 2/3] Fix s3 names --- examples/ai-finetune-inference-ts/alien.ts | 9 +++++++-- examples/ai-finetune-inference-ts/src/index.ts | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/examples/ai-finetune-inference-ts/alien.ts b/examples/ai-finetune-inference-ts/alien.ts index c13b7105e..1f647e08a 100644 --- a/examples/ai-finetune-inference-ts/alien.ts +++ b/examples/ai-finetune-inference-ts/alien.ts @@ -7,7 +7,12 @@ import * as alien from "@alienplatform/core" // The training dataset lives in the customer's object storage: // S3 on AWS, Cloud Storage on GCP, Blob Storage on Azure. The worker uploads // JSONL examples here; the tuning job reads them in-account. -const dataset = new alien.Storage("dataset").build() +// +// NOTE: S3 bucket names are GLOBALLY unique across all AWS accounts. The bucket +// is named `-`, so a generic id like "dataset" +// can collide with a bucket someone else already owns. Keep this id distinctive +// (and change it if you hit "bucket name is not available"). +const dataset = new alien.Storage("finetune-training-data").build() // A model-less AI gateway that also fine-tunes. On deploy, the resource's cloud // controller submits the provider's tuning job (Bedrock CreateModelCustomizationJob, @@ -49,7 +54,7 @@ export default new alien.Stack("ai-finetune-inference") execution: { // Read/write the dataset bucket, and both submit the tuning job and // invoke models (base + tuned) through the gateway. - dataset: ["storage/data-read", "storage/data-write"], + "finetune-training-data": ["storage/data-read", "storage/data-write"], "*": ["ai/invoke", "ai/finetune"], }, }, diff --git a/examples/ai-finetune-inference-ts/src/index.ts b/examples/ai-finetune-inference-ts/src/index.ts index e92b36ee4..c208e4f40 100644 --- a/examples/ai-finetune-inference-ts/src/index.ts +++ b/examples/ai-finetune-inference-ts/src/index.ts @@ -16,7 +16,7 @@ app.post("/dataset", async c => { if (!body.trim()) { return c.json({ error: "POST JSONL training data as the request body" }, 400) } - await storage("dataset").put(TRAINING_KEY, new TextEncoder().encode(body)) + await storage("finetune-training-data").put(TRAINING_KEY, new TextEncoder().encode(body)) const lines = body.split("\n").filter(l => l.trim()).length return c.json({ uploaded: TRAINING_KEY, examples: lines }) }) From 31d77c5708652777c7ba88ff2d6a65190879798c Mon Sep 17 00:00:00 2001 From: AB Date: Thu, 23 Jul 2026 19:31:13 +0300 Subject: [PATCH 3/3] provision better --- Cargo.lock | 2 + .../src/emitters/aws/ai.rs | 128 +++- .../tests/generator/aws_ai_tests.rs | 80 ++- crates/alien-core/src/bindings/ai.rs | 77 +++ crates/alien-core/src/bindings/mod.rs | 3 +- crates/alien-core/src/resources/ai.rs | 22 +- crates/alien-gateway/Cargo.toml | 2 + crates/alien-gateway/src/config.rs | 8 +- crates/alien-gateway/src/finetune.rs | 164 +++++ crates/alien-gateway/src/finetune/bedrock.rs | 232 +++++++ crates/alien-gateway/src/finetune/foundry.rs | 224 +++++++ crates/alien-gateway/src/finetune/vertex.rs | 243 +++++++ crates/alien-gateway/src/lib.rs | 4 + crates/alien-gateway/src/router.rs | 211 ++++++ crates/alien-gateway/tests/integration.rs | 2 + crates/alien-gateway/tests/live_bedrock.rs | 2 + .../tests/live_foundry_claude.rs | 1 + .../alien-gateway/tests/live_vertex_claude.rs | 1 + crates/alien-infra/src/ai/aws.rs | 532 ++++----------- crates/alien-infra/src/ai/aws_import.rs | 4 +- crates/alien-infra/src/ai/azure.rs | 621 +++--------------- crates/alien-infra/src/ai/azure_import.rs | 6 +- crates/alien-infra/src/ai/gcp.rs | 564 ++++------------ crates/alien-infra/src/ai/gcp_import.rs | 7 +- crates/alien-terraform/src/emitters/aws/ai.rs | 123 +++- .../tests/generator/aws_data_layer_tests.rs | 74 ++- examples/ai-finetune-inference-ts/README.md | 42 +- examples/ai-finetune-inference-ts/alien.ts | 12 +- .../ai-finetune-inference-ts/src/index.ts | 40 +- .../tests/ai-finetune-inference.test.ts | 22 +- .../alien/ai-finetune-runtime-design.md | 105 +++ .../ai-gateway/src/__tests__/client.test.ts | 89 +++ packages/ai-gateway/src/client.ts | 87 ++- packages/ai-gateway/src/index.ts | 2 + packages/core/src/generated/schemas/ai.json | 2 +- .../src/generated/schemas/finetuneSpec.json | 2 +- .../src/generated/zod/finetune-spec-schema.ts | 8 +- packages/sdk/src/index.ts | 2 + 38 files changed, 2285 insertions(+), 1465 deletions(-) create mode 100644 crates/alien-gateway/src/finetune.rs create mode 100644 crates/alien-gateway/src/finetune/bedrock.rs create mode 100644 crates/alien-gateway/src/finetune/foundry.rs create mode 100644 crates/alien-gateway/src/finetune/vertex.rs create mode 100644 internal-docs/alien/ai-finetune-runtime-design.md diff --git a/Cargo.lock b/Cargo.lock index 3af4444f7..85f1c3743 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -603,6 +603,7 @@ dependencies = [ "alien-bindings", "alien-core", "alien-error", + "async-trait", "aws-config", "aws-credential-types", "aws-sigv4", @@ -619,6 +620,7 @@ dependencies = [ "temp-env", "tokio", "tracing", + "urlencoding", ] [[package]] diff --git a/crates/alien-cloudformation/src/emitters/aws/ai.rs b/crates/alien-cloudformation/src/emitters/aws/ai.rs index a92428fb7..856e05c9e 100644 --- a/crates/alien-cloudformation/src/emitters/aws/ai.rs +++ b/crates/alien-cloudformation/src/emitters/aws/ai.rs @@ -15,23 +15,42 @@ use crate::{ emitters::aws::{ helpers::{ cf_from_json, required_logical_id, resource_config, service_account_role_id, - uniquify_iam_statement_sids, + service_trust_policy, stack_name, tags, uniquify_iam_statement_sids, }, service_account::permission_context, }, template::{CfExpression, CfResource}, }; -use alien_core::{import::EmitContext, Ai, ErrorData, PermissionProfile, PermissionSetReference, Result}; +use alien_core::{ + import::EmitContext, Ai, ErrorData, PermissionProfile, PermissionSetReference, Result, Storage, +}; use alien_error::{AlienError, Context, IntoAlienError}; use alien_permissions::{generators::AwsCloudFormationPermissionsGenerator, BindingTarget}; +/// The `ai/finetune` permission set id. When a permission profile references it on +/// this AI resource, the emitter provisions a dedicated Bedrock-trusted IAM role so +/// Bedrock can assume it to read training data and write output. +const AI_FINETUNE_PERMISSION_ID: &str = "ai/finetune"; + #[derive(Debug, Clone, Copy, Default)] pub struct AwsAiEmitter; impl CfEmitter for AwsAiEmitter { fn emit_resources(&self, ctx: &EmitContext<'_>) -> Result> { - resource_config::(ctx, Ai::RESOURCE_TYPE)?; - ai_iam_policies(ctx) + let ai = resource_config::(ctx, Ai::RESOURCE_TYPE)?; + let mut resources = ai_iam_policies(ctx)?; + + // When a permission profile references `ai/finetune` on this resource, emit a + // dedicated IAM role Bedrock can assume for model-customization jobs. The + // stack's service-account roles only trust compute principals + // (`lambda`/`codebuild`/`ec2`), so Bedrock cannot assume them — that is the + // real AccessDenied this role fixes. Its name matches the controller's + // `role_arn` (`{prefix}-{id}-finetune`) so the runtime gateway can pass it. + if resource_references_finetune(ctx) { + resources.push(finetune_role(ctx, ai.id())); + } + + Ok(resources) } fn emit_import_ref(&self, ctx: &EmitContext<'_>) -> Result { @@ -144,3 +163,104 @@ fn ai_permission_refs( refs } +/// True when any permission profile references the `ai/finetune` set on this AI +/// resource. Only then is the Bedrock-trusted finetune role needed. +fn resource_references_finetune(ctx: &EmitContext<'_>) -> bool { + ctx.stack.permission_profiles().values().any(|profile| { + ai_permission_refs(profile, ctx.resource_id) + .iter() + .any(|reference| reference.id() == AI_FINETUNE_PERMISSION_ID) + }) +} + +/// The dedicated Bedrock-trusted finetune IAM role plus its inline S3 policy. +/// +/// The role is named `${AWS::StackName}-{id}-finetune` (matching the controller's +/// `role_arn`), trusts `bedrock.amazonaws.com`, and can read training data +/// (`s3:GetObject`/`s3:ListBucket`) and write output (`s3:PutObject`) on every +/// storage bucket in the stack. +fn finetune_role(ctx: &EmitContext<'_>, ai_id: &str) -> CfResource { + let logical_id = ctx.name_for(ctx.resource_id).unwrap_or(ctx.resource_id); + let role_id = format!("{logical_id}FinetuneRole"); + + let mut role = CfResource::new(role_id, "AWS::IAM::Role".to_string()); + role.properties + .insert("RoleName".to_string(), stack_name(&format!("{ai_id}-finetune"))); + role.properties.insert( + "AssumeRolePolicyDocument".to_string(), + service_trust_policy(["bedrock.amazonaws.com"]), + ); + role.properties.insert( + "Policies".to_string(), + CfExpression::list([CfExpression::object([ + ("PolicyName", stack_name(&format!("{ai_id}-finetune-s3"))), + ( + "PolicyDocument", + CfExpression::object([ + ("Version", CfExpression::from("2012-10-17")), + ("Statement", CfExpression::list(finetune_s3_statements(ctx))), + ]), + ), + ])]), + ); + role.properties.insert("Tags".to_string(), tags(ctx)); + + // The role reads/writes the storage buckets, so it must be created after them. + for bucket_id in storage_bucket_logical_ids(ctx) { + role.depends_on.push(bucket_id); + } + + role +} + +/// S3 statements scoping the finetune role to the stack's storage buckets: +/// list/read on the bucket + objects, put on objects (training in, output out). +fn finetune_s3_statements(ctx: &EmitContext<'_>) -> Vec { + let mut bucket_arns = Vec::new(); + let mut object_arns = Vec::new(); + for bucket_id in storage_bucket_logical_ids(ctx) { + bucket_arns.push(CfExpression::get_att(&bucket_id, "Arn")); + object_arns.push(CfExpression::sub(format!("${{{bucket_id}.Arn}}/*"))); + } + + vec![ + // Read the training dataset: list the bucket and get objects. + CfExpression::object([ + ("Effect", CfExpression::from("Allow")), + ( + "Action", + CfExpression::list([ + CfExpression::from("s3:GetObject"), + CfExpression::from("s3:ListBucket"), + ]), + ), + ( + "Resource", + CfExpression::list( + bucket_arns + .iter() + .cloned() + .chain(object_arns.iter().cloned()), + ), + ), + ]), + // Write the tuning-job output back to storage. + CfExpression::object([ + ("Effect", CfExpression::from("Allow")), + ("Action", CfExpression::from("s3:PutObject")), + ("Resource", CfExpression::list(object_arns)), + ]), + ] +} + +/// Logical ids of every `Storage` (S3 bucket) resource in the stack. +fn storage_bucket_logical_ids(ctx: &EmitContext<'_>) -> Vec { + ctx.stack + .resources() + .filter_map(|(id, entry)| { + entry.config.downcast_ref::()?; + ctx.name_for(id).map(|label| label.to_string()) + }) + .collect() +} + diff --git a/crates/alien-cloudformation/tests/generator/aws_ai_tests.rs b/crates/alien-cloudformation/tests/generator/aws_ai_tests.rs index 653853941..f66b8ac67 100644 --- a/crates/alien-cloudformation/tests/generator/aws_ai_tests.rs +++ b/crates/alien-cloudformation/tests/generator/aws_ai_tests.rs @@ -2,7 +2,9 @@ use super::helpers::render_built_ins; use alien_cloudformation::RegistrationMode; -use alien_core::{Ai, PermissionProfile, ResourceLifecycle, ServiceAccount, Stack, StackSettings}; +use alien_core::{ + Ai, PermissionProfile, ResourceLifecycle, ServiceAccount, Stack, StackSettings, Storage, +}; #[test] fn aws_ai_invoke_permissions_attach_to_service_account_role() { @@ -74,4 +76,80 @@ fn aws_ai_without_permissions_emits_no_iam_policy() { !yaml.contains("AWS::IAM::Policy"), "expected no IAM policy without a permission profile:\n{yaml}" ); + // A pure inference gateway must NOT get a Bedrock-trusted finetune role. + assert!( + !yaml.contains("bedrock.amazonaws.com"), + "expected no bedrock trust policy without ai/finetune:\n{yaml}" + ); +} + +#[test] +fn aws_ai_finetune_emits_bedrock_trusted_role_with_s3_policy() { + // When a permission profile references ai/finetune on the AI resource, the + // emitter provisions a dedicated IAM role Bedrock can assume (the real fix for + // AccessDenied: service-account roles only trust compute principals) with an + // inline S3 policy over the stack's storage buckets. + let stack = Stack::new("ai-finetune".to_string()) + .permission( + "execution", + PermissionProfile::new().resource("llm", ["ai/finetune"]), + ) + .add( + ServiceAccount::new("execution-sa".to_string()).build(), + ResourceLifecycle::Frozen, + ) + .add( + Storage::new("training".to_string()).build(), + ResourceLifecycle::Frozen, + ) + .add( + Ai::new("llm".to_string()) + .finetune(alien_core::FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "training".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: alien_core::FinetuneMethod::Sft, + }) + .build(), + ResourceLifecycle::Frozen, + ) + .build(); + + let yaml = render_built_ins( + &stack, + StackSettings::default(), + RegistrationMode::OutputsFallback, + "aws ai bedrock finetune role", + ); + + // The dedicated finetune role exists and is named deterministically to match + // the controller's role_arn (`{prefix}-{id}-finetune`). + assert!( + yaml.contains("AWS::IAM::Role"), + "expected a dedicated finetune IAM::Role:\n{yaml}" + ); + assert!( + yaml.contains("${AWS::StackName}-llm-finetune"), + "expected role name ${{prefix}}-llm-finetune matching the controller role_arn:\n{yaml}" + ); + // Its trust policy allows Bedrock to assume it — the crux of the fix. + assert!( + yaml.contains("bedrock.amazonaws.com"), + "finetune role must trust bedrock.amazonaws.com:\n{yaml}" + ); + // The inline policy grants S3 read (training data) and write (output). + assert!( + yaml.contains("s3:GetObject") && yaml.contains("s3:ListBucket"), + "finetune role must read the training dataset from S3:\n{yaml}" + ); + assert!( + yaml.contains("s3:PutObject"), + "finetune role must write tuning output to S3:\n{yaml}" + ); + // The S3 grants are scoped to the stack's storage bucket, not "*". + assert!( + yaml.contains("Training") || yaml.contains("training"), + "S3 grants must reference the storage bucket by ARN:\n{yaml}" + ); } diff --git a/crates/alien-core/src/bindings/ai.rs b/crates/alien-core/src/bindings/ai.rs index f4bf118a4..48fd56c72 100644 --- a/crates/alien-core/src/bindings/ai.rs +++ b/crates/alien-core/src/bindings/ai.rs @@ -39,6 +39,38 @@ pub struct TunedModel { pub upstream_id: String, } +/// The fine-tuning capability a managed binding carries when its `Ai` resource +/// declared `.finetune(...)`. This is *not* a completed tuned model — it is the +/// declaration the gateway needs to submit a tuning job at runtime (when the app +/// calls `POST //v1/finetune`) and to rediscover the tuned model by +/// convention once the job completes. Absent for a pure inference gateway. +/// +/// The tuning job runs at runtime, not at deploy time, so this capability travels +/// on the binding while `tuned_model` is populated only after a job succeeds (via +/// rediscovery, so it may be absent even when a capability is present). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +pub struct FinetuneCapability { + /// Provider-native base-model id the job tunes. + pub base_model: String, + /// The object-storage bucket (S3 / GCS / Blob) holding the training dataset. + /// Resolved by the controller from the training-data storage dependency. + pub training_bucket: String, + /// Object key of the training file within `training_bucket`. + pub training_key: String, + /// The public model id the tuned model is served under. + pub served_model_id: String, + /// The deterministic provider-side name of the tuned model / job, used both to + /// submit and to rediscover it (e.g. the Bedrock custom-model + job name). + pub job_name: String, + /// IAM role ARN (AWS) the tuning job assumes to read/write S3. Empty on clouds + /// that submit under the ambient identity without a passed role. + #[serde(default, skip_serializing_if = "String::is_empty")] + pub role_arn: String, +} + /// AWS Bedrock AI binding configuration #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] @@ -51,6 +83,10 @@ pub struct BedrockAiBinding { /// declared a completed fine-tuning job. Absent for a pure inference gateway. #[serde(default, skip_serializing_if = "Option::is_none")] pub tuned_model: Option, + /// The fine-tuning capability the gateway uses to submit and rediscover a + /// runtime tuning job. Present when the resource declared `.finetune(...)`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finetune: Option, } /// GCP Vertex AI binding configuration @@ -67,6 +103,10 @@ pub struct VertexAiBinding { /// declared a completed fine-tuning job. Absent for a pure inference gateway. #[serde(default, skip_serializing_if = "Option::is_none")] pub tuned_model: Option, + /// The fine-tuning capability the gateway uses to submit and rediscover a + /// runtime tuning job. Present when the resource declared `.finetune(...)`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finetune: Option, } /// Azure AI Foundry binding configuration @@ -83,6 +123,10 @@ pub struct FoundryAiBinding { /// declared a completed fine-tuning job. Absent for a pure inference gateway. #[serde(default, skip_serializing_if = "Option::is_none")] pub tuned_model: Option, + /// The fine-tuning capability the gateway uses to submit and rediscover a + /// runtime tuning job. Present when the resource declared `.finetune(...)`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finetune: Option, } /// External AI provider binding configuration (BYO-key). @@ -127,6 +171,7 @@ impl AiBinding { Self::Bedrock(BedrockAiBinding { region: region.into(), tuned_model: None, + finetune: None, }) } @@ -135,6 +180,7 @@ impl AiBinding { project: project.into(), location: location.into(), tuned_model: None, + finetune: None, }) } @@ -143,6 +189,7 @@ impl AiBinding { endpoint: endpoint.into(), account: account.into(), tuned_model: None, + finetune: None, }) } @@ -171,6 +218,26 @@ impl AiBinding { } } + /// Attach a fine-tuning capability to a managed binding, so the gateway can + /// submit and rediscover a runtime tuning job. A no-op on `External`. + pub fn with_finetune(self, capability: FinetuneCapability) -> Self { + match self { + Self::Bedrock(b) => Self::Bedrock(BedrockAiBinding { + finetune: Some(capability), + ..b + }), + Self::Vertex(b) => Self::Vertex(VertexAiBinding { + finetune: Some(capability), + ..b + }), + Self::Foundry(b) => Self::Foundry(FoundryAiBinding { + finetune: Some(capability), + ..b + }), + Self::External(b) => Self::External(b), + } + } + /// The tuned model attached to this binding, if any. pub fn tuned_model(&self) -> Option<&TunedModel> { match self { @@ -181,6 +248,16 @@ impl AiBinding { } } + /// The fine-tuning capability attached to this binding, if any. + pub fn finetune(&self) -> Option<&FinetuneCapability> { + match self { + Self::Bedrock(b) => b.finetune.as_ref(), + Self::Vertex(b) => b.finetune.as_ref(), + Self::Foundry(b) => b.finetune.as_ref(), + Self::External(_) => None, + } + } + pub fn external( provider: impl Into, api_key: impl Into>, diff --git a/crates/alien-core/src/bindings/mod.rs b/crates/alien-core/src/bindings/mod.rs index 9b44528c9..eca2b1d07 100644 --- a/crates/alien-core/src/bindings/mod.rs +++ b/crates/alien-core/src/bindings/mod.rs @@ -28,7 +28,8 @@ mod vault; mod worker; pub use ai::{ - AiBinding, BedrockAiBinding, ExternalAiBinding, FoundryAiBinding, VertexAiBinding, + AiBinding, BedrockAiBinding, ExternalAiBinding, FinetuneCapability, FoundryAiBinding, + TunedModel, VertexAiBinding, }; pub use artifact_registry::{ AcrArtifactRegistryBinding, ArtifactRegistryBinding, EcrArtifactRegistryBinding, diff --git a/crates/alien-core/src/resources/ai.rs b/crates/alien-core/src/resources/ai.rs index 526d5a56f..a1474e055 100644 --- a/crates/alien-core/src/resources/ai.rs +++ b/crates/alien-core/src/resources/ai.rs @@ -36,14 +36,18 @@ impl Default for FinetuneMethod { /// Declares that an [`Ai`] resource should fine-tune a base model in the /// customer's cloud before serving it. /// -/// The training data lives in a customer-owned [`Storage`](crate::Storage) -/// bucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the -/// cloud controller submits the provider's tuning job (Bedrock +/// This is a *capability declaration*, not a deploy-time trigger: a resource +/// with a `finetune` spec provisions and is Ready immediately (no job runs at +/// deploy). The declaration flows to the gateway as a fine-tuning capability; +/// the app then starts a job at runtime by calling `ai("").finetune(...)`, +/// which the gateway submits to the provider (Bedrock /// `CreateModelCustomizationJob`, Vertex tuning job, or Foundry -/// `fine_tuning.jobs`), polls it to completion via the heartbeat loop, and -/// records the resulting artifact so the gateway can route `served_model_id` -/// to it. Base-model inference is unaffected — an `Ai` without a `finetune` -/// spec behaves exactly as before. +/// `fine_tuning.jobs`) under the workload's ambient identity, reading the +/// training data from the customer-owned [`Storage`](crate::Storage) bucket +/// (S3 / GCS / Blob) referenced by `training_data`. Once a job completes, the +/// gateway serves the tuned model under `served_model_id`, rediscovering it by +/// convention. Base-model inference is unaffected — an `Ai` without a +/// `finetune` spec behaves exactly as before. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))] @@ -51,10 +55,10 @@ impl Default for FinetuneMethod { pub struct FinetuneSpec { /// Provider-native base-model identifier to tune (e.g. an Amazon Nova model /// id on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on - /// Foundry). Validated against the target cloud when the job is submitted. + /// Foundry). Validated against the target cloud when a runtime job is submitted. pub base_model: String, - /// The storage resource holding the JSONL training dataset. The controller + /// The storage resource holding the JSONL training dataset. The gateway /// reads it from the customer bucket the storage resolves to; the data /// never leaves the customer's cloud. pub training_data: String, diff --git a/crates/alien-gateway/Cargo.toml b/crates/alien-gateway/Cargo.toml index b8a5465c3..f9c539d0f 100644 --- a/crates/alien-gateway/Cargo.toml +++ b/crates/alien-gateway/Cargo.toml @@ -23,6 +23,8 @@ reqwest = { workspace = true, features = ["json", "stream", "blocking", "rustls- futures = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +async-trait = { workspace = true } +urlencoding = { workspace = true } base64 = { workspace = true } tracing = { workspace = true } aws-config = { workspace = true } diff --git a/crates/alien-gateway/src/config.rs b/crates/alien-gateway/src/config.rs index f3734e776..c4411ea23 100644 --- a/crates/alien-gateway/src/config.rs +++ b/crates/alien-gateway/src/config.rs @@ -158,11 +158,13 @@ fn bindings_from_pairs( /// env-var key encodes it). fn gateway_binding(name: &str, binding: AiBinding) -> Option { // A managed binding may carry a tuned model the gateway serves alongside the - // static catalog. Read it before consuming the binding into its variant. + // static catalog, and/or a fine-tuning capability the control-plane routes use. + // Read both before consuming the binding into its variant. let tuned = binding.tuned_model().map(|t| TunedRoute { served_id: t.served_id.clone(), upstream_id: t.upstream_id.clone(), }); + let finetune = binding.finetune().cloned(); match binding { AiBinding::Bedrock(b) => Some(GatewayBinding { name: name.to_string(), @@ -171,6 +173,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { project: None, azure_endpoint: None, tuned, + finetune, }), AiBinding::Vertex(b) => Some(GatewayBinding { name: name.to_string(), @@ -179,6 +182,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { project: Some(b.project), azure_endpoint: None, tuned, + finetune, }), AiBinding::Foundry(b) => Some(GatewayBinding { name: name.to_string(), @@ -187,6 +191,7 @@ fn gateway_binding(name: &str, binding: AiBinding) -> Option { project: None, azure_endpoint: Some(b.endpoint), tuned, + finetune, }), // External is a BYO-key provider, not an ambient-managed cloud — not served here. AiBinding::External(_) => None, @@ -249,6 +254,7 @@ pub async fn resolve_route(binding: GatewayBinding, managed: Option<&Managed>) - cred, upstream_base_override: None, tuned: binding.tuned, + finetune: binding.finetune, }) } diff --git a/crates/alien-gateway/src/finetune.rs b/crates/alien-gateway/src/finetune.rs new file mode 100644 index 000000000..c8ffe6f0a --- /dev/null +++ b/crates/alien-gateway/src/finetune.rs @@ -0,0 +1,164 @@ +//! Runtime fine-tuning control plane for the AI gateway. +//! +//! Inference is a stateless proxy (see [`crate::router`]); fine-tuning adds a small +//! *control-plane* surface on the same gateway, reusing the same ambient credential. +//! The app triggers a job at runtime: +//! +//! - `POST //v1/finetune` → submit a job, returns `{ jobId, servedModel }` +//! - `GET //v1/finetune/` → poll it, returns `{ status, model? }` +//! +//! The gateway is per-process and stateless: a job started by one worker completes on +//! the cloud's side hours later, possibly after that worker is gone. So nothing about +//! job state is persisted. Instead the tuned model's cloud name is **deterministic** +//! from the binding's [`FinetuneCapability`], and the router rediscovers the completed +//! model by that name on the next inference request (see [`FineTuneProvider::resolve_served_model`]). +//! +//! Each cloud implements [`FineTuneProvider`]; the router dispatches on the route's cloud. + +use alien_core::bindings::FinetuneCapability; +use alien_error::AlienError; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::creds::AmbientCred; +use crate::error::{ErrorData, Result}; + +pub mod bedrock; +pub mod foundry; +pub mod vertex; + +/// The lifecycle status of a fine-tuning job, normalized across providers. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum JobStatus { + /// The job is queued or running. + Running, + /// The job completed and the tuned model is ready to serve. + Succeeded, + /// The job failed; `message` on [`JobState`] carries the reason. + Failed, +} + +/// A submitted job's identity: the provider job id the app polls with. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct JobHandle { + /// Provider job id (a Bedrock job ARN, Vertex tuning-job name, or Foundry job id). + pub job_id: String, + /// The public model id the tuned model will be served under once complete. + pub served_model: String, +} + +/// The observed state of a job, plus the tuned upstream id once it succeeds. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct JobState { + pub status: JobStatus, + /// The tuned model's provider-native upstream id (custom-model ARN / tuned + /// endpoint / deployment name), set only when `status == Succeeded`. + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + /// A human-readable failure reason, set only when `status == Failed`. + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, +} + +/// Everything a provider needs to submit one job: the declared capability plus an +/// optional per-request override of the training key (so an app can retrain on a +/// freshly-uploaded file without redeploying). +pub struct SubmitRequest<'a> { + pub capability: &'a FinetuneCapability, + /// Override for `capability.training_key`; `None` uses the declared key. + pub training_key: Option, +} + +impl SubmitRequest<'_> { + /// The training key this request uses (override or the capability's default). + pub fn training_key(&self) -> &str { + self.training_key + .as_deref() + .unwrap_or(&self.capability.training_key) + } +} + +/// A cloud's fine-tuning control plane. Implementations issue signed control-plane +/// calls via the ambient credential and translate provider responses into the +/// normalized [`JobStatus`]/[`JobState`] above. +#[async_trait] +pub trait FineTuneProvider: Send + Sync { + /// Submit a tuning job. Returns the job id to poll and the public served model id. + async fn submit(&self, request: &SubmitRequest<'_>) -> Result; + + /// Poll a previously-submitted job by its provider job id. + async fn status(&self, job_id: &str) -> Result; + + /// Rediscover the tuned model by its deterministic name, without a job id. + /// Returns the tuned upstream id if the model exists and is ready to serve, + /// `None` if it doesn't exist yet or is still being created. This is how a + /// stateless gateway routes `served_model_id` after the worker that ran the + /// job is gone. + async fn resolve_served_model(&self, capability: &FinetuneCapability) -> Result>; +} + +/// Build a per-request [`FineTuneProvider`] for a route's cloud. The provider borrows +/// the route's ambient credential and the shared HTTP client, so it is cheap to build +/// on each control-plane request. `None` for clouds/bindings the gateway does not +/// fine-tune (currently only Bedrock is wired; Vertex/Foundry follow the same trait). +/// The route location fields a provider needs to build control-plane URLs, plus the +/// optional test base-URL override. Borrowed from the router's `GatewayRoute`. +pub struct ProviderCtx<'a> { + pub cloud: alien_core::Platform, + /// AWS region / Vertex location. + pub region: Option<&'a str>, + /// GCP project id. + pub project: Option<&'a str>, + /// Azure Foundry account endpoint. + pub azure_endpoint: Option<&'a str>, + pub cred: &'a AmbientCred, + pub client: &'a reqwest::Client, + /// Test-only upstream base override (mirrors the inference proxy). + pub base_override: Option<&'a str>, +} + +pub fn provider_for<'a>(ctx: &ProviderCtx<'a>) -> Option> { + match ctx.cloud { + alien_core::Platform::Aws => ctx.region.map(|region| { + let provider = match ctx.base_override { + Some(base) => bedrock::BedrockFineTune::with_base_override( + region.to_string(), + ctx.cred, + ctx.client, + base.to_string(), + ), + None => bedrock::BedrockFineTune::new(region.to_string(), ctx.cred, ctx.client), + }; + Box::new(provider) as Box + }), + alien_core::Platform::Gcp => match (ctx.region, ctx.project) { + (Some(location), Some(project)) => Some(Box::new(vertex::VertexFineTune::new( + location.to_string(), + project.to_string(), + ctx.cred, + ctx.client, + ctx.base_override.map(str::to_string), + )) as Box), + _ => None, + }, + alien_core::Platform::Azure => ctx.azure_endpoint.map(|endpoint| { + Box::new(foundry::FoundryFineTune::new( + endpoint.to_string(), + ctx.cred, + ctx.client, + ctx.base_override.map(str::to_string), + )) as Box + }), + _ => None, + } +} + +/// Map a "no fine-tuning capability on this binding" condition to a gateway error. +pub(crate) fn no_capability(binding: &str) -> AlienError { + AlienError::new(ErrorData::InvalidRequest { + message: format!("binding '{binding}' has no fine-tuning capability declared"), + }) +} diff --git a/crates/alien-gateway/src/finetune/bedrock.rs b/crates/alien-gateway/src/finetune/bedrock.rs new file mode 100644 index 000000000..1f8b5530b --- /dev/null +++ b/crates/alien-gateway/src/finetune/bedrock.rs @@ -0,0 +1,232 @@ +//! AWS Bedrock fine-tuning provider. +//! +//! Submits `CreateModelCustomizationJob`, polls `GetModelCustomizationJob`, and +//! rediscovers the tuned model with `GetCustomModel` (which accepts the model +//! *name*, so the stateless gateway needs no stored ARN). All three are signed +//! with the workload's ambient SigV4 credential for the `bedrock` service against +//! the control-plane host `bedrock.{region}.amazonaws.com` (distinct from the +//! `bedrock-runtime` inference host). + +use alien_core::bindings::FinetuneCapability; +use alien_error::{Context, IntoAlienError}; +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::json; + +use crate::creds::AmbientCred; +use crate::error::{ErrorData, Result}; + +use super::{FineTuneProvider, JobHandle, JobState, JobStatus, SubmitRequest}; + +/// SigV4 service name for the Bedrock control plane. +const BEDROCK_SERVICE: &str = "bedrock"; + +pub struct BedrockFineTune<'a> { + region: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + /// When set, control-plane requests target this base URL instead of the + /// region-derived Bedrock host. Lets tests aim the provider at a mock upstream + /// (mirrors the inference proxy's `upstream_base_override`). + base_override: Option, +} + +impl<'a> BedrockFineTune<'a> { + pub fn new(region: String, cred: &'a AmbientCred, client: &'a reqwest::Client) -> Self { + Self { region, cred, client, base_override: None } + } + + /// Build a provider aimed at an explicit base URL (test upstream). + pub fn with_base_override( + region: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + base: String, + ) -> Self { + Self { region, cred, client, base_override: Some(base) } + } + + /// Bedrock control-plane host (not the `bedrock-runtime` inference host). + fn host(&self) -> String { + self.base_override + .clone() + .unwrap_or_else(|| format!("https://bedrock.{}.amazonaws.com", self.region)) + } + + /// Sign `req` with the ambient credential for the `bedrock` service and execute it, + /// returning the parsed JSON body on 2xx or a contextual error otherwise. + async fn send(&self, mut req: reqwest::Request) -> Result { + let url = req.url().to_string(); + self.cred.authorize(&mut req, BEDROCK_SERVICE).await?; + let resp = self + .client + .execute(req) + .await + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Bedrock control-plane request to {url} failed"), + })?; + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + if !status.is_success() { + return Err(alien_error::AlienError::new(ErrorData::UpstreamFailed { + // Bedrock returns a JSON `{message, __type}`; surface it verbatim. + message: format!("Bedrock {status} for {url}: {body}"), + })); + } + if body.is_empty() { + return Ok(serde_json::Value::Null); + } + serde_json::from_str(&body) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Bedrock returned non-JSON from {url}: {body}"), + }) + } + + fn build_json_post(&self, path: &str, body: serde_json::Value) -> Result { + self.client + .post(format!("{}{path}", self.host())) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .json(&body) + .build() + .into_alien_error() + .context(ErrorData::Other { + message: format!("could not build Bedrock POST {path}"), + }) + } + + fn build_get(&self, path: &str) -> Result { + self.client + .get(format!("{}{path}", self.host())) + .build() + .into_alien_error() + .context(ErrorData::Other { + message: format!("could not build Bedrock GET {path}"), + }) + } +} + +/// `CreateModelCustomizationJob` response — only the job ARN is needed. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct CreateJobResponse { + job_arn: String, +} + +/// `GetModelCustomizationJob` response — status plus the produced custom-model ARN. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct GetJobResponse { + status: String, + #[serde(default)] + output_model_arn: Option, + #[serde(default)] + failure_message: Option, +} + +/// `GetCustomModel` response — status plus the model ARN for rediscovery. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct GetCustomModelResponse { + model_arn: String, + model_status: String, +} + +#[async_trait] +impl FineTuneProvider for BedrockFineTune<'_> { + async fn submit(&self, request: &SubmitRequest<'_>) -> Result { + let cap = request.capability; + let training_uri = format!("s3://{}/{}", cap.training_bucket, request.training_key()); + let output_uri = format!("s3://{}/alien-finetune-output/{}/", cap.training_bucket, cap.job_name); + + // clientRequestToken makes the submit idempotent so a retried POST /finetune + // (or a transport retry) doesn't create a second job for the same run. + let body = json!({ + "jobName": cap.job_name, + "customModelName": cap.job_name, + "roleArn": cap.role_arn, + "baseModelIdentifier": cap.base_model, + "customizationType": "FINE_TUNING", + "clientRequestToken": cap.job_name, + "trainingDataConfig": { "s3Uri": training_uri }, + "outputDataConfig": { "s3Uri": output_uri }, + }); + + let req = self.build_json_post("/model-customization-jobs", body)?; + let value = self.send(req).await?; + let parsed: CreateJobResponse = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Bedrock CreateModelCustomizationJob returned no jobArn".to_string(), + })?; + + Ok(JobHandle { + job_id: parsed.job_arn, + served_model: cap.served_model_id.clone(), + }) + } + + async fn status(&self, job_id: &str) -> Result { + // jobIdentifier can be the job ARN; percent-encode it for the path. + let encoded = urlencoding::encode(job_id); + let req = self.build_get(&format!("/model-customization-jobs/{encoded}"))?; + let value = self.send(req).await?; + let parsed: GetJobResponse = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Bedrock GetModelCustomizationJob returned an unexpected body".to_string(), + })?; + + Ok(match parsed.status.as_str() { + "Completed" => JobState { + status: JobStatus::Succeeded, + model: parsed.output_model_arn, + message: None, + }, + "Failed" | "Stopped" => JobState { + status: JobStatus::Failed, + model: None, + message: parsed.failure_message, + }, + // InProgress | Stopping | anything else -> still running. + _ => JobState { status: JobStatus::Running, model: None, message: None }, + }) + } + + async fn resolve_served_model(&self, cap: &FinetuneCapability) -> Result> { + // GetCustomModel accepts the model NAME, so the deterministic job_name (also the + // custom-model name) rediscovers the model without any stored ARN. A 404 means + // the job hasn't produced the model yet. + let encoded = urlencoding::encode(&cap.job_name); + let req = self.build_get(&format!("/custom-models/{encoded}"))?; + match self.send(req).await { + Ok(value) => { + let parsed: GetCustomModelResponse = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Bedrock GetCustomModel returned an unexpected body".to_string(), + })?; + Ok(match parsed.model_status.as_str() { + "Active" => Some(parsed.model_arn), + // Creating / Failed -> not servable yet. + _ => None, + }) + } + // A ResourceNotFound (model not created yet) is not an error for rediscovery. + Err(e) if is_not_found(&e) => Ok(None), + Err(e) => Err(e), + } + } +} + +/// True if the error is a Bedrock 404 / ResourceNotFound, which for rediscovery means +/// "the tuned model doesn't exist yet", not a real failure. +fn is_not_found(err: &alien_error::AlienError) -> bool { + match &err.error { + Some(ErrorData::UpstreamFailed { message }) => { + message.contains("404") || message.contains("ResourceNotFound") + } + _ => false, + } +} diff --git a/crates/alien-gateway/src/finetune/foundry.rs b/crates/alien-gateway/src/finetune/foundry.rs new file mode 100644 index 000000000..9f65fa08e --- /dev/null +++ b/crates/alien-gateway/src/finetune/foundry.rs @@ -0,0 +1,224 @@ +//! Azure AI Foundry fine-tuning provider (data plane). +//! +//! Submits `POST {endpoint}/openai/fine_tuning/jobs`, polls `GET .../jobs/{id}`, and +//! rediscovers the tuned model as the deployment named after `served_model_id`. All +//! calls carry the workload's ambient bearer token for the cognitive-services data +//! plane; the SigV4 service name is unused. +//! +//! Note: Foundry can import training data from Blob, but that requires the storage +//! account to allow public network access — a gotcha documented in the example README. + +use alien_core::bindings::FinetuneCapability; +use alien_error::{Context, IntoAlienError}; +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::json; + +use crate::creds::AmbientCred; +use crate::error::{ErrorData, Result}; + +use super::{FineTuneProvider, JobHandle, JobState, JobStatus, SubmitRequest}; + +/// Bearer auth ignores the SigV4 service name. +const UNUSED_SERVICE: &str = "cognitiveservices"; +/// Data-plane API version for fine-tuning + deployments. +const API_VERSION: &str = "2024-10-21"; + +pub struct FoundryFineTune<'a> { + endpoint: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + base_override: Option, +} + +impl<'a> FoundryFineTune<'a> { + pub fn new( + endpoint: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + base_override: Option, + ) -> Self { + Self { endpoint, cred, client, base_override } + } + + /// The account data-plane base (or the test override), without a trailing slash. + fn base(&self) -> String { + self.base_override + .clone() + .unwrap_or_else(|| self.endpoint.clone()) + .trim_end_matches('/') + .to_string() + } + + async fn send(&self, mut req: reqwest::Request) -> Result { + let url = req.url().to_string(); + self.cred.authorize(&mut req, UNUSED_SERVICE).await?; + let resp = self + .client + .execute(req) + .await + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Foundry fine-tuning request to {url} failed"), + })?; + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + if !status.is_success() { + return Err(alien_error::AlienError::new(ErrorData::UpstreamFailed { + message: format!("Foundry {status} for {url}: {body}"), + })); + } + if body.is_empty() { + return Ok(serde_json::Value::Null); + } + serde_json::from_str(&body) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Foundry returned non-JSON from {url}: {body}"), + }) + } +} + +/// A fine-tuning job — `id` to poll, `status`, and `fine_tuned_model` once done. +#[derive(Deserialize)] +struct FineTuningJob { + #[serde(default)] + id: Option, + #[serde(default)] + status: Option, + #[serde(default)] + fine_tuned_model: Option, + #[serde(default)] + error: Option, +} + +#[derive(Deserialize)] +struct JobError { + #[serde(default)] + message: Option, +} + +/// A deployment resource — `provisioningState` tells us if the tuned model is servable. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Deployment { + #[serde(default)] + provisioning_state: Option, +} + +fn map_status(status: Option<&str>, model: Option, err: Option) -> JobState { + match status { + Some("succeeded") => JobState { + status: JobStatus::Succeeded, + model, + message: None, + }, + Some("failed") | Some("cancelled") => { + JobState { status: JobStatus::Failed, model: None, message: err } + } + _ => JobState { status: JobStatus::Running, model: None, message: None }, + } +} + +#[async_trait] +impl FineTuneProvider for FoundryFineTune<'_> { + async fn submit(&self, request: &SubmitRequest<'_>) -> Result { + let cap = request.capability; + // The training file is a Blob URL under the training bucket/container. Foundry + // requires the storage account to allow public network access for Blob import. + let training_file = format!( + "https://{}.blob.core.windows.net/{}", + cap.training_bucket, + request.training_key() + ); + let body = json!({ + "model": cap.base_model, + "training_file": training_file, + "suffix": cap.served_model_id, + }); + let url = format!("{}/openai/fine_tuning/jobs?api-version={API_VERSION}", self.base()); + let req = self + .client + .post(url) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .json(&body) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Foundry fine_tuning POST".to_string() })?; + let value = self.send(req).await?; + let parsed: FineTuningJob = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Foundry CreateFineTuningJob returned no id".to_string(), + })?; + let job_id = parsed.id.ok_or_else(|| { + alien_error::AlienError::new(ErrorData::UpstreamFailed { + message: "Foundry fine-tuning job response had no id".to_string(), + }) + })?; + Ok(JobHandle { job_id, served_model: cap.served_model_id.clone() }) + } + + async fn status(&self, job_id: &str) -> Result { + let encoded = urlencoding::encode(job_id); + let url = format!( + "{}/openai/fine_tuning/jobs/{encoded}?api-version={API_VERSION}", + self.base() + ); + let req = self + .client + .get(url) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Foundry GET job".to_string() })?; + let value = self.send(req).await?; + let parsed: FineTuningJob = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Foundry GetFineTuningJob returned an unexpected body".to_string(), + })?; + Ok(map_status( + parsed.status.as_deref(), + parsed.fine_tuned_model, + parsed.error.and_then(|e| e.message), + )) + } + + async fn resolve_served_model(&self, cap: &FinetuneCapability) -> Result> { + // The tuned model is served as a deployment named `served_model_id`; if it + // exists and is Succeeded, that deployment name is the OpenAI-path `model`. + let encoded = urlencoding::encode(&cap.served_model_id); + let url = format!( + "{}/openai/deployments/{encoded}?api-version={API_VERSION}", + self.base() + ); + let req = self + .client + .get(url) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Foundry GET deployment".to_string() })?; + match self.send(req).await { + Ok(value) => { + let dep: Deployment = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Foundry GetDeployment returned an unexpected body".to_string(), + })?; + Ok(match dep.provisioning_state.as_deref() { + Some("Succeeded") => Some(cap.served_model_id.clone()), + _ => None, + }) + } + Err(e) if is_not_found(&e) => Ok(None), + Err(e) => Err(e), + } + } +} + +fn is_not_found(err: &alien_error::AlienError) -> bool { + matches!( + &err.error, + Some(ErrorData::UpstreamFailed { message }) if message.contains("404") + ) +} diff --git a/crates/alien-gateway/src/finetune/vertex.rs b/crates/alien-gateway/src/finetune/vertex.rs new file mode 100644 index 000000000..c069e7252 --- /dev/null +++ b/crates/alien-gateway/src/finetune/vertex.rs @@ -0,0 +1,243 @@ +//! GCP Vertex AI fine-tuning provider. +//! +//! Submits `POST tuningJobs`, polls `GET tuningJobs/{id}`, and rediscovers the tuned +//! model by listing tuning jobs filtered on the deterministic display name (the +//! stateless gateway keeps no job id). All calls carry the workload's ambient bearer +//! token; the SigV4 service name is unused for bearer auth. +//! +//! Vertex supports only supervised tuning for Gemini (no user-selectable LoRA/DPO), +//! so `method` is not forwarded; the job is always supervised. + +use alien_core::bindings::FinetuneCapability; +use alien_error::{Context, IntoAlienError}; +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::json; + +use crate::creds::AmbientCred; +use crate::error::{ErrorData, Result}; + +use super::{FineTuneProvider, JobHandle, JobState, JobStatus, SubmitRequest}; + +/// Bearer auth ignores the SigV4 service name. +const UNUSED_SERVICE: &str = "aiplatform"; + +pub struct VertexFineTune<'a> { + location: String, + project: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + base_override: Option, +} + +impl<'a> VertexFineTune<'a> { + pub fn new( + location: String, + project: String, + cred: &'a AmbientCred, + client: &'a reqwest::Client, + base_override: Option, + ) -> Self { + Self { location, project, cred, client, base_override } + } + + /// Regional Vertex host (`global` uses the un-prefixed host), or the test override. + fn host(&self) -> String { + if let Some(base) = &self.base_override { + return base.clone(); + } + if self.location == "global" { + "https://aiplatform.googleapis.com".to_string() + } else { + format!("https://{}-aiplatform.googleapis.com", self.location) + } + } + + fn jobs_path(&self) -> String { + format!( + "/v1/projects/{}/locations/{}/tuningJobs", + self.project, self.location + ) + } + + async fn send(&self, mut req: reqwest::Request) -> Result { + let url = req.url().to_string(); + self.cred.authorize(&mut req, UNUSED_SERVICE).await?; + let resp = self + .client + .execute(req) + .await + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Vertex tuning request to {url} failed"), + })?; + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + if !status.is_success() { + return Err(alien_error::AlienError::new(ErrorData::UpstreamFailed { + message: format!("Vertex {status} for {url}: {body}"), + })); + } + if body.is_empty() { + return Ok(serde_json::Value::Null); + } + serde_json::from_str(&body) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: format!("Vertex returned non-JSON from {url}: {body}"), + }) + } +} + +/// A tuning job resource — `name` is the job id, `state` the lifecycle, `tunedModel` +/// the result once succeeded. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct TuningJob { + #[serde(default)] + name: Option, + #[serde(default)] + state: Option, + #[serde(default)] + tuned_model: Option, + #[serde(default)] + error: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct TunedModelRef { + /// The deployed tuned-model endpoint the OpenAI-compat path targets. + #[serde(default)] + endpoint: Option, + /// The tuned model resource name (fallback if no endpoint). + #[serde(default)] + model: Option, +} + +impl TunedModelRef { + fn upstream_id(&self) -> Option { + self.endpoint.clone().or_else(|| self.model.clone()) + } +} + +#[derive(Deserialize)] +struct StatusError { + #[serde(default)] + message: Option, +} + +/// List response for rediscovery. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct ListTuningJobs { + #[serde(default)] + tuning_jobs: Vec, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct TuningJobWithDisplay { + #[serde(default)] + tuned_model_display_name: Option, + #[serde(default)] + state: Option, + #[serde(default)] + tuned_model: Option, +} + +fn map_state(state: Option<&str>, tuned: Option<&TunedModelRef>, err: Option) -> JobState { + match state { + Some("JOB_STATE_SUCCEEDED") => JobState { + status: JobStatus::Succeeded, + model: tuned.and_then(TunedModelRef::upstream_id), + message: None, + }, + Some("JOB_STATE_FAILED") + | Some("JOB_STATE_CANCELLED") + | Some("JOB_STATE_EXPIRED") + | Some("JOB_STATE_PARTIALLY_SUCCEEDED") => { + JobState { status: JobStatus::Failed, model: None, message: err } + } + _ => JobState { status: JobStatus::Running, model: None, message: None }, + } +} + +#[async_trait] +impl FineTuneProvider for VertexFineTune<'_> { + async fn submit(&self, request: &SubmitRequest<'_>) -> Result { + let cap = request.capability; + let training_uri = format!("gs://{}/{}", cap.training_bucket, request.training_key()); + let body = json!({ + "baseModel": cap.base_model, + "supervisedTuningSpec": { "trainingDatasetUri": training_uri }, + "tunedModelDisplayName": cap.job_name, + }); + let req = self + .client + .post(format!("{}{}", self.host(), self.jobs_path())) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .json(&body) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Vertex tuningJobs POST".to_string() })?; + let value = self.send(req).await?; + let parsed: TuningJob = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Vertex CreateTuningJob returned no job name".to_string(), + })?; + let job_id = parsed.name.ok_or_else(|| { + alien_error::AlienError::new(ErrorData::UpstreamFailed { + message: "Vertex tuning job response had no name".to_string(), + }) + })?; + Ok(JobHandle { job_id, served_model: cap.served_model_id.clone() }) + } + + async fn status(&self, job_id: &str) -> Result { + // job_id is the full resource name; GET the host + "/v1/{name}". + let req = self + .client + .get(format!("{}/v1/{}", self.host(), job_id)) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Vertex GET tuningJob".to_string() })?; + let value = self.send(req).await?; + let parsed: TuningJob = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Vertex GetTuningJob returned an unexpected body".to_string(), + })?; + Ok(map_state( + parsed.state.as_deref(), + parsed.tuned_model.as_ref(), + parsed.error.and_then(|e| e.message), + )) + } + + async fn resolve_served_model(&self, cap: &FinetuneCapability) -> Result> { + // No stored job id: list tuning jobs and find the succeeded one whose display + // name matches the deterministic job_name, then take its tuned-model endpoint. + let req = self + .client + .get(format!("{}{}", self.host(), self.jobs_path())) + .build() + .into_alien_error() + .context(ErrorData::Other { message: "could not build Vertex ListTuningJobs".to_string() })?; + let value = self.send(req).await?; + let list: ListTuningJobs = serde_json::from_value(value) + .into_alien_error() + .context(ErrorData::UpstreamFailed { + message: "Vertex ListTuningJobs returned an unexpected body".to_string(), + })?; + Ok(list + .tuning_jobs + .into_iter() + .find(|j| { + j.tuned_model_display_name.as_deref() == Some(&cap.job_name) + && j.state.as_deref() == Some("JOB_STATE_SUCCEEDED") + }) + .and_then(|j| j.tuned_model.and_then(|t| t.upstream_id()))) + } +} diff --git a/crates/alien-gateway/src/lib.rs b/crates/alien-gateway/src/lib.rs index ccd15fb16..edb416d7b 100644 --- a/crates/alien-gateway/src/lib.rs +++ b/crates/alien-gateway/src/lib.rs @@ -9,6 +9,7 @@ mod config; mod creds; mod error; +mod finetune; mod router; pub use config::{bindings_from_env, bindings_from_env_map}; pub use creds::{AmbientCred, AwsSigV4Cred, BearerTokenCred}; @@ -47,6 +48,9 @@ pub struct GatewayBinding { /// A tuned model this binding serves alongside the static catalog, produced /// by a completed fine-tuning job. `None` for a pure inference gateway. pub tuned: Option, + /// The fine-tuning capability the runtime control-plane routes use. `None` + /// unless the binding declared `.finetune(...)`. + pub finetune: Option, } /// A tuned model the gateway routes to: the public id an app requests mapped to diff --git a/crates/alien-gateway/src/router.rs b/crates/alien-gateway/src/router.rs index 15d111836..3d26b8d53 100644 --- a/crates/alien-gateway/src/router.rs +++ b/crates/alien-gateway/src/router.rs @@ -49,6 +49,9 @@ pub struct GatewayRoute { /// before the catalog so a completed fine-tuning job's artifact is reachable /// under its public `served_id`. `None` for a pure inference gateway. pub tuned: Option, + /// The fine-tuning capability the runtime control-plane routes use to submit + /// and rediscover jobs. `None` if the binding declared no `.finetune(...)`. + pub finetune: Option, } struct AppState { @@ -70,9 +73,96 @@ pub fn build_router(routes: Vec) -> Router { .route("/{binding}/v1/messages", post(proxy)) .route("/{binding}/v1/responses", post(proxy_responses)) .route("/{binding}/v1/models", get(list_models)) + // Runtime fine-tuning control plane (see `crate::finetune`). + .route("/{binding}/v1/finetune", post(submit_finetune)) + .route("/{binding}/v1/finetune/{job}", get(finetune_status)) .with_state(state) } +/// `POST //v1/finetune` — submit a tuning job for the binding's declared +/// capability. Body: optional `{ "trainingKey": "..." }` to override the training +/// file. Returns `{ "jobId", "servedModel" }`. +async fn submit_finetune( + State(state): State>, + Path(binding): Path, + body: Bytes, +) -> Result { + let route = state.routes.get(&binding).ok_or_else(|| { + AlienError::new(ErrorData::UnknownBinding { binding: binding.clone() }) + })?; + let capability = route + .finetune + .as_ref() + .ok_or_else(|| crate::finetune::no_capability(&binding))?; + let provider = crate::finetune::provider_for(&crate::finetune::ProviderCtx { + cloud: route.cloud, + region: route.region.as_deref(), + project: route.project.as_deref(), + azure_endpoint: route.azure_endpoint.as_deref(), + cred: &route.cred, + client: &state.client, + base_override: route.upstream_base_override.as_deref(), + }) + .ok_or_else(|| { + AlienError::new(ErrorData::InvalidRequest { + message: format!("the gateway does not fine-tune {:?} bindings yet", route.cloud), + }) + })?; + + // An empty body is fine (use the declared training key); otherwise parse the override. + let training_key = if body.is_empty() { + None + } else { + let parsed: SubmitFinetuneBody = serde_json::from_slice(&body).map_err(|e| { + AlienError::new(ErrorData::InvalidRequest { + message: format!("invalid finetune body: {e}"), + }) + })?; + parsed.training_key + }; + + let handle = provider + .submit(&crate::finetune::SubmitRequest { capability, training_key }) + .await?; + Ok(Json(handle).into_response()) +} + +/// `GET //v1/finetune/` — poll a submitted job. Returns +/// `{ "status", "model"?, "message"? }`. +async fn finetune_status( + State(state): State>, + Path((binding, job)): Path<(String, String)>, +) -> Result { + let route = state.routes.get(&binding).ok_or_else(|| { + AlienError::new(ErrorData::UnknownBinding { binding: binding.clone() }) + })?; + let provider = crate::finetune::provider_for(&crate::finetune::ProviderCtx { + cloud: route.cloud, + region: route.region.as_deref(), + project: route.project.as_deref(), + azure_endpoint: route.azure_endpoint.as_deref(), + cred: &route.cred, + client: &state.client, + base_override: route.upstream_base_override.as_deref(), + }) + .ok_or_else(|| { + AlienError::new(ErrorData::InvalidRequest { + message: format!("the gateway does not fine-tune {:?} bindings yet", route.cloud), + }) + })?; + + let state_out = provider.status(&job).await?; + Ok(Json(state_out).into_response()) +} + +/// Optional body for `POST /finetune`. +#[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct SubmitFinetuneBody { + #[serde(default)] + training_key: Option, +} + /// Parse a proxied request body as JSON and pull out its required `model` field. /// Both the chat/completions|messages handler and the Responses handler route on /// the request's `model`, so they share this preamble. @@ -174,11 +264,36 @@ async fn proxy( // `upstream_id` is a custom-model id, on Vertex a tuned-endpoint model id, on // Foundry a deployment name — in every case it becomes the request body's // `model` against the cloud's OpenAI endpoint. + // + // Two sources, in order: + // 1. A binding that already carries a resolved `tuned` model (deploy-time + // path / test injection) routes straight to it. + // 2. Otherwise, if the request's model matches the fine-tuning capability's + // `served_model_id`, rediscover the completed tuned model by convention + // (stateless — no stored ARN). If it isn't ready yet, fall through to the + // catalog, which yields `ModelNotAvailable` for the tuned id. if let Some(tuned) = &route.tuned { if model == tuned.served_id { return proxy_openai(&state, route, tuned.upstream_id.clone(), payload).await; } } + if let Some(cap) = &route.finetune { + if model == cap.served_model_id { + if let Some(provider) = crate::finetune::provider_for(&crate::finetune::ProviderCtx { + cloud: route.cloud, + region: route.region.as_deref(), + project: route.project.as_deref(), + azure_endpoint: route.azure_endpoint.as_deref(), + cred: &route.cred, + client: &state.client, + base_override: route.upstream_base_override.as_deref(), + }) { + if let Some(upstream) = provider.resolve_served_model(cap).await? { + return proxy_openai(&state, route, upstream, payload).await; + } + } + } + } // Cloud-scoped resolution: Claude ids appear once per cloud, so a first-match // resolve would always land on another cloud's entry and fail the cloud filter. @@ -1001,6 +1116,7 @@ mod tests { cred: test_aws_cred(), upstream_base_override: Some(upstream.to_string()), tuned: None, + finetune: None, } } @@ -1014,6 +1130,7 @@ mod tests { cred: AmbientCred::Bearer(BearerTokenCred::static_token("t")), upstream_base_override: None, tuned: None, + finetune: None, } } @@ -1189,6 +1306,7 @@ mod tests { cred: AmbientCred::Bearer(BearerTokenCred::static_token("t")), upstream_base_override: None, tuned: None, + finetune: None, } } @@ -1301,6 +1419,99 @@ mod tests { mock.assert_async().await; } + fn bedrock_finetune_capability() -> alien_core::bindings::FinetuneCapability { + alien_core::bindings::FinetuneCapability { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_bucket: "my-bucket".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: "support-tuned".to_string(), + job_name: "ai-finetune-llm".to_string(), + role_arn: "arn:aws:iam::123456789012:role/ft".to_string(), + } + } + + #[tokio::test] + async fn submit_finetune_posts_customization_job() { + // POST /finetune submits CreateModelCustomizationJob with the capability's + // names, role, base model, and the S3 training/output URIs, signed for bedrock. + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(POST) + .path("/model-customization-jobs") + .body_contains("\"jobName\":\"ai-finetune-llm\"") + .body_contains("\"baseModelIdentifier\":\"amazon.nova-lite-v1:0\"") + .body_contains("s3://my-bucket/training.jsonl") + .body_contains("arn:aws:iam::123456789012:role/ft") + .header_exists("authorization"); + then.status(201) + .header("content-type", "application/json") + .body(r#"{"jobArn":"arn:aws:bedrock:us-east-2:123456789012:model-customization-job/nova.x/abc"}"#); + }) + .await; + + let mut route = aws_route(&server.base_url()); + route.finetune = Some(bedrock_finetune_capability()); + + let url = serve(build_router(vec![route])).await; + let resp = reqwest::Client::new() + .post(format!("{url}/llm/v1/finetune")) + .json(&json!({})) + .send() + .await + .expect("submit request"); + + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["servedModel"], "support-tuned"); + assert!(body["jobId"].as_str().unwrap().contains("model-customization-job")); + mock.assert_async().await; + } + + #[tokio::test] + async fn submit_finetune_without_capability_is_rejected() { + // A binding with no finetune capability rejects the submit rather than 500ing. + let server = MockServer::start_async().await; + let route = aws_route(&server.base_url()); // finetune: None + let url = serve(build_router(vec![route])).await; + let resp = reqwest::Client::new() + .post(format!("{url}/llm/v1/finetune")) + .json(&json!({})) + .send() + .await + .expect("submit request"); + assert_eq!(resp.status(), 400); + } + + #[tokio::test] + async fn finetune_status_maps_completed_to_succeeded() { + // GET /finetune/{job} maps Bedrock "Completed" + outputModelArn to succeeded+model. + let server = MockServer::start_async().await; + let mock = server + .mock_async(|when, then| { + when.method(GET).path_contains("/model-customization-jobs/"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"status":"Completed","outputModelArn":"arn:aws:bedrock:us-east-2:123456789012:custom-model/nova.x/abc"}"#); + }) + .await; + + let mut route = aws_route(&server.base_url()); + route.finetune = Some(bedrock_finetune_capability()); + let url = serve(build_router(vec![route])).await; + let resp = reqwest::Client::new() + .get(format!("{url}/llm/v1/finetune/job-123")) + .send() + .await + .expect("status request"); + + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["status"], "succeeded"); + assert!(body["model"].as_str().unwrap().contains("custom-model")); + mock.assert_async().await; + } + #[tokio::test] async fn tuned_model_routes_to_upstream_artifact() { // A request for the tuned model's public served id must be rewritten to the diff --git a/crates/alien-gateway/tests/integration.rs b/crates/alien-gateway/tests/integration.rs index 6da65209e..9a0380ea7 100644 --- a/crates/alien-gateway/tests/integration.rs +++ b/crates/alien-gateway/tests/integration.rs @@ -81,6 +81,7 @@ async fn routes_two_clouds_with_rewrite_auth_and_passthrough() { cred: aws_cred(), upstream_base_override: Some(aws_upstream.base_url()), tuned: None, + finetune: None, }, GatewayRoute { name: "azllm".to_string(), @@ -91,6 +92,7 @@ async fn routes_two_clouds_with_rewrite_auth_and_passthrough() { cred: AmbientCred::Bearer(BearerTokenCred::static_token("test-azure-token")), upstream_base_override: Some(azure_upstream.base_url()), tuned: None, + finetune: None, }, ]; diff --git a/crates/alien-gateway/tests/live_bedrock.rs b/crates/alien-gateway/tests/live_bedrock.rs index 94c215f77..c52aa2a3b 100644 --- a/crates/alien-gateway/tests/live_bedrock.rs +++ b/crates/alien-gateway/tests/live_bedrock.rs @@ -44,6 +44,7 @@ async fn live_bedrock_openai_chat() { cred, upstream_base_override: None, tuned: None, + finetune: None, }; let base = serve(build_router(vec![route])).await; @@ -100,6 +101,7 @@ async fn live_bedrock_claude_streaming() { cred, upstream_base_override: None, tuned: None, + finetune: None, }; let base = serve(build_router(vec![route])).await; diff --git a/crates/alien-gateway/tests/live_foundry_claude.rs b/crates/alien-gateway/tests/live_foundry_claude.rs index 74e467eea..c4a923b75 100644 --- a/crates/alien-gateway/tests/live_foundry_claude.rs +++ b/crates/alien-gateway/tests/live_foundry_claude.rs @@ -49,6 +49,7 @@ fn foundry_route() -> GatewayRoute { cred: AmbientCred::Bearer(BearerTokenCred::static_token(token)), upstream_base_override: None, tuned: None, + finetune: None, } } diff --git a/crates/alien-gateway/tests/live_vertex_claude.rs b/crates/alien-gateway/tests/live_vertex_claude.rs index c5486db5a..1b65fb4b9 100644 --- a/crates/alien-gateway/tests/live_vertex_claude.rs +++ b/crates/alien-gateway/tests/live_vertex_claude.rs @@ -41,6 +41,7 @@ fn vertex_route() -> GatewayRoute { cred: AmbientCred::Bearer(BearerTokenCred::static_token(token)), upstream_base_override: None, tuned: None, + finetune: None, } } diff --git a/crates/alien-infra/src/ai/aws.rs b/crates/alien-infra/src/ai/aws.rs index 755d9ad07..087bcf687 100644 --- a/crates/alien-infra/src/ai/aws.rs +++ b/crates/alien-infra/src/ai/aws.rs @@ -1,50 +1,29 @@ use std::time::Duration; -use tracing::{info, warn}; +use tracing::info; use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; -use alien_aws_clients::bedrock::{ - CreateModelCustomizationJobRequest, ModelCustomizationJobStatus, S3DataConfig, -}; use alien_core::{ - bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, - AwsBedrockAiHeartbeatData, FinetuneMethod, HeartbeatBackend, Platform, ResourceHeartbeat, - ResourceHeartbeatData, ResourceOutputs, ResourceRef, ResourceStatus, Storage, + bindings::{AiBinding, FinetuneCapability}, + Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, AwsBedrockAiHeartbeatData, HeartbeatBackend, + Platform, ResourceHeartbeat, ResourceHeartbeatData, ResourceOutputs, ResourceRef, + ResourceStatus, Storage, }; use alien_error::{AlienError, Context, IntoAlienError}; use alien_macros::controller; use chrono::Utc; -/// Poll interval while a Bedrock model-customization job runs. Matches the -/// heartbeat cadence used elsewhere in this controller. -const TUNING_POLL_INTERVAL: Duration = Duration::from_secs(30); - -/// Bedrock `customizationType` for a supervised / preference / LoRA fine-tune. -/// Bedrock's public model-customization API exposes `FINE_TUNING`; the specific -/// technique (SFT vs DPO vs LoRA) is selected per base model via hyperparameters, -/// not a distinct customizationType, so all `FinetuneMethod` variants map here. -fn bedrock_customization_type(_method: FinetuneMethod) -> &'static str { - "FINE_TUNING" -} - #[controller] pub struct AwsAiController { /// AWS region where Bedrock is accessed. None until create_start runs. pub(crate) region: Option, - /// ARN of the submitted Bedrock model-customization job, set once - /// `SubmittingTuningJob` succeeds. Used as the poll identifier. Only ever set - /// for a finetune-enabled resource. - pub(crate) tuning_job_arn: Option, - /// The completed custom-model ARN Bedrock returns. This is the `upstream_id` - /// the OpenAI chat endpoint accepts for the tuned model, and is attached to - /// the binding via `with_tuned_model`. None until the job completes. - pub(crate) tuned_model_id: Option, - /// The gateway-facing served model id the tuned model is exposed under - /// (`spec.served_model_id_or_default(&config.id)`). Captured when the job - /// completes so `get_binding_params` (which has no `ctx`) can pair it with - /// `tuned_model_id`. None for a pure inference gateway. - pub(crate) served_id: Option, + /// The fine-tuning capability this resource carries on its binding, resolved + /// during the create flow from the training-storage dependency and the + /// deterministic finetune role name. `get_binding_params` is a pure function of + /// controller state (no `ctx`), so it is captured here rather than resolved + /// live. None for a pure inference gateway. + pub(crate) finetune: Option, } #[controller] @@ -93,197 +72,21 @@ impl AwsAiController { info!(ai=%config.id, "Successfully applied resource-scoped permissions"); - // A pure inference gateway is Ready as soon as permissions are applied — - // exactly as before. A finetune-enabled resource additionally submits and - // polls a Bedrock model-customization job before serving. - let next_state = if config.finetune.is_some() { - info!(ai=%config.id, "Finetune requested; submitting Bedrock model-customization job"); - SubmittingTuningJob - } else { - Ready - }; + // A finetune-enabled resource carries a `FinetuneCapability` on its binding so + // the gateway can submit and rediscover a runtime tuning job. Resolve it here + // (where `ctx` is available) and stash it on the controller; `get_binding_params` + // is a pure function of state and reads it back. The tuning job itself is NOT + // submitted here — it is triggered at runtime via the gateway API — so the + // resource is Ready as soon as permissions are applied, whether or not it + // declares a `finetune` spec. + self.finetune = resolve_finetune_capability(ctx).await?; Ok(HandlerAction::Continue { - state: next_state, + state: Ready, suggested_delay: None, }) } - // ─────────────── FINE-TUNING FLOW ────────────────────────── - // Only entered when the resource declares a `finetune` spec. - - #[handler( - state = SubmittingTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn submitting_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let config = ctx.desired_resource_config::()?; - let spec = config.finetune.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "SubmittingTuningJob reached without a finetune spec".to_string(), - }) - })?; - let aws_config = ctx.get_aws_config()?; - - // Resolve the training bucket from the dependency's real state rather than - // re-deriving the name. The Ai resource declares its training-data Storage - // as a dependency (see `Ai::get_dependencies`), and the AWS storage - // controller stores the actual (prefixed) bucket name in `bucket_name`. - // Reading it here keeps this controller correct even if the storage naming - // scheme changes. - let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); - let storage_state = ctx - .require_dependency::(&training_ref)?; - let bucket = storage_state.bucket_name.clone().ok_or_else(|| { - AlienError::new(ErrorData::DependencyNotReady { - resource_id: config.id.clone(), - dependency_id: spec.training_data.clone(), - }) - })?; - - // Bedrock assumes an IAM role to read the training data and write output. - // The Ai resource carries no dedicated service account, so derive the - // execution role ARN deterministically from the deployment account id and - // the stack's resource-role naming convention (`{prefix}-{id}`, matching - // `service_account::aws::get_aws_role_name`). The `ai/finetune` permission - // set grants this role the Bedrock job + S3 read actions it needs. - let role_arn = format!( - "arn:aws:iam::{}:role/{}-{}", - aws_config.account_id, ctx.resource_prefix, config.id - ); - - let training_key = &spec.training_key; - let training_uri = format!("s3://{}/{}", bucket, training_key); - let output_uri = format!("s3://{}/alien-finetune-output/{}/", bucket, config.id); - - // Bedrock job + custom-model names must match `([0-9a-zA-Z][_-]?){1,63}`; - // the resource id is already constrained to `[A-Za-z0-9-_]{1,64}`. - let job_name = format!("{}-{}", ctx.resource_prefix, config.id); - let custom_model_name = format!("{}-{}", ctx.resource_prefix, config.id); - - let request = CreateModelCustomizationJobRequest::builder() - .job_name(job_name) - .custom_model_name(custom_model_name) - .role_arn(role_arn) - .base_model_identifier(spec.base_model.clone()) - .customization_type(bedrock_customization_type(spec.method).to_string()) - .training_data_config(S3DataConfig::builder().s3_uri(training_uri).build()) - .output_data_config(S3DataConfig::builder().s3_uri(output_uri).build()) - .build(); - - let client = ctx.service_provider.get_aws_bedrock_client(aws_config).await?; - let response = client - .create_model_customization_job(&request) - .await - .context(ErrorData::CloudPlatformError { - message: format!( - "Failed to submit Bedrock model-customization job for AI '{}'", - config.id - ), - resource_id: Some(config.id.clone()), - })?; - - info!(ai=%config.id, job_arn=%response.job_arn, "Submitted Bedrock model-customization job"); - self.tuning_job_arn = Some(response.job_arn); - - Ok(HandlerAction::Continue { - state: WaitingForTuningJob, - suggested_delay: Some(TUNING_POLL_INTERVAL), - }) - } - - #[handler( - state = WaitingForTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn waiting_for_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let config = ctx.desired_resource_config::()?; - let aws_config = ctx.get_aws_config()?; - - let job_arn = self.tuning_job_arn.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "Tuning job ARN not set in state".to_string(), - }) - })?; - - let client = ctx.service_provider.get_aws_bedrock_client(aws_config).await?; - let job = client - .get_model_customization_job(job_arn) - .await - .context(ErrorData::CloudPlatformError { - message: format!( - "Failed to poll Bedrock model-customization job '{}' for AI '{}'", - job_arn, config.id - ), - resource_id: Some(config.id.clone()), - })?; - - match job.status() { - ModelCustomizationJobStatus::Completed => { - // The custom-model ARN is the artifact the gateway forwards to. - let output_model_arn = job.output_model_arn.clone().ok_or_else(|| { - AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Bedrock job '{}' completed but returned no outputModelArn", - job_arn - ), - resource_id: Some(config.id.clone()), - }) - })?; - - let spec = config.finetune.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "WaitingForTuningJob completed without a finetune spec".to_string(), - }) - })?; - - info!(ai=%config.id, custom_model=%output_model_arn, "Bedrock model-customization job completed"); - self.tuned_model_id = Some(output_model_arn); - self.served_id = Some(spec.served_model_id_or_default(&config.id)); - - Ok(HandlerAction::Continue { - state: Ready, - suggested_delay: None, - }) - } - // Fail loud on any terminal failure, mirroring the Azure controller's - // WaitingForDeployments fail-fast style (don't poll a dead job forever). - status @ (ModelCustomizationJobStatus::Failed - | ModelCustomizationJobStatus::Stopped) => { - let reason = job - .failure_message - .clone() - .unwrap_or_else(|| "no failure message reported".to_string()); - Err(AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Bedrock model-customization job '{}' entered terminal state '{:?}': {}", - job_arn, status, reason - ), - resource_id: Some(config.id.clone()), - })) - } - // InProgress / Stopping / any unrecognized status: keep polling. - other => { - info!(ai=%config.id, ?other, "Bedrock model-customization job not yet complete, re-polling"); - Ok(HandlerAction::Continue { - state: WaitingForTuningJob, - suggested_delay: Some(TUNING_POLL_INTERVAL), - }) - } - } - } - // ─────────────── READY STATE ──────────────────────────────── // Loops as a heartbeat tick; Bedrock has no per-stack resource to poll. @@ -325,7 +128,7 @@ impl AwsAiController { // ─────────────── UPDATE FLOW ────────────────────────────── // Ai has no mutable inference fields, and finetune base/training are immutable // (enforced in `Ai::validate_update`), so update is a no-op that also recovers - // RefreshFailed. The tuned model id persists in controller state across updates. + // RefreshFailed. The resolved finetune capability persists in state across updates. #[flow_entry(Update, from = [Ready, RefreshFailed])] #[handler( @@ -395,21 +198,12 @@ impl AwsAiController { let mut binding = AiBinding::bedrock(region); - // Attach the tuned model only once the job has completed and produced both - // a custom-model id and its served id (both are set together in - // WaitingForTuningJob). Until then (and for a pure inference gateway) the - // base binding is emitted unchanged. - match (&self.tuned_model_id, &self.served_id) { - (Some(upstream_id), Some(served_id)) => { - binding = binding.with_tuned_model(served_id.clone(), upstream_id.clone()); - } - (Some(_), None) => { - // A tuned artifact without a served id is a controller-state - // inconsistency (they are set together). Emit the untuned binding - // rather than serve under an unknown id, and warn loudly. - warn!("Tuned model id present but served id missing; emitting untuned binding"); - } - _ => {} + // A finetune-enabled resource carries a `FinetuneCapability` so the gateway + // can submit and rediscover a runtime tuning job. No tuned model is attached + // here — the gateway rediscovers it by convention once a runtime job + // succeeds. A pure inference gateway emits the plain binding unchanged. + if let Some(capability) = &self.finetune { + binding = binding.with_finetune(capability.clone()); } Ok(Some(serde_json::to_value(binding).into_alien_error().context( @@ -421,6 +215,48 @@ impl AwsAiController { } } +/// Resolve the fine-tuning capability from the declared `finetune` spec, the +/// training-storage dependency's real bucket name, and the deterministic +/// Bedrock-trusted finetune role. Returns `None` for a pure inference gateway. +async fn resolve_finetune_capability( + ctx: &ResourceControllerContext<'_>, +) -> Result> { + let config = ctx.desired_resource_config::()?; + let Some(spec) = config.finetune.as_ref() else { + return Ok(None); + }; + let aws_config = ctx.get_aws_config()?; + + // Resolve the training bucket from the dependency's real state rather than + // re-deriving the name. The Ai resource declares its training-data Storage as a + // dependency (see `Ai::get_dependencies`), and the AWS storage controller stores + // the actual (prefixed) bucket name in `bucket_name`. + let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); + let storage_state = + ctx.require_dependency::(&training_ref)?; + let training_bucket = storage_state.bucket_name.clone().ok_or_else(|| { + AlienError::new(ErrorData::DependencyNotReady { + resource_id: config.id.clone(), + dependency_id: spec.training_data.clone(), + }) + })?; + + Ok(Some(FinetuneCapability { + base_model: spec.base_model.clone(), + training_bucket, + training_key: spec.training_key.clone(), + served_model_id: spec.served_model_id_or_default(&config.id), + job_name: format!("{}-{}", ctx.resource_prefix, config.id), + // The dedicated Bedrock-trusted finetune role emitted by the AWS emitter + // (`{prefix}-{id}-finetune`). The gateway passes this ARN as the tuning job's + // `roleArn` so Bedrock can read training data and write output. + role_arn: format!( + "arn:aws:iam::{}:role/{}-{}-finetune", + aws_config.account_id, ctx.resource_prefix, config.id + ), + })) +} + #[cfg(all(test, feature = "test-utils"))] mod tests { use super::*; @@ -428,9 +264,6 @@ mod tests { use crate::core::ResourceController; use crate::storage::AwsStorageController; use crate::MockPlatformServiceProvider; - use alien_aws_clients::bedrock::{ - CreateModelCustomizationJobResponse, GetModelCustomizationJobResponse, MockBedrockApi, - }; use alien_core::bindings::AiBinding; use alien_core::{Ai, FinetuneMethod, FinetuneSpec, Platform, ResourceStatus, Storage}; use std::sync::Arc; @@ -460,62 +293,12 @@ mod tests { Ai::new("my-ai".to_string()).build() } - /// The completed custom-model ARN the gateway forwards tuned requests to. - const CUSTOM_MODEL_ARN: &str = - "arn:aws:bedrock:us-east-1:123456789012:custom-model/amazon.nova-lite-v1:0/abcdef012345"; - const JOB_ARN: &str = - "arn:aws:bedrock:us-east-1:123456789012:model-customization-job/amazon.nova-lite-v1:0/abcdef012345"; - - fn completed_job() -> GetModelCustomizationJobResponse { - serde_json::from_value(serde_json::json!({ - "status": "Completed", - "jobArn": JOB_ARN, - "outputModelArn": CUSTOM_MODEL_ARN, - "outputModelName": "test-my-ai", - })) - .expect("valid completed job json") - } - - fn in_progress_job() -> GetModelCustomizationJobResponse { - serde_json::from_value(serde_json::json!({ "status": "InProgress" })) - .expect("valid in-progress job json") - } - - fn failed_job() -> GetModelCustomizationJobResponse { - serde_json::from_value(serde_json::json!({ - "status": "Failed", - "failureMessage": "training data validation failed", - })) - .expect("valid failed job json") - } - - /// Build a mock Bedrock client that submits the job then returns the given - /// sequence of poll responses (one per `get_model_customization_job` call). - fn mock_bedrock(polls: Vec) -> Arc { - let mut mock = MockBedrockApi::new(); - mock.expect_create_model_customization_job().returning(|_| { - Ok(CreateModelCustomizationJobResponse { - job_arn: JOB_ARN.to_string(), - }) - }); - let responses = std::sync::Mutex::new(polls.into_iter()); - mock.expect_get_model_customization_job() - .returning(move |_| { - responses - .lock() - .unwrap() - .next() - .map(Ok) - .unwrap_or_else(|| Ok(completed_job())) - }); - Arc::new(mock) - } - - fn provider_with_bedrock(mock: Arc) -> Arc { + /// A provider that must never touch Bedrock: the runtime model submits tuning + /// jobs from the gateway, never from the controller, so the controller must not + /// construct a Bedrock client at deploy time. + fn provider_without_bedrock() -> Arc { let mut provider = MockPlatformServiceProvider::new(); - provider - .expect_get_aws_bedrock_client() - .returning(move |_| Ok(mock.clone())); + provider.expect_get_aws_bedrock_client().never(); Arc::new(provider) } @@ -537,12 +320,24 @@ mod tests { serde_json::from_value(value).expect("binding deserializes") } + /// Step the executor until it reaches Running (Ready is a heartbeat loop, not a + /// terminal state) or the bound is exhausted. + async fn drive_to_running(executor: &mut SingleControllerExecutor, steps: usize) { + for _ in 0..steps { + if executor.status() == ResourceStatus::Running { + return; + } + executor.step().await.expect("step should not error"); + } + } + #[tokio::test] - async fn test_finetune_submit_poll_succeeds_and_binds_tuned_model() { - // Submit -> InProgress -> Completed -> Ready, then the binding must carry - // the tuned model with the right served + upstream ids. - let mock = mock_bedrock(vec![in_progress_job(), completed_job()]); - let provider = provider_with_bedrock(mock); + async fn test_finetune_reaches_ready_immediately_with_capability_and_no_job() { + // The runtime model: a finetune resource reaches Ready as soon as + // permissions are applied — NO Bedrock job is submitted at deploy time — and + // its binding carries a FinetuneCapability with the resolved training bucket + // and the dedicated `-finetune` role ARN. + let provider = provider_without_bedrock(); let mut executor = SingleControllerExecutor::builder() .resource(tuned_ai()) @@ -557,77 +352,48 @@ mod tests { .await .expect("executor builds"); - // Ready is not terminal (heartbeat loop), so step until Running. - for _ in 0..8 { - if executor.status() == ResourceStatus::Running { - break; - } - executor.step().await.expect("step should not error"); - } + drive_to_running(&mut executor, 6).await; assert_eq!( executor.status(), ResourceStatus::Running, - "finetune resource must reach Ready after the job completes" + "a finetune resource must reach Ready immediately (no deploy-time job)" ); let binding = current_binding(&executor); - let tuned = binding - .tuned_model() - .expect("completed finetune must attach a tuned model"); + assert!( + binding.tuned_model().is_none(), + "no tuned model is attached at deploy time; the gateway rediscovers it" + ); + + let capability = binding + .finetune() + .expect("finetune resource must carry a FinetuneCapability"); + assert_eq!(capability.base_model, "amazon.nova-lite-v1:0"); + assert_eq!( + capability.training_bucket, + training_bucket_name(), + "training bucket must come from the storage dependency's real state" + ); + assert_eq!(capability.training_key, "training.jsonl"); assert_eq!( - tuned.served_id, "my-ai-tuned", + capability.served_model_id, "my-ai-tuned", "served id must default to -tuned" ); assert_eq!( - tuned.upstream_id, CUSTOM_MODEL_ARN, - "upstream id must be the completed custom-model ARN" + capability.job_name, "test-my-ai", + "job name must be the deterministic {{prefix}}-{{id}}" ); - } - - #[tokio::test] - async fn test_finetune_job_failed_fails_fast_to_create_failed() { - // A terminal Failed status must route to CreateFailed, not poll forever. - let mock = mock_bedrock(vec![failed_job()]); - let provider = provider_with_bedrock(mock); - - let mut executor = SingleControllerExecutor::builder() - .resource(tuned_ai()) - .controller(AwsAiController::default()) - .platform(Platform::Aws) - .service_provider(provider) - .with_dependency( - training_storage(), - AwsStorageController::mock_ready(TRAINING_STORAGE_ID), - ) - .build() - .await - .expect("executor builds"); - - // A terminal Failed job must surface as a handler error (which the real - // executor routes to CreateFailed via `on_failure`), NOT an endless poll. - // The test harness's `step()` returns that error rather than applying the - // failure transition, so assert the error surfaces. Bounded so a - // poll-forever regression fails the test instead of hanging. - let mut surfaced_error = false; - for _ in 0..10 { - if executor.step().await.is_err() { - surfaced_error = true; - break; - } - } - assert!( - surfaced_error, - "a terminal Failed job must surface as an error, not a silent retry" + assert_eq!( + capability.role_arn, "arn:aws:iam::123456789012:role/test-my-ai-finetune", + "role ARN must name the dedicated Bedrock-trusted `-finetune` role" ); } #[tokio::test] - async fn test_no_finetune_reaches_ready_with_untuned_binding() { - // Regression: an inference-only resource must never touch Bedrock tuning - // and must emit an untuned binding. - let mut provider = MockPlatformServiceProvider::new(); - provider.expect_get_aws_bedrock_client().never(); - let provider = Arc::new(provider); + async fn test_no_finetune_reaches_ready_with_plain_binding() { + // Regression: an inference-only resource never touches Bedrock tuning and + // emits a plain binding with neither a tuned model nor a capability. + let provider = provider_without_bedrock(); let mut executor = SingleControllerExecutor::builder() .resource(untuned_ai()) @@ -639,12 +405,7 @@ mod tests { .await .expect("executor builds"); - for _ in 0..6 { - if executor.status() == ResourceStatus::Running { - break; - } - executor.step().await.expect("step should not error"); - } + drive_to_running(&mut executor, 6).await; assert_eq!( executor.status(), ResourceStatus::Running, @@ -656,56 +417,9 @@ mod tests { binding.tuned_model().is_none(), "inference-only binding must omit tuned_model" ); - } - - #[tokio::test] - async fn test_finetune_uses_training_bucket_and_customization_type() { - // Assert the submitted job carries the dependency's real bucket in both the - // training and output S3 URIs, and FINE_TUNING as the customization type. - let bucket = training_bucket_name(); - let expected_training_uri = format!("s3://{}/training.jsonl", bucket); - let expected_output_prefix = format!("s3://{}/alien-finetune-output/my-ai/", bucket); - - let mut mock = MockBedrockApi::new(); - mock.expect_create_model_customization_job() - .withf(move |req| { - req.customization_type == "FINE_TUNING" - && req.base_model_identifier == "amazon.nova-lite-v1:0" - && req.training_data_config.s3_uri == expected_training_uri - && req.output_data_config.s3_uri == expected_output_prefix - && req.role_arn == "arn:aws:iam::123456789012:role/test-my-ai" - }) - .returning(|_| { - Ok(CreateModelCustomizationJobResponse { - job_arn: JOB_ARN.to_string(), - }) - }); - mock.expect_get_model_customization_job() - .returning(|_| Ok(completed_job())); - - let provider = provider_with_bedrock(Arc::new(mock)); - - let mut executor = SingleControllerExecutor::builder() - .resource(tuned_ai()) - .controller(AwsAiController::default()) - .platform(Platform::Aws) - .service_provider(provider) - .with_dependency( - training_storage(), - AwsStorageController::mock_ready(TRAINING_STORAGE_ID), - ) - .build() - .await - .expect("executor builds"); - - // Reaching Running proves the withf predicate matched (a mismatch panics - // the mock, which surfaces as a step error and never reaches Running). - for _ in 0..8 { - if executor.status() == ResourceStatus::Running { - break; - } - executor.step().await.expect("step should not error"); - } - assert_eq!(executor.status(), ResourceStatus::Running); + assert!( + binding.finetune().is_none(), + "inference-only binding must omit the finetune capability" + ); } } diff --git a/crates/alien-infra/src/ai/aws_import.rs b/crates/alien-infra/src/ai/aws_import.rs index e35b3cdf2..484fb8004 100644 --- a/crates/alien-infra/src/ai/aws_import.rs +++ b/crates/alien-infra/src/ai/aws_import.rs @@ -24,9 +24,7 @@ impl ResourceImporter for AwsAiImporter { let controller = AwsAiController { state: AwsAiState::Ready, region: Some(data.region), - tuning_job_arn: None, - tuned_model_id: None, - served_id: None, + finetune: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-infra/src/ai/azure.rs b/crates/alien-infra/src/ai/azure.rs index 7f3667a4c..a6fafd592 100644 --- a/crates/alien-infra/src/ai/azure.rs +++ b/crates/alien-infra/src/ai/azure.rs @@ -3,19 +3,19 @@ use tracing::info; use crate::azure_utils::get_resource_group_name; use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; -use crate::infra_requirements::azure_utils::get_storage_account_name; +use crate::storage::AzureStorageController; use alien_azure_clients::azure::cognitive_services::{ CognitiveServicesAccountCreateParameters, CognitiveServicesAccountCreateProperties, CognitiveServicesDeploymentCreateParameters, CognitiveServicesDeploymentCreateProperties, CognitiveServicesDeploymentModel, CognitiveServicesDeploymentSku, CognitiveServicesSku, }; use alien_azure_clients::long_running_operation::OperationResult; -use alien_core::FinetuneSpec; use alien_client_core::ErrorData as CloudClientErrorData; use alien_core::{ - bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, - AzureFoundryAiHeartbeatData, HeartbeatBackend, Platform, ResourceHeartbeat, - ResourceHeartbeatData, ResourceOutputs, ResourceStatus, + bindings::{AiBinding, FinetuneCapability}, + Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, AzureFoundryAiHeartbeatData, + HeartbeatBackend, Platform, ResourceHeartbeat, ResourceHeartbeatData, ResourceOutputs, + ResourceRef, ResourceStatus, Storage, }; use alien_error::{AlienError, Context, ContextError, IntoAlienError}; use alien_macros::controller; @@ -54,15 +54,14 @@ pub struct AzureAiController { pub(crate) resource_group: Option, /// The Azure region where the account is created. pub(crate) location: Option, - /// The Foundry fine-tuning job id, set once the job is submitted. Only - /// populated when the resource declares a `finetune` spec. + /// The fine-tuning capability this gateway advertises, resolved during the + /// create flow when the resource declares a `finetune` spec. The gateway + /// submits and rediscovers the Foundry fine-tuning job at runtime from this + /// capability; the controller never starts a job itself. Captured into state + /// (rather than rebuilt in `get_binding_params`, which has no `ctx`) so the + /// binding can carry it. `None` for a pure-inference gateway. #[serde(default)] - pub(crate) tuning_job_id: Option, - /// The deployment name serving the tuned model, set once the tuned model is - /// deployed. This is the `upstream_id` the gateway forwards to over the - /// OpenAI chat path. Absent for a pure inference gateway. - #[serde(default)] - pub(crate) tuned_deployment_name: Option, + pub(crate) finetune: Option, } #[controller] @@ -469,14 +468,12 @@ impl AzureAiController { info!(account_name = %account_name, "All Azure model deployments provisioned"); - // A pure inference gateway is Ready here. When the resource declares a - // finetune spec, tune the base model and serve the result before Ready. - if config.finetune.is_some() { - return Ok(HandlerAction::Continue { - state: SubmittingTuningJob, - suggested_delay: None, - }); - } + // Fine-tuning is triggered at runtime through the gateway, not at deploy + // time, so the resource is Ready as soon as the base deployments succeed. + // When the resource declares a `finetune` spec, resolve the capability the + // gateway needs to submit/rediscover a runtime tuning job and carry it on + // the binding. + self.finetune = self.resolve_finetune_capability(ctx, &config)?; Ok(HandlerAction::Continue { state: Ready, @@ -484,249 +481,6 @@ impl AzureAiController { }) } - // ─────────────── FINE-TUNING FLOW ─────────────────────────── - - #[handler( - state = SubmittingTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn submitting_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let azure_config = ctx.get_azure_config()?; - let config = ctx.desired_resource_config::()?; - let spec = require_finetune_spec(&config)?; - - let endpoint = self.require_endpoint(&config.id)?.to_string(); - - // Foundry imports the training dataset directly from the customer's Blob - // container. We derive the blob URL deterministically the same way the - // Azure Storage controller names the container - // (`{prefix}-{training_data}`, lowercased, `_`->`-`) under the shared - // default storage account, rather than requiring a public storage state - // type. See `crates/alien-infra/src/storage/azure.rs`. - // - // GOTCHA: Foundry's Blob import path requires the AIServices account to - // allow *public network access*. A private-endpoint-only account rejects - // the blob URL at job-submit time; the job then surfaces as a terminal - // failure in WaitingForTuningJob (fail-fast), not a silent hang. - let storage_account = get_storage_account_name(ctx.state)?; - let training_file = training_blob_url( - &storage_account, - ctx.resource_prefix, - &spec.training_data, - &spec.training_key, - ); - - info!( - id = %config.id, - endpoint = %endpoint, - base_model = %spec.base_model, - training_file = %training_file, - "Submitting Azure Foundry fine-tuning job" - ); - - let finetuning_client = ctx - .service_provider - .get_azure_foundry_finetuning_client(azure_config)?; - - let job = finetuning_client - .create_fine_tuning_job(&endpoint, &spec.base_model, &training_file) - .await - .context(ErrorData::CloudPlatformError { - message: format!( - "Failed to submit Azure Foundry fine-tuning job for base model '{}'", - spec.base_model - ), - resource_id: Some(config.id.clone()), - })?; - - info!(id = %config.id, job_id = %job.id, status = %job.status, "Fine-tuning job submitted"); - self.tuning_job_id = Some(job.id); - - Ok(HandlerAction::Continue { - state: WaitingForTuningJob, - suggested_delay: Some(std::time::Duration::from_secs(30)), - }) - } - - #[handler( - state = WaitingForTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn waiting_for_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let azure_config = ctx.get_azure_config()?; - let config = ctx.desired_resource_config::()?; - - let endpoint = self.require_endpoint(&config.id)?.to_string(); - let job_id = self.tuning_job_id.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Fine-tuning job id not set in state".to_string(), - resource_id: Some(config.id.clone()), - }) - })?; - - info!(id = %config.id, job_id = %job_id, "Polling Azure Foundry fine-tuning job"); - - let finetuning_client = ctx - .service_provider - .get_azure_foundry_finetuning_client(azure_config)?; - - let job = finetuning_client - .get_fine_tuning_job(&endpoint, job_id) - .await - .context(ErrorData::CloudPlatformError { - message: format!("Failed to poll Azure Foundry fine-tuning job '{}'", job_id), - resource_id: Some(config.id.clone()), - })?; - - // Fail fast on a terminal failure rather than polling a dead job forever, - // mirroring the WaitingForDeployments style. - if job.is_terminal_failure() { - return Err(AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Azure Foundry fine-tuning job '{}' entered terminal state '{}'", - job.id, job.status - ), - resource_id: Some(config.id.clone()), - })); - } - - if !job.is_succeeded() { - info!(id = %config.id, job_id = %job.id, status = %job.status, "Fine-tuning job not yet complete, retrying"); - return Ok(HandlerAction::Continue { - state: WaitingForTuningJob, - suggested_delay: Some(std::time::Duration::from_secs(30)), - }); - } - - // Succeeded: capture the tuned model name to deploy. Fail fast if the - // provider reports success without one — we cannot serve a missing model. - let fine_tuned_model = job.fine_tuned_model.ok_or_else(|| { - AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Azure Foundry fine-tuning job '{}' succeeded but returned no fine_tuned_model", - job.id - ), - resource_id: Some(config.id.clone()), - }) - })?; - - info!(id = %config.id, fine_tuned_model = %fine_tuned_model, "Fine-tuning job succeeded"); - - // Deploy the tuned model under the served id, so the OpenAI chat path - // accepts it as a deployment target. - let spec = require_finetune_spec(&config)?; - let served_id = spec.served_model_id_or_default(&config.id); - self.deploy_tuned_model(ctx, &config.id, &served_id, &fine_tuned_model) - .await?; - self.tuned_deployment_name = Some(served_id); - - Ok(HandlerAction::Continue { - state: WaitingForTunedDeployment, - suggested_delay: Some(std::time::Duration::from_secs(10)), - }) - } - - #[handler( - state = WaitingForTunedDeployment, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn waiting_for_tuned_deployment( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let azure_config = ctx.get_azure_config()?; - let config = ctx.desired_resource_config::()?; - - let account_name = self.require_account_name(&config.id)?.to_string(); - let resource_group_name = self.require_resource_group(&config.id)?.to_string(); - let deployment_name = self.tuned_deployment_name.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Tuned deployment name not set in state".to_string(), - resource_id: Some(config.id.clone()), - }) - })?; - - let cognitive_client = ctx - .service_provider - .get_azure_cognitive_services_client(azure_config)?; - - match cognitive_client - .get_deployment(&resource_group_name, &account_name, deployment_name) - .await - { - Ok(deployment) => { - let provisioning_state = deployment - .properties - .as_ref() - .and_then(|p| p.provisioning_state.as_deref()) - .unwrap_or(""); - - if provisioning_state.eq_ignore_ascii_case("Failed") - || provisioning_state.eq_ignore_ascii_case("Canceled") - { - return Err(AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Azure tuned-model deployment '{}' entered terminal state '{}'", - deployment_name, provisioning_state - ), - resource_id: Some(config.id.clone()), - })); - } - - if !provisioning_state.eq_ignore_ascii_case("Succeeded") { - info!( - account_name = %account_name, - deployment = %deployment_name, - provisioning_state = %provisioning_state, - "Tuned-model deployment not yet ready, retrying" - ); - return Ok(HandlerAction::Continue { - state: WaitingForTunedDeployment, - suggested_delay: Some(std::time::Duration::from_secs(10)), - }); - } - - info!(account_name = %account_name, deployment = %deployment_name, "Tuned-model deployment provisioned"); - Ok(HandlerAction::Continue { - state: Ready, - suggested_delay: None, - }) - } - Err(e) - if matches!( - &e.error, - Some(CloudClientErrorData::RemoteResourceNotFound { .. }) - ) => - { - info!( - account_name = %account_name, - deployment = %deployment_name, - "Tuned-model deployment not yet visible, retrying" - ); - Ok(HandlerAction::Continue { - state: WaitingForTunedDeployment, - suggested_delay: Some(std::time::Duration::from_secs(10)), - }) - } - Err(e) => Err(e.context(ErrorData::CloudPlatformError { - message: format!( - "Failed to poll Azure tuned-model deployment '{}'", - deployment_name - ), - resource_id: Some(config.id.clone()), - })), - } - } - // ─────────────── READY STATE ──────────────────────────────── #[handler( @@ -972,13 +726,12 @@ impl AzureAiController { let mut binding = AiBinding::foundry(endpoint, account_name); - // Attach the tuned model only once its deployment exists, so a pure - // inference gateway keeps the unchanged untuned wire shape. We deploy the - // tuned model under a deployment named `served_id`, and Foundry's OpenAI - // chat path takes that same deployment name as the request-body `model`. - // So served_id and upstream_id are the one stored deployment name. - if let Some(deployment_name) = &self.tuned_deployment_name { - binding = binding.with_tuned_model(deployment_name.clone(), deployment_name.clone()); + // Carry the fine-tuning capability when the resource declared one, so the + // gateway can submit/rediscover a runtime tuning job. A pure-inference + // gateway returns the untuned binding unchanged. The controller never + // attaches a `tuned_model` — the gateway rediscovers it by convention. + if let Some(capability) = self.finetune.as_ref() { + binding = binding.with_finetune(capability.clone()); } Ok(Some(serde_json::to_value(binding).into_alien_error().context( @@ -996,95 +749,44 @@ impl AzureAiController { self.endpoint = None; self.resource_group = None; self.location = None; - self.tuning_job_id = None; - self.tuned_deployment_name = None; - } - - /// Returns the provisioned account endpoint, or a config error if unset. - fn require_endpoint(&self, resource_id: &str) -> Result<&str> { - self.endpoint.as_deref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Endpoint not set in state".to_string(), - resource_id: Some(resource_id.to_string()), - }) - }) - } - - /// Returns the AIServices account name, or a config error if unset. - fn require_account_name(&self, resource_id: &str) -> Result<&str> { - self.account_name.as_deref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Account name not set in state".to_string(), - resource_id: Some(resource_id.to_string()), - }) - }) - } - - /// Returns the resource group name, or a config error if unset. - fn require_resource_group(&self, resource_id: &str) -> Result<&str> { - self.resource_group.as_deref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Resource group not set in state".to_string(), - resource_id: Some(resource_id.to_string()), - }) - }) + self.finetune = None; } - /// Deploys the tuned model under `deployment_name`, reusing the same - /// control-plane `create_deployment` path as the base catalog so the OpenAI - /// chat endpoint accepts the deployment name as a `model`. - async fn deploy_tuned_model( + /// Builds the fine-tuning capability the binding carries when the resource + /// declares a `finetune` spec, or `None` for a pure-inference gateway. + /// + /// Resolves the training data's real Blob container from the storage + /// dependency's controller state (rather than re-deriving the prefixed name), + /// keeping it in lockstep with the storage controller's naming. Foundry + /// submits the runtime tuning job under the gateway's ambient identity, so no + /// role is passed (`role_arn` is empty). + fn resolve_finetune_capability( &self, ctx: &ResourceControllerContext<'_>, - resource_id: &str, - deployment_name: &str, - fine_tuned_model: &str, - ) -> Result<()> { - let azure_config = ctx.get_azure_config()?; - let account_name = self.require_account_name(resource_id)?.to_string(); - let resource_group_name = self.require_resource_group(resource_id)?.to_string(); - - info!( - account_name = %account_name, - deployment = %deployment_name, - fine_tuned_model = %fine_tuned_model, - "Deploying tuned model" - ); - - let cognitive_client = ctx - .service_provider - .get_azure_cognitive_services_client(azure_config)?; - - // The tuned model's format is OpenAI and its "version" the fine-tuned - // model id Foundry returned. Capacity mirrors the base deployments. - let parameters = CognitiveServicesDeploymentCreateParameters { - sku: CognitiveServicesDeploymentSku { - name: "GlobalStandard".to_string(), - capacity: DEFAULT_DEPLOYMENT_CAPACITY, - }, - properties: CognitiveServicesDeploymentCreateProperties { - model: CognitiveServicesDeploymentModel { - format: "OpenAI".to_string(), - name: fine_tuned_model.to_string(), - version: "1".to_string(), - }, - }, + config: &Ai, + ) -> Result> { + let Some(spec) = config.finetune.as_ref() else { + return Ok(None); }; - cognitive_client - .create_deployment( - &resource_group_name, - &account_name, - deployment_name, - ¶meters, - ) - .await - .context(ErrorData::CloudPlatformError { - message: format!("Failed to deploy tuned model as '{}'", deployment_name), - resource_id: Some(resource_id.to_string()), - })?; + let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); + let storage_state = ctx.require_dependency::(&training_ref)?; + let training_bucket = storage_state.container_name.ok_or_else(|| { + AlienError::new(ErrorData::DependencyNotReady { + resource_id: config.id.clone(), + dependency_id: spec.training_data.clone(), + }) + })?; - Ok(()) + Ok(Some(FinetuneCapability { + base_model: spec.base_model.clone(), + training_bucket, + training_key: spec.training_key.clone(), + served_model_id: spec.served_model_id_or_default(&config.id), + job_name: format!("{}-{}", ctx.resource_prefix, config.id), + // Foundry submits under the ambient identity; no passed role. + role_arn: String::new(), + })) } /// Creates a controller in the ready state for testing purposes. @@ -1096,8 +798,7 @@ impl AzureAiController { endpoint: Some(endpoint.to_string()), resource_group: Some("mock-rg".to_string()), location: Some("eastus".to_string()), - tuning_job_id: None, - tuned_deployment_name: None, + finetune: None, _internal_stay_count: None, } } @@ -1112,59 +813,10 @@ impl AzureAiController { endpoint: Some(endpoint.to_string()), resource_group: Some("mock-rg".to_string()), location: Some("eastus".to_string()), - tuning_job_id: None, - tuned_deployment_name: None, + finetune: None, _internal_stay_count: None, } } - - /// Creates a controller poised to submit a fine-tuning job (account and base - /// deployments already provisioned), for testing the tuning states. - #[cfg(feature = "test-utils")] - pub fn mock_submitting_tuning(account_name: &str, endpoint: &str) -> Self { - Self { - state: AzureAiState::SubmittingTuningJob, - account_name: Some(account_name.to_string()), - endpoint: Some(endpoint.to_string()), - resource_group: Some("mock-rg".to_string()), - location: Some("eastus".to_string()), - tuning_job_id: None, - tuned_deployment_name: None, - _internal_stay_count: None, - } - } -} - -/// Extracts the finetune spec from an `Ai` config, erroring if absent. Called -/// only from tuning states, which are unreachable without a spec. -fn require_finetune_spec(config: &Ai) -> Result<&FinetuneSpec> { - config.finetune.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - message: "Reached a fine-tuning state without a finetune spec".to_string(), - resource_id: Some(config.id.clone()), - }) - }) -} - -/// Builds the Blob URL Foundry imports the training dataset from. -/// -/// Derived deterministically to match the Azure Storage controller's container -/// naming (`{prefix}-{name}` lowercased with `_`->`-`) under the shared default -/// storage account, so we don't need a public storage state type. Kept as a -/// pure function so it is unit-testable. -fn training_blob_url( - storage_account: &str, - resource_prefix: &str, - training_data: &str, - training_key: &str, -) -> String { - let container = format!("{}-{}", resource_prefix, training_data) - .to_lowercase() - .replace('_', "-"); - format!( - "https://{}.blob.core.windows.net/{}/{}", - storage_account, container, training_key - ) } #[cfg(all(test, feature = "test-utils"))] @@ -1177,7 +829,6 @@ mod tests { CognitiveServicesDeployment, CognitiveServicesDeploymentProperties, MockCognitiveServicesAccountsApi, }; - use alien_azure_clients::azure::openai_finetuning::{FineTuningJob, MockFoundryFineTuningApi}; use alien_azure_clients::long_running_operation::OperationResult; use alien_client_core::ErrorData as CloudClientErrorData; use crate::storage::AzureStorageController; @@ -1215,17 +866,7 @@ mod tests { Storage::new(TRAINING_STORAGE_ID.to_string()).build() } - /// A fine-tuning job in the given status, carrying `fine_tuned_model` only - /// when succeeded. - fn job(status: &str, fine_tuned_model: Option<&str>) -> FineTuningJob { - FineTuningJob { - id: "ftjob-abc".to_string(), - status: status.to_string(), - fine_tuned_model: fine_tuned_model.map(|s| s.to_string()), - } - } - - /// A cognitive-services mock that succeeds base + tuned deployments (create + + /// A cognitive-services mock that succeeds all base deployments (create + /// get both report Succeeded). fn mock_cognitive_all_succeed() -> MockCognitiveServicesAccountsApi { let mut mock = MockCognitiveServicesAccountsApi::new(); @@ -1475,55 +1116,18 @@ mod tests { // ─────────────── FINE-TUNING TESTS ────────────────────────── - #[test] - fn training_blob_url_matches_storage_container_naming() { - // Must match the Azure Storage controller's container naming so Foundry - // imports from the container the storage controller actually created. - let url = training_blob_url("myacct", "test", "training_set", "training.jsonl"); - assert_eq!( - url, - "https://myacct.blob.core.windows.net/test-training-set/training.jsonl", - "underscores must become hyphens and the prefix must be applied, lowercased" - ); - } - #[tokio::test] - async fn test_finetune_submit_poll_succeeds_and_binds_tuned_model() { - // Full flow from the base deployments (mock_deploying) through submit -> - // pending -> succeeded -> deploy tuned -> Ready, then assert the tuned - // binding. Drives the real WaitingForDeployments -> SubmittingTuningJob - // branch (config.finetune is Some). - // One shared finetuning mock (cloned per getter call) so its poll - // sequence persists across steps: first `running`, then `succeeded`. - let mut finetuning_mock = MockFoundryFineTuningApi::new(); - finetuning_mock - .expect_create_fine_tuning_job() - .returning(|_, _, _| Ok(job("pending", None))); - let polls = std::sync::Mutex::new( - vec![ - job("running", None), - job("succeeded", Some("gpt-4o-mini.ft-abc")), - ] - .into_iter(), - ); - finetuning_mock - .expect_get_fine_tuning_job() - .returning(move |_, _| { - Ok(polls - .lock() - .unwrap() - .next() - .unwrap_or_else(|| job("succeeded", Some("gpt-4o-mini.ft-abc")))) - }); - let finetuning_mock = Arc::new(finetuning_mock); - + async fn test_finetune_reaches_ready_immediately_with_capability() { + // Fine-tuning is triggered at runtime through the gateway, so a finetune + // resource reaches Ready as soon as the base deployments succeed — no + // tuning job or tuned deployment is created at deploy time. The Foundry + // finetuning client is never wired, so a call to it would panic on the + // unset expectation, catching any regression that re-introduces + // deploy-time tuning. The binding carries the FinetuneCapability instead. let mut mock_provider = MockPlatformServiceProvider::new(); mock_provider .expect_get_azure_cognitive_services_client() .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); - mock_provider - .expect_get_azure_foundry_finetuning_client() - .returning(move |_| Ok(finetuning_mock.clone())); let mut executor = SingleControllerExecutor::builder() .resource(tuned_ai()) @@ -1543,7 +1147,7 @@ mod tests { .expect("executor builds"); // Ready is a heartbeat loop (not terminal), so step until Running. - for _ in 0..12 { + for _ in 0..8 { if executor.status() == ResourceStatus::Running { break; } @@ -1552,85 +1156,42 @@ mod tests { assert_eq!( executor.status(), ResourceStatus::Running, - "finetune resource must reach Ready after the job completes and the tuned model deploys" + "a finetune resource is Ready immediately; the gateway triggers tuning at runtime" ); let binding = current_binding(&executor); - let tuned = binding - .tuned_model() - .expect("completed finetune must attach a tuned model"); - assert_eq!( - tuned.served_id, "my-ai-tuned", - "served id must default to -tuned" - ); - assert_eq!( - tuned.upstream_id, "my-ai-tuned", - "upstream id is the tuned deployment name, which equals served_id" + assert!( + binding.tuned_model().is_none(), + "the controller must not attach a tuned model; the gateway rediscovers it" ); - } - - #[tokio::test] - async fn test_finetune_job_failed_fails_fast_to_create_failed() { - // A terminal Failed job status must route to CreateFailed, not poll forever. - let mut mock_provider = MockPlatformServiceProvider::new(); - mock_provider - .expect_get_azure_cognitive_services_client() - .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); - mock_provider - .expect_get_azure_foundry_finetuning_client() - .returning(|_| { - let mut mock = MockFoundryFineTuningApi::new(); - mock.expect_create_fine_tuning_job() - .returning(|_, _, _| Ok(job("pending", None))); - mock.expect_get_fine_tuning_job() - .returning(|_, _| Ok(job("failed", None))); - Ok(Arc::new(mock)) - }); - - let mut executor = SingleControllerExecutor::builder() - .resource(tuned_ai()) - .controller(AzureAiController::mock_submitting_tuning( - "default-storage-account", - "https://my-ai.cognitiveservices.azure.com/", - )) - .platform(Platform::Azure) - .service_provider(Arc::new(mock_provider)) - .with_test_dependencies() - .with_dependency( - training_storage(), - AzureStorageController::mock_ready(TRAINING_STORAGE_ID), - ) - .build() - .await - .expect("executor builds"); - - // A terminal Failed job must surface as a handler error (which the - // executor's on_failure routing turns into CreateFailed), not an endless - // poll. Bounded so a poll-forever regression fails instead of hanging. - let mut surfaced_error = false; - for _ in 0..8 { - if executor.step().await.is_err() { - surfaced_error = true; - break; - } - } + let cap = binding + .finetune() + .expect("finetune binding must carry the fine-tuning capability"); + assert_eq!(cap.base_model, "gpt-4o-mini"); + // Bucket resolved from the storage dependency's container name + // (test-stack-, from AzureStorageController::mock_ready), not + // re-derived here. + assert_eq!(cap.training_bucket, "test-stack-training-set"); + assert_eq!(cap.training_key, "training.jsonl"); + assert_eq!(cap.served_model_id, "my-ai-tuned"); + // Deterministic {prefix}-{id}; the executor test harness uses the + // resource prefix "test" (distinct from the storage mock's "test-stack" + // container-name prefix). + assert_eq!(cap.job_name, "test-my-ai"); assert!( - surfaced_error, - "a terminal Failed job must surface as an error, not a silent retry" + cap.role_arn.is_empty(), + "Foundry submits under the ambient identity; no role is passed" ); } #[tokio::test] - async fn test_no_finetune_reaches_ready_with_untuned_binding() { - // Regression: an inference-only resource must never call the fine-tuning - // client and must emit an untuned binding. + async fn test_no_finetune_reaches_ready_with_plain_binding() { + // Regression: an inference-only resource reaches Ready with a plain + // binding carrying neither a tuned model nor a finetune capability. let mut mock_provider = MockPlatformServiceProvider::new(); mock_provider .expect_get_azure_cognitive_services_client() .returning(|_| Ok(Arc::new(mock_cognitive_all_succeed()))); - mock_provider - .expect_get_azure_foundry_finetuning_client() - .never(); let mut executor = SingleControllerExecutor::builder() .resource(basic_ai()) @@ -1662,5 +1223,9 @@ mod tests { binding.tuned_model().is_none(), "inference-only binding must omit tuned_model" ); + assert!( + binding.finetune().is_none(), + "inference-only binding must omit the finetune capability" + ); } } diff --git a/crates/alien-infra/src/ai/azure_import.rs b/crates/alien-infra/src/ai/azure_import.rs index f9356287e..d52458e7e 100644 --- a/crates/alien-infra/src/ai/azure_import.rs +++ b/crates/alien-infra/src/ai/azure_import.rs @@ -28,9 +28,9 @@ impl ResourceImporter for AzureAiImporter { endpoint: Some(data.endpoint), resource_group: Some(data.resource_group), location: Some(data.location), - // An imported resource has no in-flight tuning job or tuned deployment. - tuning_job_id: None, - tuned_deployment_name: None, + // Import targets an existing base gateway; it advertises no + // fine-tuning capability. + finetune: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-infra/src/ai/gcp.rs b/crates/alien-infra/src/ai/gcp.rs index 775dc8bb9..a218bd759 100644 --- a/crates/alien-infra/src/ai/gcp.rs +++ b/crates/alien-infra/src/ai/gcp.rs @@ -6,43 +6,31 @@ use crate::core::{ResourceControllerContext, ResourcePermissionsHelper}; use crate::error::{ErrorData, Result}; use crate::storage::GcpStorageController; use alien_core::{ - bindings::AiBinding, Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, FinetuneMethod, - GcpVertexAiHeartbeatData, HeartbeatBackend, Platform, ResourceHeartbeat, ResourceHeartbeatData, - ResourceOutputs, ResourceRef, ResourceStatus, Storage, + bindings::{AiBinding, FinetuneCapability}, + Ai, AiHeartbeatData, AiHeartbeatStatus, AiOutputs, GcpVertexAiHeartbeatData, HeartbeatBackend, + Platform, ResourceHeartbeat, ResourceHeartbeatData, ResourceOutputs, ResourceRef, + ResourceStatus, Storage, }; use alien_error::{AlienError, Context, IntoAlienError}; -use alien_gcp_clients::aiplatform::{ - CreateTuningJobRequest, JobState, SupervisedTuningSpec, -}; use alien_gcp_clients::iam::IamPolicy; use alien_gcp_clients::resource_manager::GetPolicyOptions; use alien_macros::controller; use chrono::Utc; -/// How often to re-poll a still-running Vertex tuning job. -const TUNING_POLL_INTERVAL: Duration = Duration::from_secs(30); - #[controller] pub struct GcpAiController { /// GCP project ID. None until create_start runs. pub(crate) project: Option, /// GCP region (location) for the Vertex AI endpoint. None until create_start runs. pub(crate) location: Option, - /// Full resource name of the submitted Vertex tuning job - /// (`projects/{p}/locations/{l}/tuningJobs/{id}`). Set once - /// `SubmittingTuningJob` runs; used by `WaitingForTuningJob` to poll. - /// `None` for a pure-inference gateway (no `finetune` spec). - pub(crate) tuning_job_name: Option, - /// The tuned model's upstream artifact id (the serving endpoint / model - /// resource name) the Vertex OpenAI-compat chat path routes to. Set only - /// once the tuning job reaches `JOB_STATE_SUCCEEDED`. Attached to the - /// binding via `.with_tuned_model(..)` when present. - pub(crate) tuned_model_upstream_id: Option, - /// The public model id the gateway serves the tuned model under - /// (`spec.served_model_id_or_default(&config.id)`). Captured alongside the - /// upstream id at success so `get_binding_params` (which has no ctx) can - /// build the tuned binding. - pub(crate) tuned_model_served_id: Option, + /// The fine-tuning capability this gateway advertises, resolved during the + /// create flow when the resource declares a `finetune` spec. The gateway + /// submits and rediscovers the tuning job at runtime from this capability; + /// the controller never starts a job itself. Captured into state (rather than + /// rebuilt in `get_binding_params`, which has no `ctx`) so the binding can + /// carry it. `None` for a pure-inference gateway. + #[serde(default)] + pub(crate) finetune: Option, } #[controller] @@ -209,218 +197,19 @@ impl GcpAiController { ) .await?; - // A pure-inference gateway is ready as soon as permissions are applied. - // A finetune spec first submits and awaits a Vertex tuning job. - if config.finetune.is_none() { - return Ok(HandlerAction::Continue { - state: Ready, - suggested_delay: None, - }); - } - - info!(id = %config.id, "Finetune requested; submitting Vertex tuning job"); - Ok(HandlerAction::Continue { - state: SubmittingTuningJob, - suggested_delay: None, - }) - } - - // ─────────────── FINETUNE FLOW ───────────────────────────── - // Only reached when the Ai declares a `finetune` spec. Submits a Vertex - // supervised tuning job reading the customer's training data from GCS, then - // polls it to completion before serving the tuned model through the gateway. - - #[handler( - state = SubmittingTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn submitting_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let gcp_config = ctx.get_gcp_config()?; - let config = ctx.desired_resource_config::()?; - let spec = config.finetune.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "SubmittingTuningJob reached without a finetune spec".to_string(), - }) - })?; - - // Vertex exposes only supervised tuning for Gemini; there is no - // user-selectable LoRA/QLoRA or DPO method. Map Sft -> supervised tuning - // and reject the others loudly rather than silently mis-tuning. - match spec.method { - FinetuneMethod::Sft => {} - other => { - return Err(AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: format!( - "Vertex AI supports only supervised fine-tuning (sft); method {other:?} is not available on Vertex" - ), - })); - } - } - - // Resolve the training data's real GCS bucket from the dependency's - // controller state, rather than re-deriving the prefixed name here. The - // Ai resource declares the training Storage as a dependency - // (`Ai::get_dependencies`), and the GCS controller records the exact, - // prefixed bucket name it created in `bucket_name` — reading it keeps - // this in lockstep with the storage controller's naming. - let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); - let storage_state = ctx.require_dependency::(&training_ref)?; - let bucket_name = storage_state.bucket_name.ok_or_else(|| { - AlienError::new(ErrorData::DependencyNotReady { - resource_id: config.id.clone(), - dependency_id: spec.training_data.clone(), - }) - })?; - - let training_uri = format!("gs://{bucket_name}/{}", spec.training_key); - - let request = CreateTuningJobRequest::builder() - .base_model(spec.base_model.clone()) - .supervised_tuning_spec( - SupervisedTuningSpec::builder() - .training_dataset_uri(training_uri.clone()) - .build(), - ) - .tuned_model_display_name(spec.served_model_id_or_default(&config.id)) - .build(); - - info!( - id = %config.id, - base_model = %spec.base_model, - training_uri = %training_uri, - "Submitting Vertex supervised tuning job" - ); - - let client = ctx.service_provider.get_gcp_aiplatform_client(gcp_config)?; - let job = client - .create_tuning_job(request) - .await - .context(ErrorData::CloudPlatformError { - message: "Failed to submit Vertex tuning job".to_string(), - resource_id: Some(config.id.clone()), - })?; - - let job_name = job.name.ok_or_else(|| { - AlienError::new(ErrorData::CloudPlatformError { - message: "Vertex tuning job creation returned no resource name".to_string(), - resource_id: Some(config.id.clone()), - }) - })?; + // Fine-tuning is triggered at runtime through the gateway, not at deploy + // time, so the resource is Ready as soon as permissions are applied. When + // the resource declares a `finetune` spec, resolve the capability the + // gateway needs to submit/rediscover a runtime tuning job and carry it on + // the binding. + self.finetune = self.resolve_finetune_capability(ctx, &config)?; - info!(id = %config.id, job = %job_name, "Vertex tuning job submitted"); - self.tuning_job_name = Some(job_name); - - // Poll once immediately; only *in-progress* re-polls wait a full interval. Ok(HandlerAction::Continue { - state: WaitingForTuningJob, + state: Ready, suggested_delay: None, }) } - #[handler( - state = WaitingForTuningJob, - on_failure = CreateFailed, - status = ResourceStatus::Provisioning, - )] - async fn waiting_for_tuning_job( - &mut self, - ctx: &ResourceControllerContext<'_>, - ) -> Result { - let gcp_config = ctx.get_gcp_config()?; - let config = ctx.desired_resource_config::()?; - let spec = config.finetune.as_ref().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "WaitingForTuningJob reached without a finetune spec".to_string(), - }) - })?; - let spec_served_id = spec.served_model_id_or_default(&config.id); - let job_name = self.tuning_job_name.clone().ok_or_else(|| { - AlienError::new(ErrorData::ResourceConfigInvalid { - resource_id: Some(config.id.clone()), - message: "WaitingForTuningJob reached without a submitted job name".to_string(), - }) - })?; - - let client = ctx.service_provider.get_gcp_aiplatform_client(gcp_config)?; - let job = client - .get_tuning_job(job_name.clone()) - .await - .context(ErrorData::CloudPlatformError { - message: format!("Failed to poll Vertex tuning job '{job_name}'"), - resource_id: Some(config.id.clone()), - })?; - - let state = job.state.unwrap_or(JobState::JobStateUnspecified); - - if state.is_in_progress() { - info!(id = %config.id, job = %job_name, ?state, "Vertex tuning job still running; re-polling"); - return Ok(HandlerAction::Continue { - state: WaitingForTuningJob, - suggested_delay: Some(TUNING_POLL_INTERVAL), - }); - } - - // Terminal failure states fail loud — a fine-tuning resource whose job - // failed must not silently degrade to serving the untuned base model. - if state.is_terminal_failure() { - let detail = job - .error - .map(|e| e.message) - .unwrap_or_else(|| "no error detail returned".to_string()); - return Err(AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Vertex tuning job '{job_name}' reached terminal state {state:?}: {detail}" - ), - resource_id: Some(config.id.clone()), - })); - } - - // Success: record the upstream artifact the gateway routes to. - if state == JobState::JobStateSucceeded { - let upstream_id = job - .tuned_model - .as_ref() - .and_then(|m| m.upstream_id()) - .ok_or_else(|| { - AlienError::new(ErrorData::CloudPlatformError { - message: format!( - "Vertex tuning job '{job_name}' succeeded but returned no tuned model endpoint/model" - ), - resource_id: Some(config.id.clone()), - }) - })? - .to_string(); - - info!( - id = %config.id, - job = %job_name, - upstream_id = %upstream_id, - "Vertex tuning job succeeded; tuned model ready" - ); - self.tuned_model_upstream_id = Some(upstream_id); - self.tuned_model_served_id = Some(spec_served_id); - - return Ok(HandlerAction::Continue { - state: Ready, - suggested_delay: None, - }); - } - - // Any other (unspecified/unknown) terminal state is unexpected — fail - // rather than loop forever or serve an untuned model. - Err(AlienError::new(ErrorData::CloudPlatformError { - message: format!("Vertex tuning job '{job_name}' in unexpected state {state:?}"), - resource_id: Some(config.id.clone()), - })) - } - // ─────────────── READY STATE ──────────────────────────────── // Loops as a heartbeat tick; Vertex AI has no per-stack resource to poll. @@ -554,14 +343,12 @@ impl GcpAiController { let mut binding = AiBinding::vertex(project, location); - // Attach the tuned model only once the tuning job has completed and we - // recorded both ids; otherwise the base (untuned) binding is returned - // unchanged, matching a pure-inference gateway. - if let (Some(served_id), Some(upstream_id)) = ( - self.tuned_model_served_id.as_ref(), - self.tuned_model_upstream_id.as_ref(), - ) { - binding = binding.with_tuned_model(served_id, upstream_id); + // Carry the fine-tuning capability when the resource declared one, so the + // gateway can submit/rediscover a runtime tuning job. A pure-inference + // gateway returns the untuned binding unchanged. The controller never + // attaches a `tuned_model` — the gateway rediscovers it by convention. + if let Some(capability) = self.finetune.as_ref() { + binding = binding.with_finetune(capability.clone()); } Ok(Some(serde_json::to_value(binding).into_alien_error().context( @@ -573,24 +360,61 @@ impl GcpAiController { } } +impl GcpAiController { + /// Builds the fine-tuning capability the binding carries when the resource + /// declares a `finetune` spec, or `None` for a pure-inference gateway. + /// + /// Resolves the training data's real GCS bucket from the storage dependency's + /// controller state (rather than re-deriving the prefixed name), keeping it in + /// lockstep with the storage controller's naming. Vertex submits the runtime + /// tuning job under the gateway's ambient identity, so no role is passed + /// (`role_arn` is empty). + fn resolve_finetune_capability( + &self, + ctx: &ResourceControllerContext<'_>, + config: &Ai, + ) -> Result> { + let Some(spec) = config.finetune.as_ref() else { + return Ok(None); + }; + + let training_ref = ResourceRef::new(Storage::RESOURCE_TYPE, spec.training_data.clone()); + let storage_state = ctx.require_dependency::(&training_ref)?; + let training_bucket = storage_state.bucket_name.ok_or_else(|| { + AlienError::new(ErrorData::DependencyNotReady { + resource_id: config.id.clone(), + dependency_id: spec.training_data.clone(), + }) + })?; + + Ok(Some(FinetuneCapability { + base_model: spec.base_model.clone(), + training_bucket, + training_key: spec.training_key.clone(), + served_model_id: spec.served_model_id_or_default(&config.id), + job_name: format!("{}-{}", ctx.resource_prefix, config.id), + // Vertex submits under the ambient identity; no passed role. + role_arn: String::new(), + })) + } +} + #[cfg(test)] mod tests { //! GCP Vertex AI controller tests. //! - //! These drive the finetune state machine end-to-end against a mocked - //! `AiPlatformApi`: submit -> pending -> succeeded -> Ready, asserting the - //! resulting binding carries the tuned model. They also cover the fail-fast - //! path (job FAILED -> ProvisionFailed) and the pure-inference regression - //! (no finetune -> Ready with an untuned binding). + //! Fine-tuning is triggered at runtime through the gateway, not at deploy + //! time, so the controller reaches Ready immediately and never submits a + //! Vertex tuning job. These assert that a finetune resource reaches Ready + //! with no job created and that its binding carries the `FinetuneCapability` + //! the gateway needs, plus the pure-inference regression (no finetune -> + //! Ready with a plain binding). - use std::sync::{Arc, Mutex}; + use std::sync::Arc; use super::GcpAiController; use alien_core::bindings::AiBinding; use alien_core::{Ai, FinetuneMethod, FinetuneSpec, Platform, ResourceStatus, Storage}; - use alien_gcp_clients::aiplatform::{ - JobState, MockAiPlatformApi, TunedModelRef, TuningJob, - }; use alien_gcp_clients::iam::IamPolicy; use alien_gcp_clients::longrunning::Operation; use alien_gcp_clients::resource_manager::MockResourceManagerApi; @@ -601,7 +425,6 @@ mod tests { use crate::storage::GcpStorageController; const TRAINING_STORAGE_ID: &str = "training-set"; - const TUNED_ENDPOINT: &str = "projects/test-project-123/locations/us-central1/endpoints/9988"; // ─────────────── FIXTURES ────────────────────────────────── @@ -610,7 +433,7 @@ mod tests { Ai::new("llm".to_string()).build() } - /// An Ai that fine-tunes a Gemini base model from the training storage. + /// An Ai that declares a fine-tuning capability over the training storage. fn finetune_ai() -> Ai { Ai::new("llm".to_string()) .finetune(FinetuneSpec { @@ -623,19 +446,6 @@ mod tests { .build() } - /// An Ai requesting a method Vertex does not expose (LoRA). - fn lora_ai() -> Ai { - Ai::new("llm".to_string()) - .finetune(FinetuneSpec { - base_model: "gemini-2.0-flash-001".to_string(), - training_data: TRAINING_STORAGE_ID.to_string(), - training_key: "training.jsonl".to_string(), - served_model_id: None, - method: FinetuneMethod::Lora, - }) - .build() - } - fn training_storage() -> Storage { Storage::new(TRAINING_STORAGE_ID.to_string()).build() } @@ -661,10 +471,12 @@ mod tests { Arc::new(mock) } - /// Wires a service provider with the given aiplatform mock plus the - /// service-usage and resource-manager mocks the create flow needs. - fn provider_with(aiplatform: MockAiPlatformApi) -> Arc { - let aiplatform = Arc::new(aiplatform); + /// Wires a service provider with the service-usage and resource-manager + /// mocks the create flow needs. The controller no longer submits a Vertex + /// tuning job, so no aiplatform client is wired — a call to it would panic on + /// the unset expectation, catching any regression that re-introduces + /// deploy-time tuning. + fn provider() -> Arc { let mut provider = MockPlatformServiceProvider::new(); provider .expect_get_gcp_service_usage_client() @@ -678,76 +490,22 @@ mod tests { let m = iam_mock(); move |_| Ok(m.clone()) }); - provider - .expect_get_gcp_aiplatform_client() - .returning(move |_| Ok(aiplatform.clone())); Arc::new(provider) } - fn succeeded_job() -> TuningJob { - TuningJob::builder() - .name("projects/test-project-123/locations/us-central1/tuningJobs/42".to_string()) - .state(JobState::JobStateSucceeded) - .tuned_model( - TunedModelRef::builder() - .endpoint(TUNED_ENDPOINT.to_string()) - .build(), - ) - .build() - } - // ─────────────── TESTS ───────────────────────────────────── - /// Submit -> pending -> running -> succeeded -> Ready, and the binding - /// carries the tuned model with the right served/upstream ids. + /// A finetune resource reaches Ready immediately (no tuning job submitted at + /// deploy) and its binding carries the `FinetuneCapability` with the bucket + /// resolved from the storage dependency, no `tuned_model`, and an empty + /// `role_arn` (Vertex submits under the ambient identity). #[tokio::test] - async fn finetune_reaches_ready_with_tuned_binding() { - let mut aiplatform = MockAiPlatformApi::new(); - - // create returns a job with a name to poll. - aiplatform - .expect_create_tuning_job() - .withf(|req| { - req.base_model == "gemini-2.0-flash-001" - // The gs:// URI is built from the dependency bucket - // (test-stack-, from GcpStorageController::mock_ready) + training_key. - && req.supervised_tuning_spec.training_dataset_uri - == "gs://test-stack-training-set/training.jsonl" - && req.tuned_model_display_name.as_deref() == Some("llm-tuned") - }) - .times(1) - .returning(|_| { - Ok(TuningJob::builder() - .name( - "projects/test-project-123/locations/us-central1/tuningJobs/42" - .to_string(), - ) - .state(JobState::JobStatePending) - .build()) - }); - - // Poll: still-running once (proving the re-poll loop), then succeeded. - let poll_count = Arc::new(Mutex::new(0u32)); - aiplatform.expect_get_tuning_job().returning(move |name| { - assert_eq!( - name, - "projects/test-project-123/locations/us-central1/tuningJobs/42" - ); - let mut n = poll_count.lock().unwrap(); - *n += 1; - match *n { - 1 => Ok(TuningJob::builder().state(JobState::JobStateRunning).build()), - _ => Ok(succeeded_job()), - } - }); - - let provider = provider_with(aiplatform); - + async fn finetune_reaches_ready_immediately_with_capability() { let mut executor = SingleControllerExecutor::builder() .resource(finetune_ai()) .controller(GcpAiController::default()) .platform(Platform::Gcp) - .service_provider(provider) + .service_provider(provider()) .with_dependency( training_storage(), GcpStorageController::mock_ready(TRAINING_STORAGE_ID), @@ -764,108 +522,52 @@ mod tests { assert_eq!( executor.status(), ResourceStatus::Running, - "a completed tuning job leaves the gateway Running" + "a finetune resource is Ready immediately; the gateway triggers tuning at runtime" ); - // The controller recorded the tuned artifact. let controller = executor .internal_state::() .expect("controller downcast"); - assert_eq!( - controller.tuned_model_upstream_id.as_deref(), - Some(TUNED_ENDPOINT) - ); - assert_eq!(controller.tuned_model_served_id.as_deref(), Some("llm-tuned")); - // The binding the gateway consumes carries the tuned model. + // The binding the gateway consumes carries the finetune capability but no + // tuned model (the gateway rediscovers the tuned model by convention). let params = controller .get_binding_params() .expect("binding params ok") .expect("binding present once project/location are set"); - let binding: AiBinding = - serde_json::from_value(params).expect("binding deserializes"); - let tuned = binding - .tuned_model() - .expect("binding must carry the tuned model"); - assert_eq!(tuned.served_id, "llm-tuned"); - assert_eq!(tuned.upstream_id, TUNED_ENDPOINT); - } - - /// A terminal FAILED job fails the resource loud (no silent fall-back to base). - #[tokio::test] - async fn finetune_failed_job_reaches_provision_failed() { - let mut aiplatform = MockAiPlatformApi::new(); - aiplatform.expect_create_tuning_job().returning(|_| { - Ok(TuningJob::builder() - .name("projects/test-project-123/locations/us-central1/tuningJobs/7".to_string()) - .state(JobState::JobStatePending) - .build()) - }); - aiplatform.expect_get_tuning_job().returning(|_| { - Ok(TuningJob::builder() - .state(JobState::JobStateFailed) - .error( - alien_gcp_clients::longrunning::Status::builder() - .code(3) - .message("training data malformed".to_string()) - .build(), - ) - .build()) - }); - - let provider = provider_with(aiplatform); - - let mut executor = SingleControllerExecutor::builder() - .resource(finetune_ai()) - .controller(GcpAiController::default()) - .platform(Platform::Gcp) - .service_provider(provider) - .with_dependency( - training_storage(), - GcpStorageController::mock_ready(TRAINING_STORAGE_ID), - ) - .build() - .await - .expect("executor should build"); - - // Fail-fast: the handler surfaces the terminal-failure as an error rather - // than silently reaching Ready on the untuned base model. In production the - // executor catches this and applies `on_failure = CreateFailed` - // (ProvisionFailed); the test harness surfaces the raw error, which we - // assert on directly. - let err = executor - .run_until_terminal() - .await - .expect_err("a failed tuning job must error out, not reach Ready"); - assert_eq!(err.code, "CLOUD_PLATFORM_ERROR"); + let binding: AiBinding = serde_json::from_value(params).expect("binding deserializes"); assert!( - err.to_string().contains("JobStateFailed") - && err.to_string().contains("training data malformed"), - "error must name the terminal state and carry the provider detail: {err}" + binding.tuned_model().is_none(), + "the controller must not attach a tuned model; the gateway rediscovers it" + ); + let cap = binding + .finetune() + .expect("finetune binding must carry the fine-tuning capability"); + assert_eq!(cap.base_model, "gemini-2.0-flash-001"); + // Bucket resolved from the storage dependency (test-stack-, from + // GcpStorageController::mock_ready), not re-derived here. + assert_eq!(cap.training_bucket, "test-stack-training-set"); + assert_eq!(cap.training_key, "training.jsonl"); + assert_eq!(cap.served_model_id, "llm-tuned"); + // Deterministic {prefix}-{id}; the executor test harness uses the + // resource prefix "test" (distinct from the storage mock's "test-stack" + // bucket-name prefix). + assert_eq!(cap.job_name, "test-llm"); + assert!( + cap.role_arn.is_empty(), + "Vertex submits under the ambient identity; no role is passed" ); - - // The resource never reached Running and recorded no tuned model, so the - // gateway would not serve an untuned model as if it were tuned. - assert_ne!(executor.status(), ResourceStatus::Running); - let controller = executor - .internal_state::() - .expect("controller downcast"); - assert!(controller.tuned_model_upstream_id.is_none()); } - /// Regression: an Ai without a finetune spec reaches Ready with an untuned - /// binding and never touches the aiplatform client. + /// Regression: an Ai without a finetune spec reaches Ready with a plain + /// binding carrying neither a tuned model nor a finetune capability. #[tokio::test] - async fn no_finetune_reaches_ready_with_untuned_binding() { - // A create-tuning-job expectation would fail if the controller called it. - let aiplatform = MockAiPlatformApi::new(); - let provider = provider_with(aiplatform); - + async fn no_finetune_reaches_ready_with_plain_binding() { let mut executor = SingleControllerExecutor::builder() .resource(base_ai()) .controller(GcpAiController::default()) .platform(Platform::Gcp) - .service_provider(provider) + .service_provider(provider()) .build() .await .expect("executor should build"); @@ -881,54 +583,22 @@ mod tests { .internal_state::() .expect("controller downcast"); assert!( - controller.tuning_job_name.is_none(), - "a pure-inference gateway never submits a tuning job" + controller.finetune.is_none(), + "a pure-inference gateway advertises no fine-tuning capability" ); let params = controller .get_binding_params() .expect("binding params ok") .expect("binding present"); - let binding: AiBinding = - serde_json::from_value(params).expect("binding deserializes"); + let binding: AiBinding = serde_json::from_value(params).expect("binding deserializes"); assert!( binding.tuned_model().is_none(), - "an untuned gateway must not carry a tuned model" + "an inference-only gateway must not carry a tuned model" ); - } - - /// Vertex does not expose LoRA for Gemini supervised tuning, so submit fails - /// loud rather than silently mapping it to something else. - #[tokio::test] - async fn unsupported_method_reaches_provision_failed() { - // No create call should happen — the method check rejects before submit. - let aiplatform = MockAiPlatformApi::new(); - let provider = provider_with(aiplatform); - - let mut executor = SingleControllerExecutor::builder() - .resource(lora_ai()) - .controller(GcpAiController::default()) - .platform(Platform::Gcp) - .service_provider(provider) - .with_dependency( - training_storage(), - GcpStorageController::mock_ready(TRAINING_STORAGE_ID), - ) - .build() - .await - .expect("executor should build"); - - // The method check rejects before any client call; the handler errors out - // (mapped to `on_failure = CreateFailed` by the executor in production). - let err = executor - .run_until_terminal() - .await - .expect_err("an unsupported tuning method must fail rather than mis-tune"); - assert_eq!(err.code, "RESOURCE_CONFIG_INVALID"); assert!( - err.to_string().to_lowercase().contains("supervised"), - "error should explain Vertex only supports supervised tuning: {err}" + binding.finetune().is_none(), + "an inference-only gateway must not carry a finetune capability" ); - assert_ne!(executor.status(), ResourceStatus::Running); } } diff --git a/crates/alien-infra/src/ai/gcp_import.rs b/crates/alien-infra/src/ai/gcp_import.rs index fa4cebf3c..635d58109 100644 --- a/crates/alien-infra/src/ai/gcp_import.rs +++ b/crates/alien-infra/src/ai/gcp_import.rs @@ -29,10 +29,9 @@ impl ResourceImporter for GcpAiImporter { state: GcpAiState::Ready, project: Some(data.project_id), location: Some(data.location), - // Import targets an existing base gateway; no tuning job is involved. - tuning_job_name: None, - tuned_model_upstream_id: None, - tuned_model_served_id: None, + // Import targets an existing base gateway; it advertises no + // fine-tuning capability. + finetune: None, _internal_stay_count: None, }; make_imported_state(controller, ctx) diff --git a/crates/alien-terraform/src/emitters/aws/ai.rs b/crates/alien-terraform/src/emitters/aws/ai.rs index 97d72caaa..d1c80d991 100644 --- a/crates/alien-terraform/src/emitters/aws/ai.rs +++ b/crates/alien-terraform/src/emitters/aws/ai.rs @@ -12,19 +12,27 @@ //! `stack_permission_sets`; resource-scoped grants are emitted here. use crate::{ + block::{attr, resource_block}, emitter::{TfEmitter, TfFragment}, emitters::aws::helpers::{ aws_terraform_permission_context, downcast, emit_iam_role_policy_for_target_with_label, - iam_policy_name_sanitize, required_label, + iam_policy_name_sanitize, iam_role_name_template, jsonencode, required_label, + service_assume_role_policy, tags, }, expr, }; use alien_core::{ import::EmitContext, Ai, PermissionProfile, PermissionSetReference, Result, ServiceAccount, + Storage, }; use alien_permissions::BindingTarget; use hcl::expr::Expression; +/// The `ai/finetune` permission set id. When a permission profile references it on +/// this AI resource, the emitter provisions a dedicated Bedrock-trusted IAM role so +/// Bedrock can assume it to read training data and write output. +const AI_FINETUNE_PERMISSION_ID: &str = "ai/finetune"; + #[derive(Debug, Clone, Copy, Default)] pub struct AwsAiEmitter; @@ -59,6 +67,16 @@ impl TfEmitter for AwsAiEmitter { } } + // When a permission profile references `ai/finetune` on this resource, emit a + // dedicated IAM role Bedrock can assume for model-customization jobs. The + // stack's service-account roles only trust compute principals + // (`lambda`/`codebuild`/`ec2`), so Bedrock cannot assume them — that is the + // real AccessDenied this role fixes. Its name matches the controller's + // `role_arn` (`{prefix}-{id}-finetune`) so the runtime gateway can pass it. + if resource_references_finetune(ctx) { + emit_finetune_role(&mut fragment, ctx, ai.id()); + } + Ok(fragment) } @@ -118,3 +136,106 @@ fn sanitize_label_segment(input: &str) -> String { }) .collect() } + +/// True when any permission profile references the `ai/finetune` set on this AI +/// resource. Only then is the Bedrock-trusted finetune role needed. +fn resource_references_finetune(ctx: &EmitContext<'_>) -> bool { + ctx.stack.permission_profiles().values().any(|profile| { + ai_permission_refs(profile, ctx.resource_id) + .iter() + .any(|reference| reference.id() == AI_FINETUNE_PERMISSION_ID) + }) +} + +/// Emit the dedicated Bedrock-trusted finetune IAM role plus its inline S3 policy. +/// +/// The role is named `{prefix}-{id}-finetune` (matching the controller's +/// `role_arn`), trusts `bedrock.amazonaws.com`, and can read training data +/// (`s3:GetObject`/`s3:ListBucket`) and write output (`s3:PutObject`) on every +/// storage bucket in the stack. +fn emit_finetune_role(fragment: &mut TfFragment, ctx: &EmitContext<'_>, ai_id: &str) { + let role_label = format!("{}_finetune", sanitize_label_segment(ai_id)); + + fragment.resource_blocks.push(resource_block( + "aws_iam_role", + &role_label, + [ + attr( + "name", + iam_role_name_template(&format!("{ai_id}-finetune")), + ), + attr( + "assume_role_policy", + service_assume_role_policy(&["bedrock.amazonaws.com"]), + ), + attr("tags", tags(ctx, "ai")), + ], + )); + + let statements = finetune_s3_statements(ctx); + fragment.resource_blocks.push(resource_block( + "aws_iam_role_policy", + &format!("{role_label}_s3"), + [ + attr( + "name", + Expression::String(format!("{ai_id}-finetune-s3")), + ), + attr( + "role", + expr::traversal(["aws_iam_role", role_label.as_str(), "id"]), + ), + attr( + "policy", + jsonencode(expr::object([ + ("Version", Expression::String("2012-10-17".to_string())), + ("Statement", Expression::Array(statements)), + ])), + ), + ], + )); +} + +/// S3 statements scoping the finetune role to the stack's storage buckets: +/// list/read on the bucket + objects, put on objects (training in, output out). +fn finetune_s3_statements(ctx: &EmitContext<'_>) -> Vec { + let mut bucket_arns = Vec::new(); + let mut object_arns = Vec::new(); + for (id, entry) in ctx.stack.resources() { + if entry.config.downcast_ref::().is_none() { + continue; + } + let Some(label) = ctx.name_for(id) else { + continue; + }; + bucket_arns.push(expr::traversal(["aws_s3_bucket", label, "arn"])); + object_arns.push(expr::template(format!("${{aws_s3_bucket.{label}.arn}}/*"))); + } + + vec![ + // Read the training dataset: list the bucket and get objects. + expr::object([ + ("Effect", Expression::String("Allow".to_string())), + ( + "Action", + Expression::Array(vec![ + Expression::String("s3:GetObject".to_string()), + Expression::String("s3:ListBucket".to_string()), + ]), + ), + ( + "Resource", + Expression::Array(bucket_arns.iter().cloned().chain(object_arns.iter().cloned()).collect()), + ), + ]), + // Write the tuning-job output back to storage. + expr::object([ + ("Effect", Expression::String("Allow".to_string())), + ( + "Action", + Expression::String("s3:PutObject".to_string()), + ), + ("Resource", Expression::Array(object_arns)), + ]), + ] +} diff --git a/crates/alien-terraform/tests/generator/aws_data_layer_tests.rs b/crates/alien-terraform/tests/generator/aws_data_layer_tests.rs index a18e20f5f..e135cd63d 100644 --- a/crates/alien-terraform/tests/generator/aws_data_layer_tests.rs +++ b/crates/alien-terraform/tests/generator/aws_data_layer_tests.rs @@ -7,8 +7,8 @@ use super::helpers::{assert_terraform_valid, render, snapshot_module}; use alien_core::{ - Ai, Kv, LifecycleRule, PermissionProfile, Queue, ResourceLifecycle, ServiceAccount, Stack, - StackSettings, Storage, Vault, + Ai, FinetuneMethod, FinetuneSpec, Kv, LifecycleRule, PermissionProfile, Queue, + ResourceLifecycle, ServiceAccount, Stack, StackSettings, Storage, Vault, }; use alien_terraform::TerraformTarget; @@ -203,3 +203,73 @@ fn aws_ai_invoke_permissions_attach_to_service_account_role() { ); assert_terraform_valid(&module, "aws_ai_invoke_permissions"); } + +#[test] +fn aws_ai_finetune_emits_bedrock_trusted_role_with_s3_policy() { + // When a permission profile references ai/finetune on the AI resource, the + // emitter provisions a dedicated IAM role Bedrock can assume (the real fix for + // AccessDenied: service-account roles only trust compute principals) with an + // inline S3 policy over the stack's storage buckets. + let stack = Stack::new("acme-ai".to_string()) + .permission( + "execution", + PermissionProfile::new().resource("llm", ["ai/finetune"]), + ) + .add( + ServiceAccount::new("execution-sa".to_string()).build(), + ResourceLifecycle::Frozen, + ) + .add( + Storage::new("training".to_string()).build(), + ResourceLifecycle::Frozen, + ) + .add( + Ai::new("llm".to_string()) + .finetune(FinetuneSpec { + base_model: "amazon.nova-lite-v1:0".to_string(), + training_data: "training".to_string(), + training_key: "training.jsonl".to_string(), + served_model_id: None, + method: FinetuneMethod::Sft, + }) + .build(), + ResourceLifecycle::Frozen, + ) + .build(); + let module = render(&stack, TerraformTarget::Aws, StackSettings::default()); + let rendered = module + .iter() + .map(|(_, contents)| contents) + .collect::(); + + // The dedicated finetune role exists, named deterministically to match the + // controller's role_arn (`${resource_prefix}-llm-finetune`). + assert!( + rendered.contains("llm_finetune"), + "expected a dedicated finetune aws_iam_role:\n{rendered}" + ); + assert!( + rendered.contains("llm-finetune"), + "expected role name suffix llm-finetune matching the controller role_arn:\n{rendered}" + ); + // Its trust policy allows Bedrock to assume it — the crux of the fix. + assert!( + rendered.contains("bedrock.amazonaws.com"), + "finetune role must trust bedrock.amazonaws.com:\n{rendered}" + ); + // The inline policy grants S3 read (training data) and write (output), + // scoped to the stack's storage bucket ARN (not "*"). + assert!( + rendered.contains("s3:GetObject") && rendered.contains("s3:ListBucket"), + "finetune role must read the training dataset from S3:\n{rendered}" + ); + assert!( + rendered.contains("s3:PutObject"), + "finetune role must write tuning output to S3:\n{rendered}" + ); + assert!( + rendered.contains("aws_s3_bucket.training.arn"), + "S3 grants must reference the storage bucket ARN, not \"*\":\n{rendered}" + ); + assert_terraform_valid(&module, "aws_ai_finetune_role"); +} diff --git a/examples/ai-finetune-inference-ts/README.md b/examples/ai-finetune-inference-ts/README.md index 5b6165ffb..9a23e1c68 100644 --- a/examples/ai-finetune-inference-ts/README.md +++ b/examples/ai-finetune-inference-ts/README.md @@ -10,25 +10,37 @@ Fine-tune a base model **inside the customer's cloud**, then serve it for infere ## How it works +Fine-tuning is triggered **at runtime by the app**, not at deploy time. The +`alien.AI("llm")` resource provisions and is **Ready immediately** — it's just the +inference gateway plus a fine-tuning capability. The app starts a job by calling the +gateway; the resource never blocks on a hours-long training job. + ``` -alien.Storage("dataset") # S3 / GCS / Blob in the customer's account - │ training.jsonl uploaded by the worker - ▼ -alien.AI("llm").finetune({...}) # on deploy, the controller submits the - │ # provider's tuning job reading `dataset`, - │ # polls to completion, records the artifact - ▼ -gateway serves "support-tuned" # tuned model routed alongside the base catalog +deploy: alien.Storage("finetune-training-data") # S3 / GCS / Blob, customer's account + alien.AI("llm").finetune({...}) # Ready immediately — capability only + # (no job runs at deploy) + +runtime: POST /dataset → upload training.jsonl into the bucket + POST /finetune → ai("llm").finetune() submits the cloud job → { jobId } + GET /finetune/status → ai("llm").finetuneStatus(jobId) → running|succeeded|failed + POST /chat-tuned → inference on "support-tuned" once the job succeeded ``` -At deploy time the AI resource's cloud controller submits the tuning job (Bedrock `CreateModelCustomizationJob`, Vertex `tuningJobs`, or Foundry `fine_tuning.jobs`), polls it to completion via its heartbeat loop, and records the tuned artifact. The gateway then routes the public id `support-tuned` to that artifact — so app code calls it exactly like a base model, only the `model` string differs. +When the app calls `ai("llm").finetune(...)`, the gateway submits the provider's tuning +job (Bedrock `CreateModelCustomizationJob`, Vertex `tuningJobs`, or Foundry +`fine_tuning.jobs`) under the workload's ambient identity, reading the dataset from the +customer's bucket, and returns a job id. Inference for `support-tuned` works once the job +succeeds: the gateway **rediscovers** the completed tuned model by convention (no stored +state) and routes to it — so app code calls it exactly like a base model, only the `model` +string differs. ## API | Route | Method | Purpose | |-------|--------|---------| | `/dataset` | POST | Upload JSONL training data into the customer's bucket (body = JSONL) | -| `/finetune/status` | GET | `ready` once the tuning job has completed, else `pending` | +| `/finetune` | POST | Start a tuning job at runtime; returns `{ jobId, servedModel }` | +| `/finetune/status` | GET | Poll a job: `?jobId=` → `{ status, model?, message? }` | | `/chat` | POST | Inference against a base foundation model (`{ "message": "..." }`) | | `/chat-tuned` | POST | Inference against the fine-tuned model — same call, different model id | @@ -39,15 +51,19 @@ npm install alien dev ``` -Upload the sample dataset and query the tuned model: +Upload data, start a job, poll it, then query the tuned model: ```bash curl -X POST --data-binary @sample-training.jsonl http://localhost:8080/dataset -curl http://localhost:8080/finetune/status +JOB=$(curl -s -X POST http://localhost:8080/finetune | jq -r .jobId) +curl "http://localhost:8080/finetune/status?jobId=$JOB" +# once succeeded: curl -X POST http://localhost:8080/chat-tuned -d '{"message":"How do I reset my password?"}' ``` -On the **local** platform the AI resource is a BYO-key provider (set `OPENAI_API_KEY`); fine-tuning is a managed-cloud capability, so `/finetune/status` reports `pending` locally and the tuned route falls back to the base model. Deploy to a cloud to exercise the real tuning flow. +On the **local** platform the AI resource is a BYO-key provider (set `OPENAI_API_KEY`); +fine-tuning is a managed-cloud capability, so `POST /finetune` is rejected locally and the +tuned route falls back to the base model. Deploy to a cloud to exercise the real tuning flow. ## Picking `baseModel` per cloud diff --git a/examples/ai-finetune-inference-ts/alien.ts b/examples/ai-finetune-inference-ts/alien.ts index 1f647e08a..7e23128fb 100644 --- a/examples/ai-finetune-inference-ts/alien.ts +++ b/examples/ai-finetune-inference-ts/alien.ts @@ -14,11 +14,13 @@ import * as alien from "@alienplatform/core" // (and change it if you hit "bucket name is not available"). const dataset = new alien.Storage("finetune-training-data").build() -// A model-less AI gateway that also fine-tunes. On deploy, the resource's cloud -// controller submits the provider's tuning job (Bedrock CreateModelCustomizationJob, -// Vertex tuningJobs, or Foundry fine_tuning.jobs) reading `dataset` in the -// customer's account, polls it to completion, and serves the tuned model through -// the same gateway under `servedModelId`. Base models remain callable too. +// A model-less AI gateway with a fine-tuning CAPABILITY. `.finetune(...)` here is a +// declaration, not a deploy-time trigger: the resource provisions and is Ready +// immediately (no job runs at deploy). The app starts a tuning job at RUNTIME by +// calling `ai("llm").finetune(...)` (see src/index.ts) — the gateway then submits the +// provider's job (Bedrock CreateModelCustomizationJob, Vertex tuningJobs, or Foundry +// fine_tuning.jobs) reading `dataset` in the customer's account, and serves the tuned +// model under `servedModelId` once it completes. Base models remain callable too. // // `baseModel` is a provider-native id; pick the one that matches the cloud you // deploy to (see the README's per-provider table). The default here targets diff --git a/examples/ai-finetune-inference-ts/src/index.ts b/examples/ai-finetune-inference-ts/src/index.ts index c208e4f40..4c46b0879 100644 --- a/examples/ai-finetune-inference-ts/src/index.ts +++ b/examples/ai-finetune-inference-ts/src/index.ts @@ -7,10 +7,9 @@ const TRAINING_KEY = "training.jsonl" const app = new Hono() -// 1. Upload the JSONL training set into the customer's bucket. The tuning job -// (submitted by the AI resource's controller at deploy time) reads it from -// there — the data never leaves the customer's cloud. In a real app you'd -// seed this before deploy; exposed here so the flow is runnable end-to-end. +// 1. Upload the JSONL training set into the customer's bucket. Nothing is tuned yet — +// the AI resource is already provisioned and Ready; fine-tuning is triggered at +// runtime (step 2). The data never leaves the customer's cloud. app.post("/dataset", async c => { const body = await c.req.text() if (!body.trim()) { @@ -21,25 +20,34 @@ app.post("/dataset", async c => { return c.json({ uploaded: TRAINING_KEY, examples: lines }) }) -// 2. Fine-tune status. The tuned model shows up in the gateway's model list only -// once its job has completed, so its presence is a simple readiness signal. +// 2. Trigger fine-tuning at runtime. This is the whole point: the provisioned `llm` +// resource is always Ready, and the app itself kicks off a job by calling the +// gateway — `ai("llm").finetune(...)`. The gateway submits the cloud tuning job +// (Bedrock / Vertex / Foundry) under the workload's ambient identity and returns a +// job id to poll. Long-running (minutes to hours); this returns immediately. +app.post("/finetune", async c => { + const { jobId, servedModel } = await ai("llm").finetune({ trainingKey: TRAINING_KEY }) + return c.json({ jobId, servedModel, message: "tuning started; poll /finetune/status?jobId=" }) +}) + +// 3. Poll a job. The gateway queries the cloud live (stateless — no job state stored). app.get("/finetune/status", async c => { - const models = await ai("llm").getAvailableModels() - const ready = models.some(m => m.id === TUNED_MODEL) - return c.json({ - tunedModel: TUNED_MODEL, - status: ready ? "ready" : "pending", - availableModels: models.map(m => m.id), - }) + const jobId = c.req.query("jobId") + if (!jobId) { + return c.json({ error: "pass ?jobId= from the POST /finetune response" }, 400) + } + const state = await ai("llm").finetuneStatus(jobId) + return c.json(state) }) -// 3. Inference against a base foundation model (the per-cloud catalog). +// 4. Inference against a base foundation model (the per-cloud catalog). app.post("/chat", async c => { return chat(c, await defaultBaseModel()) }) -// 4. Inference against the fine-tuned model — same OpenAI-compatible call, just a -// different `model` id. The gateway routes it to the tuned artifact in-account. +// 5. Inference against the fine-tuned model — same OpenAI-compatible call, just a +// different `model` id. Works once the job has succeeded: the gateway rediscovers +// the completed tuned model by convention and routes to it. app.post("/chat-tuned", async c => { return chat(c, TUNED_MODEL) }) diff --git a/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts b/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts index d3696c4f1..638211f5e 100644 --- a/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts +++ b/examples/ai-finetune-inference-ts/tests/ai-finetune-inference.test.ts @@ -40,18 +40,18 @@ describe("ai-finetune-inference-ts", () => { expect(response.status).toBe(400) }) - it("reports a well-formed fine-tune status", async () => { + it("requires a jobId to poll status", async () => { + // /finetune/status needs the jobId from a POST /finetune response. const response = await fetch(`${deployment.url}/finetune/status`) - expect(response.status).toBe(200) - const body = (await response.json()) as { - tunedModel: string - status: string - availableModels: string[] - } - expect(body.tunedModel).toBe("support-tuned") - // Locally the managed tuning job never runs, so the tuned model is not ready. - expect(body.status).toBe("pending") - expect(Array.isArray(body.availableModels)).toBe(true) + expect(response.status).toBe(400) + }) + + it("rejects starting a job on the local (BYO-key) platform", async () => { + // Fine-tuning is a managed-cloud capability; the local platform serves the AI + // resource as a BYO-key provider, so triggering a job is not supported here. + // (On a real cloud deploy this returns { jobId, servedModel }.) + const response = await fetch(`${deployment.url}/finetune`, { method: "POST" }) + expect(response.status).toBeGreaterThanOrEqual(400) }) it.skipIf(!OPENAI_KEY)("answers a base-model chat when a provider key is set", async () => { diff --git a/internal-docs/alien/ai-finetune-runtime-design.md b/internal-docs/alien/ai-finetune-runtime-design.md new file mode 100644 index 000000000..7e6c71736 --- /dev/null +++ b/internal-docs/alien/ai-finetune-runtime-design.md @@ -0,0 +1,105 @@ +# AI fine-tuning: runtime, API-triggered design + +Status: implementing. Supersedes the deploy-time tuning model. + +## Why this changed + +The first implementation submitted the cloud tuning job at **deploy time**, inside the +`Ai` resource controller, and blocked the resource in a `WaitingForTuningJob` state until +the job completed. Three problems made that wrong: + +1. **Empty-bucket ordering.** The job wants training data at deploy time, but apps upload + data at runtime (there is no data in the bucket on first deploy). +2. **Hours-long deploys.** Fine-tuning takes minutes to hours; a resource that isn't + `Ready` until training finishes makes `alien release` block for hours. +3. **Retraining = redeploy.** Training on new data meant editing the resource and + redeploying. + +Fine-tuning is fundamentally an **imperative, long-running job**, not a declarative +resource state. So we move the trigger to runtime. + +## The model + +- `.finetune({...})` on the `Ai` resource is now a **capability declaration**: it says + "this gateway may fine-tune `baseModel`, reading data from `trainingData`, serving the + result as `servedModelId`", and it drives the `ai/finetune` permission grant. It does + **not** start a job. +- The `Ai` controller reaches `Ready` immediately (no tuning states). +- At runtime the app calls the gateway to start/track jobs: + - `ai("llm").finetune({ trainingKey? }) -> { jobId }` — submits the cloud tuning job. + - `ai("llm").finetuneStatus(jobId) -> { status, model? }` — polls it. +- Inference against `servedModelId` works as soon as the tuned model is `Active`, via + **rediscovery by convention** (below) — no state store. + +## Gateway control-plane surface + +New routes on the in-process gateway (alongside the inference proxy): + +- `POST //v1/finetune` — body `{ trainingKey?, baseModel?, method? }`. Submits + the provider job using the binding's ambient credential and the `finetune` capability + from the binding. Returns `{ jobId, servedModel }`. +- `GET //v1/finetune/` — returns `{ status: "pending"|"running"|"succeeded"|"failed", model? }`. + +The gateway already holds the ambient credential and signs arbitrary +host/service requests (`AmbientCred::authorize(req, service)`), so control-plane calls +(`bedrock` control host, Vertex `tuningJobs`, Foundry `fine_tuning.jobs`) reuse the same +credential path as inference. No new credential wiring. + +### Per-cloud provider trait + +```rust +#[async_trait] +trait FineTuneProvider { + async fn submit(&self, spec: &FineTuneRequest) -> Result; + async fn status(&self, job: &str) -> Result; // by job id + async fn resolve_served_model(&self, served_id: &str) -> Result>; // rediscovery +} +``` + +- **Bedrock**: submit = `CreateModelCustomizationJob`; status = `GetModelCustomizationJob`; + rediscovery = `GetCustomModel()` → if `modelStatus == Active`, its + `modelArn` is the upstream id. `GetCustomModel` accepts the model **name**, so no ARN + needs storing. +- **Vertex**: submit = `POST tuningJobs`; status = `GET tuningJobs/{id}`; rediscovery = + the tuned endpoint id derived from the job / a list filtered by display name. +- **Foundry**: submit = `POST fine_tuning/jobs`; status = `GET fine_tuning/jobs/{id}`; + on success the job carries `fine_tuned_model`; rediscovery = the deployment named after + `servedModelId` (create-on-first-success), or `GET deployments/{name}`. + +## Rediscovery by convention (stateless) + +The gateway is per-process and stateless. A job started by one worker completes on the +cloud's side hours later, possibly after that worker is gone. So the gateway does **not** +track job state across restarts. Instead: + +- The tuned model's cloud name is **deterministic** from the binding + `servedModelId`. +- On an inference request for `servedModelId`, if the route has no cached tuned upstream, + the gateway calls `resolve_served_model(servedId)`; if the provider reports the model + `Active`, it caches and routes to it. If not ready, it returns `model not available`. +- `GET /finetune/` likewise queries the cloud live each call. + +No storage dependency, no background poller. The trade-off: a completed job is only +"noticed" on the next status poll or inference request (fine — it's a pull model). + +## Role / credentials note (fixes the deploy-time role-ARN bug) + +The deploy-time version derived a job `roleArn` (`{prefix}-{id}`) that didn't reliably +exist. In the runtime model the gateway submits the job under the **workload's ambient +identity** — the same identity it uses for inference. Bedrock's `CreateModelCustomizationJob` +still needs a `roleArn` it can assume to read S3 / write output, so the gateway resolves +the workload's actual execution role (or a dedicated finetune role the `ai/finetune` +permission set provisions with a `bedrock.amazonaws.com` trust policy). This is resolved +at submit time from the binding, not guessed from a naming convention. + +## Layers touched + +- `alien-gateway`: new `finetune` module (trait + 3 providers), 2 routes, tuned-model + cache + rediscovery in the router. +- `packages/ai-gateway` (SDK): `finetune()` / `finetuneStatus()` on the `ai()` client. +- `alien-core`: binding carries the `finetune` capability (baseModel, trainingData bucket, + servedModelId, method) so the gateway can submit without a control-plane round-trip to + the resource. `FinetuneSpec` stays on the resource as the declaration. +- `alien-infra` controllers: drop `SubmittingTuningJob`/`WaitingForTuningJob`; reach + `Ready` immediately; still emit the `ai/finetune` grant and the capability in the binding. +- Example `ai-finetune-inference-ts`: `POST /dataset` → `POST /finetune` → poll + `/finetune/status` → `POST /chat-tuned`. diff --git a/packages/ai-gateway/src/__tests__/client.test.ts b/packages/ai-gateway/src/__tests__/client.test.ts index 9039f4e5e..f3c9d2096 100644 --- a/packages/ai-gateway/src/__tests__/client.test.ts +++ b/packages/ai-gateway/src/__tests__/client.test.ts @@ -278,3 +278,92 @@ describe("Ai.responses.create", () => { ) }) }) + +// ───────────────────────────────────────────────────────────────────────────── +// Ai.finetune / Ai.finetuneStatus — runtime fine-tuning (ambient gateway only) +// ───────────────────────────────────────────────────────────────────────────── + +const AMBIENT = JSON.stringify({ service: "bedrock", region: "us-east-2" }) + +describe("Ai.finetune", () => { + it("POSTs the training key to the gateway and returns { jobId, servedModel }", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + const fetchMock = stubFetch({ jobId: "job-123", servedModel: "llm-tuned" }) + + const result = await ai("llm").finetune({ trainingKey: "datasets/train.jsonl" }) + + expect(callUrl(fetchMock)).toBe(`${GATEWAY_URL}/llm/v1/finetune`) + expect(callInit(fetchMock).method).toBe("POST") + expect(callBody(fetchMock)).toEqual({ trainingKey: "datasets/train.jsonl" }) + // Ambient path injects the credential in the gateway; no client-side auth header. + const headers = (callInit(fetchMock).headers ?? {}) as Record + expect(headers.Authorization).toBeUndefined() + expect(result).toEqual({ jobId: "job-123", servedModel: "llm-tuned" }) + }) + + it("omits trainingKey from the body when not provided", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + const fetchMock = stubFetch({ jobId: "job-1", servedModel: "llm-tuned" }) + + await ai("llm").finetune() + + expect(callBody(fetchMock)).toEqual({}) + }) + + it("rejects for a BYO-key External binding without POSTing", async () => { + // EXTERNAL binding is stubbed in the top-level beforeEach. + const fetchMock = stubFetch({ jobId: "x", servedModel: "y" }) + await expect(ai("llm").finetune({ trainingKey: "k" })).rejects.toMatchObject({ + code: "INVALID_BINDING_CONFIG", + }) + expect(fetchMock).not.toHaveBeenCalled() + }) + + it("throws an AiUpstreamError on a non-2xx gateway response", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + stubFetch({ error: { message: "boom" } }, 500) + await expect(ai("llm").finetune({ trainingKey: "k" })).rejects.toThrow(AlienError) + }) +}) + +describe("Ai.finetuneStatus", () => { + it("GETs the job by id and maps a succeeded status with the tuned model", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + const fetchMock = stubFetch({ status: "succeeded", model: "llm-tuned-v1" }) + + const status = await ai("llm").finetuneStatus("job-123") + + expect(callUrl(fetchMock)).toBe(`${GATEWAY_URL}/llm/v1/finetune/job-123`) + expect(callInit(fetchMock).method).toBe("GET") + expect(status).toEqual({ status: "succeeded", model: "llm-tuned-v1" }) + }) + + it("maps a running status", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + stubFetch({ status: "running" }) + expect(await ai("llm").finetuneStatus("job-1")).toEqual({ status: "running" }) + }) + + it("URL-encodes the job id", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + const fetchMock = stubFetch({ status: "running" }) + await ai("llm").finetuneStatus("arn:aws:bedrock/job 1") + expect(callUrl(fetchMock)).toBe( + `${GATEWAY_URL}/llm/v1/finetune/${encodeURIComponent("arn:aws:bedrock/job 1")}`, + ) + }) + + it("rejects for a BYO-key External binding without a GET", async () => { + const fetchMock = stubFetch({ status: "running" }) + await expect(ai("llm").finetuneStatus("job-1")).rejects.toMatchObject({ + code: "INVALID_BINDING_CONFIG", + }) + expect(fetchMock).not.toHaveBeenCalled() + }) + + it("throws an AiUpstreamError on a non-2xx gateway response", async () => { + vi.stubEnv("ALIEN_LLM_BINDING", AMBIENT) + stubFetch({ error: { message: "not found" } }, 404) + await expect(ai("llm").finetuneStatus("missing")).rejects.toThrow(AlienError) + }) +}) diff --git a/packages/ai-gateway/src/client.ts b/packages/ai-gateway/src/client.ts index fa7cfddc5..7742f7de0 100644 --- a/packages/ai-gateway/src/client.ts +++ b/packages/ai-gateway/src/client.ts @@ -11,7 +11,12 @@ import { AlienError } from "@alienplatform/core" import { isExternalAiBinding, parseAiBinding } from "./binding.js" -import { AiTransportError, AiUpstreamError, BindingNotFoundError } from "./errors.js" +import { + AiTransportError, + AiUpstreamError, + BindingNotFoundError, + InvalidBindingConfigError, +} from "./errors.js" import type { Gateway } from "./gateway.js" // ───────────────────────────────────────────────────────────────────────────── @@ -37,6 +42,24 @@ export interface AiModel { id: string } +/** The handle returned when a runtime fine-tuning job is submitted through the gateway. */ +export interface FinetuneResult { + /** Gateway-assigned id used to poll the job's status. */ + jobId: string + /** The served model id inference should target once the job succeeds. */ + servedModel: string +} + +/** The status of a runtime fine-tuning job, as reported by the gateway. */ +export interface FinetuneJobStatus { + /** Lifecycle state of the job. */ + status: "running" | "succeeded" | "failed" + /** The tuned model id, present once the job has succeeded. */ + model?: string + /** A human-readable detail (e.g. the failure reason), when the gateway provides one. */ + message?: string +} + // Upstream base URL (no `/v1`) for a BYO-key provider. `ALIEN_AI_LOCAL_BASE_URL` overrides it so // any OpenAI-compatible provider works; unknown providers default to OpenAI's. function providerBaseUrl(provider: string): string { @@ -270,6 +293,57 @@ export class Ai { return body.data } + /** + * Submit a runtime fine-tuning job for this binding's model. Only the ambient (gateway) + * path supports fine-tuning — a BYO-key External binding has no gateway finetune endpoint, + * so this rejects rather than POSTing to the raw provider. + */ + async finetune(opts?: { trainingKey?: string }): Promise { + const { baseUrl, apiKey } = await this.connection() + if (apiKey) throw finetuneUnsupportedError() + + const url = `${baseUrl}/v1/finetune` + // Omit trainingKey entirely when not supplied; the capability already carries a default key. + const body = opts?.trainingKey === undefined ? {} : { trainingKey: opts.trainingKey } + const response = await this._fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }) + if (!response.ok) { + throw createUpstreamError(url, response.status, await extractErrorMessage(response)) + } + return this._parseJson(url, response) + } + + /** Poll a runtime fine-tuning job by its id. Ambient (gateway) path only. */ + async finetuneStatus(jobId: string): Promise { + const { baseUrl, apiKey } = await this.connection() + if (apiKey) throw finetuneUnsupportedError() + + const url = `${baseUrl}/v1/finetune/${encodeURIComponent(jobId)}` + const response = await this._fetch(url, { method: "GET" }) + if (!response.ok) { + throw createUpstreamError(url, response.status, await extractErrorMessage(response)) + } + return this._parseJson(url, response) + } + + // Parse a JSON response body, wrapping a parse failure in an AiTransportError like the + // other JSON-returning methods (getAvailableModels, _postSurface) do. + private async _parseJson(url: string, response: Response): Promise { + try { + return (await response.json()) as T + } catch (jsonError) { + throw (await AlienError.from(jsonError)).withContext( + AiTransportError.create({ + url, + reason: `Response body is not valid JSON: ${jsonError instanceof Error ? jsonError.message : String(jsonError)}`, + }), + ) + } + } + private _chatCompletionsCreate(params: ChatCompletionCreateParams): Promise { return this._postSurface("/v1/chat/completions", params) } @@ -385,6 +459,17 @@ export function createAiClient(gateway: Gateway): AiClient { // Helpers // ───────────────────────────────────────────────────────────────────────────── +// Fine-tuning is a capability of the ambient gateway; a BYO-key (External) provider is reached +// directly and has no gateway finetune endpoint, so the operation is unsupported there. +function finetuneUnsupportedError() { + return new AlienError( + InvalidBindingConfigError.create({ + message: "Fine-tuning is not supported for BYO-key (External) AI providers", + suggestion: "Use an ambient-cloud AI binding with a fine-tune capability configured", + }), + ) +} + async function extractErrorMessage(response: Response): Promise { try { const errBody = (await response.json()) as Record diff --git a/packages/ai-gateway/src/index.ts b/packages/ai-gateway/src/index.ts index 358620465..9e0a664dc 100644 --- a/packages/ai-gateway/src/index.ts +++ b/packages/ai-gateway/src/index.ts @@ -36,6 +36,8 @@ export type { AiConnection, AiModel, ChatCompletionCreateParams, + FinetuneJobStatus, + FinetuneResult, ResponseCreateParams, } from "./client.js" export { aiBindingEnvVarName, isExternalAiBinding, parseAiBinding } from "./binding.js" diff --git a/packages/core/src/generated/schemas/ai.json b/packages/core/src/generated/schemas/ai.json index a0f8fc14a..33a143f78 100644 --- a/packages/core/src/generated/schemas/ai.json +++ b/packages/core/src/generated/schemas/ai.json @@ -1 +1 @@ -{"type":"object","description":"Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.\n\nWhen `finetune` is set, the resource also tunes a base model in the\ncustomer's cloud and serves the result through the same gateway (see\n[`FinetuneSpec`]).","required":["id"],"properties":{"finetune":{"oneOf":[{"type":"null"},{"description":"Optional fine-tuning declaration. When present, the resource tunes\n`finetune.base_model` on the customer's cloud and serves the result\nalongside the base models. Absent for a pure inference gateway.","type":"object","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"}]},"id":{"type":"string","description":"Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters."}},"additionalProperties":false,"x-readme-ref-name":"Ai"} \ No newline at end of file +{"type":"object","description":"Represents an AI Gateway resource that provides a unified interface to\nmanaged AI inference services across cloud providers.\n\nBYO-key external providers (OpenAI/Anthropic) are NOT declared here. Like any\nother BYO infrastructure (e.g. external Redis for `kv`), an external AI\nprovider is supplied at deploy time as an `ExternalBinding::Ai` in the\nstack's external-bindings map; the executor then skips the cloud controller.\n\nWhen `finetune` is set, the resource also tunes a base model in the\ncustomer's cloud and serves the result through the same gateway (see\n[`FinetuneSpec`]).","required":["id"],"properties":{"finetune":{"oneOf":[{"type":"null"},{"description":"Optional fine-tuning declaration. When present, the resource tunes\n`finetune.base_model` on the customer's cloud and serves the result\nalongside the base models. Absent for a pure inference gateway.","type":"object","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when a runtime job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The gateway\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"}]},"id":{"type":"string","description":"Identifier for the AI resource. Must contain only alphanumeric characters, hyphens, and underscores ([A-Za-z0-9-_]).\nMaximum 64 characters."}},"additionalProperties":false,"x-readme-ref-name":"Ai"} \ No newline at end of file diff --git a/packages/core/src/generated/schemas/finetuneSpec.json b/packages/core/src/generated/schemas/finetuneSpec.json index 347bf463f..eef867b2d 100644 --- a/packages/core/src/generated/schemas/finetuneSpec.json +++ b/packages/core/src/generated/schemas/finetuneSpec.json @@ -1 +1 @@ -{"type":"object","description":"Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider's tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before.","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"} \ No newline at end of file +{"type":"object","description":"Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThis is a *capability declaration*, not a deploy-time trigger: a resource\nwith a `finetune` spec provisions and is Ready immediately (no job runs at\ndeploy). The declaration flows to the gateway as a fine-tuning capability;\nthe app then starts a job at runtime by calling `ai(\"\").finetune(...)`,\nwhich the gateway submits to the provider (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`) under the workload's ambient identity, reading the\ntraining data from the customer-owned [`Storage`](crate::Storage) bucket\n(S3 / GCS / Blob) referenced by `training_data`. Once a job completes, the\ngateway serves the tuned model under `served_model_id`, rediscovering it by\nconvention. Base-model inference is unaffected — an `Ai` without a\n`finetune` spec behaves exactly as before.","required":["baseModel","trainingData"],"properties":{"baseModel":{"type":"string","description":"Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when a runtime job is submitted."},"method":{"description":"The fine-tuning method. Defaults to supervised fine-tuning.","type":"string","enum":["sft","dpo","lora"],"x-readme-ref-name":"FinetuneMethod"},"servedModelId":{"type":["string","null"],"description":"The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`."},"trainingData":{"type":"string","description":"The storage resource holding the JSONL training dataset. The gateway\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."},"trainingKey":{"type":"string","description":"Object key of the training file within `training_data`.\nDefaults to `training.jsonl`."}},"additionalProperties":false,"x-readme-ref-name":"FinetuneSpec"} \ No newline at end of file diff --git a/packages/core/src/generated/zod/finetune-spec-schema.ts b/packages/core/src/generated/zod/finetune-spec-schema.ts index 68b608447..8767ce2df 100644 --- a/packages/core/src/generated/zod/finetune-spec-schema.ts +++ b/packages/core/src/generated/zod/finetune-spec-schema.ts @@ -7,16 +7,16 @@ import * as z from "zod"; import { FinetuneMethodSchema } from "./finetune-method-schema.js"; /** - * @description Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer\'s cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider\'s tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before. + * @description Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer\'s cloud before serving it.\n\nThis is a *capability declaration*, not a deploy-time trigger: a resource\nwith a `finetune` spec provisions and is Ready immediately (no job runs at\ndeploy). The declaration flows to the gateway as a fine-tuning capability;\nthe app then starts a job at runtime by calling `ai(\"\").finetune(...)`,\nwhich the gateway submits to the provider (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`) under the workload\'s ambient identity, reading the\ntraining data from the customer-owned [`Storage`](crate::Storage) bucket\n(S3 / GCS / Blob) referenced by `training_data`. Once a job completes, the\ngateway serves the tuned model under `served_model_id`, rediscovering it by\nconvention. Base-model inference is unaffected — an `Ai` without a\n`finetune` spec behaves exactly as before. */ export const FinetuneSpecSchema = z.object({ - "baseModel": z.string().describe("Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when the job is submitted."), + "baseModel": z.string().describe("Provider-native base-model identifier to tune (e.g. an Amazon Nova model\nid on Bedrock, a Gemini model on Vertex, or a gpt-4o family model on\nFoundry). Validated against the target cloud when a runtime job is submitted."), get "method"(){ return FinetuneMethodSchema.describe("The fine-tuning method applied to the base model.\n\nThe gateway-side controllers map each variant onto the provider's native\ntechnique: `Sft` is supervised fine-tuning (all three clouds), `Dpo` is\ndirect preference optimization (Bedrock/Foundry), and `Lora` requests a\nparameter-efficient adapter where the provider exposes it. Providers that\nonly implement a subset reject unsupported methods at job-submit time.").optional() }, "servedModelId": z.string().describe("The public model id apps use to invoke the tuned model through the\ngateway (the `model` field in an OpenAI-compatible request). Defaults to\n`-tuned`.").nullish(), -"trainingData": z.string().describe("The storage resource holding the JSONL training dataset. The controller\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."), +"trainingData": z.string().describe("The storage resource holding the JSONL training dataset. The gateway\nreads it from the customer bucket the storage resolves to; the data\nnever leaves the customer's cloud."), "trainingKey": z.optional(z.string().describe("Object key of the training file within `training_data`.\nDefaults to `training.jsonl`.")) - }).describe("Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThe training data lives in a customer-owned [`Storage`](crate::Storage)\nbucket (S3 / GCS / Blob), referenced by `training_data`. At deploy time the\ncloud controller submits the provider's tuning job (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`), polls it to completion via the heartbeat loop, and\nrecords the resulting artifact so the gateway can route `served_model_id`\nto it. Base-model inference is unaffected — an `Ai` without a `finetune`\nspec behaves exactly as before.") + }).describe("Declares that an [`Ai`] resource should fine-tune a base model in the\ncustomer's cloud before serving it.\n\nThis is a *capability declaration*, not a deploy-time trigger: a resource\nwith a `finetune` spec provisions and is Ready immediately (no job runs at\ndeploy). The declaration flows to the gateway as a fine-tuning capability;\nthe app then starts a job at runtime by calling `ai(\"\").finetune(...)`,\nwhich the gateway submits to the provider (Bedrock\n`CreateModelCustomizationJob`, Vertex tuning job, or Foundry\n`fine_tuning.jobs`) under the workload's ambient identity, reading the\ntraining data from the customer-owned [`Storage`](crate::Storage) bucket\n(S3 / GCS / Blob) referenced by `training_data`. Once a job completes, the\ngateway serves the tuned model under `served_model_id`, rediscovering it by\nconvention. Base-model inference is unaffected — an `Ai` without a\n`finetune` spec behaves exactly as before.") export type FinetuneSpec = z.infer \ No newline at end of file diff --git a/packages/sdk/src/index.ts b/packages/sdk/src/index.ts index 2c0d7d8c7..e1231f1fc 100644 --- a/packages/sdk/src/index.ts +++ b/packages/sdk/src/index.ts @@ -76,6 +76,8 @@ export type { AiConnection, AiModel, ChatCompletionCreateParams, + FinetuneJobStatus, + FinetuneResult, ResponseCreateParams, } from "@alienplatform/ai-gateway"