def test_normal_task(mock_client):
merge_sort_remotely = load_proto_from_file(
task_pb2.Task,
os.path.join(responses_dir, "admin.task_pb2.Task.pb"),
)
admin_task = task_models.Task.from_flyte_idl(merge_sort_remotely)
mock_client.get_task.return_value = admin_task
ft = rr.fetch_task(name="merge_sort_remotely", version="tst")
@workflow
def my_wf(numbers: typing.List[int],
run_local_at_count: int) -> typing.List[int]:
t1_node = create_node(ft,
numbers=numbers,
run_local_at_count=run_local_at_count)
return t1_node.o0
serialization_settings = flytekit.configuration.SerializationSettings(
project="project",
domain="domain",
version="version",
env=None,
image_config=ImageConfig.auto(img_name=DefaultImages.default_image()),
)
wf_spec = get_serializable(OrderedDict(), serialization_settings, my_wf)
assert wf_spec.template.nodes[
0].task_node.reference_id.name == "merge_sort_remotely"
def test_bad_configuration():
container_image = "{{.image.xyz.fqn}}:{{.image.default.version}}"
image_config = ImageConfig.auto(config_file=os.path.join(
os.path.dirname(os.path.realpath(__file__)), "configs/images.config"))
# No default image in the images.config file so nothing to pull version from
with pytest.raises(AssertionError):
get_registerable_container_image(container_image, image_config)
def test_module_loading(mock_entities, mock_entities_2):
entities = []
mock_entities.entities = entities
mock_entities_2.entities = entities
with tempfile.TemporaryDirectory() as tmp_dir:
# Create directories
top_level = os.path.join(tmp_dir, "top")
middle_level = os.path.join(top_level, "middle")
bottom_level = os.path.join(middle_level, "bottom")
os.makedirs(bottom_level)
top_level_2 = os.path.join(tmp_dir, "top2")
middle_level_2 = os.path.join(top_level_2, "middle")
os.makedirs(middle_level_2)
# Create init files
pathlib.Path(os.path.join(top_level, "__init__.py")).touch()
pathlib.Path(os.path.join(top_level, "a.py")).touch()
pathlib.Path(os.path.join(middle_level, "__init__.py")).touch()
pathlib.Path(os.path.join(middle_level, "a.py")).touch()
pathlib.Path(os.path.join(bottom_level, "__init__.py")).touch()
pathlib.Path(os.path.join(bottom_level, "a.py")).touch()
with open(os.path.join(bottom_level, "a.py"), "w") as fh:
fh.write(task_text)
pathlib.Path(os.path.join(middle_level_2, "__init__.py")).touch()
# Because they have different roots
with pytest.raises(ValueError):
find_common_root([middle_level_2, bottom_level])
# But now add one more init file
pathlib.Path(os.path.join(top_level_2, "__init__.py")).touch()
# Now it should pass
root = find_common_root([middle_level_2, bottom_level])
assert pathlib.Path(root).resolve() == pathlib.Path(tmp_dir).resolve()
# Now load them
serialization_settings = flytekit.configuration.SerializationSettings(
project="project",
domain="domain",
version="version",
env=None,
image_config=ImageConfig.auto(
img_name=DefaultImages.default_image()),
)
x = load_packages_and_modules(serialization_settings,
pathlib.Path(root), [bottom_level])
assert len(x) == 1
def test_serialization_images():
@task(container_image="{{.image.xyz.fqn}}:{{.image.xyz.version}}")
def t1(a: int) -> int:
return a
@task(container_image="{{.image.abc.fqn}}:{{.image.xyz.version}}")
def t2():
pass
@task(container_image="docker.io/org/myimage:latest")
def t4():
pass
@task(container_image="docker.io/org/myimage:{{.image.xyz.version}}")
def t5(a: int) -> int:
return a
@task(container_image="{{.image.xyz_123.fqn}}:{{.image.xyz_123.version}}")
def t6(a: int) -> int:
return a
os.environ["FLYTE_INTERNAL_IMAGE"] = "docker.io/default:version"
imgs = ImageConfig.auto(config_file=os.path.join(
os.path.dirname(os.path.realpath(__file__)), "configs/images.config"))
rs = flytekit.configuration.SerializationSettings(
project="project",
domain="domain",
version="version",
env=None,
image_config=imgs,
)
t1_spec = get_serializable(OrderedDict(), rs, t1)
assert t1_spec.template.container.image == "docker.io/xyz:latest"
t1_spec.to_flyte_idl()
t2_spec = get_serializable(OrderedDict(), rs, t2)
assert t2_spec.template.container.image == "docker.io/abc:latest"
t4_spec = get_serializable(OrderedDict(), rs, t4)
assert t4_spec.template.container.image == "docker.io/org/myimage:latest"
t5_spec = get_serializable(OrderedDict(), rs, t5)
assert t5_spec.template.container.image == "docker.io/org/myimage:latest"
t5_spec = get_serializable(OrderedDict(), rs, t6)
assert t5_spec.template.container.image == "docker.io/xyz_123:v1"
def serialize_all(
pkgs: typing.List[str] = None,
local_source_root: typing.Optional[str] = None,
folder: typing.Optional[str] = None,
mode: typing.Optional[SerializationMode] = None,
image: typing.Optional[str] = None,
flytekit_virtualenv_root: typing.Optional[str] = None,
python_interpreter: typing.Optional[str] = None,
config_file: typing.Optional[str] = None,
):
"""
This function will write to the folder specified the following protobuf types ::
flyteidl.admin.launch_plan_pb2.LaunchPlan
flyteidl.admin.workflow_pb2.WorkflowSpec
flyteidl.admin.task_pb2.TaskSpec
These can be inspected by calling (in the launch plan case) ::
flyte-cli parse-proto -f filename.pb -p flyteidl.admin.launch_plan_pb2.LaunchPlan
See :py:class:`flytekit.models.core.identifier.ResourceType` to match the trailing index in the file name with the
entity type.
:param pkgs: Dot-delimited Python packages/subpackages to look into for serialization.
:param local_source_root: Where to start looking for the code.
:param folder: Where to write the output protobuf files
:param mode: Regular vs fast
:param image: The fully qualified and versioned default image to use
:param flytekit_virtualenv_root: The full path of the virtual env in the container.
"""
if not (mode == SerializationMode.DEFAULT
or mode == SerializationMode.FAST):
raise AssertionError(f"Unrecognized serialization mode: {mode}")
serialization_settings = SerializationSettings(
image_config=ImageConfig.auto(config_file, img_name=image),
fast_serialization_settings=FastSerializationSettings(
enabled=mode == SerializationMode.FAST,
# TODO: if we want to move the destination dir as a serialization argument, we should initialize it here
),
flytekit_virtualenv_root=flytekit_virtualenv_root,
python_interpreter=python_interpreter,
)
serialize_to_folder(pkgs, serialization_settings, local_source_root,
folder)
def test_more_stuff(mock_client):
r = FlyteRemote(config=Config.auto(),
default_project="project",
default_domain="domain")
# Can't upload a folder
with pytest.raises(ValueError):
with tempfile.TemporaryDirectory() as tmp_dir:
r._upload_file(pathlib.Path(tmp_dir))
# Test that this copies the file.
with tempfile.TemporaryDirectory() as tmp_dir:
mm = MagicMock()
mm.signed_url = os.path.join(tmp_dir, "tmp_file")
mock_client.return_value.get_upload_signed_url.return_value = mm
r._upload_file(pathlib.Path(__file__))
serialization_settings = flytekit.configuration.SerializationSettings(
project="project",
domain="domain",
version="version",
env=None,
image_config=ImageConfig.auto(img_name=DefaultImages.default_image()),
)
# gives a thing
computed_v = r._version_from_hash(b"", serialization_settings)
assert len(computed_v) > 0
# gives the same thing
computed_v2 = r._version_from_hash(b"", serialization_settings)
assert computed_v2 == computed_v2
# should give a different thing
computed_v3 = r._version_from_hash(b"", serialization_settings, "hi")
assert computed_v2 != computed_v3
from flytekit.configuration import Config, ImageConfig, SerializationSettings
from flytekit.loggers import logger
from flytekit.remote.remote import FlyteRemote
from .resources import hello_wf
#####
# THESE TESTS ARE NOT RUN IN CI. THEY ARE HERE TO MAKE LOCAL TESTING EASIER.
# Update these to use these tests
IMAGE_STR = "flytecookbook:core-f7af27e23b3935a166645cf96a68583cdd263a87"
FETCH_VERSION = "a351b7c7445a8a818cdf87bf1c1cf38b63beddf1"
RELEASED_EXAMPLES_VERSION = "a351b7c7445a8a818cdf87bf1c1cf38b63beddf1"
#####
image_config = ImageConfig.auto(img_name=IMAGE_STR)
rr = FlyteRemote(
Config.for_sandbox(),
default_project="flytesnacks",
default_domain="development",
)
def get_get_version():
_VERSION_PREFIX = "sandbox_test_" + uuid.uuid4().hex[:3]
logger.warning(f"Test version prefix is {_VERSION_PREFIX}")
print(f"fdsafdsaTest version prefix is {_VERSION_PREFIX}")
def fn(suffix: str = "") -> str:
return _VERSION_PREFIX + (f"_{suffix}" if suffix else "")