Lora support in Model Ext (#8199)

* adding more enhancements

* releasing v04 preview

* updating snapshot

* adding model to registry

* Update CLI snapshots for registry.json changes

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add LoRA config support to models create and show commands

- Add LoRAConfig struct (rank, alpha, targetModules, dropout) to CustomModel
- Add --lora-rank, --lora-alpha, --lora-target-modules, --lora-dropout flags
  to 'models create' with validation when --weight-type is LoRA
- Display LoRA Configuration section in 'models show' table output
- Add LoRAConfig field to RegisterModelRequest for API serialization
- Add WeightType column to 'models list' table view
- Add tests for flag validation, JSON round-trip, and omitempty behavior

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Bump azure.ai.models to 0.0.7-preview with LoRA changelog

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Move LoRA flag validation before upload steps

Validate --lora-* flags and --weight-type consistency before
StartPendingUpload and AzCopy upload to fail fast on invalid input
without creating unnecessary upload artifacts.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* chore: trigger check-enforcer re-evaluation

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: achauhan-scc

cli/azd/extensions/azure.ai.models/CHANGELOG.md

@@ -1,6 +1,14 @@
 # Release History
 
 
+## 0.0.7-preview (Unreleased)
+
+### Features
+
+- Added LoRA adapter support to `create` command with `--lora-rank`, `--lora-alpha`, `--lora-target-modules`, and `--lora-dropout` flags for registering LoRA adapters (`--weight-type LoRA`)
+- `show` command now displays LoRA Configuration section (rank, alpha, target modules, dropout) for LoRA adapters
+- `list` command now shows Weight Type column to distinguish FullWeight and LoRA models
+
 ## 0.0.6-preview (Unreleased)
 
 ### Features

cli/azd/extensions/azure.ai.models/extension.yaml

@@ -3,7 +3,7 @@ namespace: ai.models
 displayName: Foundry Custom Models (Preview)
 description: Extension for managing custom models in Azure AI Foundry. (Preview)
 usage: azd ai models <command> [options]
-version: 0.0.6-preview
+version: 0.0.7-preview
 language: go
 capabilities:
     - custom-commands

cli/azd/extensions/azure.ai.models/internal/cmd/custom_create.go

@@ -23,16 +23,20 @@ import (
 )
 
 type customCreateFlags struct {
-	Name        string
-	Version     string
-	Source      string
-	SourceFile  string
-	Description string
-	BaseModel   string
-	WeightType  string
-	Publisher   string
-	AzcopyPath  string
-	NoWait      bool
+	Name              string
+	Version           string
+	Source            string
+	SourceFile        string
+	Description       string
+	BaseModel         string
+	WeightType        string
+	Publisher         string
+	AzcopyPath        string
+	NoWait            bool
+	LoRARank          int
+	LoRAAlpha         int
+	LoRATargetModules string
+	LoRADropout       float64
 }
 
 func newCustomCreateCommand(parentFlags *customFlags) *cobra.Command {
@@ -77,7 +81,7 @@ provide a file containing the URL instead.`,
 				return fmt.Errorf("either --source or --source-file is required")
 			}
 
-			return runCustomCreate(ctx, parentFlags, flags)
+			return runCustomCreate(ctx, cmd, parentFlags, flags)
 		},
 	}
 
@@ -91,14 +95,19 @@ provide a file containing the URL instead.`,
 	cmd.Flags().StringVar(&flags.Publisher, "publisher", "", "Model publisher ID for catalog info (e.g., Fireworks)")
 	cmd.Flags().StringVar(&flags.AzcopyPath, "azcopy-path", "", "Path to azcopy binary (auto-detected if not provided)")
 	cmd.Flags().BoolVar(&flags.NoWait, "no-wait", false, "Start async registration and return immediately with the operation URL")
+	cmd.Flags().IntVar(&flags.LoRARank, "lora-rank", 0, "LoRA rank (r) — required when --weight-type is LoRA")
+	cmd.Flags().IntVar(&flags.LoRAAlpha, "lora-alpha", 0, "LoRA scaling factor (alpha) — required when --weight-type is LoRA")
+	cmd.Flags().StringVar(&flags.LoRATargetModules, "lora-target-modules", "",
+		"Comma-separated list of target modules (e.g., \"q_proj,v_proj,k_proj,o_proj\")")
+	cmd.Flags().Float64Var(&flags.LoRADropout, "lora-dropout", 0, "LoRA dropout rate used during training (informational)")
 
 	_ = cmd.MarkFlagRequired("name")
 	_ = cmd.MarkFlagRequired("base-model")
 
 	return cmd
 }
 
-func runCustomCreate(ctx context.Context, parentFlags *customFlags, flags *customCreateFlags) error {
+func runCustomCreate(ctx context.Context, cmd *cobra.Command, parentFlags *customFlags, flags *customCreateFlags) error {
 	azdClient, err := azdext.NewAzdClient()
 	if err != nil {
 		return fmt.Errorf("failed to create azd client: %w", err)
@@ -132,6 +141,19 @@ func runCustomCreate(ctx context.Context, parentFlags *customFlags, flags *custo
 		return err
 	}
 
+	// Validate LoRA flags early, before any uploads
+	var loraConfig *models.LoRAConfig
+	if strings.EqualFold(flags.WeightType, "LoRA") {
+		var err error
+		loraConfig, err = buildLoRAConfig(cmd, flags)
+		if err != nil {
+			return err
+		}
+	} else if cmd.Flags().Changed("lora-rank") || cmd.Flags().Changed("lora-alpha") ||
+		cmd.Flags().Changed("lora-target-modules") || cmd.Flags().Changed("lora-dropout") {
+		return fmt.Errorf("--lora-* flags are only valid when --weight-type is LoRA")
+	}
+
 	// ── Step 1: Start pending upload ──
 	fmt.Printf("Creating custom model: %s (version %s)\n\n", flags.Name, flags.Version)
 
@@ -240,6 +262,11 @@ func runCustomCreate(ctx context.Context, parentFlags *customFlags, flags *custo
 		}
 	}
 
+	// Attach pre-validated LoRA config
+	if loraConfig != nil {
+		regReq.LoRAConfig = loraConfig
+	}
+
 	operationURL, err := foundryClient.RegisterModelAsync(ctx, flags.Name, flags.Version, regReq)
 	_ = regSpinner.Stop(ctx)
 	fmt.Println()
@@ -335,3 +362,35 @@ func extractVersionFromURI(uri string) string {
 	}
 	return ""
 }
+
+// buildLoRAConfig validates LoRA flags and builds a LoRAConfig for the registration request.
+func buildLoRAConfig(cmd *cobra.Command, flags *customCreateFlags) (*models.LoRAConfig, error) {
+	if flags.LoRARank <= 0 {
+		return nil, fmt.Errorf("--lora-rank is required and must be a positive integer when --weight-type is LoRA")
+	}
+	if flags.LoRAAlpha <= 0 {
+		return nil, fmt.Errorf("--lora-alpha is required and must be a positive integer when --weight-type is LoRA")
+	}
+
+	config := &models.LoRAConfig{
+		Rank:  new(flags.LoRARank),
+		Alpha: new(flags.LoRAAlpha),
+	}
+
+	if flags.LoRATargetModules != "" {
+		modules := strings.Split(flags.LoRATargetModules, ",")
+		for i := range modules {
+			modules[i] = strings.TrimSpace(modules[i])
+			if modules[i] == "" {
+				return nil, fmt.Errorf("--lora-target-modules contains an empty entry")
+			}
+		}
+		config.TargetModules = modules
+	}
+
+	if cmd.Flags().Changed("lora-dropout") {
+		config.Dropout = new(flags.LoRADropout)
+	}
+
+	return config, nil
+}

cli/azd/extensions/azure.ai.models/internal/cmd/custom_create_test.go

@@ -4,7 +4,12 @@
 package cmd
 
 import (
+	"encoding/json"
 	"testing"
+
+	"azure.ai.models/pkg/models"
+
+	"github.com/spf13/cobra"
 )
 
 func TestBuildDerivedModelURI(t *testing.T) {
@@ -100,3 +105,171 @@ func TestExtractVersionFromURI(t *testing.T) {
 		})
 	}
 }
+
+// newTestCmd creates a cobra command with lora flags registered for testing buildLoRAConfig.
+func newTestCmd(flags *customCreateFlags) *cobra.Command {
+	cmd := &cobra.Command{Use: "test"}
+	cmd.Flags().IntVar(&flags.LoRARank, "lora-rank", 0, "")
+	cmd.Flags().IntVar(&flags.LoRAAlpha, "lora-alpha", 0, "")
+	cmd.Flags().StringVar(&flags.LoRATargetModules, "lora-target-modules", "", "")
+	cmd.Flags().Float64Var(&flags.LoRADropout, "lora-dropout", 0, "")
+	return cmd
+}
+
+func TestBuildLoRAConfig(t *testing.T) {
+	tests := []struct {
+		name        string
+		args        []string
+		wantErr     string
+		wantRank    int
+		wantAlpha   int
+		wantModules []string
+		wantDropout *float64
+	}{
+		{
+			name:    "missing rank",
+			args:    []string{"--lora-alpha", "32"},
+			wantErr: "--lora-rank is required",
+		},
+		{
+			name:    "missing alpha",
+			args:    []string{"--lora-rank", "16"},
+			wantErr: "--lora-alpha is required",
+		},
+		{
+			name:      "rank and alpha only",
+			args:      []string{"--lora-rank", "16", "--lora-alpha", "32"},
+			wantRank:  16,
+			wantAlpha: 32,
+		},
+		{
+			name:        "all fields",
+			args:        []string{"--lora-rank", "8", "--lora-alpha", "16", "--lora-target-modules", "q_proj,v_proj", "--lora-dropout", "0.05"},
+			wantRank:    8,
+			wantAlpha:   16,
+			wantModules: []string{"q_proj", "v_proj"},
+			wantDropout: new(0.05),
+		},
+		{
+			name:    "empty entry in target modules",
+			args:    []string{"--lora-rank", "16", "--lora-alpha", "32", "--lora-target-modules", "q_proj,,v_proj"},
+			wantErr: "empty entry",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			flags := &customCreateFlags{}
+			cmd := newTestCmd(flags)
+			if err := cmd.ParseFlags(tt.args); err != nil {
+				t.Fatalf("ParseFlags: %v", err)
+			}
+
+			got, err := buildLoRAConfig(cmd, flags)
+			if tt.wantErr != "" {
+				if err == nil {
+					t.Fatalf("expected error containing %q, got nil", tt.wantErr)
+				}
+				if !contains(err.Error(), tt.wantErr) {
+					t.Fatalf("expected error containing %q, got %q", tt.wantErr, err.Error())
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if got.Rank == nil || *got.Rank != tt.wantRank {
+				t.Errorf("Rank = %v, want %d", got.Rank, tt.wantRank)
+			}
+			if got.Alpha == nil || *got.Alpha != tt.wantAlpha {
+				t.Errorf("Alpha = %v, want %d", got.Alpha, tt.wantAlpha)
+			}
+			if len(tt.wantModules) > 0 {
+				if len(got.TargetModules) != len(tt.wantModules) {
+					t.Errorf("TargetModules = %v, want %v", got.TargetModules, tt.wantModules)
+				}
+				for i := range tt.wantModules {
+					if i < len(got.TargetModules) && got.TargetModules[i] != tt.wantModules[i] {
+						t.Errorf("TargetModules[%d] = %q, want %q", i, got.TargetModules[i], tt.wantModules[i])
+					}
+				}
+			}
+			if tt.wantDropout != nil {
+				if got.Dropout == nil || *got.Dropout != *tt.wantDropout {
+					t.Errorf("Dropout = %v, want %v", got.Dropout, *tt.wantDropout)
+				}
+			}
+		})
+	}
+}
+
+func contains(s, substr string) bool {
+	return len(s) >= len(substr) && searchString(s, substr)
+}
+
+func searchString(s, substr string) bool {
+	for i := 0; i <= len(s)-len(substr); i++ {
+		if s[i:i+len(substr)] == substr {
+			return true
+		}
+	}
+	return false
+}
+
+func TestLoRAConfigJSONRoundTrip(t *testing.T) {
+	original := &models.CustomModel{
+		Name:       "test-lora-adapter",
+		Version:    "1",
+		WeightType: "LoRA",
+		LoRAConfig: &models.LoRAConfig{
+			Rank:          new(16),
+			Alpha:         new(32),
+			TargetModules: []string{"q_proj", "v_proj", "k_proj", "o_proj"},
+			Dropout:       new(0.05),
+		},
+	}
+
+	data, err := json.Marshal(original)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+
+	var decoded models.CustomModel
+	if err := json.Unmarshal(data, &decoded); err != nil {
+		t.Fatalf("Unmarshal: %v", err)
+	}
+
+	if decoded.LoRAConfig == nil {
+		t.Fatal("LoRAConfig is nil after round-trip")
+	}
+	if decoded.LoRAConfig.Rank == nil || *decoded.LoRAConfig.Rank != 16 {
+		t.Errorf("Rank = %v, want 16", decoded.LoRAConfig.Rank)
+	}
+	if decoded.LoRAConfig.Alpha == nil || *decoded.LoRAConfig.Alpha != 32 {
+		t.Errorf("Alpha = %v, want 32", decoded.LoRAConfig.Alpha)
+	}
+	if len(decoded.LoRAConfig.TargetModules) != 4 {
+		t.Errorf("TargetModules len = %d, want 4", len(decoded.LoRAConfig.TargetModules))
+	}
+	if decoded.LoRAConfig.Dropout == nil || *decoded.LoRAConfig.Dropout != 0.05 {
+		t.Errorf("Dropout = %v, want 0.05", decoded.LoRAConfig.Dropout)
+	}
+}
+
+func TestLoRAConfigJSONOmittedWhenNil(t *testing.T) {
+	model := &models.CustomModel{
+		Name:       "full-weight-model",
+		Version:    "1",
+		WeightType: "FullWeight",
+	}
+
+	data, err := json.Marshal(model)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+
+	jsonStr := string(data)
+	if contains(jsonStr, "loraConfig") {
+		t.Errorf("loraConfig should be omitted for FullWeight models, got: %s", jsonStr)
+	}
+}

cli/azd/extensions/azure.ai.models/internal/cmd/custom_show.go

@@ -165,6 +165,22 @@ func runCustomShow(ctx context.Context, parentFlags *customFlags, flags *customS
 			fmt.Printf("  URI: %s\n", *model.DerivedModelInformation.BaseModel)
 		}
 
+		if model.LoRAConfig != nil {
+			fmt.Println("\nLoRA Configuration:")
+			if model.LoRAConfig.Rank != nil {
+				fmt.Printf("  Rank:           %d\n", *model.LoRAConfig.Rank)
+			}
+			if model.LoRAConfig.Alpha != nil {
+				fmt.Printf("  Alpha:          %d\n", *model.LoRAConfig.Alpha)
+			}
+			if len(model.LoRAConfig.TargetModules) > 0 {
+				fmt.Printf("  Target Modules: %s\n", strings.Join(model.LoRAConfig.TargetModules, ", "))
+			}
+			if model.LoRAConfig.Dropout != nil {
+				fmt.Printf("  Dropout:        %g\n", *model.LoRAConfig.Dropout)
+			}
+		}
+
 		if model.Source != nil {
 			fmt.Println("\nSource:")
 			fmt.Printf("  Type: %s\n", model.Source.SourceType)