PyPI - flyte - Versions diffs - 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl - Mend

flyte 0.1.0py3-none-any.whl → 0.2.0a0py3-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 (219) hide show

flyte/__init__.py +78 -2
flyte/_bin/__init__.py +0 -0
flyte/_bin/runtime.py +152 -0
flyte/_build.py +26 -0
flyte/_cache/__init__.py +12 -0
flyte/_cache/cache.py +145 -0
flyte/_cache/defaults.py +9 -0
flyte/_cache/policy_function_body.py +42 -0
flyte/_code_bundle/__init__.py +8 -0
flyte/_code_bundle/_ignore.py +113 -0
flyte/_code_bundle/_packaging.py +187 -0
flyte/_code_bundle/_utils.py +323 -0
flyte/_code_bundle/bundle.py +209 -0
flyte/_context.py +152 -0
flyte/_deploy.py +243 -0
flyte/_doc.py +29 -0
flyte/_docstring.py +32 -0
flyte/_environment.py +84 -0
flyte/_excepthook.py +37 -0
flyte/_group.py +32 -0
flyte/_hash.py +23 -0
flyte/_image.py +762 -0
flyte/_initialize.py +492 -0
flyte/_interface.py +84 -0
flyte/_internal/__init__.py +3 -0
flyte/_internal/controllers/__init__.py +128 -0
flyte/_internal/controllers/_local_controller.py +193 -0
flyte/_internal/controllers/_trace.py +41 -0
flyte/_internal/controllers/remote/__init__.py +60 -0
flyte/_internal/controllers/remote/_action.py +146 -0
flyte/_internal/controllers/remote/_client.py +47 -0
flyte/_internal/controllers/remote/_controller.py +494 -0
flyte/_internal/controllers/remote/_core.py +410 -0
flyte/_internal/controllers/remote/_informer.py +361 -0
flyte/_internal/controllers/remote/_service_protocol.py +50 -0
flyte/_internal/imagebuild/__init__.py +11 -0
flyte/_internal/imagebuild/docker_builder.py +427 -0
flyte/_internal/imagebuild/image_builder.py +246 -0
flyte/_internal/imagebuild/remote_builder.py +0 -0
flyte/_internal/resolvers/__init__.py +0 -0
flyte/_internal/resolvers/_task_module.py +54 -0
flyte/_internal/resolvers/common.py +31 -0
flyte/_internal/resolvers/default.py +28 -0
flyte/_internal/runtime/__init__.py +0 -0
flyte/_internal/runtime/convert.py +342 -0
flyte/_internal/runtime/entrypoints.py +135 -0
flyte/_internal/runtime/io.py +136 -0
flyte/_internal/runtime/resources_serde.py +138 -0
flyte/_internal/runtime/task_serde.py +330 -0
flyte/_internal/runtime/taskrunner.py +191 -0
flyte/_internal/runtime/types_serde.py +54 -0
flyte/_logging.py +135 -0
flyte/_map.py +215 -0
flyte/_pod.py +19 -0
flyte/_protos/__init__.py +0 -0
flyte/_protos/common/authorization_pb2.py +66 -0
flyte/_protos/common/authorization_pb2.pyi +108 -0
flyte/_protos/common/authorization_pb2_grpc.py +4 -0
flyte/_protos/common/identifier_pb2.py +71 -0
flyte/_protos/common/identifier_pb2.pyi +82 -0
flyte/_protos/common/identifier_pb2_grpc.py +4 -0
flyte/_protos/common/identity_pb2.py +48 -0
flyte/_protos/common/identity_pb2.pyi +72 -0
flyte/_protos/common/identity_pb2_grpc.py +4 -0
flyte/_protos/common/list_pb2.py +36 -0
flyte/_protos/common/list_pb2.pyi +71 -0
flyte/_protos/common/list_pb2_grpc.py +4 -0
flyte/_protos/common/policy_pb2.py +37 -0
flyte/_protos/common/policy_pb2.pyi +27 -0
flyte/_protos/common/policy_pb2_grpc.py +4 -0
flyte/_protos/common/role_pb2.py +37 -0
flyte/_protos/common/role_pb2.pyi +53 -0
flyte/_protos/common/role_pb2_grpc.py +4 -0
flyte/_protos/common/runtime_version_pb2.py +28 -0
flyte/_protos/common/runtime_version_pb2.pyi +24 -0
flyte/_protos/common/runtime_version_pb2_grpc.py +4 -0
flyte/_protos/logs/dataplane/payload_pb2.py +100 -0
flyte/_protos/logs/dataplane/payload_pb2.pyi +177 -0
flyte/_protos/logs/dataplane/payload_pb2_grpc.py +4 -0
flyte/_protos/secret/definition_pb2.py +49 -0
flyte/_protos/secret/definition_pb2.pyi +93 -0
flyte/_protos/secret/definition_pb2_grpc.py +4 -0
flyte/_protos/secret/payload_pb2.py +62 -0
flyte/_protos/secret/payload_pb2.pyi +94 -0
flyte/_protos/secret/payload_pb2_grpc.py +4 -0
flyte/_protos/secret/secret_pb2.py +38 -0
flyte/_protos/secret/secret_pb2.pyi +6 -0
flyte/_protos/secret/secret_pb2_grpc.py +198 -0
flyte/_protos/secret/secret_pb2_grpc_grpc.py +198 -0
flyte/_protos/validate/validate/validate_pb2.py +76 -0
flyte/_protos/workflow/common_pb2.py +27 -0
flyte/_protos/workflow/common_pb2.pyi +14 -0
flyte/_protos/workflow/common_pb2_grpc.py +4 -0
flyte/_protos/workflow/environment_pb2.py +29 -0
flyte/_protos/workflow/environment_pb2.pyi +12 -0
flyte/_protos/workflow/environment_pb2_grpc.py +4 -0
flyte/_protos/workflow/node_execution_service_pb2.py +26 -0
flyte/_protos/workflow/node_execution_service_pb2.pyi +4 -0
flyte/_protos/workflow/node_execution_service_pb2_grpc.py +32 -0
flyte/_protos/workflow/queue_service_pb2.py +105 -0
flyte/_protos/workflow/queue_service_pb2.pyi +146 -0
flyte/_protos/workflow/queue_service_pb2_grpc.py +172 -0
flyte/_protos/workflow/run_definition_pb2.py +128 -0
flyte/_protos/workflow/run_definition_pb2.pyi +314 -0
flyte/_protos/workflow/run_definition_pb2_grpc.py +4 -0
flyte/_protos/workflow/run_logs_service_pb2.py +41 -0
flyte/_protos/workflow/run_logs_service_pb2.pyi +28 -0
flyte/_protos/workflow/run_logs_service_pb2_grpc.py +69 -0
flyte/_protos/workflow/run_service_pb2.py +129 -0
flyte/_protos/workflow/run_service_pb2.pyi +171 -0
flyte/_protos/workflow/run_service_pb2_grpc.py +412 -0
flyte/_protos/workflow/state_service_pb2.py +66 -0
flyte/_protos/workflow/state_service_pb2.pyi +75 -0
flyte/_protos/workflow/state_service_pb2_grpc.py +138 -0
flyte/_protos/workflow/task_definition_pb2.py +79 -0
flyte/_protos/workflow/task_definition_pb2.pyi +81 -0
flyte/_protos/workflow/task_definition_pb2_grpc.py +4 -0
flyte/_protos/workflow/task_service_pb2.py +60 -0
flyte/_protos/workflow/task_service_pb2.pyi +59 -0
flyte/_protos/workflow/task_service_pb2_grpc.py +138 -0
flyte/_resources.py +226 -0
flyte/_retry.py +32 -0
flyte/_reusable_environment.py +25 -0
flyte/_run.py +482 -0
flyte/_secret.py +61 -0
flyte/_task.py +449 -0
flyte/_task_environment.py +183 -0
flyte/_timeout.py +47 -0
flyte/_tools.py +27 -0
flyte/_trace.py +120 -0
flyte/_utils/__init__.py +26 -0
flyte/_utils/asyn.py +119 -0
flyte/_utils/async_cache.py +139 -0
flyte/_utils/coro_management.py +23 -0
flyte/_utils/file_handling.py +72 -0
flyte/_utils/helpers.py +134 -0
flyte/_utils/lazy_module.py +54 -0
flyte/_utils/org_discovery.py +57 -0
flyte/_utils/uv_script_parser.py +49 -0
flyte/_version.py +21 -0
flyte/cli/__init__.py +3 -0
flyte/cli/_abort.py +28 -0
flyte/cli/_common.py +337 -0
flyte/cli/_create.py +145 -0
flyte/cli/_delete.py +23 -0
flyte/cli/_deploy.py +152 -0
flyte/cli/_gen.py +163 -0
flyte/cli/_get.py +310 -0
flyte/cli/_params.py +538 -0
flyte/cli/_run.py +231 -0
flyte/cli/main.py +166 -0
flyte/config/__init__.py +3 -0
flyte/config/_config.py +216 -0
flyte/config/_internal.py +64 -0
flyte/config/_reader.py +207 -0
flyte/connectors/__init__.py +0 -0
flyte/errors.py +172 -0
flyte/extras/__init__.py +5 -0
flyte/extras/_container.py +263 -0
flyte/io/__init__.py +27 -0
flyte/io/_dir.py +448 -0
flyte/io/_file.py +467 -0
flyte/io/_structured_dataset/__init__.py +129 -0
flyte/io/_structured_dataset/basic_dfs.py +219 -0
flyte/io/_structured_dataset/structured_dataset.py +1061 -0
flyte/models.py +391 -0
flyte/remote/__init__.py +26 -0
flyte/remote/_client/__init__.py +0 -0
flyte/remote/_client/_protocols.py +133 -0
flyte/remote/_client/auth/__init__.py +12 -0
flyte/remote/_client/auth/_auth_utils.py +14 -0
flyte/remote/_client/auth/_authenticators/__init__.py +0 -0
flyte/remote/_client/auth/_authenticators/base.py +397 -0
flyte/remote/_client/auth/_authenticators/client_credentials.py +73 -0
flyte/remote/_client/auth/_authenticators/device_code.py +118 -0
flyte/remote/_client/auth/_authenticators/external_command.py +79 -0
flyte/remote/_client/auth/_authenticators/factory.py +200 -0
flyte/remote/_client/auth/_authenticators/pkce.py +516 -0
flyte/remote/_client/auth/_channel.py +215 -0
flyte/remote/_client/auth/_client_config.py +83 -0
flyte/remote/_client/auth/_default_html.py +32 -0
flyte/remote/_client/auth/_grpc_utils/__init__.py +0 -0
flyte/remote/_client/auth/_grpc_utils/auth_interceptor.py +288 -0
flyte/remote/_client/auth/_grpc_utils/default_metadata_interceptor.py +151 -0
flyte/remote/_client/auth/_keyring.py +143 -0
flyte/remote/_client/auth/_token_client.py +260 -0
flyte/remote/_client/auth/errors.py +16 -0
flyte/remote/_client/controlplane.py +95 -0
flyte/remote/_console.py +18 -0
flyte/remote/_data.py +159 -0
flyte/remote/_logs.py +176 -0
flyte/remote/_project.py +85 -0
flyte/remote/_run.py +970 -0
flyte/remote/_secret.py +132 -0
flyte/remote/_task.py +391 -0
flyte/report/__init__.py +3 -0
flyte/report/_report.py +178 -0
flyte/report/_template.html +124 -0
flyte/storage/__init__.py +29 -0
flyte/storage/_config.py +233 -0
flyte/storage/_remote_fs.py +34 -0
flyte/storage/_storage.py +271 -0
flyte/storage/_utils.py +5 -0
flyte/syncify/__init__.py +56 -0
flyte/syncify/_api.py +371 -0
flyte/types/__init__.py +36 -0
flyte/types/_interface.py +40 -0
flyte/types/_pickle.py +118 -0
flyte/types/_renderer.py +162 -0
flyte/types/_string_literals.py +120 -0
flyte/types/_type_engine.py +2287 -0
flyte/types/_utils.py +80 -0
flyte-0.2.0a0.dist-info/METADATA +249 -0
flyte-0.2.0a0.dist-info/RECORD +218 -0
{flyte-0.1.0.dist-info → flyte-0.2.0a0.dist-info}/WHEEL +2 -1
flyte-0.2.0a0.dist-info/entry_points.txt +3 -0
flyte-0.2.0a0.dist-info/top_level.txt +1 -0
flyte-0.1.0.dist-info/METADATA +0 -6
flyte-0.1.0.dist-info/RECORD +0 -5

flyte/report/_template.html ADDED Viewed

@@ -0,0 +1,124 @@
+<!doctype html>
+<html lang="">
+<head>
+    <meta charset="utf-8">
+    <title>User Content</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <link href="https://fonts.googleapis.com/css?family=Lato:300,400,700%7COpen+Sans:400,700" rel="stylesheet">
+    <style>
+        ol, ul {
+            list-style: none;
+        }
+        table {
+            border-collapse: collapse;
+            border-spacing: 0;
+        }
+        #flyte-frame-nav {
+            display: flex;
+            width: 100%;
+        }
+        #flyte-frame-tabs {
+            display: flex;
+            width: 100%;
+            justify-content: center;
+            margin-block: 0;
+            padding-inline-start: 0;
+        }
+        #flyte-frame-tabs li {
+            cursor: pointer;
+            padding: 8px;
+            margin: 0;
+            margin-right: 12px;
+            font-size: 14px;
+            line-height: 20px;
+            font-weight: 700;
+            font-style: normal;
+            font-family: Open Sans, helvetica, arial, sans-serif;
+            color: #666666;
+            width: 126px;
+            text-align: center;
+        }
+        #flyte-frame-tabs li:last-child {
+            margin-right: 0;
+        }
+        #flyte-frame-tabs li.active {
+            border-bottom: 4px solid rgb(163, 26, 255);
+            color: #333333;
+        }
+        #flyte-frame-container {
+            width: auto;
+        }
+        #flyte-frame-container > div {
+            display: None;
+        }
+        #flyte-frame-container > div.active {
+            display: block;
+            padding: 2rem 2rem;
+        }
+    </style>
+</head>
+<body>
+    <nav id="flyte-frame-nav">
+        <ul id="flyte-frame-tabs">
+            $NAV_HTML
+        </ul>
+    </nav>
+    <div id="flyte-frame-container">
+        $BODY_HTML
+    </div>
+</body>
+<script>
+    const setTabs = index => {
+        const container = document.getElementById('flyte-frame-tabs')
+        for (let i = 0; i < container.children.length; i++) {
+            const tabIndex = container.children[i].getAttribute('link_index')
+            if (tabIndex === index) {
+                container.children[i].classList.add('active')
+            } else {
+                container.children[i].className = ''
+            }
+        }
+    }
+    const setContent = index => {
+        const container = document.getElementById('flyte-frame-container')
+        for (let i = 0; i < container.children.length; i++) {
+            const tabIndex = container.children[i].getAttribute('link_index')
+            if (tabIndex === index) {
+                container.children[i].classList.add('active')
+            } else {
+                container.children[i].className = ''
+            }
+        }
+    }
+    const setLinkIndex = index => {
+        setTabs(index)
+        setContent(index)
+    }
+    const handleLinkClick = e => {
+        const linkIndex = e.getAttribute('link_index');
+        setLinkIndex(linkIndex)
+    }
+    const tabs = document.getElementById('flyte-frame-tabs');
+    const containers = document.getElementById('flyte-frame-container');
+    for(var i = 0; i < tabs.children.length; i++) {
+        if (i === 0) {
+            tabs.children[i].classList.add('active')
+            containers.children[i].classList.add('active')
+        }
+        tabs.children[i].setAttribute("link_index", i+1)
+        containers.children[i].setAttribute("link_index", i+1)
+    }
+</script>
+</html>

flyte/storage/__init__.py ADDED Viewed

@@ -0,0 +1,29 @@
+__all__ = [
+    "ABFS",
+    "GCS",
+    "S3",
+    "Storage",
+    "get",
+    "get_random_local_directory",
+    "get_random_local_path",
+    "get_stream",
+    "get_underlying_filesystem",
+    "is_remote",
+    "join",
+    "put",
+    "put_stream",
+    "put_stream",
+]
+from ._config import ABFS, GCS, S3, Storage
+from ._storage import (
+    get,
+    get_random_local_directory,
+    get_random_local_path,
+    get_stream,
+    get_underlying_filesystem,
+    is_remote,
+    join,
+    put,
+    put_stream,
+)

flyte/storage/_config.py ADDED Viewed

@@ -0,0 +1,233 @@
+from __future__ import annotations
+import datetime
+import os
+import typing
+from dataclasses import dataclass
+from typing import ClassVar
+from flyte.config import set_if_exists
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class Storage(object):
+    """
+    Data storage configuration that applies across any provider.
+    """
+    retries: int = 3
+    backoff: datetime.timedelta = datetime.timedelta(seconds=5)
+    enable_debug: bool = False
+    attach_execution_metadata: bool = True
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "enable_debug": "UNION_STORAGE_DEBUG",
+        "retries": "UNION_STORAGE_RETRIES",
+        "backoff": "UNION_STORAGE_BACKOFF_SECONDS",
+    }
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        """
+        Returns the configuration as kwargs for constructing an fsspec filesystem.
+        """
+        return {}
+    @classmethod
+    def _auto_as_kwargs(cls) -> typing.Dict[str, typing.Any]:
+        retries = os.getenv(cls._KEY_ENV_VAR_MAPPING["retries"])
+        backoff = os.getenv(cls._KEY_ENV_VAR_MAPPING["backoff"])
+        enable_debug = os.getenv(cls._KEY_ENV_VAR_MAPPING["enable_debug"])
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "enable_debug", enable_debug)
+        kwargs = set_if_exists(kwargs, "retries", retries)
+        kwargs = set_if_exists(kwargs, "backoff", backoff)
+        return kwargs
+    @classmethod
+    def auto(cls) -> Storage:
+        """
+        Construct the config object automatically from environment variables.
+        """
+        return cls(**cls._auto_as_kwargs())
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class S3(Storage):
+    """
+    S3 specific configuration
+    """
+    endpoint: typing.Optional[str] = None
+    access_key_id: typing.Optional[str] = None
+    secret_access_key: typing.Optional[str] = None
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "endpoint": "FLYTE_AWS_ENDPOINT",
+        "access_key_id": "FLYTE_AWS_ACCESS_KEY_ID",
+        "secret_access_key": "FLYTE_AWS_SECRET_ACCESS_KEY",
+    } | Storage._KEY_ENV_VAR_MAPPING
+    # Refer to https://github.com/developmentseed/obstore/blob/33654fc37f19a657689eb93327b621e9f9e01494/obstore/python/obstore/store/_aws.pyi#L11
+    # for key and secret
+    _CONFIG_KEY_FSSPEC_S3_KEY_ID: ClassVar = "access_key_id"
+    _CONFIG_KEY_FSSPEC_S3_SECRET: ClassVar = "secret_access_key"
+    _CONFIG_KEY_ENDPOINT: ClassVar = "endpoint_url"
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+    @classmethod
+    def auto(cls) -> S3:
+        """
+        :return: Config
+        """
+        endpoint = os.getenv(cls._KEY_ENV_VAR_MAPPING["endpoint"], None)
+        access_key_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["access_key_id"], None)
+        secret_access_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["secret_access_key"], None)
+        kwargs = super()._auto_as_kwargs()
+        kwargs = set_if_exists(kwargs, "endpoint", endpoint)
+        kwargs = set_if_exists(kwargs, "access_key_id", access_key_id)
+        kwargs = set_if_exists(kwargs, "secret_access_key", secret_access_key)
+        return S3(**kwargs)
+    @classmethod
+    def for_sandbox(cls) -> S3:
+        """
+        :return:
+        """
+        kwargs = super()._auto_as_kwargs()
+        final_kwargs = kwargs | {
+            "endpoint": "http://localhost:4566",
+            "access_key_id": "minio",
+            "secret_access_key": "miniostorage",
+        }
+        return S3(**final_kwargs)
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        # Construct the config object
+        kwargs.pop("anonymous", None)  # Remove anonymous if it exists, as we handle it separately
+        config: typing.Dict[str, typing.Any] = {}
+        if self._CONFIG_KEY_FSSPEC_S3_KEY_ID in kwargs or self.access_key_id:
+            config[self._CONFIG_KEY_FSSPEC_S3_KEY_ID] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_KEY_ID, self.access_key_id
+            )
+        if self._CONFIG_KEY_FSSPEC_S3_SECRET in kwargs or self.secret_access_key:
+            config[self._CONFIG_KEY_FSSPEC_S3_SECRET] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_SECRET, self.secret_access_key
+            )
+        if self._CONFIG_KEY_ENDPOINT in kwargs or self.endpoint:
+            config["endpoint_url"] = kwargs.pop(self._CONFIG_KEY_ENDPOINT, self.endpoint)
+        retries = kwargs.pop("retries", self.retries)
+        backoff = kwargs.pop("backoff", self.backoff)
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+        retry_config = {
+            "max_retries": retries,
+            "backoff": {
+                "base": 2,
+                "init_backoff": backoff,
+                "max_backoff": datetime.timedelta(seconds=16),
+            },
+            "retry_timeout": datetime.timedelta(minutes=3),
+        }
+        client_options = {"timeout": "99999s", "allow_http": True}
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options or None
+        kwargs["retry_config"] = retry_config or None
+        return kwargs
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class GCS(Storage):
+    """
+    Any GCS specific configuration.
+    """
+    gsutil_parallelism: bool = False
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "gsutil_parallelism": "GCP_GSUTIL_PARALLELISM",
+    }
+    @classmethod
+    def auto(cls) -> GCS:
+        gsutil_parallelism = os.getenv(cls._KEY_ENV_VAR_MAPPING["gsutil_parallelism"], None)
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "gsutil_parallelism", gsutil_parallelism)
+        return GCS(**kwargs)
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        return kwargs
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class ABFS(Storage):
+    """
+    Any Azure Blob Storage specific configuration.
+    """
+    account_name: typing.Optional[str] = None
+    account_key: typing.Optional[str] = None
+    tenant_id: typing.Optional[str] = None
+    client_id: typing.Optional[str] = None
+    client_secret: typing.Optional[str] = None
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "account_name": "AZURE_STORAGE_ACCOUNT_NAME",
+        "account_key": "AZURE_STORAGE_ACCOUNT_KEY",
+        "tenant_id": "AZURE_TENANT_ID",
+        "client_id": "AZURE_CLIENT_ID",
+        "client_secret": "AZURE_CLIENT_SECRET",
+    }
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+    @classmethod
+    def auto(cls) -> ABFS:
+        account_name = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_name"], None)
+        account_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_key"], None)
+        tenant_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["tenant_id"], None)
+        client_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_id"], None)
+        client_secret = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_secret"], None)
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "account_name", account_name)
+        kwargs = set_if_exists(kwargs, "account_key", account_key)
+        kwargs = set_if_exists(kwargs, "tenant_id", tenant_id)
+        kwargs = set_if_exists(kwargs, "client_id", client_id)
+        kwargs = set_if_exists(kwargs, "client_secret", client_secret)
+        return ABFS(**kwargs)
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        config: typing.Dict[str, typing.Any] = {}
+        if "account_name" in kwargs or self.account_name:
+            config["account_name"] = kwargs.get("account_name", self.account_name)
+        if "account_key" in kwargs or self.account_key:
+            config["account_key"] = kwargs.get("account_key", self.account_key)
+        if "client_id" in kwargs or self.client_id:
+            config["client_id"] = kwargs.get("client_id", self.client_id)
+        if "client_secret" in kwargs or self.client_secret:
+            config["client_secret"] = kwargs.get("client_secret", self.client_secret)
+        if "tenant_id" in kwargs or self.tenant_id:
+            config["tenant_id"] = kwargs.get("tenant_id", self.tenant_id)
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+        client_options = {"timeout": "99999s", "allow_http": "true"}
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options
+        return kwargs

flyte/storage/_remote_fs.py ADDED Viewed

@@ -0,0 +1,34 @@
+from __future__ import annotations
+import threading
+import typing
+# This file system is not really a filesystem, so users aren't really able to specify the remote path,
+# at least not yet.
+REMOTE_PLACEHOLDER = "flyte://data"
+HashStructure = typing.Dict[str, typing.Tuple[bytes, int]]
+class RemoteFSPathResolver:
+    protocol = "flyte://"
+    _flyte_path_to_remote_map: typing.ClassVar[typing.Dict[str, str]] = {}
+    _lock = threading.Lock()
+    @classmethod
+    def resolve_remote_path(cls, flyte_uri: str) -> typing.Optional[str]:
+        """
+        Given a flyte uri, return the remote path if it exists or was created in current session, otherwise return None
+        """
+        with cls._lock:
+            if flyte_uri in cls._flyte_path_to_remote_map:
+                return cls._flyte_path_to_remote_map[flyte_uri]
+            return None
+    @classmethod
+    def add_mapping(cls, flyte_uri: str, remote_path: str):
+        """
+        Thread safe method to dd a mapping from a flyte uri to a remote path
+        """
+        with cls._lock:
+            cls._flyte_path_to_remote_map[flyte_uri] = remote_path

flyte/storage/_storage.py ADDED Viewed

@@ -0,0 +1,271 @@
+import os
+import pathlib
+import random
+import tempfile
+import typing
+from typing import AsyncIterator, Optional
+from uuid import UUID
+import fsspec
+from fsspec.asyn import AsyncFileSystem
+from fsspec.utils import get_protocol
+from obstore.exceptions import GenericError
+from obstore.fsspec import register
+from flyte._initialize import get_storage
+from flyte._logging import logger
+def is_remote(path: typing.Union[pathlib.Path | str]) -> bool:
+    """
+    Let's find a replacement
+    """
+    protocol = get_protocol(str(path))
+    if protocol is None:
+        return False
+    return protocol != "file"
+def strip_file_header(path: str) -> str:
+    """
+    Drops file:// if it exists from the file
+    """
+    if path.startswith("file://"):
+        return path.replace("file://", "", 1)
+    return path
+def get_random_local_path(file_path_or_file_name: pathlib.Path | str | None = None) -> pathlib.Path:
+    """
+    Use file_path_or_file_name, when you want a random directory, but want to preserve the leaf file name
+    """
+    local_tmp = pathlib.Path(tempfile.mkdtemp(prefix="flyte-tmp-"))
+    key = UUID(int=random.getrandbits(128)).hex
+    tmp_folder = local_tmp / key
+    tail = ""
+    if file_path_or_file_name:
+        _, tail = os.path.split(file_path_or_file_name)
+    if tail:
+        tmp_folder.mkdir(parents=True, exist_ok=True)
+        return tmp_folder / tail
+    local_tmp.mkdir(parents=True, exist_ok=True)
+    return tmp_folder
+def get_random_local_directory() -> pathlib.Path:
+    """
+    :return: a random directory
+    :rtype: pathlib.Path
+    """
+    _dir = get_random_local_path(None)
+    pathlib.Path(_dir).mkdir(parents=True, exist_ok=True)
+    return _dir
+def get_underlying_filesystem(
+    protocol: typing.Optional[str] = None,
+    anonymous: bool = False,
+    path: typing.Optional[str] = None,
+    **kwargs,
+) -> fsspec.AbstractFileSystem:
+    if protocol is None:
+        # If protocol is None, get it from the path
+        protocol = get_protocol(path)
+    storage_config = get_storage()
+    if storage_config:
+        kwargs = storage_config.get_fsspec_kwargs(anonymous, **kwargs)
+    elif protocol:
+        match protocol:
+            case "s3":
+                # If the protocol is s3, we can use the s3 filesystem
+                from flyte.storage import S3
+                kwargs = S3.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case "gs":
+                # If the protocol is gs, we can use the gs filesystem
+                from flyte.storage import GCS
+                kwargs = GCS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case "abfs" | "abfss":
+                # If the protocol is abfs or abfss, we can use the abfs filesystem
+                from flyte.storage import ABFS
+                kwargs = ABFS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case _:
+                pass
+    return fsspec.filesystem(protocol, **kwargs)
+def _get_anonymous_filesystem(from_path):
+    """Get the anonymous file system if needed."""
+    return get_underlying_filesystem(get_protocol(from_path), anonymous=True, asynchronous=True)
+async def get(from_path: str, to_path: Optional[str | pathlib.Path] = None, recursive: bool = False, **kwargs) -> str:
+    if not to_path:
+        name = pathlib.Path(from_path).name
+        to_path = get_random_local_path(file_path_or_file_name=name)
+        logger.debug(f"Storing file from {from_path} to {to_path}")
+    file_system = get_underlying_filesystem(path=from_path)
+    try:
+        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
+    except (OSError, GenericError) as oe:
+        logger.debug(f"Error in getting {from_path} to {to_path} rec {recursive} {oe}")
+        if isinstance(file_system, AsyncFileSystem):
+            try:
+                exists = await file_system._exists(from_path)  # pylint: disable=W0212
+            except GenericError:
+                # for obstore, as it does not raise FileNotFoundError in fsspec but GenericError
+                # force it to try get_filesystem(anonymous=True)
+                exists = True
+        else:
+            exists = file_system.exists(from_path)
+        if not exists:
+            # TODO: update exception to be more specific
+            raise AssertionError(f"Unable to load data from {from_path}")
+        file_system = _get_anonymous_filesystem(from_path)
+        logger.debug(f"Attempting anonymous get with {file_system}")
+        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
+async def _get_from_filesystem(
+    file_system: fsspec.AbstractFileSystem,
+    from_path: str | pathlib.Path,
+    to_path: str | pathlib.Path,
+    recursive: bool,
+    **kwargs,
+):
+    if isinstance(file_system, AsyncFileSystem):
+        dst = await file_system._get(from_path, to_path, recursive=recursive, **kwargs)  # pylint: disable=W0212
+    else:
+        dst = file_system.get(from_path, to_path, recursive=recursive, **kwargs)
+    if isinstance(dst, (str, pathlib.Path)):
+        return dst
+    return to_path
+async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs) -> str:
+    if not to_path:
+        from flyte._context import internal_ctx
+        ctx = internal_ctx()
+        name = pathlib.Path(from_path).name if not recursive else None  # don't pass a name for folders
+        to_path = ctx.raw_data.get_random_remote_path(file_name=name)
+    file_system = get_underlying_filesystem(path=to_path)
+    from_path = strip_file_header(from_path)
+    if isinstance(file_system, AsyncFileSystem):
+        dst = await file_system._put(from_path, to_path, recursive=recursive, **kwargs)  # pylint: disable=W0212
+    else:
+        dst = file_system.put(from_path, to_path, recursive=recursive, **kwargs)
+    if isinstance(dst, (str, pathlib.Path)):
+        return str(dst)
+    else:
+        return to_path
+async def put_stream(
+    data_iterable: typing.AsyncIterable[bytes] | bytes, *, name: str | None = None, to_path: str | None = None, **kwargs
+) -> str:
+    """
+    Put a stream of data to a remote location. This is useful for streaming data to a remote location.
+    Example usage:
+    ```python
+    import flyte.storage as storage
+    storage.put_stream(iter([b'hello']), name="my_file.txt")
+    OR
+    storage.put_stream(iter([b'hello']), to_path="s3://my_bucket/my_file.txt")
+    ```
+    :param data_iterable: Iterable of bytes to be streamed.
+    :param name: Name of the file to be created. If not provided, a random name will be generated.
+    :param to_path: Path to the remote location where the data will be stored.
+    :param kwargs: Additional arguments to be passed to the underlying filesystem.
+    :rtype: str
+    :return: The path to the remote location where the data was stored.
+    """
+    if not to_path:
+        from flyte._context import internal_ctx
+        ctx = internal_ctx()
+        to_path = ctx.raw_data.get_random_remote_path(file_name=name)
+    fs = get_underlying_filesystem(path=to_path)
+    file_handle = None
+    if isinstance(fs, AsyncFileSystem):
+        try:
+            file_handle = await fs.open_async(to_path, "wb", **kwargs)
+            if isinstance(data_iterable, bytes):
+                await file_handle.write(data_iterable)
+            else:
+                async for data in data_iterable:
+                    await file_handle.write(data)
+            return str(to_path)
+        except NotImplementedError:
+            logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync")
+        finally:
+            if file_handle is not None:
+                await file_handle.close()
+    with fs.open(to_path, "wb", **kwargs) as f:
+        if isinstance(data_iterable, bytes):
+            f.write(data_iterable)
+        else:
+            # If data_iterable is async iterable, iterate over it and write each chunk to the file
+            async for data in data_iterable:
+                f.write(data)
+    return str(to_path)
+async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncIterator[bytes]:
+    """
+    Get a stream of data from a remote location.
+    This is useful for downloading streaming data from a remote location.
+    Example usage:
+    ```python
+    import flyte.storage as storage
+    obj = storage.get_stream(path="s3://my_bucket/my_file.txt")
+    ```
+    :param path: Path to the remote location where the data will be downloaded.
+    :param kwargs: Additional arguments to be passed to the underlying filesystem.
+    :param chunk_size: Size of each chunk to be read from the file.
+    :return: An async iterator that yields chunks of data.
+    """
+    fs = get_underlying_filesystem(path=path, **kwargs)
+    file_size = fs.info(path)["size"]
+    total_read = 0
+    file_handle = None
+    try:
+        if isinstance(fs, AsyncFileSystem):
+            file_handle = await fs.open_async(path, "rb")
+            while chunk := await file_handle.read(min(chunk_size, file_size - total_read)):
+                total_read += len(chunk)
+                yield chunk
+            return
+    except NotImplementedError:
+        logger.debug(f"{fs} doesn't implement 'open_async', falling back to sync")
+    finally:
+        if file_handle is not None:
+            file_handle.close()
+    # Sync fallback
+    with fs.open(path, "rb") as file_handle:
+        while chunk := file_handle.read(min(chunk_size, file_size - total_read)):
+            total_read += len(chunk)
+            yield chunk
+def join(*paths: str) -> str:
+    """
+    Join multiple paths together. This is a wrapper around os.path.join.
+    # TODO replace with proper join with fsspec root etc
+    :param paths: Paths to be joined.
+    """
+    return str(os.path.join(*paths))
+register(["s3", "gs", "abfs", "abfss"], asynchronous=True)

flyte/storage/_utils.py ADDED Viewed

@@ -0,0 +1,5 @@
+import os
+# This is the default chunk size flyte will use for writing to S3 and GCS. This is set to 25MB by default and is
+# configurable by the user if needed. This is used when put() is called on filesystems.
+_WRITE_SIZE_CHUNK_BYTES = int(os.environ.get("_F_P_WRITE_CHUNK_SIZE", "26214400"))  # 25 * 2**20

flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

Potentially problematic release.

flyte 0.1.0py3-none-any.whl → 0.2.0a0py3-none-any.whl