SpecialInstanceTypeRouter implements CellRouter { @Override public Optional<Cell> routeKey(JobDescriptor jobDescriptor) { return selector.select(instanceTypeRouteKeyResolver.apply(jobDescriptor), c -> true); } SpecialInstanceTypeRouter(CellInfoResolver cellInfoResolver,
                                     Function<JobDescriptor, String> instanceTypeRouteKeyResolver,
                                     TitusFederationConfiguration federationConfiguration); @Override Optional<Cell> routeKey(JobDescriptor jobDescriptor); static SpecialInstanceTypeRouter getGpuInstanceTypeRouter(CellInfoResolver cellInfoResolver,
                                                                     TitusFederationConfiguration federationConfiguration);  }
@Test public void testGpuCellIsSelectedForGpuJob() { JobDescriptor gpuJobDescriptor = JobDescriptor.newBuilder() .setContainer(Container.newBuilder() .setResources(ContainerResources.newBuilder() .setGpu(5) ) ) .build(); assertThat(router.routeKey(gpuJobDescriptor)).contains(CELL_2); }
@Test public void testNoCellIsSelectedForNonGpuJob() { assertThat(router.routeKey(JobDescriptor.getDefaultInstance())).isEmpty(); }