ob-metaflow-extensions 1.1.142__py2.py3-none-any.whl → 1.4.33__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py

@@ -7,9 +7,10 @@ from metaflow.metaflow_config import SERVICE_URL
 from metaflow.metaflow_config_funcs import init_config
 from typing import Dict
 from os import environ
-
+import sys
 import json
 import requests
+import random
 import time
 
 
@@ -75,7 +76,9 @@ def get_snowflake_token(user: str = "", role: str = "", integration: str = "") -
     }
     json_payload = json.dumps(payload)
     headers = provisioner.get_service_auth_header()
-    response = requests.get(snowflake_token_url, data=json_payload, headers=headers)
+    response = _api_server_get(
+        snowflake_token_url, data=json_payload, headers=headers, conn_error_retries=5
+    )
     response.raise_for_status()
     return response.json()["token"]
 
@@ -150,6 +153,39 @@ def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
         )
 
 
+def _api_server_get(*args, conn_error_retries=2, **kwargs):
+    """
+    There are two categories of errors that we need to handle when dealing with any API server.
+    1. HTTP errors. These are are errors that are returned from the API server.
+        - How to handle retries for this case will be application specific.
+    2. Errors when the API server may not be reachable (DNS resolution / network issues)
+        - In this scenario, we know that something external to the API server is going wrong causing the issue.
+        - Failing pre-maturely in the case might not be the best course of action since critical user jobs might crash on intermittent issues.
+        - So in this case, we can just planely retry the request.
+
+    This function handles the second case. It's a simple wrapper to handle the retry logic for connection errors.
+    If this function is provided a `conn_error_retries` of 5, then the last retry will have waited 32 seconds.
+    Generally this is a safe enough number of retries after which we can assume that something is really broken. Until then,
+    there can be intermittent issues that would resolve themselves if we retry gracefully.
+    """
+    _num_retries = 0
+    noise = random.uniform(-0.5, 0.5)
+    while _num_retries < conn_error_retries:
+        try:
+            return requests.get(*args, **kwargs)
+        except requests.exceptions.ConnectionError:
+            if _num_retries <= conn_error_retries - 1:
+                # Exponential backoff with 2^(_num_retries+1) seconds
+                time.sleep((2 ** (_num_retries + 1)) + noise)
+                _num_retries += 1
+            else:
+                print(
+                    "[@snowflake] Failed to connect to the API server. ",
+                    file=sys.stderr,
+                )
+                raise
+
+
 class Snowflake:
     def __init__(
         self, user: str = "", role: str = "", integration: str = "", **kwargs
@@ -273,7 +309,9 @@ class SnowflakeIntegrationProvisioner:
         retryable_status_codes = [409]
         json_payload = json.dumps(payload)
         for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
-            response = requests.get(url, data=json_payload, headers=request_headers)
+            response = _api_server_get(
+                url, data=json_payload, headers=request_headers, conn_error_retries=5
+            )
             if response.status_code not in retryable_status_codes:
                 break
 
@@ -281,7 +319,9 @@ class SnowflakeIntegrationProvisioner:
                 sleep_time = 0.5 * (attempt + 1)
                 time.sleep(sleep_time)
 
-        response = requests.get(url, data=json_payload, headers=request_headers)
+        response = _api_server_get(
+            url, data=json_payload, headers=request_headers, conn_error_retries=5
+        )
         self._handle_error_response(response)
         return response.json()
 

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py

@@ -27,9 +27,12 @@ class SnowparkClient(object):
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
                 "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by executing - "
-                "%s -m pip install snowflake\n"
+                "Python packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\n"
                 "or equivalent through your favorite Python package manager."
                 % sys.executable
             )

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py

@@ -42,10 +42,13 @@ class Snowflake(object):
             return session
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
-                "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by using the @pypi decorator - "
-                "Eg: @pypi(packages={'snowflake': '0.11.0'})\n"
+                "Could not import module 'snowflake'.\n\n"
+                "Install required Snowflake packages using the @pypi decorator:\n"
+                "@pypi(packages={\n"
+                "    'snowflake': '1.8.0',\n"
+                "    'snowflake-connector-python': '3.18.0',\n"
+                "    'snowflake-snowpark-python': '1.40.0'\n"
+                "})\n"
             )
 
 
@@ -143,9 +146,12 @@ class SnowparkDecorator(StepDecorator):
         except (NameError, ImportError, ModuleNotFoundError):
             raise SnowflakeException(
                 "Could not import module 'snowflake'.\n\nInstall Snowflake "
-                "Python package (https://pypi.org/project/snowflake/) first.\n"
-                "You can install the module by executing - "
-                "%s -m pip install snowflake\n"
+                "Python packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\n"
                 "or equivalent through your favorite Python package manager."
                 % sys.executable
             )

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -199,11 +199,17 @@ class RunningJob(object):
 
     @property
     def status(self):
-        return self.status_obj()[0].get("status")
+        status_list = self.status_obj()
+        if not status_list:
+            return "UNKNOWN"
+        return status_list[0].get("status", "UNKNOWN")
 
     @property
     def message(self):
-        return self.status_obj()[0].get("message")
+        status_list = self.status_obj()
+        if not status_list:
+            return None
+        return status_list[0].get("message")
 
     @property
     def is_waiting(self):

metaflow_extensions/outerbounds/plugins/torchtune/__init__.py

@@ -0,0 +1,163 @@
+from queue import Queue, Empty
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional, List, Dict
+import subprocess
+import shutil
+import sys
+from metaflow import current
+
+__mf_promote_submodules__ = ["plugins.torchtune"]
+
+
+class TorchTune:
+    def __init__(
+        self,
+        use_multi_node_config: bool = False,
+        config_overrides: Optional[Dict] = None,
+    ):
+        """
+        Initialize the Tune launcher.
+
+        :param use_multi_node_config: If True, attempt to build a distributed configuration
+                                      from current.torch.torchrun_args.
+        :param config_overrides: Optional dictionary of config overrides for tune run.
+        """
+        self.multi_node_config = {}
+        if use_multi_node_config:
+            if getattr(current, "torch", None):
+                print(
+                    "[Metaflow Tune] Since @torchrun is used, multi-node config can be used to launch the job."
+                )
+                # For distributed torchtune launches, we use similar parameters as torchrun.
+                # (You might need to adjust the keys according to your environment.)
+                self.multi_node_config = {
+                    "nnodes": current.torch.torchrun_args["nnodes"],
+                    "master_addr": current.torch.torchrun_args["master_addr"],
+                    "master_port": int(current.torch.torchrun_args["master_port"]),
+                    "node_rank": current.torch.torchrun_args["node_rank"],
+                    "nproc_per_node": current.torch.torchrun_args["nproc_per_node"],
+                    "num_processes": current.torch.torchrun_args["nproc_per_node"]
+                    * current.torch.torchrun_args["nnodes"],
+                }
+                if config_overrides:
+                    self.multi_node_config.update(config_overrides)
+                print(
+                    f"[Metaflow Tune] Discovered multi-node config for torchrun: {self.multi_node_config}"
+                )
+            else:
+                print(
+                    "[Metaflow Tune] Since @torchrun is not used, default multi-node config cannot be used to launch the job."
+                )
+
+    def run(
+        self,
+        recipe: str,
+        config_dict: Dict,
+        additional_cli_options: Optional[List[str]] = None,
+    ):
+        """
+        Launch the torchtune job via its CLI.
+
+        :param recipe: The path to the recipe (or name of the recipe) to run.
+        :param config_dict: Optional dictionary that will be dumped to a YAML file and passed via --config.
+        :param additional_cli_options: Optional list of additional CLI options.
+        :raises: subprocess.CalledProcessError if the subprocess returns a nonzero exit code.
+        """
+        import yaml
+        import tempfile
+        import os
+
+        _temp_dir = tempfile.mkdtemp()
+        try:
+            config_path = os.path.join(_temp_dir, "config.yaml")
+            with open(config_path, "w") as f:
+                yaml.dump(config_dict, f)
+
+            additional_options = (
+                additional_cli_options if additional_cli_options else []
+            )
+
+            # Build the command. Here we use "tune run" as the base command.
+            cmd = ["tune", "run"]
+
+            # If distributed configuration is present, add torchrun–style flags.
+            if self.multi_node_config:
+                cmd.extend(
+                    [
+                        "--nnodes",
+                        str(self.multi_node_config.get("nnodes")),
+                        "--nproc-per-node",
+                        str(self.multi_node_config.get("nproc_per_node")),
+                        # "--rdzv_conf", f"rdzv_endpoint={self.multi_node_config.get('master_addr')}:{self.multi_node_config.get('master_port')}"
+                        "--rdzv-backend",
+                        "c10d",
+                        "--rdzv-endpoint",
+                        f"{self.multi_node_config.get('master_addr')}:{self.multi_node_config.get('master_port')}",
+                        "--rdzv-id",
+                        "1234567890",
+                        "--node-rank",
+                        str(self.multi_node_config.get("node_rank")),
+                        # TODO: should there be a masterip/port here ?
+                    ]
+                )
+
+            cmd.extend(additional_options)
+
+            cmd.append(recipe)
+            # If a recipe configuration was provided, pass it via the --config flag.
+            cmd.extend(["--config", config_path])
+
+            # Append any additional CLI options.
+
+            # Launch the subprocess.
+            print(f"[Metaflow tune] {' '.join(cmd)}")
+            process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                universal_newlines=True,
+            )
+
+            # Stream the output in real-time.
+            for out_line, err_line in read_popen_pipes(process):
+                print(out_line, end="", flush=True)
+                print(err_line, end="", file=sys.stderr, flush=True)
+
+            process.wait()
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(process.returncode, cmd)
+        finally:
+            shutil.rmtree(_temp_dir)
+
+
+def enqueue_output(file, queue):
+    for line in iter(file.readline, ""):
+        queue.put(line)
+    file.close()
+
+
+def read_popen_pipes(p):
+
+    with ThreadPoolExecutor(2) as pool:
+        q_stdout, q_stderr = Queue(), Queue()
+
+        pool.submit(enqueue_output, p.stdout, q_stdout)
+        pool.submit(enqueue_output, p.stderr, q_stderr)
+
+        while True:
+
+            if p.poll() is not None and q_stdout.empty() and q_stderr.empty():
+                break
+
+            out_line = err_line = ""
+
+            try:
+                out_line = q_stdout.get_nowait()
+            except Empty:
+                pass
+            try:
+                err_line = q_stderr.get_nowait()
+            except Empty:
+                pass
+
+            yield (out_line, err_line)

metaflow_extensions/outerbounds/plugins/vllm/__init__.py

@@ -0,0 +1,255 @@
+from metaflow.decorators import StepDecorator
+from metaflow import current
+import functools
+from enum import Enum
+import threading
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
+from metaflow.metaflow_config import from_conf
+
+from .vllm_manager import VLLMOpenAIManager, VLLMPyManager
+from .status_card import VLLMStatusCard, CardDecoratorInjector
+
+__mf_promote_submodules__ = ["plugins.vllm"]
+
+
+### The following classes are used to store the vLLM information in the current environment.
+# Then, Metaflow users can access the vLLM information through the current environment.
+class OpenAIAPIInfo:
+    def __init__(self, local_endpoint, local_api_key):
+        self.local_endpoint = local_endpoint
+        self.local_api_key = local_api_key
+
+
+class VLLM:
+    def __init__(self, llm):
+        self.llm = llm
+
+
+class VLLMDecorator(StepDecorator, CardDecoratorInjector):
+    """
+    This decorator is used to run vllm APIs as Metaflow task sidecars.
+
+    User code call
+    --------------
+    @vllm(
+        model="...",
+        ...
+    )
+
+    Valid backend options
+    ---------------------
+    - 'local': Run as a separate process on the local task machine.
+
+    Valid model options
+    -------------------
+    Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
+
+    NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
+    If you need multiple models, you must create multiple @vllm decorators.
+
+    Parameters
+    ----------
+    model: str
+        HuggingFace model identifier to be served by vLLM.
+    backend: str
+        Determines where and how to run the vLLM process.
+    openai_api_server: bool
+        Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
+        Default is False (uses native engine).
+        Set to True for backward compatibility with existing code.
+    debug: bool
+        Whether to turn on verbose debugging logs.
+    card_refresh_interval: int
+        Interval in seconds for refreshing the vLLM status card.
+        Only used when openai_api_server=True.
+    max_retries: int
+        Maximum number of retries checking for vLLM server startup.
+        Only used when openai_api_server=True.
+    retry_alert_frequency: int
+        Frequency of alert logs for vLLM server startup retries.
+        Only used when openai_api_server=True.
+    engine_args : dict
+        Additional keyword arguments to pass to the vLLM engine.
+        For example, `tensor_parallel_size=2`.
+    """
+
+    name = "vllm"
+    defaults = {
+        "model": None,
+        "backend": "local",
+        "openai_api_server": False,  # Default to native engine
+        "debug": False,
+        "stream_logs_to_card": False,
+        "card_refresh_interval": 10,
+        "max_retries": 60,
+        "retry_alert_frequency": 5,
+        "engine_args": {},
+    }
+
+    def step_init(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        super().step_init(
+            flow, graph, step_name, decorators, environment, flow_datastore, logger
+        )
+
+        # Validate that a model is specified
+        if not self.attributes["model"]:
+            raise ValueError(
+                f"@vllm decorator on step '{step_name}' requires a 'model' parameter. "
+                f"Example: @vllm(model='meta-llama/Llama-3.2-1B')"
+            )
+
+        # Attach the vllm status card only for API server mode
+        if self.attributes["openai_api_server"]:
+            self.attach_card_decorator(
+                flow,
+                step_name,
+                "vllm_status",
+                "blank",
+                refresh_interval=self.attributes["card_refresh_interval"],
+            )
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def vllm_wrapper():
+            # FIXME: Kind of ugly branch. Causing branching elsewhere.
+            # Other possibile code paths:
+            # - OpenAI batch API
+            # - Embedding
+            # - Special types of models
+            if self.attributes["openai_api_server"]:
+                # API Server mode (existing functionality)
+                self._run_api_server_mode(step_func)
+            else:
+                # Native engine mode (new functionality)
+                self._run_native_engine_mode(step_func)
+
+        return vllm_wrapper
+
+    def _run_api_server_mode(self, step_func):
+        """Run vLLM in API server mode (subprocess, existing functionality)"""
+        self.vllm_manager = None
+        self.status_card = None
+        self.card_monitor_thread = None
+
+        try:
+            self.status_card = VLLMStatusCard(
+                refresh_interval=self.attributes["card_refresh_interval"]
+            )
+
+            def monitor_card():
+                try:
+                    self.status_card.on_startup(current.card["vllm_status"])
+
+                    while not getattr(self.card_monitor_thread, "_stop_event", False):
+                        try:
+                            self.status_card.on_update(
+                                current.card["vllm_status"], None
+                            )
+                            import time
+
+                            time.sleep(self.attributes["card_refresh_interval"])
+                        except Exception as e:
+                            if self.attributes["debug"]:
+                                print(f"[@vllm] Card monitoring error: {e}")
+                            break
+                except Exception as e:
+                    if self.attributes["debug"]:
+                        print(f"[@vllm] Card monitor thread error: {e}")
+                    self.status_card.on_error(current.card["vllm_status"], str(e))
+
+            self.card_monitor_thread = threading.Thread(
+                target=monitor_card, daemon=True
+            )
+            self.card_monitor_thread._stop_event = False
+            self.card_monitor_thread.start()
+            self.vllm_manager = VLLMOpenAIManager(
+                model=self.attributes["model"],
+                backend=self.attributes["backend"],
+                debug=self.attributes["debug"],
+                status_card=self.status_card,
+                max_retries=self.attributes["max_retries"],
+                retry_alert_frequency=self.attributes["retry_alert_frequency"],
+                stream_logs_to_card=self.attributes["stream_logs_to_card"],
+                **self.attributes["engine_args"],
+            )
+            current._update_env(
+                dict(
+                    vllm=OpenAIAPIInfo(
+                        local_endpoint=f"http://127.0.0.1:{self.vllm_manager.port}/v1",
+                        local_api_key="token123",
+                    )
+                )
+            )
+
+            if self.attributes["debug"]:
+                print("[@vllm] API server mode initialized.")
+
+        except Exception as e:
+            if self.status_card:
+                self.status_card.add_event("error", f"Initialization failed: {str(e)}")
+                try:
+                    self.status_card.on_error(current.card["vllm_status"], str(e))
+                except:
+                    pass
+            print(f"[@vllm] Error initializing API server mode: {e}")
+            raise
+
+        try:
+            if self.status_card:
+                self.status_card.add_event("info", "Starting user step function")
+            step_func()
+            if self.status_card:
+                self.status_card.add_event(
+                    "success", "User step function completed successfully"
+                )
+        finally:
+            if self.vllm_manager:
+                self.vllm_manager.terminate_models()
+
+            if self.card_monitor_thread and self.status_card:
+                import time
+
+                try:
+                    self.status_card.on_update(current.card["vllm_status"], None)
+                except Exception as e:
+                    if self.attributes["debug"]:
+                        print(f"[@vllm] Final card update error: {e}")
+                time.sleep(2)
+
+            if self.card_monitor_thread:
+                self.card_monitor_thread._stop_event = True
+                self.card_monitor_thread.join(timeout=5)
+                if self.attributes["debug"]:
+                    print("[@vllm] Card monitoring thread stopped.")
+
+    def _run_native_engine_mode(self, step_func):
+        """Run vLLM in native engine mode (direct LLM API access)"""
+        self.vllm = None
+
+        try:
+            if self.attributes["debug"]:
+                print("[@vllm] Initializing native engine mode")
+
+            self.vllm = VLLMPyManager(
+                model=self.attributes["model"],
+                debug=self.attributes["debug"],
+                **self.attributes["engine_args"],
+            )
+            current._update_env(dict(vllm=VLLM(llm=self.vllm.engine)))
+
+            if self.attributes["debug"]:
+                print("[@vllm] Native engine mode initialized.")
+
+        except Exception as e:
+            print(f"[@vllm] Error initializing native engine mode: {e}")
+            raise
+
+        try:
+            step_func()
+        finally:
+            if self.vllm:
+                self.vllm.terminate_engine()

metaflow_extensions/outerbounds/plugins/vllm/constants.py

@@ -0,0 +1 @@
+VLLM_SUFFIX = "mf.vllm"

metaflow_extensions/outerbounds/plugins/vllm/exceptions.py

@@ -0,0 +1 @@
+from metaflow.exception import MetaflowException