From f66a860535db21075b51b116448294d9c636d3b5 Mon Sep 17 00:00:00 2001 From: Saanika Gupta Date: Sat, 9 May 2026 10:33:01 +0530 Subject: [PATCH 1/2] Implement show services --- .../azure.ai.training/internal/cmd/job.go | 1 + .../internal/cmd/job_show_services.go | 192 ++++++++++++++++++ .../pkg/client/serviceinstances.go | 60 ++++++ 3 files changed, 253 insertions(+) create mode 100644 cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services.go diff --git a/cli/azd/extensions/azure.ai.training/internal/cmd/job.go b/cli/azd/extensions/azure.ai.training/internal/cmd/job.go index 71721767309..8f100acd39d 100644 --- a/cli/azd/extensions/azure.ai.training/internal/cmd/job.go +++ b/cli/azd/extensions/azure.ai.training/internal/cmd/job.go @@ -39,6 +39,7 @@ func newJobCommand() *cobra.Command { cmd.AddCommand(newJobConnectSSHCommand()) cmd.AddCommand(newJobSSHProxyCommand()) cmd.AddCommand(newJobDownloadCommand()) + cmd.AddCommand(newJobShowServicesCommand()) return cmd } diff --git a/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services.go b/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services.go new file mode 100644 index 00000000000..ad40aeee9d7 --- /dev/null +++ b/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services.go @@ -0,0 +1,192 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "bytes" + "encoding/json" + "fmt" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/spf13/cobra" +) + +func newJobShowServicesCommand() *cobra.Command { + var name string + var nodeIndex int + + cmd := &cobra.Command{ + Use: "show-services", + Short: "Show services of a training job per node (e.g. SSH, JupyterLab, TensorBoard)", + Long: "Show the services running on a specific node of a training job. Output is JSON.\n\n" + + "Example:\n" + + " azd ai training job show-services --name my-job --node-index 0", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + ctx := azdext.WithAccessToken(cmd.Context()) + + if name == "" { + return fmt.Errorf("--name is required") + } + if nodeIndex < 0 { + return fmt.Errorf("--node-index must be >= 0") + } + + apiClient, err := buildJobAPIClient(ctx) + if err != nil { + return err + } + + // Get job → tracking endpoint (reuses existing GetJob and helper) + job, err := apiClient.GetJob(ctx, name) + if err != nil { + return fmt.Errorf("failed to get job %q: %w", name, err) + } + + trackingEndpoint := extractServiceEndpointStr(job.Properties.Services, "Tracking") + if trackingEndpoint == "" { + return fmt.Errorf("job %q has no tracking endpoint yet; ensure the job has started", name) + } + + raw, err := apiClient.GetServiceInstanceRaw(ctx, trackingEndpoint, name, nodeIndex) + if err != nil { + return fmt.Errorf("failed to get services for node %d of job %q: %w", nodeIndex, name, err) + } + + out, empty, err := transformServiceInstanceResponse(raw) + if err != nil { + return fmt.Errorf("failed to parse services response: %w", err) + } + if empty { + return fmt.Errorf("no services found for node %d of job %q", nodeIndex, name) + } + + // Pretty-print the transformed payload (parity with AML CLI shape). + var pretty bytes.Buffer + if err := json.Indent(&pretty, out, "", " "); err != nil { + fmt.Println(string(out)) + return nil + } + fmt.Println(pretty.String()) + return nil + }, + } + + cmd.Flags().StringVar(&name, "name", "", "Job name (required)") + cmd.Flags().IntVar(&nodeIndex, "node-index", 0, + "Zero-based index of the node in a multi-node job (default 0)") + + return cmd +} + +// outputService is the per-service shape we emit. Field order here is the +// JSON key order in the output. All fields are RawMessage so we can default +// missing ones to literal `null` rather than relying on omitempty. +type outputService struct { + Type json.RawMessage `json:"type"` + Port json.RawMessage `json:"port"` + Status json.RawMessage `json:"status"` + Error json.RawMessage `json:"error"` + Endpoint json.RawMessage `json:"endpoint"` + Properties json.RawMessage `json:"properties"` +} + +// transformServiceInstanceResponse reshapes the AML history serviceinstances +// response to match the AML CLI output: +// - drop the top-level `instances` envelope +// - flatten `error` to just the inner message string (else null) +// - guarantee all 6 fields are present per service (null for unset) +// - sort service-name keys alphabetically (Go marshals map keys sorted) +// +// Returns (out, empty, err). `empty` is true when the response had no services. +func transformServiceInstanceResponse(raw json.RawMessage) ([]byte, bool, error) { + if len(raw) == 0 { + return nil, true, nil + } + + var envelope struct { + Instances map[string]json.RawMessage `json:"instances"` + } + if err := json.Unmarshal(raw, &envelope); err != nil { + return nil, false, err + } + if len(envelope.Instances) == 0 { + return nil, true, nil + } + + null := json.RawMessage("null") + transformed := make(map[string]outputService, len(envelope.Instances)) + for svcName, svcRaw := range envelope.Instances { + var fields map[string]json.RawMessage + if err := json.Unmarshal(svcRaw, &fields); err != nil { + return nil, false, fmt.Errorf("service %q has unexpected shape: %w", svcName, err) + } + pick := func(key string) json.RawMessage { + if v, ok := fields[key]; ok && len(v) > 0 { + return v + } + return null + } + transformed[svcName] = outputService{ + Type: pick("type"), + Port: pick("port"), + Status: pick("status"), + Error: flattenServiceError(fields["error"]), + Endpoint: pick("endpoint"), + Properties: pick("properties"), + } + } + + out, err := marshalNoEscape(transformed) + if err != nil { + return nil, false, err + } + return out, false, nil +} + +// marshalNoEscape is like json.Marshal but does not HTML-escape <, >, & in +// strings, so endpoint URLs containing literal angle brackets render cleanly. +func marshalNoEscape(v any) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if err := enc.Encode(v); err != nil { + return nil, err + } + // Encoder appends a trailing newline; trim it so json.Indent output is clean. + return bytes.TrimRight(buf.Bytes(), "\n"), nil +} + +// flattenServiceError extracts the inner error message from the AML error +// envelope. Returns a JSON string literal of the message, or `null` if no +// usable message is present. +// +// Input shape (when set): +// +// { "error": { "message": "...", ... }, "time": "...", ... } +// +// Output: "..." (or null) +func flattenServiceError(raw json.RawMessage) json.RawMessage { + null := json.RawMessage("null") + if len(raw) == 0 || string(bytes.TrimSpace(raw)) == "null" { + return null + } + + var envelope struct { + Error *struct { + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(raw, &envelope); err != nil { + return null + } + if envelope.Error == nil || envelope.Error.Message == "" { + return null + } + encoded, err := marshalNoEscape(envelope.Error.Message) + if err != nil { + return null + } + return encoded +} diff --git a/cli/azd/extensions/azure.ai.training/pkg/client/serviceinstances.go b/cli/azd/extensions/azure.ai.training/pkg/client/serviceinstances.go index 5e494a97793..4f590c77fcc 100644 --- a/cli/azd/extensions/azure.ai.training/pkg/client/serviceinstances.go +++ b/cli/azd/extensions/azure.ai.training/pkg/client/serviceinstances.go @@ -7,6 +7,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "net/url" "strconv" @@ -75,6 +76,65 @@ func (c *Client) GetServiceInstance( return &result, nil } +// GetServiceInstanceRaw is like GetServiceInstance but returns the raw JSON +// response body so the caller can pass it through to output without losing +// any fields (e.g. nullable values, fields not modeled in Go structs). +// Returns nil with no error when the node does not exist (404). +func (c *Client) GetServiceInstanceRaw( + ctx context.Context, + trackingEndpoint string, + runID string, + nodeIndex int, +) (json.RawMessage, 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/serviceinstances/%s", + baseURL, + workspacePath, + url.PathEscape(runID), + strconv.Itoa(nodeIndex), + ) + + 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) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read service instance response: %w", err) + } + + return json.RawMessage(body), nil +} + // GetARMToken returns a bearer token scoped for ARM (management.azure.com). // Used for the WebSocket tunnel auth header. func (c *Client) GetARMToken(ctx context.Context) (string, error) { From 3ceb6b500a7d994fe4b234a8f43a83d9b06d73f4 Mon Sep 17 00:00:00 2001 From: Saanika Gupta Date: Sat, 9 May 2026 11:18:48 +0530 Subject: [PATCH 2/2] Add UTs --- .../internal/cmd/job_show_services_test.go | 217 ++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services_test.go diff --git a/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services_test.go b/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services_test.go new file mode 100644 index 00000000000..368a13579d4 --- /dev/null +++ b/cli/azd/extensions/azure.ai.training/internal/cmd/job_show_services_test.go @@ -0,0 +1,217 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestTransformServiceInstanceResponse_EmptyCases(t *testing.T) { + tests := []struct { + name string + raw string + }{ + {name: "nil/empty raw", raw: ""}, + {name: "missing instances key", raw: `{}`}, + {name: "null instances", raw: `{"instances":null}`}, + {name: "empty instances map", raw: `{"instances":{}}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + out, empty, err := transformServiceInstanceResponse(json.RawMessage(tt.raw)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !empty { + t.Fatalf("expected empty=true, got out=%s", string(out)) + } + }) + } +} + +func TestTransformServiceInstanceResponse_InvalidJSON(t *testing.T) { + _, _, err := transformServiceInstanceResponse(json.RawMessage(`{not json`)) + if err == nil { + t.Fatal("expected error for invalid JSON") + } +} + +// TestTransformServiceInstanceResponse_FullSample exercises the AML-parity +// transformations against a representative response: envelope strip, error +// flatten (string + null), all 6 fields always present (null for unset), +// snake_case keys, alphabetical service ordering, and no HTML escaping +// of <, >, & in endpoint URLs. +func TestTransformServiceInstanceResponse_FullSample(t *testing.T) { + input := `{ + "instances": { + "my_ssh": { + "type": "SSH", + "port": 8705, + "status": "Running", + "error": null, + "endpoint": "", + "properties": {"ProxyEndpoint": "wss://ssh-host"} + }, + "tensorboard": { + "type": "TensorBoard", + "port": 6006, + "status": "Failed", + "error": { + "error": { + "code": null, + "message": "failed to start endpoint tensorboard" + }, + "time": "0001-01-01T00:00:00+00:00" + }, + "endpoint": "https://tnsrb-host", + "properties": {} + }, + "vscode": { + "type": "VSCode", + "status": "Running", + "endpoint": "vscode://x?a=1&b=2", + "properties": {"ProxyEndpoint": "https://-host"} + }, + "grafana": { + "type": "Grafana", + "port": 3000, + "status": "Running", + "error": null, + "endpoint": "https://3000-host/", + "properties": {} + } + } +}` + + out, empty, err := transformServiceInstanceResponse(json.RawMessage(input)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if empty { + t.Fatal("expected non-empty output") + } + + // 1. No top-level instances envelope. + if strings.Contains(string(out), `"instances"`) { + t.Errorf("output should not contain 'instances' envelope: %s", out) + } + + // 2. Decode and inspect each service. + var got map[string]map[string]json.RawMessage + if err := json.Unmarshal(out, &got); err != nil { + t.Fatalf("failed to decode output: %v", err) + } + + // 3. Service ordering — Go marshals map keys sorted alphabetically. + wantOrder := []string{`"grafana"`, `"my_ssh"`, `"tensorboard"`, `"vscode"`} + prev := -1 + for _, key := range wantOrder { + idx := strings.Index(string(out), key) + if idx < 0 { + t.Fatalf("expected key %s in output: %s", key, out) + } + if idx <= prev { + t.Errorf("services not sorted alphabetically: %s appears at idx %d (previous %d)", key, idx, prev) + } + prev = idx + } + + // 4. All 6 fields per service, in order; missing values should be `null`. + requiredKeys := []string{"type", "port", "status", "error", "endpoint", "properties"} + for svcName, svc := range got { + for _, k := range requiredKeys { + v, ok := svc[k] + if !ok { + t.Errorf("service %q missing required field %q", svcName, k) + continue + } + if len(v) == 0 { + t.Errorf("service %q field %q is empty raw", svcName, k) + } + } + if len(svc) != len(requiredKeys) { + t.Errorf("service %q has %d fields, want %d (extra: %v)", svcName, len(svc), len(requiredKeys), svc) + } + } + + // 5. vscode.port should be the literal `null` (input omitted it). + if got["vscode"] == nil { + t.Fatal("missing vscode entry") + } + if string(got["vscode"]["port"]) != "null" { + t.Errorf("vscode.port = %s, want null", got["vscode"]["port"]) + } + // vscode.error should also be null (input omitted it). + if string(got["vscode"]["error"]) != "null" { + t.Errorf("vscode.error = %s, want null", got["vscode"]["error"]) + } + + // 6. tensorboard.error flattened to the message string. + wantErr := `"failed to start endpoint tensorboard"` + if string(got["tensorboard"]["error"]) != wantErr { + t.Errorf("tensorboard.error = %s, want %s", got["tensorboard"]["error"], wantErr) + } + + // 7. my_ssh.error explicit null preserved as null. + if string(got["my_ssh"]["error"]) != "null" { + t.Errorf("my_ssh.error = %s, want null", got["my_ssh"]["error"]) + } + + // 8. No HTML-escaping of <, >, & in URLs. + for _, bad := range []string{`\u003c`, `\u003e`, `\u0026`} { + if strings.Contains(string(out), bad) { + t.Errorf("output contains HTML-escape sequence %q (should be literal): %s", bad, out) + } + } + // Positive: literal & must be present in vscode endpoint. + if !strings.Contains(string(out), `"vscode://x?a=1&b=2"`) { + t.Errorf("vscode endpoint not preserved verbatim: %s", out) + } +} + +func TestFlattenServiceError(t *testing.T) { + tests := []struct { + name string + in string + want string + }{ + {name: "nil raw", in: "", want: "null"}, + {name: "literal null", in: "null", want: "null"}, + {name: "valid envelope with message", in: `{"error":{"message":"boom"},"time":"0001-01-01T00:00:00+00:00"}`, want: `"boom"`}, + {name: "envelope with empty message", in: `{"error":{"message":""}}`, want: "null"}, + {name: "missing inner error", in: `{"time":"x"}`, want: "null"}, + {name: "inner error null", in: `{"error":null}`, want: "null"}, + {name: "malformed json", in: `{not json`, want: "null"}, + {name: "message with html chars not escaped", in: `{"error":{"message":"a < b & c > d"}}`, want: `"a < b & c > d"`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := string(flattenServiceError(json.RawMessage(tt.in))) + if got != tt.want { + t.Errorf("flattenServiceError(%s) = %s, want %s", tt.in, got, tt.want) + } + }) + } +} + +func TestMarshalNoEscape(t *testing.T) { + in := map[string]string{"u": "ac&d"} + out, err := marshalNoEscape(in) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + got := string(out) + if strings.Contains(got, `\u003c`) || strings.Contains(got, `\u003e`) || strings.Contains(got, `\u0026`) { + t.Errorf("output contains HTML escapes: %s", got) + } + if !strings.Contains(got, `"ac&d"`) { + t.Errorf("output missing literal value: %s", got) + } + // Must not have a trailing newline. + if strings.HasSuffix(got, "\n") { + t.Errorf("output should not end with newline: %q", got) + } +}