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

flyte/__init__.py

@@ -1,32 +1,85 @@
 """
-Flyte SDK for authoring Compound AI applications, services and workflows.
+Flyte SDK for authoring compound AI applications, services and workflows.
+"""
 
-## Environments
+from __future__ import annotations
 
-TaskEnvironment class to define a new environment for a set of tasks.
+import sys
 
-Example usage:
+from ._build import build
+from ._cache import Cache, CachePolicy, CacheRequest
+from ._context import ctx
+from ._custom_context import custom_context, get_custom_context
+from ._deploy import build_images, deploy
+from ._environment import Environment
+from ._excepthook import custom_excepthook
+from ._group import group
+from ._image import Image
+from ._initialize import current_domain, init, init_from_api_key, init_from_config, init_in_cluster, init_passthrough
+from ._link import Link
+from ._logging import logger
+from ._map import map
+from ._pod import PodTemplate
+from ._resources import AMD_GPU, GPU, HABANA_GAUDI, TPU, Device, DeviceClass, Neuron, Resources
+from ._retry import RetryStrategy
+from ._reusable_environment import ReusePolicy
+from ._run import run, with_runcontext
+from ._secret import Secret, SecretRequest
+from ._serve import serve, with_servecontext
+from ._task_environment import TaskEnvironment
+from ._timeout import Timeout, TimeoutType
+from ._trace import trace
+from ._trigger import Cron, FixedRate, Trigger, TriggerTime
+from ._version import __version__
 
-```python
-env = flyte.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))
+sys.excepthook = custom_excepthook
+
+
+def _silence_grpc_warnings():
+    """
+    Silences gRPC warnings that can clutter the output.
+    """
+    import os
+
+    # Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+    # before importing grpc
+    if "GRPC_VERBOSITY" not in os.environ:
+        os.environ["GRPC_VERBOSITY"] = "ERROR"
+        os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+        # Disable fork support (stops "skipping fork() handlers")
+        os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+        # Reduce absl/glog verbosity
+        os.environ["GLOG_minloglevel"] = "2"
+        os.environ["ABSL_LOG"] = "0"
+
+
+_silence_grpc_warnings()
+
+
+def version() -> str:
+    """
+    Returns the version of the Flyte SDK.
+    """
+    return __version__
 
-@env.task
-async def my_task():
-    pass
-```
-"""
 
 __all__ = [
-    "ABFS",
-    "GCS",
+    "AMD_GPU",
     "GPU",
-    "S3",
+    "HABANA_GAUDI",
     "TPU",
     "Cache",
     "CachePolicy",
     "CacheRequest",
+    "Cron",
     "Device",
+    "DeviceClass",
+    "Environment",
+    "FixedRate",
     "Image",
+    "Link",
+    "Neuron",
+    "PodTemplate",
     "Resources",
     "RetryStrategy",
     "ReusePolicy",
@@ -35,28 +88,28 @@ __all__ = [
     "TaskEnvironment",
     "Timeout",
     "TimeoutType",
+    "Trigger",
+    "TriggerTime",
     "__version__",
+    "build",
+    "build_images",
     "ctx",
+    "current_domain",
+    "custom_context",
     "deploy",
+    "get_custom_context",
     "group",
     "init",
+    "init_from_api_key",
+    "init_from_config",
+    "init_in_cluster",
+    "init_passthrough",
+    "logger",
+    "map",
     "run",
+    "serve",
     "trace",
+    "version",
     "with_runcontext",
+    "with_servecontext",
 ]
-
-from ._cache import Cache, CachePolicy, CacheRequest
-from ._context import ctx
-from ._deploy import deploy
-from ._group import group
-from ._image import Image
-from ._initialize import ABFS, GCS, S3, init
-from ._resources import GPU, TPU, Device, Resources
-from ._retry import RetryStrategy
-from ._reusable_environment import ReusePolicy
-from ._run import run, with_runcontext
-from ._secret import Secret, SecretRequest
-from ._task_environment import TaskEnvironment
-from ._timeout import Timeout, TimeoutType
-from ._trace import trace
-from ._version import __version__

flyte/_bin/connect.py

@@ -0,0 +1,61 @@
+from typing import List
+
+import click
+
+
+@click.group()
+def _connect():
+    """Debug commands for Flyte."""
+
+
+@_connect.command()
+@click.option(
+    "--port",
+    default="8000",
+    is_flag=False,
+    type=int,
+    help="Grpc port for the connector service. Defaults to 8000",
+)
+@click.option(
+    "--prometheus_port",
+    default="9090",
+    is_flag=False,
+    type=int,
+    help="Prometheus port for the connector service. Defaults to 9090",
+)
+@click.option(
+    "--worker",
+    default="10",
+    is_flag=False,
+    type=int,
+    help="Number of workers for the grpc server",
+)
+@click.option(
+    "--timeout",
+    default=None,
+    is_flag=False,
+    type=int,
+    help="It will wait for the specified number of seconds before shutting down grpc server. It should only be used "
+    "for testing.",
+)
+@click.option(
+    "--modules",
+    required=False,
+    multiple=True,
+    type=str,
+    help="List of additional files or module that defines the connector",
+)
+@click.pass_context
+def main(
+    _: click.Context, port: int, prometheus_port: int, worker: int, timeout: int | None, modules: List[str] | None
+):
+    """
+    Start a grpc server for the connector service.
+    """
+    from flyte.connectors import ConnectorService
+
+    ConnectorService.run(port, prometheus_port, worker, timeout, modules)
+
+
+if __name__ == "__main__":
+    _connect()

flyte/_bin/debug.py

@@ -0,0 +1,38 @@
+import click
+
+
+@click.group()
+def _debug():
+    """Debug commands for Flyte."""
+
+
+@_debug.command("resume")
+@click.option("--pid", "-m", type=int, required=True, help="PID of the vscode server.")
+def resume(pid):
+    """
+    Resume a Flyte task for debugging purposes.
+
+    Args:
+        pid (int): PID of the vscode server.
+    """
+    import os
+    import signal
+
+    print("Terminating server and resuming task.")
+    answer = (
+        input(
+            "This operation will kill the server. All unsaved data will be lost,"
+            " and you will no longer be able to connect to it. Do you really want to terminate? (Y/N): "
+        )
+        .strip()
+        .upper()
+    )
+    if answer == "Y":
+        os.kill(pid, signal.SIGTERM)
+        print("The server has been terminated and the task has been resumed.")
+    else:
+        print("Operation canceled.")
+
+
+if __name__ == "__main__":
+    _debug()

flyte/_bin/runtime.py

@@ -12,23 +12,28 @@ from typing import List
 
 import click
 
-# Todo: work with pvditt to make these the names
-# ACTION_NAME = "_U_ACTION_NAME"
-# RUN_NAME = "_U_RUN_NAME"
-# PROJECT_NAME = "_U_PROJECT_NAME"
-# DOMAIN_NAME = "_U_DOMAIN_NAME"
-# ORG_NAME = "_U_ORG_NAME"
+from flyte.models import PathRewrite
 
 ACTION_NAME = "ACTION_NAME"
 RUN_NAME = "RUN_NAME"
-PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
-DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
+INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
+FLYTE_ENABLE_VSCODE_KEY = "_F_E_VS"
 
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+_F_PATH_REWRITE = "_F_PATH_REWRITE"
 
-@click.command("a0")
+
+@click.group()
+def _pass_through():
+    pass
+
+
+@_pass_through.command("a0")
 @click.option("--inputs", "-i", required=True)
 @click.option("--outputs-path", "-o", required=True)
 @click.option("--version", "-v", required=True)
@@ -41,6 +46,8 @@ RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
 @click.option("--project", envvar=PROJECT_NAME, required=False)
 @click.option("--domain", envvar=DOMAIN_NAME, required=False)
 @click.option("--org", envvar=ORG_NAME, required=False)
+@click.option("--debug", envvar=FLYTE_ENABLE_VSCODE_KEY, type=click.BOOL, required=False)
+@click.option("--interactive-mode", type=click.BOOL, required=False)
 @click.option("--image-cache", required=False)
 @click.option("--tgz", required=False)
 @click.option("--pkl", required=False)
@@ -51,12 +58,16 @@ RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
     type=click.UNPROCESSED,
     nargs=-1,
 )
+@click.pass_context
 def main(
+    ctx: click.Context,
     run_name: str,
     name: str,
     project: str,
     domain: str,
     org: str,
+    debug: bool,
+    interactive_mode: bool,
     image_cache: str,
     version: str,
     inputs: str,
@@ -73,13 +84,24 @@ def main(
 ):
     sys.path.insert(0, ".")
 
+    import faulthandler
+    import signal
+
+    import flyte
     import flyte._utils as utils
-    from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
-    from flyte._initialize import S3, initialize_in_cluster
+    import flyte.errors
+    import flyte.storage as storage
+    from flyte._initialize import init_in_cluster
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte._logging import logger
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
+
+    logger.info("Registering faulthandler for SIGUSR1")
+    faulthandler.register(signal.SIGUSR1)
 
+    logger.info(f"Initializing flyte runtime - version {flyte.__version__}")
     assert org, "Org is required for now"
     assert project, "Project is required"
     assert domain, "Domain is required"
@@ -91,21 +113,42 @@ def main(
     if name.startswith("{{"):
         name = os.getenv("ACTION_NAME", "")
 
-    ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+    logger.warning(f"Flyte runtime started for action {name} with run name {run_name}")
+
+    if debug and name == "a0":
+        from flyte._debug.vscode import _start_vscode_server
 
-    bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
-    # TODO configure storage correctly for cluster
-    initialize_in_cluster(storage=S3.auto())
-    controller = create_controller(ct="remote", endpoint=ep, insecure=True)
+        asyncio.run(_start_vscode_server(ctx))
+
+    bundle = None
+    if tgz or pkl:
+        bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
+
+    controller_kwargs = init_in_cluster(org=org, project=project, domain=domain)
+    # Controller is created with the same kwargs as init, so that it can be used to run tasks
+    controller = create_controller(ct="remote", **controller_kwargs)
 
     ic = ImageCache.from_transport(image_cache) if image_cache else None
 
+    path_rewrite_cfg = os.getenv(_F_PATH_REWRITE, None)
+    path_rewrite = None
+    if path_rewrite_cfg:
+        potential_path_rewrite = PathRewrite.from_str(path_rewrite_cfg)
+        if storage.exists_sync(potential_path_rewrite.new_prefix):
+            path_rewrite = potential_path_rewrite
+            logger.info(f"Path rewrite configured for {path_rewrite.new_prefix}")
+        else:
+            logger.error(
+                f"Path rewrite failed for path {potential_path_rewrite.new_prefix}, "
+                f"not found, reverting to original path {potential_path_rewrite.old_prefix}"
+            )
+
     # Create a coroutine to load the task and run it
     task_coroutine = load_and_run_task(
         resolver=resolver,
         resolver_args=resolver_args,
         action=ActionID(name=name, run_name=run_name, project=project, domain=domain, org=org),
-        raw_data_path=RawDataPath(path=raw_data_path),
+        raw_data_path=RawDataPath(path=raw_data_path, path_rewrite=path_rewrite),
         checkpoints=Checkpoints(checkpoint_path, prev_checkpoint),
         code_bundle=bundle,
         input_path=inputs,
@@ -114,13 +157,38 @@ def main(
         version=version,
         controller=controller,
         image_cache=ic,
+        interactive_mode=interactive_mode or debug,
     )
     # Create a coroutine to watch for errors
     controller_failure = controller.watch_for_errors()
 
     # Run both coroutines concurrently and wait for first to finish and cancel the other
     async def _run_and_stop():
-        await utils.run_coros(controller_failure, task_coroutine)
-        await controller.stop()
+        loop = asyncio.get_event_loop()
+        loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
+        try:
+            await utils.run_coros(controller_failure, task_coroutine)
+            await controller.stop()
+        except flyte.errors.RuntimeSystemError as e:
+            logger.error(f"Runtime system error: {e}")
+            from flyte._internal.runtime.convert import convert_from_native_to_error
+            from flyte._internal.runtime.io import upload_error
+
+            logger.error(f"Flyte runtime failed for action {name} with run name {run_name}, error: {e}")
+            err = convert_from_native_to_error(e)
+            path = await upload_error(err.err, outputs_path)
+            logger.error(f"Run {run_name} Action {name} failed with error: {err}. Uploaded error to {path}")
+            await controller.stop()
 
     asyncio.run(_run_and_stop())
+    logger.warning(f"Flyte runtime completed for action {name} with run name {run_name}")
+    for h in logger.handlers:
+        h.flush()
+    sys.stdout.flush()
+    # We os._exit here to ensure that grpc does not block the exiting! grpc currently has a graceful shutdown system
+    # that blocks the process from exiting
+    os._exit(0)
+
+
+if __name__ == "__main__":
+    _pass_through()