Utils { public static Map<String, String> getContainerEnvForDocker(Configuration tonyConf, String jobType) { Map<String, String> containerEnv = new HashMap<>(); if (tonyConf.getBoolean(TonyConfigurationKeys.DOCKER_ENABLED, TonyConfigurationKeys.DEFAULT_DOCKER_ENABLED)) { String imagePath = tonyConf.get(TonyConfigurationKeys.getContainerDockerKey()); String jobImagePath = tonyConf.get(TonyConfigurationKeys.getDockerImageKey(jobType)); if (jobImagePath != null) { imagePath = jobImagePath; } if (imagePath == null) { LOG.error("Docker is enabled but " + TonyConfigurationKeys.getContainerDockerKey() + " is not set."); return containerEnv; } else { Class containerRuntimeClass = null; Class dockerRuntimeClass = null; try { containerRuntimeClass = Class.forName(Constants.CONTAINER_RUNTIME_CONSTANTS_CLASS); dockerRuntimeClass = Class.forName(Constants.DOCKER_LINUX_CONTAINER_RUNTIME_CLASS); } catch (ClassNotFoundException e) { LOG.error("Docker runtime classes not found in this version (" + org.apache.hadoop.util.VersionInfo.getVersion() + ") of Hadoop.", e); } if (dockerRuntimeClass != null) { try { String containerMounts = tonyConf.get(TonyConfigurationKeys.getContainerDockerMountKey()); if (StringUtils.isNotEmpty(containerMounts)) { String envDockerContainerMounts = (String) dockerRuntimeClass.getField(Constants.ENV_DOCKER_CONTAINER_MOUNTS).get(null); containerEnv.put(envDockerContainerMounts, containerMounts); } String envContainerType = (String) containerRuntimeClass.getField(Constants.ENV_CONTAINER_TYPE).get(null); String envDockerImage = (String) dockerRuntimeClass.getField(Constants.ENV_DOCKER_CONTAINER_IMAGE).get(null); containerEnv.put(envContainerType, "docker"); containerEnv.put(envDockerImage, imagePath); } catch (NoSuchFieldException e) { LOG.error("Field " + Constants.ENV_CONTAINER_TYPE + " or " + Constants.ENV_DOCKER_CONTAINER_IMAGE + " or " + Constants.ENV_DOCKER_CONTAINER_MOUNTS + " not found in " + containerRuntimeClass.getName() + " or " + dockerRuntimeClass.getName(), e); } catch (IllegalAccessException e) { LOG.error("Unable to access " + Constants.ENV_CONTAINER_TYPE + " or " + Constants.ENV_DOCKER_CONTAINER_IMAGE + " or " + Constants.ENV_DOCKER_CONTAINER_MOUNTS + " fields.", e); } } } } return containerEnv; } private  Utils(); static boolean poll(Callable<Boolean> func, int interval, int timeout); static T pollTillNonNull(Callable<T> func, int interval, int timeout); static String parseMemoryString(String memory); static void zipFolder(java.nio.file.Path sourceFolderPath, java.nio.file.Path zipPath); static void unzipArchive(String src, String dst); static void setCapabilityGPU(Resource resource, int gpuCount); static String constructUrl(String urlString); static String constructContainerUrl(Container container); static String constructContainerUrl(String nodeAddress, ContainerId containerId); static void printTaskUrl(TaskInfo taskInfo, Log log); static void printTonyPortalUrl(String portalUrl, String appId, Log log); static String getTonySrcZipName(String appId); static Map<String, String> parseKeyValue(String[] keyValues); static Options getCommonOptions(); static int executeShell(String taskCommand, long timeout, Map<String, String> env); static String getCurrentHostName(); static String getHostNameOrIpFromTokenConf(Configuration conf); static void addEnvironmentForResource(LocalResource resource, FileSystem fs, String envPrefix,
                                               Map<String, String> env); static Map<String, JobContainerRequest> parseContainerRequests(Configuration conf); static Set<String> getAllJobTypes(Configuration conf); static int getNumTotalTasks(Configuration conf); static String getTaskType(String confKey); static boolean isArchive(String path); static boolean renameFile(String oldName, String newName); static String constructTFConfig(String clusterSpec, String jobName, int taskIndex); static String getClientResourcesPath(String appId, String fileName); static void cleanupHDFSPath(Configuration hdfsConf, Path path); static void addResources(String[] resources, Map<String, LocalResource> resourcesMap, FileSystem fs); static void addResource(String path, Map<String, LocalResource> resourcesMap, FileSystem fs); static String buildRMUrl(Configuration yarnConf, String appId); static void printCompletedTrackedTasks(int completedTrackedTasks, int totalTrackedTasks); static String parseClusterSpecForPytorch(String clusterSpec); static String[] parseClusterSpecForMXNet(String clusterSpec); static String resolveNameToIpAddress(String hostname); static void createDirIfNotExists(FileSystem fs, Path dir, FsPermission permission); static String[] getUntrackedJobTypes(Configuration conf); static boolean isJobTypeTracked(String taskName, Configuration tonyConf); static String[] getStopOnFailureJobTypes(Configuration conf); static void uploadFileAndSetConfResources(Path hdfsPath, Path filePath, String fileName,
                                                   Configuration tonyConf, FileSystem fs,
                                                   LocalResourceType resourceType, String resourceKey); static String[] splitAddressPort(String hostname); static void appendConfResources(String key, String resource, Configuration tonyConf); static void initYarnConf(Configuration yarnConf); static void initHdfsConf(Configuration hdfsConf); static void extractResources(String appId); static Map<String, String> getContainerEnvForDocker(Configuration tonyConf, String jobType); static Map<String, String> linksToBeDisplayedOnPage(String jobId);  }
@Test public void testGetContainerEnvForDocker() { Configuration conf = mock(Configuration.class); when(conf.getBoolean(TonyConfigurationKeys.DOCKER_ENABLED, TonyConfigurationKeys.DEFAULT_DOCKER_ENABLED)) .thenReturn(true); assertEquals(Utils.getContainerEnvForDocker(conf, "tony.worker.gpus"), new HashMap<>()); when(conf.get(TonyConfigurationKeys .getDockerImageKey("tony.worker.gpus"))).thenReturn("foo"); assertEquals(Utils.getContainerEnvForDocker(conf, "tony.worker.gpus"), new HashMap<>()); }