#4853 [BUG] Task to fails to import modules in `pyflyte-map-execute` mode but works in `pyflyte-execute` mode Issue created by debajyoti-truefoundry ### Describe the bug

~/work/pipelines/..         │I   1 import os     ■ Missing module docstring
    .git                    │    2 from functools import partial
    flyte                   │    3 from pathlib import Path
      installation          │    4 from typing import List, Tuple
      my_workflow           │    5
        __pycache__         │E   6 import flytekit     ■ Unable to import 'flytekit'
        my_module           │E   7 from flytekit import (     ■ Unable to import 'flytekit'
          __pycache__       │    8     ImageSpec,
           __init__.py       │    9     PodTemplate,
           my_module.py      │   10     Resources,
         __init__.py         │   11     conditional,
         register.sh         │   12     map_task,
         run.sh              │   13     task,
         train_model_workflow│   14     workflow,
       .gitignore            │   15 )
       config.yaml           │E  16 from flytekit.types.directory import FlyteDirectory     ■ Unable to import 'flytekit.types.directory'
       requirements.txt      │E  17 from kubernetes import client     ■ Unable to import 'kubernetes'
                              │   18 from my_workflow.my_module.my_module import random

This is what my current folder structure looks like, and on line 18, I have imported a module. My workflow file looks like this, import os from functools import partial from pathlib import Path from typing import List, Tuple import flytekit from flytekit import ( ImageSpec, PodTemplate, Resources, conditional, map_task, task, workflow, ) from flytekit.types.directory import FlyteDirectory from kubernetes import client from my_workflow.my_module.my_module import random nvml_image_spec = ImageSpec( base_image="python:3.9-slim", packages=["flytekit==1.10.3", "pynvml==11.5.0"], registry="redacted", name="redacted-pytorch", source_root="..", ) normal_image = ImageSpec( base_image="python:3.9-slim", packages=["flytekit==1.10.3"], registry="redacted", name="redacted-normal", source_root="..", ) cpu_resource = Resources(mem="300Mi", cpu="0.1") gpu_resource = Resources(mem="900Mi", cpu="0.1", gpu="1") gpu_pod_template = PodTemplate( pod_spec=client.V1PodSpec( containers=[], affinity=client.V1Affinity( node_affinity=client.V1NodeAffinity( required_during_scheduling_ignored_during_execution=client.V1NodeSelector( node_selector_terms=[ client.V1NodeSelectorTerm( match_expressions=[ client.V1NodeSelectorRequirement( key="karpenter.k8s.aws/instance-family", operator="In", values=["g4dn"], ) ] ) ] ) ) ), ) ) @task(requests=cpu_resource, limits=cpu_resource, container_image=normal_image) def should_train_tokenizer(tokenizer: str) -> bool: return not bool(tokenizer) @task(requests=cpu_resource, limits=cpu_resource, container_image=normal_image) def train_tokenizer() -> str: return "trained_tokenizer" @task( container_image=nvml_image_spec, requests=gpu_resource, limits=gpu_resource, pod_template=gpu_pod_template, # accelerator=flytekit.extras.accelerators.T4, ) def train_model(tokenizer: str) -> FlyteDirectory: from pynvml import nvmlInit, nvmlDeviceGetCount nvmlInit() assert nvmlDeviceGetCount() > 0 random.random() working_dir = flytekit.current_context().working_directory local_dir = Path(os.path.join(working_dir, "csv_files")) local_dir.mkdir(exist_ok=True) with open(os.path.join(local_dir, "model"), "w", encoding="utf-8") as f: f.write(tokenizer) return FlyteDirectory(path=str(local_dir)) @task(requests=cpu_resource, limits=cpu_resource, container_image=normal_image) def validate_model(model: FlyteDirectory, tokenizer: str, validation_data: str) -> bool: print(validation_data) model_path = os.path.join(model, "model") with open(model_path, "r", encoding="utf-8") as f: return f.read() == tokenizer @task(requests=cpu_resource, limits=cpu_resource, container_image=normal_image) def all_good(validations: List[bool]) -> bool: return all(validations) @workflow def train_tokenizer_and_model() -> Tuple[FlyteDirectory, str]: tokenizer = train_tokenizer() model = train_model(tokenizer=tokenizer) return model, tokenizer @workflow def just_train_model(tokenizer: str) -> Tuple[FlyteDirectory, str]: model = train_model(tokenizer=tokenizer) return model, tokenizer @workflow def train(tokenizer: str = "") -> bool: stt = should_train_tokenizer(tokenizer=tokenizer) model, t = ( conditional("train_tokenizer") .if_(stt.is_true()) .then(train_tokenizer_and_model()) .else_() .then(just_train_model(tokenizer=tokenizer)) ) validation_task = partial(validate_model, model=model, tokenizer=t) validations = map_task( validation_task, concurrency=2, )(validation_data=["foo", "bar", "baz"]) return all_good(validations=validations) I used the following command to run the workflow remotely.

pyflyte -v --config ../config.yaml run --remote train_model_workflow.py train --tokenizer Deb

[image](https://private-user-images.githubusercontent.com/97010181/302980970-c46c8414-298c-44d8-9f21-65baa78a174a.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NDgxMzIxOTIsIm5iZiI6MTc0ODEzMTg5MiwicGF0aCI6Ii85NzAxMDE4MS8zMDI5ODA5NzAtYzQ2Yzg0MTQtMjk4Yy00NGQ4LTlmMjEtNjViYWE3OGExNzRhLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNTA1MjUlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjUwNTI1VDAwMTEzMlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWI3ZjUyZWZlOTMyMjNiNzRiNjA4NDgwNjc2ZjA1NTUwZmQyOTdiMWQ1OGM0Yjk2M2RhMDA1YTcwNWJmOGY0MTEmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.ujX48V7XuMMBDZlXUClmEr9vyN1AySMZWUqwVT7Swlk) As you can see in the above screenshot, the validate model map task fails. This is the log I can see in the corresponding pod ``` tar: Removing leading `/' from member names ╭───────────────────── Traceback (most recent call last) ──────────────────────╮ │ /usr/local/bin/pyflyte-map-execute:8 in <module> │ │ │ │ ❱ 8 │ sys.exit(map_execute_task_cmd()) │ │ │ │ /usr/local/lib/python3.9/site-packages/click/core.py:1157 in call │ │ │ │ ❱ 1157 │ │ return self.main(*args, **kwargs) │ │ … flyteorg/flyte