# Get available cuda tags

cuda_tags_cmd = "./get_cuda_tags.sh"

process = subprocess.run(

cuda_tags_cmd.split(), stdout=subprocess.PIPE, encoding="utf-8"

)

os_name = "ubuntu"

for tag in process.stdout.split("\n"):

os_version = tag[tag.find(os_name) + len(os_name) :]

cuda_version = tag[: tag.find("-devel")]

if os_version in cuda_dict:

cuda_dict[os_version].append(cuda_version)

# Get User info

user_id = os.geteuid()

# Check if sudo was used

if user_id == 0:

print(red_wrapper("DON'T RUN THIS AS SUDO!!!"))

exit(0)

group_id = os.getegid()

if args.user:

user_name = os.environ["USER"]

group_name = grp.getgrgid(group_id).gr_name

else:

user_name = "docker_user"

group_name = "docker_user"

if args.default:

ubuntu_input = default_ubuntu_vers

cuda_input = default_cuda_vers

data_input = default_data_path

workspace_input = default_workspace_path

else:

ubuntu_input = args.ubuntu

cuda_input = args.cuda

data_input = args.data_path

workspace_input = args.workspace_path

# Get Ubuntu Version from user

if ubuntu_input == "" or not ubuntu_input in ubuntu_dict:

question_str = "Version of Ubuntu would you like to use (Default: {}): ".format(default_ubuntu_vers)

question_str = blue_wrapper(question_str)

# print("What version of Ubuntu would you like to use?")

ubuntu_input = input(question_str)

while ubuntu_input and not ubuntu_input in ubuntu_dict:

print("\"{}\" is not available from {}".format(ubuntu_input,

list(ubuntu_dict.keys())))

ubuntu_input = input(question_str)

if not ubuntu_input:

ubuntu_input = default_ubuntu_vers

# Get CUDA version from user

if cuda_input == "" or not cuda_input in cuda_dict[ubuntu_input]:

question_str = "Version of CUDA would you like to use (Default: {}): ".format(

default_cuda_vers

)

question_str = blue_wrapper(question_str)

# print("What version of cuda would you like to use?")

cuda_input = input(question_str)

while cuda_input and not cuda_input in cuda_dict[ubuntu_input]:

print(

'"{}" is not available from {}'.format(

cuda_input, cuda_dict[ubuntu_input]

)

)

cuda_input = input(question_str)

if not cuda_input:

cuda_input = default_cuda_vers

# Get data folder from user

if data_input == default_data_path:

question_str = "What folder would you like to be `/data` in the Docker Container (Default: {})? ".format(

default_data_path

)

question_str = blue_wrapper(question_str)

data_input = input(question_str)

if not data_input:

data_input = default_data_path

# Get workspace folder from user

if workspace_input == default_workspace_path:

question_str = "What folder would you like to be `~/workspaces` in the Docker Container (Default: {})? ".format(

default_workspace_path

)

question_str = blue_wrapper(question_str)

workspace_input = input(question_str)

if not workspace_input:

workspace_input = default_workspace_path

workspace_input = os.path.expanduser(workspace_input)

data_input = os.path.expanduser(data_input)

create_folder(data_input)

create_folder(workspace_input)

env_file = (

"################ CAN BE OVERWRITTEN BY create_cuda_docker.py ################\n"

+ "USER_NAME={}\n".format(user_name)

+ "GROUP_NAME={}\n".format(group_name)

+ "USER_ID={}\n".format(user_id)

+ "GROUP_ID={}\n".format(group_id)

+ "UBUNTU_VERSION={}\n".format(ubuntu_input)

+ "UBUNTU_NAME={}\n".format(ubuntu_dict[ubuntu_input])

+ "CUDA_IMG_VERSION={}\n".format(cuda_input)

+ "ROS_DISTRO_NAME={}\n".format(ros_dict[ubuntu_input])

+ "HOME_DIRECTORY={}\n".format(os.path.expanduser("~"))

+ "WORKSPACE_DIRECTORY={}\n".format(workspace_input)

+ "DATA_DIRECTORY={}\n".format(data_input)

+ "DOCKER_USER_HOME_DIRECTORY={}\n".format(args.docker_user_home_directory)

)

with open(".env", "w") as file:

file.write(env_file)

print(blue_wrapper(".env has been rewritten to:"))

print(env_file)

docker_cmd = "docker compose up --no-start --build"

if args.build:

subprocess.run(docker_cmd.split())

print(blue_wrapper("Finished building docker container"))

docker_clean_cmd = "docker image prune"

else: