private static ExecutableStage createExecutableStage(Collection<SideInputReference> sideInputs) { Components components = Components.getDefaultInstance(); Environment environment = Environment.getDefaultInstance(); PCollectionNode inputCollection = PipelineNode.pCollection("collection-id", RunnerApi.PCollection.getDefaultInstance()); return ImmutableExecutableStage.of( components, environment, inputCollection, sideInputs, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); }
ParDoPayload.class.getSimpleName(), Environment.class.getSimpleName()); if (!env.get().equals(environment)) {
/** * Return a map of IDs to {@link PTransform} which are executed by an SDK Harness. * * <p>The transforms that are present in the returned map are the {@link RunnerApi.PTransform} * versions of the {@link ExecutableStage ExecutableStages} returned in {@link #getFusedStages()}. * The IDs of the returned transforms will not collide with any transform ID present in {@link * #getComponents()}. */ private Map<String, PTransform> getEnvironmentExecutedTransforms() { Map<String, PTransform> topLevelTransforms = new HashMap<>(); for (ExecutableStage stage : getFusedStages()) { String baseName = String.format( "%s/%s", stage.getInputPCollection().getPCollection().getUniqueName(), stage.getEnvironment().getUrn()); Set<String> usedNames = Sets.union(topLevelTransforms.keySet(), getComponents().getTransformsMap().keySet()); String uniqueId = SyntheticComponents.uniqueId(baseName, usedNames::contains); topLevelTransforms.put(uniqueId, stage.toPTransform(uniqueId)); } return topLevelTransforms; } }
ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(),
ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(),
ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(),
ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(),
Environments.createOrGetDefaultEnvironment(Environments.ENVIRONMENT_DOCKER, "java"), is( Environment.newBuilder() .setUrl("java") .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.DOCKER)) "{\"os\": \"linux\", \"arch\": \"amd64\", \"command\": \"run.sh\", \"env\":{\"k1\": \"v1\", \"k2\": \"v2\"} }"), is( Environment.newBuilder() .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.PROCESS)) .setPayload( Environments.ENVIRONMENT_PROCESS, "{\"command\": \"run.sh\"}"), is( Environment.newBuilder() .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.PROCESS)) .setPayload(ProcessPayload.newBuilder().setCommand("run.sh").build().toByteString())
/** This method doesn't do any pruning for reachability, but this may not require a test. */ @Test public void retainOnlyPrimitivesIgnoresUnreachableNodes() { Pipeline p = Pipeline.create(); p.apply( new org.apache.beam.sdk.transforms.PTransform<PBegin, PCollection<Long>>() { @Override public PCollection<Long> expand(PBegin input) { return input .apply(GenerateSequence.from(2L)) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(5L)))) .apply(MapElements.into(TypeDescriptors.longs()).via(l -> l + 1)); } }); Components augmentedComponents = PipelineTranslation.toProto(p) .getComponents() .toBuilder() .putCoders("extra-coder", RunnerApi.Coder.getDefaultInstance()) .putWindowingStrategies( "extra-windowing-strategy", RunnerApi.WindowingStrategy.getDefaultInstance()) .putEnvironments("extra-env", RunnerApi.Environment.getDefaultInstance()) .putPcollections("extra-pc", RunnerApi.PCollection.getDefaultInstance()) .build(); Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(augmentedComponents); } }
private static Environment getStageEnvironment( QueryablePipeline pipeline, Set<PTransformNode> initialNodes) { Supplier<IllegalArgumentException> missingEnv = () -> new IllegalArgumentException( String.format( "%s must be populated on all %s in a %s", Environment.class.getSimpleName(), PTransformNode.class.getSimpleName(), GreedyStageFuser.class.getSimpleName())); Environment env = pipeline.getEnvironment(initialNodes.iterator().next()).orElseThrow(missingEnv); initialNodes.forEach( transformNode -> checkArgument( env.equals(pipeline.getEnvironment(transformNode).orElseThrow(missingEnv)), "All %s in a %s must be the same. Got %s and %s", Environment.class.getSimpleName(), ExecutableStage.class.getSimpleName(), env, pipeline.getEnvironment(transformNode).get() /* will throw above if absent. */)); return env; }
public static Environment createProcessEnvironment( String os, String arch, String command, Map<String, String> env) { ProcessPayload.Builder builder = ProcessPayload.newBuilder(); if (!Strings.isNullOrEmpty(os)) { builder.setOs(os); } if (!Strings.isNullOrEmpty(arch)) { builder.setArch(arch); } if (!Strings.isNullOrEmpty(command)) { builder.setCommand(command); } if (env != null) { builder.putAllEnv(env); } return Environment.newBuilder() .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.PROCESS)) .setPayload(builder.build().toByteString()) .build(); }
private static ExecutableStage createExecutableStage(Collection<SideInputReference> sideInputs) { Components components = Components.getDefaultInstance(); Environment environment = Environment.getDefaultInstance(); PCollectionNode inputCollection = PipelineNode.pCollection("collection-id", RunnerApi.PCollection.getDefaultInstance()); return ImmutableExecutableStage.of( components, environment, inputCollection, sideInputs, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); }
/** * A WindowInto can be fused into a stage if it executes in the same Environment as that stage. */ private static boolean canFuseAssignWindows( PTransformNode window, Environment environmemnt, @SuppressWarnings("unused") PCollectionNode candidate, @SuppressWarnings("unused") Collection<PCollectionNode> stagePCollections, QueryablePipeline pipeline) { // WindowInto transforms may not have an environment Optional<Environment> windowEnvironment = pipeline.getEnvironment(window); return environmemnt.equals(windowEnvironment.orElse(null)); }
/** * Registers the provided {@link Environment} into this {@link SdkComponents}, returning a unique * ID for the {@link Environment}. Multiple registrations of the same {@link Environment} will * return the same unique ID. */ public String registerEnvironment(Environment env) { String existing = environmentIds.get(env); if (existing != null) { return existing; } String name = uniqify(env.getUrn(), environmentIds.values()); environmentIds.put(env, name); componentsBuilder.putEnvironments(name, env); return name; }
private static Environment createEmbeddedEnvironment(String config) { return Environment.newBuilder() .setUrn(ENVIRONMENT_EMBEDDED) .setPayload(ByteString.copyFromUtf8(MoreObjects.firstNonNull(config, ""))) .build(); }
public static Environment createDockerEnvironment(String dockerImageUrl) { return Environment.newBuilder() .setUrl(dockerImageUrl) .setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.DOCKER)) .setPayload( DockerPayload.newBuilder().setContainerImage(dockerImageUrl).build().toByteString()) .build(); }