flyte 2.0.0b22__py3-none-any.whl → 2.0.0b30__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.
Files changed (197) hide show
  1. flyte/__init__.py +18 -2
  2. flyte/_bin/runtime.py +43 -5
  3. flyte/_cache/cache.py +4 -2
  4. flyte/_cache/local_cache.py +216 -0
  5. flyte/_code_bundle/_ignore.py +1 -1
  6. flyte/_code_bundle/_packaging.py +4 -4
  7. flyte/_code_bundle/_utils.py +14 -8
  8. flyte/_code_bundle/bundle.py +13 -5
  9. flyte/_constants.py +1 -0
  10. flyte/_context.py +4 -1
  11. flyte/_custom_context.py +73 -0
  12. flyte/_debug/constants.py +0 -1
  13. flyte/_debug/vscode.py +6 -1
  14. flyte/_deploy.py +223 -59
  15. flyte/_environment.py +5 -0
  16. flyte/_excepthook.py +1 -1
  17. flyte/_image.py +144 -82
  18. flyte/_initialize.py +95 -12
  19. flyte/_interface.py +2 -0
  20. flyte/_internal/controllers/_local_controller.py +65 -24
  21. flyte/_internal/controllers/_trace.py +1 -1
  22. flyte/_internal/controllers/remote/_action.py +13 -11
  23. flyte/_internal/controllers/remote/_client.py +1 -1
  24. flyte/_internal/controllers/remote/_controller.py +9 -4
  25. flyte/_internal/controllers/remote/_core.py +16 -16
  26. flyte/_internal/controllers/remote/_informer.py +4 -4
  27. flyte/_internal/controllers/remote/_service_protocol.py +7 -7
  28. flyte/_internal/imagebuild/docker_builder.py +139 -84
  29. flyte/_internal/imagebuild/image_builder.py +7 -13
  30. flyte/_internal/imagebuild/remote_builder.py +65 -13
  31. flyte/_internal/imagebuild/utils.py +51 -3
  32. flyte/_internal/resolvers/_task_module.py +5 -38
  33. flyte/_internal/resolvers/default.py +2 -2
  34. flyte/_internal/runtime/convert.py +42 -20
  35. flyte/_internal/runtime/entrypoints.py +24 -1
  36. flyte/_internal/runtime/io.py +21 -8
  37. flyte/_internal/runtime/resources_serde.py +20 -6
  38. flyte/_internal/runtime/reuse.py +1 -1
  39. flyte/_internal/runtime/rusty.py +20 -5
  40. flyte/_internal/runtime/task_serde.py +33 -27
  41. flyte/_internal/runtime/taskrunner.py +10 -1
  42. flyte/_internal/runtime/trigger_serde.py +160 -0
  43. flyte/_internal/runtime/types_serde.py +1 -1
  44. flyte/_keyring/file.py +39 -9
  45. flyte/_logging.py +79 -12
  46. flyte/_map.py +31 -12
  47. flyte/_module.py +70 -0
  48. flyte/_pod.py +2 -2
  49. flyte/_resources.py +213 -31
  50. flyte/_run.py +107 -41
  51. flyte/_task.py +66 -10
  52. flyte/_task_environment.py +96 -24
  53. flyte/_task_plugins.py +4 -2
  54. flyte/_trigger.py +1000 -0
  55. flyte/_utils/__init__.py +2 -1
  56. flyte/_utils/asyn.py +3 -1
  57. flyte/_utils/docker_credentials.py +173 -0
  58. flyte/_utils/module_loader.py +17 -2
  59. flyte/_version.py +3 -3
  60. flyte/cli/_abort.py +3 -3
  61. flyte/cli/_build.py +1 -3
  62. flyte/cli/_common.py +78 -7
  63. flyte/cli/_create.py +178 -3
  64. flyte/cli/_delete.py +23 -1
  65. flyte/cli/_deploy.py +49 -11
  66. flyte/cli/_get.py +79 -34
  67. flyte/cli/_params.py +8 -6
  68. flyte/cli/_plugins.py +209 -0
  69. flyte/cli/_run.py +127 -11
  70. flyte/cli/_serve.py +64 -0
  71. flyte/cli/_update.py +37 -0
  72. flyte/cli/_user.py +17 -0
  73. flyte/cli/main.py +30 -4
  74. flyte/config/_config.py +2 -0
  75. flyte/config/_internal.py +1 -0
  76. flyte/config/_reader.py +3 -3
  77. flyte/connectors/__init__.py +11 -0
  78. flyte/connectors/_connector.py +270 -0
  79. flyte/connectors/_server.py +197 -0
  80. flyte/connectors/utils.py +135 -0
  81. flyte/errors.py +10 -1
  82. flyte/extend.py +8 -1
  83. flyte/extras/_container.py +6 -1
  84. flyte/git/_config.py +11 -9
  85. flyte/io/__init__.py +2 -0
  86. flyte/io/_dataframe/__init__.py +2 -0
  87. flyte/io/_dataframe/basic_dfs.py +1 -1
  88. flyte/io/_dataframe/dataframe.py +12 -8
  89. flyte/io/_dir.py +551 -120
  90. flyte/io/_file.py +538 -141
  91. flyte/models.py +57 -12
  92. flyte/remote/__init__.py +6 -1
  93. flyte/remote/_action.py +18 -16
  94. flyte/remote/_client/_protocols.py +39 -4
  95. flyte/remote/_client/auth/_channel.py +10 -6
  96. flyte/remote/_client/controlplane.py +17 -5
  97. flyte/remote/_console.py +3 -2
  98. flyte/remote/_data.py +4 -3
  99. flyte/remote/_logs.py +3 -3
  100. flyte/remote/_run.py +47 -7
  101. flyte/remote/_secret.py +26 -17
  102. flyte/remote/_task.py +21 -9
  103. flyte/remote/_trigger.py +306 -0
  104. flyte/remote/_user.py +33 -0
  105. flyte/storage/__init__.py +6 -1
  106. flyte/storage/_parallel_reader.py +274 -0
  107. flyte/storage/_storage.py +185 -103
  108. flyte/types/__init__.py +16 -0
  109. flyte/types/_interface.py +2 -2
  110. flyte/types/_pickle.py +17 -4
  111. flyte/types/_string_literals.py +8 -9
  112. flyte/types/_type_engine.py +26 -19
  113. flyte/types/_utils.py +1 -1
  114. {flyte-2.0.0b22.data → flyte-2.0.0b30.data}/scripts/runtime.py +43 -5
  115. {flyte-2.0.0b22.dist-info → flyte-2.0.0b30.dist-info}/METADATA +8 -1
  116. flyte-2.0.0b30.dist-info/RECORD +192 -0
  117. flyte/_protos/__init__.py +0 -0
  118. flyte/_protos/common/authorization_pb2.py +0 -66
  119. flyte/_protos/common/authorization_pb2.pyi +0 -108
  120. flyte/_protos/common/authorization_pb2_grpc.py +0 -4
  121. flyte/_protos/common/identifier_pb2.py +0 -99
  122. flyte/_protos/common/identifier_pb2.pyi +0 -120
  123. flyte/_protos/common/identifier_pb2_grpc.py +0 -4
  124. flyte/_protos/common/identity_pb2.py +0 -48
  125. flyte/_protos/common/identity_pb2.pyi +0 -72
  126. flyte/_protos/common/identity_pb2_grpc.py +0 -4
  127. flyte/_protos/common/list_pb2.py +0 -36
  128. flyte/_protos/common/list_pb2.pyi +0 -71
  129. flyte/_protos/common/list_pb2_grpc.py +0 -4
  130. flyte/_protos/common/policy_pb2.py +0 -37
  131. flyte/_protos/common/policy_pb2.pyi +0 -27
  132. flyte/_protos/common/policy_pb2_grpc.py +0 -4
  133. flyte/_protos/common/role_pb2.py +0 -37
  134. flyte/_protos/common/role_pb2.pyi +0 -53
  135. flyte/_protos/common/role_pb2_grpc.py +0 -4
  136. flyte/_protos/common/runtime_version_pb2.py +0 -28
  137. flyte/_protos/common/runtime_version_pb2.pyi +0 -24
  138. flyte/_protos/common/runtime_version_pb2_grpc.py +0 -4
  139. flyte/_protos/imagebuilder/definition_pb2.py +0 -60
  140. flyte/_protos/imagebuilder/definition_pb2.pyi +0 -153
  141. flyte/_protos/imagebuilder/definition_pb2_grpc.py +0 -4
  142. flyte/_protos/imagebuilder/payload_pb2.py +0 -32
  143. flyte/_protos/imagebuilder/payload_pb2.pyi +0 -21
  144. flyte/_protos/imagebuilder/payload_pb2_grpc.py +0 -4
  145. flyte/_protos/imagebuilder/service_pb2.py +0 -29
  146. flyte/_protos/imagebuilder/service_pb2.pyi +0 -5
  147. flyte/_protos/imagebuilder/service_pb2_grpc.py +0 -66
  148. flyte/_protos/logs/dataplane/payload_pb2.py +0 -100
  149. flyte/_protos/logs/dataplane/payload_pb2.pyi +0 -177
  150. flyte/_protos/logs/dataplane/payload_pb2_grpc.py +0 -4
  151. flyte/_protos/secret/definition_pb2.py +0 -49
  152. flyte/_protos/secret/definition_pb2.pyi +0 -93
  153. flyte/_protos/secret/definition_pb2_grpc.py +0 -4
  154. flyte/_protos/secret/payload_pb2.py +0 -62
  155. flyte/_protos/secret/payload_pb2.pyi +0 -94
  156. flyte/_protos/secret/payload_pb2_grpc.py +0 -4
  157. flyte/_protos/secret/secret_pb2.py +0 -38
  158. flyte/_protos/secret/secret_pb2.pyi +0 -6
  159. flyte/_protos/secret/secret_pb2_grpc.py +0 -198
  160. flyte/_protos/secret/secret_pb2_grpc_grpc.py +0 -198
  161. flyte/_protos/validate/validate/validate_pb2.py +0 -76
  162. flyte/_protos/workflow/common_pb2.py +0 -27
  163. flyte/_protos/workflow/common_pb2.pyi +0 -14
  164. flyte/_protos/workflow/common_pb2_grpc.py +0 -4
  165. flyte/_protos/workflow/environment_pb2.py +0 -29
  166. flyte/_protos/workflow/environment_pb2.pyi +0 -12
  167. flyte/_protos/workflow/environment_pb2_grpc.py +0 -4
  168. flyte/_protos/workflow/node_execution_service_pb2.py +0 -26
  169. flyte/_protos/workflow/node_execution_service_pb2.pyi +0 -4
  170. flyte/_protos/workflow/node_execution_service_pb2_grpc.py +0 -32
  171. flyte/_protos/workflow/queue_service_pb2.py +0 -111
  172. flyte/_protos/workflow/queue_service_pb2.pyi +0 -168
  173. flyte/_protos/workflow/queue_service_pb2_grpc.py +0 -172
  174. flyte/_protos/workflow/run_definition_pb2.py +0 -123
  175. flyte/_protos/workflow/run_definition_pb2.pyi +0 -352
  176. flyte/_protos/workflow/run_definition_pb2_grpc.py +0 -4
  177. flyte/_protos/workflow/run_logs_service_pb2.py +0 -41
  178. flyte/_protos/workflow/run_logs_service_pb2.pyi +0 -28
  179. flyte/_protos/workflow/run_logs_service_pb2_grpc.py +0 -69
  180. flyte/_protos/workflow/run_service_pb2.py +0 -137
  181. flyte/_protos/workflow/run_service_pb2.pyi +0 -185
  182. flyte/_protos/workflow/run_service_pb2_grpc.py +0 -446
  183. flyte/_protos/workflow/state_service_pb2.py +0 -67
  184. flyte/_protos/workflow/state_service_pb2.pyi +0 -76
  185. flyte/_protos/workflow/state_service_pb2_grpc.py +0 -138
  186. flyte/_protos/workflow/task_definition_pb2.py +0 -82
  187. flyte/_protos/workflow/task_definition_pb2.pyi +0 -88
  188. flyte/_protos/workflow/task_definition_pb2_grpc.py +0 -4
  189. flyte/_protos/workflow/task_service_pb2.py +0 -60
  190. flyte/_protos/workflow/task_service_pb2.pyi +0 -59
  191. flyte/_protos/workflow/task_service_pb2_grpc.py +0 -138
  192. flyte-2.0.0b22.dist-info/RECORD +0 -250
  193. {flyte-2.0.0b22.data → flyte-2.0.0b30.data}/scripts/debug.py +0 -0
  194. {flyte-2.0.0b22.dist-info → flyte-2.0.0b30.dist-info}/WHEEL +0 -0
  195. {flyte-2.0.0b22.dist-info → flyte-2.0.0b30.dist-info}/entry_points.txt +0 -0
  196. {flyte-2.0.0b22.dist-info → flyte-2.0.0b30.dist-info}/licenses/LICENSE +0 -0
  197. {flyte-2.0.0b22.dist-info → flyte-2.0.0b30.dist-info}/top_level.txt +0 -0
flyte/io/_dir.py CHANGED
@@ -4,12 +4,14 @@ import os
4
4
  from pathlib import Path
5
5
  from typing import AsyncIterator, Dict, Generic, Iterator, List, Optional, Type, TypeVar, Union
6
6
 
7
- from flyteidl.core import literals_pb2, types_pb2
7
+ from flyteidl2.core import literals_pb2, types_pb2
8
8
  from fsspec.asyn import AsyncFileSystem
9
+ from fsspec.utils import get_protocol
9
10
  from mashumaro.types import SerializableType
10
11
  from pydantic import BaseModel, model_validator
11
12
 
12
13
  import flyte.storage as storage
14
+ from flyte._context import internal_ctx
13
15
  from flyte.io._file import File
14
16
  from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
15
17
 
@@ -20,28 +22,176 @@ T = TypeVar("T")
20
22
  class Dir(BaseModel, Generic[T], SerializableType):
21
23
  """
22
24
  A generic directory class representing a directory with files of a specified format.
23
- Provides both async and sync interfaces for directory operations.
24
- Users are responsible for handling all I/O - the type transformer for Dir does not do any automatic uploading
25
- or downloading of files.
25
+ Provides both async and sync interfaces for directory operations. All methods without _sync suffix are async.
26
+
27
+ The class should be instantiated using one of the class methods. The constructor should only be used to
28
+ instantiate references to existing remote directories.
26
29
 
27
30
  The generic type T represents the format of the files in the directory.
28
31
 
29
- Example:
32
+ Important methods:
33
+ - `from_existing_remote`: Create a Dir object referencing an existing remote directory.
34
+ - `from_local` / `from_local_sync`: Upload a local directory to remote storage.
35
+
36
+ **Asynchronous methods**:
37
+ - `walk`: Asynchronously iterate through files in the directory.
38
+ - `list_files`: Asynchronously get a list of all files (non-recursive).
39
+ - `download`: Asynchronously download the entire directory to a local path.
40
+ - `exists`: Asynchronously check if the directory exists.
41
+ - `get_file`: Asynchronously get a specific file from the directory by name.
42
+
43
+ **Synchronous methods** (suffixed with `_sync`):
44
+ - `walk_sync`: Synchronously iterate through files in the directory.
45
+ - `list_files_sync`: Synchronously get a list of all files (non-recursive).
46
+ - `download_sync`: Synchronously download the entire directory to a local path.
47
+ - `exists_sync`: Synchronously check if the directory exists.
48
+ - `get_file_sync`: Synchronously get a specific file from the directory by name.
49
+
50
+ Example: Walk through directory files recursively (Async).
51
+
52
+ ```python
53
+ @env.task
54
+ async def process_all_files(d: Dir) -> int:
55
+ file_count = 0
56
+ async for file in d.walk(recursive=True):
57
+ async with file.open("rb") as f:
58
+ content = await f.read()
59
+ # Process content
60
+ file_count += 1
61
+ return file_count
62
+ ```
63
+
64
+ Example: Walk through directory files recursively (Sync).
65
+
66
+ ```python
67
+ @env.task
68
+ def process_all_files_sync(d: Dir) -> int:
69
+ file_count = 0
70
+ for file in d.walk_sync(recursive=True):
71
+ with file.open_sync("rb") as f:
72
+ content = f.read()
73
+ # Process content
74
+ file_count += 1
75
+ return file_count
76
+ ```
77
+
78
+ Example: List files in directory (Async).
79
+
80
+ ```python
81
+ @env.task
82
+ async def count_files(d: Dir) -> int:
83
+ files = await d.list_files()
84
+ return len(files)
85
+ ```
86
+
87
+ Example: List files in directory (Sync).
88
+
89
+ ```python
90
+ @env.task
91
+ def count_files_sync(d: Dir) -> int:
92
+ files = d.list_files_sync()
93
+ return len(files)
94
+ ```
95
+
96
+ Example: Get a specific file from directory (Async).
97
+
30
98
  ```python
31
- # Async usage
32
- from pandas import DataFrame
33
- data_dir = Dir[DataFrame](path="s3://my-bucket/data/")
34
-
35
- # Walk through files
36
- async for file in data_dir.walk():
37
- async with file.open() as f:
38
- content = await f.read()
39
-
40
- # Sync alternative
41
- for file in data_dir.walk_sync():
42
- with file.open_sync() as f:
43
- content = f.read()
99
+ @env.task
100
+ async def read_config_file(d: Dir) -> str:
101
+ config_file = await d.get_file("config.json")
102
+ if config_file:
103
+ async with config_file.open("rb") as f:
104
+ return (await f.read()).decode("utf-8")
105
+ return "Config not found"
44
106
  ```
107
+
108
+ Example: Get a specific file from directory (Sync).
109
+
110
+ ```python
111
+ @env.task
112
+ def read_config_file_sync(d: Dir) -> str:
113
+ config_file = d.get_file_sync("config.json")
114
+ if config_file:
115
+ with config_file.open_sync("rb") as f:
116
+ return f.read().decode("utf-8")
117
+ return "Config not found"
118
+ ```
119
+
120
+ Example: Upload a local directory to remote storage (Async).
121
+
122
+ ```python
123
+ @env.task
124
+ async def upload_directory() -> Dir:
125
+ # Create local directory with files
126
+ os.makedirs("/tmp/my_data", exist_ok=True)
127
+ with open("/tmp/my_data/file1.txt", "w") as f:
128
+ f.write("data1")
129
+ # Upload to remote storage
130
+ return await Dir.from_local("/tmp/my_data/")
131
+ ```
132
+
133
+ Example: Upload a local directory to remote storage (Sync).
134
+
135
+ ```python
136
+ @env.task
137
+ def upload_directory_sync() -> Dir:
138
+ # Create local directory with files
139
+ os.makedirs("/tmp/my_data", exist_ok=True)
140
+ with open("/tmp/my_data/file1.txt", "w") as f:
141
+ f.write("data1")
142
+ # Upload to remote storage
143
+ return Dir.from_local_sync("/tmp/my_data/")
144
+ ```
145
+
146
+ Example: Download a directory to local storage (Async).
147
+
148
+ ```python
149
+ @env.task
150
+ async def download_directory(d: Dir) -> str:
151
+ local_path = await d.download()
152
+ # Process files in local directory
153
+ return local_path
154
+ ```
155
+
156
+ Example: Download a directory to local storage (Sync).
157
+
158
+ ```python
159
+ @env.task
160
+ def download_directory_sync(d: Dir) -> str:
161
+ local_path = d.download_sync()
162
+ # Process files in local directory
163
+ return local_path
164
+ ```
165
+
166
+ Example: Reference an existing remote directory.
167
+
168
+ ```python
169
+ @env.task
170
+ async def process_existing_dir() -> int:
171
+ d = Dir.from_existing_remote("s3://my-bucket/data/")
172
+ files = await d.list_files()
173
+ return len(files)
174
+ ```
175
+
176
+ Example: Check if directory exists (Async).
177
+
178
+ ```python
179
+ @env.task
180
+ async def check_directory(d: Dir) -> bool:
181
+ return await d.exists()
182
+ ```
183
+
184
+ Example: Check if directory exists (Sync).
185
+
186
+ ```python
187
+ @env.task
188
+ def check_directory_sync(d: Dir) -> bool:
189
+ return d.exists_sync()
190
+ ```
191
+
192
+ Args:
193
+ path: The path to the directory (can be local or remote)
194
+ name: Optional name for the directory (defaults to basename of path)
45
195
  """
46
196
 
47
197
  # Represents either a local or remote path.
@@ -56,20 +206,24 @@ class Dir(BaseModel, Generic[T], SerializableType):
56
206
  @model_validator(mode="before")
57
207
  @classmethod
58
208
  def pre_init(cls, data):
209
+ """Internal: Pydantic validator to set default name from path. Not intended for direct use."""
59
210
  if data.get("name") is None:
60
211
  data["name"] = Path(data["path"]).name
61
212
  return data
62
213
 
63
214
  def _serialize(self) -> Dict[str, Optional[str]]:
215
+ """Internal: Serialize Dir to dictionary. Not intended for direct use."""
64
216
  pyd_dump = self.model_dump()
65
217
  return pyd_dump
66
218
 
67
219
  @classmethod
68
220
  def _deserialize(cls, file_dump: Dict[str, Optional[str]]) -> Dir:
221
+ """Internal: Deserialize Dir from dictionary. Not intended for direct use."""
69
222
  return cls.model_validate(file_dump)
70
223
 
71
224
  @classmethod
72
225
  def schema_match(cls, incoming: dict):
226
+ """Internal: Check if incoming schema matches Dir schema. Not intended for direct use."""
73
227
  this_schema = cls.model_json_schema()
74
228
  current_required = this_schema.get("required")
75
229
  incoming_required = incoming.get("required")
@@ -86,19 +240,48 @@ class Dir(BaseModel, Generic[T], SerializableType):
86
240
  """
87
241
  Asynchronously walk through the directory and yield File objects.
88
242
 
89
- Args:
90
- recursive: If True, recursively walk subdirectories
91
- max_depth: Maximum depth for recursive walking
243
+ Use this to iterate through all files in a directory. Each yielded File can be read directly without
244
+ downloading.
92
245
 
93
- Yields:
94
- File objects for each file found in the directory
246
+ Example (Async - Recursive):
247
+
248
+ ```python
249
+ @env.task
250
+ async def list_all_files(d: Dir) -> list[str]:
251
+ file_names = []
252
+ async for file in d.walk(recursive=True):
253
+ file_names.append(file.name)
254
+ return file_names
255
+ ```
256
+
257
+ Example (Async - Non-recursive):
95
258
 
96
- Example:
97
259
  ```python
98
- async for file in directory.walk():
99
- local_path = await file.download()
100
- # Process the file
260
+ @env.task
261
+ async def list_top_level_files(d: Dir) -> list[str]:
262
+ file_names = []
263
+ async for file in d.walk(recursive=False):
264
+ file_names.append(file.name)
265
+ return file_names
101
266
  ```
267
+
268
+ Example (Async - With max depth):
269
+
270
+ ```python
271
+ @env.task
272
+ async def list_files_max_depth(d: Dir) -> list[str]:
273
+ file_names = []
274
+ async for file in d.walk(recursive=True, max_depth=2):
275
+ file_names.append(file.name)
276
+ return file_names
277
+ ```
278
+
279
+ Args:
280
+ recursive: If True, recursively walk subdirectories. If False, only list files in the top-level directory.
281
+ max_depth: Maximum depth for recursive walking. If None, walk through all subdirectories.
282
+
283
+ Yields:
284
+ File objects for each file found in the directory
102
285
  """
103
286
  fs = storage.get_underlying_filesystem(path=self.path)
104
287
  if recursive is False:
@@ -125,20 +308,48 @@ class Dir(BaseModel, Generic[T], SerializableType):
125
308
  """
126
309
  Synchronously walk through the directory and yield File objects.
127
310
 
128
- Args:
129
- recursive: If True, recursively walk subdirectories
130
- file_pattern: Glob pattern to filter files
131
- max_depth: Maximum depth for recursive walking
311
+ Use this in non-async tasks to iterate through all files in a directory.
132
312
 
133
- Yields:
134
- File objects for each file found in the directory
313
+ Example (Sync - Recursive):
314
+
315
+ ```python
316
+ @env.task
317
+ def list_all_files_sync(d: Dir) -> list[str]:
318
+ file_names = []
319
+ for file in d.walk_sync(recursive=True):
320
+ file_names.append(file.name)
321
+ return file_names
322
+ ```
323
+
324
+ Example (Sync - With file pattern):
325
+
326
+ ```python
327
+ @env.task
328
+ def list_text_files(d: Dir) -> list[str]:
329
+ file_names = []
330
+ for file in d.walk_sync(recursive=True, file_pattern="*.txt"):
331
+ file_names.append(file.name)
332
+ return file_names
333
+ ```
334
+
335
+ Example (Sync - Non-recursive with max depth):
135
336
 
136
- Example:
137
337
  ```python
138
- for file in directory.walk_sync():
139
- local_path = file.download_sync()
140
- # Process the file
338
+ @env.task
339
+ def list_files_limited(d: Dir) -> list[str]:
340
+ file_names = []
341
+ for file in d.walk_sync(recursive=True, max_depth=2):
342
+ file_names.append(file.name)
343
+ return file_names
141
344
  ```
345
+
346
+ Args:
347
+ recursive: If True, recursively walk subdirectories. If False, only list files in the top-level directory.
348
+ file_pattern: Glob pattern to filter files (e.g., "*.txt", "*.csv"). Default is "*" (all files).
349
+ max_depth: Maximum depth for recursive walking. If None, walk through all subdirectories.
350
+
351
+ Yields:
352
+ File objects for each file found in the directory
142
353
  """
143
354
  fs = storage.get_underlying_filesystem(path=self.path)
144
355
  for parent, _, files in fs.walk(self.path, maxdepth=max_depth):
@@ -153,14 +364,32 @@ class Dir(BaseModel, Generic[T], SerializableType):
153
364
  """
154
365
  Asynchronously get a list of all files in the directory (non-recursive).
155
366
 
367
+ Use this when you need a list of all files in the top-level directory at once.
368
+
156
369
  Returns:
157
- A list of File objects
370
+ A list of File objects for files in the top-level directory
371
+
372
+ Example (Async):
158
373
 
159
- Example:
160
374
  ```python
161
- files = await directory.list_files()
162
- for file in files:
163
- # Process the file
375
+ @env.task
376
+ async def count_files(d: Dir) -> int:
377
+ files = await d.list_files()
378
+ return len(files)
379
+ ```
380
+
381
+ Example (Async - Process files):
382
+
383
+ ```python
384
+ @env.task
385
+ async def process_all_files(d: Dir) -> list[str]:
386
+ files = await d.list_files()
387
+ contents = []
388
+ for file in files:
389
+ async with file.open("rb") as f:
390
+ content = await f.read()
391
+ contents.append(content.decode("utf-8"))
392
+ return contents
164
393
  ```
165
394
  """
166
395
  # todo: this should probably also just defer to fsspec.find()
@@ -173,14 +402,32 @@ class Dir(BaseModel, Generic[T], SerializableType):
173
402
  """
174
403
  Synchronously get a list of all files in the directory (non-recursive).
175
404
 
405
+ Use this in non-async tasks when you need a list of all files in the top-level directory at once.
406
+
176
407
  Returns:
177
- A list of File objects
408
+ A list of File objects for files in the top-level directory
409
+
410
+ Example (Sync):
178
411
 
179
- Example:
180
412
  ```python
181
- files = directory.list_files_sync()
182
- for file in files:
183
- # Process the file
413
+ @env.task
414
+ def count_files_sync(d: Dir) -> int:
415
+ files = d.list_files_sync()
416
+ return len(files)
417
+ ```
418
+
419
+ Example (Sync - Process files):
420
+
421
+ ```python
422
+ @env.task
423
+ def process_all_files_sync(d: Dir) -> list[str]:
424
+ files = d.list_files_sync()
425
+ contents = []
426
+ for file in files:
427
+ with file.open_sync("rb") as f:
428
+ content = f.read()
429
+ contents.append(content.decode("utf-8"))
430
+ return contents
184
431
  ```
185
432
  """
186
433
  return list(self.walk_sync(recursive=False))
@@ -189,19 +436,43 @@ class Dir(BaseModel, Generic[T], SerializableType):
189
436
  """
190
437
  Asynchronously download the entire directory to a local path.
191
438
 
192
- Args:
193
- local_path: The local path to download the directory to. If None, a temporary
194
- directory will be used.
439
+ Use this when you need to download all files in a directory to your local filesystem for processing.
195
440
 
196
- Returns:
197
- The path to the downloaded directory
441
+ Example (Async):
442
+
443
+ ```python
444
+ @env.task
445
+ async def download_directory(d: Dir) -> str:
446
+ local_dir = await d.download()
447
+ # Process files in the local directory
448
+ return local_dir
449
+ ```
450
+
451
+ Example (Async - Download to specific path):
198
452
 
199
- Example:
200
453
  ```python
201
- local_dir = await directory.download('/tmp/my_data/')
454
+ @env.task
455
+ async def download_to_path(d: Dir) -> str:
456
+ local_dir = await d.download("/tmp/my_data/")
457
+ return local_dir
202
458
  ```
459
+
460
+ Args:
461
+ local_path: The local path to download the directory to. If None, a temporary
462
+ directory will be used and a path will be generated.
463
+
464
+ Returns:
465
+ The absolute path to the downloaded directory
203
466
  """
204
- local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
467
+ # If no local_path specified, create a unique path + append source directory name
468
+ if local_path is None:
469
+ unique_path = storage.get_random_local_path()
470
+ source_dirname = Path(self.path).name # will need to be updated for windows
471
+ local_dest = str(Path(unique_path) / source_dirname)
472
+ else:
473
+ # If local_path is specified, use it directly (contents go into it)
474
+ local_dest = str(local_path)
475
+
205
476
  if not storage.is_remote(self.path):
206
477
  if not local_path or local_path == self.path:
207
478
  # Skip copying
@@ -216,25 +487,49 @@ class Dir(BaseModel, Generic[T], SerializableType):
216
487
  await loop.run_in_executor(None, lambda: shutil.copytree(self.path, local_dest, dirs_exist_ok=True))
217
488
 
218
489
  await copy_tree()
490
+ return local_dest
219
491
  return await storage.get(self.path, local_dest, recursive=True)
220
492
 
221
493
  def download_sync(self, local_path: Optional[Union[str, Path]] = None) -> str:
222
494
  """
223
495
  Synchronously download the entire directory to a local path.
224
496
 
225
- Args:
226
- local_path: The local path to download the directory to. If None, a temporary
227
- directory will be used.
497
+ Use this in non-async tasks when you need to download all files in a directory to your local filesystem.
228
498
 
229
- Returns:
230
- The path to the downloaded directory
499
+ Example (Sync):
500
+
501
+ ```python
502
+ @env.task
503
+ def download_directory_sync(d: Dir) -> str:
504
+ local_dir = d.download_sync()
505
+ # Process files in the local directory
506
+ return local_dir
507
+ ```
508
+
509
+ Example (Sync - Download to specific path):
231
510
 
232
- Example:
233
511
  ```python
234
- local_dir = directory.download_sync('/tmp/my_data/')
512
+ @env.task
513
+ def download_to_path_sync(d: Dir) -> str:
514
+ local_dir = d.download_sync("/tmp/my_data/")
515
+ return local_dir
235
516
  ```
517
+ Args:
518
+ local_path: The local path to download the directory to. If None, a temporary
519
+ directory will be used and a path will be generated.
520
+
521
+ Returns:
522
+ The absolute path to the downloaded directory
236
523
  """
237
- local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
524
+ # If no local_path specified, create a unique path + append source directory name
525
+ if local_path is None:
526
+ unique_path = storage.get_random_local_path()
527
+ source_dirname = Path(self.path).name
528
+ local_dest = str(Path(unique_path) / source_dirname)
529
+ else:
530
+ # If local_path is specified, use it directly (contents go into it)
531
+ local_dest = str(local_path)
532
+
238
533
  if not storage.is_remote(self.path):
239
534
  if not local_path or local_path == self.path:
240
535
  # Skip copying
@@ -244,80 +539,188 @@ class Dir(BaseModel, Generic[T], SerializableType):
244
539
  import shutil
245
540
 
246
541
  shutil.copytree(self.path, local_dest, dirs_exist_ok=True)
542
+ return local_dest
247
543
 
248
- # Figure this out when we figure out the final sync story
249
- raise NotImplementedError("Sync download is not implemented for remote paths")
544
+ fs = storage.get_underlying_filesystem(path=self.path)
545
+ fs.get(self.path, local_dest, recursive=True)
546
+ return local_dest
250
547
 
251
548
  @classmethod
252
549
  async def from_local(
253
550
  cls,
254
551
  local_path: Union[str, Path],
255
- remote_path: Optional[str] = None,
552
+ remote_destination: Optional[str] = None,
256
553
  dir_cache_key: Optional[str] = None,
257
554
  ) -> Dir[T]:
258
555
  """
259
- Asynchronously create a new Dir by uploading a local directory to the configured remote store.
556
+ Asynchronously create a new Dir by uploading a local directory to remote storage.
557
+
558
+ Use this in async tasks when you have a local directory that needs to be uploaded to remote storage.
260
559
 
560
+ Example (Async):
561
+
562
+ ```python
563
+ @env.task
564
+ async def upload_local_directory() -> Dir:
565
+ # Create a local directory with files
566
+ os.makedirs("/tmp/data_dir", exist_ok=True)
567
+ with open("/tmp/data_dir/file1.txt", "w") as f:
568
+ f.write("data1")
569
+
570
+ # Upload to remote storage
571
+ remote_dir = await Dir.from_local("/tmp/data_dir/")
572
+ return remote_dir
573
+ ```
574
+
575
+ Example (Async - With specific destination):
576
+
577
+ ```python
578
+ @env.task
579
+ async def upload_to_specific_path() -> Dir:
580
+ remote_dir = await Dir.from_local("/tmp/data_dir/", "s3://my-bucket/data/")
581
+ return remote_dir
582
+ ```
583
+
584
+ Example (Async - With cache key):
585
+
586
+ ```python
587
+ @env.task
588
+ async def upload_with_cache_key() -> Dir:
589
+ remote_dir = await Dir.from_local("/tmp/data_dir/", dir_cache_key="my_cache_key_123")
590
+ return remote_dir
591
+ ```
261
592
  Args:
262
593
  local_path: Path to the local directory
263
- remote_path: Optional path to store the directory remotely. If None, a path will be generated.
264
- dir_cache_key: If you have a precomputed hash value you want to use when computing cache keys for
265
- discoverable tasks that this File is an input to.
594
+ remote_destination: Optional remote path to store the directory. If None, a path will be automatically
595
+ generated.
596
+ dir_cache_key: Optional precomputed hash value to use for cache key computation when this Dir is used
597
+ as an input to discoverable tasks. If not specified, the cache key will be based on
598
+ directory attributes.
266
599
 
267
600
  Returns:
268
601
  A new Dir instance pointing to the uploaded directory
269
-
270
- Example:
271
- ```python
272
- remote_dir = await Dir[DataFrame].from_local('/tmp/data_dir/', 's3://bucket/data/')
273
- # With a known hash value you want to use for cache key calculation
274
- remote_dir = await Dir[DataFrame].from_local('/tmp/data_dir/', 's3://bucket/data/', dir_cache_key='abc123')
275
- ```
276
602
  """
277
603
  local_path_str = str(local_path)
278
604
  dirname = os.path.basename(os.path.normpath(local_path_str))
605
+ resolved_remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path(dirname)
606
+ protocol = get_protocol(resolved_remote_path)
607
+
608
+ # Shortcut for local, don't copy and just return
609
+ if "file" in protocol and remote_destination is None:
610
+ output_path = str(Path(local_path).absolute())
611
+ return cls(path=output_path, name=dirname, hash=dir_cache_key)
279
612
 
280
- output_path = await storage.put(from_path=local_path_str, to_path=remote_path, recursive=True)
613
+ # todo: in the future, mirror File and set the file to_path here
614
+ output_path = await storage.put(from_path=local_path_str, to_path=remote_destination, recursive=True)
281
615
  return cls(path=output_path, name=dirname, hash=dir_cache_key)
282
616
 
283
617
  @classmethod
284
- def from_existing_remote(cls, remote_path: str, dir_cache_key: Optional[str] = None) -> Dir[T]:
618
+ def from_local_sync(
619
+ cls,
620
+ local_path: Union[str, Path],
621
+ remote_destination: Optional[str] = None,
622
+ dir_cache_key: Optional[str] = None,
623
+ ) -> Dir[T]:
285
624
  """
286
- Create a Dir reference from an existing remote directory.
625
+ Synchronously create a new Dir by uploading a local directory to remote storage.
287
626
 
288
- Args:
289
- remote_path: The remote path to the existing directory
290
- dir_cache_key: Optional hash value to use for cache key computation. If not specified,
291
- the cache key will be computed based on this object's attributes.
627
+ Use this in non-async tasks when you have a local directory that needs to be uploaded to remote storage.
628
+
629
+ Example (Sync):
292
630
 
293
- Example:
294
631
  ```python
295
- remote_dir = Dir.from_existing_remote("s3://bucket/data/")
296
- # With a known hash
297
- remote_dir = Dir.from_existing_remote("s3://bucket/data/", dir_cache_key="abc123")
632
+ @env.task
633
+ def upload_local_directory_sync() -> Dir:
634
+ # Create a local directory with files
635
+ os.makedirs("/tmp/data_dir", exist_ok=True)
636
+ with open("/tmp/data_dir/file1.txt", "w") as f:
637
+ f.write("data1")
638
+
639
+ # Upload to remote storage
640
+ remote_dir = Dir.from_local_sync("/tmp/data_dir/")
641
+ return remote_dir
298
642
  ```
299
- """
300
- return cls(path=remote_path, hash=dir_cache_key)
301
643
 
302
- @classmethod
303
- def from_local_sync(cls, local_path: Union[str, Path], remote_path: Optional[str] = None) -> Dir[T]:
304
- """
305
- Synchronously create a new Dir by uploading a local directory to the configured remote store.
644
+ Example (Sync - With specific destination):
645
+
646
+ ```python
647
+ @env.task
648
+ def upload_to_specific_path_sync() -> Dir:
649
+ remote_dir = Dir.from_local_sync("/tmp/data_dir/", "s3://my-bucket/data/")
650
+ return remote_dir
651
+ ```
652
+
653
+ Example (Sync - With cache key):
654
+
655
+ ```python
656
+ @env.task
657
+ def upload_with_cache_key_sync() -> Dir:
658
+ remote_dir = Dir.from_local_sync("/tmp/data_dir/", dir_cache_key="my_cache_key_123")
659
+ return remote_dir
660
+ ```
306
661
 
307
662
  Args:
308
663
  local_path: Path to the local directory
309
- remote_path: Optional path to store the directory remotely. If None, a path will be generated.
664
+ remote_destination: Optional remote path to store the directory. If None, a path will be automatically
665
+ generated.
666
+ dir_cache_key: Optional precomputed hash value to use for cache key computation when this Dir is used
667
+ as an input to discoverable tasks. If not specified, the cache key will be based on
668
+ directory attributes.
310
669
 
311
670
  Returns:
312
671
  A new Dir instance pointing to the uploaded directory
672
+ """
673
+ local_path_str = str(local_path)
674
+ dirname = os.path.basename(os.path.normpath(local_path_str))
675
+
676
+ resolved_remote_path = remote_destination or internal_ctx().raw_data.get_random_remote_path(dirname)
677
+ protocol = get_protocol(resolved_remote_path)
678
+
679
+ # Shortcut for local, don't copy and just return
680
+ if "file" in protocol and remote_destination is None:
681
+ output_path = str(Path(local_path).absolute())
682
+ return cls(path=output_path, name=dirname, hash=dir_cache_key)
683
+
684
+ fs = storage.get_underlying_filesystem(path=resolved_remote_path)
685
+ fs.put(local_path_str, resolved_remote_path, recursive=True)
686
+ return cls(path=resolved_remote_path, name=dirname, hash=dir_cache_key)
687
+
688
+ @classmethod
689
+ def from_existing_remote(cls, remote_path: str, dir_cache_key: Optional[str] = None) -> Dir[T]:
690
+ """
691
+ Create a Dir reference from an existing remote directory.
692
+
693
+ Use this when you want to reference a directory that already exists in remote storage without uploading it.
313
694
 
314
695
  Example:
696
+
315
697
  ```python
316
- remote_dir = Dir[DataFrame].from_local_sync('/tmp/data_dir/', 's3://bucket/data/')
698
+ @env.task
699
+ async def process_existing_directory() -> int:
700
+ d = Dir.from_existing_remote("s3://my-bucket/data/")
701
+ files = await d.list_files()
702
+ return len(files)
317
703
  ```
704
+
705
+ Example (With cache key):
706
+
707
+ ```python
708
+ @env.task
709
+ async def process_with_cache_key() -> int:
710
+ d = Dir.from_existing_remote("s3://my-bucket/data/", dir_cache_key="abc123")
711
+ files = await d.list_files()
712
+ return len(files)
713
+ ```
714
+
715
+ Args:
716
+ remote_path: The remote path to the existing directory
717
+ dir_cache_key: Optional hash value to use for cache key computation. If not specified,
718
+ the cache key will be computed based on the directory's attributes.
719
+
720
+ Returns:
721
+ A new Dir instance pointing to the existing remote directory
318
722
  """
319
- # Implement this after we figure out the final sync story
320
- raise NotImplementedError("Sync upload is not implemented for remote paths")
723
+ return cls(path=remote_path, hash=dir_cache_key)
321
724
 
322
725
  async def exists(self) -> bool:
323
726
  """
@@ -326,10 +729,15 @@ class Dir(BaseModel, Generic[T], SerializableType):
326
729
  Returns:
327
730
  True if the directory exists, False otherwise
328
731
 
329
- Example:
732
+ Example (Async):
733
+
330
734
  ```python
331
- if await directory.exists():
332
- # Process the directory
735
+ @env.task
736
+ async def check_directory(d: Dir) -> bool:
737
+ if await d.exists():
738
+ print("Directory exists!")
739
+ return True
740
+ return False
333
741
  ```
334
742
  """
335
743
  fs = storage.get_underlying_filesystem(path=self.path)
@@ -342,13 +750,20 @@ class Dir(BaseModel, Generic[T], SerializableType):
342
750
  """
343
751
  Synchronously check if the directory exists.
344
752
 
753
+ Use this in non-async tasks or when you need synchronous directory existence checking.
754
+
345
755
  Returns:
346
756
  True if the directory exists, False otherwise
347
757
 
348
- Example:
758
+ Example (Sync):
759
+
349
760
  ```python
350
- if directory.exists_sync():
351
- # Process the directory
761
+ @env.task
762
+ def check_directory_sync(d: Dir) -> bool:
763
+ if d.exists_sync():
764
+ print("Directory exists!")
765
+ return True
766
+ return False
352
767
  ```
353
768
  """
354
769
  fs = storage.get_underlying_filesystem(path=self.path)
@@ -356,20 +771,28 @@ class Dir(BaseModel, Generic[T], SerializableType):
356
771
 
357
772
  async def get_file(self, file_name: str) -> Optional[File[T]]:
358
773
  """
359
- Asynchronously get a specific file from the directory.
774
+ Asynchronously get a specific file from the directory by name.
775
+
776
+ Use this when you know the name of a specific file in the directory you want to access.
777
+
778
+ Example (Async):
779
+
780
+ ```python
781
+ @env.task
782
+ async def read_specific_file(d: Dir) -> str:
783
+ file = await d.get_file("data.csv")
784
+ if file:
785
+ async with file.open("rb") as f:
786
+ content = await f.read()
787
+ return content.decode("utf-8")
788
+ return "File not found"
789
+ ```
360
790
 
361
791
  Args:
362
792
  file_name: The name of the file to get
363
793
 
364
794
  Returns:
365
795
  A File instance if the file exists, None otherwise
366
-
367
- Example:
368
- ```python
369
- file = await directory.get_file("data.csv")
370
- if file:
371
- # Process the file
372
- ```
373
796
  """
374
797
  fs = storage.get_underlying_filesystem(path=self.path)
375
798
  file_path = fs.sep.join([self.path, file_name])
@@ -381,20 +804,28 @@ class Dir(BaseModel, Generic[T], SerializableType):
381
804
 
382
805
  def get_file_sync(self, file_name: str) -> Optional[File[T]]:
383
806
  """
384
- Synchronously get a specific file from the directory.
807
+ Synchronously get a specific file from the directory by name.
808
+
809
+ Use this in non-async tasks when you know the name of a specific file in the directory you want to access.
810
+
811
+ Example (Sync):
812
+
813
+ ```python
814
+ @env.task
815
+ def read_specific_file_sync(d: Dir) -> str:
816
+ file = d.get_file_sync("data.csv")
817
+ if file:
818
+ with file.open_sync("rb") as f:
819
+ content = f.read()
820
+ return content.decode("utf-8")
821
+ return "File not found"
822
+ ```
385
823
 
386
824
  Args:
387
825
  file_name: The name of the file to get
388
826
 
389
827
  Returns:
390
828
  A File instance if the file exists, None otherwise
391
-
392
- Example:
393
- ```python
394
- file = directory.get_file_sync("data.csv")
395
- if file:
396
- # Process the file
397
- ```
398
829
  """
399
830
  file_path = os.path.join(self.path, file_name)
400
831
  file = File[T](path=file_path)