Merge branch 'master' into fg91/feat/kf-pytorch-reasonable-defaults

.github/workflows/pythonbuild.yml

@@ -316,6 +316,7 @@ jobs:
           - flytekit-aws-batch
           - flytekit-aws-sagemaker
           - flytekit-bigquery
+          - flytekit-comet-ml
           - flytekit-dask
           - flytekit-data-fsspec
           - flytekit-dbt

.gitignore

@@ -35,6 +35,7 @@ docs/source/_tags/
 .hypothesis
 .npm
 /**/target
+coverage.xml
 
 # Version file is auto-generated by setuptools_scm
 flytekit/_version.py

CODEOWNERS

@@ -1,3 +1,4 @@
 # These owners will be the default owners for everything in
 # the repo. Unless a later match takes precedence.
-*       @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla
+*       @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla @thomasjpfan @future-outlier
+plugins/flytekit-kf-pytorch @fg91 @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla @thomasjpfan @future-outlier

Dockerfile

@@ -23,7 +23,6 @@ ARG DOCKER_IMAGE
 RUN apt-get update && apt-get install build-essential -y \
     && pip install uv \
     && uv pip install --system --no-cache-dir -U flytekit==$VERSION \
-        flytekitplugins-deck-standard==$VERSION \
     && apt-get clean autoclean \
     && apt-get autoremove --yes \
     && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

dev-requirements.in

@@ -20,7 +20,6 @@ IPython
 keyrings.alt
 setuptools_scm
 pytest-icdiff
-jinja2
 
 # Tensorflow is not available for python 3.12 yet: https://github.com/tensorflow/tensorflow/issues/62003
 tensorflow; python_version<'3.12'

dev-requirements.txt

@@ -44,7 +44,7 @@ botocore==1.34.106
     # via aiobotocore
 cachetools==5.3.3
     # via google-auth
-certifi==2024.2.2
+certifi==2024.7.4
     # via
     #   kubernetes
     #   requests
@@ -76,6 +76,7 @@ cryptography==42.0.7
     #   azure-storage-blob
     #   msal
     #   pyjwt
+    #   secretstorage
 dataclasses-json==0.5.9
     # via flytekit
 decorator==5.1.1
@@ -181,9 +182,7 @@ iniconfig==2.0.0
 ipython==8.25.0
     # via -r dev-requirements.in
 isodate==0.6.1
-    # via
-    #   azure-storage-blob
-    #   flytekit
+    # via azure-storage-blob
 jaraco-classes==3.4.0
     # via
     #   keyring
@@ -196,10 +195,10 @@ jaraco-functools==4.0.1
     # via keyring
 jedi==0.19.1
     # via ipython
-jinja2==3.1.4
+jeepney==0.8.0
     # via
-    #   -r dev-requirements.in
-    #   flytekit
+    #   keyring
+    #   secretstorage
 jmespath==1.0.1
     # via botocore
 joblib==1.4.2
@@ -221,8 +220,6 @@ markdown-it-py==3.0.0
     # via
     #   flytekit
     #   rich
-markupsafe==2.1.5
-    # via jinja2
 marshmallow==3.21.2
     # via
     #   dataclasses-json
@@ -314,6 +311,7 @@ proto-plus==1.23.0
     #   google-cloud-bigquery-storage
 protobuf==4.25.3
     # via
+    #   -r dev-requirements.in
     #   flyteidl
     #   flytekit
     #   google-api-core
@@ -423,6 +421,8 @@ scikit-learn==1.5.0
     # via -r dev-requirements.in
 scipy==1.13.1
     # via scikit-learn
+secretstorage==3.3.3
+    # via keyring
 setuptools-scm==8.1.0
     # via -r dev-requirements.in
 six==1.16.0
@@ -450,7 +450,7 @@ types-decorator==5.1.8.20240310
     # via -r dev-requirements.in
 types-mock==5.1.0.20240425
     # via -r dev-requirements.in
-types-protobuf==5.26.0.20240422
+types-protobuf==4.25.0.20240417
     # via -r dev-requirements.in
 types-requests==2.32.0.20240523
     # via -r dev-requirements.in

docs/source/_templates/file_types.rst

@@ -2,7 +2,7 @@
 
 .. currentmodule:: {{ module }}
 
-{% if objname == 'FlyteFile' %}
+{% if objname == 'FlyteFile' or  objname == 'FlyteDirectory' %}
 
 .. autoclass:: {{ objname }}
 

flytekit/_ast/parser.py

@@ -0,0 +1,25 @@
+import ast
+import inspect
+import typing
+
+
+def get_function_param_location(func: typing.Callable, param_name: str) -> (int, int):
+    """
+    Get the line and column number of the parameter in the source code of the function definition.
+    """
+    # Get source code of the function
+    source_lines, start_line = inspect.getsourcelines(func)
+    source_code = "".join(source_lines)
+
+    # Parse the source code into an AST
+    module = ast.parse(source_code)
+
+    # Traverse the AST to find the function definition
+    for node in ast.walk(module):
+        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
+            for i, arg in enumerate(node.args.args):
+                if arg.arg == param_name:
+                    # Calculate the line and column number of the parameter
+                    line_number = start_line + node.lineno - 1
+                    column_offset = arg.col_offset
+                    return line_number, column_offset

flytekit/clis/sdk_in_container/backfill.py

@@ -3,9 +3,9 @@
 
 import rich_click as click
 
-from flytekit import WorkflowFailurePolicy
 from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context
 from flytekit.clis.sdk_in_container.utils import domain_option_dec, project_option_dec
+from flytekit.core.workflow import WorkflowFailurePolicy
 from flytekit.interaction.click_types import DateTimeType, DurationParamType
 
 _backfill_help = """

flytekit/clis/sdk_in_container/register.py

@@ -179,24 +179,21 @@ def register(
     # Create and save FlyteRemote,
     remote = get_and_save_remote_with_click_context(ctx, project, domain, data_upload_location="flyte://data")
     click.secho(f"Registering against {remote.config.platform.endpoint}")
-    try:
-        repo.register(
-            project,
-            domain,
-            image_config,
-            output,
-            destination_dir,
-            service_account,
-            raw_data_prefix,
-            version,
-            deref_symlinks,
-            fast=not non_fast,
-            package_or_module=package_or_module,
-            remote=remote,
-            env=env,
-            dry_run=dry_run,
-            activate_launchplans=activate_launchplans,
-            skip_errors=skip_errors,
-        )
-    except Exception as e:
-        raise e
+    repo.register(
+        project,
+        domain,
+        image_config,
+        output,
+        destination_dir,
+        service_account,
+        raw_data_prefix,
+        version,
+        deref_symlinks,
+        fast=not non_fast,
+        package_or_module=package_or_module,
+        remote=remote,
+        env=env,
+        dry_run=dry_run,
+        activate_launchplans=activate_launchplans,
+        skip_errors=skip_errors,
+    )

flytekit/clis/sdk_in_container/run.py

@@ -7,10 +7,10 @@
 import tempfile
 import typing
 from dataclasses import dataclass, field, fields
-from typing import cast, get_args
+from typing import get_args
 
 import rich_click as click
-from dataclasses_json import DataClassJsonMixin
+from mashumaro.codecs.json import JSONEncoder
 from rich.progress import Progress
 
 from flytekit import Annotations, FlyteContext, FlyteContextManager, Labels, Literal
@@ -395,7 +395,8 @@ def to_click_option(
             if type(default_val) == dict or type(default_val) == list:
                 default_val = json.dumps(default_val)
             else:
-                default_val = cast(DataClassJsonMixin, default_val).to_json()
+                encoder = JSONEncoder(python_type)
+                default_val = encoder.encode(default_val)
         if literal_var.type.metadata:
             description_extra = f": {json.dumps(literal_var.type.metadata)}"
 
@@ -528,7 +529,8 @@ def _run(*args, **kwargs):
         # By the time we get to this function, all the loading has already happened
 
         run_level_params: RunLevelParams = ctx.obj
-        logger.debug(f"Running {entity.name} with {kwargs} and run_level_params {run_level_params}")
+        entity_type = "workflow" if isinstance(entity, PythonFunctionWorkflow) else "task"
+        logger.debug(f"Running {entity_type} {entity.name} with input {kwargs}")
 
         click.secho(f"Running Execution on {'Remote' if run_level_params.is_remote else 'local'}.", fg="cyan")
         try:

flytekit/clis/sdk_in_container/serve.py

@@ -8,6 +8,8 @@
     add_AsyncAgentServiceServicer_to_server,
     add_SyncAgentServiceServicer_to_server,
 )
+from rich.console import Console
+from rich.table import Table
 
 
 @click.group("serve")
@@ -55,8 +57,8 @@ def agent(_: click.Context, port, worker, timeout):
 async def _start_grpc_server(port: int, worker: int, timeout: int):
     from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService
 
+    click.secho("🚀 Starting the agent service...")
     _start_http_server()
-    click.secho("Starting the agent service...", fg="blue")
     print_agents_metadata()
 
     server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker))
@@ -75,7 +77,7 @@ def _start_http_server():
     try:
         from prometheus_client import start_http_server
 
-        click.secho("Starting up the server to expose the prometheus metrics...", fg="blue")
+        click.secho("Starting up the server to expose the prometheus metrics...")
         start_http_server(9090)
     except ImportError as e:
         click.secho(f"Failed to start the prometheus server with error {e}", fg="red")
@@ -104,7 +106,17 @@ def print_agents_metadata():
     from flytekit.extend.backend.base_agent import AgentRegistry
 
     agents = AgentRegistry.list_agents()
-    for agent in agents:
-        name = agent.name
-        metadata = [category.name for category in agent.supported_task_categories]
-        click.secho(f"Starting {name} that supports task categories {metadata}", fg="blue")
+
+    table = Table(title="Agent Metadata")
+    table.add_column("Agent Name", style="cyan", no_wrap=True)
+    table.add_column("Support Task Types", style="cyan")
+    table.add_column("Is Sync", style="green")
+
+    for a in agents:
+        categories = ""
+        for c in a.supported_task_categories:
+            categories += f"{c.name} (v{c.version}) "
+        table.add_row(a.name, categories, str(a.is_sync))
+
+    console = Console()
+    console.print(table)