add openai batch API agent (flyteorg#2353)

* add openai batch endpoint agent Signed-off-by: Samhita Alla <[email protected]> * add jsonl type transformer, modify agent Signed-off-by: Samhita Alla <[email protected]> * revert format changes Signed-off-by: Samhita Alla <[email protected]> * revert iterator edits Signed-off-by: Samhita Alla <[email protected]> * typealias Signed-off-by: Samhita Alla <[email protected]> * add jsonlines Signed-off-by: Samhita Alla <[email protected]> * remove typealias Signed-off-by: Samhita Alla <[email protected]> * update readthedocs python version Signed-off-by: Samhita Alla <[email protected]> * update docs python version Signed-off-by: Samhita Alla <[email protected]> * replace dict with enum Signed-off-by: Samhita Alla <[email protected]> * modify json type; add a check to validate if iterator's empty Signed-off-by: Samhita Alla <[email protected]> * ignore mypy check Signed-off-by: Samhita Alla <[email protected]> * modify JSON type Signed-off-by: Samhita Alla <[email protected]> * move to openai folder Signed-off-by: Samhita Alla <[email protected]> * update readme Signed-off-by: Samhita Alla <[email protected]> * update setup.py Signed-off-by: Samhita Alla <[email protected]> * batch_api to batch, add plugin to setup.py Signed-off-by: Samhita Alla <[email protected]> * modify return type Signed-off-by: Samhita Alla <[email protected]> * modify return type Signed-off-by: Samhita Alla <[email protected]> * fix lint error Signed-off-by: Samhita Alla <[email protected]> * remove guess_python_type Signed-off-by: Samhita Alla <[email protected]> * modify json type Signed-off-by: Samhita Alla <[email protected]> * modify tests Signed-off-by: Samhita Alla <[email protected]> * json to iterator[json] Signed-off-by: Samhita Alla <[email protected]> * update plugin readme Signed-off-by: Samhita Alla <[email protected]> * replace flytefile with jsonlfile Signed-off-by: Samhita Alla <[email protected]> * modify output type of batch; add typealias Signed-off-by: Samhita Alla <[email protected]> * nit Signed-off-by: Samhita Alla <[email protected]> * nit Signed-off-by: Samhita Alla <[email protected]> * replace openai-batch-endpoint with openai-batch Signed-off-by: Samhita Alla <[email protected]> * update readme Signed-off-by: Samhita Alla <[email protected]> * revert to secrets Signed-off-by: Samhita Alla <[email protected]> * fix openai batch code; add json iterator click type Signed-off-by: Samhita Alla <[email protected]> * update readme Signed-off-by: Samhita Alla <[email protected]> * fix types Signed-off-by: Samhita Alla <[email protected]> * add shim tasks Signed-off-by: Samhita Alla <[email protected]> * dict to dataclass; add container image; add guess_python_type to json iterator Signed-off-by: Samhita Alla <[email protected]> * json iterator check in type engine and flytefile Signed-off-by: Samhita Alla <[email protected]> * update image version Signed-off-by: Samhita Alla <[email protected]> * add secret; update json iterator Signed-off-by: Samhita Alla <[email protected]> * add secret to shim task init method Signed-off-by: Samhita Alla <[email protected]> * fix secret Signed-off-by: Samhita Alla <[email protected]> * fix secret Signed-off-by: Samhita Alla <[email protected]> * add secret to dict Signed-off-by: Samhita Alla <[email protected]> * fix logging error; remove iterator copy; remove in flyte entity names Signed-off-by: Samhita Alla <[email protected]> * openai tests Signed-off-by: Samhita Alla <[email protected]> * lint and remove auto spec in openai tests Signed-off-by: Samhita Alla <[email protected]> * fix test Signed-off-by: Samhita Alla <[email protected]> * json key check Signed-off-by: Samhita Alla <[email protected]> * modify input type of upload json data to jsonlfile only Signed-off-by: Samhita Alla <[email protected]> * add jsonl to mime type Signed-off-by: Samhita Alla <[email protected]> * change mime type Signed-off-by: Samhita Alla <[email protected]> * change mime type and fix tests Signed-off-by: Samhita Alla <[email protected]> * add json mime type Signed-off-by: Samhita Alla <[email protected]> * lint Signed-off-by: Samhita Alla <[email protected]> * lint Signed-off-by: Samhita Alla <[email protected]> * incorporate kevin's suggestion Signed-off-by: Samhita Alla <[email protected]> * requests 2.32.2 doesn't work either Signed-off-by: Samhita Alla <[email protected]> --------- Signed-off-by: Samhita Alla <[email protected]>
austin362667 · Jun 15, 2024 · 4f5a468 · 4f5a468
1 parent cf3b44b
commit 4f5a468
Show file tree

Hide file tree

Showing 34 changed files with 1,209 additions and 106 deletions.
diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml
@@ -147,6 +147,29 @@ jobs:
           file: ./plugins/flytekit-sqlalchemy/Dockerfile
           cache-from: type=gha
           cache-to: type=gha,mode=max
+      - name: Prepare OpenAI Batch Image Names
+        id: openai-batch-names
+        uses: docker/metadata-action@v3
+        with:
+          images: |
+            ghcr.io/${{ github.repository_owner }}/flytekit
+          tags: |
+            py${{ matrix.python-version }}-openai-batch-latest
+            py${{ matrix.python-version }}-openai-batch-${{ github.sha }}
+            py${{ matrix.python-version }}-openai-batch-${{ needs.deploy.outputs.version }}
+      - name: Push OpenAI Batch Image to GitHub Registry
+        uses: docker/build-push-action@v2
+        with:
+          context: "./plugins/flytekit-openai/"
+          platforms: linux/arm64, linux/amd64
+          push: ${{ github.event_name == 'release' }}
+          tags: ${{ steps.openai-batch-names.outputs.tags }}
+          build-args: |
+            VERSION=${{ needs.deploy.outputs.version }}
+            PYTHON_VERSION=${{ matrix.python-version }}
+          file: ./plugins/flytekit-openai/Dockerfile.batch
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
 
   build-and-push-flyteagent-images:
     runs-on: ubuntu-latest

diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -1,15 +1,11 @@
 #
-# This file is autogenerated by pip-compile with Python 3.11
+# This file is autogenerated by pip-compile with Python 3.12
 # by the following command:
 #
 #    pip-compile dev-requirements.in
 #
 -e file:.#egg=flytekit
     # via -r dev-requirements.in
-absl-py==2.1.0
-    # via
-    #   tensorboard
-    #   tensorflow-macos
 adlfs==2023.9.0
     # via flytekit
 aiobotocore==2.5.4
@@ -26,12 +22,11 @@ aiosignal==1.3.1
     # via aiohttp
 asttokens==2.4.1
     # via stack-data
-astunparse==1.6.3
-    # via tensorflow-macos
 attrs==23.2.0
     # via
     #   aiohttp
     #   hypothesis
+    #   jsonlines
 autoflake==2.2.1
     # via -r dev-requirements.in
 azure-core==1.30.0
@@ -100,11 +95,7 @@ execnet==2.0.2
 executing==2.0.1
     # via stack-data
 filelock==3.13.1
-    # via
-    #   torch
-    #   virtualenv
-flatbuffers==23.5.26
-    # via tensorflow-macos
+    # via virtualenv
 flyteidl @ git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl
     # via
     #   -r dev-requirements.in
@@ -119,9 +110,6 @@ fsspec==2023.9.2
     #   flytekit
     #   gcsfs
     #   s3fs
-    #   torch
-gast==0.5.4
-    # via tensorflow-macos
 gcsfs==2023.9.2
     # via flytekit
 google-api-core[grpc]==2.16.2
@@ -138,11 +126,8 @@ google-auth==2.27.0
     #   google-cloud-core
     #   google-cloud-storage
     #   kubernetes
-    #   tensorboard
 google-auth-oauthlib==1.2.0
-    # via
-    #   gcsfs
-    #   tensorboard
+    # via gcsfs
 google-cloud-bigquery==3.17.1
     # via -r dev-requirements.in
 google-cloud-bigquery-storage==2.24.0
@@ -157,8 +142,6 @@ google-crc32c==1.5.0
     # via
     #   google-cloud-storage
     #   google-resumable-media
-google-pasta==0.2.0
-    # via tensorflow-macos
 google-resumable-media==2.7.0
     # via
     #   google-cloud-bigquery
@@ -175,14 +158,10 @@ grpcio==1.60.1
     #   flytekit
     #   google-api-core
     #   grpcio-status
-    #   tensorboard
-    #   tensorflow-macos
 grpcio-status==1.60.1
     # via
     #   flytekit
     #   google-api-core
-h5py==3.10.0
-    # via tensorflow-macos
 hypothesis==6.98.2
     # via -r dev-requirements.in
 icdiff==2.0.7
@@ -194,9 +173,7 @@ idna==3.6
     #   requests
     #   yarl
 importlib-metadata==7.0.1
-    # via
-    #   flytekit
-    #   keyring
+    # via flytekit
 iniconfig==2.0.0
     # via pytest
 ipython==8.21.0
@@ -212,38 +189,30 @@ jaraco-classes==3.3.0
 jedi==0.19.1
     # via ipython
 jinja2==3.1.3
-    # via
-    #   -r dev-requirements.in
-    #   torch
+    # via -r dev-requirements.in
 jmespath==1.0.1
     # via botocore
 joblib==1.3.2
     # via
     #   -r dev-requirements.in
     #   flytekit
     #   scikit-learn
+jsonlines==4.0.0
+    # via flytekit
 jsonpickle==3.0.2
     # via flytekit
-keras==2.15.0
-    # via tensorflow-macos
 keyring==24.3.0
     # via flytekit
 keyrings-alt==5.0.0
     # via -r dev-requirements.in
 kubernetes==29.0.0
     # via -r dev-requirements.in
-libclang==16.0.6
-    # via tensorflow-macos
-markdown==3.5.2
-    # via tensorboard
 markdown-it-py==3.0.0
     # via
     #   flytekit
     #   rich
 markupsafe==2.1.5
-    # via
-    #   jinja2
-    #   werkzeug
+    # via jinja2
 marshmallow==3.20.2
     # via
     #   dataclasses-json
@@ -261,14 +230,10 @@ matplotlib-inline==0.1.6
     # via ipython
 mdurl==0.1.2
     # via markdown-it-py
-ml-dtypes==0.2.0
-    # via tensorflow-macos
 mock==5.1.0
     # via -r dev-requirements.in
 more-itertools==10.2.0
     # via jaraco-classes
-mpmath==1.3.0
-    # via sympy
 msal==1.26.0
     # via
     #   azure-datalake-store
@@ -286,28 +251,19 @@ mypy-extensions==1.0.0
     # via
     #   mypy
     #   typing-inspect
-networkx==3.2.1
-    # via torch
 nodeenv==1.8.0
     # via pre-commit
 numpy==1.26.4
     # via
     #   -r dev-requirements.in
-    #   h5py
-    #   ml-dtypes
-    #   opt-einsum
     #   pandas
     #   pyarrow
     #   scikit-learn
     #   scipy
-    #   tensorboard
-    #   tensorflow-macos
 oauthlib==3.2.2
     # via
     #   kubernetes
     #   requests-oauthlib
-opt-einsum==3.3.0
-    # via tensorflow-macos
 orjson==3.9.12
     # via -r dev-requirements.in
 packaging==23.2
@@ -318,7 +274,6 @@ packaging==23.2
     #   msal-extensions
     #   pytest
     #   setuptools-scm
-    #   tensorflow-macos
 pandas==2.2.0
     # via -r dev-requirements.in
 parso==0.8.3
@@ -353,8 +308,6 @@ protobuf==4.23.4
     #   grpcio-status
     #   proto-plus
     #   protoc-gen-openapiv2
-    #   tensorboard
-    #   tensorflow-macos
 protoc-gen-openapiv2==0.0.1
     # via flyteidl
 ptyprocess==0.7.0
@@ -379,7 +332,9 @@ pygments==2.17.2
     #   ipython
     #   rich
 pyjwt[crypto]==2.8.0
-    # via msal
+    # via
+    #   msal
+    #   pyjwt
 pytest==7.4.4
     # via
     #   -r dev-requirements.in
@@ -436,7 +391,6 @@ requests==2.31.0
     #   kubernetes
     #   msal
     #   requests-oauthlib
-    #   tensorboard
 requests-oauthlib==1.3.1
     # via
     #   google-auth-oauthlib
@@ -460,40 +414,18 @@ setuptools-scm==8.0.4
 six==1.16.0
     # via
     #   asttokens
-    #   astunparse
     #   azure-core
-    #   google-pasta
     #   isodate
     #   kubernetes
     #   python-dateutil
-    #   tensorboard
-    #   tensorflow-macos
 sortedcontainers==2.4.0
     # via hypothesis
 stack-data==0.6.3
     # via ipython
 statsd==3.3.0
     # via flytekit
-sympy==1.12
-    # via torch
-tensorboard==2.15.1
-    # via tensorflow-macos
-tensorboard-data-server==0.7.2
-    # via tensorboard
-tensorflow==2.15.0 ; python_version < "3.12"
-    # via -r dev-requirements.in
-tensorflow-estimator==2.15.0
-    # via tensorflow-macos
-tensorflow-io-gcs-filesystem==0.34.0
-    # via tensorflow-macos
-tensorflow-macos==2.15.0
-    # via tensorflow
-termcolor==2.4.0
-    # via tensorflow-macos
 threadpoolctl==3.2.0
     # via scikit-learn
-torch==2.2.0 ; python_version < "3.12"
-    # via -r dev-requirements.in
 traitlets==5.14.1
     # via
     #   ipython
@@ -519,8 +451,6 @@ typing-extensions==4.9.0
     #   mypy
     #   rich-click
     #   setuptools-scm
-    #   tensorflow-macos
-    #   torch
     #   typing-inspect
 typing-inspect==0.9.0
     # via dataclasses-json
@@ -541,14 +471,8 @@ websocket-client==1.7.0
     # via
     #   docker
     #   kubernetes
-werkzeug==3.0.1
-    # via tensorboard
-wheel==0.42.0
-    # via astunparse
 wrapt==1.14.1
-    # via
-    #   aiobotocore
-    #   tensorflow-macos
+    # via aiobotocore
 yarl==1.9.4
     # via aiohttp
 zipp==3.17.0

diff --git a/docs/source/plugins/index.rst b/docs/source/plugins/index.rst
@@ -31,6 +31,7 @@ Plugin API reference
 * :ref:`MLflow <mlflow>` - MLflow API reference
 * :ref:`DuckDB <duckdb>` - DuckDB API reference
 * :ref:`SageMaker Inference <awssagemaker_inference>` - SageMaker Inference API reference
+* :ref:`OpenAI <openai>` - OpenAI API reference
 
 .. toctree::
    :maxdepth: 2
@@ -63,3 +64,4 @@ Plugin API reference
    MLflow <mlflow>
    DuckDB <duckdb>
    SageMaker Inference <awssagemaker_inference>
+   OpenAI <openai>
diff --git a/docs/source/plugins/openai.rst b/docs/source/plugins/openai.rst
@@ -0,0 +1,12 @@
+.. _openai:
+
+################
+OpenAI reference
+################
+
+.. tags:: Integration, OpenAI
+
+.. automodule:: flytekitplugins.openai
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:
diff --git a/docs/source/types.builtins.iterator.rst b/docs/source/types.builtins.iterator.rst
@@ -0,0 +1,4 @@
+.. automodule:: flytekit.types.iterator
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:
diff --git a/docs/source/types.extend.rst b/docs/source/types.extend.rst
@@ -14,6 +14,7 @@ Refer to the :ref:`extensibility contribution guide <cookbook:advanced_custom_ty
    types.builtins.structured
    types.builtins.file
    types.builtins.directory
+   types.builtins.iterator
    extras.pytorch
    extras.tensorflow
    extras.sklearn