Python requests_gpu Examples

Example #1

File: test_accelerators.py Project: zhaohb/polyaxon

 def test_requests_gpu(self):
     assert (requests_gpu(
         k8s_schemas.V1ResourceRequirements(limits={"cpu": 1})) is False)
     assert (requests_gpu(
         k8s_schemas.V1ResourceRequirements(limits={"amd.com/gpu": 1})) is
             True)
     assert (requests_gpu(
         k8s_schemas.V1ResourceRequirements(requests={"nvidia.com/gpu": 1}))
             is True)

Example #2

File: env_vars.py Project: opentechfn/polyaxon

def get_resources_env_vars(
    resources: k8s_schemas.V1ResourceRequirements,
) -> List[k8s_schemas.V1EnvVar]:
    env_vars = []
    has_gpu = requests_gpu(resources)

    # Fix https://github.com/kubernetes/kubernetes/issues/59629
    # When resources.gpu.limits is not set or set to 0, we explicitly
    # pass NVIDIA_VISIBLE_DEVICES=none into container to avoid exposing GPUs.
    if not has_gpu:
        env_vars.append(
            k8s_schemas.V1EnvVar(name="NVIDIA_VISIBLE_DEVICES", value="none"))

    return env_vars