Add . to the PYTHONPATH in ImageSpec

[pingsutw]

Tracking issue

https://flyte-org.slack.com/archives/CP2HDHKE1/p1716911881652929

Why are the changes needed?

Module not found while running on a remote cluster since Python can not find the local module unless you explicitly add the local directory to the PYTHONPATH

Note: when you are using non-default base_image, the WORKDIR might be others (/ or /home, etc).

What changes were proposed in this pull request?

Add . to PYTHONPATH

How was this patch tested?

from flytekit import task, workflow, ImageSpec

new_flytekit = "git+https://github.com/flyteorg/flytekit@554dd6dd4882d703839dc4b846285b7b2f9f66dc"

image_spec = ImageSpec(
    name="helloworld",
    base_image="python:3.11.5-slim-bullseye",
    packages=[new_flytekit, "mypy"],
    apt_packages=["git"],
    registry="localhost:30000",
)


@task(container_image=image_spec)
def say_hello() -> str:
    return "Hello, World!"


@workflow
def wf() -> str:
    res = say_hello()
    return res


if __name__ == "__main__":
    print(f"Running hello_world_wf() {wf()}")

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

NA

Docs link

NA

[thomasjpfan]

For custom base images, should ImageSpec change the WORKDIR to /root?

[thomasjpfan]

The "." is easy to miss. I think using PWD works too?

Suggested change

	env = {"PYTHONPATH": "/root:.", _F_IMG_ID: image_spec.image_name()}
	env = {"PYTHONPATH": "/root:$PWD", _F_IMG_ID: image_spec.image_name()}

[wild-endeavor]

you don't need . you just need the :

[eapolinario]

how does this work? I was under the impression that the current working directory was already part of sys.path when the python interpreter starts.

[pingsutw]

For custom base images, should ImageSpec change the WORKDIR to /root?

we don't need to. It should use the workdir in the base image.

[eapolinario]

Can we add a description of the problem to this PR? I'd love to understand why this fixes the problem.