This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0aacd0375f0 Add helpers to interact with pipeline options in boot
entrypoints (#39595)
0aacd0375f0 is described below
commit 0aacd0375f0b7558e845d126334ad9461368ff3d
Author: tvalentyn <[email protected]>
AuthorDate: Fri Aug 7 10:44:36 2026 -0700
Add helpers to interact with pipeline options in boot entrypoints (#39595)
* Create helpers to parse pipeline options in boot entrypoints that work
for Dataflow and Portable runners.
* Better handle the go SDK case.
* Support SDK namespaces for more deterministic parsing.
* fix test.
* Remove the error signature in ParseOptionsFromProto
* Simplify the tests.
* Simplify profiler setting logic in Go's boot.go
* Also support retrieving sliced options as a single commma-separated
string.
* Support float.
* Add a note on v1.
* Move the feature to a new release.
---
CHANGES.md | 3 +-
sdks/go/container/boot.go | 40 ++--
sdks/go/container/boot_test.go | 17 +-
sdks/go/container/tools/pipeline_options.go | 218 ++++++++++++++++++
sdks/go/container/tools/pipeline_options_test.go | 243 +++++++++++++++++++++
sdks/go/pkg/beam/artifact/options.go | 48 ----
sdks/go/pkg/beam/artifact/options_test.go | 78 -------
sdks/java/container/boot.go | 4 +-
.../python/apache_beam/options/pipeline_options.py | 1 +
sdks/python/container/boot.go | 57 +----
sdks/python/container/profiler.go | 69 ++++--
sdks/python/container/profiler_test.go | 19 +-
sdks/typescript/container/boot.go | 4 +-
13 files changed, 579 insertions(+), 222 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 81f7d493677..6f61ef415b4 100644
--- a/CHANGES.md
+++ b/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
@@ -184,7 +184,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
diff --git a/sdks/go/container/boot.go b/sdks/go/container/boot.go
index 469285821f7..aeeb87cc811 100644
--- a/sdks/go/container/boot.go
+++ b/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()))
}
}
diff --git a/sdks/go/container/boot_test.go b/sdks/go/container/boot_test.go
index bb94aca36be..57aa083245f 100644
--- a/sdks/go/container/boot_test.go
+++ b/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 {
diff --git a/sdks/go/container/tools/pipeline_options.go
b/sdks/go/container/tools/pipeline_options.go
index 026fb31b099..dff51540311 100644
--- a/sdks/go/container/tools/pipeline_options.go
+++ b/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
+}
diff --git a/sdks/go/container/tools/pipeline_options_test.go
b/sdks/go/container/tools/pipeline_options_test.go
index 7a0d7ebd5f0..c220a501756 100644
--- a/sdks/go/container/tools/pipeline_options_test.go
+++ b/sdks/go/container/tools/pipeline_options_test.go
@@ -16,10 +16,28 @@
package tools
import (
+ "encoding/json"
"os"
"testing"
+
+ structpb "google.golang.org/protobuf/types/known/structpb"
)
+func parseProtoForTest(t *testing.T, options string) *structpb.Struct {
+ if options == "" {
+ options = "{}"
+ }
+ var raw map[string]interface{}
+ if err := json.Unmarshal([]byte(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)
+ }
+ return st
+}
+
func TestMakePipelineOptionsFileAndEnvVar(t *testing.T) {
tests := []struct {
name string
@@ -56,3 +74,228 @@ func TestMakePipelineOptionsFileAndEnvVar(t *testing.T) {
}
os.Remove("pipeline_options.json")
}
+
+func TestParseOptionsFromProto_NestedOptionsNoNamespace(t *testing.T) {
+ p := parseProtoForTest(t, `{
+ "options": {
+ "profiler_agent": "memray",
+ "profile_upload_interval_sec": 10,
+ "profiler_stop_after_crash": true,
+ "profile_sample_rate": 0.5,
+ "experiments": ["beam_fn_api",
"pip_use_build_isolation"]
+ }
+ }`)
+ po := ParseOptionsFromProto(p, "")
+
+ if got, err := po.GetString("profiler_agent"); err != nil || got !=
"memray" {
+ t.Errorf("GetString(profiler_agent) = (%q, %v), want
(\"memray\", nil)", got, err)
+ }
+ if got, err := po.GetInt("profile_upload_interval_sec"); err != nil ||
got != 10 {
+ t.Errorf("GetInt(profile_upload_interval_sec) = (%d, %v), want
(10, nil)", got, err)
+ }
+ if got, err := po.GetBool("profiler_stop_after_crash"); err != nil ||
got != true {
+ t.Errorf("GetBool(profiler_stop_after_crash) = (%t, %v), want
(true, nil)", got, err)
+ }
+ // Sample float option in Beam. Unused for memray in practice.
+ if got, err := po.GetFloat64("profile_sample_rate"); err != nil || got
!= 0.5 {
+ t.Errorf("GetFloat64(profile_sample_rate) = (%f, %v), want
(0.5, nil)", got, err)
+ }
+ experiments, err := po.GetStringSlice("experiments")
+ if err != nil || len(experiments) != 2 || experiments[0] !=
"beam_fn_api" || experiments[1] != "pip_use_build_isolation" {
+ t.Errorf("GetStringSlice(experiments) = (%v, %v), want
([beam_fn_api, pip_use_build_isolation], nil)", experiments, err)
+ }
+ if !po.HasExperiment("beam_fn_api") ||
!po.HasExperiment("pip_use_build_isolation") {
+ t.Errorf("expected experiments beam_fn_api and
pip_use_build_isolation to be present, experiments map: %+v", po.experiments)
+ }
+}
+
+func TestParseOptionsFromProto_FlatOptionsWithURN(t *testing.T) {
+ p := parseProtoForTest(t, `{
+ "beam:option:profiler_agent:v1": "memray",
+ "beam:option:profile_upload_interval_sec:v1": "10",
+ "beam:option:profiler_stop_after_crash:v1": "true",
+ "beam:option:profile_sample_rate:v1": "0.5",
+ "beam:option:experiments:v1": ["beam_fn_api", "another_exp"]
+ }`)
+ po := ParseOptionsFromProto(p, "")
+
+ if got, err := po.GetString("profiler_agent"); err != nil || got !=
"memray" {
+ t.Errorf("GetString(profiler_agent) = (%q, %v), want
(\"memray\", nil)", got, err)
+ }
+ if got, err := po.GetInt("profile_upload_interval_sec"); err != nil ||
got != 10 {
+ t.Errorf("GetInt(profile_upload_interval_sec) = (%d, %v), want
(10, nil)", got, err)
+ }
+ if got, err := po.GetBool("profiler_stop_after_crash"); err != nil ||
got != true {
+ t.Errorf("GetBool(profiler_stop_after_crash) = (%t, %v), want
(true, nil)", got, err)
+ }
+ if got, err := po.GetFloat64("profile_sample_rate"); err != nil || got
!= 0.5 {
+ t.Errorf("GetFloat64(profile_sample_rate) = (%f, %v), want
(0.5, nil)", got, err)
+ }
+ if got, err := po.GetString("experiments"); err != nil || got !=
"beam_fn_api,another_exp" {
+ t.Errorf("GetString(experiments) = (%q, %v), want
(\"beam_fn_api,another_exp\", nil)", got, err)
+ }
+ experiments, err := po.GetStringSlice("experiments")
+ if err != nil || len(experiments) != 2 || experiments[0] !=
"beam_fn_api" || experiments[1] != "another_exp" {
+ t.Errorf("GetStringSlice(experiments) = (%v, %v), want
([beam_fn_api, another_exp], nil)", experiments, err)
+ }
+ if !po.HasExperiment("beam_fn_api") || !po.HasExperiment("another_exp")
{
+ t.Errorf("expected experiments beam_fn_api and another_exp to
be present, experiments map: %+v", po.experiments)
+ }
+}
+
+func TestParseOptionsFromProto_CommaSeparatedExperiments(t *testing.T) {
+ p := parseProtoForTest(t, `{
+ "beam:option:go_options:v1": {
+ "options": {
+ "experiments": "exp1,exp2,exp3",
+ "dataflow_service_options": "opt1"
+ }
+ }
+ }`)
+ po := ParseOptionsFromProto(p, "go_options")
+
+ experiments, err := po.GetStringSlice("experiments")
+ if err != nil || len(experiments) != 3 || experiments[0] != "exp1" ||
experiments[1] != "exp2" || experiments[2] != "exp3" {
+ t.Errorf("GetStringSlice(experiments) = (%v, %v), want ([exp1,
exp2, exp3], nil)", experiments, err)
+ }
+ if !po.HasExperiment("exp1") || !po.HasExperiment("exp2") ||
!po.HasExperiment("exp3") {
+ t.Errorf("expected experiments exp1, exp2, and exp3 to be
present, experiments map: %+v", po.experiments)
+ }
+ serviceOpts, err := po.GetStringSlice("dataflow_service_options")
+ if err != nil || len(serviceOpts) != 1 || serviceOpts[0] != "opt1" {
+ t.Errorf("GetStringSlice(dataflow_service_options) = (%v, %v),
want ([opt1], nil)", serviceOpts, err)
+ }
+}
+
+func TestParseOptionsFromProto_MalformedOptions(t *testing.T) {
+ t.Run("malformed integer", func(t *testing.T) {
+ p := parseProtoForTest(t, `{"options":
{"profile_upload_interval_sec": "invalid"}}`)
+ po := ParseOptionsFromProto(p, "")
+ _, err := po.GetInt("profile_upload_interval_sec")
+ if err == nil {
+ t.Errorf("expected error, got nil")
+ }
+ })
+
+ t.Run("malformed bool", func(t *testing.T) {
+ p := parseProtoForTest(t, `{"options":
{"profiler_stop_after_crash": "maybe"}}`)
+ po := ParseOptionsFromProto(p, "")
+ _, err := po.GetBool("profiler_stop_after_crash")
+ if err == nil {
+ t.Errorf("expected error, got nil")
+ }
+ })
+
+ t.Run("type mismatch int expected got bool", func(t *testing.T) {
+ p := parseProtoForTest(t, `{"options":
{"profile_upload_interval_sec": true}}`)
+ po := ParseOptionsFromProto(p, "")
+ _, err := po.GetInt("profile_upload_interval_sec")
+ if err == nil {
+ t.Errorf("expected error, got nil")
+ }
+ })
+
+ t.Run("missing key returns error", func(t *testing.T) {
+ p := parseProtoForTest(t, `{}`)
+ po := ParseOptionsFromProto(p, "")
+ _, errInt := po.GetInt("profile_upload_interval_sec")
+ _, errBool := po.GetBool("profiler_stop_after_crash")
+ _, errFloat := po.GetFloat64("profile_sample_rate")
+ if errInt == nil || errBool == nil || errFloat == nil {
+ t.Errorf("expected error for missing keys, got:
intErr=%v, boolErr=%v, floatErr=%v", errInt, errBool, errFloat)
+ }
+ })
+}
+
+func TestPipelineOptions_HasOption(t *testing.T) {
+ p := parseProtoForTest(t, `{"options": {"profile_upload_interval_sec":
10}}`)
+ po := ParseOptionsFromProto(p, "")
+ if !po.HasOption("profile_upload_interval_sec") {
+ t.Errorf("HasOption(profile_upload_interval_sec) = false, want
true")
+ }
+ if po.HasOption("profiler_stop_after_crash") {
+ t.Errorf("HasOption(profiler_stop_after_crash) = true, want
false")
+ }
+}
+
+func TestPipelineOptions_HasExperiment(t *testing.T) {
+ p := parseProtoForTest(t, `{"options": {"experiments": ["exp1",
"exp2=val2"]}}`)
+ po := ParseOptionsFromProto(p, "")
+ if !po.HasExperiment("exp1") {
+ t.Errorf("HasExperiment(exp1) = false, want true")
+ }
+ if !po.HasExperiment("exp2") {
+ t.Errorf("HasExperiment(exp2) = false, want true")
+ }
+ if po.HasExperiment("exp3") {
+ t.Errorf("HasExperiment(exp3) = true, want false")
+ }
+}
+
+func TestPipelineOptions_LookupExperiment(t *testing.T) {
+ p := parseProtoForTest(t, `{"options": {"experiments": ["exp1",
"exp2=val2", "exp3=val3=val4"]}}`)
+ po := ParseOptionsFromProto(p, "")
+
+ val, ok := po.LookupExperiment("exp1")
+ if !ok || val != "" {
+ t.Errorf("LookupExperiment(exp1) = (%q, %t), want (\"\",
true)", val, ok)
+ }
+
+ val, ok = po.LookupExperiment("exp2")
+ if !ok || val != "val2" {
+ t.Errorf("LookupExperiment(exp2) = (%q, %t), want (\"val2\",
true)", val, ok)
+ }
+
+ val, ok = po.LookupExperiment("exp3")
+ if !ok || val != "val3=val4" {
+ t.Errorf("LookupExperiment(exp3) = (%q, %t), want
(\"val3=val4\", true)", val, ok)
+ }
+
+ val, ok = po.LookupExperiment("exp4")
+ if ok || val != "" {
+ t.Errorf("LookupExperiment(exp4) = (%q, %t), want (\"\",
false)", val, ok)
+ }
+}
+
+func TestParseOptionsFromProto_SDKOptionsPromotion(t *testing.T) {
+ optionsStruct, err := structpb.NewStruct(map[string]interface{}{
+ "options": map[string]interface{}{
+ "region": "us-central1",
+ },
+ "beam:option:experiments:v1": []interface{}{"expA", "expB"},
+ "beam:option:go_options:v1": map[string]interface{}{
+ "options": map[string]interface{}{
+ "dataflow_service_options":
"enable_google_cloud_profiler,enable_new_custom_feature",
+ },
+ },
+ })
+ if err != nil {
+ t.Fatalf("failed to create proto Struct: %v", err)
+ }
+
+ po := ParseOptionsFromProto(optionsStruct, "go_options")
+ if got, err := po.GetString("region"); err != nil || got !=
"us-central1" {
+ t.Errorf("GetString(region) = (%q, %v), want (\"us-central1\",
nil)", got, err)
+ }
+ if !po.HasExperiment("expA") || !po.HasExperiment("expB") {
+ t.Errorf("expected experiments expA and expB to be present,
options: %+v", po.options)
+ }
+ if got, err := po.GetString("dataflow_service_options"); err != nil ||
got != "enable_google_cloud_profiler,enable_new_custom_feature" {
+ t.Errorf("GetString(dataflow_service_options) = (%q, %v), want
(\"enable_google_cloud_profiler,enable_new_custom_feature\", nil)", got, err)
+ }
+ serviceOpts, err := po.GetStringSlice("dataflow_service_options")
+ if err != nil || len(serviceOpts) != 2 || serviceOpts[0] !=
"enable_google_cloud_profiler" || serviceOpts[1] != "enable_new_custom_feature"
{
+ t.Errorf("GetStringSlice(dataflow_service_options) = (%v, %v),
want ([enable_google_cloud_profiler, enable_new_custom_feature], nil)",
serviceOpts, err)
+ }
+ goOpts, ok := po.options["go_options"].(map[string]any)
+ if !ok {
+ t.Errorf("expected go_options map to be present, options: %+v",
po.options)
+ }
+ nestedOpts, ok := goOpts["options"].(map[string]any)
+ if !ok {
+ t.Errorf("expected nested options map inside go_options, got:
%+v", goOpts)
+ }
+ if got := nestedOpts["dataflow_service_options"]; got !=
"enable_google_cloud_profiler,enable_new_custom_feature" {
+ t.Errorf("got dataflow_service_options = %v, want
enable_google_cloud_profiler,enable_new_custom_feature", got)
+ }
+}
diff --git a/sdks/go/pkg/beam/artifact/options.go
b/sdks/go/pkg/beam/artifact/options.go
deleted file mode 100644
index 47356433161..00000000000
--- a/sdks/go/pkg/beam/artifact/options.go
+++ /dev/null
@@ -1,48 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one or more
-// contributor license agreements. See the NOTICE file distributed with
-// this work for additional information regarding copyright ownership.
-// The ASF licenses this file to You under the Apache License, Version 2.0
-// (the "License"); you may not use this file except in compliance with
-// the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package artifact
-
-import (
- structpb "google.golang.org/protobuf/types/known/structpb"
-)
-
-// GetExperiments extracts a list of experiments from the pipeline options.
-func GetExperiments(options *structpb.Struct) []string {
- if options == nil {
- return nil
- }
-
- var exps []string
- // Try legacy style
- for _, v := range
options.GetFields()["options"].GetStructValue().GetFields()["experiments"].GetListValue().GetValues()
{
- exps = append(exps, v.GetStringValue())
- }
- // Try URN style
- for _, v := range
options.GetFields()["beam:option:experiments:v1"].GetListValue().GetValues() {
- exps = append(exps, v.GetStringValue())
- }
- return exps
-}
-
-// HasExperiment checks if a specific experiment is enabled in the pipeline
options.
-func HasExperiment(options *structpb.Struct, experiment string) bool {
- for _, exp := range GetExperiments(options) {
- if exp == experiment {
- return true
- }
- }
- return false
-}
diff --git a/sdks/go/pkg/beam/artifact/options_test.go
b/sdks/go/pkg/beam/artifact/options_test.go
deleted file mode 100644
index a9f0e4bb7e3..00000000000
--- a/sdks/go/pkg/beam/artifact/options_test.go
+++ /dev/null
@@ -1,78 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one or more
-// contributor license agreements. See the NOTICE file distributed with
-// this work for additional information regarding copyright ownership.
-// The ASF licenses this file to You under the Apache License, Version 2.0
-// (the "License"); you may not use this file except in compliance with
-// the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package artifact
-
-import (
- "testing"
-
- structpb "google.golang.org/protobuf/types/known/structpb"
-)
-
-func TestGetExperiments_Nil(t *testing.T) {
- if got := GetExperiments(nil); got != nil {
- t.Errorf("GetExperiments(nil) = %v, want nil", got)
- }
-}
-
-func TestGetExperiments_Legacy(t *testing.T) {
- options, _ := structpb.NewStruct(map[string]interface{}{
- "options": map[string]interface{}{
- "experiments": []interface{}{"exp1", "exp2"},
- },
- })
- exps := GetExperiments(options)
- if len(exps) != 2 || exps[0] != "exp1" || exps[1] != "exp2" {
- t.Errorf("GetExperiments() = %v, want [exp1 exp2]", exps)
- }
-}
-
-func TestGetExperiments_URN(t *testing.T) {
- urnOptions, _ := structpb.NewStruct(map[string]interface{}{
- "beam:option:experiments:v1": []interface{}{"expA", "expB"},
- })
- expsURN := GetExperiments(urnOptions)
- if len(expsURN) != 2 || expsURN[0] != "expA" || expsURN[1] != "expB" {
- t.Errorf("GetExperiments() = %v, want [expA expB]", expsURN)
- }
-}
-
-func TestHasExperiment(t *testing.T) {
- options, _ := structpb.NewStruct(map[string]interface{}{
- "options": map[string]interface{}{
- "experiments": []interface{}{"exp1", "exp2"},
- },
- })
-
- if !HasExperiment(options, "exp1") {
- t.Errorf("HasExperiment(exp1) = false, want true")
- }
- if HasExperiment(options, "exp3") {
- t.Errorf("HasExperiment(exp3) = true, want false")
- }
-}
-
-func TestGetExperiments_Combined(t *testing.T) {
- options, _ := structpb.NewStruct(map[string]interface{}{
- "options": map[string]interface{}{
- "experiments": []interface{}{"exp1", "exp2"},
- },
- "beam:option:experiments:v1": []interface{}{"expA", "expB"},
- })
- exps := GetExperiments(options)
- if len(exps) != 4 || exps[0] != "exp1" || exps[1] != "exp2" || exps[2]
!= "expA" || exps[3] != "expB" {
- t.Errorf("GetExperiments() = %v, want [exp1 exp2 expA expB]",
exps)
- }
-}
diff --git a/sdks/java/container/boot.go b/sdks/java/container/boot.go
index ad29f8d940a..6322c3004e5 100644
--- a/sdks/java/container/boot.go
+++ b/sdks/java/container/boot.go
@@ -105,8 +105,10 @@ func main() {
logger.Fatalf(ctx, "Failed to convert pipeline options: %v",
err)
}
+ po := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "")
+
// 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 user jars. We ignore any disk limit,
// because the staged jars are mandatory.
diff --git a/sdks/python/apache_beam/options/pipeline_options.py
b/sdks/python/apache_beam/options/pipeline_options.py
index ee7e14f3de2..2ab470e8beb 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -612,6 +612,7 @@ class PipelineOptions(HasDisplayData):
def from_runner_api(cls, proto_options, original_options=None):
def from_urn(key):
assert key.startswith('beam:option:')
+ # Update sdks/go/container/tools/pipeline_options.go if :v1 part changes.
assert key.endswith(':v1')
return key[12:-3]
diff --git a/sdks/python/container/boot.go b/sdks/python/container/boot.go
index 5a8d6da46ab..1e912ebebb7 100644
--- a/sdks/python/container/boot.go
+++ b/sdks/python/container/boot.go
@@ -132,32 +132,6 @@ func main() {
// ],
// }
// }
-type PipelineOptionsData struct {
- Options OptionsData `json:"options"`
-}
-
-type OptionsData struct {
- Experiments []string `json:"experiments"`
- ProfilerAgent string `json:"profiler_agent"`
- ProfilerExtraArgs []string `json:"profiler_extra_args"`
- ProfilerExtraEnvVars []string `json:"profiler_extra_env_vars"`
- ProfileLocation string `json:"profile_location"`
- ProfileTempLocation string `json:"profile_temp_location"`
- ProfileUploadIntervalSec int
`json:"profile_upload_interval_sec"`
- ProfilerStopAfterSec int `json:"profiler_stop_after_sec"`
- ProfilerStopAfterCrash bool
`json:"profiler_stop_after_crash"`
- ProfilePostprocessIntervalSec int
`json:"profile_postprocess_interval_sec"`
- JobId string `json:"jobId,omitempty"`
-}
-
-func getExperiments(options string) []string {
- var opts PipelineOptionsData
- err := json.Unmarshal([]byte(options), &opts)
- if err != nil {
- return nil
- }
- return opts.Options.Experiments
-}
func launchSDKProcess() error {
ctx := grpcx.WriteWorkerID(context.Background(), *id)
@@ -193,31 +167,20 @@ func launchSDKProcess() error {
// (1) Obtain the pipeline options
- options, err := tools.ProtoToJSON(info.GetPipelineOptions())
- if err != nil {
- logger.Fatalf(ctx, "Failed to convert pipeline options: %v",
err)
- }
+ po := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "")
+ logger.Printf(ctx, "Parsed options in boot entrypoint: %v", po)
// 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"))
- experiments := getExperiments(options)
- logger.Printf(ctx, "Experiments=%v", experiments)
-
- pipNoBuildIsolation = true
- if slices.Contains(experiments, "pip_use_build_isolation") {
- pipNoBuildIsolation = false
- logger.Printf(ctx, "Build isolation enabled when installing
packages with pip")
- } else {
+ pipNoBuildIsolation = !po.HasExperiment("pip_use_build_isolation")
+ if pipNoBuildIsolation {
logger.Printf(ctx, "Build isolation disabled when installing
packages with pip")
+ } else {
+ logger.Printf(ctx, "Build isolation enabled when installing
packages with pip")
}
- var opts PipelineOptionsData
- if err := json.Unmarshal([]byte(options), &opts); err != nil {
- logger.Warnf(ctx, "Failed to unmarshal pipeline options for
profiling config: %v", err)
- }
-
- ctx = setupProfilerConfig(ctx, logger, &opts)
+ ctx = setupProfilerConfig(ctx, logger, po)
startProfilerBackgroundTasks(ctx, logger)
// (2) Retrieve and install the staged packages.
@@ -286,6 +249,10 @@ func launchSDKProcess() error {
// (3) Invoke python
// Write the JSON string of pipeline options into a file to prevent
"argument list too long" error.
+ options, err := tools.ProtoToJSON(info.GetPipelineOptions())
+ if err != nil {
+ logger.Fatalf(ctx, "Failed to convert pipeline options: %v",
err)
+ }
if err := tools.MakePipelineOptionsFileAndEnvVar(options); err != nil {
logger.Fatalf(ctx, "Failed to load pipeline options to worker:
%v", err)
}
diff --git a/sdks/python/container/profiler.go
b/sdks/python/container/profiler.go
index d19923f912c..ff88fc9e64d 100644
--- a/sdks/python/container/profiler.go
+++ b/sdks/python/container/profiler.go
@@ -60,19 +60,25 @@ type ProfilerConfig struct {
GcloudAvailable bool
}
-// setupProfilerConfig parses PipelineOptionsData and stores a resolved
ProfilerConfig in the context.
-func setupProfilerConfig(ctx context.Context, logger *tools.Logger, opts
*PipelineOptionsData) context.Context {
- agent := opts.Options.ProfilerAgent
- if agent == "" {
+// setupProfilerConfig parses PipelineOptions and stores a resolved
ProfilerConfig in the context.
+func setupProfilerConfig(ctx context.Context, logger *tools.Logger, po
*tools.PipelineOptions) context.Context {
+ agent, err := po.GetString("profiler_agent")
+ if err != nil || agent == "" {
return ctx
}
- baseTempDir := opts.Options.ProfileTempLocation
- if baseTempDir == "" {
+ baseTempDir, err := po.GetString("profile_temp_location")
+ if err != nil || baseTempDir == "" {
baseTempDir = filepath.Join(*semiPersistDir, "profiles")
}
- jobId := opts.Options.JobId
+ jobId, err := po.GetString("jobId")
+ if err != nil || jobId == "" {
+ jobId = os.Getenv("JOB_ID")
+ }
+ if jobId == "" {
+ jobId = os.Getenv("JOB_NAME")
+ }
if jobId == "" {
jobId = "BEAM_JOB"
}
@@ -86,8 +92,11 @@ func setupProfilerConfig(ctx context.Context, logger
*tools.Logger, opts *Pipeli
var gcsDestPath string
gcloudAvailable := false
- if strings.HasPrefix(opts.Options.ProfileLocation, "gs://") {
- gcsDestPath = strings.TrimSuffix(opts.Options.ProfileLocation,
"/")
+ profileLocation, err := po.GetString("profile_location")
+ if err != nil || profileLocation == "" {
+ logger.Printf(ctx, "profile_location not specified, profiles
will only be stored locally.")
+ } else if strings.HasPrefix(profileLocation, "gs://") {
+ gcsDestPath = strings.TrimSuffix(profileLocation, "/")
if _, err := exec.LookPath("gcloud"); err == nil {
gcloudAvailable = true
} else {
@@ -95,20 +104,48 @@ func setupProfilerConfig(ctx context.Context, logger
*tools.Logger, opts *Pipeli
}
}
+ profilerExtraArgs, err := po.GetStringSlice("profiler_extra_args")
+ if err != nil {
+ profilerExtraArgs = []string{}
+ }
+ profilerExtraEnvVars, err :=
po.GetStringSlice("profiler_extra_env_vars")
+ if err != nil {
+ profilerExtraEnvVars = []string{}
+ }
+
+ profileUploadIntervalSec, err :=
po.GetInt("profile_upload_interval_sec")
+ if err != nil {
+ profileUploadIntervalSec = 300
+ logger.Printf(ctx, "Using default profile_upload_interval_sec:
%v", profileUploadIntervalSec)
+ }
+ profilerStopAfterSec, err := po.GetInt("profiler_stop_after_sec")
+ if err != nil {
+ profilerStopAfterSec = 0
+ }
+ profilerStopAfterCrash, err := po.GetBool("profiler_stop_after_crash")
+ if err != nil {
+ profilerStopAfterCrash = false
+ }
+ profilePostprocessIntervalSec, err :=
po.GetInt("profile_postprocess_interval_sec")
+ if err != nil {
+ profilePostprocessIntervalSec = 600
+ logger.Printf(ctx, "Using default
profile_postprocess_interval_sec: %v", profilePostprocessIntervalSec)
+ }
+
config := &ProfilerConfig{
Enabled: true,
Agent: agent,
- ExtraArgs: opts.Options.ProfilerExtraArgs,
- ExtraEnvVars: opts.Options.ProfilerExtraEnvVars,
- Location: opts.Options.ProfileLocation,
+ ExtraArgs: profilerExtraArgs,
+ ExtraEnvVars: profilerExtraEnvVars,
+ Location: profileLocation,
BaseTempDir: baseTempDir,
TempLocation: tempLocation,
StopSentinelPath: sentinelPath,
GcsDestPath: gcsDestPath,
- UploadIntervalSec: opts.Options.ProfileUploadIntervalSec,
- StopAfterSec: opts.Options.ProfilerStopAfterSec,
- StopAfterCrash: opts.Options.ProfilerStopAfterCrash,
- PostprocessIntervalSec:
opts.Options.ProfilePostprocessIntervalSec,
+ UploadIntervalSec: profileUploadIntervalSec,
+ StopAfterSec: profilerStopAfterSec,
+ StopAfterCrash: profilerStopAfterCrash,
+ PostprocessIntervalSec: profilePostprocessIntervalSec,
GcloudAvailable: gcloudAvailable,
}
diff --git a/sdks/python/container/profiler_test.go
b/sdks/python/container/profiler_test.go
index 27abf8a2ab3..15f93a2a319 100644
--- a/sdks/python/container/profiler_test.go
+++ b/sdks/python/container/profiler_test.go
@@ -20,6 +20,9 @@ import (
"os"
"path/filepath"
"testing"
+
+ "github.com/apache/beam/sdks/v2/go/container/tools"
+ "google.golang.org/protobuf/types/known/structpb"
)
func TestActivePidsRegistry(t *testing.T) {
@@ -48,13 +51,19 @@ func TestActivePidsRegistry(t *testing.T) {
}
func TestSetupProfilerConfig(t *testing.T) {
- opts := &PipelineOptionsData{
- Options: OptionsData{
- ProfilerAgent: "coredump",
- JobId: "test-job",
+ st, err := structpb.NewStruct(map[string]interface{}{
+ "options": map[string]interface{}{
+ "profiler_agent": "coredump",
+ "jobId": "test-job",
},
+ })
+ if err != nil {
+ t.Fatalf("Failed to create structpb: %v", err)
}
- ctx := setupProfilerConfig(context.Background(), nil, opts)
+
+ po := tools.ParseOptionsFromProto(st, "")
+
+ ctx := setupProfilerConfig(context.Background(), &tools.Logger{}, po)
pcfg := getProfilerConfig(ctx)
if pcfg == nil {
t.Fatal("ProfilerConfig was nil")
diff --git a/sdks/typescript/container/boot.go
b/sdks/typescript/container/boot.go
index 95e26124fac..666b101c225 100644
--- a/sdks/typescript/container/boot.go
+++ b/sdks/typescript/container/boot.go
@@ -91,8 +91,10 @@ func main() {
logger.Fatalf(ctx, "Failed to convert pipeline options: %v",
err)
}
+ po := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "")
+
// 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 and install the staged packages.
