https://flyte.org
Join Slack
How are people overriding RayJob and PyTorch task_configs today? I want to modify PyTorch Replicas a...
# flyte-support
p
How are people overriding RayJob and PyTorch task_configs today? I want to modify PyTorch Replicas and RayJob Worker Replicas but there are only checks for valid PythonTasks and ArrayTasks
@clean-glass-36808, @cool-lifeguard-49380 any ideas?
c
You can configure a pod template on the RayJob workers to customize things, not 100% sure what you're asking about tho
ie.
Copy code
head_pod_spec = {
    "containers": [
        {
            "name": "ray-head",
            "resources": {"requests": {"cpu": "1", "memory": "4Gi"}, "limits": {"cpu": "1", "memory": "4Gi"}},
        }
    ]
}
worker_a_pod_spec = {
    "containers": [
        {
            "name": "ray-worker",
            "resources": {"requests": {"cpu": "1", "memory": "3Gi"}, "limits": {"cpu": "1", "memory": "3Gi"}},
        }
    ]
}

worker_b_pod_spec = {
    "containers": [
        {
            "name": "ray-worker",
            "resources": {"requests": {"cpu": "1", "memory": "2Gi"}, "limits": {"cpu": "1", "memory": "2Gi"}},
        }
    ]
}

ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(ray_start_params={"log-color": "True"}, k8s_pod=K8sPod(pod_spec=head_pod_spec)),
    worker_node_config=[
        WorkerNodeConfig(group_name="ray-group-a", replicas=1, k8s_pod=K8sPod(pod_spec=worker_a_pod_spec)),
        WorkerNodeConfig(group_name="ray-group-b", replicas=2, k8s_pod=K8sPod(pod_spec=worker_b_pod_spec)),
    ],
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=3600,
)
p
I mean configuring using 
with_overrides()
c
Oh, I'm not familiar with that function or what it does
p
What I'm trying to do is dynamically change the RayJobConfig based on user inputs
Copy code
@task(
    container_image=custom_image,
    task_config=RayJobConfig(
        head_node_config=xxxx,
        worker_node_config=yyyy,
        shutdown_after_job_finishes=True,
    ),
    pod_template_name=<>,
)
def task_fun()
I am trying to override this as follows, during the task call:
Copy code
custom_task_config = RayJobConfig(
            head_node_config=xxxx,
            worker_node_config=zzzz,
            shutdown_after_job_finishes=True,
)
  task_fun(task_params).with_overrides(task_config=custom_task_config)
Another idea came to mind, just doing this:
Copy code
with open('config.yaml', 'r', encoding='utf-8') as f:
    config = yaml.safe_load(f)

@task(
    container_image=RL_IMAGE,
    task_config=RayJobConfig(
        head_node_config=config['head_config'],
        worker_node_config=config['worker_config'],
        shutdown_after_job_finishes=True,
    ),
I guess I'm more asking how developers are currently dynamically setting task configurations
f
you can only adjust config in dynamic
p
I tried using dynamic as well, but it didn't change anything. It works for PythonTasks, but not for RayTasks or PyTorchTasks
c
Yes, overriding the task_config via 
with_overrides
is unfortunately not supported yet.
you can only adjust config in dynamic
A dynamic task wrapping a single pytorch task with a single with_overrides. Not multiple uses of with_overrides.
f
Ohh ya- @purple-father-70173 we are working on something new a major update that addresses all these usecases
17 Views
Open in Slack
PreviousNext
Powered by

Flyte

Flyte enables production-grade orchestration for machine learning workflows and data processing created to accelerate local workflows to production.

Powered by