ob-metaflow-extensions 1.1.66__py2.py3-none-any.whl → 1.1.68__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -240,3 +240,5 @@ class ObpGcpAuthProvider(object):
 
 
 GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]
+
+FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]

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

@@ -1,7 +1,7 @@
 from functools import partial
 from metaflow.decorators import FlowDecorator
 from metaflow import current
-from .nim import NimManager
+from .nim_manager import NimManager
 
 
 class NimDecorator(FlowDecorator):
@@ -55,6 +55,3 @@ class NimDecorator(FlowDecorator):
                 )
             }
         )
-
-
-FLOW_DECORATORS_DESC = [("nim", ".NimDecorator")]

metaflow_extensions/outerbounds/plugins/nim/nim_manager.py

@@ -0,0 +1,189 @@
+import os
+import time
+import json
+import requests
+from urllib.parse import urlparse
+from metaflow.metaflow_config import SERVICE_URL
+from metaflow.metaflow_config_funcs import init_config
+
+NVCF_URL = "https://api.nvcf.nvidia.com"
+NVCF_SUBMIT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/functions"
+NVCF_RESULT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/status"
+
+COMMON_HEADERS = {"accept": "application/json", "Content-Type": "application/json"}
+POLL_INTERVAL = 1
+
+
+class NimMetadata(object):
+    def __init__(self):
+        self._nvcf_chat_completion_models = []
+        self._coreweave_chat_completion_models = []
+
+        conf = init_config()
+
+        if "OBP_AUTH_SERVER" in conf:
+            auth_host = conf["OBP_AUTH_SERVER"]
+        else:
+            auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
+
+        nim_info_url = "https://" + auth_host + "/generate/nim"
+
+        if "METAFLOW_SERVICE_AUTH_KEY" in conf:
+            headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
+            res = requests.get(nim_info_url, headers=headers)
+        else:
+            headers = json.loads(os.environ.get("METAFLOW_SERVICE_HEADERS"))
+            res = requests.get(nim_info_url, headers=headers)
+
+        res.raise_for_status()
+        self._ngc_api_key = res.json()["nvcf"]["api_key"]
+
+        for model in res.json()["nvcf"]["functions"]:
+            self._nvcf_chat_completion_models.append(
+                {
+                    "name": model["model_key"],
+                    "function-id": model["id"],
+                    "version-id": model["version"],
+                }
+            )
+        for model in res.json()["coreweave"]["containers"]:
+            self._coreweave_chat_completion_models.append(
+                {"name": model["nim_name"], "ip-address": model["ip_addr"]}
+            )
+
+    def get_nvcf_chat_completion_models(self):
+        return self._nvcf_chat_completion_models
+
+    def get_coreweave_chat_completion_models(self):
+        return self._coreweave_chat_completion_models
+
+    def get_headers_for_nvcf_request(self):
+        return {**COMMON_HEADERS, "Authorization": f"Bearer {self._ngc_api_key}"}
+
+    def get_headers_for_coreweave_request(self):
+        return COMMON_HEADERS
+
+
+class NimManager(object):
+    def __init__(self, models, backend):
+        nim_metadata = NimMetadata()
+        if backend == "managed":
+            nvcf_models = [
+                m["name"] for m in nim_metadata.get_nvcf_chat_completion_models()
+            ]
+            cw_models = [
+                m["name"] for m in nim_metadata.get_coreweave_chat_completion_models()
+            ]
+
+            self.models = {}
+            for m in models:
+                if m in nvcf_models:
+                    self.models[m] = NimChatCompletion(
+                        model=m, provider="NVCF", nim_metadata=nim_metadata
+                    )
+                elif m in cw_models:
+                    self.models[m] = NimChatCompletion(
+                        model=m, provider="CoreWeave", nim_metadata=nim_metadata
+                    )
+                else:
+                    raise ValueError(
+                        f"Model {m} not supported by the Outerbounds @nim offering."
+                        f"\nYou can choose from these options: {nvcf_models + cw_models}\n\n"
+                        "Reach out to Outerbounds if there are other models you'd like supported."
+                    )
+        else:
+            raise ValueError(
+                f"Backend {backend} not supported by the Outerbounds @nim offering. Please reach out to Outerbounds."
+            )
+
+
+class NimChatCompletion(object):
+    def __init__(
+        self,
+        model="meta/llama3-8b-instruct",
+        provider="CoreWeave",
+        nim_metadata=None,
+        **kwargs,
+    ):
+        if nim_metadata is None:
+            raise ValueError(
+                "NimMetadata object is required to initialize NimChatCompletion object."
+            )
+
+        self._nim_metadata = nim_metadata
+        self.compute_provider = provider
+        self.invocations = []
+
+        if self.compute_provider == "CoreWeave":
+            cw_model_names = [
+                m["name"]
+                for m in self._nim_metadata.get_coreweave_chat_completion_models()
+            ]
+            self.model = model
+            self.ip_address = self._nim_metadata.get_coreweave_chat_completion_models()[
+                cw_model_names.index(model)
+            ]["ip-address"]
+            self.endpoint = f"http://{self.ip_address}:8000/v1/chat/completions"
+
+        elif self.compute_provider == "NVCF":
+            nvcf_model_names = [
+                m["name"] for m in self._nim_metadata.get_nvcf_chat_completion_models()
+            ]
+            self.model = model
+            self.function_id = self._nim_metadata.get_nvcf_chat_completion_models()[
+                nvcf_model_names.index(model)
+            ]["function-id"]
+            self.version_id = self._nim_metadata.get_nvcf_chat_completion_models()[
+                nvcf_model_names.index(model)
+            ]["version-id"]
+
+    def __call__(self, **kwargs):
+
+        if self.compute_provider == "CoreWeave":
+            request_data = {"model": self.model, **kwargs}
+            response = requests.post(
+                self.endpoint,
+                headers=self._nim_metadata.get_headers_for_coreweave_request(),
+                json=request_data,
+            )
+            response.raise_for_status()
+            return response.json()
+
+        elif self.compute_provider == "NVCF":
+
+            request_data = {"model": self.model, **kwargs}
+            request_url = f"{NVCF_SUBMIT_ENDPOINT}/{self.function_id}"
+
+            response = requests.post(
+                request_url,
+                headers=self._nim_metadata.get_headers_for_nvcf_request(),
+                json=request_data,
+            )
+            response.raise_for_status()
+            if response.status_code == 202:
+                invocation_id = response.headers.get("NVCF-REQID")
+                self.invocations.append(invocation_id)
+            elif response.status_code == 200:
+                return response.json()
+
+            def _poll():
+                poll_request_url = f"{NVCF_RESULT_ENDPOINT}/{invocation_id}"
+                poll_response = requests.get(
+                    poll_request_url,
+                    headers=self._nim_metadata.get_headers_for_nvcf_request(),
+                )
+                poll_response.raise_for_status()
+                if poll_response.status_code == 200:
+                    return poll_response.json()
+                elif poll_response.status_code == 202:
+                    return 202
+                else:
+                    raise Exception(
+                        f"NVCF returned {poll_response.status_code} status code. Please contact Outerbounds."
+                    )
+
+            while True:
+                data = _poll()
+                if data and data != 202:
+                    return data
+                time.sleep(POLL_INTERVAL)

{ob_metaflow_extensions-1.1.66.dist-info → ob_metaflow_extensions-1.1.68.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.66.dist-info → ob_metaflow_extensions-1.1.68.dist-info}/RECORD

@@ -1,15 +1,13 @@
-metaflow_extensions/nim/plugins/__init__.py,sha256=P4zm938Ie4KWxEVKGF1ZQZfKMr5KPP7uAJh7mNT4ra0,1587
-metaflow_extensions/nim/plugins/nim.py,sha256=8DWBQPN3NHak3E5zTO-evIPDsfPluUP6smGNXk8Q6F0,5432
-metaflow_extensions/nim/toplevel/mfextinit_nim.py,sha256=-5aNJ1ffBDfAoWwQKleAiB5jHMWLnVz44EUtavaUYK8,26
-metaflow_extensions/nim/toplevel/nim_toplevel.py,sha256=7lZ5qYSN0Dk-qm4IZ9iwzJx2jE1iqxWNrMaiujGzLgA,48
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=HPFH4e3ZK3p-wS5HlS75fhR8_2avdD1AHQIZl2KnjeQ,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=mYo95obHU1IE1wbPkeVz_pfTzNqlNabp1QBEMTGllbE,112
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=rK-EPg7107su6dY-yAx1IK3wAzyvQIIg6y4um_F_BXc,9343
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=oR3krG5x3-W4g1sm5ygNPe9KVLBmxg7KEtzJsonQo_4,9398
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=JhlMFcR7SPSfR1C9w6GlqJq-NYNhOfISmHl2PdkYUok,2212
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=z8tSAkWtiITB-JtSQS7fkhlBwvxSxeTgEwFjahAzv-U,2238
 metaflow_extensions/outerbounds/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/plugins/kubernetes/kubernetes_client.py,sha256=gj6Iaz26bGbZm3aQuNS18Mqh_80iJp5PgFwFSlJRcn8,1968
+metaflow_extensions/outerbounds/plugins/nim/__init__.py,sha256=wzeYo7TX6ueIifGMt4lmJ6_LEsrUYNubWqACgHLwTis,1543
+metaflow_extensions/outerbounds/plugins/nim/nim_manager.py,sha256=l8WDfVtsMt7aZaOaeIPT5ySidxfxXU8gmwLoKUP3f04,7044
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
 metaflow_extensions/outerbounds/profilers/gpu.py,sha256=a5YZAepujuP0uDqG9UpXBlZS3wjUt4Yv8CjybXqeT2c,24342
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
@@ -17,7 +15,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
-ob_metaflow_extensions-1.1.66.dist-info/METADATA,sha256=0IWe2LQCZOEkBzEtSBei8G8hbJFWfN8lgtHgueolyvs,519
-ob_metaflow_extensions-1.1.66.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.66.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.66.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.68.dist-info/METADATA,sha256=SJgAqyqfY268HCnd9Nn2XnGkRi918pmao-fFdQnp7tc,519
+ob_metaflow_extensions-1.1.68.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.68.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.68.dist-info/RECORD,,

metaflow_extensions/nim/plugins/nim.py

@@ -1,137 +0,0 @@
-import os
-import time
-import json
-import requests
-from urllib.parse import urlparse
-from metaflow.metaflow_config import SERVICE_URL
-from metaflow.metaflow_config_funcs import init_config
-
-auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
-nim_info_url = "https://" + auth_host + "/generate/nim"
-
-NVCF_URL = "https://api.nvcf.nvidia.com"
-NVCF_SUBMIT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/functions"
-NVCF_RESULT_ENDPOINT = f"{NVCF_URL}/v2/nvcf/pexec/status"
-
-NVCF_CHAT_COMPLETION_MODELS = []
-COREWEAVE_CHAT_COMPLETION_MODELS = []
-
-conf = init_config()
-if "METAFLOW_SERVICE_AUTH_KEY" in conf:
-    headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
-    res = requests.get(nim_info_url, headers=headers)
-else:
-    headers = json.loads(os.environ.get("METAFLOW_SERVICE_HEADERS"))
-    res = requests.get(nim_info_url, headers=headers)
-
-res.raise_for_status()
-NGC_API_KEY = res.json()["nvcf"]["api_key"]
-for model in res.json()["nvcf"]["functions"]:
-    NVCF_CHAT_COMPLETION_MODELS.append(
-        {
-            "name": model["model_key"],
-            "function-id": model["id"],
-            "version-id": model["version"],
-        }
-    )
-for model in res.json()["coreweave"]["containers"]:
-    COREWEAVE_CHAT_COMPLETION_MODELS.append(
-        {"name": model["nim_name"], "ip-address": model["ip_addr"]}
-    )
-
-COMMON_HEADERS = {"accept": "application/json", "Content-Type": "application/json"}
-COREWEAVE_HEADERS = COMMON_HEADERS
-NVCF_HEADERS = {**COMMON_HEADERS, "Authorization": f"Bearer {NGC_API_KEY}"}
-POLL_INTERVAL = 1
-
-
-class NimManager(object):
-    def __init__(self, models, backend):
-        if backend == "managed":
-            nvcf_models = [m["name"] for m in NVCF_CHAT_COMPLETION_MODELS]
-            cw_models = [m["name"] for m in COREWEAVE_CHAT_COMPLETION_MODELS]
-            for m in models:
-                if m not in nvcf_models and m not in cw_models:
-                    raise ValueError(
-                        f"Model {m} not supported by the Outerbounds @nim offering."
-                        f"\nYou can choose from these options: {nvcf_models + cw_models}\n\n"
-                        "Reach out to Outerbounds if there are other models you'd like supported."
-                    )
-            self.models = {}
-            for m in models:
-                if m in nvcf_models:
-                    self.models[m] = NimChatCompletion(model=m, provider="NVCF")
-                elif m in cw_models:
-                    self.models[m] = NimChatCompletion(model=m, provider="CoreWeave")
-        else:
-            raise ValueError(
-                f"Backend {backend} not supported by the Outerbounds @nim offering. Please reach out to Outerbounds."
-            )
-
-
-class NimChatCompletion(object):
-    def __init__(self, model="meta/llama3-8b-instruct", provider="CoreWeave", **kwargs):
-        self.compute_provider = provider
-        self.invocations = []
-
-        if self.compute_provider == "CoreWeave":
-            cw_model_names = [m["name"] for m in COREWEAVE_CHAT_COMPLETION_MODELS]
-            self.model = model
-            self.ip_address = COREWEAVE_CHAT_COMPLETION_MODELS[
-                cw_model_names.index(model)
-            ]["ip-address"]
-            self.endpoint = f"http://{self.ip_address}:8000/v1/chat/completions"
-
-        elif self.compute_provider == "NVCF":
-            nvcf_model_names = [m["name"] for m in NVCF_CHAT_COMPLETION_MODELS]
-            self.model = model
-            self.function_id = NVCF_CHAT_COMPLETION_MODELS[
-                nvcf_model_names.index(model)
-            ]["function-id"]
-            self.version_id = NVCF_CHAT_COMPLETION_MODELS[
-                nvcf_model_names.index(model)
-            ]["version-id"]
-
-    def __call__(self, **kwargs):
-
-        if self.compute_provider == "CoreWeave":
-            request_data = {"model": self.model, **kwargs}
-            response = requests.post(
-                self.endpoint, headers=COREWEAVE_HEADERS, json=request_data
-            )
-            response.raise_for_status()
-            return response.json()
-
-        elif self.compute_provider == "NVCF":
-
-            request_data = {"model": self.model, **kwargs}
-            request_url = f"{NVCF_SUBMIT_ENDPOINT}/{self.function_id}"
-
-            response = requests.post(
-                request_url, headers=NVCF_HEADERS, json=request_data
-            )
-            response.raise_for_status()
-            if response.status_code == 202:
-                invocation_id = response.headers.get("NVCF-REQID")
-                self.invocations.append(invocation_id)
-            elif response.status_code == 200:
-                return response.json()
-
-            def _poll():
-                poll_request_url = f"{NVCF_RESULT_ENDPOINT}/{invocation_id}"
-                poll_response = requests.get(poll_request_url, headers=NVCF_HEADERS)
-                poll_response.raise_for_status()
-                if poll_response.status_code == 200:
-                    return poll_response.json()
-                elif poll_response.status_code == 202:
-                    return 202
-                else:
-                    raise Exception(
-                        f"NVCF returned {poll_response.status_code} status code. Please contact Outerbounds."
-                    )
-
-            while True:
-                data = _poll()
-                if data and data != 202:
-                    return data
-                time.sleep(POLL_INTERVAL)

metaflow_extensions/nim/toplevel/mfextinit_nim.py

@@ -1 +0,0 @@
-toplevel = "nim_toplevel"

metaflow_extensions/nim/toplevel/nim_toplevel.py

@@ -1,2 +0,0 @@
-__mf_extensions__ = "nim"
-__version__ = "0.0.1"