Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#Environment

private RunnerApi.Pipeline resolveDependencies(RunnerApi.Pipeline pipeline, String stagingToken) {
  Map<String, List<RunnerApi.ArtifactInformation>> resolvedDependencies =
      stagingService.getService().getStagedArtifacts(stagingToken);
  Map<String, RunnerApi.Environment> newEnvironments = new HashMap<>();
  for (Map.Entry<String, RunnerApi.Environment> entry :
      pipeline.getComponents().getEnvironmentsMap().entrySet()) {
    if (entry.getValue().getDependenciesCount() > 0 && resolvedDependencies == null) {
      throw new RuntimeException(
          "Artifact dependencies provided but not staged for " + entry.getKey());
    }
    newEnvironments.put(
        entry.getKey(),
        entry.getValue().getDependenciesCount() == 0
            ? entry.getValue()
            : entry
                .getValue()
                .toBuilder()
                .clearDependencies()
                .addAllDependencies(resolvedDependencies.get(entry.getKey()))
                .build());
  }
  RunnerApi.Pipeline.Builder builder = pipeline.toBuilder();
  builder.getComponentsBuilder().clearEnvironments().putAllEnvironments(newEnvironments);
  return builder.build();
}
 

private static byte[] serializeWindowingStrategy(
    WindowingStrategy<?, ?> windowingStrategy, PipelineOptions options) {
  try {
    SdkComponents sdkComponents = SdkComponents.create();

    String workerHarnessContainerImageURL =
        DataflowRunner.getContainerImageForJob(options.as(DataflowPipelineOptions.class));
    RunnerApi.Environment defaultEnvironmentForDataflow =
        Environments.createDockerEnvironment(workerHarnessContainerImageURL);
    sdkComponents.registerEnvironment(defaultEnvironmentForDataflow);

    return WindowingStrategyTranslation.toMessageProto(windowingStrategy, sdkComponents)
        .toByteArray();
  } catch (Exception e) {
    throw new RuntimeException(
        String.format("Unable to format windowing strategy %s as bytes", windowingStrategy), e);
  }
}
 

static com.google.api.services.dataflow.model.Source translateIOToCloudSource(
    BoundedSource<?> io, DataflowPipelineOptions options) throws Exception {
  DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options);
  Pipeline p = Pipeline.create(options);
  p.begin().apply(Read.from(io));

  DataflowRunner runner = DataflowRunner.fromOptions(options);
  SdkComponents sdkComponents = SdkComponents.create();
  RunnerApi.Environment defaultEnvironmentForDataflow =
      Environments.createDockerEnvironment("dummy-image-url");
  sdkComponents.registerEnvironment(defaultEnvironmentForDataflow);
  RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);

  Job workflow =
      translator
          .translate(p, pipelineProto, sdkComponents, runner, new ArrayList<DataflowPackage>())
          .getJob();
  Step step = workflow.getSteps().get(0);

  return stepToCloudSource(step);
}
 

@Override
public RunnerApi.Environment createEnvironment() throws IOException {
	Map<String, String> env = constructEnvironmentVariables();

	if (dependencyInfo.getRequirementsFilePath().isPresent()) {
		LOG.info("Trying to pip install the Python requirements...");
		PythonEnvironmentManagerUtils.pipInstallRequirements(
			dependencyInfo.getRequirementsFilePath().get(),
			dependencyInfo.getRequirementsCacheDir().orElse(null),
			requirementsDirectory,
			dependencyInfo.getPythonExec(),
			env);
	}
	String runnerScript = PythonEnvironmentManagerUtils.getPythonUdfRunnerScript(dependencyInfo.getPythonExec(), env);

	return Environments.createProcessEnvironment(
		"",
		"",
		runnerScript,
		env);
}
 

/**
 * Stages all dependencies in pipeline into the jar file at outputStream, returning a new pipeline
 * that references these artifacts as classpath artifacts.
 */
@VisibleForTesting
protected Pipeline writeArtifacts(Pipeline pipeline, String jobName) throws IOException {
  Pipeline.Builder result = pipeline.toBuilder();
  for (Map.Entry<String, RunnerApi.Environment> env :
      pipeline.getComponents().getEnvironmentsMap().entrySet()) {
    result
        .getComponentsBuilder()
        .putEnvironments(env.getKey(), writeArtifacts(env.getValue(), jobName));
  }
  return result.build();
}
 

/**
 * Stages all dependencies in environment into the jar file at outputStream, returning a new
 * environment that references these artifacts as classpath artifacts.
 */
private RunnerApi.Environment writeArtifacts(RunnerApi.Environment environment, String jobName)
    throws IOException {
  RunnerApi.Environment.Builder result = environment.toBuilder();
  result.clearDependencies();
  for (RunnerApi.ArtifactInformation artifact : environment.getDependenciesList()) {
    result.addDependencies(writeArtifact(artifact, jobName));
  }
  return result.build();
}
 

@Override
public RunnerApi.Pipeline resolveArtifacts(RunnerApi.Pipeline pipeline) {
  ImmutableMap.Builder<String, RunnerApi.Environment> environmentMapBuilder =
      ImmutableMap.builder();
  for (Map.Entry<String, RunnerApi.Environment> entry :
      pipeline.getComponents().getEnvironmentsMap().entrySet()) {
    List<RunnerApi.ArtifactInformation> resolvedDependencies =
        entry
            .getValue()
            .getDependenciesList()
            .parallelStream()
            .flatMap(resolver)
            .collect(Collectors.toList());
    environmentMapBuilder.put(
        entry.getKey(),
        entry
            .getValue()
            .toBuilder()
            .clearDependencies()
            .addAllDependencies(resolvedDependencies)
            .build());
  }
  return pipeline
      .toBuilder()
      .setComponents(
          pipeline.getComponents().toBuilder().putAllEnvironments(environmentMapBuilder.build()))
      .build();
}
 

private RunnerApi.Environment resolveArtifacts(
    ArtifactRetrievalServiceGrpc.ArtifactRetrievalServiceBlockingStub retrievalStub,
    RunnerApi.Environment environment)
    throws IOException {
  return environment
      .toBuilder()
      .clearDependencies()
      .addAllDependencies(resolveArtifacts(retrievalStub, environment.getDependenciesList()))
      .build();
}
 

private SdkComponents createSdkComponents(PipelineOptions options) {
  SdkComponents sdkComponents = SdkComponents.create();

  String workerHarnessContainerImageURL =
      DataflowRunner.getContainerImageForJob(options.as(DataflowPipelineOptions.class));
  RunnerApi.Environment defaultEnvironmentForDataflow =
      Environments.createDockerEnvironment(workerHarnessContainerImageURL);

  sdkComponents.registerEnvironment(defaultEnvironmentForDataflow);
  return sdkComponents;
}
 

public static RemoteEnvironment create(
    ProcessManager processManager,
    RunnerApi.Environment environment,
    String workerId,
    InstructionRequestHandler instructionHandler) {
  return new ProcessEnvironment(processManager, environment, workerId, instructionHandler);
}
 

private ProcessEnvironment(
    ProcessManager processManager,
    RunnerApi.Environment environment,
    String workerId,
    InstructionRequestHandler instructionHandler) {

  this.processManager = processManager;
  this.environment = environment;
  this.workerId = workerId;
  this.instructionHandler = instructionHandler;
}
 

public void registerEnvironment(String workerId, RunnerApi.Environment environment) {
  environments.put(workerId, environment);
}
 

/** Creates an active {@link Environment} and returns a handle to it. */
RemoteEnvironment createEnvironment(RunnerApi.Environment environment, String workerId)
    throws Exception;
 

@Override
public RunnerApi.Environment getEnvironment() {
  return environment;
}
 

/**
 * Creates the Environment object used in Apache Beam Fn API.
 *
 * @return The Environment object which represents the environment(process, docker, etc) the python worker would run
 *         in.
 */
RunnerApi.Environment createEnvironment() throws Exception;
 

/**
 * Creates a specification which specifies the portability Python execution environment.
 * It's used by Beam's portability framework to creates the actual Python execution environment.
 */
protected RunnerApi.Environment createPythonExecutionEnvironment() throws Exception {
	return environmentManager.createEnvironment();
}