This is an automated email from the ASF dual-hosted git repository.
jrmccluskey 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 32d5037d336 [Go] Add parsing for cpu_count resource hint (#39977)
32d5037d336 is described below
commit 32d5037d3369d942ce3b0ce791f2e72eac4e54f3
Author: Jack McCluskey <[email protected]>
AuthorDate: Wed Sep 2 14:25:56 2026 -0400
[Go] Add parsing for cpu_count resource hint (#39977)
* Add parsing for cpu_count resource hint
* Fix formatting bug
---
sdks/go/pkg/beam/options/jobopts/options.go | 2 +
sdks/go/pkg/beam/options/jobopts/options_test.go | 4 +-
sdks/go/pkg/beam/options/resource/hint.go | 17 ++++++-
sdks/go/pkg/beam/options/resource/hint_test.go | 63 ++++++++++++++++++++++++
4 files changed, 84 insertions(+), 2 deletions(-)
diff --git a/sdks/go/pkg/beam/options/jobopts/options.go
b/sdks/go/pkg/beam/options/jobopts/options.go
index 327f3895b11..8456c8fd061 100644
--- a/sdks/go/pkg/beam/options/jobopts/options.go
+++ b/sdks/go/pkg/beam/options/jobopts/options.go
@@ -207,6 +207,8 @@ func GetPipelineResourceHints() resource.Hints {
h = resource.ParseMinRAM(val)
case "accelerator", "beam:resources:accelerator:v1":
h = resource.Accelerator(val)
+ case "cpu_count", "beam:resources:cpu_count:v1":
+ h = resource.ParseCPUCount(val)
default:
if strings.HasPrefix(name, "beam:resources:") {
h = stringHint{urn: name, value: val}
diff --git a/sdks/go/pkg/beam/options/jobopts/options_test.go
b/sdks/go/pkg/beam/options/jobopts/options_test.go
index 89e2d5a721d..5bcdf39ea59 100644
--- a/sdks/go/pkg/beam/options/jobopts/options_test.go
+++ b/sdks/go/pkg/beam/options/jobopts/options_test.go
@@ -149,9 +149,11 @@ func TestGetPipelineResourceHints(t *testing.T) {
hints.Set("accelerator=pedal_to_the_metal")
hints.Set("beam:resources:novel_execution:v1=jaguar")
hints.Set("min_ram=1GB")
+ hints.Set("cpu_count=1")
+ hints.Set("beam:resources:cpu_count:v1=4")
ResourceHints = hints
- want := resource.NewHints(resource.ParseMinRAM("1GB"),
resource.Accelerator("pedal_to_the_metal"), stringHint{
+ want := resource.NewHints(resource.ParseMinRAM("1GB"),
resource.Accelerator("pedal_to_the_metal"), resource.ParseCPUCount("4"),
stringHint{
urn: "beam:resources:novel_execution:v1",
value: "jaguar",
})
diff --git a/sdks/go/pkg/beam/options/resource/hint.go
b/sdks/go/pkg/beam/options/resource/hint.go
index 5efaca80ede..1d3f1cd37ec 100644
--- a/sdks/go/pkg/beam/options/resource/hint.go
+++ b/sdks/go/pkg/beam/options/resource/hint.go
@@ -246,6 +246,21 @@ func (h acceleratorHint) String() string {
return fmt.Sprintf("accelerator=%v", h.value)
}
+// ParseCPUCount converts a number in string form into a hint.
+// An invalid format will cause ParseCPUCount to panic.
+//
+// Hints are advisory only and runners may not respect them.
+//
+// See https://beam.apache.org/documentation/runtime/resource-hints/ for more
information about
+// resource hints.
+func ParseCPUCount(v string) Hint {
+ b, err := strconv.ParseUint(v, 10, 64)
+ if err != nil {
+ panic(fmt.Sprintf("resource.ParseCPUCount: unable to parse %q:
%v", v, err))
+ }
+ return CPUCount(uint64(b))
+}
+
// CPUCount hints that this scope should be put in a machine with at least
this many CPUs or vCPUs.
//
// Hints are advisory only and runners may not respect them.
@@ -280,5 +295,5 @@ func (h CPUCountHint) MergeWithOuter(outer Hint) Hint {
}
func (h CPUCountHint) String() string {
- return fmt.Sprintf("cpu_count=%v", humanize.Bytes(uint64(h.value)))
+ return fmt.Sprintf("cpu_count=%v", h.value)
}
diff --git a/sdks/go/pkg/beam/options/resource/hint_test.go
b/sdks/go/pkg/beam/options/resource/hint_test.go
index 2c4c6506e3a..35dc5fa0dd7 100644
--- a/sdks/go/pkg/beam/options/resource/hint_test.go
+++ b/sdks/go/pkg/beam/options/resource/hint_test.go
@@ -143,6 +143,69 @@ func TestCPUCountHint_Payload(t *testing.T) {
}
}
+func TestCPUCountHint_String(t *testing.T) {
+ tests := []struct {
+ value uint64
+ want string
+ }{
+ {0, "cpu_count=0"},
+ {1, "cpu_count=1"},
+ {4, "cpu_count=4"},
+ {128, "cpu_count=128"},
+ }
+
+ for _, test := range tests {
+ h := CPUCountHint{value: test.value}
+ if got, want := h.String(), test.want; got != want {
+ t.Errorf("%v.String() = %v, want %v", h, got, want)
+ }
+ }
+}
+
+func TestParseCPUCount(t *testing.T) {
+ tests := []struct {
+ value string
+ payload string
+ }{
+ {"0", "0"},
+ {"1", "1"},
+ {"2", "2"},
+ {"4", "4"},
+ {"11", "11"},
+ {"2003", "2003"},
+ {"12000000", "12000000"},
+ {"18446744073709551615", "18446744073709551615"},
+ }
+
+ for _, test := range tests {
+ h := ParseCPUCount(test.value)
+ if got, want := h.Payload(), []byte(test.payload);
!bytes.Equal(got, want) {
+ t.Errorf("%v.Payload() = %v, want %v", h, string(got),
string(want))
+ }
+ }
+}
+
+func TestParseCPUCount_panic(t *testing.T) {
+ tests := []string{
+ "a bad cpu string",
+ "-1",
+ "1.5",
+ "",
+ "18446744073709551616",
+ }
+
+ for _, test := range tests {
+ t.Run(test, func(t *testing.T) {
+ defer func() {
+ if r := recover(); r == nil {
+ t.Errorf("want ParseCPUCount(%q) to
panic", test)
+ }
+ }()
+ ParseCPUCount(test)
+ })
+ }
+}
+
func TestMaxActiveBundlesPerWorkerHint_MergeWith(t *testing.T) {
low := maxActiveBundlesPerWorkerHint{value: 2}
high := maxActiveBundlesPerWorkerHint{value: 4}
