Refactor/browser module

runongpu/browser.py

@@ -0,0 +1,55 @@
+import os
+import socket
+import time
+from pathlib import Path
+from urllib.error import URLError
+from urllib.request import urlopen
+
+# 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.
+CHROME_EXE = (
+    Path(os.environ["PROGRAMFILES"])
+    / "Google"
+    / "Chrome"
+    / "Application"
+    / "chrome.exe"
+)
+
+
+# Dedicated Chrome profile for RunOnGPU.
+# This keeps Colab login/session data persistent without touching the user's
+# everyday Chrome profile.
+RUNONGPU_PROFILE_DIR = Path.home() / ".runongpu" / "chrome-profile"
+
+
+def get_debug_port() -> int:
+    """Return an available port for Chrome remote debugging."""
+    preferred_ports = [9223, 9224, 9225, 9226, 9227, 9228]
+    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 preferred ports are busy, let the OS pick an available port.
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        #returns port number
+        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.monotonic()
+
+    while time.monotonic() - start_time < timeout_seconds:
+        try:
+            # Chrome uses this local endpoint after remote debugging starts.
+            with urlopen(f"http://127.0.0.1:{port}/json/version", timeout=1):
+                return
+        except URLError:
+            # Chrome can take a moment to launch, so retry briefly instead of failing immediately.
+            time.sleep(0.5)
+
+    raise RuntimeError(f"Chrome did not open remote debugging port {port}.")

runongpu/colab.py

@@ -1,117 +1,66 @@
-import os
 import subprocess
-import time
-from pathlib import Path
-from urllib.error import URLError
-from urllib.request import urlopen
 
 from playwright.sync_api import (
     TimeoutError as PlaywrightTimeoutError,
+)
+from playwright.sync_api import (
     sync_playwright,
 )
-from runongpu.config import load_config
-
 from rich.console import Console
-import socket
-
-console = Console()
-
-
-# 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.
-CHROME_EXE = (
-    Path(os.environ["PROGRAMFILES"])
-    / "Google"
-    / "Chrome"
-    / "Application"
-    / "chrome.exe"
+from runongpu.browser import (
+    CHROME_EXE,
+    RUNONGPU_PROFILE_DIR,
+    get_debug_port,
+    wait_for_debug_port,
 )
+from runongpu.config import load_config
 
-# Dedicated Chrome profile for RunOnGPU.
-# This keeps Colab login/session data persistent without touching the user's
-# everyday Chrome profile.
-RUNONGPU_PROFILE_DIR = Path.home() / ".runongpu" / "chrome-profile"
-
-# Local Chrome DevTools Protocol port.
-# Playwright connects to this port to control the real Chrome window.
-
-# Shared starter notebook used only when the user does not already have a saved
-# RunOnGPU notebook URL.
+console = Console()
+# Shared notebook used only when the user does not already have a saved.
+# Ensures the automation tool can correctly pick out locators without any issue
 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()
-
-    while time.time() - start_time < timeout_seconds:
-        try:
-            # Chrome exposes this local endpoint after remote debugging starts.
-            with urlopen(f"http://127.0.0.1:{port}/json/version", timeout=1):
-                return
-        except URLError:
-            # Chrome can take a moment to launch, so retry briefly instead of failing immediately.
-            time.sleep(0.5)
-
-
-    raise RuntimeError(
-        f"Chrome did not open remote debugging port {port}. "
-        f"Trying another port. Please give it some time..."
-    )
-
-
 def open_colab(notebook_url: str = "", project_config: dict | None = None) -> str:
     """Open a saved Colab notebook, or copy the template notebook on first run."""
-    DEBUG_PORT = 9222
+    debug_port = 9222
     target_url = notebook_url or TEMPLATE_URL
 
     # Launch real Chrome with remote debugging enabled so Playwright can attach.
     # The custom profile lets users sign into Colab once and reuse that session.
     try:
-        subprocess.Popen([
-            str(CHROME_EXE),
-            f"--remote-debugging-port={DEBUG_PORT}",
-            f"--user-data-dir={RUNONGPU_PROFILE_DIR}",
-            "--no-first-run",
-            "--no-default-browser-check",
-            target_url,
-        ])
+        subprocess.Popen(
+            [
+                str(CHROME_EXE),
+                f"--remote-debugging-port={debug_port}",
+                f"--user-data-dir={RUNONGPU_PROFILE_DIR}",
+                "--no-first-run",
+                "--no-default-browser-check",
+                target_url,
+            ]
+        )
 
         # Avoid connecting before Chrome has opened its debugging endpoint.
-        wait_for_debug_port(port=DEBUG_PORT)
+        wait_for_debug_port(port=debug_port)
+
     except RuntimeError:
-        DEBUG_PORT = get_debug_port()
-        subprocess.Popen([
-            str(CHROME_EXE),
-            f"--remote-debugging-port={DEBUG_PORT}",
-            f"--user-data-dir={RUNONGPU_PROFILE_DIR}",
-            "--no-first-run",
-            "--no-default-browser-check",
-            target_url,
-        ])
-        wait_for_debug_port(port=DEBUG_PORT)
+        debug_port = get_debug_port()
+        subprocess.Popen(
+            [
+                str(CHROME_EXE),
+                f"--remote-debugging-port={debug_port}",
+                f"--user-data-dir={RUNONGPU_PROFILE_DIR}",
+                "--no-first-run",
+                "--no-default-browser-check",
+                target_url,
+            ]
+        )
+        wait_for_debug_port(port=debug_port)
+
     with sync_playwright() as playwright:
         # Attach to the already-open real Chrome window instead of launching a new browser.
-        browser = playwright.chromium.connect_over_cdp(
-            f"http://127.0.0.1:{DEBUG_PORT}"
-        )
+        browser = playwright.chromium.connect_over_cdp(f"http://127.0.0.1:{debug_port}")
 
-        # Use the active browser context and newest tab opened by RunOnGPU.
         context = browser.contexts[0]
         page = context.pages[-1]
 
@@ -120,36 +69,33 @@ def open_colab(notebook_url: str = "", project_config: dict | None = None) -> st
                 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:
+                    with context.expect_page(timeout=7_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
+                    # No new tab opened so colab may have 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...")
+                    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 "colab.research.google.com" in page.url
+                    page.url != target_url and "colab.research.google.com" in page.url
                 )
 
                 if copied_successfully:
                     break
         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("[cyan]Writing RunOnGPU cell...[/cyan]")
         write_runongpu_cell(page, project_config)
 
@@ -164,6 +110,10 @@ def open_colab(notebook_url: str = "", project_config: dict | None = None) -> st
 
 
 def write_runongpu_cell(page, project_config: dict):
+    """
+    Paste the code that clones the configured repository and executes its setup,
+    build, test, and run commands from our config file.
+    """
     saved_config = load_config()
 
     # Close popups such as Gemini/Colab assistant before selecting the target cell.
@@ -192,7 +142,7 @@ def write_runongpu_cell(page, project_config: dict):
     folder_name = saved_config["folder_name"]
     repo_url = saved_config["repo_url"]
 
-    code =  f"""
+    code = f"""
 # Paste your code here:
 folder_name = "{folder_name}"
 github_repo_url = "{repo_url}"
@@ -208,35 +158,34 @@ def write_runongpu_cell(page, project_config: dict):
 
 
 def set_t4_gpu_and_run_all(page) -> None:
-    # Open Colab runtime settings so the notebook uses a GPU runtime.
+    """Configure the Colab runtime to use a T4 GPU and run all notebook cells."""
+
     console.print("[cyan]Opening Runtime menu...[/cyan]")
     page.get_by_role("button", name="Runtime", exact=True).click()
 
     console.print("[cyan]Opening Change runtime type...[/cyan]")
     page.get_by_role("menuitem", name="Change runtime type", exact=True).click()
 
-    # Select the free T4 GPU option available in Colab.
     console.print("[cyan]Selecting T4 GPU...[/cyan]")
     page.get_by_role("radio", name="T4 GPU").click()
 
     console.print("[cyan]Saving runtime settings...[/cyan]")
     page.get_by_role("button", name="Save").click()
 
-    # Give Colab time to close the runtime dialog before sending keyboard shortcuts.
     console.print("[cyan]Waiting for runtime dialog to close...[/cyan]")
     page.wait_for_timeout(3000)
     page.keyboard.press("Escape")
 
-    # 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]")
-
 
-# Deals with two options for logging in to Goolge Account
-# Helps implement restart option
+
 def check_google_sign_in(page) -> bool:
-
+    """
+    Return whether Colab redirected to or displayed a Google sign-in prompt.
+    """
+
     # Case 1: the page itself redirected to Google's login site.
     if "accounts.google.com" in page.url:
         return True
@@ -247,8 +196,5 @@ def check_google_sign_in(page) -> bool:
         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
-    )
+
+    return any(locator.first.is_visible(timeout=1000) for locator in sign_in_elements)