diff --git a/runongpu/cli.py b/runongpu/cli.py index c4d5744..82a3a4e 100644 --- a/runongpu/cli.py +++ b/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: diff --git a/runongpu/colab.py b/runongpu/colab.py index 4dfa45e..3ef1074 100644 --- a/runongpu/colab.py +++ b/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]") \ No newline at end of file + 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 + ) \ No newline at end of file