diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py index e3558058a1..d16c6856b1 100644 --- a/flytekit/extend/backend/base_agent.py +++ b/flytekit/extend/backend/base_agent.py @@ -215,7 +215,7 @@ async def _get(self, resource_meta: bytes) -> GetTaskResponse: else: res = self._agent.get(grpc_ctx, resource_meta) state = res.resource.state - logger.info(f"Task state: {state}") + logger.info(f"Task state: {state}, State message: {res.resource.message}") return res def signal_handler(self, resource_meta: bytes, signum: int, frame: FrameType) -> typing.Any: diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index 727ce85143..88425c012f 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1,4 +1,4 @@ -flyteidl>=1.3.0 +flyteidl>=1.10.0 -e file:../../.#egg=flytekitplugins-pod&subdirectory=plugins/flytekit-k8s-pod -e file:../../.#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark -e file:../../.#egg=flytekitplugins-awsbatch&subdirectory=plugins/flytekit-aws-batch diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt index beabff648a..96edd53601 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with Python 3.10 +# This file is autogenerated by pip-compile with Python 3.9 # by the following command: # # pip-compile dev-requirements.in @@ -10,63 +10,95 @@ # via -r dev-requirements.in -e file:../../.#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark # via -r dev-requirements.in -aiohttp==3.8.5 - # via flytekitplugins-spark +adlfs==2023.10.0 + # via flytekit +aiobotocore==2.5.4 + # via s3fs +aiohttp==3.9.0 + # via + # adlfs + # aiobotocore + # flytekitplugins-spark + # gcsfs + # s3fs +aioitertools==0.11.0 + # via aiobotocore aiosignal==1.3.1 # via aiohttp -arrow==1.2.3 - # via jinja2-time +arrow==1.3.0 + # via cookiecutter async-timeout==4.0.3 # via aiohttp attrs==23.1.0 # via aiohttp +azure-core==1.29.5 + # via + # adlfs + # azure-identity + # azure-storage-blob +azure-datalake-store==0.0.53 + # via adlfs +azure-identity==1.15.0 + # via adlfs +azure-storage-blob==12.19.0 + # via adlfs binaryornot==0.4.4 # via cookiecutter -cachetools==5.2.0 +botocore==1.31.17 + # via aiobotocore +cachetools==5.3.2 # via google-auth -certifi==2022.12.7 +certifi==2023.11.17 # via # kubernetes # requests -cffi==1.15.1 - # via cryptography -chardet==5.1.0 - # via binaryornot -charset-normalizer==2.1.1 +cffi==1.16.0 # via - # aiohttp - # requests -click==8.1.3 + # azure-datalake-store + # cryptography +chardet==5.2.0 + # via binaryornot +charset-normalizer==3.3.2 + # via requests +click==8.1.7 # via # cookiecutter # flytekit -cloudpickle==2.2.0 + # rich-click +cloudpickle==3.0.0 # via flytekit -cookiecutter==2.1.1 +cookiecutter==2.4.0 # via flytekit -croniter==1.3.8 +croniter==2.0.1 # via flytekit -cryptography==38.0.4 - # via pyopenssl -dataclasses-json==0.5.7 +cryptography==41.0.5 + # via + # azure-identity + # azure-storage-blob + # msal + # pyjwt + # pyopenssl + # secretstorage +dataclasses-json==0.5.9 # via flytekit decorator==5.1.1 - # via retry -deprecated==1.2.13 + # via gcsfs +deprecated==1.2.14 # via flytekit -diskcache==5.4.0 +diskcache==5.6.3 # via flytekit -docker==6.0.1 +docker==6.1.3 # via flytekit docker-image-py==0.1.12 # via flytekit docstring-parser==0.15 # via flytekit -flyteidl==1.3.0 +flyteidl==1.10.0 # via # -r dev-requirements.in # flytekit -flytekit==1.3.0b2 + # flytekitplugins-spark +flytekit==1.10.1 # via # flytekitplugins-awsbatch # flytekitplugins-pod @@ -75,181 +107,271 @@ frozenlist==1.4.0 # via # aiohttp # aiosignal -google-auth==2.15.0 - # via kubernetes -googleapis-common-protos==1.57.0 +fsspec==2023.9.2 + # via + # adlfs + # flytekit + # gcsfs + # s3fs +gcsfs==2023.9.2 + # via flytekit +gitdb==4.0.11 + # via gitpython +gitpython==3.1.40 + # via flytekit +google-api-core==2.14.0 + # via + # google-cloud-core + # google-cloud-storage +google-auth==2.23.4 + # via + # gcsfs + # google-api-core + # google-auth-oauthlib + # google-cloud-core + # google-cloud-storage + # kubernetes +google-auth-oauthlib==1.1.0 + # via gcsfs +google-cloud-core==2.3.3 + # via google-cloud-storage +google-cloud-storage==2.13.0 + # via gcsfs +google-crc32c==1.5.0 + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.6.0 + # via google-cloud-storage +googleapis-common-protos==1.61.0 # via # flyteidl + # flytekit + # google-api-core # grpcio-status -grpcio==1.51.1 +grpcio==1.59.3 # via # flytekit # grpcio-status -grpcio-status==1.51.1 +grpcio-status==1.59.3 # via flytekit idna==3.4 # via # requests # yarl -importlib-metadata==5.1.0 +importlib-metadata==6.8.0 # via # flytekit # keyring -jaraco-classes==3.2.3 +isodate==0.6.1 + # via azure-storage-blob +jaraco-classes==3.3.0 # via keyring -jinja2==3.1.2 +jeepney==0.8.0 # via - # cookiecutter - # jinja2-time -jinja2-time==0.2.0 + # keyring + # secretstorage +jinja2==3.1.2 # via cookiecutter -joblib==1.2.0 +jmespath==1.0.1 + # via botocore +joblib==1.3.2 # via flytekit -keyring==23.11.0 +jsonpickle==3.0.2 # via flytekit -kubernetes==25.3.0 - # via flytekitplugins-pod -markupsafe==2.1.1 +keyring==24.3.0 + # via flytekit +kubernetes==28.1.0 + # via + # flytekit + # flytekitplugins-pod +markdown-it-py==3.0.0 + # via rich +markupsafe==2.1.3 # via jinja2 -marshmallow==3.19.0 +marshmallow==3.20.1 # via # dataclasses-json # marshmallow-enum # marshmallow-jsonschema marshmallow-enum==1.5.1 - # via dataclasses-json + # via + # dataclasses-json + # flytekit marshmallow-jsonschema==0.13.0 # via flytekit -more-itertools==9.0.0 +mashumaro==3.10 + # via flytekit +mdurl==0.1.2 + # via markdown-it-py +more-itertools==10.1.0 # via jaraco-classes +msal==1.25.0 + # via + # azure-datalake-store + # azure-identity + # msal-extensions +msal-extensions==1.0.0 + # via azure-identity multidict==6.0.4 # via # aiohttp # yarl -mypy-extensions==0.4.3 +mypy-extensions==1.0.0 # via typing-inspect -natsort==8.2.0 +natsort==8.4.0 # via flytekit -numpy==1.23.5 +numpy==1.24.4 # via + # flytekit # pandas # pyarrow oauthlib==3.2.2 - # via requests-oauthlib -packaging==22.0 + # via + # kubernetes + # requests-oauthlib +packaging==23.2 # via # docker # marshmallow -pandas==1.5.2 +pandas==1.5.3 # via flytekit -protobuf==4.21.11 +portalocker==2.8.2 + # via msal-extensions +protobuf==4.24.4 # via # flyteidl + # flytekit + # google-api-core # googleapis-common-protos # grpcio-status # protoc-gen-swagger protoc-gen-swagger==0.1.0 # via flyteidl -py==1.11.0 - # via retry -py4j==0.10.9.5 +py4j==0.10.9.7 # via pyspark pyarrow==10.0.1 # via flytekit -pyasn1==0.4.8 +pyasn1==0.5.1 # via # pyasn1-modules # rsa -pyasn1-modules==0.2.8 +pyasn1-modules==0.3.0 # via google-auth pycparser==2.21 # via cffi -pyopenssl==22.1.0 +pygments==2.17.1 + # via rich +pyjwt[crypto]==2.8.0 + # via + # msal + # pyjwt +pyopenssl==23.3.0 # via flytekit -pyspark==3.3.0 +pyspark==3.5.0 # via flytekitplugins-spark python-dateutil==2.8.2 # via # arrow + # botocore # croniter # flytekit # kubernetes # pandas -python-json-logger==2.0.4 +python-json-logger==2.0.7 # via flytekit -python-slugify==7.0.0 +python-slugify==8.0.1 # via cookiecutter pytimeparse==1.1.8 # via flytekit -pytz==2022.6 +pytz==2023.3.post1 # via + # croniter # flytekit # pandas -pyyaml==6.0 +pyyaml==6.0.1 # via # cookiecutter # flytekit # kubernetes -regex==2022.10.31 +regex==2023.10.3 # via docker-image-py -requests==2.28.1 +requests==2.31.0 # via + # azure-core + # azure-datalake-store # cookiecutter # docker # flytekit + # gcsfs + # google-api-core + # google-cloud-storage # kubernetes + # msal # requests-oauthlib - # responses requests-oauthlib==1.3.1 - # via kubernetes -responses==0.22.0 - # via flytekit -retry==0.9.2 + # via + # google-auth-oauthlib + # kubernetes +rich==13.7.0 + # via + # cookiecutter + # flytekit + # rich-click +rich-click==1.7.1 # via flytekit rsa==4.9 # via google-auth +s3fs==2023.9.2 + # via flytekit +secretstorage==3.3.3 + # via keyring six==1.16.0 # via - # google-auth + # azure-core + # isodate # kubernetes # python-dateutil +smmap==5.0.1 + # via gitdb sortedcontainers==2.4.0 # via flytekit statsd==3.3.0 # via flytekit text-unidecode==1.3 # via python-slugify -toml==0.10.2 - # via responses -types-toml==0.10.8.1 - # via responses -typing-extensions==4.4.0 +types-python-dateutil==2.8.19.14 + # via arrow +typing-extensions==4.8.0 # via + # aioitertools + # azure-core + # azure-storage-blob # flytekit + # mashumaro + # rich-click # typing-inspect -typing-inspect==0.8.0 +typing-inspect==0.9.0 # via dataclasses-json -urllib3==1.26.13 +urllib3==1.26.18 # via + # botocore # docker # flytekit # kubernetes # requests - # responses -websocket-client==1.4.2 +websocket-client==1.6.4 # via # docker # kubernetes -wheel==0.38.4 +wheel==0.41.3 # via flytekit -wrapt==1.14.1 +wrapt==1.16.0 # via + # aiobotocore # deprecated # flytekit -yarl==1.9.2 +yarl==1.9.3 # via aiohttp -zipp==3.11.0 +zipp==3.17.0 # via importlib-metadata - -# The following packages are considered to be unsafe in a requirements file: -# setuptools diff --git a/plugins/flytekit-spark/flytekitplugins/spark/agent.py b/plugins/flytekit-spark/flytekitplugins/spark/agent.py index ce0148ad2f..4822844a18 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/agent.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/agent.py @@ -70,10 +70,15 @@ async def async_get(self, context: grpc.ServicerContext, resource_meta: bytes) - response = await resp.json() cur_state = PENDING - if response.get("state") and response["state"].get("result_state"): - cur_state = convert_to_flyte_state(response["state"]["result_state"]) - - return GetTaskResponse(resource=Resource(state=cur_state)) + message = "" + state = response.get("state") + if state: + if state.get("result_state"): + cur_state = convert_to_flyte_state(state["result_state"]) + if state.get("state_message"): + message = state["state_message"] + + return GetTaskResponse(resource=Resource(state=cur_state, message=message)) async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteTaskResponse: metadata = pickle.loads(resource_meta) diff --git a/plugins/flytekit-spark/setup.py b/plugins/flytekit-spark/setup.py index c0f9d8c8b4..65fa580170 100644 --- a/plugins/flytekit-spark/setup.py +++ b/plugins/flytekit-spark/setup.py @@ -4,7 +4,7 @@ microlib_name = f"flytekitplugins-{PLUGIN_NAME}" -plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0", "aiohttp"] +plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0", "aiohttp", "flyteidl>=1.10.0"] __version__ = "0.0.0+develop" diff --git a/plugins/flytekit-spark/tests/test_agent.py b/plugins/flytekit-spark/tests/test_agent.py index 1b3941d7f6..bb4db3ccf5 100644 --- a/plugins/flytekit-spark/tests/test_agent.py +++ b/plugins/flytekit-spark/tests/test_agent.py @@ -113,7 +113,7 @@ async def test_databricks_agent(): ) mock_create_response = {"run_id": "123"} - mock_get_response = {"run_id": "123", "state": {"result_state": "SUCCESS"}} + mock_get_response = {"run_id": "123", "state": {"result_state": "SUCCESS", "state_message": "OK"}} mock_delete_response = {} create_url = "https://test-account.cloud.databricks.com/api/2.0/jobs/runs/submit" get_url = "https://test-account.cloud.databricks.com/api/2.0/jobs/runs/get?run_id=123" @@ -127,6 +127,7 @@ async def test_databricks_agent(): res = await agent.async_get(ctx, metadata_bytes) assert res.resource.state == SUCCEEDED assert res.resource.outputs == literals.LiteralMap({}).to_flyte_idl() + assert res.resource.message == "OK" mocked.post(delete_url, status=200, payload=mock_delete_response) await agent.async_delete(ctx, metadata_bytes)