Skip to content

Commit

Permalink
Merge remote-tracking branch 'upstream/master' into torchjit
Browse files Browse the repository at this point in the history
  • Loading branch information
AnandInguva committed Feb 10, 2023
2 parents 5637f39 + bf3710c commit fbdfe59
Show file tree
Hide file tree
Showing 70 changed files with 2,899 additions and 206 deletions.
2 changes: 1 addition & 1 deletion playground/backend/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ require (
github.com/rs/cors v1.8.2
github.com/spf13/viper v1.14.0
github.com/stretchr/testify v1.8.1
github.com/google/go-cmp v0.5.9
go.uber.org/goleak v1.2.0
google.golang.org/grpc v1.51.0
google.golang.org/protobuf v1.28.1
Expand All @@ -49,7 +50,6 @@ require (
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.2 // indirect
github.com/golang/snappy v0.0.4 // indirect
github.com/google/go-cmp v0.5.9 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.2.0 // indirect
github.com/googleapis/gax-go/v2 v2.6.0 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"context"
"encoding/json"
"fmt"
"github.com/google/go-cmp/cmp"
"io/fs"
"os"
"os/exec"
Expand Down Expand Up @@ -626,10 +627,18 @@ func Test_getRunOrTestCmd(t *testing.T) {
want: wantTestExec,
},
}

execComparer := cmp.Comparer(func(a exec.Cmd, b exec.Cmd) bool {
return a.Path == b.Path &&
cmp.Equal(a.Args, b.Args) &&
cmp.Equal(a.Env, b.Env) &&
a.Dir == b.Dir
})

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := getExecuteCmd(tt.args.isUnitTest, tt.args.executor, tt.args.ctxWithTimeout); !reflect.DeepEqual(got, tt.want) {
t.Errorf("getExecuteCmd() = %v, want %v", got, tt.want)
if got := getExecuteCmd(tt.args.isUnitTest, tt.args.executor, tt.args.ctxWithTimeout); !cmp.Equal(got, tt.want, execComparer) {
t.Errorf("getExecuteCmd() = '%v', want '%v', diff = %v", got, tt.want, cmp.Diff(got, tt.want, execComparer))
}
})
}
Expand Down
3 changes: 2 additions & 1 deletion playground/infrastructure/grpc_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ class GRPCClient:

def __init__(self, wait_for_ready=True):
use_webgrpc = os.getenv("BEAM_USE_WEBGRPC", False)
timeout = os.getenv("GRPC_TIMEOUT", 10)
timeout = int(os.getenv("GRPC_TIMEOUT", 30))
logging.info("grpc timeout: %d", timeout)
if use_webgrpc:
self._channel = sonora.aio.insecure_web_channel(Config.SERVER_ADDRESS)
else:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1465,7 +1465,7 @@ public DataflowPipelineJob run(Pipeline pipeline) {
return dataflowPipelineJob;
}

private static String getContainerImageFromEnvironmentId(
private static EnvironmentInfo getEnvironmentInfoFromEnvironmentId(
String environmentId, RunnerApi.Pipeline pipelineProto) {
RunnerApi.Environment environment =
pipelineProto.getComponents().getEnvironmentsMap().get(environmentId);
Expand All @@ -1481,30 +1481,31 @@ private static String getContainerImageFromEnvironmentId(
} catch (InvalidProtocolBufferException e) {
throw new RuntimeException("Error parsing docker payload.", e);
}
return dockerPayload.getContainerImage();
return EnvironmentInfo.create(
environmentId, dockerPayload.getContainerImage(), environment.getCapabilitiesList());
}

@AutoValue
abstract static class EnvironmentInfo {
static EnvironmentInfo create(String environmentId, String containerUrl) {
return new AutoValue_DataflowRunner_EnvironmentInfo(environmentId, containerUrl);
static EnvironmentInfo create(
String environmentId, String containerUrl, List<String> capabilities) {
return new AutoValue_DataflowRunner_EnvironmentInfo(
environmentId, containerUrl, capabilities);
}

abstract String environmentId();

abstract String containerUrl();

abstract List<String> capabilities();
}

private static List<EnvironmentInfo> getAllEnvironmentInfo(RunnerApi.Pipeline pipelineProto) {
return pipelineProto.getComponents().getTransformsMap().values().stream()
.map(transform -> transform.getEnvironmentId())
.filter(environmentId -> !environmentId.isEmpty())
.distinct()
.map(
environmentId ->
EnvironmentInfo.create(
environmentId,
getContainerImageFromEnvironmentId(environmentId, pipelineProto)))
.map(environmentId -> getEnvironmentInfoFromEnvironmentId(environmentId, pipelineProto))
.collect(Collectors.toList());
}

Expand All @@ -1520,6 +1521,7 @@ static void configureSdkHarnessContainerImages(
if (environmentInfo.containerUrl().toLowerCase().contains("python")) {
image.setUseSingleCorePerContainer(true);
}
image.setCapabilities(environmentInfo.capabilities());
return image;
})
.collect(Collectors.toList());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1683,7 +1683,9 @@ private void verifySdkHarnessConfiguration(DataflowPipelineOptions options) thro
SdkComponents sdkComponents = SdkComponents.create();
RunnerApi.Environment defaultEnvironmentForDataflow =
Environments.createDockerEnvironment(defaultSdkContainerImage);
sdkComponents.registerEnvironment(defaultEnvironmentForDataflow.toBuilder().build());
RunnerApi.Environment.Builder envBuilder =
defaultEnvironmentForDataflow.toBuilder().addCapabilities("my_dummy_capability");
sdkComponents.registerEnvironment(envBuilder.build());

RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);

Expand Down Expand Up @@ -1729,6 +1731,7 @@ private void verifySdkHarnessConfiguration(DataflowPipelineOptions options) thro
Collectors.toMap(
SdkHarnessContainerImage::getEnvironmentId,
SdkHarnessContainerImage::getContainerImage)));
assertTrue(sdks.get(0).getCapabilities().contains("my_dummy_capability"));
}

@Test
Expand Down
19 changes: 8 additions & 11 deletions sdks/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,11 @@ go 1.18

require (
cloud.google.com/go/bigquery v1.45.0
cloud.google.com/go/bigtable v1.18.1
cloud.google.com/go/datastore v1.10.0
cloud.google.com/go/profiler v0.3.1
cloud.google.com/go/pubsub v1.28.0
cloud.google.com/go/spanner v1.43.0
cloud.google.com/go/storage v1.29.0
github.com/aws/aws-sdk-go-v2 v1.17.3
github.com/aws/aws-sdk-go-v2/config v1.18.11
Expand All @@ -46,27 +48,23 @@ require (
github.com/proullon/ramsql v0.0.0-20211120092837-c8d0a408b939
github.com/spf13/cobra v1.6.1
github.com/testcontainers/testcontainers-go v0.15.0
github.com/tetratelabs/wazero v1.0.0-pre.7
github.com/xitongsys/parquet-go v1.6.2
github.com/xitongsys/parquet-go-source v0.0.0-20220315005136-aec0fe3e777c
go.mongodb.org/mongo-driver v1.11.1
golang.org/x/exp v0.0.0-20230206171751-46f607a40771
golang.org/x/net v0.5.0
golang.org/x/oauth2 v0.0.0-20221014153046-6fdb5e3db783
golang.org/x/sync v0.1.0
golang.org/x/sys v0.4.0
golang.org/x/text v0.6.0
google.golang.org/api v0.108.0
google.golang.org/api v0.109.0
google.golang.org/genproto v0.0.0-20230110181048-76db0878b65f
google.golang.org/grpc v1.52.3
google.golang.org/protobuf v1.28.1
gopkg.in/retry.v1 v1.0.3
gopkg.in/yaml.v2 v2.4.0
)

require cloud.google.com/go/spanner v1.43.0

require (
cloud.google.com/go/bigtable v1.18.1
github.com/tetratelabs/wazero v1.0.0-pre.7
gopkg.in/yaml.v3 v3.0.1
)

require (
Expand Down Expand Up @@ -136,9 +134,8 @@ require (
github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d // indirect
go.opencensus.io v0.24.0 // indirect
golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d // indirect
golang.org/x/tools v0.1.12 // indirect
golang.org/x/tools v0.2.0 // indirect
golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
google.golang.org/appengine v1.6.7 // indirect
gopkg.in/linkedin/goavro.v1 v1.0.5 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
)
12 changes: 7 additions & 5 deletions sdks/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -911,6 +911,8 @@ golang.org/x/exp v0.0.0-20191227195350-da58074b4299/go.mod h1:2RIsYlXP63K8oxa1u0
golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u096TMicItID8zy7Y6sNkU49FU4=
golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM=
golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU=
golang.org/x/exp v0.0.0-20230206171751-46f607a40771 h1:xP7rWLUr1e1n2xkK5YB4LI0hPEy3LJC6Wk+D4pGlOJg=
golang.org/x/exp v0.0.0-20230206171751-46f607a40771/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc=
golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
Expand Down Expand Up @@ -1137,8 +1139,8 @@ golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roY
golang.org/x/tools v0.0.0-20200916195026-c9a70fc28ce3/go.mod h1:z6u4i615ZeAfBE4XtMziQW1fSVJXACjjbWkB/mvPzlU=
golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
golang.org/x/tools v0.1.12 h1:VveCTK38A2rkS8ZqFY25HIDFscX5X9OoEhJd3quQmXU=
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
golang.org/x/tools v0.2.0 h1:G6AHpWxTMGY1KyEYoAQ5WTtIekUUvDNjan3ugu60JvE=
golang.org/x/tools v0.2.0/go.mod h1:y4OqIKeOV/fWJetJ8bXPU1sEVniLMIyDAZWeHdV+NTA=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
Expand All @@ -1157,8 +1159,8 @@ google.golang.org/api v0.17.0/go.mod h1:BwFmGc8tA3vsd7r/7kR8DY7iEEGSU04BFxCo5jP/
google.golang.org/api v0.18.0/go.mod h1:BwFmGc8tA3vsd7r/7kR8DY7iEEGSU04BFxCo5jP/sfE=
google.golang.org/api v0.20.0/go.mod h1:BwFmGc8tA3vsd7r/7kR8DY7iEEGSU04BFxCo5jP/sfE=
google.golang.org/api v0.22.0/go.mod h1:BwFmGc8tA3vsd7r/7kR8DY7iEEGSU04BFxCo5jP/sfE=
google.golang.org/api v0.108.0 h1:WVBc/faN0DkKtR43Q/7+tPny9ZoLZdIiAyG5Q9vFClg=
google.golang.org/api v0.108.0/go.mod h1:2Ts0XTHNVWxypznxWOYUeI4g3WdP9Pk2Qk58+a/O9MY=
google.golang.org/api v0.109.0 h1:sW9hgHyX497PP5//NUM7nqfV8D0iDfBApqq7sOh1XR8=
google.golang.org/api v0.109.0/go.mod h1:2Ts0XTHNVWxypznxWOYUeI4g3WdP9Pk2Qk58+a/O9MY=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
google.golang.org/appengine v1.5.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
Expand Down Expand Up @@ -1324,4 +1326,4 @@ sigs.k8s.io/structured-merge-diff/v4 v4.0.1/go.mod h1:bJZC9H9iH24zzfZ/41RGcq60oK
sigs.k8s.io/structured-merge-diff/v4 v4.0.2/go.mod h1:bJZC9H9iH24zzfZ/41RGcq60oK1F7G282QMXDPYydCw=
sigs.k8s.io/structured-merge-diff/v4 v4.0.3/go.mod h1:bJZC9H9iH24zzfZ/41RGcq60oK1F7G282QMXDPYydCw=
sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=
sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc=
sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc=
Loading

0 comments on commit fbdfe59

Please sign in to comment.