I am currently trying to combine dynamic overrides with ray to dynamically determine the ray cluster size when running a task. But when the task under the dynamic runs, it is considered a python task and never creates a ray cluster.

@dynamic
def request_train_resources_task(
    compressed_dataset: FlyteFile,
    train_parameters: TrainParameters,
) -> TrainOutputs:
    task_config = RayJobConfig(
        head_node_config=HeadNodeConfig(
            ray_start_params={
                "block": "true",
            }
        ),
        worker_node_config=[
            WorkerNodeConfig(
                group_name="ray-group",
                replicas=train_parameters.num_nodes - 1,
            )
        ],
    )
    resources_request = Resources(
        cpu=str(train_parameters.cpus_per_gpu * train_parameters.num_gpus),
        mem=f"{train_parameters.memory_per_gpu * train_parameters.num_gpus}G",
        gpu=str(train_parameters.num_gpus),
        ephemeral_storage=f"{train_parameters.ephemeral_storage}G",
    )
    print(f"task_config: {task_config}")
    print(f"resources_request: {resources_request}")
    return pl_mnist_train_task(
        compressed_dataset=compressed_dataset,
        train_parameters=train_parameters,
    ).with_overrides(
        task_config=task_config,
        requests=resources_request,
    )

we’ll follow up tomorrow a bit more, but the

task_config

isn’t actually read by the with_overrides call.

I noticed something interesting and im guessing it is a bug, but when a ray task requests for resources it overwrites the tolerations set in values_gcp.yaml

we’ll get back to you tomorrow… it’s a bit tricky. the issue is that typically, the normal resources like gpu/cpu are part of the task node overrides

This is the code change i made to test it

@dynamic(requests=Resources(ephemeral_storage="50G"))
def request_train_resources_task(
    compressed_dataset: FlyteFile,
    train_parameters: TrainParameters,
) -> TrainOutputs:
    if train_parameters.use_ray:
        task_config = RayJobConfig(
            head_node_config=HeadNodeConfig(
                ray_start_params={
                    "block": "true",
                }
            ),
            worker_node_config=[
                WorkerNodeConfig(
                    group_name="ray-group",
                    replicas=train_parameters.num_nodes - 1,
                )
            ],
        )
    else:
        task_config = None
    resources_request = Resources(
        cpu=str(train_parameters.cpus_per_gpu * train_parameters.num_gpus),
        mem=f"{train_parameters.memory_per_gpu * train_parameters.num_gpus}G",
        gpu=str(train_parameters.num_gpus),
        ephemeral_storage=f"{train_parameters.ephemeral_storage}G",
    )
    print(f"task_config: {task_config}")
    print(f"resources_request: {resources_request}")
    if train_parameters.use_ray:
        return pl_mnist_train_ray_task(
            compressed_dataset=compressed_dataset,
            train_parameters=train_parameters,
        ).with_overrides(
            task_config=task_config,
            requests=resources_request,
        )
    else:
        return pl_mnist_train_python_task(
            compressed_dataset=compressed_dataset,
            train_parameters=train_parameters,
        ).with_overrides(
            task_config=task_config,
            requests=resources_request,
        )

when changing the ray cluster resources, you’re now changing part of the task definition. this is okay, but let us get back to you on the best way to do it

Alright no rush

yes so this is not possible, not directly at least -

"changed the task config for the under lying task to ray"

So what I mean for that is i made 2 tasks one with a default task type of ray with basically a dummy config for it which will be replaced later and another task whose default task type is python. I did this to just test how dynamically calling the task will work.

@task(
    task_config=RayJobConfig(
        head_node_config=HeadNodeConfig(
            ray_start_params={
                "block": "true",
            }
        ),
        worker_node_config=[
            WorkerNodeConfig(
                group_name="ray-group",
                replicas=0,
            )
        ],
    ),
    secret_requests=[Secret(group="wandb-secrets", key="API_KEY")],
)
def pl_mnist_train_ray_task(...

@task(
    secret_requests=[Secret(group="wandb-secrets", key="API_KEY")],
)
def pl_mnist_train_python_task(

tolerations are a separate concern right?

yes

you’re saying the tolerations

Tolerations:     gpu-type=t4:NoSchedule
                 <http://node.kubernetes.io/not-ready:NoExecute|node.kubernetes.io/not-ready:NoExecute> op=Exists for 300s
                 <http://node.kubernetes.io/unreachable:NoExecute|node.kubernetes.io/unreachable:NoExecute> op=Exists for 300s
                 num-gpus=1:NoSchedule
                 <http://nvidia.com/gpu:NoSchedule|nvidia.com/gpu:NoSchedule> op=Exists

are not correct?

That is the correct toleration

but they’ve been set resource tolerations?

correct

kevin will take a look, maybe the plugin is not calling that somehow.

Another thing I found is that when running a ray task, they cannot find the kube secrets, but the python task can

I found that ray-plugin in propeller doesn’t update the toleration and secret. I will update it, and create a pr.

thanks, so I think the only issue now is the dynamic resources. I tried having a dynamic number of replica workers and only the the dummy config amount of workers is created which is 0 in my case and when I changed the placeholder value to 1 it created a worker. So it would need to be able to dynamically create the correct number of workers

@Tarmily Wen I’m working on this today, will get back to you

sorry I was away. Where is the PR?

The pr looks good to me, is there any way for me to test the plugin before merging?

you could update the plugin version in the flytepropeller, and then build an image

make docker_build

. after that, update the propeller image in your cluster.

I think it is good to merge in

sorry we had a couple tickets come up on the databricks integration we’re working on.

@Tarmily Wen have you tested it. If not, I’ve built a propeller image, you can test it on your side.

pingsutw/flytepropeller:34c4e43574a7fded751c362733508b9f11b8637c

@Tarmily Wen how widespread will this be?

so in the override(task_config=...) that doesn't do the same thing?

it should

So these are my tasks

@task(
    task_config=RayJobConfig(
        head_node_config=HeadNodeConfig(
            ray_start_params={
                "block": "true",
            }
        ),
        worker_node_config=[
            WorkerNodeConfig(
                group_name="ray-group",
                replicas=0,
            )
        ],
    ),
    secret_requests=[Secret(group="wandb-secrets", key="API_KEY")],
    requests=Resources(cpu="22", mem="80G", gpu="1", ephemeral_storage="700G"),
)
def pl_mnist_train_ray_task(
...

@dynamic(requests=Resources(ephemeral_storage="50G"))
def request_train_resources_task(
    compressed_dataset: FlyteFile,
    train_parameters: TrainParameters,
) -> TrainOutputs:
    if train_parameters.use_ray:
        task_config = RayJobConfig(
            head_node_config=HeadNodeConfig(
                ray_start_params={
                    "block": "true",
                }
            ),
            worker_node_config=[
                WorkerNodeConfig(
                    group_name="ray-group",
                    replicas=train_parameters.num_nodes - 1,
                )
            ],
        )
    else:
        task_config = None
    resources_request = Resources(
        cpu=str(train_parameters.cpus_per_gpu * train_parameters.num_gpus),
        mem=f"{train_parameters.memory_per_gpu * train_parameters.num_gpus}G",
        gpu=str(train_parameters.num_gpus),
        ephemeral_storage=f"{train_parameters.ephemeral_storage}G",
    )
    print(f"task_config: {task_config}")
    print(f"resources_request: {resources_request}")
    if train_parameters.use_ray:
        return pl_mnist_train_ray_task(
            compressed_dataset=compressed_dataset,
            train_parameters=train_parameters,
        ).with_overrides(
            task_config=task_config,
            requests=resources_request,
        )

I just give the task a ray config so that it will show up as a ray task. I give it a dummy value of 0 worker replicas. But when it is actually executed from the dynamic it still have 0 workers even when the override task_config specifies workers=1. I have tested this the other way around as well where the dummy value is replicas=1 and the dynamic is set to 0. 1 worker is created but never used

@Tarmily Wen one of the option is that override

task_config

, like

@dynamic
def dynamic_ray(n: int) -> typing.List[int]:
    if n > 2:
        ray_task._task_config.worker_node_config[0].replicas = 6
    else:
        ray_task._task_config.worker_node_config[0].replicas = 8
    return ray_task(n=5)

Okay that workaround works for me just tested it. When do you think this will be fixed?

I just created a pr to support it. https://github.com/flyteorg/flytekit/pull/1410

Thank you!