flyte 2.0.0b25__py3-none-any.whl → 2.0.0b28__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of flyte might be problematic. Click here for more details.

Files changed (54) hide show
  1. flyte/__init__.py +2 -0
  2. flyte/_bin/runtime.py +8 -0
  3. flyte/_code_bundle/_utils.py +4 -4
  4. flyte/_code_bundle/bundle.py +1 -1
  5. flyte/_constants.py +1 -0
  6. flyte/_deploy.py +0 -1
  7. flyte/_excepthook.py +1 -1
  8. flyte/_initialize.py +10 -0
  9. flyte/_interface.py +2 -0
  10. flyte/_internal/imagebuild/docker_builder.py +3 -1
  11. flyte/_internal/imagebuild/remote_builder.py +3 -1
  12. flyte/_internal/resolvers/_task_module.py +4 -37
  13. flyte/_internal/runtime/convert.py +3 -2
  14. flyte/_internal/runtime/entrypoints.py +24 -1
  15. flyte/_internal/runtime/rusty.py +3 -3
  16. flyte/_internal/runtime/task_serde.py +19 -4
  17. flyte/_internal/runtime/trigger_serde.py +2 -2
  18. flyte/_map.py +2 -35
  19. flyte/_module.py +68 -0
  20. flyte/_resources.py +38 -0
  21. flyte/_run.py +23 -6
  22. flyte/_task.py +1 -2
  23. flyte/_task_plugins.py +4 -2
  24. flyte/_trigger.py +623 -5
  25. flyte/_utils/__init__.py +2 -1
  26. flyte/_utils/asyn.py +3 -1
  27. flyte/_utils/docker_credentials.py +173 -0
  28. flyte/_utils/module_loader.py +15 -0
  29. flyte/_version.py +3 -3
  30. flyte/cli/_common.py +15 -3
  31. flyte/cli/_create.py +100 -3
  32. flyte/cli/_deploy.py +38 -4
  33. flyte/cli/_plugins.py +208 -0
  34. flyte/cli/_run.py +69 -6
  35. flyte/cli/_serve.py +154 -0
  36. flyte/cli/main.py +6 -0
  37. flyte/connectors/__init__.py +3 -0
  38. flyte/connectors/_connector.py +270 -0
  39. flyte/connectors/_server.py +183 -0
  40. flyte/connectors/utils.py +26 -0
  41. flyte/models.py +13 -4
  42. flyte/remote/_client/auth/_channel.py +9 -5
  43. flyte/remote/_console.py +3 -2
  44. flyte/remote/_secret.py +6 -4
  45. flyte/remote/_trigger.py +2 -2
  46. flyte/types/_type_engine.py +1 -2
  47. {flyte-2.0.0b25.data → flyte-2.0.0b28.data}/scripts/runtime.py +8 -0
  48. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/METADATA +6 -2
  49. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/RECORD +54 -46
  50. {flyte-2.0.0b25.data → flyte-2.0.0b28.data}/scripts/debug.py +0 -0
  51. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/WHEEL +0 -0
  52. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/entry_points.txt +0 -0
  53. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/licenses/LICENSE +0 -0
  54. {flyte-2.0.0b25.dist-info → flyte-2.0.0b28.dist-info}/top_level.txt +0 -0
flyte/__init__.py CHANGED
@@ -16,6 +16,7 @@ from ._excepthook import custom_excepthook
16
16
  from ._group import group
17
17
  from ._image import Image
18
18
  from ._initialize import current_domain, init, init_from_config
19
+ from ._logging import logger
19
20
  from ._map import map
20
21
  from ._pod import PodTemplate
21
22
  from ._resources import AMD_GPU, GPU, HABANA_GAUDI, TPU, Device, DeviceClass, Neuron, Resources
@@ -97,6 +98,7 @@ __all__ = [
97
98
  "group",
98
99
  "init",
99
100
  "init_from_config",
101
+ "logger",
100
102
  "map",
101
103
  "run",
102
104
  "trace",
flyte/_bin/runtime.py CHANGED
@@ -12,6 +12,7 @@ from typing import Any, List
12
12
 
13
13
  import click
14
14
 
15
+ from flyte._utils.helpers import str2bool
15
16
  from flyte.models import PathRewrite
16
17
 
17
18
  # Todo: work with pvditt to make these the names
@@ -27,6 +28,7 @@ PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
27
28
  DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
28
29
  ORG_NAME = "_U_ORG_NAME"
29
30
  ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
31
+ INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
30
32
  RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
31
33
  FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
32
34
 
@@ -139,6 +141,12 @@ def main(
139
141
  controller_kwargs["insecure"] = True
140
142
  logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
141
143
 
144
+ # Check for insecure_skip_verify override (e.g. for self-signed certs)
145
+ insecure_skip_verify_str = os.getenv(INSECURE_SKIP_VERIFY_OVERRIDE, "")
146
+ if str2bool(insecure_skip_verify_str):
147
+ controller_kwargs["insecure_skip_verify"] = True
148
+ logger.info("SSL certificate verification disabled (insecure_skip_verify=True)")
149
+
142
150
  bundle = None
143
151
  if tgz or pkl:
144
152
  bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
@@ -193,10 +193,10 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
193
193
  def _file_is_in_directory(file: str, directory: str) -> bool:
194
194
  """Return True if file is in directory and in its children."""
195
195
  try:
196
- return os.path.commonpath([file, directory]) == directory
197
- except ValueError as e:
198
- # ValueError is raised by windows if the paths are not from the same drive
199
- logger.debug(f"{file} and {directory} are not in the same drive: {e!s}")
196
+ return pathlib.Path(file).resolve().is_relative_to(pathlib.Path(directory).resolve())
197
+ except OSError as e:
198
+ # OSError can be raised if paths cannot be resolved (permissions, broken symlinks, etc.)
199
+ logger.debug(f"Failed to resolve paths for {file} and {directory}: {e!s}")
200
200
  return False
201
201
 
202
202
 
@@ -104,7 +104,7 @@ async def build_pkl_bundle(
104
104
  import shutil
105
105
 
106
106
  # Copy the bundle to the given path
107
- shutil.copy(dest, copy_bundle_to)
107
+ shutil.copy(dest, copy_bundle_to, follow_symlinks=True)
108
108
  local_path = copy_bundle_to / dest.name
109
109
  return CodeBundle(pkl=str(local_path), computed_version=str_digest)
110
110
  return CodeBundle(pkl=str(dest), computed_version=str_digest)
flyte/_constants.py ADDED
@@ -0,0 +1 @@
1
+ FLYTE_SYS_PATH = "_F_SYS_PATH" # The paths that will be appended to sys.path at runtime
flyte/_deploy.py CHANGED
@@ -255,7 +255,6 @@ async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] |
255
255
 
256
256
  for env_name, image_uri in final_images:
257
257
  logger.warning(f"Built Image for environment {env_name}, image: {image_uri}")
258
- env = deployment.envs[env_name]
259
258
  image_identifier_map[env_name] = image_uri
260
259
 
261
260
  return ImageCache(image_lookup=image_identifier_map)
flyte/_excepthook.py CHANGED
@@ -33,5 +33,5 @@ def custom_excepthook(exc_type, exc_value, exc_tb):
33
33
  filtered_tb = [frame for frame in tb_list if should_include_frame(frame)]
34
34
  # Print the filtered version (custom format)
35
35
  print("Filtered traceback (most recent call last):")
36
- traceback.print_tb(filtered_tb)
36
+ traceback.print_list(filtered_tb)
37
37
  print(f"{exc_type.__name__}: {exc_value}\n")
flyte/_initialize.py CHANGED
@@ -34,6 +34,7 @@ class CommonInit:
34
34
  domain: str | None = None
35
35
  batch_size: int = 1000
36
36
  source_config_path: Optional[Path] = None # Only used for documentation
37
+ sync_local_sys_paths: bool = True
37
38
 
38
39
 
39
40
  @dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
@@ -144,6 +145,7 @@ async def init(
144
145
  image_builder: ImageBuildEngine.ImageBuilderType = "local",
145
146
  images: typing.Dict[str, str] | None = None,
146
147
  source_config_path: Optional[Path] = None,
148
+ sync_local_sys_paths: bool = True,
147
149
  ) -> None:
148
150
  """
149
151
  Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -180,6 +182,8 @@ async def init(
180
182
  :param image_builder: Optional image builder configuration, if not provided, the default image builder will be used.
181
183
  :param images: Optional dict of images that can be used by referencing the image name.
182
184
  :param source_config_path: Optional path to the source configuration file (This is only used for documentation)
185
+ :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
186
+ into the remote container (default: True).
183
187
  :return: None
184
188
  """
185
189
  from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
@@ -230,6 +234,7 @@ async def init(
230
234
  image_builder=image_builder,
231
235
  images=images or {},
232
236
  source_config_path=source_config_path,
237
+ sync_local_sys_paths=sync_local_sys_paths,
233
238
  )
234
239
 
235
240
 
@@ -240,6 +245,7 @@ async def init_from_config(
240
245
  log_level: int | None = None,
241
246
  storage: Storage | None = None,
242
247
  images: tuple[str, ...] | None = None,
248
+ sync_local_sys_paths: bool = True,
243
249
  ) -> None:
244
250
  """
245
251
  Initialize the Flyte system using a configuration file or Config object. This method should be called before any
@@ -253,6 +259,9 @@ async def init_from_config(
253
259
  :param log_level: Optional logging level for the framework logger,
254
260
  default is set using the default initialization policies
255
261
  :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
262
+ :param images: List of image strings in format "imagename=imageuri" or just "imageuri".
263
+ :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
264
+ into the remote container (default: True).
256
265
  :return: None
257
266
  """
258
267
  from rich.highlighter import ReprHighlighter
@@ -306,6 +315,7 @@ async def init_from_config(
306
315
  images=cfg.image.image_refs,
307
316
  storage=storage,
308
317
  source_config_path=cfg_path,
318
+ sync_local_sys_paths=sync_local_sys_paths,
309
319
  )
310
320
 
311
321
 
flyte/_interface.py CHANGED
@@ -51,6 +51,8 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
51
51
  Note that Options 1 and 2 are identical, just syntactic sugar. In the NamedTuple case, we'll use the names in the
52
52
  definition. In all other cases, we'll automatically generate output names, indexed starting at 0.
53
53
  """
54
+ if isinstance(return_annotation, str):
55
+ raise TypeError("String return annotations are not supported.")
54
56
 
55
57
  # Handle Option 6
56
58
  # We can think about whether we should add a default output name with type None in the future.
@@ -47,7 +47,7 @@ FLYTE_DOCKER_BUILDER_CACHE_TO = "FLYTE_DOCKER_BUILDER_CACHE_TO"
47
47
 
48
48
  UV_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE = Template("""\
49
49
  RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
50
- --mount=type=bind,target=uv.lock,src=$UV_LOCK_PATH \
50
+ --mount=type=bind,target=uv.lock,src=$UV_LOCK_PATH,rw \
51
51
  --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
52
52
  $SECRET_MOUNT \
53
53
  uv sync --active --inexact $PIP_INSTALL_ARGS
@@ -271,6 +271,8 @@ class UVProjectHandler:
271
271
  pip_install_args = " ".join(layer.get_pip_install_args())
272
272
  if "--no-install-project" not in pip_install_args:
273
273
  pip_install_args += " --no-install-project"
274
+ if "--no-sources" not in pip_install_args:
275
+ pip_install_args += " --no-sources"
274
276
  # Only Copy pyproject.yaml and uv.lock.
275
277
  pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
276
278
  uvlock_dst = copy_files_to_context(layer.uvlock, context_path)
@@ -134,7 +134,7 @@ class RemoteImageBuilder(ImageBuilder):
134
134
  if run_details.action_details.raw_phase == run_definition_pb2.PHASE_SUCCEEDED:
135
135
  logger.warning(f"[bold green]✅ Build completed in {elapsed}![/bold green]")
136
136
  else:
137
- raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at [cyan]{run.url}[/cyan]")
137
+ raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {run.url}")
138
138
 
139
139
  outputs = await run_details.outputs()
140
140
  return _get_fully_qualified_image_name(outputs)
@@ -272,6 +272,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
272
272
  pip_options.extra_args += " --no-install-project"
273
273
  else:
274
274
  pip_options.extra_args = " --no-install-project"
275
+ if "--no-sources" not in pip_options.extra_args:
276
+ pip_options.extra_args += " --no-sources"
275
277
  else:
276
278
  # Copy the entire project
277
279
  docker_ignore_patterns = get_and_list_dockerignore(image)
@@ -1,9 +1,7 @@
1
- import inspect
2
- import os
3
1
  import pathlib
4
- import sys
5
2
  from typing import Tuple
6
3
 
4
+ from flyte._module import extract_obj_module
7
5
  from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
8
6
 
9
7
 
@@ -15,40 +13,9 @@ def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None =
15
13
  :param source_dir: The source directory to use for relative paths.
16
14
  :return: A tuple containing the entity name, module
17
15
  """
18
- entity_name = task.name
19
16
  if isinstance(task, AsyncFunctionTaskTemplate):
20
- entity_module = inspect.getmodule(task.func)
21
- if entity_module is None:
22
- raise ValueError(f"Task {entity_name} has no module.")
23
-
24
- fp = entity_module.__file__
25
- if fp is None:
26
- raise ValueError(f"Task {entity_name} has no module.")
27
-
28
- file_path = pathlib.Path(fp)
29
- # Get the relative path to the current directory
30
- # Will raise ValueError if the file is not in the source directory
31
- relative_path = file_path.relative_to(str(source_dir))
32
-
33
- if relative_path == pathlib.Path("."):
34
- entity_module_name = entity_module.__name__
35
- else:
36
- # Replace file separators with dots and remove the '.py' extension
37
- dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
38
- entity_module_name = dotted_path
39
-
40
17
  entity_name = task.func.__name__
18
+ entity_module_name = extract_obj_module(task.func, source_dir)
19
+ return entity_name, entity_module_name
41
20
  else:
42
- raise NotImplementedError(f"Task module {entity_name} not implemented.")
43
-
44
- if entity_module_name == "__main__":
45
- """
46
- This case is for the case in which the task is run from the main module.
47
- """
48
- fp = sys.modules["__main__"].__file__
49
- if fp is None:
50
- raise ValueError(f"Task {entity_name} has no module.")
51
- main_path = pathlib.Path(fp)
52
- entity_module_name = main_path.stem
53
-
54
- return entity_name, entity_module_name
21
+ raise NotImplementedError(f"Task module {task.name} not implemented.")
@@ -143,13 +143,14 @@ async def convert_from_native_to_inputs(
143
143
  (default_value is not None and default_value is not inspect.Signature.empty)
144
144
  or (default_value is None and is_optional_type(input_type))
145
145
  or input_type is None
146
+ or input_type is type(None)
146
147
  ):
147
148
  if default_value == NativeInterface.has_default:
148
149
  if interface._remote_defaults is None or input_name not in interface._remote_defaults:
149
150
  raise ValueError(f"Input '{input_name}' has a default value but it is not set in the interface.")
150
151
  already_converted_kwargs[input_name] = interface._remote_defaults[input_name]
151
- elif input_type is None:
152
- # If the type is None, we assume it's a placeholder for no type
152
+ elif input_type is None or input_type is type(None):
153
+ # If the type is 'None' or 'class<None>', we assume it's a placeholder for no type
153
154
  kwargs[input_name] = None
154
155
  type_hints[input_name] = NoneType
155
156
  else:
@@ -1,4 +1,6 @@
1
1
  import importlib
2
+ import os
3
+ import traceback
2
4
  from typing import List, Optional, Tuple, Type
3
5
 
4
6
  import flyte.errors
@@ -10,6 +12,7 @@ from flyte._logging import log, logger
10
12
  from flyte._task import TaskTemplate
11
13
  from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
12
14
 
15
+ from ..._utils import adjust_sys_path
13
16
  from .convert import Error, Inputs, Outputs
14
17
  from .taskrunner import (
15
18
  convert_and_run,
@@ -72,7 +75,26 @@ def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
72
75
  """
73
76
  resolver_class = load_class(resolver)
74
77
  resolver_instance = resolver_class()
75
- return resolver_instance.load_task(resolver_args)
78
+ try:
79
+ return resolver_instance.load_task(resolver_args)
80
+ except ModuleNotFoundError as e:
81
+ cwd = os.getcwd()
82
+ files = []
83
+ try:
84
+ for root, dirs, filenames in os.walk(cwd):
85
+ for name in dirs + filenames:
86
+ rel_path = os.path.relpath(os.path.join(root, name), cwd)
87
+ files.append(rel_path)
88
+ except Exception as list_err:
89
+ files = [f"(Failed to list directory: {list_err})"]
90
+
91
+ msg = (
92
+ "\n\nFull traceback:\n" + "".join(traceback.format_exc()) + f"\n[ImportError Diagnostics]\n"
93
+ f"Module '{e.name}' not found in either the Python virtual environment or the current working directory.\n"
94
+ f"Current working directory: {cwd}\n"
95
+ f"Files found under current directory:\n" + "\n".join(f" - {f}" for f in files)
96
+ )
97
+ raise ModuleNotFoundError(msg) from e
76
98
 
77
99
 
78
100
  def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
@@ -100,6 +122,7 @@ async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
100
122
  :param code_bundle: The code bundle to download.
101
123
  :return: The code bundle with the downloaded path.
102
124
  """
125
+ adjust_sys_path()
103
126
  logger.debug(f"Downloading {code_bundle}")
104
127
  downloaded_path = await download_bundle(code_bundle)
105
128
  return code_bundle.with_downloaded_path(downloaded_path)
@@ -1,5 +1,4 @@
1
1
  import asyncio
2
- import sys
3
2
  import time
4
3
  from typing import Any, List, Tuple
5
4
 
@@ -11,6 +10,7 @@ from flyte._internal.runtime.entrypoints import download_code_bundle, load_pkl_t
11
10
  from flyte._internal.runtime.taskrunner import extract_download_run_upload
12
11
  from flyte._logging import logger
13
12
  from flyte._task import TaskTemplate
13
+ from flyte._utils import adjust_sys_path
14
14
  from flyte.models import ActionID, Checkpoints, CodeBundle, PathRewrite, RawDataPath
15
15
 
16
16
 
@@ -23,7 +23,7 @@ async def download_tgz(destination: str, version: str, tgz: str) -> CodeBundle:
23
23
  :return: The CodeBundle object.
24
24
  """
25
25
  logger.info(f"[rusty] Downloading tgz code bundle from {tgz} to {destination} with version {version}")
26
- sys.path.insert(0, ".")
26
+ adjust_sys_path()
27
27
 
28
28
  code_bundle = CodeBundle(
29
29
  tgz=tgz,
@@ -42,7 +42,7 @@ async def download_load_pkl(destination: str, version: str, pkl: str) -> Tuple[C
42
42
  :return: The CodeBundle object.
43
43
  """
44
44
  logger.info(f"[rusty] Downloading pkl code bundle from {pkl} to {destination} with version {version}")
45
- sys.path.insert(0, ".")
45
+ adjust_sys_path()
46
46
 
47
47
  code_bundle = CodeBundle(
48
48
  pkl=pkl,
@@ -221,10 +221,25 @@ def _get_urun_container(
221
221
  # This computes the image uri, computing hashes as necessary so can fail if done remotely.
222
222
  img_uri = task_template.image.uri
223
223
  elif serialize_context.image_cache and env_name not in serialize_context.image_cache.image_lookup:
224
- img_uri = task_template.image.uri
225
-
226
- logger.warning(
227
- f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}."
224
+ raise flyte.errors.RuntimeUserError(
225
+ "MissingEnvironment",
226
+ f"Environment '{env_name}' not found in image cache.\n\n"
227
+ "💡 To fix this:\n"
228
+ " 1. If your parent environment calls a task in another environment,"
229
+ " declare that dependency using 'depends_on=[...]'.\n"
230
+ " Example:\n"
231
+ " env1 = flyte.TaskEnvironment(\n"
232
+ " name='outer',\n"
233
+ " image=flyte.Image.from_debian_base().with_pip_packages('requests'),\n"
234
+ " depends_on=[env2, env3],\n"
235
+ " )\n"
236
+ " 2. If you're using os.getenv() to set the environment name,"
237
+ " make sure the runtime environment has the same environment variable defined.\n"
238
+ " Example:\n"
239
+ " env = flyte.TaskEnvironment(\n"
240
+ ' name=os.getenv("my-name"),\n'
241
+ ' env_vars={"my-name": os.getenv("my-name")},\n'
242
+ " )\n",
228
243
  )
229
244
  else:
230
245
  img_uri = serialize_context.image_cache.image_lookup[env_name]
@@ -12,7 +12,7 @@ from flyte import Cron, FixedRate, Trigger, TriggerTime
12
12
  def _to_schedule(m: Union[Cron, FixedRate], kickoff_arg_name: str | None = None) -> common_pb2.Schedule:
13
13
  if isinstance(m, Cron):
14
14
  return common_pb2.Schedule(
15
- cron_expression=m.expression,
15
+ cron_expression=m.timezone_expression,
16
16
  kickoff_time_input_arg=kickoff_arg_name,
17
17
  )
18
18
  elif isinstance(m, FixedRate):
@@ -147,7 +147,7 @@ async def to_task_trigger(
147
147
  inputs=common_pb2.Inputs(literals=literals),
148
148
  ),
149
149
  automation_spec=common_pb2.TriggerAutomationSpec(
150
- type=common_pb2.TriggerAutomationSpec.Type.TYPE_SCHEDULE,
150
+ type=common_pb2.TriggerAutomationSpecType.TYPE_SCHEDULE,
151
151
  schedule=automation,
152
152
  ),
153
153
  )
flyte/_map.py CHANGED
@@ -189,7 +189,7 @@ class _Mapper(Generic[P, R]):
189
189
  group_name: str | None = None,
190
190
  concurrency: int = 0,
191
191
  return_exceptions: bool = True,
192
- ) -> Iterator[R]: ...
192
+ ) -> Iterator[Union[R, Exception]]: ...
193
193
 
194
194
  def __call__(
195
195
  self,
@@ -309,37 +309,4 @@ async def _map(
309
309
  yield result
310
310
 
311
311
 
312
- @overload
313
- def map(
314
- func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
315
- *args: Iterable[Any],
316
- group_name: str | None = None,
317
- concurrency: int = 0,
318
- ) -> Iterator[R]: ...
319
-
320
-
321
- @overload
322
- def map(
323
- func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
324
- *args: Iterable[Any],
325
- group_name: str | None = None,
326
- concurrency: int = 0,
327
- return_exceptions: bool = True,
328
- ) -> Iterator[Union[R, Exception]]: ...
329
-
330
-
331
- def map(
332
- func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
333
- *args: Iterable[Any],
334
- group_name: str | None = None,
335
- concurrency: int = 0,
336
- return_exceptions: bool = True,
337
- ) -> Iterator[Union[R, Exception]]:
338
- map: _Mapper = _Mapper()
339
- return map(
340
- func,
341
- *args,
342
- group_name=group_name,
343
- concurrency=concurrency,
344
- return_exceptions=return_exceptions,
345
- )
312
+ map: _Mapper = _Mapper()
flyte/_module.py ADDED
@@ -0,0 +1,68 @@
1
+ import inspect
2
+ import os
3
+ import pathlib
4
+ import sys
5
+
6
+
7
+ def extract_obj_module(obj: object, /, source_dir: pathlib.Path | None = None) -> str:
8
+ """
9
+ Extract the module from the given object. If source_dir is provided, the module will be relative to the source_dir.
10
+
11
+ Args:
12
+ obj: The object to extract the module from.
13
+ source_dir: The source directory to use for relative paths.
14
+
15
+ Returns:
16
+ The module name as a string.
17
+ """
18
+ # Get the module containing the object
19
+ entity_module = inspect.getmodule(obj)
20
+ if entity_module is None:
21
+ obj_name = getattr(obj, "__name__", str(obj))
22
+ raise ValueError(f"Object {obj_name} has no module.")
23
+
24
+ fp = entity_module.__file__
25
+ if fp is None:
26
+ obj_name = getattr(obj, "__name__", str(obj))
27
+ raise ValueError(f"Object {obj_name} has no module.")
28
+
29
+ file_path = pathlib.Path(fp)
30
+ try:
31
+ # Get the relative path to the current directory
32
+ # Will raise ValueError if the file is not in the source directory
33
+ relative_path = file_path.relative_to(str(source_dir))
34
+
35
+ if relative_path == pathlib.Path("_internal/resolvers"):
36
+ entity_module_name = entity_module.__name__
37
+ else:
38
+ # Replace file separators with dots and remove the '.py' extension
39
+ dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
40
+ entity_module_name = dotted_path
41
+ except ValueError:
42
+ # If source_dir is not provided or file is not in source_dir, fallback to module name
43
+ # File is not relative to source_dir - check if it's an installed package
44
+ file_path_str = str(file_path)
45
+ if "site-packages" in file_path_str or "dist-packages" in file_path_str:
46
+ # It's an installed package - use the module's __name__ directly
47
+ # This will be importable via importlib.import_module()
48
+ entity_module_name = entity_module.__name__
49
+ else:
50
+ # File is not in source_dir and not in site-packages - re-raise the error
51
+ obj_name = getattr(obj, "__name__", str(obj))
52
+ raise ValueError(
53
+ f"Object {obj_name} module file {file_path} is not relative to "
54
+ f"source directory {source_dir} and is not an installed package."
55
+ )
56
+
57
+ if entity_module_name == "__main__":
58
+ """
59
+ This case is for the case in which the object is run from the main module.
60
+ """
61
+ fp = sys.modules["__main__"].__file__
62
+ if fp is None:
63
+ obj_name = getattr(obj, "__name__", str(obj))
64
+ raise ValueError(f"Object {obj_name} has no module.")
65
+ main_path = pathlib.Path(fp)
66
+ entity_module_name = main_path.stem
67
+
68
+ return entity_module_name
flyte/_resources.py CHANGED
@@ -148,16 +148,54 @@ Accelerators = Literal[
148
148
  "T4:8",
149
149
  # Trn1
150
150
  "Trn1:1",
151
+ "Trn1:4",
152
+ "Trn1:8",
153
+ "Trn1:16",
151
154
  # Trn1n
152
155
  "Trn1n:1",
156
+ "Trn1n:4",
157
+ "Trn1n:8",
158
+ "Trn1n:16",
153
159
  # Trn2
154
160
  "Trn2:1",
161
+ "Trn2:4",
162
+ "Trn2:8",
163
+ "Trn2:16",
155
164
  # Trn2u
156
165
  "Trn2u:1",
166
+ "Trn2u:4",
167
+ "Trn2u:8",
168
+ "Trn2u:16",
157
169
  # Inf1
158
170
  "Inf1:1",
171
+ "Inf1:2",
172
+ "Inf1:3",
173
+ "Inf1:4",
174
+ "Inf1:5",
175
+ "Inf1:6",
176
+ "Inf1:7",
177
+ "Inf1:8",
178
+ "Inf1:9",
179
+ "Inf1:10",
180
+ "Inf1:11",
181
+ "Inf1:12",
182
+ "Inf1:13",
183
+ "Inf1:14",
184
+ "Inf1:15",
185
+ "Inf1:16",
159
186
  # Inf2
160
187
  "Inf2:1",
188
+ "Inf2:2",
189
+ "Inf2:3",
190
+ "Inf2:4",
191
+ "Inf2:5",
192
+ "Inf2:6",
193
+ "Inf2:7",
194
+ "Inf2:8",
195
+ "Inf2:9",
196
+ "Inf2:10",
197
+ "Inf2:11",
198
+ "Inf2:12",
161
199
  # MI100
162
200
  "MI100:1",
163
201
  # MI210
flyte/_run.py CHANGED
@@ -2,6 +2,7 @@ from __future__ import annotations
2
2
 
3
3
  import asyncio
4
4
  import pathlib
5
+ import sys
5
6
  import uuid
6
7
  from dataclasses import dataclass
7
8
  from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
@@ -29,6 +30,8 @@ from flyte.models import (
29
30
  )
30
31
  from flyte.syncify import syncify
31
32
 
33
+ from ._constants import FLYTE_SYS_PATH
34
+
32
35
  if TYPE_CHECKING:
33
36
  from flyte.remote import Run
34
37
  from flyte.remote._task import LazyEntity
@@ -131,7 +134,7 @@ class _Runner:
131
134
  async def _run_remote(self, obj: TaskTemplate[P, R, F] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
132
135
  import grpc
133
136
  from flyteidl2.common import identifier_pb2
134
- from flyteidl2.core import literals_pb2
137
+ from flyteidl2.core import literals_pb2, security_pb2
135
138
  from flyteidl2.task import run_pb2
136
139
  from flyteidl2.workflow import run_definition_pb2, run_service_pb2
137
140
  from google.protobuf import wrappers_pb2
@@ -172,9 +175,7 @@ class _Runner:
172
175
  if not self._dry_run:
173
176
  image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
174
177
  else:
175
- from ._internal.imagebuild.image_builder import ImageCache
176
-
177
- image_cache = ImageCache(image_lookup={})
178
+ image_cache = None
178
179
 
179
180
  if self._interactive_mode:
180
181
  code_bundle = await build_pkl_bundle(
@@ -220,6 +221,12 @@ class _Runner:
220
221
  else:
221
222
  env["LOG_LEVEL"] = str(logger.getEffectiveLevel())
222
223
 
224
+ # These paths will be appended to sys.path at runtime.
225
+ if cfg.sync_local_sys_paths:
226
+ env[FLYTE_SYS_PATH] = ":".join(
227
+ f"./{pathlib.Path(p).relative_to(cfg.root_dir)}" for p in sys.path if p.startswith(str(cfg.root_dir))
228
+ )
229
+
223
230
  if not self._dry_run:
224
231
  if get_client() is None:
225
232
  # This can only happen, if the user forces flyte.run(mode="remote") without initializing the client
@@ -264,6 +271,14 @@ class _Runner:
264
271
  env_kv = run_pb2.Envs(values=kv_pairs)
265
272
  annotations = run_pb2.Annotations(values=self._annotations)
266
273
  labels = run_pb2.Labels(values=self._labels)
274
+ raw_data_storage = (
275
+ run_pb2.RawDataStorage(raw_data_prefix=self._raw_data_path) if self._raw_data_path else None
276
+ )
277
+ security_context = (
278
+ security_pb2.SecurityContext(run_as=security_pb2.Identity(k8s_service_account=self._service_account))
279
+ if self._service_account
280
+ else None
281
+ )
267
282
 
268
283
  try:
269
284
  resp = await get_client().run_service.CreateRun(
@@ -281,6 +296,8 @@ class _Runner:
281
296
  labels=labels,
282
297
  envs=env_kv,
283
298
  cluster=self._queue or task.queue,
299
+ raw_data_storage=raw_data_storage,
300
+ security_context=security_context,
284
301
  ),
285
302
  ),
286
303
  )
@@ -613,8 +630,8 @@ def with_runcontext(
613
630
  :param interactive_mode: Optional, can be forced to True or False.
614
631
  If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
615
632
  interactive mode, while scripts are not. This is used to determine how the code bundle is created.
616
- :param raw_data_path: Use this path to store the raw data for the run. Currently only supported for local runs,
617
- and can be used to store raw data in specific locations. TODO coming soon for remote runs as well.
633
+ :param raw_data_path: Use this path to store the raw data for the run for local and remote, and can be used to
634
+ store raw data in specific locations.
618
635
  :param run_base_dir: Optional The base directory to use for the run. This is used to store the metadata for the run,
619
636
  that is passed between tasks.
620
637
  :param overwrite_cache: Optional If true, the cache will be overwritten for the run