ci: Add Python 3.9 to build matrix

.github/workflows/pythonbuild.yml

@@ -28,7 +28,7 @@ jobs:
           if [[ ${{ github.event_name }} == "schedule" ]]; then
             echo "python_versions=[\"3.8\",\"3.9\",\"3.10\",\"3.11\",\"3.12\"]" >> $GITHUB_ENV
           else
-            echo "python_versions=[\"3.12\"]" >> $GITHUB_ENV
+            echo "python_versions=[\"3.9\", \"3.12\"]" >> $GITHUB_ENV
           fi
 
   build:
@@ -128,6 +128,8 @@ jobs:
             pandas: "pandas<2.0.0"
           - numpy: "numpy<2.0.0"
             pandas: "pandas>=2.0.0"
+          - numpy: "numpy>=2.0.0"
+            python-version: "3.8"
 
     steps:
       - uses: actions/checkout@v4
@@ -248,6 +250,8 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache  # sandbox has disk pressure, so we need to clear the cache to get more disk space.
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
@@ -354,6 +358,10 @@ jobs:
           - flytekit-vaex
           - flytekit-whylogs
         exclude:
+          - python-version: 3.8
+            plugin-names: "flytekit-aws-sagemaker"
+          - python-version: 3.9
+            plugin-names: "flytekit-aws-sagemaker"
           # flytekit-modin depends on ray which does not have a 3.11 wheel yet.
           # Issue tracked in https://github.com/ray-project/ray/issues/27881
           - python-version: 3.11

Makefile

@@ -119,5 +119,6 @@ build-dev: export PLATFORM ?= linux/arm64
 build-dev: export REGISTRY ?= localhost:30000
 build-dev: export PYTHON_VERSION ?= 3.12
 build-dev: export PSEUDO_VERSION ?= $(shell python -m setuptools_scm)
+build-dev: export TAG ?= dev
 build-dev:
 	docker build --platform ${PLATFORM} --push . -f Dockerfile.dev -t ${REGISTRY}/flytekit:${TAG} --build-arg PYTHON_VERSION=${PYTHON_VERSION} --build-arg PSEUDO_VERSION=${PSEUDO_VERSION}

dev-requirements.in

@@ -22,7 +22,7 @@ setuptools_scm
 pytest-icdiff
 
 # Tensorflow is not available for python 3.12 yet: https://github.com/tensorflow/tensorflow/issues/62003
-tensorflow; python_version<'3.12'
+tensorflow<=2.15.1; python_version<'3.12'
 # Newer versions of torch bring in nvidia dependencies that are not present in windows, so
 # we put this constraint while we do not have per-environment requirements files
 torch<=1.12.1; python_version<'3.11'

flytekit/core/array_node_map_task.py

@@ -4,10 +4,10 @@
 import logging
 import math
 import os  # TODO: use flytekit logger
-import typing
 from contextlib import contextmanager
 from typing import Any, Dict, List, Optional, Set, Union, cast
 
+import typing_extensions
 from flyteidl.core import tasks_pb2
 
 from flytekit.configuration import SerializationSettings
@@ -70,7 +70,7 @@ def __init__(
             transformer = TypeEngine.get_transformer(v)
             if isinstance(transformer, FlytePickleTransformer):
                 if is_annotated(v):
-                    for annotation in typing.get_args(v)[1:]:
+                    for annotation in typing_extensions.get_args(v)[1:]:
                         if isinstance(annotation, pickle.BatchSize):
                             raise ValueError("Choosing a BatchSize for map tasks inputs is not supported.")
 

plugins/flytekit-airflow/tests/test_agent.py

@@ -75,6 +75,7 @@ async def test_airflow_agent():
         "This is deprecated!",
         True,
         "A",
+        None
     )
 
     interfaces = interface_models.TypedInterface(inputs={}, outputs={})

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py

@@ -15,7 +15,7 @@
 import flytekit
 from flytekit import PythonFunctionTask, Resources, lazy_module
 from flytekit.configuration import SerializationSettings
-from flytekit.core.context_manager import FlyteContextManager, OutputMetadata
+from flytekit.core.context_manager import OutputMetadata
 from flytekit.core.pod_template import PodTemplate
 from flytekit.core.resources import convert_resources_to_resource_model
 from flytekit.exceptions.user import FlyteRecoverableException
@@ -465,7 +465,7 @@ def fn_partial():
         # Rank 0 returns the result of the task function
         if 0 in out:
             # For rank 0, we transfer the decks created in the worker process to the parent process
-            ctx = FlyteContextManager.current_context()
+            ctx = flytekit.current_context()
             for deck in out[0].decks:
                 if not isinstance(deck, flytekit.deck.deck.TimeLineDeck):
                     ctx.decks.append(deck)

plugins/flytekit-mlflow/tests/test_mlflow_tracking.py

@@ -29,4 +29,4 @@ def train_model(epochs: int):
 
 def test_local_exec():
     train_model(epochs=1)
-    assert len(flytekit.current_context().decks) == 5  # mlflow metrics, params, timeline, input, and output
+    assert len(flytekit.current_context().decks) == 7  # mlflow metrics, params, timeline, input, and output, source code, dependencies

plugins/flytekit-ray/flytekitplugins/ray/task.py

@@ -46,7 +46,6 @@ class RayJobConfig:
     address: typing.Optional[str] = None
     shutdown_after_job_finishes: bool = False
     ttl_seconds_after_finished: typing.Optional[int] = None
-    excludes_working_dir: typing.Optional[typing.List[str]] = None
 
 
 class RayFunctionTask(PythonFunctionTask):
@@ -76,10 +75,6 @@ def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
                 "excludes": ["script_mode.tar.gz", "fast*.tar.gz"],
             }
 
-            cfg = self._task_config
-            if cfg.excludes_working_dir:
-                init_params["runtime_env"]["excludes"].extend(cfg.excludes_working_dir)
-
         ray.init(**init_params)
         return user_params
 

plugins/flytekit-ray/tests/test_ray.py

@@ -72,4 +72,4 @@ def t1(a: int) -> str:
     ]
 
     assert t1(a=3) == "5"
-    assert not ray.is_initialized()
+    assert ray.is_initialized()

tests/flytekit/integration/remote/workflows/basic/array_map.py

@@ -1,3 +1,4 @@
+import typing
 from functools import partial
 
 from flytekit import map_task, task, workflow
@@ -9,6 +10,6 @@ def fn(x: int, y: int) -> int:
 
 
 @workflow
-def workflow_with_maptask(data: list[int], y: int) -> list[int]:
+def workflow_with_maptask(data: typing.List[int], y: int) -> typing.List[int]:
     partial_fn = partial(fn, y=y)
     return map_task(partial_fn)(x=data)

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -26,7 +26,7 @@
 )
 from flytekit.interaction.click_types import DirParamType, FileParamType
 from flytekit.remote import FlyteRemote
-from typing import Iterator
+from typing import Iterator, List
 from flytekit.types.iterator import JSON
 from flytekit import workflow
 
@@ -276,6 +276,7 @@ def test_union_type_with_invalid_input():
     assert result.exit_code == 2
 
 
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="listing entities requires python>=3.9")
 @pytest.mark.parametrize(
     "workflow_file",
     [
@@ -521,11 +522,11 @@ def tk(x: Iterator[JSON] = jsons()) -> Iterator[JSON]:
         return t1(x=x)
 
     @task
-    def t2(x: list[int]) -> list[int]:
+    def t2(x: List[int]) -> List[int]:
         return x
 
     @workflow
-    def tk_list(x: list[int] = [1, 2, 3]) -> list[int]:
+    def tk_list(x: List[int] = [1, 2, 3]) -> List[int]:
         return t2(x=x)
 
     @task

tests/flytekit/unit/core/test_array_node_map_task.py

@@ -2,6 +2,7 @@
 import typing
 from collections import OrderedDict
 from typing import List
+from typing_extensions import Annotated
 
 import pytest
 
@@ -78,7 +79,7 @@ def say_hello(name: str) -> str:
 
 def test_map_task_with_pickle():
     @task
-    def say_hello(name: typing.Annotated[typing.Any, BatchSize(10)]) -> str:
+    def say_hello(name: Annotated[typing.Any, BatchSize(10)]) -> str:
         return f"hello {name}!"
 
     with pytest.raises(ValueError, match="Choosing a BatchSize for map tasks inputs is not supported."):

tests/flytekit/unit/core/test_dataclass.py

@@ -5,7 +5,8 @@
 import sys
 import tempfile
 from dataclasses import dataclass
-from typing import Annotated, List, Dict, Optional
+from typing import List, Dict, Optional
+from typing_extensions import Annotated
 from flytekit.types.schema import FlyteSchema
 from flytekit.core.type_engine import TypeEngine
 from flytekit.core.context_manager import FlyteContextManager

tests/flytekit/unit/core/test_serialization.py

@@ -525,7 +525,7 @@ def wf_with_input() -> int:
         return t1(a=input_val)
 
     @workflow
-    def wf_with_sub_wf() -> tuple[int, int]:
+    def wf_with_sub_wf() -> typing.Tuple[int, int]:
         return (wf_no_input(), wf_with_input())
 
     wf_no_input_spec = get_serializable(OrderedDict(), serialization_settings, wf_no_input)
@@ -564,7 +564,7 @@ def wf_with_input() -> str:
         return t1(a=input_val)
 
     @workflow
-    def wf_with_sub_wf() -> tuple[str, str]:
+    def wf_with_sub_wf() -> typing.Tuple[str, str]:
         return (wf_no_input(), wf_with_input())
 
     wf_no_input_spec = get_serializable(OrderedDict(), serialization_settings, wf_no_input)
@@ -603,7 +603,7 @@ def wf_with_input() -> typing.Optional[int]:
         return t1(a=input_val)
 
     @workflow
-    def wf_with_sub_wf() -> tuple[typing.Optional[int], typing.Optional[int]]:
+    def wf_with_sub_wf() -> typing.Tuple[typing.Optional[int], typing.Optional[int]]:
         return (wf_no_input(), wf_with_input())
 
     wf_no_input_spec = get_serializable(OrderedDict(), serialization_settings, wf_no_input)
@@ -673,7 +673,7 @@ def wf_with_input() -> typing.Optional[int]:
         return t1(a=input_val)
 
     @workflow
-    def wf_with_sub_wf() -> tuple[typing.Optional[int], typing.Optional[int]]:
+    def wf_with_sub_wf() -> typing.Tuple[typing.Optional[int], typing.Optional[int]]:
         return (wf_no_input(), wf_with_input())
 
     wf_no_input_spec = get_serializable(OrderedDict(), serialization_settings, wf_no_input)
@@ -764,15 +764,15 @@ def test_default_args_task_list_type():
     input_val = [1, 2, 3]
 
     @task
-    def t1(a: list[int] = []) -> list[int]:
+    def t1(a: typing.List[int] = []) -> typing.List[int]:
         return a
 
     @workflow
-    def wf_no_input() -> list[int]:
+    def wf_no_input() -> typing.List[int]:
         return t1()
 
     @workflow
-    def wf_with_input() -> list[int]:
+    def wf_with_input() -> typing.List[int]:
         return t1(a=input_val)
 
     with pytest.raises(FlyteAssertion, match="Cannot use non-hashable object as default argument"):
@@ -799,15 +799,15 @@ def test_default_args_task_dict_type():
     input_val = {"a": 1, "b": 2}
 
     @task
-    def t1(a: dict[str, int] = {}) -> dict[str, int]:
+    def t1(a: typing.Dict[str, int] = {}) -> typing.Dict[str, int]:
         return a
 
     @workflow
-    def wf_no_input() -> dict[str, int]:
+    def wf_no_input() -> typing.Dict[str, int]:
         return t1()
 
     @workflow
-    def wf_with_input() -> dict[str, int]:
+    def wf_with_input() -> typing.Dict[str, int]:
         return t1(a=input_val)
 
     with pytest.raises(FlyteAssertion, match="Cannot use non-hashable object as default argument"):
@@ -846,7 +846,7 @@ def wf_with_input() -> typing.Optional[typing.List[int]]:
         return t1(a=input_val)
 
     @workflow
-    def wf_with_sub_wf() -> tuple[typing.Optional[typing.List[int]], typing.Optional[typing.List[int]]]:
+    def wf_with_sub_wf() -> typing.Tuple[typing.Optional[typing.List[int]], typing.Optional[typing.List[int]]]:
         return (wf_no_input(), wf_with_input())
 
     wf_no_input_spec = get_serializable(OrderedDict(), serialization_settings, wf_no_input)

tests/flytekit/unit/core/test_type_engine.py

@@ -2389,17 +2389,28 @@ class Result(DataClassJsonMixin):
     schema: TestSchema  # type: ignore
 
 
-@pytest.mark.parametrize(
-    "t",
-    [
+def get_unsupported_complex_literals_tests():
+    if sys.version_info < (3, 9):
+        return [
+        typing_extensions.Annotated[typing.Dict[int, str], FlyteAnnotation({"foo": "bar"})],
+        typing_extensions.Annotated[typing.Dict[str, str], FlyteAnnotation({"foo": "bar"})],
+        typing_extensions.Annotated[Color, FlyteAnnotation({"foo": "bar"})],
+        typing_extensions.Annotated[Result, FlyteAnnotation({"foo": "bar"})],
+    ]
+    return [
         typing_extensions.Annotated[dict, FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[dict[int, str], FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[typing.Dict[int, str], FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[dict[str, str], FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[typing.Dict[str, str], FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[Color, FlyteAnnotation({"foo": "bar"})],
         typing_extensions.Annotated[Result, FlyteAnnotation({"foo": "bar"})],
-    ],
+    ]
+
+
+@pytest.mark.parametrize(
+    "t",
+    get_unsupported_complex_literals_tests(),
 )
 def test_unsupported_complex_literals(t):
     with pytest.raises(ValueError):
@@ -3006,7 +3017,7 @@ def test_dataclass_encoder_and_decoder_registry():
     class Datum:
         x: int
         y: str
-        z: dict[int, int]
+        z: typing.Dict[int, int]
         w: List[int]
 
     @task

tests/flytekit/unit/core/test_type_hints.py

@@ -1526,6 +1526,7 @@ def t2() -> dict:
     assert output_lm.literals["o0"].scalar.generic == expected_struct
 
 
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="Use of dict hints is only supported in Python 3.9+")
 def test_guess_dict4():
     @dataclass
     class Foo(DataClassJsonMixin):

tests/flytekit/unit/extend/test_agent.py

@@ -15,7 +15,7 @@
     GetTaskRequest,
     ListAgentsRequest,
     ListAgentsResponse,
-    TaskCategory,
+    TaskCategory, DeleteTaskResponse,
 )
 from flyteidl.core.execution_pb2 import TaskExecution, TaskLog
 from flyteidl.core.identifier_pb2 import ResourceType
@@ -223,7 +223,7 @@ async def test_async_agent_service(agent, consume_metadata):
     res = await service.GetTask(GetTaskRequest(task_category=task_category, resource_meta=metadata_bytes), ctx)
     assert res.resource.phase == TaskExecution.SUCCEEDED
     res = await service.DeleteTask(DeleteTaskRequest(task_category=task_category, resource_meta=metadata_bytes), ctx)
-    assert res is None
+    assert res == DeleteTaskResponse()
 
     agent_metadata = AgentRegistry.get_agent_metadata(agent.name)
     assert agent_metadata.supported_task_types[0] == agent.task_category.name
@@ -399,16 +399,3 @@ def sample_agents():
         name="ChatGPT Agent", is_sync=True, supported_task_categories=[TaskCategory(name="chatgpt", version=0)]
     )
     return [async_agent, sync_agent]
-
-
-@patch("flytekit.clis.sdk_in_container.serve.click.secho")
-@patch("flytekit.extend.backend.base_agent.AgentRegistry.list_agents")
-def test_print_agents_metadata_output(list_agents_mock, mock_secho, sample_agents):
-    list_agents_mock.return_value = sample_agents
-    print_agents_metadata()
-    expected_calls = [
-        (("Starting Sensor that supports task categories ['sensor']",), {"fg": "blue"}),
-        (("Starting ChatGPT Agent that supports task categories ['chatgpt']",), {"fg": "blue"}),
-    ]
-    mock_secho.assert_has_calls(expected_calls, any_order=True)
-    assert mock_secho.call_count == len(expected_calls)

tests/flytekit/unit/extras/tasks/test_shell.py

@@ -43,7 +43,7 @@ def test_shell_task_access_to_result():
     t()
 
     assert t.result.returncode == 0
-    assert t.result.output == "Hello World!"  # ShellTask strips carriage returns
+    assert "Hello World!" in t.result.output  # ShellTask strips carriage returns
     assert t.result.error == ""
 
 

tests/flytekit/unit/interaction/test_click_types.py

@@ -259,8 +259,8 @@ def test_dataclass_type():
     class Datum:
         x: int
         y: str
-        z: dict[int, str]
-        w: list[int]
+        z: typing.Dict[int, str]
+        w: typing.List[int]
 
     t = JsonParamType(Datum)
     value = '{ "x": 1, "y": "2", "z": { "1": "one", "2": "two" }, "w": [1, 2, 3] }'