diff --git a/cli/azd/extensions/azure.ai.customtraining/extension.yaml b/cli/azd/extensions/azure.ai.customtraining/extension.yaml index 8874c1a9b8a..aaf39c21014 100644 --- a/cli/azd/extensions/azure.ai.customtraining/extension.yaml +++ b/cli/azd/extensions/azure.ai.customtraining/extension.yaml @@ -1,7 +1,7 @@ id: azure.ai.customtraining namespace: ai.training displayName: Azure AI Custom Training (Preview) -description: Extension for Azure AI Foundry custom training jobs. (Preview) +description: Extension for Microsoft Foundry custom training jobs. (Preview) usage: azd ai training [options] version: 0.0.1-preview language: go diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go index 45baf7101b7..ccdac9a8f16 100644 --- a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go +++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job.go @@ -35,6 +35,7 @@ func newJobCommand() *cobra.Command { cmd.AddCommand(newJobDeleteCommand()) cmd.AddCommand(newJobCancelCommand()) cmd.AddCommand(newJobValidateCommand()) + cmd.AddCommand(newJobStreamCommand()) return cmd } diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go new file mode 100644 index 00000000000..34e7ff4e790 --- /dev/null +++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/job_stream.go @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "fmt" + + "azure.ai.customtraining/internal/service" + "azure.ai.customtraining/internal/utils" + "azure.ai.customtraining/pkg/client" + + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/spf13/cobra" +) + +func newJobStreamCommand() *cobra.Command { + var name string + + cmd := &cobra.Command{ + Use: "stream", + Short: "Stream logs from a running training job", + Long: "Stream live log output from a training job. Polls for new log content\nuntil the job reaches a terminal state.\n\nExample:\n azd ai training job stream --name my-job-123", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := azdext.WithAccessToken(cmd.Context()) + + if name == "" { + return fmt.Errorf("--name is required") + } + + azdClient, err := azdext.NewAzdClient() + if err != nil { + return fmt.Errorf("failed to create azd client: %w", err) + } + defer azdClient.Close() + + envValues, err := utils.GetEnvironmentValues(ctx, azdClient) + if err != nil { + return fmt.Errorf("failed to get environment values: %w", err) + } + + accountName := envValues[utils.EnvAzureAccountName] + projectName := envValues[utils.EnvAzureProjectName] + tenantID := envValues[utils.EnvAzureTenantID] + + if accountName == "" || projectName == "" { + return fmt.Errorf("environment not configured. Run 'azd ai training init' first") + } + + credential, err := azidentity.NewAzureDeveloperCLICredential( + &azidentity.AzureDeveloperCLICredentialOptions{ + TenantID: tenantID, + AdditionallyAllowedTenants: []string{"*"}, + }, + ) + if err != nil { + return fmt.Errorf("failed to create azure credential: %w", err) + } + + endpoint := buildProjectEndpoint(accountName, projectName) + apiClient, err := client.NewClient(endpoint, credential) + if err != nil { + return fmt.Errorf("failed to create API client: %w", err) + } + apiClient.SetDebugBody(rootFlags.Debug) + + streamSvc := service.NewStreamService(apiClient) + result, err := streamSvc.StreamJobLogs(ctx, name) + if err != nil { + return err + } + + fmt.Println() + fmt.Println("Execution Summary") + fmt.Printf("RunId: %s\n", result.JobName) + fmt.Printf("Status: %s\n", result.Status) + if result.StudioURL != "" { + fmt.Printf("Web View: %s\n", result.StudioURL) + } + + return nil + }, + } + + cmd.Flags().StringVar(&name, "name", "", "Job name/ID (required)") + + return cmd +} diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go index 25bc52d4e4f..bfc10f8f485 100644 --- a/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go +++ b/cli/azd/extensions/azure.ai.customtraining/internal/cmd/root.go @@ -17,7 +17,7 @@ var rootFlags rootFlagsDefinition func NewRootCommand() *cobra.Command { rootCmd := &cobra.Command{ Use: "training [options]", - Short: "Extension for Azure AI Foundry custom training jobs. (Preview)", + Short: "Extension for Microsoft Foundry custom training jobs. (Preview)", SilenceUsage: true, SilenceErrors: true, CompletionOptions: cobra.CompletionOptions{ diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go new file mode 100644 index 00000000000..69cc232bf23 --- /dev/null +++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service.go @@ -0,0 +1,313 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package service + +import ( + "context" + "fmt" + "math" + "os" + "regexp" + "sort" + "strings" + "time" + + "azure.ai.customtraining/pkg/client" +) + +// StreamService handles polling and streaming log output from a running job. +type StreamService struct { + client *client.Client +} + +// NewStreamService creates a new stream service. +func NewStreamService(apiClient *client.Client) *StreamService { + return &StreamService{client: apiClient} +} + +// Log file patterns matching the Azure ML SDK behavior. +// Primary: Common Runtime user logs. +// Fallback: Legacy azureml-logs for older compute targets. +var ( + commonRuntimeLogPattern = regexp.MustCompile(`user_logs/std_log[\D]*[0]*(?:_ps)?\.txt`) + legacyLogPattern = regexp.MustCompile(`azureml-logs/[\d]{2}.+\.txt`) +) + +// terminalStates are job statuses that indicate the job has finished. +var terminalStates = map[string]bool{ + "Completed": true, + "Failed": true, + "Canceled": true, + "NotResponding": true, + "Paused": true, +} + +// activeStates are job statuses where streaming is applicable. +var activeStates = map[string]bool{ + "NotStarted": true, + "Queued": true, + "Preparing": true, + "Provisioning": true, + "Starting": true, + "Running": true, + "Finalizing": true, +} + +// StreamResult contains the final state of a streamed job. +type StreamResult struct { + JobName string + Status string + StudioURL string +} + +// StreamJobLogs polls the job and streams log output until the job reaches a terminal state. +func (s *StreamService) StreamJobLogs(ctx context.Context, jobName string) (*StreamResult, error) { + fmt.Fprintf(os.Stderr, "Streaming logs for job: %s\n\n", jobName) + + // Line-count tracking per file, matching the Azure ML SDK approach. + // Each poll downloads full content, skips already-printed lines, prints the rest. + processedLines := make(map[string]int) + + const maxConsecutiveErrs = 3 + + startTime := time.Now() + consecutiveErrs := 0 + trackingEndpoint := "" + firstPoll := true + multiFile := false // latches to true once >1 file is seen + + for { + if ctx.Err() != nil { + return nil, ctx.Err() + } + + // Check job status every poll + job, err := s.client.GetJob(ctx, jobName) + if err != nil { + consecutiveErrs++ + if consecutiveErrs >= maxConsecutiveErrs { + return nil, fmt.Errorf("failed to get job status after %d retries: %w", maxConsecutiveErrs, err) + } + if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil { + return nil, err + } + continue + } + consecutiveErrs = 0 + + jobStatus := job.Properties.Status + studioURL := extractServiceEndpoint(job.Properties.Services, "Studio") + + if trackingEndpoint == "" { + trackingEndpoint = extractServiceEndpoint(job.Properties.Services, "Tracking") + } + + if terminalStates[jobStatus] { + if trackingEndpoint != "" && !firstPoll { + s.flushLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile) + } + return &StreamResult{ + JobName: jobName, + Status: jobStatus, + StudioURL: studioURL, + }, nil + } + + if !activeStates[jobStatus] { + fmt.Fprintf(os.Stderr, "Job status: %s, waiting...\n", jobStatus) + if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil { + return nil, err + } + firstPoll = false + continue + } + + // Stream logs if tracking endpoint is available + if trackingEndpoint != "" { + var newMultiFile bool + _, newMultiFile, err = s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile) + if newMultiFile { + multiFile = true + } + if err != nil { + consecutiveErrs++ + if consecutiveErrs >= maxConsecutiveErrs { + return nil, fmt.Errorf("failed to stream logs after %d retries: %w", maxConsecutiveErrs, err) + } + } else { + consecutiveErrs = 0 + } + } else { + fmt.Fprintf(os.Stderr, "Waiting for job to initialize...\n") + } + + firstPoll = false + if err := sleepWithContext(ctx, pollInterval(startTime)); err != nil { + return nil, err + } + } +} + +// sleepWithContext sleeps for the given duration, returning early if ctx is cancelled. +func sleepWithContext(ctx context.Context, d time.Duration) error { + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(d): + return nil + } +} + +// pollInterval returns the polling interval using a sigmoid curve from 2s → 60s +// based on elapsed time since streaming started. Matches the Azure ML SDK approach. +// +// duration = MAX / (1.0 + 100 * exp(-elapsed_seconds / 20.0)) +// return max(MIN, duration) +func pollInterval(startTime time.Time) time.Duration { + const ( + minInterval = 2 * time.Second + maxInterval = 60 * time.Second + ) + elapsed := time.Since(startTime).Seconds() + durationSec := 60.0 / (1.0 + 100.0*math.Exp(-elapsed/20.0)) + duration := time.Duration(durationSec * float64(time.Second)) + if duration < minInterval { + return minInterval + } + return duration +} + +// filterLogFiles selects streamable log files from the history details. +// Matches Common Runtime user logs first; falls back to legacy azureml-logs. +// Returns matched file names sorted alphabetically. +func filterLogFiles(logFiles map[string]string) []string { + var matched []string + for name := range logFiles { + if commonRuntimeLogPattern.MatchString(name) { + matched = append(matched, name) + } + } + if len(matched) == 0 { + // Fallback to legacy log pattern for older compute targets + for name := range logFiles { + if legacyLogPattern.MatchString(name) { + matched = append(matched, name) + } + } + } + sort.Strings(matched) + return matched +} + +// printFileHeader prints a visual separator for a log file, matching the Azure ML SDK style. +func printFileHeader(fileName string) { + fmt.Println() + fmt.Printf("Streaming %s\n", fileName) + fmt.Println(strings.Repeat("=", len("Streaming ")+len(fileName))) + fmt.Println() +} + +// pollAndPrintLogs fetches run history details and prints only new log lines. +// Returns whether new content was printed and whether multiple files were seen. +func (s *StreamService) pollAndPrintLogs( + ctx context.Context, + trackingEndpoint string, + jobName string, + processedLines map[string]int, + multiFile bool, +) (bool, bool, error) { + details, err := s.client.GetRunHistoryDetails(ctx, trackingEndpoint, jobName) + if err != nil { + return false, multiFile, err + } + if details == nil || len(details.LogFiles) == 0 { + return false, multiFile, nil + } + + fileNames := filterLogFiles(details.LogFiles) + if len(fileNames) == 0 { + return false, multiFile, nil + } + + // Latch multiFile once we see >1 file — formatting stays consistent thereafter + if len(fileNames) > 1 { + multiFile = true + } + + hasNewContent := false + for _, fileName := range fileNames { + sasURI := details.LogFiles[fileName] + + content, _, err := s.client.GetBlobContent(ctx, sasURI) + if err != nil { + fmt.Fprintf(os.Stderr, "Warning: failed to read log file %s: %v\n", fileName, err) + continue + } + if content == "" { + continue + } + + lines := strings.Split(content, "\n") + // Remove trailing empty element from final newline + if len(lines) > 0 && lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] + } + + previousLines := processedLines[fileName] + if len(lines) <= previousLines { + continue + } + + // Print file header: always on first encounter, or on every switch for multi-file jobs + if previousLines == 0 || multiFile { + printFileHeader(fileName) + } + + for _, line := range lines[previousLines:] { + fmt.Println(line) + } + hasNewContent = true + processedLines[fileName] = len(lines) + } + + return hasNewContent, multiFile, nil +} + +// flushLogs does a final poll to capture any remaining log output. +func (s *StreamService) flushLogs( + ctx context.Context, + trackingEndpoint string, + jobName string, + processedLines map[string]int, + multiFile bool, +) { + _, _, err := s.pollAndPrintLogs(ctx, trackingEndpoint, jobName, processedLines, multiFile) + if err != nil { + fmt.Fprintf(os.Stderr, "Warning: failed to flush final logs: %v\n", err) + } +} + +// extractServiceEndpoint extracts the endpoint URL from the job services map. +func extractServiceEndpoint(services map[string]interface{}, serviceName string) string { + if services == nil { + return "" + } + svc, ok := services[serviceName] + if !ok { + return "" + } + svcMap, ok := svc.(map[string]interface{}) + if !ok { + return "" + } + endpoint, ok := svcMap["endpoint"] + if !ok { + return "" + } + str, ok := endpoint.(string) + if !ok { + return "" + } + return str +} diff --git a/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go new file mode 100644 index 00000000000..cfdea80e757 --- /dev/null +++ b/cli/azd/extensions/azure.ai.customtraining/internal/service/stream_service_test.go @@ -0,0 +1,318 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package service + +import ( + "context" + "testing" + "time" +) + +func TestPollInterval_SigmoidBoundaries(t *testing.T) { + tests := []struct { + name string + elapsedSec float64 + wantMinSec float64 + wantMaxSec float64 + description string + }{ + { + name: "start returns minimum", + elapsedSec: 0, + wantMinSec: 2.0, + wantMaxSec: 2.0, + description: "at t=0 the sigmoid value is ~0.59s, clamped to 2s minimum", + }, + { + name: "10 seconds still near minimum", + elapsedSec: 10, + wantMinSec: 2.0, + wantMaxSec: 3.0, + description: "early in the curve, interval stays low", + }, + { + name: "60 seconds mid range", + elapsedSec: 60, + wantMinSec: 5.0, + wantMaxSec: 15.0, + description: "at 60s the sigmoid is in the transition zone", + }, + { + name: "120 seconds approaching max", + elapsedSec: 120, + wantMinSec: 45.0, + wantMaxSec: 60.0, + description: "at 120s the sigmoid approaches the 60s asymptote", + }, + { + name: "300 seconds at max", + elapsedSec: 300, + wantMinSec: 59.0, + wantMaxSec: 60.0, + description: "well past the inflection, saturated at max", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + startTime := time.Now().Add(-time.Duration(tt.elapsedSec * float64(time.Second))) + got := pollInterval(startTime) + gotSec := got.Seconds() + if gotSec < tt.wantMinSec || gotSec > tt.wantMaxSec { + t.Errorf("pollInterval(elapsed=%vs) = %vs, want [%v, %v]s", + tt.elapsedSec, gotSec, tt.wantMinSec, tt.wantMaxSec) + } + }) + } +} + +func TestPollInterval_NeverBelowMinimum(t *testing.T) { + for elapsed := 0; elapsed <= 300; elapsed += 5 { + startTime := time.Now().Add(-time.Duration(elapsed) * time.Second) + got := pollInterval(startTime) + if got < 2*time.Second { + t.Errorf("pollInterval(elapsed=%ds) = %v, below 2s minimum", elapsed, got) + } + if got > 60*time.Second { + t.Errorf("pollInterval(elapsed=%ds) = %v, above 60s maximum", elapsed, got) + } + } +} + +func TestSleepWithContext_CancelledImmediately(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel before sleeping + + start := time.Now() + err := sleepWithContext(ctx, 10*time.Second) + elapsed := time.Since(start) + + if err == nil { + t.Error("expected error from cancelled context, got nil") + } + if elapsed > 500*time.Millisecond { + t.Errorf("sleepWithContext took %v, should have returned immediately", elapsed) + } +} + +func TestSleepWithContext_CompletesNormally(t *testing.T) { + ctx := context.Background() + start := time.Now() + err := sleepWithContext(ctx, 50*time.Millisecond) + elapsed := time.Since(start) + + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if elapsed < 50*time.Millisecond { + t.Errorf("sleepWithContext returned too early: %v", elapsed) + } +} + +func TestFilterLogFiles_PrimaryPattern(t *testing.T) { + tests := []struct { + name string + logFiles map[string]string + want []string + }{ + { + name: "empty map returns nil", + logFiles: map[string]string{}, + want: nil, + }, + { + name: "matches single user_logs std_log", + logFiles: map[string]string{ + "user_logs/std_log.txt": "https://blob.core/std_log.txt?sas", + }, + want: []string{"user_logs/std_log.txt"}, + }, + { + name: "matches std_log_ps variant", + logFiles: map[string]string{ + "user_logs/std_log_ps.txt": "https://blob.core/std_log_ps.txt?sas", + }, + want: []string{"user_logs/std_log_ps.txt"}, + }, + { + name: "excludes activity logs and system files", + logFiles: map[string]string{ + "user_logs/std_log.txt": "https://blob.core/std_log.txt?sas", + "system_logs/virtualcluster_activity_log.txt": "https://blob.core/activity.txt?sas", + "system_logs/cs_capability/cs-capability-0.txt": "https://blob.core/cs.txt?sas", + }, + want: []string{"user_logs/std_log.txt"}, + }, + { + name: "multiple user_logs sorted alphabetically", + logFiles: map[string]string{ + "user_logs/std_log.txt": "https://blob.core/std_log.txt?sas", + "user_logs/std_log_ps.txt": "https://blob.core/std_log_ps.txt?sas", + }, + want: []string{"user_logs/std_log.txt", "user_logs/std_log_ps.txt"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := filterLogFiles(tt.logFiles) + if !slicesEqual(got, tt.want) { + t.Errorf("filterLogFiles() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestFilterLogFiles_FallbackToLegacy(t *testing.T) { + tests := []struct { + name string + logFiles map[string]string + want []string + }{ + { + name: "falls back to azureml-logs when no user_logs match", + logFiles: map[string]string{ + "azureml-logs/70_driver_log.txt": "https://blob.core/70.txt?sas", + "azureml-logs/55_azureml-log.txt": "https://blob.core/55.txt?sas", + }, + want: []string{"azureml-logs/55_azureml-log.txt", "azureml-logs/70_driver_log.txt"}, + }, + { + name: "prefers user_logs over azureml-logs when both exist", + logFiles: map[string]string{ + "user_logs/std_log.txt": "https://blob.core/std_log.txt?sas", + "azureml-logs/70_driver_log.txt": "https://blob.core/70.txt?sas", + }, + want: []string{"user_logs/std_log.txt"}, + }, + { + name: "no matching files returns nil", + logFiles: map[string]string{ + "system_logs/something.txt": "https://blob.core/sys.txt?sas", + "random_file.log": "https://blob.core/random.txt?sas", + }, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := filterLogFiles(tt.logFiles) + if !slicesEqual(got, tt.want) { + t.Errorf("filterLogFiles() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestExtractServiceEndpoint(t *testing.T) { + tests := []struct { + name string + services map[string]any + serviceName string + want string + }{ + { + name: "nil services", + services: nil, + serviceName: "Studio", + want: "", + }, + { + name: "missing service", + services: map[string]any{"Tracking": map[string]any{"endpoint": "https://tracking"}}, + serviceName: "Studio", + want: "", + }, + { + name: "service not a map", + services: map[string]any{"Studio": "not-a-map"}, + serviceName: "Studio", + want: "", + }, + { + name: "missing endpoint key", + services: map[string]any{"Studio": map[string]any{"other": "value"}}, + serviceName: "Studio", + want: "", + }, + { + name: "endpoint not a string", + services: map[string]any{"Studio": map[string]any{"endpoint": 42}}, + serviceName: "Studio", + want: "", + }, + { + name: "valid Studio endpoint", + services: map[string]any{ + "Studio": map[string]any{ + "endpoint": "https://ml.azure.com/runs/my-job?wsid=/sub/rg/ws", + }, + }, + serviceName: "Studio", + want: "https://ml.azure.com/runs/my-job?wsid=/sub/rg/ws", + }, + { + name: "valid Tracking endpoint", + services: map[string]any{ + "Tracking": map[string]any{ + "endpoint": "azureml://eastus.api.azureml.ms/mlflow/v1.0/sub/rg/ws", + }, + }, + serviceName: "Tracking", + want: "azureml://eastus.api.azureml.ms/mlflow/v1.0/sub/rg/ws", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractServiceEndpoint(tt.services, tt.serviceName) + if got != tt.want { + t.Errorf("extractServiceEndpoint(%q) = %q, want %q", tt.serviceName, got, tt.want) + } + }) + } +} + +func TestTerminalAndActiveStates(t *testing.T) { + // Verify terminal states don't overlap with active states + for state := range terminalStates { + if activeStates[state] { + t.Errorf("state %q is both terminal and active", state) + } + } + + // Verify expected terminal states + expectedTerminal := []string{"Completed", "Failed", "Canceled", "NotResponding", "Paused"} + for _, s := range expectedTerminal { + if !terminalStates[s] { + t.Errorf("expected %q to be a terminal state", s) + } + } + + // Verify expected active states + expectedActive := []string{"NotStarted", "Queued", "Preparing", "Provisioning", "Starting", "Running", "Finalizing"} + for _, s := range expectedActive { + if !activeStates[s] { + t.Errorf("expected %q to be an active state", s) + } + } +} + +// slicesEqual compares two string slices including nil vs empty. +func slicesEqual(a, b []string) bool { + if len(a) == 0 && len(b) == 0 { + // Treat nil and empty as equal + return (a == nil) == (b == nil) || (len(a) == 0 && len(b) == 0) + } + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go new file mode 100644 index 00000000000..568d9d94b43 --- /dev/null +++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/blob.go @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package client + +import ( + "context" + "fmt" + "io" + "net/http" +) + +// GetBlobContent downloads content from a SAS URI. +// Fetches the full content on each call. No authentication is needed since the +// URL contains a SAS token. The caller is responsible for tracking incremental +// progress (e.g., line counts for log streaming). +func (c *Client) GetBlobContent(ctx context.Context, sasURI string) (string, int64, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, sasURI, nil) + if err != nil { + return "", 0, fmt.Errorf("failed to create blob request: %w", err) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return "", 0, fmt.Errorf("blob request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", 0, fmt.Errorf("blob request returned status %d", resp.StatusCode) + } + + // Cap read to 1MB per call to avoid memory issues with very large blobs + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return "", 0, fmt.Errorf("failed to read blob content: %w", err) + } + + return string(body), int64(len(body)), nil +} diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go index 8784444cdae..e22b513e906 100644 --- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go +++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/client.go @@ -88,7 +88,9 @@ func (c *Client) doDataPlaneWithVersion(ctx context.Context, method, path, apiVe reqURL += fmt.Sprintf("&%s=%s", queryParams[i], url.QueryEscape(queryParams[i+1])) } - fmt.Printf("[DEBUG] %s %s\n", method, reqURL) + if c.debugBody { + fmt.Printf("[DEBUG] %s %s\n", method, reqURL) + } var bodyReader io.Reader if body != nil { diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go index 7836f62cf0d..acfe9eb04d4 100644 --- a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go +++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history.go @@ -8,6 +8,8 @@ import ( "encoding/json" "fmt" "net/http" + "net/url" + "strings" "azure.ai.customtraining/pkg/models" ) @@ -36,3 +38,84 @@ func (c *Client) GetRunHistory(ctx context.Context, runID string) (*models.RunHi return &result, nil } + +// GetRunHistoryDetails retrieves detailed run information including log file SAS URIs. +// This calls the AML history service directly using the tracking endpoint from the job response. +// GET https://{region}.api.azureml.ms/history/v1.0/{workspace}/runs/{runId}/details +func (c *Client) GetRunHistoryDetails(ctx context.Context, trackingEndpoint string, runID string) (*models.RunHistoryDetails, error) { + baseURL, workspacePath, err := parseTrackingEndpoint(trackingEndpoint) + if err != nil { + return nil, fmt.Errorf("failed to parse tracking endpoint: %w", err) + } + + reqURL := fmt.Sprintf("%s/history/v1.0%s/runs/%s/details", baseURL, workspacePath, url.PathEscape(runID)) + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + if err := c.addAuth(ctx, req, DataPlaneScope); err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + + if c.debugBody { + fmt.Printf("[DEBUG] GET %s\n", reqURL) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return nil, nil + } + + if resp.StatusCode != http.StatusOK { + return nil, c.HandleError(resp) + } + + var result models.RunHistoryDetails + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, fmt.Errorf("failed to decode run history details response: %w", err) + } + + return &result, nil +} + +// parseTrackingEndpoint extracts the base URL and workspace path from a tracking endpoint. +// Input format: azureml://{region}.api.azureml.ms/mlflow/v1.0/{workspace-path}?... +// Returns: https://{region}.api.azureml.ms, /{workspace-path} +func parseTrackingEndpoint(trackingEndpoint string) (string, string, error) { + // Remove the azureml:// prefix + endpoint := strings.TrimPrefix(trackingEndpoint, "azureml://") + // Remove any trailing query params + if idx := strings.Index(endpoint, "?"); idx != -1 { + endpoint = endpoint[:idx] + } + + // Parse: {host}/mlflow/v1.0/{subscription-path} + parsed, err := url.Parse("https://" + endpoint) + if err != nil { + return "", "", fmt.Errorf("invalid tracking endpoint: %w", err) + } + + host := parsed.Host + if host == "" { + return "", "", fmt.Errorf("invalid tracking endpoint: missing host") + } + + // Extract workspace path: everything after /mlflow/v1.0/ or /mlflow/v2.0/ + path := parsed.Path + for _, prefix := range []string{"/mlflow/v1.0", "/mlflow/v2.0"} { + if strings.HasPrefix(path, prefix) { + workspacePath := strings.TrimPrefix(path, prefix) + return "https://" + host, workspacePath, nil + } + } + + return "", "", fmt.Errorf("invalid tracking endpoint: expected mlflow path prefix in %q", path) +} diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go new file mode 100644 index 00000000000..52bfdaa540f --- /dev/null +++ b/cli/azd/extensions/azure.ai.customtraining/pkg/client/history_test.go @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package client + +import ( + "testing" +) + +func TestParseTrackingEndpoint(t *testing.T) { + tests := []struct { + name string + endpoint string + wantBaseURL string + wantWorkPath string + wantErr bool + }{ + { + name: "standard v1.0 endpoint", + endpoint: "azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws", + wantBaseURL: "https://eastus.api.azureml.ms", + wantWorkPath: "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws", + }, + { + name: "v2.0 endpoint", + endpoint: "azureml://westus2.api.azureml.ms/mlflow/v2.0/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws", + wantBaseURL: "https://westus2.api.azureml.ms", + wantWorkPath: "/subscriptions/sub-id/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws", + }, + { + name: "endpoint with query params stripped", + endpoint: "azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/sub/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws?extra=param", + wantBaseURL: "https://eastus.api.azureml.ms", + wantWorkPath: "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.MachineLearningServices/workspaces/ws", + }, + { + name: "missing mlflow prefix", + endpoint: "azureml://eastus.api.azureml.ms/other/v1.0/subscriptions/sub/resourceGroups/rg", + wantErr: true, + }, + { + name: "empty string", + endpoint: "", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + baseURL, workPath, err := parseTrackingEndpoint(tt.endpoint) + if tt.wantErr { + if err == nil { + t.Errorf("parseTrackingEndpoint(%q) expected error, got baseURL=%q workPath=%q", tt.endpoint, baseURL, workPath) + } + return + } + if err != nil { + t.Fatalf("parseTrackingEndpoint(%q) unexpected error: %v", tt.endpoint, err) + } + if baseURL != tt.wantBaseURL { + t.Errorf("baseURL = %q, want %q", baseURL, tt.wantBaseURL) + } + if workPath != tt.wantWorkPath { + t.Errorf("workPath = %q, want %q", workPath, tt.wantWorkPath) + } + }) + } +} diff --git a/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go b/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go index 5a709bf224e..23a17f31b8a 100644 --- a/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go +++ b/cli/azd/extensions/azure.ai.customtraining/pkg/models/history.go @@ -5,25 +5,25 @@ package models // RunHistory represents the run history details for a job. type RunHistory struct { - RunID string `json:"runId,omitempty"` - RunUUID string `json:"runUuid,omitempty"` - RootRunID string `json:"rootRunId,omitempty"` - Status string `json:"status,omitempty"` - StartTimeUTC string `json:"startTimeUtc,omitempty"` - EndTimeUTC string `json:"endTimeUtc,omitempty"` - Duration string `json:"duration,omitempty"` - ComputeDuration string `json:"computeDuration,omitempty"` - CreatedUTC string `json:"createdUtc,omitempty"` - LastModifiedUTC string `json:"lastModifiedUtc,omitempty"` - DisplayName string `json:"displayName,omitempty"` - Description string `json:"description,omitempty"` - Target string `json:"target,omitempty"` - RunType string `json:"runType,omitempty"` - Error *RunHistoryError `json:"error,omitempty"` - CreatedBy *RunHistoryUser `json:"createdBy,omitempty"` - Compute *RunHistoryCompute `json:"compute,omitempty"` - Properties map[string]string `json:"properties,omitempty"` - Tags map[string]string `json:"tags,omitempty"` + RunID string `json:"runId,omitempty"` + RunUUID string `json:"runUuid,omitempty"` + RootRunID string `json:"rootRunId,omitempty"` + Status string `json:"status,omitempty"` + StartTimeUTC string `json:"startTimeUtc,omitempty"` + EndTimeUTC string `json:"endTimeUtc,omitempty"` + Duration string `json:"duration,omitempty"` + ComputeDuration string `json:"computeDuration,omitempty"` + CreatedUTC string `json:"createdUtc,omitempty"` + LastModifiedUTC string `json:"lastModifiedUtc,omitempty"` + DisplayName string `json:"displayName,omitempty"` + Description string `json:"description,omitempty"` + Target string `json:"target,omitempty"` + RunType string `json:"runType,omitempty"` + Error *RunHistoryError `json:"error,omitempty"` + CreatedBy *RunHistoryUser `json:"createdBy,omitempty"` + Compute *RunHistoryCompute `json:"compute,omitempty"` + Properties map[string]string `json:"properties,omitempty"` + Tags map[string]string `json:"tags,omitempty"` Inputs map[string]RunHistoryAsset `json:"inputs,omitempty"` Outputs map[string]RunHistoryAsset `json:"outputs,omitempty"` } @@ -68,3 +68,10 @@ type RunHistoryList struct { Value []RunHistory `json:"value"` NextLink string `json:"nextLink,omitempty"` } + +// RunHistoryDetails represents the details response from the AML history service. +// GET .../runs/{runId}/details +type RunHistoryDetails struct { + RunID string `json:"runId,omitempty"` + LogFiles map[string]string `json:"logFiles,omitempty"` // filename → SAS URI +}