#1601 Add is_inside func to ImageSpec Pull request opened by pingsutw TL;DR

is_inside

will check if current image is built from image_spec when running remotely. This function always returns true when running locally, which means it will import all the package in local execution.

from flytekit import task, workflow, Resources
from flytekit.image_spec import ImageSpec

new_flytekit = "git+<https://github.com/flyteorg/flytekit@db448b054533de52e3a2d7fdb08be31e7b8a5c7d>"
tensorflow_image_spec = ImageSpec(packages=["tensorflow", new_flytekit], apt_packages=["git"], registry="pingsutw")
torch_image_spec = ImageSpec(packages=["torch", new_flytekit], apt_packages=["git"], registry="pingsutw")

if tensorflow_image_spec.in_container():
    print("import tensorflow")
    import tensorflow  # import tensorflow only when running t1

if torch_image_spec.in_container():
    print("import torch")
    import torch  # import torch only when running t2


@task(container_image=tensorflow_image_spec, requests=Resources(mem="900Mi"))
def t1():
    tensorflow.constant(4)
    print("tensorflow")


@task(container_image=torch_image_spec, requests=Resources(mem="900Mi"))
def t2():
    torch.tensor([[1., -1.], [1., -1.]])
    print("torch")


@workflow
def wf():
    t1()
    t2()


if __name__ == '__main__':
    wf()

from flytekit import task, workflow, Resources
from flytekit.image_spec import ImageSpec

new_flytekit = "git+<https://github.com/flyteorg/flytekit@6a1dfaea2e2521c9f9a009cf1e9cefdc97c27aed>"
tensorflow_image_spec = ImageSpec(packages=["tensorflow", new_flytekit], apt_packages=["git"], registry="pingsutw")
torch_image_spec = ImageSpec(packages=["torch", new_flytekit], apt_packages=["git"], registry="pingsutw")

if tensorflow_image_spec.is_container():
    print("import tensorflow")
    import tensorflow


    @task(container_image=tensorflow_image_spec, requests=Resources(mem="900Mi"))
    def t1() -> tensorflow.constant:
        t = tensorflow.constant(4)
        print("tensorflow")
        return t

if torch_image_spec.is_container():
    print("import torch")
    import torch

    @task(container_image=torch_image_spec, requests=Resources(mem="900Mi"))
    def t2() -> torch.tensor:
        t = torch.tensor([[1., -1.], [1., -1.]])
        print("torch")
        return t


@workflow
def wf():
    t1()
    t2()


if __name__ == '__main__':
    wf()

Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☐ Code completed ☐ Smoke tested ☐ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description ^^^ Tracking Issue NA Follow-up issue NA flyteorg/flytekit ✅ All checks have passed 30/30 successful checks