ob-metaflow-extensions 1.1.158__tar.gz → 1.1.160rc0__tar.gz

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.158
+Version: 1.1.160rc0
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -347,4 +347,4 @@ SECRETS_PROVIDERS_DESC = [
     ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
 ]
 # Adding an override here so the library can be imported at the metaflow.plugins level
-__mf_promote_submodules__ = ["snowflake", "ollama"]
+__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune"]

ob-metaflow-extensions-1.1.160rc0/metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -0,0 +1,110 @@
+import threading
+import time
+import sys
+from typing import Dict, Optional, Any, Callable
+from functools import partial
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import FAST_BAKERY_URL
+
+from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
+from .docker_environment import cache_request
+
+BAKERY_METAFILE = ".imagebakery-cache"
+
+
+class BakerException(MetaflowException):
+    headline = "Ran into an error while baking image"
+
+    def __init__(self, msg):
+        super(BakerException, self).__init__(msg)
+
+
+def bake_image(
+    cache_file_path: str,
+    ref: Optional[str] = None,
+    python: Optional[str] = None,
+    pypi_packages: Optional[Dict[str, str]] = None,
+    conda_packages: Optional[Dict[str, str]] = None,
+    base_image: Optional[str] = None,
+    logger: Optional[Callable[[str], Any]] = None,
+) -> FastBakeryApiResponse:
+    """
+    Bakes a Docker image with the specified dependencies.
+
+    Args:
+        cache_file_path: Path to the cache file
+        ref: Reference identifier for this bake (for logging purposes)
+        python: Python version to use
+        pypi_packages: Dictionary of PyPI packages and versions
+        conda_packages: Dictionary of Conda packages and versions
+        base_image: Base Docker image to use
+        logger: Optional logger function to output progress
+
+    Returns:
+        FastBakeryApiResponse: The response from the bakery service
+
+    Raises:
+        BakerException: If the baking process fails
+    """
+    # Default logger if none provided
+    if logger is None:
+        logger = partial(print, file=sys.stderr)
+
+    # Thread lock for logging
+    logger_lock = threading.Lock()
+    images_baked = 0
+
+    @cache_request(cache_file_path)
+    def _cached_bake(
+        ref=None,
+        python=None,
+        pypi_packages=None,
+        conda_packages=None,
+        base_image=None,
+    ):
+        try:
+            bakery = FastBakery(url=FAST_BAKERY_URL)
+            bakery._reset_payload()
+            bakery.python_version(python)
+            bakery.pypi_packages(pypi_packages)
+            bakery.conda_packages(conda_packages)
+            bakery.base_image(base_image)
+            # bakery.ignore_cache()
+
+            with logger_lock:
+                logger(f"🍳 Baking [{ref}] ...")
+                logger(f"     🐍 Python: {python}")
+
+                if pypi_packages:
+                    logger(f"     📦 PyPI packages:")
+                    for package, version in pypi_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                if conda_packages:
+                    logger(f"     📦 Conda packages:")
+                    for package, version in conda_packages.items():
+                        logger(f"        🔧 {package}: {version}")
+
+                logger(f"     🏗️  Base image: {base_image}")
+
+            start_time = time.time()
+            res = bakery.bake()
+            # TODO: Get actual bake time from bakery
+            bake_time = time.time() - start_time
+
+            with logger_lock:
+                logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
+            nonlocal images_baked
+            images_baked += 1
+            return res
+        except FastBakeryException as ex:
+            raise BakerException(f"Bake [{ref}] failed: {str(ex)}")
+
+    # Call the cached bake function with the provided parameters
+    return _cached_bake(
+        ref=ref,
+        python=python,
+        pypi_packages=pypi_packages,
+        conda_packages=conda_packages,
+        base_image=base_image,
+    )

ob-metaflow-extensions-1.1.160rc0/metaflow_extensions/outerbounds/plugins/nvcf/constants.py

@@ -0,0 +1,3 @@
+SUPPORTABLE_GPU_TYPES = ["L40", "L40S", "L40G", "H100", "NEBIUS_H100"]
+DEFAULT_GPU_TYPE = "H100"
+MAX_N_GPU_BY_TYPE = {"L40": 1, "L40S": 1, "L40G": 1, "H100": 4, "NEBIUS_H100": 8}

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/metaflow_extensions/outerbounds/plugins/nvct/nvct_decorator.py

@@ -65,6 +65,28 @@ SUPPORTABLE_GPU_TYPES = {
             "backend": "gcp-asia-se-1a",
         },
     ],
+    "NEBIUS_H100": [
+        {
+            "n_gpus": 1,
+            "instance_type": "ON-PREM.GPU.H100_1x",
+            "backend": "default-project-eu-north1",
+        },
+        {
+            "n_gpus": 2,
+            "instance_type": "ON-PREM.GPU.H100_2x",
+            "backend": "default-project-eu-north1",
+        },
+        {
+            "n_gpus": 4,
+            "instance_type": "ON-PREM.GPU.H100_4x",
+            "backend": "default-project-eu-north1",
+        },
+        {
+            "n_gpus": 8,
+            "instance_type": "ON-PREM.GPU.H100_8x",
+            "backend": "default-project-eu-north1",
+        },
+    ],
 }
 
 
@@ -154,6 +176,8 @@ class NvctDecorator(StepDecorator):
 
         self.attributes["instance_type"] = valid_config["instance_type"]
         self.attributes["gpu_type"] = requested_gpu_type
+        if self.attributes["gpu_type"] == "NEBIUS_H100":
+            self.attributes["gpu_type"] = "H100"
         self.attributes["backend"] = valid_config["backend"]
 
     def runtime_init(self, flow, graph, package, run_id):

ob-metaflow-extensions-1.1.160rc0/metaflow_extensions/outerbounds/plugins/torchtune/__init__.py

@@ -0,0 +1,159 @@
+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,
+    ):
+        """
+        Initialize the Tune launcher.
+
+        :param use_multi_node_config: If True, attempt to build a distributed configuration
+                                      from current.torch.torchrun_args.
+        """
+        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"],
+                }
+                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)

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -53,3 +53,4 @@ def S3(*args, **kwargs):
 from .. import profilers
 from ..plugins.snowflake import Snowflake
 from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
+from . import ob_internal

ob-metaflow-extensions-1.1.160rc0/metaflow_extensions/outerbounds/toplevel/ob_internal.py

@@ -0,0 +1 @@
+from ..plugins.fast_bakery.baker import bake_image

ob-metaflow-extensions-1.1.160rc0/metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.torchtune"]

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.158
+Version: 1.1.160rc0
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -19,6 +19,7 @@ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py
+metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py
@@ -63,15 +64,18 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py
 metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py
+metaflow_extensions/outerbounds/plugins/torchtune/__init__.py
 metaflow_extensions/outerbounds/profilers/__init__.py
 metaflow_extensions/outerbounds/profilers/gpu.py
 metaflow_extensions/outerbounds/toplevel/__init__.py
 metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py
+metaflow_extensions/outerbounds/toplevel/ob_internal.py
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py
+metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py
 ob_metaflow_extensions.egg-info/PKG-INFO
 ob_metaflow_extensions.egg-info/SOURCES.txt
 ob_metaflow_extensions.egg-info/dependency_links.txt

{ob-metaflow-extensions-1.1.158 → ob-metaflow-extensions-1.1.160rc0}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.158"
+version = "1.1.160rc0"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 

ob-metaflow-extensions-1.1.158/metaflow_extensions/outerbounds/plugins/nvcf/constants.py

@@ -1,3 +0,0 @@
-SUPPORTABLE_GPU_TYPES = ["L40", "L40S", "L40G", "H100"]
-DEFAULT_GPU_TYPE = "H100"
-MAX_N_GPU_BY_TYPE = {"L40": 1, "L40S": 1, "L40G": 1, "H100": 4}