Add helpers to interact with pipeline options in boot entrypoints

CHANGES.md

@@ -80,7 +80,7 @@
 
 ## Bugfixes
 
-* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)).
 
 ## Security Fixes
 
@@ -180,7 +180,6 @@
 * (Python) Typehints of dataclass fields are honored during type inferences. To restore the behavior of fallback-to-any,
   use pipeline option `--exclude_infer_dataclass_field_type` ([#38797](https://github.com/apache/beam/issues/38797)).
   However fixing forward is recommended.
-* X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Bugfixes
 

sdks/go/container/boot.go

@@ -17,7 +17,6 @@ package main
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -31,7 +30,6 @@ import (
 	"github.com/apache/beam/sdks/v2/go/container/pool"
 	"github.com/apache/beam/sdks/v2/go/container/tools"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
-	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime"
 
 	// Import gcs filesystem so that it can be used to upload heap dumps
 	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/io/filesystem/gcs"
@@ -61,23 +59,16 @@ const (
 	workerPoolIdEnv                 = "BEAM_GO_WORKER_POOL_ID"
 )
 
-func configureGoogleCloudProfilerEnvVars(ctx context.Context, logger *tools.Logger, metadata map[string]string, options string) error {
+func configureGoogleCloudProfilerEnvVars(ctx context.Context, logger *tools.Logger, metadata map[string]string, po *tools.PipelineOptions) error {
 	const profilerKey = "enable_google_cloud_profiler="
 
-	var parsed map[string]interface{}
-	if err := json.Unmarshal([]byte(options), &parsed); err != nil {
-		panic(err)
-	}
-
 	var profilerServiceName string
 
-	// Try from "beam:option:go_options:v1" -> "options" -> "dataflow_service_options"
-	if goOpts, ok := parsed["beam:option:go_options:v1"].(map[string]interface{}); ok {
-		if options, ok := goOpts["options"].(map[string]interface{}); ok {
-			if profilerServiceNameRaw, ok := options["dataflow_service_options"].(string); ok {
-				if strings.HasPrefix(profilerServiceNameRaw, profilerKey) {
-					profilerServiceName = strings.TrimPrefix(profilerServiceNameRaw, profilerKey)
-				}
+	if serviceOpts, err := po.GetStringSlice("dataflow_service_options"); err == nil {
+		for _, opt := range serviceOpts {
+			if strings.HasPrefix(opt, profilerKey) {
+				profilerServiceName = strings.TrimPrefix(opt, profilerKey)
+				break
 			}
 		}
 	}
@@ -159,8 +150,11 @@ func main() {
 		logger.Fatalf(ctx, "Failed to convert pipeline options: %v", err)
 	}
 
+	// Go SDK wraps pipeline options inside the URN namespace: "beam:option:go_options:v1".
+	po := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "go_options")
+
 	// Inject artifact validation enabled state into context
-	ctx = artifact.WithArtifactValidation(ctx, !artifact.HasExperiment(info.GetPipelineOptions(), "disable_staged_file_integrity_checks"))
+	ctx = artifact.WithArtifactValidation(ctx, !po.HasExperiment("disable_staged_file_integrity_checks"))
 
 	// (2) Retrieve the staged files.
 	//
@@ -210,9 +204,11 @@ func main() {
 		os.Setenv("RUNNER_CAPABILITIES", strings.Join(info.GetRunnerCapabilities(), " "))
 	}
 
-	enableGoogleCloudProfiler := strings.Contains(options, enableGoogleCloudProfilerOption)
+	// Go SDK models multi-value list flags (like dataflow_service_options) as comma-separated strings.
+	serviceOpts, _ := po.GetString("dataflow_service_options")
+	enableGoogleCloudProfiler := strings.Contains(serviceOpts, "enable_google_cloud_profiler")
 	if enableGoogleCloudProfiler {
-		err := configureGoogleCloudProfilerEnvVars(ctx, logger, info.Metadata, options)
+		err := configureGoogleCloudProfilerEnvVars(ctx, logger, info.Metadata, po)
 		if err != nil {
 			logger.Printf(ctx, "could not configure Google Cloud Profiler variables, got %v", err)
 		}
@@ -221,12 +217,8 @@ func main() {
 	err = execx.Execute(prog, args...)
 
 	if err != nil {
-		var opt runtime.RawOptionsWrapper
-		err := json.Unmarshal([]byte(options), &opt)
-		if err == nil {
-			if tempLocation, ok := opt.Options.Options["temp_location"]; ok {
-				diagnostics.UploadHeapProfile(ctx, fmt.Sprintf("%v/heapProfiles/profile-%v-%d", strings.TrimSuffix(tempLocation, "/"), *id, time.Now().Unix()))
-			}
+		if tempLocation, err := po.GetString("temp_location"); err == nil && tempLocation != "" {
+			diagnostics.UploadHeapProfile(ctx, fmt.Sprintf("%v/heapProfiles/profile-%v-%d", strings.TrimSuffix(tempLocation, "/"), *id, time.Now().Unix()))
 		}
 	}
 

sdks/go/container/boot_test.go

@@ -21,11 +21,14 @@ import (
 	"path/filepath"
 	"testing"
 
+	"encoding/json"
+
 	"github.com/apache/beam/sdks/v2/go/container/tools"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
 	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
 	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 func TestEnsureEndpointsSet_AllSet(t *testing.T) {
@@ -224,7 +227,7 @@ func TestConfigureGoogleCloudProfilerEnvVars(t *testing.T) {
 			options: `{
 				"beam:option:go_options:v1": {
 					"options": {
-						"dataflow_service_options": "enable_google_cloud_profiler=custom_profiler"
+						"dataflow_service_options": "enable_google_cloud_profiler=custom_profiler,another_option"
 					}
 				}
 			}`,
@@ -287,7 +290,17 @@ func TestConfigureGoogleCloudProfilerEnvVars(t *testing.T) {
 			clearEnvVars()
 			ctx := context.Background()
 
-			err := configureGoogleCloudProfilerEnvVars(ctx, &tools.Logger{}, tt.metadata, tt.options)
+			var raw map[string]interface{}
+			if err := json.Unmarshal([]byte(tt.options), &raw); err != nil {
+				t.Fatalf("failed to unmarshal JSON for test: %v", err)
+			}
+			st, err := structpb.NewStruct(raw)
+			if err != nil {
+				t.Fatalf("failed to create structpb for test: %v", err)
+			}
+			po := tools.ParseOptionsFromProto(st, "go_options")
+
+			err = configureGoogleCloudProfilerEnvVars(ctx, &tools.Logger{}, tt.metadata, po)
 
 			if tt.expectingError {
 				if err == nil {

sdks/go/container/tools/pipeline_options.go

@@ -19,6 +19,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
+	"strings"
+
+	structpb "google.golang.org/protobuf/types/known/structpb"
 )
 
 // MakePipelineOptionsFileAndEnvVar writes the pipeline options to a file.
@@ -42,3 +46,217 @@ func MakePipelineOptionsFileAndEnvVar(options string) error {
 	os.Setenv("PIPELINE_OPTIONS_FILE", f.Name())
 	return nil
 }
+
+// PipelineOptions represents parsed pipeline options as a normalized map.
+type PipelineOptions struct {
+	options     map[string]any
+	experiments map[string]string
+}
+
+// ParseOptionsFromProto creates normalized PipelineOptions directly from a protobuf Struct.
+func ParseOptionsFromProto(opt *structpb.Struct, sdkNamespace string) *PipelineOptions {
+	if opt == nil {
+		return &PipelineOptions{options: make(map[string]any), experiments: make(map[string]string)}
+	}
+	raw := opt.AsMap()
+	flat := make(map[string]any)
+
+	// 1. Extract nested options if present (Dataflow runner uses this structure)
+	if optsVal, ok := raw["options"]; ok {
+		if optsMap, ok := optsVal.(map[string]any); ok {
+			for k, v := range optsMap {
+				flat[k] = v
+			}
+		}
+	}
+
+	// 2. Extract standard URN keys (Portable runners use this structure)
+	for k, v := range raw {
+		if k == "options" || k == "display_data" {
+			continue
+		}
+		if strings.HasPrefix(k, "beam:option:") && strings.HasSuffix(k, ":v1") {
+			name := strings.TrimPrefix(k, "beam:option:")
+			name = strings.TrimSuffix(name, ":v1")
+			flat[name] = v
+		}
+	}
+
+	// 3. Promote specified SDK namespace options (Highest precedence, may overwrite earlier entries).
+	// Beam Go SDK uses this structure.
+	if sdkNamespace != "" {
+		sdkURN := fmt.Sprintf("beam:option:%s:v1", sdkNamespace)
+		if sdkVal, ok := raw[sdkURN]; ok {
+			if urnMap, ok := sdkVal.(map[string]any); ok {
+				if nestedOpts, ok := urnMap["options"].(map[string]any); ok {
+					for nk, nv := range nestedOpts {
+						flat[nk] = nv
+					}
+				}
+			}
+		}
+	}
+
+	po := &PipelineOptions{
+		options:     flat,
+		experiments: make(map[string]string),
+	}
+	if exps, err := po.GetStringSlice("experiments"); err == nil {
+		po.experiments = parseExperiments(exps)
+	}
+	return po
+}
+
+func parseExperiments(slice []string) map[string]string {
+	res := make(map[string]string)
+	for _, item := range slice {
+		if strings.Contains(item, "=") {
+			parts := strings.SplitN(item, "=", 2)
+			res[parts[0]] = parts[1]
+		} else {
+			res[item] = ""
+		}
+	}
+	return res
+}
+
+// HasOption returns true if the option is defined and not nil.
+func (po *PipelineOptions) HasOption(name string) bool {
+	val, ok := po.options[name]
+	return ok && val != nil
+}
+
+// GetString returns the value of an option as a string.
+// As a convenience and to maintain compatibility with Go SDK's flags serialization style,
+// if the option is stored as a string slice/array, GetString will conjoin the elements
+// into a single comma-separated string (e.g. ["opt1", "opt2"] -> "opt1,opt2").
+func (po *PipelineOptions) GetString(name string) (string, error) {
+	val, ok := po.options[name]
+	if !ok || val == nil {
+		return "", fmt.Errorf("option %q not defined", name)
+	}
+	if str, ok := val.(string); ok {
+		return str, nil
+	}
+	if slice, ok := val.([]any); ok {
+		var parts []string
+		for _, item := range slice {
+			if str, ok := item.(string); ok {
+				parts = append(parts, str)
+			} else {
+				return "", fmt.Errorf("option %q: expected string slice element, got type %T", name, item)
+			}
+		}
+		return strings.Join(parts, ","), nil
+	}
+	return "", fmt.Errorf("option %q: expected string, got type %T", name, val)
+}
+
+// GetStringSlice returns the value of an option as a string slice.
+// As a convenience and to maintain compatibility with Go SDK's flags serialization style,
+// if the option is stored as a single comma-separated string (such as experiments
+// or dataflow_service_options), GetStringSlice will parse it by splitting the string
+// by comma (e.g. "opt1,opt2" -> ["opt1", "opt2"]).
+func (po *PipelineOptions) GetStringSlice(name string) ([]string, error) {
+	val, ok := po.options[name]
+	if !ok || val == nil {
+		return nil, fmt.Errorf("option %q not defined", name)
+	}
+	if slice, ok := val.([]any); ok {
+		var res []string
+		for _, item := range slice {
+			if str, ok := item.(string); ok {
+				res = append(res, str)
+			} else {
+				return nil, fmt.Errorf("option %q: expected string slice element, got type %T", name, item)
+			}
+		}
+		return res, nil
+	}
+	if str, ok := val.(string); ok {
+		// Go SDK models multi-value list flags (like experiments or dataflow_service_options)
+		// as comma-separated string values.
+		if str == "" {
+			return nil, nil
+		}
+		return strings.Split(str, ","), nil
+	}
+	return nil, fmt.Errorf("option %q: expected string slice, got type %T", name, val)
+}
+
+// GetInt returns the value of an option as an integer.
+func (po *PipelineOptions) GetInt(name string) (int, error) {
+	val, ok := po.options[name]
+	if !ok || val == nil {
+		return 0, fmt.Errorf("option %q not defined", name)
+	}
+	switch v := val.(type) {
+	case float64:
+		return int(v), nil
+	case string:
+		res, err := strconv.Atoi(v)
+		if err == nil {
+			return res, nil
+		}
+		return 0, fmt.Errorf("option %q: failed to parse %q as int: %w", name, v, err)
+	default:
+		return 0, fmt.Errorf("option %q: expected int (represented as number or string), got type %T", name, val)
+	}
+}
+
+// GetBool returns the value of an option as a boolean.
+func (po *PipelineOptions) GetBool(name string) (bool, error) {
+	val, ok := po.options[name]
+	if !ok || val == nil {
+		return false, fmt.Errorf("option %q not defined", name)
+	}
+	switch v := val.(type) {
+	case bool:
+		return v, nil
+	case string:
+		res, err := strconv.ParseBool(v)
+		if err != nil {
+			return false, fmt.Errorf("option %q: failed to parse %q as bool: %w", name, v, err)
+		}
+		return res, nil
+	case float64:
+		return v != 0, nil
+	default:
+		return false, fmt.Errorf("option %q: expected bool, got type %T", name, val)
+	}
+}
+
+// GetFloat64 returns the value of an option as a float64.
+func (po *PipelineOptions) GetFloat64(name string) (float64, error) {
+	val, ok := po.options[name]
+	if !ok || val == nil {
+		return 0, fmt.Errorf("option %q not defined", name)
+	}
+	switch v := val.(type) {
+	case float64:
+		return v, nil
+	case string:
+		res, err := strconv.ParseFloat(v, 64)
+		if err == nil {
+			return res, nil
+		}
+		return 0, fmt.Errorf("option %q: failed to parse %q as float64: %w", name, v, err)
+	default:
+		return 0, fmt.Errorf("option %q: expected float64 (represented as number or string), got type %T", name, val)
+	}
+}
+
+// LookupExperiment returns the value of an experiment option if present.
+// - If the experiment is present but has no value (e.g., --experiments=foo), it returns "", true.
+// - If the experiment is present as a key-value pair (e.g., --experiments=foo=bar), it returns "bar", true.
+// - If the experiment is not present, it returns "", false.
+func (po *PipelineOptions) LookupExperiment(key string) (string, bool) {
+	val, ok := po.experiments[key]
+	return val, ok
+}
+
+// HasExperiment returns true if the specified experiment is present in the options (either as a flag or key-value pair).
+func (po *PipelineOptions) HasExperiment(name string) bool {
+	_, ok := po.LookupExperiment(name)
+	return ok
+}