Extending the Ray Serve integration to allow attributes for Serve deployments

[zoltan-fedor]

Related Issue(s): closes #2917

We should be able to set Ray Serve attributes for the nodes of pipelines, like amount of GPU to use, max_concurrent_queries, etc.

Now this is possible from the pipeline yaml file for each node of the pipeline.

Proposed changes:

• Extending the Ray Serve integration to allow attributes to be provided for the Serve deployments. Attributes like GPU usage (incl partial GPU usage), num_concurrent_queries and so on. See full list of attributes for the ray.serve.deployment() method in the Ray Serve API documentation

Pre-flight checklist

• I have read the contributors guidelines
• I have enabled actions on my fork
• If this is a code change, I added tests or updated existing ones
• If this is a code change, I updated the docstrings

[ZanSara]

Hey @zoltan-fedor! Thank you for this PR. I will review it shortly 🙂

[ZanSara]

There are a couple of catches here.

The core issue is that you added the deployment_kwargs block to the general pipeline schema, without restricting it to Ray pipelines. This needs two actions to be fixed:

1. We need a pipeline test that verifies that Ray attributes cannot be added to base pipelines. It could be something as simple as adding the following test in tests/pipelines/test_pipeline.yaml.py, approx line 996:

def test_load_yaml_ray_args_in_pipeline(tmp_path):
    with pytest.raises(PipelineConfigError) as e:
      pipeline = Pipeline.load_from_yaml(
          SAMPLES_PATH / "pipeline" / "ray.haystack-pipeline.yml", pipeline_name="ray_query_pipeline",
      )

3. We need to restrict this argument to Ray pipelines. This is rather easy, you just need to add deployment_kwargs alongside replicas in this part of the schema:

   haystack/haystack/nodes/_json_schema.py

   Lines 312 to 321 in b78db1c

   	"oneOf": [
   	{
   	"not": {"required": ["extras"]},
   	"properties": {
   	"pipelines": {
   	"title": "Pipelines",
   	"items": {"properties": {"nodes": {"items": {"not": {"required": ["replicas"]}}}}},
   	}
   	},
   	},

I hope it's all clear! Let me know if there's any question 🙂

---

One last thing: how about renaming the new param to serve_deployment_kwargs? Just to clarify it's related to Ray Serve. Otherwise deployment_kwargs sounds very generic.

[zoltan-fedor]

HI @ZanSara ,
As requested, I have renamed deployment_kwargs to serve_deployment_kwargs.

I have also added serve_deployment_kwargs to :

"oneOf": [ 
     { 
         "not": {"required": ["extras"]}, 
         "properties": { 
             "pipelines": { 
                 "title": "Pipelines", 
                 "items": {"properties": {"nodes": {"items": {"not": {"required": ["replicas", "serve_deployment_kwargs"]}}}}}, 
             } 
         }, 
     }, 

But that will not make the exception tested in the test to be thrown, so this test will fail:

def test_load_yaml_ray_args_in_pipeline(tmp_path):
    with pytest.raises(PipelineConfigError) as e:
      pipeline = Pipeline.load_from_yaml(
          SAMPLES_PATH / "pipeline" / "ray.haystack-pipeline.yml", pipeline_name="ray_query_pipeline",
      )

Any recommendations?

[ZanSara]

You're right about the test! We should add a check in the pipeline loading for the extras: ray. I'm going to add a comment about where I'd like to see this check added, hold on

[ZanSara]

Ok that's not very easy to explain 😅 I'll add a commit myself to do that 👍

[zoltan-fedor]

Let me finish the removal of the replicas

[zoltan-fedor]

I have made the replicas change by moving it below the new serve_deployment_kwargs, but I don't think I am allowed to add the breaking label to this PR. Would you mind adding that?

[ZanSara]

Let me finish the removal of the replicas

Sorry! I haven't seen this. I'm done anyway, I'll be back in a while to see how the tests look like.

[ZanSara]

Great work, thanks for it! 😊 I have no remarks.