Apply output_schema only to the last replaced job

[jenniferpeschel]

I just came across line 1264 in core/job.py where it states that:

# only apply output schema if there is no replace.

I was wondering why this design was chosen. Given a simple example:

from pydantic import BaseModel
from jobflow import job, run_locally, Response

class MySchema(BaseModel):
    a: int

@job(output_schema=MySchema)
def job1():
    return Response(3, replace=job2())

@job
def job2():
    return {"a": 3}

run_locally(job1())

I would feel that this should fail because the first job gives an invalid output back but it actually runs without problems. No Exception is raised.

[jenniferpeschel]

I think I understand that any replaced jobs should have the same output schema as an outside user shouldn't be aware of the replace steps but just gets a single output which should conform to the original schema.

This is ensured when replacing with a flow (see function prepare_replace). But a replaced job CAN have a different output schema (line 1386). This feels unintuitive again

[FabiPi3]

If it is a feature, it would be great if there would be an option to do that for Flow as well. Adding an output_schema to the Flow class or so