ob-metaflow-extensions 1.1.96__tar.gz → 1.1.98__tar.gz

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/PKG-INFO

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

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -317,6 +317,7 @@ STEP_DECORATORS_DESC = [
         ".fast_bakery.fast_bakery_decorator.InternalFastBakeryDecorator",
     ),
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
+    ("tensorboard", ".tensorboard.TensorboardDecorator"),
 ]
 FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]
 TOGGLE_STEP_DECORATOR = [

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/metaflow_extensions/outerbounds/plugins/nvcf/nvcf.py

@@ -59,11 +59,11 @@ class Nvcf(object):
             code_package_url, code_package_ds
         )
         init_expr = " && ".join(init_cmds)
-        heartbeat_expr = f'python -m metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store "$MAIN_PID" {code_package_ds} nvcf_heartbeats 1>> $MFLOG_STDOUT 2>> $MFLOG_STDERR'
+        heartbeat_expr = f'python -m metaflow_extensions.outerbounds.plugins.nvcf.heartbeat_store "$MAIN_PID" {code_package_ds} nvcf_heartbeats & HEARTBEAT_PID=$!;'
         step_expr = bash_capture_logs(
             " && ".join(
                 self.environment.bootstrap_commands(step_name, code_package_ds)
-                + [step_cli + " & MAIN_PID=$!; " + heartbeat_expr]
+                + [step_cli + " & MAIN_PID=$!; " + heartbeat_expr + " wait $MAIN_PID"]
             )
         )
 
@@ -84,7 +84,10 @@ class Nvcf(object):
         #
         # Note that if step_expr OOMs, this tail expression is never executed.
         # We lose the last logs in this scenario.
-        cmd_str += "c=$?; %s; exit $c" % BASH_SAVE_LOGS
+        cmd_str += (
+            "c=$?; kill $HEARTBEAT_PID; wait $HEARTBEAT_PID; %s; exit $c"
+            % BASH_SAVE_LOGS
+        )
         cmd_str = (
             '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
             % cmd_str
@@ -222,10 +225,10 @@ class Job(object):
             self._invocation_id = response.headers.get("NVCF-REQID")
             if response.getcode() == 200:
                 data = json.loads(response.read())
-                if data["status"].startswith("Oops"):
-                    self._status = JobStatus.FAILED
-                else:
+                if data.get("exit_code") == 0:
                     self._status = JobStatus.SUCCESSFUL
+                else:
+                    self._status = JobStatus.FAILED
                 self._result = data
             elif response.getcode() == 202:
                 self._status = JobStatus.SUBMITTED
@@ -270,11 +273,11 @@ class Job(object):
             response = urlopen(request)
             if response.getcode() == 200:
                 data = json.loads(response.read())
-                if data["status"].startswith("Oops"):
-                    # TODO: Propagate the internal error forward
-                    self._status = JobStatus.FAILED
-                else:
+                # TODO: Propagate the internal error forward
+                if data.get("exit_code") == 0:
                     self._status = JobStatus.SUCCESSFUL
+                else:
+                    self._status = JobStatus.FAILED
                 self._result = data
             elif response.getcode() in [400, 500]:
                 self._status = JobStatus.FAILED

ob-metaflow-extensions-1.1.98/metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py

@@ -0,0 +1,50 @@
+import functools
+from metaflow.decorators import StepDecorator
+
+
+class TensorboardDecorator(StepDecorator):
+    name = "tensorboard"
+    defaults = {}
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def tb_wrapper():
+            import sys, os
+            from metaflow import metaflow_config, current
+
+            try:
+                from torch.utils.tensorboard import SummaryWriter
+            except:
+                print(
+                    "[@tensorboard] Torch and tensorboard not found - logging disabled!",
+                    file=sys.stderr,
+                )
+                step_func()
+            else:
+                tb_root = os.path.join(metaflow_config.DATATOOLS_S3ROOT, "tb")
+                pathspec = current.pathspec
+                try:
+                    log_dir = os.path.join(tb_root, current.project_flow_name, pathspec)
+                except:
+                    log_dir = os.path.join(tb_root, pathspec)
+                comps = log_dir[len(tb_root) + 1 :].split("/")
+                run_level = "/".join(comps[:-2])
+                flow_level = "/".join(comps[:-3])
+
+                print("[@tensorboard] -- INSPECTING RESULTS")
+                print(
+                    "[@tensorboard] -- Execute one of these commands on your workstation:"
+                )
+                print(f"[@tensorboard] Compare tasks of this run: obtb {run_level}")
+                print(f"[@tensorboard] Compare across runs: obtb {flow_level}")
+                writer = SummaryWriter(log_dir=log_dir)
+                setattr(flow, "obtb", writer)
+                try:
+                    step_func()
+                finally:
+                    writer.flush()
+                    delattr(flow, "obtb")
+
+        return tb_wrapper

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/ob_metaflow_extensions.egg-info/PKG-INFO

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

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -28,6 +28,7 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py
 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/profilers/__init__.py
 metaflow_extensions/outerbounds/profilers/gpu.py
 metaflow_extensions/outerbounds/toplevel/__init__.py

ob-metaflow-extensions-1.1.98/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.12.25.1

{ob-metaflow-extensions-1.1.96 → ob-metaflow-extensions-1.1.98}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.96"
+version = "1.1.98"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.12.24.1",
+        "ob-metaflow == 2.12.25.1",
     ],
 )

ob-metaflow-extensions-1.1.96/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.12.24.1