diff --git a/CHANGES.md b/CHANGES.md index bda50ac6cd18..091d1142d24a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -105,6 +105,7 @@ * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). * Improved Java pipeline performance by avoiding repeated `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). * (Python) Fixed a memory leak in Python SDK caused by storing exceptions with potentially large stack frames in a cache ([#39406](https://github.com/apache/beam/issues/39406)). +* (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). ## Security Fixes diff --git a/sdks/go/container/boot.go b/sdks/go/container/boot.go index 469285821f7e..15b1ac6a1607 100644 --- a/sdks/go/container/boot.go +++ b/sdks/go/container/boot.go @@ -159,8 +159,14 @@ 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, err := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "go_options") + if err != nil { + logger.Fatalf(ctx, "Failed to parse pipeline options: %v", err) + } + // 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. // diff --git a/sdks/go/container/tools/pipeline_options.go b/sdks/go/container/tools/pipeline_options.go index 026fb31b0991..087bd0e5b3b8 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,181 @@ 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, error) { + if opt == nil { + return &PipelineOptions{options: make(map[string]any), experiments: make(map[string]string)}, nil + } + 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 { + if expsMap, err := parseExperiments(exps); err == nil { + po.experiments = expsMap + } + } + return po, nil +} + +func parseExperiments(slice []string) (map[string]string, error) { + 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, nil +} + +// 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. +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 + } + return "", fmt.Errorf("option %q: expected string, got type %T", name, val) +} + +// GetStringSlice returns the value of an option as a string slice. +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 flags rather than JSON arrays. + 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) + } +} + +// 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 7a0d7ebd5f09..df9d03c3ca8e 100644 --- a/sdks/go/container/tools/pipeline_options_test.go +++ b/sdks/go/container/tools/pipeline_options_test.go @@ -16,10 +16,32 @@ package tools import ( + "encoding/json" "os" "testing" + + structpb "google.golang.org/protobuf/types/known/structpb" ) +func parseOptionsForTest(t *testing.T, options string) *PipelineOptions { + 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) + } + po, err := ParseOptionsFromProto(st, "go_options") + if err != nil { + t.Fatalf("ParseOptionsFromProto failed: %v", err) + } + return po +} + func TestMakePipelineOptionsFileAndEnvVar(t *testing.T) { tests := []struct { name string @@ -56,3 +78,217 @@ func TestMakePipelineOptionsFileAndEnvVar(t *testing.T) { } os.Remove("pipeline_options.json") } + +func TestPipelineOptions(t *testing.T) { + tests := []struct { + name string + inputOptions string + validate func(t *testing.T, po *PipelineOptions) + }{ + { + "nested options", + `{ + "options": { + "profiler_agent": "memray", + "profile_upload_interval_sec": 10, + "profiler_stop_after_crash": true, + "experiments": ["beam_fn_api", "use_build_isolation"] + } + }`, + func(t *testing.T, po *PipelineOptions) { + 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) + } + experiments, err := po.GetStringSlice("experiments") + if err != nil || len(experiments) != 2 || experiments[0] != "beam_fn_api" || experiments[1] != "use_build_isolation" { + t.Errorf("GetStringSlice(experiments) = (%v, %v), want ([beam_fn_api, use_build_isolation], nil)", experiments, err) + } + }, + }, + { + "flat URN options with string values", + `{ + "beam:option:profiler_agent:v1": "memray", + "beam:option:profile_upload_interval_sec:v1": "10", + "beam:option:profiler_stop_after_crash:v1": "true", + "beam:option:experiments:v1": ["beam_fn_api"] + }`, + func(t *testing.T, po *PipelineOptions) { + 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) + } + experiments, err := po.GetStringSlice("experiments") + if err != nil || len(experiments) != 1 || experiments[0] != "beam_fn_api" { + t.Errorf("GetStringSlice(experiments) = (%v, %v), want ([beam_fn_api], nil)", experiments, err) + } + }, + }, + { + "comma-separated string options parsed as slice (Go SDK style)", + `{ + "options": { + "experiments": "exp1,exp2,exp3", + "dataflow_service_options": "opt1" + } + }`, + func(t *testing.T, po *PipelineOptions) { + 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) + } + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + po := parseOptionsForTest(t, test.inputOptions) + test.validate(t, po) + }) + } +} + +func TestPipelineOptions_Errors(t *testing.T) { + t.Run("malformed integer", func(t *testing.T) { + po := parseOptionsForTest(t, `{"options": {"profile_upload_interval_sec": "invalid"}}`) + _, err := po.GetInt("profile_upload_interval_sec") + if err == nil { + t.Errorf("expected error, got nil") + } + }) + + t.Run("malformed bool", func(t *testing.T) { + po := parseOptionsForTest(t, `{"options": {"profiler_stop_after_crash": "maybe"}}`) + _, 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) { + po := parseOptionsForTest(t, `{"options": {"profile_upload_interval_sec": true}}`) + _, 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) { + po := parseOptionsForTest(t, `{}`) + _, errInt := po.GetInt("profile_upload_interval_sec") + _, errBool := po.GetBool("profiler_stop_after_crash") + if errInt == nil || errBool == nil { + t.Errorf("expected error for missing keys, got: intErr=%v, boolErr=%v", errInt, errBool) + } + }) + + t.Run("HasOption", func(t *testing.T) { + po := parseOptionsForTest(t, `{"options": {"profile_upload_interval_sec": 10}}`) + 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") + } + }) + + t.Run("HasExperiment", func(t *testing.T) { + po := parseOptionsForTest(t, `{"options": {"experiments": ["exp1", "exp2=val2"]}}`) + 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") + } + }) + + t.Run("LookupExperiment", func(t *testing.T) { + po := parseOptionsForTest(t, `{"options": {"experiments": ["exp1", "exp2=val2", "exp3=val3=val4"]}}`) + + 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) + } + }) + + + + t.Run("ParseOptionsFromProto", func(t *testing.T) { + optionsStruct, err := structpb.NewStruct(map[string]interface{}{ + "options": map[string]interface{}{ + "profiler_agent": "memray", + }, + "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=custom_profiler", + }, + }, + }) + if err != nil { + t.Fatalf("failed to create proto Struct: %v", err) + } + + po, err := ParseOptionsFromProto(optionsStruct, "go_options") + if err != nil { + t.Fatalf("ParseOptionsFromProto failed: %v", err) + } + if got, err := po.GetString("profiler_agent"); err != nil || got != "memray" { + t.Errorf("GetString(profiler_agent) = (%q, %v), want (\"memray\", 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=custom_profiler" { + t.Errorf("GetString(dataflow_service_options) = (%q, %v), want (\"enable_google_cloud_profiler=custom_profiler\", nil)", got, 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=custom_profiler" { + t.Errorf("got dataflow_service_options = %v, want enable_google_cloud_profiler=custom_profiler", got) + } + }) +} diff --git a/sdks/go/pkg/beam/artifact/options.go b/sdks/go/pkg/beam/artifact/options.go deleted file mode 100644 index 47356433161c..000000000000 --- 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 a9f0e4bb7e35..000000000000 --- 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 ad29f8d940ac..3f2b4b7dd8a5 100644 --- a/sdks/java/container/boot.go +++ b/sdks/java/container/boot.go @@ -105,8 +105,13 @@ func main() { logger.Fatalf(ctx, "Failed to convert pipeline options: %v", err) } + po, err := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "") + if err != nil { + logger.Fatalf(ctx, "Failed to parse pipeline options: %v", err) + } + // 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/container/boot.go b/sdks/python/container/boot.go index 5a8d6da46ab9..2dfcab2ec3d4 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,23 @@ func launchSDKProcess() error { // (1) Obtain the pipeline options - options, err := tools.ProtoToJSON(info.GetPipelineOptions()) + po, err := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "") if err != nil { - logger.Fatalf(ctx, "Failed to convert pipeline options: %v", err) + logger.Fatalf(ctx, "Failed to parse pipeline options: %v", err) } + 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")) - - experiments := getExperiments(options) - logger.Printf(ctx, "Experiments=%v", experiments) + ctx = artifact.WithArtifactValidation(ctx, !po.HasExperiment("disable_staged_file_integrity_checks")) - 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 +252,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 d19923f912c1..ff88fc9e64d0 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 27abf8a2ab3b..401c7d1bbec4 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,22 @@ 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, err := tools.ParseOptionsFromProto(st, "") + if err != nil { + t.Fatalf("Failed to parse pipeline options: %v", err) + } + + 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 95e26124facc..41708aa3c77c 100644 --- a/sdks/typescript/container/boot.go +++ b/sdks/typescript/container/boot.go @@ -91,8 +91,13 @@ func main() { logger.Fatalf(ctx, "Failed to convert pipeline options: %v", err) } + po, err := tools.ParseOptionsFromProto(info.GetPipelineOptions(), "") + if err != nil { + logger.Fatalf(ctx, "Failed to parse pipeline options: %v", err) + } + // 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.