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

metaflow_extensions/outerbounds/plugins/nim/{__init__.py → nim_decorator.py}

@@ -1,64 +1,31 @@
-from functools import partial
-from uuid import uuid4
-import os, time
-from metaflow.decorators import StepDecorator
+import os
+import time
 from metaflow import current
-
+from .utils import get_storage_path, NIM_MONITOR_LOCAL_STORAGE_ROOT
 from .nim_manager import NimManager
+from metaflow.decorators import StepDecorator
 from .card import NimMetricsRefresher
-from .utilities import get_storage_path, NIM_MONITOR_LOCAL_STORAGE_ROOT
-from ..card_utilities.async_cards import AsyncPeriodicRefresher
 from ..card_utilities.injector import CardDecoratorInjector
+from ..card_utilities.async_cards import AsyncPeriodicRefresher
 
 
 class NimDecorator(StepDecorator, CardDecoratorInjector):
-    """
-    This decorator is used to run NIM containers in Metaflow tasks as sidecars.
-
-    User code call
-    -----------
-    @nim(
-        models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
-        backend='managed'
-    )
-
-    Valid backend options
-    ---------------------
-    - 'managed': Outerbounds selects a compute provider based on the model.
-
-    Valid model options
-    ----------------
-        - 'meta/llama3-8b-instruct': 8B parameter model
-        - 'meta/llama3-70b-instruct': 70B parameter model
-        - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
-
-    Parameters
-    ----------
-    models: list[NIM]
-        List of NIM containers running models in sidecars.
-    backend: str
-        Compute provider to run the NIM container.
-    queue_timeout : int
-        Time to keep the job in NVCF's queue.
-    """
-
     name = "nim"
+
     defaults = {
         "models": [],
-        "backend": "managed",
         "monitor": True,
         "persist_db": False,
-        "queue_timeout": 5 * 24 * 3600,  # Default 5 days in seconds
     }
 
-    def step_init(
-        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
-    ):
-
+    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
+    # to understand where these functions are invoked in the lifecycle of a
+    # Metaflow flow.
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         if self.attributes["monitor"]:
             self.attach_card_decorator(
                 flow,
-                step_name,
+                step,
                 NimMetricsRefresher.CARD_ID,
                 "blank",
                 refresh_interval=4.0,
@@ -68,11 +35,9 @@ class NimDecorator(StepDecorator, CardDecoratorInjector):
             {
                 "nim": NimManager(
                     models=self.attributes["models"],
-                    backend=self.attributes["backend"],
                     flow=flow,
-                    step_name=step_name,
+                    step_name=step,
                     monitor=self.attributes["monitor"],
-                    queue_timeout=self.attributes["queue_timeout"],
                 )
             }
         )
@@ -81,15 +46,14 @@ class NimDecorator(StepDecorator, CardDecoratorInjector):
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
     ):
         if self.attributes["monitor"]:
-
             import sqlite3
-            from metaflow import current
 
             file_path = get_storage_path(current.task_id)
             if os.path.exists(file_path):
                 os.remove(file_path)
             os.makedirs(NIM_MONITOR_LOCAL_STORAGE_ROOT, exist_ok=True)
             conn = sqlite3.connect(file_path)
+
             cursor = conn.cursor()
             cursor.execute(
                 """