Fix

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

src/nemo_run/core/execution/slurm.py

@@ -42,8 +42,14 @@
 from nemo_run.core.packaging.base import Packager
 from nemo_run.core.packaging.git import GitArchivePackager
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
-from nemo_run.core.tunnel.callback import Callback
-from nemo_run.core.tunnel.client import LocalTunnel, PackagingJob, SSHConfigFile, SSHTunnel, Tunnel
+from nemo_run.core.tunnel.client import (
+    Callback,
+    LocalTunnel,
+    PackagingJob,
+    SSHConfigFile,
+    SSHTunnel,
+    Tunnel,
+)
 from nemo_run.core.tunnel.server import TunnelMetadata, server_dir
 from nemo_run.devspace.base import DevSpace
 

src/nemo_run/core/tunnel/callback.py

@@ -24,7 +24,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import Callable, Optional
 
 import paramiko
 import paramiko.ssh_exception
@@ -35,9 +35,6 @@
 from nemo_run.config import NEMORUN_HOME, ConfigurableMixin
 from nemo_run.core.frontend.console.api import CONSOLE
 
-if TYPE_CHECKING:
-    from nemo_run.core.tunnel.callback import Callback
-
 logger: logging.Logger = logging.getLogger(__name__)
 TUNNEL_DIR = ".tunnels"
 TUNNEL_FILE_SUBPATH = os.path.join(NEMORUN_HOME, TUNNEL_DIR)
@@ -383,3 +380,29 @@ def remove_entry(self, name: str):
                     file.writelines(lines)
 
             print(f"Removed SSH config entry for {host}.")
+
+
+class Callback:
+    def setup(self, tunnel: "Tunnel"):
+        """Called when the tunnel is setup."""
+        self.tunnel = tunnel
+
+    def on_start(self):
+        """Called when the keep_alive loop starts."""
+        pass
+
+    def on_interval(self):
+        """Called at each interval during the keep_alive loop."""
+        pass
+
+    def on_stop(self):
+        """Called when the keep_alive loop stops."""
+        pass
+
+    def on_error(self, error: Exception):
+        """Called when an error occurs during the keep_alive loop.
+
+        Args:
+            error (Exception): The exception that was raised.
+        """
+        pass

src/nemo_run/core/tunnel/client.py

@@ -19,7 +19,7 @@
 import fiddle as fdl
 
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
-from nemo_run.core.tunnel.callback import Callback
+from nemo_run.core.tunnel.client import Callback
 
 if TYPE_CHECKING:
     from nemo_run.core.execution.base import Executor