hpccm.config

# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Global configuration
#
# Import this as
# import hpccm.config
#
# And access variables as
# hpccm.config.var


from packaging.version import Version
import archspec.cpu
import logging
import platform
import sys

from hpccm.common import cpu_arch
from hpccm.common import container_type
from hpccm.common import linux_distro

# Global variables
g_cpu_arch = cpu_arch.X86_64         # CPU architecture
if platform.machine() == 'aarch64':
  g_cpu_arch = cpu_arch.AARCH64
g_cpu_target = None                  # CPU optimization target
g_ctype = container_type.DOCKER      # Container type
g_linux_distro = linux_distro.UBUNTU # Linux distribution
g_linux_version = Version('16.04') # Linux distribution version
g_singularity_version = Version('2.6') # Singularity version
g_wd = '/var/tmp' # Working directory
g_singularity_tmp_fallback = True    # Singularity / Apptainer behavior flags

def get_cpu_architecture():
[docs]
  """Return the architecture string for the currently configured CPU
  architecture, e.g., `aarch64` or `x86_64`.

  """

  this = sys.modules[__name__]
  if this.g_cpu_arch == cpu_arch.AARCH64:
    return 'aarch64'
  elif this.g_cpu_arch == cpu_arch.X86_64:
    return 'x86_64'
  else: # pragma: no cover
    raise RuntimeError('Unrecognized processor architecture')

def get_cpu_optimization_flags(compiler, version='9999'):
[docs]
  """Return the CPU optimization flags for the target and compiler
  combination.

  Args:
    compiler: A compiler family string recognized by archspec.
    version: The version of the compiler.  The default version is
      `9999`, i.e., assume the compiler supports the latest optimization
      flags.
  """
  this = sys.modules[__name__]

  if not this.g_cpu_target:
    return None

  if this.g_cpu_target not in archspec.cpu.TARGETS:
    logging.warning('unrecognized CPU target "{}"'.format(this.g_cpu_target))
    return None

  try:
    return archspec.cpu.TARGETS[this.g_cpu_target].optimization_flags(compiler, version)
  except Exception as e:
    logging.warning('get_cpu_optimization_flags: {}'.format(e))

  return None

def get_format():
[docs]
  """Return the container format string for the currently configured
  format, e.g., `bash`, `docker`, or `singularity`."""

  this = sys.modules[__name__]

  if this.g_ctype == container_type.BASH:
    return 'bash'
  elif this.g_ctype == container_type.DOCKER:
    return 'docker'
  elif this.g_ctype == container_type.SINGULARITY:
    return 'singularity'
  else: # pragma: no cover
    raise RuntimeError('Unrecognized format')

def set_container_format(ctype):
[docs]
  """Set the container format

  Args:
    ctype (string): 'docker' to specify the Dockerfile format, or
      'singularity' to specify the Singularity definition file format

  Raises:
    RuntimeError: invalid container type argument
  """
  this = sys.modules[__name__]
  if ctype == 'docker':
    this.g_ctype = container_type.DOCKER
  elif ctype == 'singularity':
    this.g_ctype = container_type.SINGULARITY
  else:
    raise RuntimeError('Unrecognized container format: {}'.format(ctype))

def set_cpu_architecture(arch):
[docs]
  """Set the CPU architecture

  In most cases, the `baseimage` primitive should be relied upon to
  set the CPU architecture.  Only use this function if you really
  know what you are doing.

  Args:
    arch (string): Value values are `aarch64`, and `x86_64`.
      `arm` and `arm64v8` are aliases for `aarch64`, and `amd64` and 
      `x86` are aliases for `x86_64`.
  """
  this = sys.modules[__name__]
  if arch == 'aarch64' or arch == 'arm' or arch == 'arm64v8':
    this.g_cpu_arch = cpu_arch.AARCH64
  elif arch == 'x86_64' or arch == 'amd64' or arch == 'x86':
    this.g_cpu_arch = cpu_arch.X86_64
  else:
    logging.warning('Unable to determine the CPU architecture, defaulting to x86_64')
    this.g_cpu_arch = cpu_arch.X86_64

def set_cpu_target(target):
[docs]
  """Set the CPU optimization target

  Args:
    target (string): A CPU microarchitecture string recognized by
      archspec.
  """
  this = sys.modules[__name__]
  this.g_cpu_target = target

def set_linux_distro(distro):
[docs]
  """Set the Linux distribution and version

  In most cases, the `baseimage` primitive should be relied upon to
  set the Linux distribution.  Only use this function if you really
  know what you are doing.

  Args:
    distro (string): Valid values are `centos7`, `centos8`, `rhel7`,
      `rhel8`, `rockylinux8`, `rockylinux9`, `rockylinux10`, `ubuntu16`,
      `ubuntu18`, `ubuntu20`, `ubuntu22`, `ubuntu24`, and `ubuntu26`.
      `ubuntu` is an alias for `ubuntu16`, `centos` is an alias for `centos7`,
      and `rhel` is an alias for `rhel7`.
  """
  this = sys.modules[__name__]
  if distro == 'centos':
    this.g_linux_distro = linux_distro.CENTOS
    this.g_linux_version = Version('7.0')
  elif distro == 'centos7':
    this.g_linux_distro = linux_distro.CENTOS
    this.g_linux_version = Version('7.0')
  elif distro == 'centos8':
    this.g_linux_distro = linux_distro.CENTOS
    this.g_linux_version = Version('8.0')
  elif distro == 'rhel':
    this.g_linux_distro = linux_distro.RHEL
    this.g_linux_version = Version('7.0')
  elif distro == 'rhel7':
    this.g_linux_distro = linux_distro.RHEL
    this.g_linux_version = Version('7.0')
  elif distro == 'rhel8':
    this.g_linux_distro = linux_distro.RHEL
    this.g_linux_version = Version('8.0')
  elif distro == 'rockylinux8':
    this.g_linux_distro = linux_distro.ROCKYLINUX
    this.g_linux_version = Version('8.0')
  elif distro == 'rockylinux9':
    this.g_linux_distro = linux_distro.ROCKYLINUX
    this.g_linux_version = Version('9.0')
  elif distro == 'rockylinux10':
    this.g_linux_distro = linux_distro.ROCKYLINUX
    this.g_linux_version = Version('10.0')
  elif distro == 'ubuntu':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('16.04')
  elif distro == 'ubuntu16':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('16.04')
  elif distro == 'ubuntu18':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('18.04')
  elif distro == 'ubuntu20':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('20.04')
  elif distro == 'ubuntu22':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('22.04')
  elif distro == 'ubuntu24':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('24.04')
  elif distro == 'ubuntu26':
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('26.04')
  else:
    logging.warning('Unable to determine the Linux distribution, defaulting to Ubuntu')
    this.g_linux_distro = linux_distro.UBUNTU
    this.g_linux_version = Version('16.04')

def set_singularity_version(ver):
[docs]
  """Set the Singularity definition file format version

  The Singularity definition file format was extended in version 3.2
  to enable multi-stage builds.  However, these changes are not
  backwards compatible.

  Args:
    ver (string): Singularity definition file format version.
  """
  this = sys.modules[__name__]
  this.g_singularity_version = Version(ver)

def set_working_directory(wd):
[docs]
  """Set the working directory to use for staging inside the container

  Args:
    wd (string): working directory path
  """
  this = sys.modules[__name__]
  this.g_wd = wd

def set_singularity_tmp_fallback(enable=True):
[docs]
  """Enable or disable the automatic %setup fallback for /tmp and /var/tmp
  destinations on Singularity >= 3.6.

  Args:
    enable (bool): True to enable the fallback (default), False to disable.
  """
  this = sys.modules[__name__]
  this.g_singularity_tmp_fallback = enable

def test_cpu_feature_flag(flag):
[docs]
  """Return True or False depending on whether the CPU supports the
  given feature flag

  Args:
    flag: A CPU feature flag, e.g., `avx`.
  """
  this = sys.modules[__name__]

  if this.g_cpu_target not in archspec.cpu.TARGETS:
    logging.warning('unrecognized CPU target "{}"'.format(this.g_cpu_target))
    return False

  if this.g_cpu_target:
    try:
      return flag in archspec.cpu.TARGETS[this.g_cpu_target]
    except Exception as e:
      logging.warning('get_cpu_optimization_flags: {}'.format(e))

  return False