https://flyte.org logo
Powered by
Title
f

Felix Ruess

10/19/2022, 7:21 PM
Hi, how can I run a ContainerTask on a GPU enabled node and set 
runtimeClassName
to 
nvidia
so that it can actually use the GPU? I added 
requests=Resources(gpu="1")
but can I add the runtimeClassName?
k

Ketan (kumare3)

10/19/2022, 8:52 PM
@Felix Ruess this is interesting. We do not support runtimeClassName at the “container” level, but there are 2 options 1. Default runtime class name using pod templates. These are the templates that are used to run all container “raw or other” tasks 2. Use Pod plugin
m

Maarten de Jong

10/20/2022, 6:35 AM
We did it in the following way, not sure if it's most ideal (since you don't get to pick a specific GPU type), but it seems to work for us: 
pod_resources = Resources(cpu="3", mem="20Gi", gpu="1")
export_terrain_texture_container_task = CustomContainerTask(
    requests=pod_resources,
    limits=pod_resources,
)

class CustomContainerTask(ContainerTask):
    def __init__(
        self,
        requests: Optional[Resources] = None,
        limits: Optional[Resources] = None,
        **kwargs: Any,
    ):
        super().__init__(
            <<stuff>>
            requests=requests,
            limits=limits,
            **kwargs,
        )

    def get_container(self, settings: SerializationSettings) -> _task_model.Container:
        env = {**settings.env, **self.environment} if self.environment else settings.env
        return _get_container_definition(
            image=self._image,
            command=self._cmd,
            args=self._args,
            data_loading_config=_task_model.DataLoadingConfig(
                input_path=self._input_data_dir,
                output_path=self._output_data_dir,
                format=self._md_format.value,
                enabled=True,
                io_strategy=self._io_strategy.value if self._io_strategy else None,
            ),
            environment=env,
            cpu_request=self.resources.requests.cpu,
            cpu_limit=self.resources.limits.cpu,
            memory_request=self.resources.requests.mem,
            memory_limit=self.resources.limits.mem,
            gpu_request=self.resources.requests.gpu,
            gpu_limit=self.resources.limits.gpu,
            ephemeral_storage_request=self.resources.requests.ephemeral_storage,
            ephemeral_storage_limit=self.resources.limits.ephemeral_storage,
        )
Where imo the important part is the override for 
get_container
, since the original does not include the passthrough of the GPU resources. i.e. these lines: 
gpu_request=self.resources.requests.gpu,
            gpu_limit=self.resources.limits.gpu,
f

Felix Ruess

10/20/2022, 7:44 AM
@Maarten de Jong thanks for the hint about passing on gpu resources, that sounds like a bug though.
@Ketan (kumare3) I'll look into pod templates, but what I need in the end is that when I request a GPU, limit 
<http://nvidia.com/gpu|nvidia.com/gpu>: "1"
and 
runtimeClassName: nvidia
are added to the container, and if I don't request a GPU, neither are.
k

Ketan (kumare3)

10/20/2022, 2:53 PM
@Felix Ruess also do propose a better UX, remember one goal, we want to make most tasks simple - without need of kubernetes pod specifics. This makes it possible to do fun optimizations (hopefully next year)
f

Felix Ruess

10/20/2022, 4:16 PM
should I open an issue for the workaround that @Maarten de Jong posted? IMHO this really should not be needed and is a bug..
For translating requests with GPU to kubernetes resource requests, tolerations, etc.. how is that currently done? I didn't get how a gpu request is turned into 
nvidia/gpu
or is that "hardcoded" in favor of nvidia gpus? Regarding tolerations, I already setup the 
ExtendedResourceToleration
admission controller which will add tolerations automatically. So in my case I don't have to handle that in flyte separately. One other idea for the 
runtimeClassName
could be to setup something similar which will add that automatically in k8s, basically a customized ExtendedResourceToleration admission controller or just an additional one...
@Maarten de Jong @Ketan (kumare3) https://github.com/flyteorg/flytekit/pull/1249
@Ketan (kumare3) thanks for the pointers! I got it to work with the above fix for flytekit and a custom PodTemplate that sets 
runtimeClassName: nvidia
This works for now and I can run GPU tasks, but also means that flyte tasks will always run with nvidia runtime (and hence only be schedules to nodes which have this runtime and hence a GPU).
flyte-nvidia-pod-template.yaml
k

Ketan (kumare3)

10/20/2022, 8:36 PM
Hmm ya this is not good. Cc @Sören Brunk & @Felix Ruess and @Alireza we should file an issue for this. This definitely seems something we can do in the higher level container api to make things better
5 Views
#ask-the-communityJoin Slack
Powered by