Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
245 changes: 74 additions & 171 deletions services/core/jobs/jobs-launcher/cmd/otel.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,8 @@ import (
"fmt"
"log/slog"
"net/http"
"net/url"
"os"
"strconv"
"strings"
"sync/atomic"
"time"

"github.com/NVIDIA-NeMo/nemo-platform/services/core/jobs/jobs-launcher/nmpclient"
Expand All @@ -33,15 +30,9 @@ const (
NEMO_JOB_ATTEMPT_ID_ENV = "NEMO_JOB_ATTEMPT_ID"
NEMO_JOB_STEP_NAME_ENV = "NEMO_JOB_STEP"
NEMO_JOB_TASK_ID_ENV = "NEMO_JOB_TASK"
serviceJobsPrincipal = "service:jobs"

launcherLogsExporterEnv = "NMP_JOB_LAUNCHER_LOGS_EXPORTER"
launcherOTLPLogsEndpointEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT"
launcherOTLPLogsHeadersEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS"
launcherOTLPLogsProtocolEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL"
launcherOTLPLogsTimeoutEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_TIMEOUT"
launcherOTLPLogsCompressEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_COMPRESSION"
launcherOTLPHTTPProto = "http/protobuf"
defaultLauncherLogsExporter = "console"
otlpHTTPLogExportTimeout = 10 * time.Second
)

Expand Down Expand Up @@ -122,129 +113,93 @@ func newLoggerProvider(ctx context.Context, res *resource.Resource) (*log.Logger
}

func newLogExporter(ctx context.Context) (log.Exporter, error) {
switch strings.ToLower(strings.TrimSpace(os.Getenv(launcherLogsExporterEnv))) {
case "", defaultLauncherLogsExporter, "stdout":
endpointURL := strings.TrimSpace(os.Getenv(launcherOTLPLogsEndpointEnv))
if endpointURL == "" {
return stdoutlog.New()
case "none":
return noopLogExporter{}, nil
case "otlp":
return newLauncherOTLPLogExporter(ctx)
default:
return nil, fmt.Errorf("unsupported %s value %q", launcherLogsExporterEnv, os.Getenv(launcherLogsExporterEnv))
}
return newLauncherOTLPLogExporter(ctx, endpointURL)
}

func newLauncherOTLPLogExporter(ctx context.Context) (log.Exporter, error) {
protocol := strings.TrimSpace(os.Getenv(launcherOTLPLogsProtocolEnv))
if protocol == "" {
protocol = launcherOTLPHTTPProto
func newLauncherOTLPLogExporter(ctx context.Context, endpointURL string) (log.Exporter, error) {
authConfig, err := newLauncherOTLPLogAuthConfig(ctx)
if err != nil {
return nil, err
}
if protocol != launcherOTLPHTTPProto {
return nil, fmt.Errorf("%s must be %q, got %q", launcherOTLPLogsProtocolEnv, launcherOTLPHTTPProto, protocol)
if _, err := authConfig.source.AuthorizationHeader(ctx); err != nil {
return nil, fmt.Errorf("configure %s for OTLP logs: %w", authConfig.description, err)
}
logOTLPLogAuthMechanism(authConfig.mechanism, authConfig.reason)

endpointURL := strings.TrimSpace(os.Getenv(launcherOTLPLogsEndpointEnv))
if endpointURL == "" {
return nil, fmt.Errorf("%s is required when %s=otlp", launcherOTLPLogsEndpointEnv, launcherLogsExporterEnv)
httpClient := launcherOTLPHTTPClient()
if httpClient == nil {
httpClient = &http.Client{Timeout: otlpHTTPLogExportTimeout}
}

// If a workload identity token file is available, use a refreshable auth
// transport that exchanges the projected SA token for platform credentials.
if os.Getenv(workloadIdentityTokenFileEnv) != "" {
tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx)
if err != nil {
return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err)
}
return newRefreshableAuthLogExporter(ctx, endpointURL, tokenSource, otlpHTTPLogExporter)
httpClient.Transport = &authHeaderTransport{
source: authConfig.source,
authDescription: authConfig.description,
base: httpClient.Transport,
}

options := []otlploghttp.Option{
otlploghttp.WithEndpointURL(endpointURL),
otlploghttp.WithHeaders(parseLauncherOTLPHeaders()),
}

timeout, timeoutSet, err := parseLauncherOTLPTimeout()
if err != nil {
return nil, err
}
if httpClient := launcherOTLPHTTPClient(timeout, timeoutSet); httpClient != nil {
options = append(options, otlploghttp.WithHTTPClient(httpClient))
}
if timeoutSet {
options = append(options, otlploghttp.WithTimeout(timeout))
}
if compression, ok, err := parseLauncherOTLPCompression(); err != nil {
return nil, err
} else if ok {
options = append(options, otlploghttp.WithCompression(compression))
// Platform job log upload is intentionally independent from OTEL_* header
// env vars. User workloads may use OTEL_EXPORTER_OTLP_LOGS_HEADERS for
// third-party telemetry; application logs construct Authorization in Go.
// Passing explicit empty headers prevents the upstream exporter from
// applying OTEL_EXPORTER_OTLP_LOGS_HEADERS here.
otlploghttp.WithHeaders(map[string]string{}),
otlploghttp.WithHTTPClient(httpClient),
}

return otlploghttp.New(ctx, options...)
}

type authHeaderSource interface {
AuthorizationHeader(context.Context) (string, error)
return otlpHTTPLogExporter(ctx, options...)
}

type logExporterFactory func(context.Context, ...otlploghttp.Option) (log.Exporter, error)

type refreshableAuthLogExporter struct {
exporter log.Exporter
stopped atomic.Bool
}

func newRefreshableAuthLogExporter(
ctx context.Context,
endpoint string,
authSource authHeaderSource,
newExporter logExporterFactory,
) (log.Exporter, error) {
if _, err := authSource.AuthorizationHeader(ctx); err != nil {
return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err)
}

exporter, err := newExporter(
ctx,
otlploghttp.WithEndpointURL(endpoint),
otlploghttp.WithHTTPClient(&http.Client{
Transport: &authHeaderTransport{source: authSource},
Timeout: otlpHTTPLogExportTimeout,
}),
)
if err != nil {
return nil, err
}
return &refreshableAuthLogExporter{exporter: exporter}, nil
type launcherOTLPLogAuthConfig struct {
source authHeaderSource
description string
mechanism string
reason string
}

func (e *refreshableAuthLogExporter) Export(ctx context.Context, records []log.Record) error {
if e.stopped.Load() {
return nil
func newLauncherOTLPLogAuthConfig(ctx context.Context) (launcherOTLPLogAuthConfig, error) {
if os.Getenv(workloadIdentityTokenFileEnv) != "" {
// True auth: exchange the workload subject token and send the resulting bearer token.
tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx)
if err != nil {
return launcherOTLPLogAuthConfig{}, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err)
}
return launcherOTLPLogAuthConfig{
source: tokenSource,
description: "workload identity auth",
mechanism: "workload_identity_token_exchange",
reason: "workload_token_file_configured",
}, nil
}

return e.exporter.Export(ctx, records)
}

func (e *refreshableAuthLogExporter) Shutdown(ctx context.Context) error {
if e.stopped.Swap(true) {
return nil
}
return e.exporter.Shutdown(ctx)
// Service identity default: when workload identity is not enabled, identify the
// launcher as service:jobs without touching user OTEL_* headers.
return launcherOTLPLogAuthConfig{
source: serviceIdentityBearerTokenSource{},
description: "service identity bearer auth",
mechanism: "service_identity_bearer_token",
reason: "workload_token_file_not_configured",
}, nil
}

func (e *refreshableAuthLogExporter) ForceFlush(ctx context.Context) error {
return e.exporter.ForceFlush(ctx)
type authHeaderSource interface {
AuthorizationHeader(context.Context) (string, error)
}

type authHeaderTransport struct {
source authHeaderSource
base http.RoundTripper
source authHeaderSource
authDescription string
base http.RoundTripper
}

func (t *authHeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
authHeader, err := t.source.AuthorizationHeader(req.Context())
if err != nil {
return nil, fmt.Errorf("refresh workload identity auth for OTLP logs: %w", err)
return nil, fmt.Errorf("refresh %s for OTLP logs: %w", t.authDescription, err)
}

clonedReq := req.Clone(req.Context())
Expand All @@ -260,80 +215,28 @@ func (t *authHeaderTransport) baseTransport() http.RoundTripper {
return http.DefaultTransport
}

func otlpHTTPLogExporter(ctx context.Context, opts ...otlploghttp.Option) (log.Exporter, error) {
return otlploghttp.New(ctx, opts...)
func logOTLPLogAuthMechanism(mechanism string, reason string) {
slog.Info(
"using OTLP log authentication",
"auth_mechanism", mechanism,
"reason", reason,
)
}

func launcherOTLPHTTPClient(timeout time.Duration, timeoutSet bool) *http.Client {
endpoint, err := nmpclient.ResolvePlatformEndpointFromEnv()
if err != nil || endpoint.Transport != nmpclient.TransportUDS {
return nil
}
httpClient := endpoint.HTTPClient()
if timeoutSet {
httpClient.Timeout = timeout
}
return httpClient
}
type serviceIdentityBearerTokenSource struct{}

func parseLauncherOTLPHeaders() map[string]string {
raw := strings.TrimSpace(os.Getenv(launcherOTLPLogsHeadersEnv))
if raw == "" {
return nil
}
headers := map[string]string{}
for _, item := range strings.Split(raw, ",") {
key, value, ok := strings.Cut(strings.TrimSpace(item), "=")
if !ok || key == "" {
continue
}
if decoded, err := url.PathUnescape(value); err == nil {
value = decoded
}
headers[key] = value
}
return headers
func (serviceIdentityBearerTokenSource) AuthorizationHeader(context.Context) (string, error) {
return "Bearer " + serviceJobsPrincipal, nil
}

func parseLauncherOTLPTimeout() (time.Duration, bool, error) {
raw := strings.TrimSpace(os.Getenv(launcherOTLPLogsTimeoutEnv))
if raw == "" {
return 0, false, nil
}
duration, err := time.ParseDuration(raw)
if err == nil {
return duration, true, nil
}
milliseconds, intErr := strconv.Atoi(raw)
if intErr == nil {
return time.Duration(milliseconds) * time.Millisecond, true, nil
}
return 0, false, fmt.Errorf("invalid %s value %q: %w", launcherOTLPLogsTimeoutEnv, raw, err)
func otlpHTTPLogExporter(ctx context.Context, opts ...otlploghttp.Option) (log.Exporter, error) {
return otlploghttp.New(ctx, opts...)
}

func parseLauncherOTLPCompression() (otlploghttp.Compression, bool, error) {
switch strings.ToLower(strings.TrimSpace(os.Getenv(launcherOTLPLogsCompressEnv))) {
case "":
return otlploghttp.NoCompression, false, nil
case "none":
return otlploghttp.NoCompression, true, nil
case "gzip":
return otlploghttp.GzipCompression, true, nil
default:
return otlploghttp.NoCompression, false, fmt.Errorf("unsupported %s value %q", launcherOTLPLogsCompressEnv, os.Getenv(launcherOTLPLogsCompressEnv))
func launcherOTLPHTTPClient() *http.Client {
endpoint, err := nmpclient.ResolvePlatformEndpointFromEnv()
if err != nil || endpoint.Transport != nmpclient.TransportUDS {
return nil
}
}

type noopLogExporter struct{}

func (noopLogExporter) Export(context.Context, []log.Record) error {
return nil
}

func (noopLogExporter) Shutdown(context.Context) error {
return nil
}

func (noopLogExporter) ForceFlush(context.Context) error {
return nil
return endpoint.HTTPClient()
}
23 changes: 4 additions & 19 deletions services/core/jobs/jobs-launcher/cmd/otel_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,27 +3,12 @@

package cmd

import (
"testing"
"time"
)
import "testing"

func TestLauncherOTLPHTTPClientSetsUDSTimeout(t *testing.T) {
func TestLauncherOTLPHTTPClientUsesUDSEndpoint(t *testing.T) {
t.Setenv("NMP_BASE_URL", "unix:///tmp/nemo-platform.sock")

client := launcherOTLPHTTPClient(250*time.Millisecond, true)
if client == nil {
t.Fatal("expected UDS HTTP client")
}
if client.Timeout != 250*time.Millisecond {
t.Fatalf("expected UDS HTTP client timeout 250ms, got %s", client.Timeout)
}
}

func TestLauncherOTLPHTTPClientLeavesUDSTimeoutUnset(t *testing.T) {
t.Setenv("NMP_BASE_URL", "unix:///tmp/nemo-platform.sock")

client := launcherOTLPHTTPClient(250*time.Millisecond, false)
client := launcherOTLPHTTPClient()
if client == nil {
t.Fatal("expected UDS HTTP client")
}
Expand All @@ -35,7 +20,7 @@ func TestLauncherOTLPHTTPClientLeavesUDSTimeoutUnset(t *testing.T) {
func TestLauncherOTLPHTTPClientSkipsTCP(t *testing.T) {
t.Setenv("NMP_BASE_URL", "http://127.0.0.1:8080")

if client := launcherOTLPHTTPClient(250*time.Millisecond, true); client != nil {
if client := launcherOTLPHTTPClient(); client != nil {
t.Fatal("expected no custom HTTP client for TCP endpoint")
}
}
25 changes: 0 additions & 25 deletions services/core/jobs/jobs-launcher/cmd/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ import (
"io"
"log/slog"
"net/http"
"net/url"
"os"
"os/exec"
"os/signal"
Expand Down Expand Up @@ -182,30 +181,6 @@ func runExecWithStdin(args []string) (exitCode int, err error) {
return runExec(args, os.Stdin)
}

func configureOTELHeadersFromWorkloadToken() {
token := os.Getenv("NEMO_WORKLOAD_TOKEN")
if token == "" {
return
}

const headersEnv = launcherOTLPLogsHeadersEnv
headers := os.Getenv(headersEnv)
for _, item := range strings.Split(headers, ",") {
key, _, _ := strings.Cut(strings.TrimSpace(item), "=")
if strings.EqualFold(key, "authorization") {
return
}
}

authHeader := "Authorization=" + url.PathEscape("Bearer "+token)
if headers == "" {
os.Setenv(headersEnv, authHeader)
return
}
os.Setenv(headersEnv, headers+","+authHeader)
}


// runExec runs the specified command with arguments, injecting secrets as environment variables if specified
func runExec(args []string, stdinReader io.Reader) (int, error) {
// Command and arguments
Expand Down
Loading