ob-metaflow-extensions 1.1.171rc1__py2.py3-none-any.whl → 1.4.39__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py

@@ -10,14 +10,15 @@ from metaflow.exception import MetaflowException
 class SnowparkClient(object):
     def __init__(
         self,
-        account: str,
-        user: str,
-        password: str,
-        role: str,
-        database: str,
-        warehouse: str,
-        schema: str,
+        account: str = None,
+        user: str = None,
+        password: str = None,
+        role: str = None,
+        database: str = None,
+        warehouse: str = None,
+        schema: str = None,
         autocommit: bool = True,
+        integration: str = None,
     ):
         try:
             from snowflake.core import Root
@@ -27,22 +28,48 @@ 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
             )
 
+        if integration:
+            # Use OAuth authentication via Outerbounds integration
+            from metaflow_extensions.outerbounds.plugins.snowflake.snowflake import (
+                get_oauth_connection_params,
+            )
+
+            self.connection_parameters = get_oauth_connection_params(
+                user=user or "",
+                role=role or "",
+                integration=integration,
+                schema=schema or "",
+                account=account,
+                warehouse=warehouse,
+                database=database,
+            )
+            self.connection_parameters["autocommit"] = autocommit
+        else:
+            # Password-based authentication
+            self.connection_parameters = {
+                "account": account,
+                "user": user,
+                "password": password,
+                "role": role,
+                "warehouse": warehouse,
+                "database": database,
+                "schema": schema,
+                "autocommit": autocommit,
+            }
+
+        # Remove None values from connection parameters
         self.connection_parameters = {
-            "account": account,
-            "user": user,
-            "password": password,
-            "role": role,
-            "warehouse": warehouse,
-            "database": database,
-            "schema": schema,
-            "autocommit": autocommit,
+            k: v for k, v in self.connection_parameters.items() if v is not None
         }
 
         try:

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"
             )
 
 
@@ -68,6 +71,7 @@ class SnowparkDecorator(StepDecorator):
         "cpu": None,
         "gpu": None,
         "memory": None,
+        "integration": None,  # Outerbounds OAuth integration name
     }
 
     package_url = None
@@ -77,12 +81,11 @@ class SnowparkDecorator(StepDecorator):
     def __init__(self, attributes=None, statically_defined=False):
         super(SnowparkDecorator, self).__init__(attributes, statically_defined)
 
+        # Set defaults from config (user can override via decorator or integration)
         if not self.attributes["account"]:
             self.attributes["account"] = SNOWPARK_ACCOUNT
         if not self.attributes["user"]:
             self.attributes["user"] = SNOWPARK_USER
-        if not self.attributes["password"]:
-            self.attributes["password"] = SNOWPARK_PASSWORD
         if not self.attributes["role"]:
             self.attributes["role"] = SNOWPARK_ROLE
         if not self.attributes["database"]:
@@ -91,6 +94,9 @@ class SnowparkDecorator(StepDecorator):
             self.attributes["warehouse"] = SNOWPARK_WAREHOUSE
         if not self.attributes["schema"]:
             self.attributes["schema"] = SNOWPARK_SCHEMA
+        # Only use password from config if not using integration (OAuth)
+        if not self.attributes["integration"] and not self.attributes["password"]:
+            self.attributes["password"] = SNOWPARK_PASSWORD
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:
@@ -143,9 +149,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

@@ -12,9 +12,9 @@ from .snowpark_exceptions import SnowparkException
 mapping = str.maketrans("0123456789", "abcdefghij")
 
 
-# keep only alpha numeric characters and underscores..
+# keep only alpha numeric characters and dashes..
 def sanitize_name(job_name: str):
-    return "".join(char for char in job_name if char.isalnum() or char == "_")
+    return "".join(char for char in job_name if char.isalnum() or char == "-")
 
 
 # this is not a decorator since the exception imports need to be inside
@@ -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

@@ -13,12 +13,14 @@ 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:
@@ -37,6 +39,8 @@ class TorchTune:
                     "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}"
                 )

metaflow_extensions/outerbounds/plugins/vllm/__init__.py

@@ -1,17 +1,30 @@
 from metaflow.decorators import StepDecorator
 from metaflow import current
 import functools
-import os
+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 VLLMManager
+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.
@@ -40,11 +53,23 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
         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.
-    kwargs : Any
-        Any other keyword arguments are passed directly to the vLLM engine.
-        This allows for flexible configuration of vLLM server settings.
+    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`.
     """
 
@@ -52,9 +77,12 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
     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": {},
     }
 
@@ -72,106 +100,156 @@ class VLLMDecorator(StepDecorator, CardDecoratorInjector):
                 f"Example: @vllm(model='meta-llama/Llama-3.2-1B')"
             )
 
-        # Attach the vllm status card
-        self.attach_card_decorator(
-            flow,
-            step_name,
-            "vllm_status",
-            "blank",
-            refresh_interval=self.attributes["card_refresh_interval"],
-        )
+        # 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():
-            self.vllm_manager = None
-            self.status_card = None
-            self.card_monitor_thread = None
+            # 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)
 
-            try:
-                self.status_card = VLLMStatusCard(
-                    refresh_interval=self.attributes["card_refresh_interval"]
-                )
+        return vllm_wrapper
 
-                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
+    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",
+                    )
                 )
-                self.card_monitor_thread._stop_event = False
-                self.card_monitor_thread.start()
-                self.vllm_manager = VLLMManager(
-                    model=self.attributes["model"],
-                    backend=self.attributes["backend"],
-                    debug=self.attributes["debug"],
-                    status_card=self.status_card,
-                    stream_logs_to_card=self.attributes["stream_logs_to_card"],
-                    **self.attributes["engine_args"],
+            )
+
+            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"
                 )
-                if self.attributes["debug"]:
-                    print("[@vllm] VLLMManager initialized.")
+        finally:
+            if self.vllm_manager:
+                self.vllm_manager.terminate_models()
 
-            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 VLLMManager: {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.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("[@vllm] Card monitoring thread stopped.")
+                        print(f"[@vllm] Final card update error: {e}")
+                time.sleep(2)
 
-        return vllm_wrapper
+            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/status_card.py

@@ -160,10 +160,10 @@ class VLLMStatusCard(CardRefresher):
             # Keep only last 10 events
             self.status_data["events"] = self.status_data["events"][:10]
 
-    def get_circuit_breaker_emoji(self, state):
-        """Get status emoji for circuit breaker state"""
-        emoji_map = {"CLOSED": "🟢", "OPEN": "🔴", "HALF_OPEN": "🟡"}
-        return emoji_map.get(state, "⚪")
+    # def get_circuit_breaker_emoji(self, state):
+    #     """Get status emoji for circuit breaker state"""
+    #     emoji_map = {"CLOSED": "🟢", "OPEN": "🔴", "HALF_OPEN": "🟡"}
+    #     return emoji_map.get(state, "⚪")
 
     def get_uptime_string(self, start_time):
         """Calculate uptime string"""
@@ -252,11 +252,11 @@ class VLLMStatusCard(CardRefresher):
             )
 
         # Simplified monitoring note
-        current_card.append(
-            Markdown(
-                "## 🔧 Monitoring\n**Advanced Features:** Disabled (Circuit Breaker, Request Interception)"
-            )
-        )
+        # current_card.append(
+        #     Markdown(
+        #         "## 🔧 Monitoring\n**Advanced Features:** Disabled (Circuit Breaker, Request Interception)"
+        #     )
+        # )
 
         # Performance metrics
         perf_data = data["performance"]