ob-metaflow-extensions 1.1.130__py2.py3-none-any.whl → 1.5.1__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,27 +76,39 @@ 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"]
 
 
-def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+def get_oauth_connection_params(
+    user: str = "", role: str = "", integration: str = "", **kwargs
+) -> Dict:
     """
-    Connect to snowflake using the token minted by Outerbounds
+    Get OAuth connection parameters for Snowflake authentication using Outerbounds integration.
+
+    This is a helper function that returns connection parameters dict that can be used
+    with both snowflake-connector-python and snowflake-snowpark-python.
+
     user: str
         The user name used to authenticate with snowflake
     role: str
-        The role to request when connect with snowflake
+        The role to request when connecting with snowflake
     integration: str
-        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+        The name of the snowflake integration to use. If not set, an existing integration
+        will be used provided that only one exists in the current perimeter.
     kwargs: dict
-        Additional arguments to pass to the python snowflake connector
+        Additional arguments to include in the connection parameters
+
+    Returns:
+        Dict with connection parameters including OAuth token
     """
     # ensure password is not set
     if "password" in kwargs:
         raise OuterboundsSnowflakeConnectorException(
-            "Password should not be set when using Outerbounds snowflake connector."
+            "Password should not be set when using Outerbounds OAuth authentication."
         )
 
     provisioner = SnowflakeIntegrationProvisioner(integration)
@@ -134,11 +147,31 @@ def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
     kwargs["role"] = role
     kwargs["user"] = user
 
+    return kwargs
+
+
+def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+    """
+    Connect to snowflake using the token minted by Outerbounds
+    user: str
+        The user name used to authenticate with snowflake
+    role: str
+        The role to request when connect with snowflake
+    integration: str
+        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+    kwargs: dict
+        Additional arguments to pass to the python snowflake connector
+    """
+    # Get OAuth connection params using the helper
+    connection_params = get_oauth_connection_params(
+        user=user, role=role, integration=integration, **kwargs
+    )
+
     # connect to snowflake
     try:
         from snowflake.connector import connect
 
-        cn = connect(**kwargs)
+        cn = connect(**connection_params)
         return cn
     except ImportError as ie:
         raise OuterboundsSnowflakeConnectorException(
@@ -150,6 +183,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 +339,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 +349,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.py

@@ -1,9 +1,11 @@
 import os
+import re
 import shlex
 import atexit
 import json
 import math
 import time
+import hashlib
 
 from metaflow import util
 
@@ -57,21 +59,29 @@ class Snowpark(object):
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
     def _job_name(self, user, flow_name, run_id, step_name, task_id, retry_count):
-        return "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
-            user=user,
-            flow_name=flow_name,
-            run_id=str(run_id) if run_id is not None else "",
-            step_name=step_name,
-            task_id=str(task_id) if task_id is not None else "",
-            retry_count=str(retry_count) if retry_count is not None else "",
+        unique_str = (
+            "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
+                user=user,
+                flow_name=flow_name,
+                run_id=str(run_id) if run_id is not None else "",
+                step_name=step_name,
+                task_id=str(task_id) if task_id is not None else "",
+                retry_count=str(retry_count) if retry_count is not None else "",
+            )
         )
+        unique_hash = hashlib.md5(unique_str.encode("utf-8")).hexdigest()[:8]
+        raw_prefix = f"{flow_name}-{step_name}"
+        safe_prefix = re.sub(r"[^a-z0-9]", "-", raw_prefix.lower())
+        safe_prefix = safe_prefix[:54]
+        safe_prefix = safe_prefix.lstrip("-")
+        return f"{safe_prefix}-{unique_hash}"
 
     def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
         mflog_expr = export_mflog_env_vars(
             datastore_type=self.datastore.TYPE,
             stdout_path=STDOUT_PATH,
             stderr_path=STDERR_PATH,
-            **task_spec
+            **task_spec,
         )
         init_cmds = environment.get_package_commands(
             code_package_url, self.datastore.TYPE

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py

@@ -66,6 +66,10 @@ def snowpark():
     "--schema",
     help="Schema for Snowpark Container Services.",
 )
+@click.option(
+    "--integration",
+    help="Outerbounds OAuth integration name for Snowpark Container Services. When set, uses OAuth authentication instead of password.",
+)
 @click.option(
     "--image",
     help="Docker image requirement for Snowpark Container Services. In name:version format.",
@@ -119,6 +123,7 @@ def step(
     database=None,
     warehouse=None,
     schema=None,
+    integration=None,
     image=None,
     stage=None,
     compute_pool=None,
@@ -235,6 +240,7 @@ def step(
                 "database": database,
                 "warehouse": warehouse,
                 "schema": schema,
+                "integration": integration,
             },
         )
         with ctx.obj.monitor.measure("metaflow.snowpark.launch_job"):

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

@@ -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)