org.apache.reef.tang.formats.ConfigurationModule Java Examples

public static ConfigurationModule getDriverConfiguration() {
  return DriverConfiguration.CONF
      .setMultiple(DriverConfiguration.GLOBAL_LIBRARIES, EnvironmentUtils.getAllClasspathJars())
      .set(DriverConfiguration.ON_EVALUATOR_ALLOCATED, JobDriver.AllocatedEvaluatorHandler.class)
      .set(DriverConfiguration.ON_EVALUATOR_FAILED, JobDriver.FailedEvaluatorHandler.class)
      .set(DriverConfiguration.ON_CONTEXT_ACTIVE, JobDriver.ActiveContextHandler.class)
      .set(DriverConfiguration.ON_CONTEXT_CLOSED, JobDriver.ClosedContextHandler.class)
      .set(DriverConfiguration.ON_CONTEXT_FAILED, JobDriver.FailedContextHandler.class)
      .set(DriverConfiguration.ON_CONTEXT_MESSAGE, JobDriver.ContextMessageHandler.class)
      .set(DriverConfiguration.ON_TASK_MESSAGE, JobDriver.TaskMessageHandler.class)
      .set(DriverConfiguration.ON_TASK_FAILED, JobDriver.FailedTaskHandler.class)
      .set(DriverConfiguration.ON_TASK_RUNNING, JobDriver.RunningTaskHandler.class)
      .set(DriverConfiguration.ON_TASK_COMPLETED, JobDriver.CompletedTaskHandler.class)
      .set(DriverConfiguration.ON_DRIVER_STARTED, JobDriver.StartHandler.class)
      .set(DriverConfiguration.ON_TASK_SUSPENDED, JobDriver.SuspendedTaskHandler.class)
      .set(DriverConfiguration.ON_EVALUATOR_COMPLETED, JobDriver.CompletedEvaluatorHandler.class);
}
 

private Configuration getLocalConfiguration(final URI jobFolder,
                                            final String clientRemoteId,
                                            final String jobId) {

  ConfigurationModule localModule = LocalDriverConfiguration.CONF
          .set(LocalDriverConfiguration.MAX_NUMBER_OF_EVALUATORS, this.maxEvaluators)
          // ROOT FOLDER will point to the current runtime directory
          .set(LocalDriverConfiguration.ROOT_FOLDER, ".")
          .set(LocalDriverConfiguration.JVM_HEAP_SLACK, this.jvmSlack)
          .set(LocalDriverConfiguration.CLIENT_REMOTE_IDENTIFIER, clientRemoteId)
          .set(LocalDriverConfiguration.JOB_IDENTIFIER, jobId)
          .set(LocalDriverConfiguration.RUNTIME_NAMES,
                  org.apache.reef.runtime.local.driver.RuntimeIdentifier.RUNTIME_NAME);
  for (final String rackName : rackNames) {
    localModule = localModule.set(LocalDriverConfiguration.RACK_NAMES, rackName);
  }

  return localModule.build();
}
 

/**
 * Assembles the driver configuration.
 *
 * @param jobFolder                The folder in which the local runtime will execute this job.
 * @param clientRemoteId           the remote identifier of the client. It is used by the Driver to establish a
 *                                 connection back to the client.
 * @param jobId                    The identifier of the job.
 * @param applicationConfiguration The configuration of the application, e.g. a filled out DriverConfiguration
 * @return The Driver configuration to be used to instantiate the Driver.
 */
@Override
public Configuration getDriverConfiguration(final URI jobFolder,
                                            final String clientRemoteId,
                                            final String jobId,
                                            final Configuration applicationConfiguration) {
  AvroMultiRuntimeDefinition runtimeDefinitions = this.definitionGenerator.getMultiRuntimeDefinition(
          jobFolder,
          clientRemoteId,
          jobId);
  ConfigurationModule conf = MultiRuntimeDriverConfiguration.CONF;

  for(AvroRuntimeDefinition runtimeDefinition : runtimeDefinitions.getRuntimes()){
    conf = conf.set(MultiRuntimeDriverConfiguration.RUNTIME_NAMES, runtimeDefinition.getRuntimeName().toString());
  }

  final Configuration mainConfiguration = this.mainRuntimeConfigGenerator.getMainConfiguration();

  return Configurations.merge(mainConfiguration, applicationConfiguration,
                  conf
                  .set(MultiRuntimeDriverConfiguration.JOB_IDENTIFIER, jobId)
                  .set(MultiRuntimeDriverConfiguration.CLIENT_REMOTE_IDENTIFIER, clientRemoteId)
                  .set(MultiRuntimeDriverConfiguration.SERIALIZED_RUNTIME_DEFINITION,
                          this.runtimeDefinitionSerializer.toString(runtimeDefinitions))
                  .build());
}
 

private Configuration getDriverConfiguration(final URI jobFolder,
                                             final String clientRemoteId,
                                             final String jobId) {
  ConfigurationModule configModule = StandaloneDriverConfiguration.CONF
      .set(StandaloneDriverConfiguration.ROOT_FOLDER, jobFolder.getPath())
      .set(StandaloneDriverConfiguration.NODE_FOLDER, this.nodeFolder)
      .set(StandaloneDriverConfiguration.NODE_LIST_FILE_PATH, this.nodeListFilePath)
      .set(StandaloneDriverConfiguration.SSH_PORT_NUM, this.sshPortNum)
      .set(StandaloneDriverConfiguration.JVM_HEAP_SLACK, this.jvmHeapSlack)
      .set(StandaloneDriverConfiguration.CLIENT_REMOTE_IDENTIFIER, clientRemoteId)
      .set(StandaloneDriverConfiguration.JOB_IDENTIFIER, jobId);
  for (final String nodeInfo : nodeInfoSet) {
    configModule = configModule.set(StandaloneDriverConfiguration.NODE_INFO_SET, nodeInfo);
  }
  return configModule.build();
}
 

@Override
public Configuration getRuntimeConfiguration(
    final ClientProtocol.DriverClientConfiguration driverClientConfiguration) {
  ConfigurationModule localRuntimeCM = LocalRuntimeConfiguration.CONF;
  if (driverClientConfiguration.getLocalRuntime().getMaxNumberOfEvaluators() > 0) {
    localRuntimeCM = localRuntimeCM.set(LocalRuntimeConfiguration.MAX_NUMBER_OF_EVALUATORS,
        driverClientConfiguration.getLocalRuntime().getMaxNumberOfEvaluators());
  }
  if (StringUtils.isNotEmpty(driverClientConfiguration.getLocalRuntime().getRuntimeRootFolder())) {
    localRuntimeCM = localRuntimeCM.set(LocalRuntimeConfiguration.RUNTIME_ROOT_FOLDER,
        driverClientConfiguration.getLocalRuntime().getRuntimeRootFolder());
  }
  if (driverClientConfiguration.getLocalRuntime().getJvmHeapSlack() > 0.0) {
    localRuntimeCM = localRuntimeCM.set(LocalRuntimeConfiguration.JVM_HEAP_SLACK,
        driverClientConfiguration.getLocalRuntime().getJvmHeapSlack());
  }
  if (StringUtils.isNotEmpty(driverClientConfiguration.getDriverJobSubmissionDirectory())) {
    localRuntimeCM = localRuntimeCM.set(LocalRuntimeConfiguration.RUNTIME_ROOT_FOLDER,
        driverClientConfiguration.getDriverJobSubmissionDirectory());
  }
  return localRuntimeCM.build();
}
 

private Configuration getDriverRestartConfiguration(
    final ClientProtocol.DriverClientConfiguration driverConfiguration) {
  final ConfigurationModule restartConfModule = DriverRestartConfiguration.CONF
      .set(DriverRestartConfiguration.ON_DRIVER_RESTARTED,
          DriverServiceHandlers.DriverRestartHandler.class)
      .set(DriverRestartConfiguration.ON_DRIVER_RESTART_CONTEXT_ACTIVE,
          DriverServiceHandlers.DriverRestartActiveContextHandler.class)
      .set(DriverRestartConfiguration.ON_DRIVER_RESTART_TASK_RUNNING,
          DriverServiceHandlers.DriverRestartRunningTaskHandler.class)
      .set(DriverRestartConfiguration.ON_DRIVER_RESTART_COMPLETED,
          DriverServiceHandlers.DriverRestartCompletedHandler.class)
      .set(DriverRestartConfiguration.ON_DRIVER_RESTART_EVALUATOR_FAILED,
          DriverServiceHandlers.DriverRestartFailedEvaluatorHandler.class);
  return driverConfiguration.getDriverRestartEvaluatorRecoverySeconds() > 0 ?
      restartConfModule
          .set(DriverRestartConfiguration.DRIVER_RESTART_EVALUATOR_RECOVERY_SECONDS,
              driverConfiguration.getDriverRestartEvaluatorRecoverySeconds())
          .build() :
      restartConfModule.build();
}
 

private Configuration getDriverConfiguration(
    final ClientProtocol.DriverClientConfiguration driverConfiguration) {
  ConfigurationModule driverServiceConfigurationModule = DriverServiceConfiguration.STATIC_DRIVER_CONF_MODULE
      .set(DriverConfiguration.DRIVER_IDENTIFIER, driverConfiguration.getJobid());

  // Set file dependencies
  final List<String> localLibraries = new ArrayList<>();
  localLibraries.add(EnvironmentUtils.getClassLocation(GRPCDriverService.class));
  localLibraries.addAll(driverConfiguration.getLocalLibrariesList());
  driverServiceConfigurationModule = driverServiceConfigurationModule
      .setMultiple(DriverConfiguration.LOCAL_LIBRARIES, localLibraries);
  driverServiceConfigurationModule = driverServiceConfigurationModule
      .setMultiple(DriverConfiguration.GLOBAL_LIBRARIES,
          driverConfiguration.getGlobalLibrariesList());
  driverServiceConfigurationModule = driverServiceConfigurationModule
      .setMultiple(DriverConfiguration.LOCAL_FILES,
          driverConfiguration.getLocalFilesList());
  driverServiceConfigurationModule = driverServiceConfigurationModule
      .setMultiple(DriverConfiguration.GLOBAL_FILES,
          driverConfiguration.getGlobalFilesList());
  // Setup driver resources
  if (driverConfiguration.getCpuCores() > 0) {
    driverServiceConfigurationModule = driverServiceConfigurationModule
        .set(DriverConfiguration.DRIVER_CPU_CORES, driverConfiguration.getCpuCores());
  }
  if (driverConfiguration.getMemoryMb() > 0) {
    driverServiceConfigurationModule = driverServiceConfigurationModule
        .set(DriverConfiguration.DRIVER_MEMORY, driverConfiguration.getMemoryMb());
  }
  // Job submission directory
  if (StringUtils.isNotEmpty(driverConfiguration.getDriverJobSubmissionDirectory())) {
    driverServiceConfigurationModule = driverServiceConfigurationModule
        .set(DriverConfiguration.DRIVER_JOB_SUBMISSION_DIRECTORY,
            driverConfiguration.getDriverJobSubmissionDirectory());
  }
  return !driverConfiguration.getEnableHttpDriver() ? driverServiceConfigurationModule.build() :
      Configurations.merge(DriverServiceConfiguration.HTTP_AND_NAMESERVER, driverServiceConfigurationModule.build());
}
 

public static LauncherStatus runHelloCLR(final Configuration runtimeConf, final int timeOut, final File clrFolder)
    throws BindException, InjectionException {

  final ConfigurationModule driverConf =
      addAll(DriverConfiguration.CONF, DriverConfiguration.GLOBAL_FILES, clrFolder)
          .set(DriverConfiguration.GLOBAL_LIBRARIES, EnvironmentUtils.getClassLocation(HelloDriver.class))
          .set(DriverConfiguration.DRIVER_IDENTIFIER, "HelloCLR")
          .set(DriverConfiguration.ON_DRIVER_STARTED, HelloDriver.StartHandler.class)
          .set(DriverConfiguration.ON_EVALUATOR_ALLOCATED, HelloDriver.EvaluatorAllocatedHandler.class);

  return DriverLauncher.getLauncher(runtimeConf).run(driverConf.build(), timeOut);
}
 

private static ConfigurationModule addAll(final ConfigurationModule conf,
                                          final OptionalParameter<String> param,
                                          final File folder) {
  ConfigurationModule result = conf;
  final File[] files = folder.listFiles();
  if (files != null) {
    for (final File f : files) {
      if (f.canRead() && f.exists() && f.isFile()) {
        result = result.set(param, f.getAbsolutePath());
      }
    }
  }
  return result;
}
 

/**
 * Get or create a {@link ConfigurationModule} for the Azure Batch runtime.
 *
 * @param isWindows true if Azure Batch pool nodes run Windows, false otherwise.
 * @return the configuration module object.
 */
public static ConfigurationModule getOrCreateAzureBatchRuntimeConfiguration(
    final boolean isWindows) {

  if (AzureBatchRuntimeConfigurationCreator.conf == null) {
    ConfigurationModuleBuilder builder = AzureBatchRuntimeConfigurationStatic.CONF;

    if (isWindows) {
      builder = builder.bindImplementation(CommandBuilder.class, WindowsCommandBuilder.class);
    } else {
      builder = builder.bindImplementation(CommandBuilder.class, LinuxCommandBuilder.class);
    }

    AzureBatchRuntimeConfigurationCreator.conf = new AzureBatchRuntimeConfiguration()
        .merge(builder.build())
        .bindNamedParameter(AzureBatchAccountName.class, AzureBatchRuntimeConfiguration.AZURE_BATCH_ACCOUNT_NAME)
        .bindNamedParameter(AzureBatchAccountUri.class, AzureBatchRuntimeConfiguration.AZURE_BATCH_ACCOUNT_URI)
        .bindNamedParameter(AzureBatchAccountKey.class, AzureBatchRuntimeConfiguration.AZURE_BATCH_ACCOUNT_KEY)
        .bindNamedParameter(AzureBatchPoolId.class, AzureBatchRuntimeConfiguration.AZURE_BATCH_POOL_ID)
        .bindNamedParameter(AzureStorageAccountName.class, AzureBatchRuntimeConfiguration.AZURE_STORAGE_ACCOUNT_NAME)
        .bindNamedParameter(AzureStorageAccountKey.class, AzureBatchRuntimeConfiguration.AZURE_STORAGE_ACCOUNT_KEY)
        .bindNamedParameter(ContainerRegistryServer.class, AzureBatchRuntimeConfiguration.CONTAINER_REGISTRY_SERVER)
        .bindNamedParameter(
            ContainerRegistryUsername.class, AzureBatchRuntimeConfiguration.CONTAINER_REGISTRY_USERNAME)
        .bindNamedParameter(
            ContainerRegistryPassword.class, AzureBatchRuntimeConfiguration.CONTAINER_REGISTRY_PASSWORD)
        .bindNamedParameter(ContainerImageName.class, AzureBatchRuntimeConfiguration.CONTAINER_IMAGE_NAME)
        .bindNamedParameter(
            AzureStorageContainerName.class, AzureBatchRuntimeConfiguration.AZURE_STORAGE_CONTAINER_NAME)
        .bindSetEntry(TcpPortSet.class, AzureBatchRuntimeConfiguration.TCP_PORT_SET)
        .build();
  }

  return AzureBatchRuntimeConfigurationCreator.conf;
}
 

public static ConfigurationModule getConfigurationModule(final boolean includeContainerConfiguration) {
  ConfigurationModuleBuilder shimConfigurationBuilder = EvaluatorShimConfiguration.CONF.getBuilder();

  // If using docker containers, then use a different set of bindings
  if (includeContainerConfiguration) {
    shimConfigurationBuilder = shimConfigurationBuilder
        .bindImplementation(LocalAddressProvider.class, ContainerBasedLocalAddressProvider.class)
        .bindImplementation(TcpPortProvider.class, SetTcpPortProvider.class);
  }

  return shimConfigurationBuilder.build();
}
 

/**
 * Fills out a ConfigurationModule.
 *
 * @param input           The ConfigurationModule to start with.
 * @param globalFileField the field on which to set() the global files.
 * @param globalLibField  the field on which to set() the global libraries.
 * @param localFileField  the field on which to set() the local files.
 * @param localLibField   the field on which to set() the local libraries.
 * @return a copy of input with files and libraries added to the given fields.
 */
public ConfigurationModule addNamesTo(final ConfigurationModule input,
                                      final OptionalParameter<String> globalFileField,
                                      final OptionalParameter<String> globalLibField,
                                      final OptionalParameter<String> localFileField,
                                      final OptionalParameter<String> localLibField) {
  ConfigurationModule result = input;
  result = this.globalFiles.addNamesTo(result, globalFileField);
  result = this.globalLibs.addNamesTo(result, globalLibField);
  result = this.localFiles.addNamesTo(result, localFileField);
  result = this.localLibs.addNamesTo(result, localLibField);
  return result;
}
 

private Configuration getDriverConfiguration(final URI jobFolder,
                                             final String clientRemoteId,
                                             final String jobId) {
  ConfigurationModule configModule = LocalDriverConfiguration.CONF
      .set(LocalDriverConfiguration.MAX_NUMBER_OF_EVALUATORS, this.maxEvaluators)
      .set(LocalDriverConfiguration.ROOT_FOLDER, jobFolder.getPath())
      .set(LocalDriverConfiguration.JVM_HEAP_SLACK, this.jvmHeapSlack)
      .set(LocalDriverConfiguration.CLIENT_REMOTE_IDENTIFIER, clientRemoteId)
      .set(LocalDriverConfiguration.JOB_IDENTIFIER, jobId)
      .set(LocalDriverConfiguration.RUNTIME_NAMES, RuntimeIdentifier.RUNTIME_NAME);
  for (final String rackName : rackNames) {
    configModule = configModule.set(LocalDriverConfiguration.RACK_NAMES, rackName);
  }
  return configModule.build();
}
 

/**
 * Adds the file names of this FileSet to the given field of the given ConfigurationModule.
 *
 * @param input the ConfigurationModule to fill out
 * @param field the field to add the files in this set to.
 * @return the filled out ConfigurationModule
 */
ConfigurationModule addNamesTo(final ConfigurationModule input, final OptionalParameter<String> field) {
  ConfigurationModule result = input;
  for (final String fileName : this.fileNames()) {
    result = result.set(field, fileName);
  }
  return result;
}
 

/**
 * Assembles the configuration based on TestDriverConfiguration.
 *
 * @param theFiles
 * @return
 * @throws BindException
 * @throws IOException
 */
private static Configuration getTestDriverConfiguration(final Set<File> theFiles) throws BindException, IOException {
  ConfigurationModule testDriverConfigurationModule = FileResourceTestDriverConfiguration.CONF;
  for (final File f : theFiles) {
    LOG.log(Level.FINEST, "Adding a file to the TestDriverConfiguration: " + f.getName());
    testDriverConfigurationModule =
        testDriverConfigurationModule.set(FileResourceTestDriverConfiguration.EXPECTED_FILE_NAME, f.getName());
  }

  final Configuration testDriverConfiguration = testDriverConfigurationModule.build();
  return testDriverConfiguration;
}
 

/**
 * Assembles the driver configuration using the DriverConfiguration class.
 *
 * @param theFiles
 * @return
 * @throws BindException
 */
private static Configuration getDriverConfiguration(final Set<File> theFiles) throws BindException {
  ConfigurationModule driverConfigurationModule = DriverConfiguration.CONF
      .set(DriverConfiguration.GLOBAL_LIBRARIES, EnvironmentUtils.getClassLocation(FileResourceTestDriver.class))
      .set(DriverConfiguration.DRIVER_IDENTIFIER, "TEST_FileResourceTest")
      .set(DriverConfiguration.ON_DRIVER_STARTED, FileResourceTestDriver.StartHandler.class)
      .set(DriverConfiguration.ON_EVALUATOR_ALLOCATED, FileResourceTestDriver.EvaluatorAllocatedHandler.class);

  for (final File f : theFiles) {
    LOG.log(Level.FINEST, "Adding a file to the DriverConfiguration: " + f.getAbsolutePath());
    driverConfigurationModule = driverConfigurationModule.set(DriverConfiguration.LOCAL_FILES, f.getAbsolutePath());
  }
  return driverConfigurationModule.build();
}
 

public DataLoadingRequestBuilder setDriverConfigurationModule(
    final ConfigurationModule driverConfigurationModule) {
  this.driverConfigurationModule = driverConfigurationModule;
  return this;
}
 

private Configuration getMultiRuntimeDriverConfiguration(
        final AvroYarnJobSubmissionParameters yarnJobSubmissionParams,
        final AvroMultiRuntimeAppSubmissionParameters multiruntimeAppSubmissionParams) {

  if (multiruntimeAppSubmissionParams.getLocalRuntimeAppParameters() == null &&
          multiruntimeAppSubmissionParams.getYarnRuntimeAppParameters() == null){
    throw new IllegalArgumentException("At least on execution runtime has to be provided");
  }

  // read yarn job submission parameters
  final AvroJobSubmissionParameters jobSubmissionParameters =
          yarnJobSubmissionParams.getSharedJobSubmissionParameters();

  // generate multi runtime definition
  final MultiRuntimeDefinitionBuilder multiRuntimeDefinitionBuilder = new MultiRuntimeDefinitionBuilder();

  if (multiruntimeAppSubmissionParams.getLocalRuntimeAppParameters() != null){
    addLocalRuntimeDefinition(
            multiruntimeAppSubmissionParams.getLocalRuntimeAppParameters(),
            jobSubmissionParameters, multiRuntimeDefinitionBuilder);
  }

  if (multiruntimeAppSubmissionParams.getYarnRuntimeAppParameters() != null){
    addYarnRuntimeDefinition(
            yarnJobSubmissionParams,
            jobSubmissionParameters,
            multiRuntimeDefinitionBuilder);
  } else {
    addDummyYarnRuntimeDefinition(
            yarnJobSubmissionParams,
            jobSubmissionParameters,
            multiRuntimeDefinitionBuilder);
  }

  multiRuntimeDefinitionBuilder.setDefaultRuntimeName(
          multiruntimeAppSubmissionParams.getDefaultRuntimeName().toString());

  // generate multi runtime configuration
  ConfigurationModule multiRuntimeDriverConfiguration = MultiRuntimeDriverConfiguration.CONF
          .set(MultiRuntimeDriverConfiguration.JOB_IDENTIFIER, jobSubmissionParameters.getJobId().toString())
          .set(MultiRuntimeDriverConfiguration.CLIENT_REMOTE_IDENTIFIER, ClientRemoteIdentifier.NONE)
          .set(MultiRuntimeDriverConfiguration.SERIALIZED_RUNTIME_DEFINITION,
                  this.runtimeDefinitionSerializer.toString(multiRuntimeDefinitionBuilder.build()));

  for (final CharSequence runtimeName : multiruntimeAppSubmissionParams.getRuntimes()){
    multiRuntimeDriverConfiguration = multiRuntimeDriverConfiguration.set(
            MultiRuntimeDriverConfiguration.RUNTIME_NAMES, runtimeName.toString());
  }

  final AvroAppSubmissionParameters appSubmissionParams =
          multiruntimeAppSubmissionParams.getSharedAppSubmissionParameters();

  // generate yarn related driver configuration
  final Configuration providerConfig = Tang.Factory.getTang().newConfigurationBuilder()
          .bindSetEntry(DriverConfigurationProviders.class, TcpPortConfigurationProvider.class)
          .bindNamedParameter(TcpPortRangeBegin.class, Integer.toString(appSubmissionParams.getTcpBeginPort()))
          .bindNamedParameter(TcpPortRangeCount.class, Integer.toString(appSubmissionParams.getTcpRangeCount()))
          .bindNamedParameter(TcpPortRangeTryCount.class, Integer.toString(appSubmissionParams.getTcpTryCount()))
          .bindNamedParameter(JobSubmissionDirectoryPrefix.class,
                  yarnJobSubmissionParams.getJobSubmissionDirectoryPrefix().toString())
          .bindImplementation(RuntimeClasspathProvider.class, YarnClasspathProvider.class)
          .bindConstructor(YarnConfiguration.class, YarnConfigurationConstructor.class)
          .build();

  final Configuration driverConfiguration = Configurations.merge(
          Constants.DRIVER_CONFIGURATION_WITH_HTTP_AND_NAMESERVER,
          multiRuntimeDriverConfiguration.build(),
          providerConfig);

  // add restart configuration if needed
  if (multiruntimeAppSubmissionParams.getYarnRuntimeAppParameters() != null &&
          multiruntimeAppSubmissionParams.getYarnRuntimeAppParameters().getDriverRecoveryTimeout() > 0) {
    LOG.log(Level.FINE, "Driver restart is enabled.");

    final Configuration yarnDriverRestartConfiguration =
            YarnDriverRestartConfiguration.CONF.build();

    final Configuration driverRestartConfiguration =
            DriverRestartConfiguration.CONF
                    .set(DriverRestartConfiguration.ON_DRIVER_RESTARTED, JobDriver.RestartHandler.class)
                    .set(DriverRestartConfiguration.ON_DRIVER_RESTART_CONTEXT_ACTIVE,
                            JobDriver.DriverRestartActiveContextHandler.class)
                    .set(DriverRestartConfiguration.ON_DRIVER_RESTART_TASK_RUNNING,
                            JobDriver.DriverRestartRunningTaskHandler.class)
                    .set(DriverRestartConfiguration.DRIVER_RESTART_EVALUATOR_RECOVERY_SECONDS,
                            multiruntimeAppSubmissionParams.getYarnRuntimeAppParameters().getDriverRecoveryTimeout())
                    .set(DriverRestartConfiguration.ON_DRIVER_RESTART_COMPLETED,
                            JobDriver.DriverRestartCompletedHandler.class)
                    .set(DriverRestartConfiguration.ON_DRIVER_RESTART_EVALUATOR_FAILED,
                            JobDriver.DriverRestartFailedEvaluatorHandler.class)
                    .build();

    return Configurations.merge(driverConfiguration, yarnDriverRestartConfiguration, driverRestartConfiguration);
  }

  return driverConfiguration;
}
 

/**
 * Make configuration for Evaluator.
 * @param contextConfiguration
 * @param evaluatorConfiguration
 * @param serviceConfiguration
 * @param taskConfiguration
 * @return Configuration
 */
private Configuration makeEvaluatorConfiguration(final String contextConfiguration,
                                                 final Optional<String> evaluatorConfiguration,
                                                 final Optional<String> serviceConfiguration,
                                                 final Optional<String> taskConfiguration) {

  final ConfigurationModule evaluatorConfigModule;
  if (this.evaluatorManager.getEvaluatorDescriptor().getProcess() instanceof CLRProcess) {
    evaluatorConfigModule = EvaluatorConfiguration.CONFCLR;
  } else if (this.evaluatorManager.getEvaluatorDescriptor().getProcess() instanceof DotNetProcess) {
    evaluatorConfigModule = EvaluatorConfiguration.CONFDOTNET;
  } else {
    evaluatorConfigModule = EvaluatorConfiguration.CONF;
  }
  ConfigurationModule evaluatorConfigurationModule = evaluatorConfigModule
      .set(EvaluatorConfiguration.APPLICATION_IDENTIFIER, this.jobIdentifier)
      .set(EvaluatorConfiguration.DRIVER_REMOTE_IDENTIFIER, this.remoteID)
      .set(EvaluatorConfiguration.EVALUATOR_IDENTIFIER, this.getId())
      .set(EvaluatorConfiguration.ROOT_CONTEXT_CONFIGURATION, contextConfiguration);

  // Add the (optional) service configuration
  if (evaluatorConfiguration.isPresent()) {
    evaluatorConfigurationModule = evaluatorConfigurationModule
        .set(EvaluatorConfiguration.EVALUATOR_CONFIGURATION, evaluatorConfiguration.get());
  }

  // Add the (optional) service configuration
  if (serviceConfiguration.isPresent()) {
    evaluatorConfigurationModule = evaluatorConfigurationModule
        .set(EvaluatorConfiguration.ROOT_SERVICE_CONFIGURATION, serviceConfiguration.get());
  } else {
    evaluatorConfigurationModule = evaluatorConfigurationModule
        .set(EvaluatorConfiguration.ROOT_SERVICE_CONFIGURATION,
            this.configurationSerializer.toString(Tang.Factory.getTang().newConfigurationBuilder().build()));
  }

  // Add the (optional) task configuration
  if (taskConfiguration.isPresent()) {
    evaluatorConfigurationModule = evaluatorConfigurationModule
        .set(EvaluatorConfiguration.TASK_CONFIGURATION, taskConfiguration.get());
  }

  // Create the evaluator configuration.
  return evaluatorConfigurationModule.build();
}
 

@Override
public void onNext(final AllocatedEvaluator eval) {

  try {

    taskId = failTaskName + "_" + eval.getId();
    LOG.log(Level.INFO, "Submit task: {0}", taskId);

    final Configuration contextConfig =
        ContextConfiguration.CONF.set(ContextConfiguration.IDENTIFIER, taskId).build();

    ConfigurationModule taskConfig =
        TaskConfiguration.CONF.set(TaskConfiguration.IDENTIFIER, taskId);

    switch (failTaskName) {
    case "FailTask":
      taskConfig = taskConfig.set(TaskConfiguration.TASK, FailTask.class);
      break;
    case "FailTaskCall":
      taskConfig = taskConfig.set(TaskConfiguration.TASK, FailTaskCall.class);
      break;
    case "FailTaskMsg":
      taskConfig = taskConfig
            .set(TaskConfiguration.TASK, FailTaskMsg.class)
            .set(TaskConfiguration.ON_MESSAGE, FailTaskMsg.class);
      break;
    case "FailTaskSuspend":
      taskConfig = taskConfig
            .set(TaskConfiguration.TASK, FailTaskSuspend.class)
            .set(TaskConfiguration.ON_SUSPEND, FailTaskSuspend.class);
      break;
    case "FailTaskStart":
      taskConfig = taskConfig
            .set(TaskConfiguration.TASK, FailTaskStart.class)
            .set(TaskConfiguration.ON_TASK_STARTED, FailTaskStart.class);
      break;
    case "FailTaskStop":
      taskConfig = taskConfig
            .set(TaskConfiguration.TASK, FailTaskStop.class)
            .set(TaskConfiguration.ON_TASK_STOP, FailTaskStop.class)
            .set(TaskConfiguration.ON_CLOSE, FailTaskStop.CloseEventHandler.class);
      break;
    case "FailTaskClose":
      taskConfig = taskConfig
            .set(TaskConfiguration.TASK, FailTaskClose.class)
            .set(TaskConfiguration.ON_CLOSE, FailTaskClose.class);
      break;
    default:
      break;
    }

    eval.submitContextAndTask(contextConfig, taskConfig.build());

  } catch (final BindException ex) {
    LOG.log(Level.WARNING, "Configuration error", ex);
    throw new DriverSideFailure("Configuration error", ex);
  }
}