flyte 0.1.0__py3-none-any.whl → 0.2.0b1__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 (205) hide show
  1. flyte/__init__.py +62 -2
  2. flyte/_api_commons.py +3 -0
  3. flyte/_bin/__init__.py +0 -0
  4. flyte/_bin/runtime.py +126 -0
  5. flyte/_build.py +25 -0
  6. flyte/_cache/__init__.py +12 -0
  7. flyte/_cache/cache.py +146 -0
  8. flyte/_cache/defaults.py +9 -0
  9. flyte/_cache/policy_function_body.py +42 -0
  10. flyte/_cli/__init__.py +0 -0
  11. flyte/_cli/_common.py +299 -0
  12. flyte/_cli/_create.py +42 -0
  13. flyte/_cli/_delete.py +23 -0
  14. flyte/_cli/_deploy.py +140 -0
  15. flyte/_cli/_get.py +235 -0
  16. flyte/_cli/_params.py +538 -0
  17. flyte/_cli/_run.py +174 -0
  18. flyte/_cli/main.py +98 -0
  19. flyte/_code_bundle/__init__.py +8 -0
  20. flyte/_code_bundle/_ignore.py +113 -0
  21. flyte/_code_bundle/_packaging.py +187 -0
  22. flyte/_code_bundle/_utils.py +339 -0
  23. flyte/_code_bundle/bundle.py +178 -0
  24. flyte/_context.py +146 -0
  25. flyte/_datastructures.py +342 -0
  26. flyte/_deploy.py +202 -0
  27. flyte/_doc.py +29 -0
  28. flyte/_docstring.py +32 -0
  29. flyte/_environment.py +43 -0
  30. flyte/_group.py +31 -0
  31. flyte/_hash.py +23 -0
  32. flyte/_image.py +757 -0
  33. flyte/_initialize.py +643 -0
  34. flyte/_interface.py +84 -0
  35. flyte/_internal/__init__.py +3 -0
  36. flyte/_internal/controllers/__init__.py +115 -0
  37. flyte/_internal/controllers/_local_controller.py +118 -0
  38. flyte/_internal/controllers/_trace.py +40 -0
  39. flyte/_internal/controllers/pbhash.py +39 -0
  40. flyte/_internal/controllers/remote/__init__.py +40 -0
  41. flyte/_internal/controllers/remote/_action.py +141 -0
  42. flyte/_internal/controllers/remote/_client.py +43 -0
  43. flyte/_internal/controllers/remote/_controller.py +361 -0
  44. flyte/_internal/controllers/remote/_core.py +402 -0
  45. flyte/_internal/controllers/remote/_informer.py +361 -0
  46. flyte/_internal/controllers/remote/_service_protocol.py +50 -0
  47. flyte/_internal/imagebuild/__init__.py +11 -0
  48. flyte/_internal/imagebuild/docker_builder.py +416 -0
  49. flyte/_internal/imagebuild/image_builder.py +241 -0
  50. flyte/_internal/imagebuild/remote_builder.py +0 -0
  51. flyte/_internal/resolvers/__init__.py +0 -0
  52. flyte/_internal/resolvers/_task_module.py +54 -0
  53. flyte/_internal/resolvers/common.py +31 -0
  54. flyte/_internal/resolvers/default.py +28 -0
  55. flyte/_internal/runtime/__init__.py +0 -0
  56. flyte/_internal/runtime/convert.py +205 -0
  57. flyte/_internal/runtime/entrypoints.py +135 -0
  58. flyte/_internal/runtime/io.py +136 -0
  59. flyte/_internal/runtime/resources_serde.py +138 -0
  60. flyte/_internal/runtime/task_serde.py +210 -0
  61. flyte/_internal/runtime/taskrunner.py +190 -0
  62. flyte/_internal/runtime/types_serde.py +54 -0
  63. flyte/_logging.py +124 -0
  64. flyte/_protos/__init__.py +0 -0
  65. flyte/_protos/common/authorization_pb2.py +66 -0
  66. flyte/_protos/common/authorization_pb2.pyi +108 -0
  67. flyte/_protos/common/authorization_pb2_grpc.py +4 -0
  68. flyte/_protos/common/identifier_pb2.py +71 -0
  69. flyte/_protos/common/identifier_pb2.pyi +82 -0
  70. flyte/_protos/common/identifier_pb2_grpc.py +4 -0
  71. flyte/_protos/common/identity_pb2.py +48 -0
  72. flyte/_protos/common/identity_pb2.pyi +72 -0
  73. flyte/_protos/common/identity_pb2_grpc.py +4 -0
  74. flyte/_protos/common/list_pb2.py +36 -0
  75. flyte/_protos/common/list_pb2.pyi +69 -0
  76. flyte/_protos/common/list_pb2_grpc.py +4 -0
  77. flyte/_protos/common/policy_pb2.py +37 -0
  78. flyte/_protos/common/policy_pb2.pyi +27 -0
  79. flyte/_protos/common/policy_pb2_grpc.py +4 -0
  80. flyte/_protos/common/role_pb2.py +37 -0
  81. flyte/_protos/common/role_pb2.pyi +53 -0
  82. flyte/_protos/common/role_pb2_grpc.py +4 -0
  83. flyte/_protos/common/runtime_version_pb2.py +28 -0
  84. flyte/_protos/common/runtime_version_pb2.pyi +24 -0
  85. flyte/_protos/common/runtime_version_pb2_grpc.py +4 -0
  86. flyte/_protos/logs/dataplane/payload_pb2.py +96 -0
  87. flyte/_protos/logs/dataplane/payload_pb2.pyi +168 -0
  88. flyte/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
  89. flyte/_protos/secret/definition_pb2.py +49 -0
  90. flyte/_protos/secret/definition_pb2.pyi +93 -0
  91. flyte/_protos/secret/definition_pb2_grpc.py +4 -0
  92. flyte/_protos/secret/payload_pb2.py +62 -0
  93. flyte/_protos/secret/payload_pb2.pyi +94 -0
  94. flyte/_protos/secret/payload_pb2_grpc.py +4 -0
  95. flyte/_protos/secret/secret_pb2.py +38 -0
  96. flyte/_protos/secret/secret_pb2.pyi +6 -0
  97. flyte/_protos/secret/secret_pb2_grpc.py +198 -0
  98. flyte/_protos/secret/secret_pb2_grpc_grpc.py +198 -0
  99. flyte/_protos/validate/validate/validate_pb2.py +76 -0
  100. flyte/_protos/workflow/node_execution_service_pb2.py +26 -0
  101. flyte/_protos/workflow/node_execution_service_pb2.pyi +4 -0
  102. flyte/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
  103. flyte/_protos/workflow/queue_service_pb2.py +106 -0
  104. flyte/_protos/workflow/queue_service_pb2.pyi +141 -0
  105. flyte/_protos/workflow/queue_service_pb2_grpc.py +172 -0
  106. flyte/_protos/workflow/run_definition_pb2.py +128 -0
  107. flyte/_protos/workflow/run_definition_pb2.pyi +310 -0
  108. flyte/_protos/workflow/run_definition_pb2_grpc.py +4 -0
  109. flyte/_protos/workflow/run_logs_service_pb2.py +41 -0
  110. flyte/_protos/workflow/run_logs_service_pb2.pyi +28 -0
  111. flyte/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
  112. flyte/_protos/workflow/run_service_pb2.py +133 -0
  113. flyte/_protos/workflow/run_service_pb2.pyi +175 -0
  114. flyte/_protos/workflow/run_service_pb2_grpc.py +412 -0
  115. flyte/_protos/workflow/state_service_pb2.py +58 -0
  116. flyte/_protos/workflow/state_service_pb2.pyi +71 -0
  117. flyte/_protos/workflow/state_service_pb2_grpc.py +138 -0
  118. flyte/_protos/workflow/task_definition_pb2.py +72 -0
  119. flyte/_protos/workflow/task_definition_pb2.pyi +65 -0
  120. flyte/_protos/workflow/task_definition_pb2_grpc.py +4 -0
  121. flyte/_protos/workflow/task_service_pb2.py +44 -0
  122. flyte/_protos/workflow/task_service_pb2.pyi +31 -0
  123. flyte/_protos/workflow/task_service_pb2_grpc.py +104 -0
  124. flyte/_resources.py +226 -0
  125. flyte/_retry.py +32 -0
  126. flyte/_reusable_environment.py +25 -0
  127. flyte/_run.py +410 -0
  128. flyte/_secret.py +61 -0
  129. flyte/_task.py +367 -0
  130. flyte/_task_environment.py +200 -0
  131. flyte/_timeout.py +47 -0
  132. flyte/_tools.py +27 -0
  133. flyte/_trace.py +128 -0
  134. flyte/_utils/__init__.py +20 -0
  135. flyte/_utils/asyn.py +119 -0
  136. flyte/_utils/coro_management.py +25 -0
  137. flyte/_utils/file_handling.py +72 -0
  138. flyte/_utils/helpers.py +108 -0
  139. flyte/_utils/lazy_module.py +54 -0
  140. flyte/_utils/uv_script_parser.py +49 -0
  141. flyte/_version.py +21 -0
  142. flyte/config/__init__.py +168 -0
  143. flyte/config/_config.py +196 -0
  144. flyte/config/_internal.py +64 -0
  145. flyte/connectors/__init__.py +0 -0
  146. flyte/errors.py +143 -0
  147. flyte/extras/__init__.py +5 -0
  148. flyte/extras/_container.py +273 -0
  149. flyte/io/__init__.py +11 -0
  150. flyte/io/_dataframe.py +0 -0
  151. flyte/io/_dir.py +448 -0
  152. flyte/io/_file.py +468 -0
  153. flyte/io/pickle/__init__.py +0 -0
  154. flyte/io/pickle/transformer.py +117 -0
  155. flyte/io/structured_dataset/__init__.py +129 -0
  156. flyte/io/structured_dataset/basic_dfs.py +219 -0
  157. flyte/io/structured_dataset/structured_dataset.py +1061 -0
  158. flyte/remote/__init__.py +25 -0
  159. flyte/remote/_client/__init__.py +0 -0
  160. flyte/remote/_client/_protocols.py +131 -0
  161. flyte/remote/_client/auth/__init__.py +12 -0
  162. flyte/remote/_client/auth/_authenticators/__init__.py +0 -0
  163. flyte/remote/_client/auth/_authenticators/base.py +397 -0
  164. flyte/remote/_client/auth/_authenticators/client_credentials.py +73 -0
  165. flyte/remote/_client/auth/_authenticators/device_code.py +118 -0
  166. flyte/remote/_client/auth/_authenticators/external_command.py +79 -0
  167. flyte/remote/_client/auth/_authenticators/factory.py +200 -0
  168. flyte/remote/_client/auth/_authenticators/pkce.py +516 -0
  169. flyte/remote/_client/auth/_channel.py +184 -0
  170. flyte/remote/_client/auth/_client_config.py +83 -0
  171. flyte/remote/_client/auth/_default_html.py +32 -0
  172. flyte/remote/_client/auth/_grpc_utils/__init__.py +0 -0
  173. flyte/remote/_client/auth/_grpc_utils/auth_interceptor.py +288 -0
  174. flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +151 -0
  175. flyte/remote/_client/auth/_keyring.py +143 -0
  176. flyte/remote/_client/auth/_token_client.py +260 -0
  177. flyte/remote/_client/auth/errors.py +16 -0
  178. flyte/remote/_client/controlplane.py +95 -0
  179. flyte/remote/_console.py +18 -0
  180. flyte/remote/_data.py +155 -0
  181. flyte/remote/_logs.py +116 -0
  182. flyte/remote/_project.py +86 -0
  183. flyte/remote/_run.py +873 -0
  184. flyte/remote/_secret.py +132 -0
  185. flyte/remote/_task.py +227 -0
  186. flyte/report/__init__.py +3 -0
  187. flyte/report/_report.py +178 -0
  188. flyte/report/_template.html +124 -0
  189. flyte/storage/__init__.py +24 -0
  190. flyte/storage/_remote_fs.py +34 -0
  191. flyte/storage/_storage.py +251 -0
  192. flyte/storage/_utils.py +5 -0
  193. flyte/types/__init__.py +13 -0
  194. flyte/types/_interface.py +25 -0
  195. flyte/types/_renderer.py +162 -0
  196. flyte/types/_string_literals.py +120 -0
  197. flyte/types/_type_engine.py +2211 -0
  198. flyte/types/_utils.py +80 -0
  199. flyte-0.2.0b1.dist-info/METADATA +179 -0
  200. flyte-0.2.0b1.dist-info/RECORD +204 -0
  201. {flyte-0.1.0.dist-info → flyte-0.2.0b1.dist-info}/WHEEL +2 -1
  202. flyte-0.2.0b1.dist-info/entry_points.txt +3 -0
  203. flyte-0.2.0b1.dist-info/top_level.txt +1 -0
  204. flyte-0.1.0.dist-info/METADATA +0 -6
  205. flyte-0.1.0.dist-info/RECORD +0 -5
flyte/_image.py ADDED
@@ -0,0 +1,757 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ import hashlib
5
+ import sys
6
+ from abc import abstractmethod
7
+ from dataclasses import asdict, dataclass, field
8
+ from functools import cached_property
9
+ from pathlib import Path
10
+ from typing import Callable, ClassVar, Dict, List, Literal, Optional, Tuple, TypeVar, Union
11
+
12
+ import rich.repr
13
+
14
+ # Supported Python versions
15
+ PYTHON_3_10 = (3, 10)
16
+ PYTHON_3_11 = (3, 11)
17
+ PYTHON_3_12 = (3, 12)
18
+ PYTHON_3_13 = (3, 13)
19
+
20
+ # 0 is a file, 1 is a directory
21
+ CopyConfigType = Literal[0, 1]
22
+
23
+ T = TypeVar("T")
24
+
25
+
26
+ def _ensure_tuple(val: Union[T, List[T], Tuple[T, ...]]) -> Tuple[T] | Tuple[T, ...]:
27
+ """
28
+ Ensure that the input is a tuple. If it is a string, convert it to a tuple with one element.
29
+ If it is a list, convert it to a tuple.
30
+ """
31
+ if isinstance(val, list):
32
+ return tuple(val)
33
+ elif isinstance(val, tuple):
34
+ return val
35
+ else:
36
+ return (val,)
37
+
38
+
39
+ @rich.repr.auto
40
+ @dataclass(frozen=True, repr=True, kw_only=True)
41
+ class Layer:
42
+ """
43
+ This is an abstract representation of Container Image Layers, which can be used to create
44
+ layered images programmatically.
45
+ """
46
+
47
+ _compute_identifier: Callable[[Layer], str] = field(default=lambda x: x.__str__(), init=True)
48
+
49
+ @abstractmethod
50
+ def update_hash(self, hasher: hashlib._Hash):
51
+ """
52
+ This method should be implemented by subclasses to provide a hash representation of the layer.
53
+
54
+ :param hasher: The hash object to update with the layer's data.
55
+ """
56
+ ...
57
+
58
+ def validate(self):
59
+ """
60
+ Raise any validation errors for the layer
61
+ :return:
62
+ """
63
+
64
+
65
+ @rich.repr.auto
66
+ @dataclass(kw_only=True, frozen=True, repr=True)
67
+ class PipPackages(Layer):
68
+ packages: Optional[Tuple[str, ...]] = None
69
+ index_url: Optional[str] = None
70
+ extra_index_urls: Optional[Tuple[str] | Tuple[str, ...] | List[str]] = None
71
+ pre: bool = False
72
+ extra_args: Optional[str] = None
73
+
74
+ # todo: to be implemented
75
+ # secret_mounts: Optional[List[Tuple[str, str]]] = None
76
+
77
+ def update_hash(self, hasher: hashlib._Hash):
78
+ """
79
+ Update the hash with the pip packages
80
+ """
81
+ hash_input = ""
82
+ if self.packages:
83
+ for package in self.packages:
84
+ hash_input += package
85
+ if self.index_url:
86
+ hash_input += self.index_url
87
+ if self.extra_index_urls:
88
+ for url in self.extra_index_urls:
89
+ hash_input += url
90
+ if self.pre:
91
+ hash_input += str(self.pre)
92
+ if self.extra_args:
93
+ hash_input += self.extra_args
94
+
95
+ hasher.update(hash_input.encode("utf-8"))
96
+
97
+
98
+ @rich.repr.auto
99
+ @dataclass(kw_only=True, frozen=True, repr=True)
100
+ class Requirements(PipPackages):
101
+ file: Path
102
+
103
+ def update_hash(self, hasher: hashlib._Hash):
104
+ from ._utils import filehash_update
105
+
106
+ super().update_hash(hasher)
107
+ filehash_update(self.file, hasher)
108
+
109
+
110
+ @rich.repr.auto
111
+ @dataclass(frozen=True, repr=True)
112
+ class AptPackages(Layer):
113
+ packages: Tuple[str, ...]
114
+
115
+ def update_hash(self, hasher: hashlib._Hash):
116
+ hasher.update("".join(self.packages).encode("utf-8"))
117
+
118
+
119
+ @rich.repr.auto
120
+ @dataclass(frozen=True, repr=True)
121
+ class Commands(Layer):
122
+ commands: Tuple[str, ...]
123
+
124
+ def update_hash(self, hasher: hashlib._Hash):
125
+ hasher.update("".join(self.commands).encode("utf-8"))
126
+
127
+
128
+ @rich.repr.auto
129
+ @dataclass(frozen=True, repr=True)
130
+ class WorkDir(Layer):
131
+ workdir: str
132
+
133
+ def update_hash(self, hasher: hashlib._Hash):
134
+ hasher.update(self.workdir.encode("utf-8"))
135
+
136
+
137
+ @rich.repr.auto
138
+ @dataclass(frozen=True, repr=True)
139
+ class CopyConfig(Layer):
140
+ path_type: CopyConfigType
141
+ context_source: Path
142
+ image_dest: str = "."
143
+
144
+ def validate(self):
145
+ if not self.context_source.exists():
146
+ raise ValueError(f"Source folder {self.context_source.absolute()} does not exist")
147
+ if not self.context_source.is_dir() and self.path_type == 1:
148
+ raise ValueError(f"Source folder {self.context_source.absolute()} is not a directory")
149
+ if not self.context_source.is_file() and self.path_type == 0:
150
+ raise ValueError(f"Source file {self.context_source.absolute()} is not a file")
151
+
152
+ def update_hash(self, hasher: hashlib._Hash):
153
+ from ._utils import update_hasher_for_source
154
+
155
+ update_hasher_for_source(self.context_source, hasher)
156
+ if self.image_dest:
157
+ hasher.update(self.image_dest.encode("utf-8"))
158
+
159
+
160
+ @rich.repr.auto
161
+ @dataclass(frozen=True, repr=True)
162
+ class UVProject(Layer):
163
+ pyproject: Path
164
+ uvlock: Path
165
+
166
+ def update_hash(self, hasher: hashlib._Hash):
167
+ from ._utils import filehash_update
168
+
169
+ filehash_update(self.uvlock, hasher)
170
+
171
+
172
+ @rich.repr.auto
173
+ @dataclass(frozen=True, repr=True)
174
+ class _DockerLines(Layer):
175
+ """
176
+ This is an internal class and should only be used by the default images. It is not supported by most
177
+ builders so please don't use it.
178
+ """
179
+
180
+ lines: Tuple[str, ...]
181
+
182
+ def update_hash(self, hasher: hashlib._Hash):
183
+ hasher.update("".join(self.lines).encode("utf-8"))
184
+
185
+
186
+ @rich.repr.auto
187
+ @dataclass(frozen=True, repr=True)
188
+ class Env(Layer):
189
+ """
190
+ This is an internal class and should only be used by the default images. It is not supported by most
191
+ builders so please don't use it.
192
+ """
193
+
194
+ env_vars: Tuple[Tuple[str, str], ...] = field(default_factory=tuple)
195
+
196
+ def update_hash(self, hasher: hashlib._Hash):
197
+ txt = [f"{k}={v}" for k, v in self.env_vars]
198
+ hasher.update(" ".join(txt).encode("utf-8"))
199
+
200
+ @classmethod
201
+ def from_dict(cls, envs: Dict[str, str]) -> Env:
202
+ return cls(env_vars=tuple((k, v) for k, v in envs.items()))
203
+
204
+
205
+ Architecture = Literal["linux/amd64", "linux/arm64"]
206
+
207
+ _BASE_REGISTRY = "ghcr.io/unionai-oss"
208
+ _DEFAULT_IMAGE_NAME = "flyte"
209
+
210
+
211
+ def _detect_python_version() -> Tuple[int, int]:
212
+ """
213
+ Detect the current Python version.
214
+ :return: Tuple of major and minor version
215
+ """
216
+ return sys.version_info.major, sys.version_info.minor
217
+
218
+
219
+ @dataclass(frozen=True, repr=True, eq=True)
220
+ class Image:
221
+ """
222
+ This is a representation of Container Images, which can be used to create layered images programmatically.
223
+
224
+ Use by first calling one of the base constructor methods. These all begin with `from` or `default_`
225
+ The image can then be amended with additional layers using the various `with_*` methods.
226
+
227
+ Invariant for this class: The construction of Image objects must be doable everywhere. That is, if a
228
+ user has a custom image that is not accessible, calling .with_source_file on a file that doesn't exist, the
229
+ instantiation of the object itself must still go through. Further, the .identifier property of the image must
230
+ also still go through. This is because it may have been already built somewhere else.
231
+ Use validate() functions to check each layer for actual errors. These are invoked at actual
232
+ build time. See self.id for more information
233
+ """
234
+
235
+ # These are base properties of an image
236
+ base_image: Optional[str] = field(default=None)
237
+ dockerfile: Optional[Path] = field(default=None)
238
+ registry: Optional[str] = field(default=None)
239
+ name: Optional[str] = field(default=None)
240
+ platform: Tuple[Architecture, ...] = field(default=("linux/amd64",))
241
+ tag: Optional[str] = field(default=None)
242
+ python_version: Tuple[int, int] = field(default_factory=_detect_python_version)
243
+
244
+ # For .auto() images. Don't compute an actual identifier.
245
+ _identifier_override: Optional[str] = field(default=None, init=False)
246
+ # This is set on default images. These images are built from the base Dockerfile in this library and shouldn't be
247
+ # modified with additional layers.
248
+ is_final: bool = field(default=False)
249
+
250
+ # Layers to be added to the image. In init, because frozen, but users shouldn't access, so underscore.
251
+ _layers: Tuple[Layer, ...] = field(default_factory=tuple)
252
+
253
+ _DEFAULT_IMAGE_PREFIXES: ClassVar = {
254
+ PYTHON_3_10: "py3.10-",
255
+ PYTHON_3_11: "py3.11-",
256
+ PYTHON_3_12: "py3.12-",
257
+ PYTHON_3_13: "py3.13-",
258
+ }
259
+
260
+ @cached_property
261
+ def identifier(self) -> str:
262
+ """
263
+ This identifier is a hash of the layers and properties of the image. It is used to look up previously built
264
+ images. Why is this useful? For example, if a user has Image.from_uv_base().with_source_file("a/local/file"),
265
+ it's not necessarily the case that that file exists within the image (further commands may have removed/changed
266
+ it), and certainly not the case that the path to the file, inside the image (which is used as part of the layer
267
+ hash computation), is the same. That is, inside the image when a task runs, as we come across the same Image
268
+ declaration, we need a way of identifying the image and its uri, without hashing all the layers again. This
269
+ is what this identifier is for. See the ImageCache object for additional information.
270
+
271
+ :return: A unique identifier of the Image
272
+ """
273
+ if self._identifier_override:
274
+ return self._identifier_override
275
+
276
+ # Only get the non-None values in the ImageSpec to ensure the hash is consistent
277
+ # across different SDK versions.
278
+ # Can potentially add a second hashing function to the Layer protocol, but relying on just asdict/str
279
+ # representation for now.
280
+ image_dict = asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if v is not None and k != "_layers"})
281
+ layers_str_repr = "".join([layer._compute_identifier(layer) for layer in self._layers])
282
+ image_dict["layers"] = layers_str_repr
283
+ spec_bytes = image_dict.__str__().encode("utf-8")
284
+ return base64.urlsafe_b64encode(hashlib.md5(spec_bytes).digest()).decode("ascii").rstrip("=")
285
+
286
+ def validate(self):
287
+ for layer in self._layers:
288
+ layer.validate()
289
+
290
+ @classmethod
291
+ def _get_default_image_for(cls, python_version: Tuple[int, int], flyte_version: Optional[str] = None) -> Image:
292
+ # Would love a way to move this outside of this class (but still needs to be accessible via Image.auto())
293
+ # this default image definition may need to be updated once there is a released pypi version
294
+ from flyte._version import __version__
295
+
296
+ dev_mode = cls._is_editable_install() or (__version__ and "dev" in __version__)
297
+ if flyte_version is None:
298
+ flyte_version = __version__.replace("+", "-")
299
+ preset_tag = flyte_version if flyte_version.startswith("v") else f"v{flyte_version}"
300
+ preset_tag = f"py{python_version[0]}.{python_version[1]}-{preset_tag}"
301
+ image = Image(
302
+ base_image=f"python:{python_version[0]}.{python_version[1]}-slim-bookworm",
303
+ registry=_BASE_REGISTRY,
304
+ name=_DEFAULT_IMAGE_NAME,
305
+ tag=preset_tag,
306
+ platform=("linux/amd64", "linux/arm64"),
307
+ )
308
+ labels_and_user = _DockerLines(
309
+ (
310
+ "LABEL org.opencontainers.image.authors='Union.AI <sales@union.ai>'",
311
+ "LABEL org.opencontainers.image.source=https://github.com/unionai/unionv2",
312
+ "RUN useradd --create-home --shell /bin/bash flytekit &&"
313
+ " chown -R flytekit /root && chown -R flytekit /home",
314
+ "WORKDIR /root",
315
+ )
316
+ )
317
+ image = image.clone(addl_layer=labels_and_user)
318
+ image = image.with_env_vars(
319
+ {
320
+ "VIRTUAL_ENV": "/opt/venv",
321
+ "PATH": "/opt/venv/bin:$PATH",
322
+ "PYTHONPATH": "/root",
323
+ "UV_LINK_MODE": "copy",
324
+ }
325
+ )
326
+ image = image.with_apt_packages(["build-essential", "ca-certificates"])
327
+
328
+ base_packages = ["kubernetes", "msgpack", "mashumaro"]
329
+
330
+ # Add in flyte library
331
+ if dev_mode:
332
+ image = image.with_pip_packages(base_packages)
333
+ image = image.with_local_v2()
334
+ else:
335
+ base_packages.append(f"flyte=={flyte_version}")
336
+ image = image.with_pip_packages(base_packages)
337
+
338
+ return image
339
+
340
+ @staticmethod
341
+ def _is_editable_install():
342
+ """Internal hacky function to see if the current install is editable or not."""
343
+ curr = Path(__file__)
344
+ pyproject = curr.parent.parent.parent / "pyproject.toml"
345
+ return pyproject.exists()
346
+
347
+ @classmethod
348
+ def from_uv_debian(
349
+ cls,
350
+ registry: str,
351
+ name: str,
352
+ tag: Optional[str] = None,
353
+ python_version: Optional[Tuple[int, int]] = None,
354
+ arch: Union[Architecture, Tuple[Architecture, ...]] = "linux/amd64",
355
+ ) -> Image:
356
+ """
357
+ This creates a new debian-based base image.
358
+ If using the Union or docker builders, image will have uv available and a virtualenv created at /opt/venv.
359
+
360
+ :param registry: Registry to use for the image
361
+ :param name: Name of the image
362
+ :param tag: Tag to use for the image
363
+ :param python_version: Python version to use for the image
364
+ :param arch: Architecture to use for the image, default is linux/amd64
365
+ :return: Image
366
+ """
367
+ base_image = "debian:bookworm-slim"
368
+ plat = arch if isinstance(arch, tuple) else (arch,)
369
+ if python_version is None:
370
+ python_version = _detect_python_version()
371
+ img = cls(
372
+ base_image=base_image, name=name, registry=registry, tag=tag, platform=plat, python_version=python_version
373
+ )
374
+ return img
375
+
376
+ @classmethod
377
+ def auto(
378
+ cls,
379
+ python_version: Optional[Tuple[int, int]] = None,
380
+ flyte_version: Optional[str] = None,
381
+ registry: Optional[str] = None,
382
+ name: Optional[str] = None,
383
+ ) -> Image:
384
+ """
385
+ Use this method to start using the default base image, built from this library's base Dockerfile
386
+ Default images are multi-arch amd/arm64
387
+
388
+ :param python_version: If not specified, will use the current Python version
389
+ :param flyte_version: Union version to use
390
+ :param registry: Registry to use for the image
391
+ :param name: Name of the image if you want to override the default name
392
+
393
+ :return: Image
394
+ """
395
+ if python_version is None:
396
+ python_version = _detect_python_version()
397
+
398
+ base_image = cls._get_default_image_for(python_version=python_version, flyte_version=flyte_version)
399
+ if name is not None and registry is None:
400
+ raise ValueError("Both name and registry must be specified to override the default image name.")
401
+
402
+ if registry and name:
403
+ return base_image.clone(registry=registry, name=name)
404
+
405
+ # Set this to auto for all auto images because the meaning of "auto" can change (based on logic inside
406
+ # _get_default_image_for, acts differently in a running task container) so let's make sure it stays auto.
407
+ object.__setattr__(base_image, "_identifier_override", "auto")
408
+ return base_image
409
+
410
+ @classmethod
411
+ def from_prebuilt(cls, image_uri: str) -> Image:
412
+ """
413
+ Use this method to start with a pre-built base image. This image must already exist in the registry of course.
414
+
415
+ :param image_uri: The full URI of the image, in the format <registry>/<name>:<tag>
416
+ :return:
417
+ """
418
+ img = cls(base_image=image_uri)
419
+ return img
420
+
421
+ @classmethod
422
+ def from_uv_script(
423
+ cls,
424
+ script: Path | str,
425
+ *,
426
+ name: str,
427
+ registry: str | None = None,
428
+ python_version: Optional[Tuple[int, int]] = None,
429
+ arch: Union[Architecture, Tuple[Architecture, ...]] = "linux/amd64",
430
+ ) -> Image:
431
+ """
432
+ Use this method to create a new image with the specified uv script.
433
+ It uses the header of the script to determine the python version, dependencies to install.
434
+ The script must be a valid uv script, otherwise an error will be raised.
435
+
436
+ Usually the header of the script will look like this:
437
+ Example:
438
+ ```python
439
+ #!/usr/bin/env -S uv run --script
440
+ # /// script
441
+ # requires-python = ">=3.12"
442
+ # dependencies = ["httpx"]
443
+ # ///
444
+ ```
445
+
446
+ For more information on the uv script format, see the documentation:
447
+ <href="https://docs.astral.sh/uv/guides/scripts/#declaring-script-dependencies">
448
+ UV: Declaring script dependencies</href>
449
+
450
+ :param name: name of the image
451
+ :param registry: registry to use for the image
452
+ :param script: path to the uv script
453
+ :param arch: architecture to use for the image, default is linux/amd64, use tuple for multiple values
454
+
455
+ :return: Image
456
+ """
457
+ from ._utils import parse_uv_script_file
458
+
459
+ if isinstance(script, str):
460
+ script = Path(script)
461
+ if not script.exists():
462
+ raise FileNotFoundError(f"UV script {script} does not exist")
463
+ if not script.is_file():
464
+ raise ValueError(f"UV script {script} is not a file")
465
+ if not script.suffix == ".py":
466
+ raise ValueError(f"UV script {script} must have a .py extension")
467
+ header = parse_uv_script_file(script)
468
+ if registry is None:
469
+ raise ValueError("registry must be specified")
470
+ img = cls.from_uv_debian(registry=registry, name=name, arch=arch, python_version=python_version)
471
+ if header.dependencies:
472
+ return img.with_pip_packages(header.dependencies)
473
+ # todo: override the _identifier_override to be the script name or a hash of the script contents
474
+ # This is needed because inside the image, the identifier will be computed to be something different.
475
+ return img
476
+
477
+ def clone(
478
+ self, registry: Optional[str] = None, name: Optional[str] = None, addl_layer: Optional[Layer] = None
479
+ ) -> Image:
480
+ """
481
+ Use this method to clone the current image and change the registry and name
482
+
483
+ :param registry: Registry to use for the image
484
+ :param name: Name of the image
485
+
486
+ :return:
487
+ """
488
+ registry = registry if registry else self.registry
489
+ name = name if name else self.name
490
+ new_layers = (*self._layers, addl_layer) if addl_layer else self._layers
491
+ img = Image(
492
+ base_image=self.base_image,
493
+ dockerfile=self.dockerfile,
494
+ registry=registry,
495
+ name=name,
496
+ tag=self.tag,
497
+ platform=self.platform,
498
+ python_version=self.python_version,
499
+ is_final=self.is_final,
500
+ _layers=new_layers,
501
+ )
502
+
503
+ return img
504
+
505
+ @classmethod
506
+ def from_dockerfile(cls, file: Path, registry: str, name: str, tag: Optional[str] = None) -> Image:
507
+ """
508
+ Use this method to create a new image with the specified dockerfile
509
+
510
+ :param file: path to the dockerfile
511
+ :param name: name of the image
512
+ :param registry: registry to use for the image
513
+ :param tag: tag to use for the image
514
+
515
+ :return:
516
+ """
517
+ tag = tag or "latest"
518
+ img = cls(dockerfile=file, registry=registry, name=name, tag=tag)
519
+
520
+ return img
521
+
522
+ def _get_hash_digest(self) -> str:
523
+ """
524
+ Returns the hash digest of the image, which is a combination of all the layers and properties of the image
525
+ """
526
+ import hashlib
527
+
528
+ from ._utils import filehash_update
529
+
530
+ hasher = hashlib.md5()
531
+ if self.dockerfile:
532
+ # Note the location of the dockerfile shouldn't matter, only the contents
533
+ filehash_update(self.dockerfile, hasher)
534
+ if self._layers:
535
+ for layer in self._layers:
536
+ layer.update_hash(hasher)
537
+ return hasher.hexdigest()
538
+
539
+ @property
540
+ def _final_tag(self) -> str:
541
+ t = self.tag if self.tag else self._get_hash_digest()
542
+ return t or "latest"
543
+
544
+ @cached_property
545
+ def uri(self) -> str:
546
+ """
547
+ Returns the URI of the image in the format <registry>/<name>:<tag>
548
+ """
549
+ if self.registry and self.name:
550
+ tag = self._final_tag
551
+ return f"{self.registry}/{self.name}:{tag}"
552
+ elif self.name:
553
+ return f"{self.name}:{self._final_tag}"
554
+ elif self.base_image:
555
+ return self.base_image
556
+
557
+ raise ValueError("Image is not fully defined. Please set registry, name and tag.")
558
+
559
+ def with_workdir(self, workdir: str) -> Image:
560
+ """
561
+ Use this method to create a new image with the specified working directory
562
+ This will override any existing working directory
563
+
564
+ :param workdir: working directory to use
565
+ :return:
566
+ """
567
+ new_image = self.clone(addl_layer=WorkDir(workdir=workdir))
568
+ return new_image
569
+
570
+ def with_requirements(self, file: Path) -> Image:
571
+ """
572
+ Use this method to create a new image with the specified requirements file layered on top of the current image
573
+ Cannot be used in conjunction with conda
574
+
575
+ :param file: path to the requirements file, must be a .txt file
576
+ :return:
577
+ """
578
+ if not file.exists():
579
+ raise FileNotFoundError(f"Requirements file {file} does not exist")
580
+ if not file.is_file():
581
+ raise ValueError(f"Requirements file {file} is not a file")
582
+ if file.suffix != ".txt":
583
+ raise ValueError(f"Requirements file {file} must have a .txt extension")
584
+ new_image = self.clone(addl_layer=Requirements(file=file))
585
+ return new_image
586
+
587
+ def with_pip_packages(
588
+ self,
589
+ packages: Union[str, List[str], Tuple[str, ...]],
590
+ index_url: Optional[str] = None,
591
+ extra_index_urls: Union[str, List[str], Tuple[str, ...], None] = None,
592
+ pre: bool = False,
593
+ extra_args: Optional[str] = None,
594
+ ) -> Image:
595
+ """
596
+ Use this method to create a new image with the specified pip packages layered on top of the current image
597
+ Cannot be used in conjunction with conda
598
+
599
+ Example:
600
+ ```python
601
+ @flyte.task(image=(flyte.Image
602
+ .ubuntu_python()
603
+ .with_pip_packages(["requests", "numpy"])))
604
+ def my_task(x: int) -> int:
605
+ import numpy as np
606
+ return np.sum([x, 1])
607
+ ```
608
+
609
+ :param packages: list of pip packages to install, follows pip install syntax
610
+ :param index_url: index url to use for pip install, default is None
611
+ :param extra_index_urls: extra index urls to use for pip install, default is None
612
+ :param pre: whether to allow pre-release versions, default is False
613
+ :param extra_args: extra arguments to pass to pip install, default is None
614
+ # :param secret_mounts: todo
615
+ :param extra_args: extra arguments to pass to pip install, default is None
616
+ :return: Image
617
+ """
618
+
619
+ new_packages: Optional[Tuple] = _ensure_tuple(packages) if packages else None
620
+ new_extra_index_urls: Optional[Tuple] = _ensure_tuple(extra_index_urls) if extra_index_urls else None
621
+
622
+ ll = PipPackages(
623
+ packages=new_packages,
624
+ index_url=index_url,
625
+ extra_index_urls=new_extra_index_urls,
626
+ pre=pre,
627
+ extra_args=extra_args,
628
+ )
629
+ new_image = self.clone(addl_layer=ll)
630
+ return new_image
631
+
632
+ def with_env_vars(self, env_vars: Dict[str, str]) -> Image:
633
+ """
634
+ Use this method to create a new image with the specified environment variables layered on top of
635
+ the current image. Cannot be used in conjunction with conda
636
+
637
+ :param env_vars: dictionary of environment variables to set
638
+ :return: Image
639
+ """
640
+ new_image = self.clone(addl_layer=Env.from_dict(env_vars))
641
+ return new_image
642
+
643
+ def with_source_folder(self, context_source: Path, image_dest: Optional[str] = None) -> Image:
644
+ """
645
+ Use this method to create a new image with the specified local directory layered on top of the current image.
646
+ If dest is not specified, it will be copied to the working directory of the image
647
+
648
+ :param context_source: root folder of the source code from the build context to be copied
649
+ :param image_dest: destination folder in the image
650
+ :return: Image
651
+ """
652
+ image_dest = image_dest if image_dest else "."
653
+ new_image = self.clone(addl_layer=CopyConfig(path_type=1, context_source=context_source, image_dest=image_dest))
654
+ return new_image
655
+
656
+ def with_source_file(self, context_source: Path, image_dest: Optional[str] = None) -> Image:
657
+ """
658
+ Use this method to create a new image with the specified local file layered on top of the current image.
659
+ If dest is not specified, it will be copied to the working directory of the image
660
+
661
+ :param context_source: root folder of the source code from the build context to be copied
662
+ :param image_dest: destination folder in the image
663
+ :return: Image
664
+ """
665
+ image_dest = image_dest if image_dest else "."
666
+ new_image = self.clone(addl_layer=CopyConfig(path_type=0, context_source=context_source, image_dest=image_dest))
667
+ return new_image
668
+
669
+ def with_uv_project(self, pyproject_file: Path) -> Image:
670
+ """
671
+ Use this method to create a new image with the specified uv.lock file layered on top of the current image
672
+ Must have a corresponding pyproject.toml file in the same directory
673
+ Cannot be used in conjunction with conda
674
+ In the Union builders, using this will change the virtual env to /root/.venv
675
+
676
+ :param pyproject_file: path to the pyproject.toml file, needs to have a corresponding uv.lock file
677
+ :return:
678
+ """
679
+ if not pyproject_file.exists():
680
+ raise FileNotFoundError(f"UVLock file {pyproject_file} does not exist")
681
+ if not pyproject_file.is_file():
682
+ raise ValueError(f"UVLock file {pyproject_file} is not a file")
683
+ lock = pyproject_file.parent / "uv.lock"
684
+ if not lock.exists():
685
+ raise ValueError(f"UVLock file {lock} does not exist")
686
+ new_image = self.clone(addl_layer=UVProject(pyproject=pyproject_file, uvlock=lock))
687
+ return new_image
688
+
689
+ def with_apt_packages(self, packages: Union[str, List[str], Tuple[str, ...]]) -> Image:
690
+ """
691
+ Use this method to create a new image with the specified apt packages layered on top of the current image
692
+
693
+ :param packages: list of apt packages to install
694
+ :return: Image
695
+ """
696
+ pkgs = _ensure_tuple(packages)
697
+ new_image = self.clone(addl_layer=AptPackages(packages=pkgs))
698
+ return new_image
699
+
700
+ def with_commands(self, commands: List[str]) -> Image:
701
+ """
702
+ Use this method to create a new image with the specified commands layered on top of the current image
703
+ Be sure not to use RUN in your command.
704
+
705
+ :param commands: list of commands to run
706
+ :return: Image
707
+ """
708
+ new_commands: Tuple = _ensure_tuple(commands)
709
+ new_image = self.clone(addl_layer=Commands(commands=new_commands))
710
+ return new_image
711
+
712
+ def with_local_v2(self) -> Image:
713
+ """
714
+ Use this method to create a new image with the local v2 builder
715
+ This will override any existing builder
716
+
717
+ :return: Image
718
+ """
719
+ dist_folder = Path(__file__).parent.parent.parent / "dist"
720
+ # Manually declare the CopyConfig instead of using with_source_folder so we can set the hashing
721
+ # used to compute the identifier. Can remove if we ever decide to expose the lambda in with_ commands
722
+ with_dist = self.clone(
723
+ addl_layer=CopyConfig(
724
+ path_type=1, context_source=dist_folder, image_dest=".", _compute_identifier=lambda x: "/dist"
725
+ )
726
+ )
727
+
728
+ return with_dist.with_commands(
729
+ [
730
+ "--mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv"
731
+ " --mount=from=uv,source=/uv,target=/usr/bin/uv"
732
+ " --mount=source=dist,target=/dist,type=bind"
733
+ " uv pip install --python $VIRTUALENV $(ls /dist/*whl)"
734
+ ]
735
+ )
736
+
737
+ def __img_str__(self) -> str:
738
+ """
739
+ For the current image only, print all the details if they are not None
740
+ """
741
+ details = []
742
+ if self.base_image:
743
+ details.append(f"Base Image: {self.base_image}")
744
+ elif self.dockerfile:
745
+ details.append(f"Dockerfile: {self.dockerfile}")
746
+ if self.registry:
747
+ details.append(f"Registry: {self.registry}")
748
+ if self.name:
749
+ details.append(f"Name: {self.name}")
750
+ if self.platform:
751
+ details.append(f"Platform: {self.platform}")
752
+
753
+ if self.__getattribute__("_layers"):
754
+ for layer in self._layers:
755
+ details.append(f"Layer: {layer}")
756
+
757
+ return "\n".join(details)