Merge branch 'master' into feature/array-node-workflow-parallelism

Signed-off-by: Paul Dittamo <[email protected]>

.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:
@@ -147,6 +149,41 @@ jobs:
           fail_ci_if_error: false
           files: coverage.xml
 
+  test-hypothesis:
+    needs:
+      - detect-python-versions
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache pip
+        uses: actions/cache@v3
+        with:
+          # This path is specific to Ubuntu
+          path: ~/.cache/pip
+          # Look to see if there is a cache hit for the corresponding requirements files
+          key: ${{ format('{0}-pip-{1}', runner.os, hashFiles('dev-requirements.in', 'requirements.in')) }}
+      - name: Install dependencies
+        run: make setup && pip freeze
+      - name: Test with coverage
+        env:
+          FLYTEKIT_HYPOTHESIS_PROFILE: ci
+        run: |
+          make unit_test_hypothesis
+      - name: Codecov
+        uses: codecov/[email protected]
+        with:
+          fail_ci_if_error: false
+          files: coverage.xml
+
   test-serialization:
     needs:
       - detect-python-versions
@@ -190,7 +227,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 +293,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
@@ -315,10 +350,6 @@ jobs:
             plugin-names: "flytekit-onnx-scikitlearn"
           - python-version: 3.11
             plugin-names: "flytekit-onnx-tensorflow"
-          # numba, a dependency of mlflow, doesn't support python 3.11
-          # https://github.com/numba/numba/issues/8304
-          - python-version: 3.11
-            plugin-names: "flytekit-mlflow"
           # vaex currently doesn't support python 3.11
           - python-version: 3.11
             plugin-names: "flytekit-vaex"
@@ -348,9 +379,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:
@@ -391,10 +422,10 @@ jobs:
     steps:
       - name: Fetch the code
         uses: actions/checkout@v4
-      - name: Set up Python 3.8
+      - name: Set up Python 3.12
         uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: 3.12
       - uses: actions/cache@v3
         with:
           path: ~/.cache/pip

.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}/ \

Makefile

@@ -62,10 +62,13 @@ unit_test_extras_codecov:
 unit_test:
 	# Skip all extra tests and run them with the necessary env var set so that a working (albeit slower)
 	# library is used to serialize/deserialize protobufs is used.
-	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
+	$(PYTEST_AND_OPTS) -m "not (serial or sandbox_test or hypothesis)" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 	# Run serial tests without any parallelism
 	$(PYTEST) -m "serial" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/ --ignore=tests/flytekit/unit/models --ignore=tests/flytekit/unit/extend ${CODECOV_OPTS}
 
+.PHONY: unit_test_hypothesis
+unit_test_hypothesis:
+	$(PYTEST_AND_OPTS) -m "hypothesis" tests/flytekit/unit/experimental ${CODECOV_OPTS}
 
 .PHONY: unit_test_extras
 unit_test_extras:

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/bin/entrypoint.py

@@ -249,7 +249,7 @@ def setup_execution(
             domain=exe_domain,
             name=exe_name,
         ),
-        execution_date=_datetime.datetime.utcnow(),
+        execution_date=_datetime.datetime.now(_datetime.timezone.utc),
         stats=_get_stats(
             cfg=StatsConfig.auto(),
             # Stats metric path will be:

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)