feat: support docarray v2

[samsja]

Context

this PR is a PoC to support docarray v2 in jina.

This pr make jina compatible with v2 while keeping v1 support as well.

PS: Some of the features are yet not compatible with v2, like filtering etc, but this is an entry point for letting people play with v2.

what this pr do

allow to do

from jina import Flow, Executor, requests
from docarray import BaseDocument, DocumentArray
from docarray.documents import Image
from docarray.typing import AnyTensor

class InputDoc(BaseDocument):
    img: Image

class OutputDoc(BaseDocument):
    embedding: AnyTensor

class MyExec(Executor):
    @requests(on='/bar')
    def bar(
        self, docs: DocumentArray[InputDoc], **kwargs
    ) -> DocumentArray[OutputDoc]:

        docs_return = DocumentArray[OutputDoc](
            [OutputDoc(embedding=np.zeros((100, 1))) for _ in range(len(docs))]
        )

        return docs_return

with Flow().add(uses=MyExec) as f:
    docs = f.post(
        on='/bar',
        inputs=InputDoc(img=Image(tensor=np.zeros((3, 224, 224)))),
        return_type=DocumentArray[OutputDoc],
    )
    assert docs[0].embedding.shape == (100, 1)
    assert docs.__class__.document_type == OutputDoc

FastAPI:

import pytest

from fastapi import FastAPI
from httpx import AsyncClient

from docarray import BaseDocument
from docarray.base_document import DocumentResponse
from docarray.documents import Image
from docarray.typing import AnyTensor


import numpy as np


class InputDoc(BaseDocument):
    img: Image


class OutputDoc(BaseDocument):
    embedding: AnyTensor


input_doc = InputDoc(img=Image(tensor=np.zeros((3, 224, 224))))

app = FastAPI()


@app.post("/embed/", response_model=OutputDoc, response_class=DocumentResponse)
async def embed(doc: InputDoc) -> OutputDoc:
    doc = OutputDoc(embedding=np.zeros((100, 1)))
    return doc


async with AsyncClient(app=app, base_url="http://test") as ac:
    response = await ac.post("/embed/", data=input_doc.json())

doc = OutputDoc.parse_raw(response.content.decode())
assert doc.embedding.shape == (100, 1)

[alexcg1]

A few things I missed on the first round

[alexcg1]

and...what?

[alexcg1]

Or was this an editing glitch on my part from earlier?

[samsja]

euh yeah I think so ?

[samsja]

okay actually it is on the next line nothing wrong

[alexcg1]

Uh, I still see the dangling and .

[github-actions]

📝 Docs are deployed on https://feat-docarray-v2--jina-docs.netlify.app 🎉

[alexcg1]

LGTM 👍

[mitkotak]

Hey coming here after reading this. Was wondering if the subtype approach that you implemented works well with TorchDynamo ?

For context, I have a tuple that I wanna attach to my torch.Tensor and be able to slice my tensor based on the tuple inside torch.compile fullgraph=True. Was wondering if you hit on any sharp edges that would be good to know ?

Thanks !

[samsja]

I did not try torch compile / torch dynamo related things with this approach. I would think that it is working