New RunLLM response generated: >>>*# Question* How do we specify resources (i.e. @task(requests=Resources(mem="1Gi", cpu="1", gpu="1"))) when using raw ContainerTask? # Answer To specify resources for a

ContainerTask

, you can set the

requests

and

limits

parameters using the

Resources

object. This can be done by extending the

ContainerTask

class and overriding the

get_container

method to include GPU resources. Here is an example:

from flytekit import Resources, ContainerTask

class CustomContainerTask(ContainerTask):
    def __init__(self, requests: Resources = None, limits: Resources = None, **kwargs):
        super().__init__(requests=requests, limits=limits, **kwargs)

    def get_container(self, settings):
        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,
        )

pod_resources = Resources(cpu="1", mem="1Gi", gpu="1")
my_container_task = CustomContainerTask(
    requests=pod_resources,
    limits=pod_resources,
)

This example ensures that the GPU resources are passed correctly to the container. Sources: - Flyte Discussion on GPU Resources - Flyte Optimizing Tasks Slack Link https://flyte-org.slack.com/archives/CP2HDHKE1/p1719475838.767279 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.