flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/extras/_container.py

@@ -2,27 +2,23 @@ import os
 import pathlib
 from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
 from flyte import Image, storage
-from flyte._datastructures import NativeInterface, SerializationContext
 from flyte._logging import logger
 from flyte._task import TaskTemplate
+from flyte.io import Dir, File
+from flyte.models import NativeInterface, SerializationContext
 
-_PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
 
-
-def _extract_command_key(cmd: str, **kwargs) -> Any:
+def _extract_command_key(cmd: str, **kwargs) -> List[Any] | None:
     """
     Extract the key from the command using regex.
     """
     import re
 
-    input_regex = r"^\{\{\s*\.inputs\.(.*?)\s*\}\}$"
-    match = re.match(input_regex, cmd)
-    if match:
-        return match.group(1)
-    return None
+    input_regex = r"\{\{\.inputs\.([a-zA-Z0-9_]+)\}\}"
+    return re.findall(input_regex, cmd)
 
 
 def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Optional[str]:
@@ -32,8 +28,9 @@ def _extract_path_command_key(cmd: str, input_data_dir: Optional[str]) -> Option
     import re
 
     input_data_dir = input_data_dir or ""
-    input_regex = rf"{re.escape(input_data_dir)}/(.+)$"
-    match = re.match(input_regex, cmd)
+    input_regex = rf"{re.escape(input_data_dir)}/([\w\-.]+)"  # captures file or dir names
+
+    match = re.search(input_regex, cmd)
     if match:
         return match.group(1)
     return None
@@ -70,7 +67,7 @@ class ContainerTask(TaskTemplate):
         input_data_dir: str | pathlib.Path = "/var/inputs",
         output_data_dir: str | pathlib.Path = "/var/outputs",
         metadata_format: MetadataFormat = "JSON",
-        local_logs: bool = False,
+        local_logs: bool = True,
         **kwargs,
     ):
         super().__init__(
@@ -83,9 +80,14 @@ class ContainerTask(TaskTemplate):
         self._image = image
         if isinstance(image, str):
             if image == "auto":
-                self._image = Image.auto()
+                self._image = Image.from_debian_base()
             else:
-                self._image = Image.from_prebuilt(image)
+                self._image = Image.from_base(image)
+
+        if command and any(not isinstance(c, str) for c in command):
+            raise ValueError("All elements in the command list must be strings.")
+        if arguments and any(not isinstance(a, str) for a in arguments):
+            raise ValueError("All elements in the arguments list must be strings.")
         self._cmd = command
         self._args = arguments
         self._input_data_dir = input_data_dir
@@ -106,32 +108,34 @@ class ContainerTask(TaskTemplate):
         For FlyteFile and FlyteDirectory commands, e.g., "/var/inputs/inputs", we extract the key from strings that
          begin with the specified `input_data_dir`.
         """
-        # from flytekit.types.directory import FlyteDirectory
-        # from flytekit.types.file import FlyteFile
+        from flyte.io import Dir, File
 
         volume_binding: Dict[str, Dict[str, str]] = {}
         path_k = _extract_path_command_key(cmd, str(self._input_data_dir))
-        k = path_k if path_k else _extract_command_key(cmd)
-
-        if k:
-            input_val = kwargs.get(k)
-            # TODO: Add support file and directory transformer first
-            # if type(input_val) in [FlyteFile, FlyteDirectory]:
-            #     if not path_k:
-            #         raise AssertionError(
-            #             "FlyteFile and FlyteDirectory commands should not use the template syntax like this:
-            #             {{.inputs.infile}}\n"
-            #             "Please use a path-like syntax, such as: /var/inputs/infile.\n"
-            #             "This requirement is due to how Flyte Propeller processes template syntax inputs."
-            #         )
-            #     local_flyte_file_or_dir_path = str(input_val)
-            #     remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
-            #     volume_binding[local_flyte_file_or_dir_path] = {
-            #         "bind": remote_flyte_file_or_dir_path,
-            #         "mode": "rw",
-            #     }
-            #     command = remote_flyte_file_or_dir_path
-            command = str(input_val)
+        keys = [path_k] if path_k else _extract_command_key(cmd)
+
+        command = cmd
+
+        if keys:
+            for k in keys:
+                input_val = kwargs.get(k)
+                # TODO: Add support file and directory transformer first
+                if input_val and type(input_val) in [File, Dir]:
+                    if not path_k:
+                        raise AssertionError(
+                            "File and Directory commands should not use the template syntax "
+                            "like this: {{.inputs.infile}}\n"
+                            "Please use a path-like syntax, such as: /var/inputs/infile.\n"
+                            "This requirement is due to how Flyte Propeller processes template syntax inputs."
+                        )
+                    local_flyte_file_or_dir_path = input_val.path
+                    remote_flyte_file_or_dir_path = os.path.join(self._input_data_dir, k)  # type: ignore
+                    volume_binding[local_flyte_file_or_dir_path] = {
+                        "bind": remote_flyte_file_or_dir_path,
+                        "mode": "rw",
+                    }
+                else:
+                    command = command.replace(f"{{{{.inputs.{k}}}}}", str(input_val))
         else:
             command = cmd
 
@@ -193,7 +197,9 @@ class ContainerTask(TaskTemplate):
             microseconds=microseconds,
         )
 
-    def _convert_output_val_to_correct_type(self, output_val: Any, output_type: Type) -> Any:
+    async def _convert_output_val_to_correct_type(
+        self, output_path: pathlib.Path, output_val: Any, output_type: Type
+    ) -> Any:
         import datetime
 
         if issubclass(output_type, bool):
@@ -202,20 +208,31 @@ class ContainerTask(TaskTemplate):
             return datetime.datetime.fromisoformat(output_val)
         elif issubclass(output_type, datetime.timedelta):
             return self._string_to_timedelta(output_val)
+        elif issubclass(output_type, File):
+            return await File.from_local(output_path)
+        elif issubclass(output_type, Dir):
+            return await Dir.from_local(output_path)
         else:
             return output_type(output_val)
 
-    def _get_output_dict(self, output_directory: pathlib.Path) -> Dict[str, Any]:
-        output_dict = {}
+    async def _get_output(self, output_directory: pathlib.Path) -> Tuple[Any]:
+        output_items = []
         if self._outputs:
             for k, output_type in self._outputs.items():
                 output_path = output_directory / k
-                with output_path.open("r") as f:
-                    output_val = f.read()
-                output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
-        return output_dict
-
-    def execute(self, **kwargs) -> Any:
+                if os.path.isfile(output_path):
+                    with output_path.open("r") as f:
+                        output_val = f.read()
+                else:
+                    output_val = None
+                parsed = await self._convert_output_val_to_correct_type(output_path, output_val, output_type)
+                output_items.append(parsed)
+        # return a tuple so that each element is treated as a separate output.
+        # this allows flyte to map the user-defined output types (dict) to individual values.
+        # if we returned a list instead, it would be treated as a single output.
+        return tuple(output_items)
+
+    async def execute(self, **kwargs) -> Any:
         try:
             import docker
         except ImportError:
@@ -235,6 +252,7 @@ class ContainerTask(TaskTemplate):
             raise AssertionError(f"Only Image objects are supported, not strings. Got {self._image} instead.")
         uri = self._image.uri
         self._pull_image_if_not_exists(client, uri)
+        print(f"Command: {commands!r}")
 
         container = client.containers.run(uri, command=commands, remove=True, volumes=volume_bindings, detach=True)
 
@@ -247,8 +265,8 @@ class ContainerTask(TaskTemplate):
 
         container.wait()
 
-        output_dict = self._get_output_dict(output_directory)
-        return output_dict
+        output = await self._get_output(output_directory)
+        return output
 
     def data_loading_config(self, sctx: SerializationContext) -> tasks_pb2.DataLoadingConfig:
         literal_to_protobuf = {
@@ -258,16 +276,11 @@ class ContainerTask(TaskTemplate):
         }
 
         return tasks_pb2.DataLoadingConfig(
-            input_path=self._input_data_dir,
-            output_path=self._output_data_dir,
+            input_path=str(self._input_data_dir) if self._input_data_dir else None,
+            output_path=str(self._output_data_dir) if self._output_data_dir else None,
             enabled=True,
             format=literal_to_protobuf.get(self._metadata_format, "JSON"),
         )
 
     def container_args(self, sctx: SerializationContext) -> List[str]:
         return self._cmd + (self._args if self._args else [])
-
-    def config(self, sctx: SerializationContext) -> Dict[str, str]:
-        if self.pod_template is None:
-            return {}
-        return {_PRIMARY_CONTAINER_NAME_FIELD: self.primary_container_name}

flyte/git/__init__.py

@@ -0,0 +1,3 @@
+from ._config import GitStatus, config_from_root
+
+__all__ = ["GitStatus", "config_from_root"]

flyte/git/_config.py

@@ -0,0 +1,279 @@
+import pathlib
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Protocol
+
+import flyte.config
+from flyte._logging import logger
+
+
+class GitUrlBuilder(Protocol):
+    @staticmethod
+    def build_url(remote_url: str, file_path: str, commit_sha: str, line_number: int, is_tree_clean: bool) -> str: ...
+
+
+class GithubUrlBuilder(GitUrlBuilder):
+    host_name = "github.com"
+
+    @staticmethod
+    def build_url(remote_url: str, file_path: str, commit_sha: str, line_number: int, is_tree_clean: bool) -> str:
+        url = f"{remote_url}/blob/{commit_sha}/{file_path}"
+        if is_tree_clean:
+            url += f"#L{line_number}"
+        return url
+
+
+class GitlabUrlBuilder(GitUrlBuilder):
+    host_name = "gitlab.com"
+
+    @staticmethod
+    def build_url(remote_url: str, file_path: str, commit_sha: str, line_number: int, is_tree_clean: bool) -> str:
+        url = f"{remote_url}/-/blob/{commit_sha}/{file_path}"
+        if is_tree_clean:
+            url += f"#L{line_number}"
+        return url
+
+
+GIT_URL_BUILDER_REGISTRY: Dict[str, GitUrlBuilder] = {
+    GithubUrlBuilder.host_name: GithubUrlBuilder,
+    GitlabUrlBuilder.host_name: GitlabUrlBuilder,
+}
+
+
+@dataclass(init=True, frozen=True)
+class GitStatus:
+    """
+    A class representing the status of a git repository.
+
+    :param is_valid: Whether git repository is valid
+    :param is_tree_clean: Whether working tree is clean
+    :param remote_url: Remote URL in HTTPS format
+    :param repo_dir: Repository root directory
+    :param commit_sha: Current commit SHA
+    """
+
+    is_valid: bool = False
+    is_tree_clean: bool = False
+    remote_url: str = ""
+    repo_dir: Path = Path()
+    commit_sha: str = ""
+
+    @classmethod
+    def from_current_repo(cls) -> "GitStatus":
+        """Discover git information from the current repository.
+
+        If Git is not installed or .git does not exist, returns GitStatus with is_valid=False.
+
+        :return: GitStatus instance with discovered git information
+        """
+        try:
+            # Check if we're in a git repository and get the root directory
+            result = subprocess.run(
+                ["git", "rev-parse", "--show-toplevel"],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            if result.returncode != 0:
+                logger.warning("Not in a git repository or git is not installed")
+                return cls()
+
+            repo_dir = Path(result.stdout.strip())
+
+            # Get current commit SHA
+            result = subprocess.run(
+                ["git", "rev-parse", "HEAD"],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode == 0:
+                commit_sha = result.stdout.strip()
+            else:
+                logger.warning("Failed to get current commit SHA")
+                return cls(repo_dir=repo_dir)
+
+            # Check if working tree is clean
+            result = subprocess.run(
+                ["git", "status", "--porcelain"],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode == 0:
+                is_tree_clean = len(result.stdout.strip()) == 0
+            else:
+                logger.warning("Failed to check if working tree is clean")
+                return cls(repo_dir=repo_dir, commit_sha=commit_sha)
+
+            # Get remote URL
+            instance = cls(repo_dir=repo_dir, commit_sha=commit_sha, is_tree_clean=is_tree_clean)
+            remote_url = instance._get_remote_url()
+            if not remote_url:
+                logger.warning("Failed to get remote URL")
+                return cls(repo_dir=repo_dir, commit_sha=commit_sha, is_tree_clean=is_tree_clean)
+
+            return cls(
+                is_valid=True,
+                is_tree_clean=is_tree_clean,
+                remote_url=remote_url,
+                repo_dir=repo_dir,
+                commit_sha=commit_sha,
+            )
+
+        except Exception as e:
+            logger.debug(f"Failed to discover git repository: {e}")
+            return cls()
+
+    def _get_remote_url(self) -> str:
+        """Get the remote push URL.
+
+        Returns the 'origin' remote push URL if it exists, otherwise returns
+        the first remote alphabetically. Converts SSH/Git protocol URLs to HTTPS format.
+
+        :return: The remote push URL in HTTPS format, or empty string if not found
+        """
+        try:
+            # Try to get origin push remote first
+            result = subprocess.run(
+                ["git", "remote", "get-url", "--push", "origin"],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            if result.returncode == 0:
+                url = result.stdout.strip()
+                return self._normalize_url_to_https(url)
+
+            # If origin doesn't exist, get all remotes
+            result = subprocess.run(
+                ["git", "remote"],
+                check=False,
+                capture_output=True,
+                text=True,
+            )
+
+            if result.returncode == 0:
+                remotes = result.stdout.strip().split("\n")
+                if remotes:
+                    # Sort alphabetically and get the first one
+                    remotes.sort()
+                    first_remote = remotes[0]
+
+                    # Get push URL for this remote
+                    result = subprocess.run(
+                        ["git", "remote", "get-url", "--push", first_remote],
+                        check=False,
+                        capture_output=True,
+                        text=True,
+                    )
+                    if result.returncode == 0:
+                        url = result.stdout.strip()
+                        return self._normalize_url_to_https(url)
+
+            return ""
+
+        except Exception:
+            return ""
+
+    def _normalize_url_to_https(self, url: str) -> str:
+        """Convert SSH or Git protocol URLs to HTTPS format.
+
+        Examples:
+            git@github.com:user/repo.git -> https://github.com/user/repo
+            https://github.com/user/repo.git -> https://github.com/user/repo
+
+        :param url: The Git URL to normalize
+        :return: The normalized HTTPS URL
+        """
+        # Remove .git suffix first
+        url = url.removesuffix(".git")
+
+        # Handle SSH format: git@host:path or user@host:path
+        if url.startswith("git@"):
+            parts = url.split("@", 1)
+            if len(parts) == 2:
+                host_and_path = parts[1].replace(":", "/", 1)
+                return f"https://{host_and_path}"
+
+        return url
+
+    def _get_remote_host(self, url: str) -> str:
+        """Get the remote host name from a normalized HTTPS URL.
+
+        :param url: URL that has been normalized to HTTPS format by _normalize_url_to_https
+        :return: The host name (e.g., "github.com", "gitlab.com")
+        """
+        parts = url.split("//", 1)
+        if len(parts) < 2:
+            return ""
+
+        # Get everything after "//" and split by "/"
+        host_and_path = parts[1]
+        parts = host_and_path.split("/", 1)
+        if len(parts) < 2:
+            return ""
+        host = host_and_path.split("/")[0]
+
+        return host
+
+    def _get_file_path(self, path: Path | str) -> str:
+        """Get the path relative to the repository root directory.
+
+        :param path: Absolute or relative path to a file
+        :return: Path relative to repo_dir as string, or empty string if failed
+        """
+        try:
+            path_obj = Path(path).resolve()
+            relative_path = path_obj.relative_to(self.repo_dir)
+            return str(relative_path)
+        except Exception as e:
+            logger.warning(f"Failed to get relative path for {path}: {e}")
+            return ""
+
+    def build_url(self, path: Path | str, line_number: int) -> str:
+        """Build a git URL for the given path.
+
+        :param path: Path to a file
+        :param line_number: Line number of the code file
+        :return: Path relative to repo_dir
+        """
+        if not self.is_valid:
+            logger.warning("GitConfig is not valid, cannot build URL")
+            return ""
+        host_name = self._get_remote_host(self.remote_url)
+        git_file_path = self._get_file_path(path)
+        if not host_name:
+            logger.warning(f"Failed to extract host name from remote URL: {self.remote_url}")
+            return ""
+        if not git_file_path:
+            return ""
+        builder = GIT_URL_BUILDER_REGISTRY.get(host_name)
+        if not builder:
+            logger.warning(f"URL builder for {host_name} is not implemented")
+            return ""
+        url = builder.build_url(self.remote_url, git_file_path, self.commit_sha, line_number, self.is_tree_clean)
+        return url
+
+
+def config_from_root(path: pathlib.Path | str = ".flyte/config.yaml") -> flyte.config.Config | None:
+    """Get the config file from the git root directory.
+
+    By default, the config file is expected to be in `.flyte/config.yaml` in the git root directory.
+
+    :param path: Path to the config file relative to git root directory (default: ".flyte/config.yaml")
+    :return: Config object if found, None otherwise
+    """
+    try:
+        result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
+        if result.returncode != 0:
+            return None
+        root = pathlib.Path(result.stdout.strip())
+        if not (root / path).exists():
+            return None
+        return flyte.config.auto(root / path)
+    except Exception:
+        return None

flyte/io/__init__.py

@@ -3,9 +3,16 @@
 
 This package contains additional data types beyond the primitive data types in python to abstract data flow
 of large datasets in Union.
+
 """
 
-__all__ = ["Dir", "File"]
+__all__ = [
+    "PARQUET",
+    "DataFrame",
+    "Dir",
+    "File",
+]
 
+from ._dataframe import PARQUET, DataFrame
 from ._dir import Dir
 from ._file import File

flyte/io/{structured_dataset → _dataframe}/__init__.py

@@ -1,15 +1,15 @@
 """
-Flytekit StructuredDataset
+Flytekit DataFrame
 ==========================================================
-.. currentmodule:: flytekit.types.structured
+.. currentmodule:: flyte.io._dataframe
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
 
-    StructuredDataset
-    StructuredDatasetDecoder
-    StructuredDatasetEncoder
+    DataFrame
+    DataFrameDecoder
+    DataFrameEncoder
 """
 
 import functools
@@ -17,12 +17,13 @@ import functools
 from flyte._logging import logger
 from flyte._utils.lazy_module import is_imported
 
-from .structured_dataset import (
+from .dataframe import (
+    PARQUET,
+    DataFrame,
+    DataFrameDecoder,
+    DataFrameEncoder,
+    DataFrameTransformerEngine,
     DuplicateHandlerError,
-    StructuredDataset,
-    StructuredDatasetDecoder,
-    StructuredDatasetEncoder,
-    StructuredDatasetTransformerEngine,
 )
 
 
@@ -30,8 +31,8 @@ from .structured_dataset import (
 def register_csv_handlers():
     from .basic_dfs import CSVToPandasDecodingHandler, PandasToCSVEncodingHandler
 
-    StructuredDatasetTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)
 
 
 @functools.lru_cache(maxsize=None)
@@ -42,9 +43,9 @@ def register_pandas_handlers():
 
     from .basic_dfs import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler
 
-    StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
+    DataFrameTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
 
 
 @functools.lru_cache(maxsize=None)
@@ -55,9 +56,9 @@ def register_arrow_handlers():
 
     from .basic_dfs import ArrowToParquetEncodingHandler, ParquetToArrowDecodingHandler
 
-    StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register_renderer(pa.Table, ArrowRenderer())
+    DataFrameTransformerEngine.register(ArrowToParquetEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(ParquetToArrowDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register_renderer(pa.Table, ArrowRenderer())
 
 
 @functools.lru_cache(maxsize=None)
@@ -70,10 +71,10 @@ def register_bigquery_handlers():
             PandasToBQEncodingHandlers,
         )
 
-        StructuredDatasetTransformerEngine.register(PandasToBQEncodingHandlers())
-        StructuredDatasetTransformerEngine.register(BQToPandasDecodingHandler())
-        StructuredDatasetTransformerEngine.register(ArrowToBQEncodingHandlers())
-        StructuredDatasetTransformerEngine.register(BQToArrowDecodingHandler())
+        DataFrameTransformerEngine.register(PandasToBQEncodingHandlers())
+        DataFrameTransformerEngine.register(BQToPandasDecodingHandler())
+        DataFrameTransformerEngine.register(ArrowToBQEncodingHandlers())
+        DataFrameTransformerEngine.register(BQToArrowDecodingHandler())
     except ImportError:
         logger.info(
             "We won't register bigquery handler for structured dataset because "
@@ -86,8 +87,8 @@ def register_snowflake_handlers():
     try:
         from .snowflake import PandasToSnowflakeEncodingHandlers, SnowflakeToPandasDecodingHandler
 
-        StructuredDatasetTransformerEngine.register(SnowflakeToPandasDecodingHandler())
-        StructuredDatasetTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
+        DataFrameTransformerEngine.register(SnowflakeToPandasDecodingHandler())
+        DataFrameTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
 
     except ImportError:
         logger.info(
@@ -96,7 +97,7 @@ def register_snowflake_handlers():
         )
 
 
-def lazy_import_structured_dataset_handler():
+def lazy_import_dataframe_handler():
     if is_imported("pandas"):
         try:
             register_pandas_handlers()
@@ -121,9 +122,10 @@ def lazy_import_structured_dataset_handler():
 
 
 __all__ = [
-    "StructuredDataset",
-    "StructuredDatasetDecoder",
-    "StructuredDatasetEncoder",
-    "StructuredDatasetTransformerEngine",
-    "lazy_import_structured_dataset_handler",
+    "PARQUET",
+    "DataFrame",
+    "DataFrameDecoder",
+    "DataFrameEncoder",
+    "DataFrameTransformerEngine",
+    "lazy_import_dataframe_handler",
 ]