diff --git a/runongpu/browser.py b/runongpu/browser.py new file mode 100644 index 0000000..77848a0 --- /dev/null +++ b/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}.") diff --git a/runongpu/colab.py b/runongpu/colab.py index 3ef1074..5c17e41 100644 --- a/runongpu/colab.py +++ b/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 - ) \ No newline at end of file + + return any(locator.first.is_visible(timeout=1000) for locator in sign_in_elements)