Merge branch 'master' into remove-dataclass_json

.github/actions/clear-action-cache/action.yml

@@ -0,0 +1,11 @@
+name: 'Clear action cache'
+description: 'As suggested by GitHub to prevent low disk space: https://github.com/actions/runner-images/issues/2840#issuecomment-790492173'
+runs:
+  using: 'composite'
+  steps:
+    - shell: bash
+      run: |
+        rm -rf /usr/share/dotnet
+        rm -rf /opt/ghc
+        rm -rf "/usr/local/share/boost"
+        rm -rf "$AGENT_TOOLSDIRECTORY"

.github/workflows/pythonbuild.yml

@@ -2,14 +2,13 @@ name: Build
 
 # Schedule runs to run twice a day
 
-
 on:
   push:
     branches:
       - master
   pull_request:
   schedule:
-      - cron: '0 13 * * *' # This schedule runs at 1pm UTC every day
+    - cron: "0 13 * * *" # This schedule runs at 1pm UTC every day
 
 env:
   FLYTE_SDK_LOGGING_LEVEL: 10 # debug
@@ -42,8 +41,9 @@ jobs:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
     steps:
-      - uses: insightsengineering/disk-space-reclaimer@v1
       - uses: actions/checkout@v4
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
@@ -79,8 +79,9 @@ jobs:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
     steps:
-      - uses: insightsengineering/disk-space-reclaimer@v1
       - uses: actions/checkout@v4
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
@@ -116,12 +117,13 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-latest ]
+        os: [ubuntu-latest]
         python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
-        pandas: [ "pandas<2.0.0", "pandas>=2.0.0" ]
+        pandas: ["pandas<2.0.0", "pandas>=2.0.0"]
     steps:
-      - uses: insightsengineering/disk-space-reclaimer@v1
       - uses: actions/checkout@v4
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
@@ -190,7 +192,6 @@ jobs:
         os: [ubuntu-latest]
         python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
     steps:
-      - uses: insightsengineering/disk-space-reclaimer@v1
       # As described in https://github.com/pypa/setuptools_scm/issues/414, SCM needs git history
       # and tags to work.
       - uses: actions/checkout@v4
@@ -257,8 +258,7 @@ jobs:
           - flytekit-async-fsspec
           - flytekit-aws-athena
           - flytekit-aws-batch
-          # TODO: uncomment this when the sagemaker agent is implemented: https://github.com/flyteorg/flyte/issues/4079
-          # - flytekit-aws-sagemaker
+          - flytekit-aws-sagemaker
           - flytekit-bigquery
           - flytekit-dask
           - flytekit-data-fsspec
@@ -348,9 +348,9 @@ jobs:
           - python-version: 3.12
             plugin-names: "flytekit-kf-pytorch"
     steps:
-      - uses: insightsengineering/disk-space-reclaimer@v1
-        if: ${{ matrix.plugin-names == 'flytekit-envd' }}
       - uses: actions/checkout@v4
+      - name: 'Clear action cache'
+        uses: ./.github/actions/clear-action-cache
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:

.pre-commit-config.yaml

@@ -1,11 +1,11 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.6
+    rev: v0.2.2
     hooks:
       # Run the linter.
       - id: ruff
-        args: [--fix, --show-fixes, --show-source]
+        args: [--fix, --show-fixes, --output-format=full]
       # Run the formatter.
       - id: ruff-format
   - repo: https://github.com/pre-commit/pre-commit-hooks
@@ -15,7 +15,7 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/shellcheck-py/shellcheck-py
-    rev: v0.8.0.4
+    rev: v0.10.0.1
     hooks:
       - id: shellcheck
   - repo: https://github.com/conorfalvey/check_pdb_hook

CODEOWNERS

@@ -1,3 +1,3 @@
 # These owners will be the default owners for everything in
 # the repo. Unless a later match takes precedence.
-*       @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy
+*       @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla

Dockerfile.agent

@@ -11,8 +11,9 @@ RUN pip install prometheus-client grpcio-health-checking
 RUN pip install --no-cache-dir -U flytekit==$VERSION \
   flytekitplugins-airflow==$VERSION \
   flytekitplugins-bigquery==$VERSION \
-  flytekitplugins-chatgpt==$VERSION \
+  flytekitplugins-openai==$VERSION \
   flytekitplugins-snowflake==$VERSION \
+  flytekitplugins-awssagemaker==$VERSION \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \
   && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

dev-requirements.in

@@ -38,6 +38,7 @@ python-magic; (platform_system=='Darwin' or platform_system=='Linux')
 
 types-protobuf
 types-croniter
+types-decorator
 types-mock
 autoflake
 

docs/source/plugins/awssagemaker_inference.rst

@@ -0,0 +1,12 @@
+.. _awssagemaker_inference:
+
+#####################################
+AWS Sagemaker Inference API reference
+#####################################
+
+.. tags:: Integration, MachineLearning, AWS
+
+.. automodule:: flytekitplugins.awssagemaker_inference
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:

docs/source/plugins/index.rst

@@ -30,6 +30,7 @@ Plugin API reference
 * :ref:`Vaex <vaex>` - Vaex API reference
 * :ref:`MLflow <mlflow>` - MLflow API reference
 * :ref:`DuckDB <duckdb>` - DuckDB API reference
+* :ref:`SageMaker Inference <awssagemaker_inference>` - SageMaker Inference API reference
 
 .. toctree::
    :maxdepth: 2
@@ -61,3 +62,4 @@ Plugin API reference
    Vaex <vaex>
    MLflow <mlflow>
    DuckDB <duckdb>
+   SageMaker Inference <awssagemaker_inference>

flytekit/__init__.py

@@ -179,6 +179,16 @@
 
    HashMethod
 
+Artifacts
+=========
+
+.. autosummary::
+   :nosignatures:
+   :template: custom.rst
+   :toctree: generated/
+
+   Artifact
+
 Documentation
 =============
 
@@ -207,6 +217,7 @@
 
 from flytekit._version import __version__
 from flytekit.core.array_node_map_task import map_task
+from flytekit.core.artifact import Artifact
 from flytekit.core.base_sql_task import SQLTask
 from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes
 from flytekit.core.checkpointer import Checkpoint

flytekit/clients/auth/authenticator.py

@@ -302,7 +302,9 @@ def refresh_credentials(self):
             self._verify,
             self._session,
         )
-        text = f"To Authenticate, navigate in a browser to the following URL: {click.style(resp.verification_uri, fg='blue', underline=True)} and enter code: {click.style(resp.user_code, fg='blue')}"
+
+        full_uri = f"{resp.verification_uri}?user_code={resp.user_code}"
+        text = f"To Authenticate, navigate in a browser to the following URL: {click.style(full_uri, fg='blue', underline=True)}"
         click.secho(text)
         try:
             # Currently the refresh token is not retrieved. We may want to add support for refreshTokens so that

flytekit/clis/sdk_in_container/executions.py

@@ -0,0 +1,83 @@
+import rich_click as click
+
+from flytekit.clis.sdk_in_container.constants import CTX_DOMAIN, CTX_PROJECT
+from flytekit.clis.sdk_in_container.helpers import FLYTE_REMOTE_INSTANCE_KEY, get_and_save_remote_with_click_context
+from flytekit.clis.sdk_in_container.utils import (
+    domain_option_dec,
+    project_option_dec,
+)
+from flytekit.interfaces import cli_identifiers
+from flytekit.remote import FlyteRemote
+
+EXECUTION_ID = "execution_id"
+
+
+@click.command("relaunch", help="Relaunch a failed execution")
+@click.pass_context
+def relaunch(ctx: click.Context):
+    """
+    Relaunches an execution.
+    """
+    project = ctx.obj[CTX_PROJECT]
+    domain = ctx.obj[CTX_DOMAIN]
+    remote: FlyteRemote = ctx.obj[FLYTE_REMOTE_INSTANCE_KEY]
+    execution_identifier = cli_identifiers.WorkflowExecutionIdentifier(
+        project=project, domain=domain, name=ctx.obj[EXECUTION_ID]
+    )
+    execution_identifier_resp = remote.client.relaunch_execution(id=execution_identifier)
+    execution_identifier = cli_identifiers.WorkflowExecutionIdentifier.promote_from_model(execution_identifier_resp)
+    click.secho("Launched execution: {}".format(execution_identifier), fg="blue")
+    click.echo("")
+
+
+@click.command("recover", help="Recover a failed execution")
+@click.pass_context
+def recover(ctx: click.Context):
+    """
+    Recovers an execution.
+    """
+    project = ctx.obj[CTX_PROJECT]
+    domain = ctx.obj[CTX_DOMAIN]
+    remote: FlyteRemote = ctx.obj[FLYTE_REMOTE_INSTANCE_KEY]
+    execution_identifier = cli_identifiers.WorkflowExecutionIdentifier(
+        project=project, domain=domain, name=ctx.obj[EXECUTION_ID]
+    )
+    execution_identifier_resp = remote.client.recover_execution(id=execution_identifier)
+    execution_identifier = cli_identifiers.WorkflowExecutionIdentifier.promote_from_model(execution_identifier_resp)
+    click.secho("Launched execution: {}".format(execution_identifier), fg="blue")
+    click.echo("")
+
+
+execution_help = """
+The execution command allows you to interact with Flyte's execution system,
+such as recovering/relaunching a failed execution.
+"""
+
+
+@click.group("execution", help=execution_help)
+@project_option_dec
+@domain_option_dec
+@click.option(
+    EXECUTION_ID,
+    "--execution-id",
+    required=True,
+    type=str,
+    help="The execution id",
+)
+@click.pass_context
+def execute(
+    ctx: click.Context,
+    project: str,
+    domain: str,
+    execution_id: str,
+):
+    # save remote instance in ctx.obj
+    get_and_save_remote_with_click_context(ctx, project, domain)
+    if ctx.obj is None:
+        ctx.obj = {}
+
+    ctx.obj.update(ctx.params)
+
+
+execute.add_command(recover)
+execute.add_command(relaunch)

flytekit/clis/sdk_in_container/pyflyte.py

@@ -7,6 +7,7 @@
 from flytekit.clis.sdk_in_container.backfill import backfill
 from flytekit.clis.sdk_in_container.build import build
 from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_PACKAGES, CTX_VERBOSE
+from flytekit.clis.sdk_in_container.executions import execute
 from flytekit.clis.sdk_in_container.fetch import fetch
 from flytekit.clis.sdk_in_container.get import get
 from flytekit.clis.sdk_in_container.init import init
@@ -94,6 +95,7 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: int):
 main.add_command(fetch)
 main.add_command(info)
 main.add_command(get)
+main.add_command(execute)
 main.epilog
 
 get_plugin().configure_pyflyte_cli(main)

flytekit/clis/sdk_in_container/run.py

@@ -242,7 +242,6 @@ class RunLevelParams(PyFlyteParams):
             required=False,
             type=int,
             default=50,
-            hidden=True,
             show_default=True,
             help="Use this to limit number of entities to fetch",
         )
@@ -661,14 +660,6 @@ def __init__(self, command_name: str):
         super().__init__(
             name=command_name,
             help=f"Retrieve {command_name} from a remote flyte instance and execute them.",
-            params=[
-                click.Option(
-                    ["--limit", "limit"],
-                    help=f"Limit the number of {command_name}'s to retrieve.",
-                    default=50,
-                    show_default=True,
-                )
-            ],
         )
         self._command_name = command_name
         self._entities = []

flytekit/core/array_node_map_task.py

@@ -113,9 +113,9 @@ def python_interface(self):
 
     def construct_node_metadata(self) -> NodeMetadata:
         # TODO: add support for other Flyte entities
-        return NodeMetadata(
-            name=self.name,
-        )
+        nm = super().construct_node_metadata()
+        nm._name = self.name
+        return nm
 
     @property
     def min_success_ratio(self) -> Optional[float]:

flytekit/core/artifact.py

@@ -15,6 +15,7 @@
 from flytekit.loggers import logger
 
 TIME_PARTITION_KWARG = "time_partition"
+MAX_PARTITIONS = 10
 
 
 class InputsBase(object):
@@ -337,6 +338,9 @@ def __init__(
             self._partitions = Partitions(p)
             self._partitions.set_reference_artifact(self)
 
+        if self.partition_keys and len(self.partition_keys) > MAX_PARTITIONS:
+            raise ValueError("There is a hard limit of 10 partition keys per artifact currently.")
+
     def __call__(self, *args, **kwargs) -> ArtifactIDSpecification:
         """
         This __call__ should only ever happen in the context of a task or workflow's output, to be