Feature/8 add restart option

runongpu/cli.py

@@ -16,6 +16,7 @@
 
 @app.command()
 def doctor():
+
     # Verifies that the user's local RunOnGPU setup has the required pieces installed.
     console.print("[bold cyan]Checking RunOnGPU setup...[/bold cyan]")
 
@@ -87,7 +88,7 @@ def run():
         return
 
     try:
-        console.print("[yellow]Parsing runongpu.txt. [/yellow]")
+        console.print("[yellow] Parsing runongpu.txt. [/yellow]")
         project_config = parse_config()
         console.print("[green] Successfully parsed through runongpu.txt")
     except ValueError as error:

runongpu/colab.py

@@ -5,8 +5,10 @@
 from urllib.error import URLError
 from urllib.request import urlopen
 
-from playwright.sync_api import sync_playwright
-
+from playwright.sync_api import (
+    TimeoutError as PlaywrightTimeoutError,
+    sync_playwright,
+)
 from runongpu.config import load_config
 
 from rich.console import Console
@@ -15,22 +17,6 @@
 console = Console()
 
 
-
-#Helper function to get another port incase one port fails
-
-def get_debug_port() -> int:
-    preferred_ports = [9223, 9224, 9225, 9226, 9227, 9228] # Start from 9223
-    for port in preferred_ports:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            result = sock.connect_ex(('127.0.0.1', port))
-            if result != 0:
-                return port
-
-    #if all ports are busy, let OS pick a random number
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.bind(('127.0.0.1', 0))
-        return sock.getsockname()[1]
-
 # Path to the real Chrome executable on Windows.
 # RunOnGPU uses real Chrome because Colab/Google login is more reliable there
 # than in Playwright's bundled browser.
@@ -55,6 +41,20 @@ def get_debug_port() -> int:
 TEMPLATE_URL = "https://colab.research.google.com/drive/1pB8iVjR4-tPVSEBFjY8ow6N_F34bcMwi?usp=sharing"
 
 
+#Helper function to get another port incase one port fails
+def get_debug_port() -> int:
+    preferred_ports = [9223, 9224, 9225, 9226, 9227, 9228] # Start from 9223
+    for port in preferred_ports:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            result = sock.connect_ex(('127.0.0.1', port))
+            if result != 0:
+                return port
+
+    #if all ports are busy, let OS pick a random number
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(('127.0.0.1', 0))
+        return sock.getsockname()[1]
+
 def wait_for_debug_port(port: int, timeout_seconds: int = 15) -> None:
     """Wait until Chrome is ready for Playwright to connect."""
     start_time = time.time()
@@ -71,7 +71,7 @@ def wait_for_debug_port(port: int, timeout_seconds: int = 15) -> None:
 
     raise RuntimeError(
         f"Chrome did not open remote debugging port {port}. "
-        "Close Chrome and try again, or use a different debug port."
+        f"Trying another port. Please give it some time..."
     )
 
 
@@ -117,33 +117,39 @@ def open_colab(notebook_url: str = "", project_config: dict | None = None) -> st
 
         if not notebook_url:
             while True:
-                # Saving a copy usually opens a new Colab tab. expect_page captures
-                # that tab directly instead of guessing with context.pages[-1].
-                with context.expect_page() as new_page_info:
-                    page.get_by_role("button", name="File", exact=True).click()
-                    page.get_by_text("Save a copy in Drive").click()
-
-                # Switch automation to the copied notebook tab.
-                page = new_page_info.value
-                page.wait_for_load_state("domcontentloaded")
+                try:
+                    # Saving a copy usually opens a new Colab tab. expect_page captures
+                    # that tab directly
+                    with context.expect_page(timeout=5_000) as new_page_info:
+                        page.get_by_role("button", name="File", exact=True).click()
+                        page.get_by_text("Save a copy in Drive").click()
+
+                    # Switch automation to the copied notebook tab.
+                    page = new_page_info.value
+                    page.wait_for_load_state("domcontentloaded")
+
+                except PlaywrightTimeoutError:
+                    # No new tab opened so colab may hae shown a sign-in popup
+                    pass
+                if check_google_sign_in(page):
+                    input("Please sign into Colab, then press Enter to try again. If already signed in, press enter...")
+                    page.goto(target_url)
+                    page.wait_for_load_state("domcontentloaded")
+                    continue
 
                 copied_successfully = (
                     page.url != target_url
-                    and "accounts.google.com" not in page.url
+                    and "colab.research.google.com" in page.url
                 )
 
                 if copied_successfully:
                     break
-
-                input("Please sign into Colab, then press Enter to try again. If already signed in, press enter...")
-
         current_url = page.url
 
         if project_config is None:
             raise RuntimeError("project_config was not passed into open_colab().")
 
         # project_config contains the parsed setup/build/test/run commands from runongpu.txt.
-        console.print(f"[cyan]project_config:[/cyan] {project_config}")
         console.print("[cyan]Writing RunOnGPU cell...[/cyan]")
         write_runongpu_cell(page, project_config)
 
@@ -224,4 +230,25 @@ def set_t4_gpu_and_run_all(page) -> None:
     # Start the notebook after the runtime is configured.
     console.print("[cyan]Running all cells...[/cyan]")
     page.keyboard.press("Control+F9")
-    console.print("[green]✓ Runtime set and cells started[/green]")
+    console.print("[green]✓ Runtime set and cells started[/green]")
+
+
+# Deals with two options for logging in to Goolge Account
+# Helps implement restart option
+def check_google_sign_in(page) -> bool:
+
+    # Case 1: the page itself redirected to Google's login site.
+    if "accounts.google.com" in page.url:
+        return True
+
+    # Case 2: Colab shows a sign-in dialog or button on the current page.
+    sign_in_elements = [
+        page.get_by_role("button", name="Sign in"),
+        page.get_by_text("Sign in to continue"),
+        page.get_by_text("Sign in with Google"),
+    ]
+
+    return any(
+        locator.first.is_visible()
+        for locator in sign_in_elements
+    )