TensorFlowPsComponent extends AbstractComponent { @Override public Component createComponent() throws IOException { TensorFlowRunJobParameters tensorFlowParams = (TensorFlowRunJobParameters) this.parameters; Objects.requireNonNull(tensorFlowParams.getPsResource(), "PS resource must not be null!"); if (tensorFlowParams.getNumPS() < 1) { throw new IllegalArgumentException("Number of PS should be at least 1!"); } Component component = new Component(); component.setName(role.getComponentName()); component.setNumberOfContainers((long) tensorFlowParams.getNumPS()); component.setRestartPolicy(Component.RestartPolicyEnum.NEVER); component.setResource( SubmarineResourceUtils.convertYarnResourceToServiceResource(tensorFlowParams.getPsResource())); if (tensorFlowParams.getPsDockerImage() != null) { component.setArtifact( DockerUtilities.getDockerArtifact(tensorFlowParams.getPsDockerImage())); } TensorFlowCommons.addCommonEnvironments(component, role); generateLaunchCommand(component); return component; } TensorFlowPsComponent(FileSystemOperations fsOperations,
                               RemoteDirectoryManager remoteDirectoryManager,
                               TensorFlowLaunchCommandFactory launchCommandFactory,
                               RunJobParameters parameters,
                               Configuration yarnConfig); @Override Component createComponent();  }
@Test public void testPSComponentWithNullResource() throws IOException { TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters(); parameters.setPsResource(null); TensorFlowPsComponent psComponent = createPsComponent(parameters); expectedException.expect(NullPointerException.class); expectedException.expectMessage("PS resource must not be null"); psComponent.createComponent(); }
@Test public void testPSComponentWithNullJobName() throws IOException { TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters(); parameters.setPsResource(testCommons.resource); parameters.setNumPS(1); parameters.setName(null); TensorFlowPsComponent psComponent = createPsComponent(parameters); expectedException.expect(NullPointerException.class); expectedException.expectMessage("Job name must not be null"); psComponent.createComponent(); }
@Test public void testPSComponentZeroNumberOfPS() throws IOException { testCommons.yarnConfig.set("hadoop.registry.dns.domain-name", "testDomain"); TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters(); parameters.setPsResource(testCommons.resource); parameters.setName("testJobName"); parameters.setPsDockerImage("testPSDockerImage"); parameters.setNumPS(0); TensorFlowPsComponent psComponent = createPsComponent(parameters); expectedException.expect(IllegalArgumentException.class); expectedException.expectMessage("Number of PS should be at least 1!"); psComponent.createComponent(); }
@Test public void testPSComponentNumPSIsOne() throws IOException { testCommons.yarnConfig.set("hadoop.registry.dns.domain-name", "testDomain"); TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters(); parameters.setPsResource(testCommons.resource); parameters.setName("testJobName"); parameters.setNumPS(1); parameters.setPsDockerImage("testPSDockerImage"); TensorFlowPsComponent psComponent = createPsComponent(parameters); Component component = psComponent.createComponent(); assertEquals(1L, (long) component.getNumberOfContainers()); verifyCommons(component); }
@Test public void testPSComponentNumPSIsTwo() throws IOException { testCommons.yarnConfig.set("hadoop.registry.dns.domain-name", "testDomain"); TensorFlowRunJobParameters parameters = new TensorFlowRunJobParameters(); parameters.setPsResource(testCommons.resource); parameters.setName("testJobName"); parameters.setNumPS(2); parameters.setPsDockerImage("testPSDockerImage"); TensorFlowPsComponent psComponent = createPsComponent(parameters); Component component = psComponent.createComponent(); assertEquals(2L, (long) component.getNumberOfContainers()); verifyCommons(component); }