-
Notifications
You must be signed in to change notification settings - Fork 4.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Browse files
Browse the repository at this point in the history
* [#23893] Do Annotation plumbing to the graph. * fmt * Groundwork for DisplayData and Deps support * rm env change * Add plubming and simpler handling and test! * Docs and details. * Missed a rename. * pipeline doc --------- Co-authored-by: lostluck <13907733+lostluck@users.noreply.github.com>
- Loading branch information
Showing
7 changed files
with
341 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,120 @@ | ||
| // 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 contextreg contains the global registrations of functions for extracting | ||
| // ptransform annotations or environment resource hints from context.Context attached to | ||
| // scopes. | ||
| // | ||
| // For beam internal use only. API subject to change. | ||
| package contextreg | ||
|
|
||
| import ( | ||
| "context" | ||
| "maps" | ||
| "sync" | ||
| ) | ||
|
|
||
| var defaultReg = &Registry{} | ||
|
|
||
| // Default is the default registry for context extractors. | ||
| func Default() *Registry { | ||
| return defaultReg | ||
| } | ||
|
|
||
| // Registry contains a set of registrations for extracting annotations and hints from a context.Context. | ||
| // | ||
| // This type is exported to allow simpler testing of new extractors, and their interaction with the registry. | ||
| type Registry struct { | ||
| mu sync.Mutex | ||
| transforms []func(context.Context) TransformMetadata | ||
| envs []func(context.Context) EnvironmentMetadata | ||
| } | ||
|
|
||
| // TransformMetadata represents additional information on transforms to be added to the Pipeline proto graph. | ||
| type TransformMetadata struct { | ||
| Annotations map[string][]byte | ||
| // DisplayData []*pipepb.DisplayData | ||
| } | ||
|
|
||
| // EnvironmentMetadata represent additional information on environmental requirements to be added to the Pipeline | ||
| // proto graph. | ||
| type EnvironmentMetadata struct { | ||
| ResourceHints map[string][]byte | ||
| // DisplayData []*pipepb.DisplayData | ||
| // Dependencies []*pipepb.ArtifactInformation | ||
| } | ||
|
|
||
| // TransformExtractor registers a transform metadata extractor to this registry. | ||
| // These will be set on the current composite transform scope. | ||
| // They are accessible to runners via the transform hypergraph. | ||
| func (r *Registry) TransformExtractor(ext func(context.Context) TransformMetadata) { | ||
| r.mu.Lock() | ||
| r.transforms = append(r.transforms, ext) | ||
| r.mu.Unlock() | ||
| } | ||
|
|
||
| // EnvExtrator registers an environment metadata extractor to this registry. | ||
| // When non-empty extraction occurs, a new environment will be derived from the parent scopes environment. | ||
| func (r *Registry) EnvExtrator(ext func(context.Context) EnvironmentMetadata) { | ||
| r.mu.Lock() | ||
| r.envs = append(r.envs, ext) | ||
| r.mu.Unlock() | ||
| } | ||
|
|
||
| // ExtractTransformMetadata runs all registered transform extractors on the provided context, | ||
| // and returns the resulting metadata. | ||
| // | ||
| // A metadata field will be nil if there's no data. A nil context bypasses extractor execution. | ||
| func (r *Registry) ExtractTransformMetadata(ctx context.Context) TransformMetadata { | ||
| r.mu.Lock() | ||
| defer r.mu.Unlock() | ||
| if ctx == nil { | ||
| return TransformMetadata{} | ||
| } | ||
| ret := TransformMetadata{ | ||
| Annotations: map[string][]byte{}, | ||
| } | ||
| for _, ext := range r.transforms { | ||
| k := ext(ctx) | ||
| maps.Copy(ret.Annotations, k.Annotations) | ||
| } | ||
| if len(ret.Annotations) == 0 { | ||
| ret.Annotations = nil | ||
| } | ||
| return ret | ||
| } | ||
|
|
||
| // ExtractEnvironmentMetadata runs all registered environment extractors on the provided context, | ||
| // and returns the resulting metadata. | ||
| // | ||
| // A metadata field will be nil if there's no data. A nil context bypasses extractor execution. | ||
| func (r *Registry) ExtractEnvironmentMetadata(ctx context.Context) EnvironmentMetadata { | ||
| r.mu.Lock() | ||
| defer r.mu.Unlock() | ||
| if ctx == nil { | ||
| return EnvironmentMetadata{} | ||
| } | ||
| ret := EnvironmentMetadata{ | ||
| ResourceHints: map[string][]byte{}, | ||
| } | ||
| for _, ext := range r.envs { | ||
| k := ext(ctx) | ||
| maps.Copy(ret.ResourceHints, k.ResourceHints) | ||
| } | ||
| if len(ret.ResourceHints) == 0 { | ||
| ret.ResourceHints = nil | ||
| } | ||
| return ret | ||
| } |
108 changes: 108 additions & 0 deletions
108
sdks/go/pkg/beam/core/runtime/contextreg/contextreg_test.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,108 @@ | ||
| // 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 contextreg | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
| ) | ||
|
|
||
| func TestPTransformExtractor(t *testing.T) { | ||
| reg := &Registry{} | ||
|
|
||
| type keyType string | ||
| key1 := keyType("annotation1") | ||
| key2 := keyType("annotation2") | ||
| key3 := keyType("annotation3") | ||
|
|
||
| reg.TransformExtractor(func(ctx context.Context) TransformMetadata { | ||
| v := ctx.Value(key1).(string) | ||
| return TransformMetadata{ | ||
| Annotations: map[string][]byte{ | ||
| "beam:test:annotation": []byte(v), | ||
| }, | ||
| } | ||
| }) | ||
| reg.TransformExtractor(func(ctx context.Context) TransformMetadata { | ||
| v := ctx.Value(key2).(string) | ||
| return TransformMetadata{ | ||
| Annotations: map[string][]byte{ | ||
| "beam:test:annotation2": []byte(v), | ||
| }, | ||
| } | ||
| }) | ||
| // Override the extaction for result annotation to use the last set version. | ||
| reg.TransformExtractor(func(ctx context.Context) TransformMetadata { | ||
| v := ctx.Value(key3).(string) | ||
| return TransformMetadata{ | ||
| Annotations: map[string][]byte{ | ||
| "beam:test:annotation": []byte(v), | ||
| }, | ||
| } | ||
| }) | ||
|
|
||
| ctx := context.Background() | ||
| // Set all 3 distinct context values. | ||
| ctx = context.WithValue(ctx, key1, "never seen") | ||
| want2 := "want_value2" | ||
| ctx = context.WithValue(ctx, key2, want2) | ||
| want3 := "want_value3" | ||
| ctx = context.WithValue(ctx, key3, want3) | ||
|
|
||
| ptrans := reg.ExtractTransformMetadata(ctx) | ||
|
|
||
| key := "beam:test:annotation" | ||
| if got, want := string(ptrans.Annotations[key]), want3; got != want { | ||
| t.Errorf("extracted annotation %q = %q, want %q", key, got, want) | ||
| } | ||
| key = "beam:test:annotation2" | ||
| if got, want := string(ptrans.Annotations[key]), want2; got != want { | ||
| t.Errorf("extracted annotation %q = %q, want %q", key, got, want) | ||
| } | ||
| if got, want := len(ptrans.Annotations), 2; got != want { | ||
| t.Errorf("extracted annotation %q = %q, want %q - have %v", key, got, want, ptrans) | ||
| } | ||
| } | ||
|
|
||
| func TestHintExtractor(t *testing.T) { | ||
| reg := &Registry{} | ||
|
|
||
| type keyType string | ||
| hintKey := keyType("hint") | ||
|
|
||
| reg.EnvExtrator(func(ctx context.Context) EnvironmentMetadata { | ||
| v := ctx.Value(hintKey).(string) | ||
| return EnvironmentMetadata{ | ||
| ResourceHints: map[string][]byte{ | ||
| "beam:test:hint": []byte(v), | ||
| }, | ||
| } | ||
| }) | ||
|
|
||
| ctx := context.Background() | ||
| wantedHint := "hint" | ||
| ctx = context.WithValue(ctx, hintKey, wantedHint) | ||
|
|
||
| env := reg.ExtractEnvironmentMetadata(ctx) | ||
|
|
||
| key := "beam:test:hint" | ||
| if got, want := string(env.ResourceHints[key]), wantedHint; got != want { | ||
| t.Errorf("extracted annotation %q = %q, want %q", key, got, want) | ||
| } | ||
| if got, want := len(env.ResourceHints), 1; got != want { | ||
| t.Errorf("extracted annotation %q = %q, want %q - have %v", key, got, want, env) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.