gitarsenal-cli 1.1.2 → 1.1.4

package/python/gitarsenal.py

@@ -59,6 +59,55 @@ def check_modal_auth():
         print(f"⚠️ Error checking Modal authentication: {e}")
         return False
 
+def check_proxy_config():
+    """Check if Modal proxy configuration exists and is valid"""
+    try:
+        from gitarsenal_proxy_client import GitArsenalProxyClient
+        
+        # Initialize client
+        client = GitArsenalProxyClient()
+        
+        # Check if configuration exists
+        config = client.load_config()
+        
+        if not config.get("proxy_url") or not config.get("api_key"):
+            print("⚠️ Modal proxy not configured. Setting up with default values...")
+            
+            # Set default proxy URL to the ngrok URL
+            default_url = "https://e74889c63199.ngrok-free.app"
+            
+            # Update configuration with default URL
+            config["proxy_url"] = default_url
+            client.save_config(config)
+            client.base_url = default_url
+            
+            print(f"✅ Set default proxy URL: {default_url}")
+            
+            # If API key is missing, prompt for configuration
+            if not config.get("api_key"):
+                print("⚠️ API key is still missing. Please configure:")
+                client.configure(interactive=True)
+            
+            return False
+        
+        # Verify connection to proxy
+        health = client.health_check()
+        if not health["success"]:
+            print(f"⚠️ Could not connect to Modal proxy at {config.get('proxy_url')}")
+            print(f"   Error: {health.get('error', 'Unknown error')}")
+            print("   Run './gitarsenal.py proxy configure' to update configuration.")
+            return False
+        
+        print(f"✅ Connected to Modal proxy at {config.get('proxy_url')}")
+        return True
+        
+    except ImportError:
+        print("⚠️ GitArsenalProxyClient module not found")
+        return False
+    except Exception as e:
+        print(f"⚠️ Error checking proxy configuration: {e}")
+        return False
+
 def main():
     parser = argparse.ArgumentParser(description="GitArsenal CLI - GPU-accelerated cloud environments")
     subparsers = parser.add_subparsers(dest="command", help="Command to run")
@@ -96,8 +145,10 @@ def main():
     sandbox_parser.add_argument("--repo-name", type=str, help="Repository name override")
     sandbox_parser.add_argument("--setup-commands", type=str, nargs="+", help="Setup commands to run")
     sandbox_parser.add_argument("--volume-name", type=str, help="Name of the Modal volume for persistent storage")
+    sandbox_parser.add_argument("--use-api", action="store_true", help="Use API for setup commands")
+    sandbox_parser.add_argument("--use-proxy", action="store_true", help="Use Modal proxy service instead of direct Modal access")
     
-    # SSH container command
+    # SSH command
     ssh_parser = subparsers.add_parser("ssh", help="Create a Modal SSH container")
     ssh_parser.add_argument("--gpu", type=str, default="A10G", choices=["A10G", "A100", "H100", "T4", "V100"],
                           help="GPU type (default: A10G)")
@@ -106,123 +157,437 @@ def main():
     ssh_parser.add_argument("--setup-commands", type=str, nargs="+", help="Setup commands to run")
     ssh_parser.add_argument("--volume-name", type=str, help="Name of the Modal volume for persistent storage")
     ssh_parser.add_argument("--timeout", type=int, default=60, help="Container timeout in minutes (default: 60)")
+    ssh_parser.add_argument("--use-api", action="store_true", help="Use API for setup commands")
+    ssh_parser.add_argument("--use-proxy", action="store_true", help="Use Modal proxy service instead of direct Modal access")
+    
+    # Proxy command
+    proxy_parser = subparsers.add_parser("proxy", help="Configure and use Modal proxy service")
+    proxy_subparsers = proxy_parser.add_subparsers(dest="proxy_command", help="Proxy command")
+    
+    # Proxy configure command
+    proxy_configure_parser = proxy_subparsers.add_parser("configure", help="Configure Modal proxy service")
+    
+    # Proxy status command
+    proxy_status_parser = proxy_subparsers.add_parser("status", help="Check Modal proxy service status")
+    
+    # Proxy sandbox command
+    proxy_sandbox_parser = proxy_subparsers.add_parser("sandbox", help="Create a sandbox through Modal proxy")
+    proxy_sandbox_parser.add_argument("--gpu", type=str, default="A10G", choices=["A10G", "A100", "H100", "T4", "V100"],
+                                    help="GPU type (default: A10G)")
+    proxy_sandbox_parser.add_argument("--repo-url", type=str, help="Repository URL to clone")
+    proxy_sandbox_parser.add_argument("--repo-name", type=str, help="Repository name override")
+    proxy_sandbox_parser.add_argument("--setup-commands", type=str, nargs="+", help="Setup commands to run")
+    proxy_sandbox_parser.add_argument("--volume-name", type=str, help="Name of the Modal volume for persistent storage")
+    proxy_sandbox_parser.add_argument("--wait", action="store_true", help="Wait for sandbox to be ready")
+    
+    # Proxy ssh command
+    proxy_ssh_parser = proxy_subparsers.add_parser("ssh", help="Create an SSH container through Modal proxy")
+    proxy_ssh_parser.add_argument("--gpu", type=str, default="A10G", choices=["A10G", "A100", "H100", "T4", "V100"],
+                                help="GPU type (default: A10G)")
+    proxy_ssh_parser.add_argument("--repo-url", type=str, help="Repository URL to clone")
+    proxy_ssh_parser.add_argument("--repo-name", type=str, help="Repository name override")
+    proxy_ssh_parser.add_argument("--setup-commands", type=str, nargs="+", help="Setup commands to run")
+    proxy_ssh_parser.add_argument("--volume-name", type=str, help="Name of the Modal volume for persistent storage")
+    proxy_ssh_parser.add_argument("--timeout", type=int, default=60, help="Container timeout in minutes (default: 60)")
+    proxy_ssh_parser.add_argument("--wait", action="store_true", help="Wait for container to be ready")
     
     args = parser.parse_args()
     
-    # If no command is provided, show help
     if not args.command:
         parser.print_help()
-        return 1
-    
-    # Handle credentials commands
-    if args.command == "credentials":
-        cred_args = []
-        
-        if not args.cred_command:
-            cred_parser.print_help()
-            return 1
-            
-        cred_args.append(args.cred_command)
-        
-        if args.cred_command == "set" or args.cred_command == "get":
-            cred_args.append(args.key)
-        elif args.cred_command == "clear" and args.key != "all":
-            cred_args.append(args.key)
-            
-        return run_script("manage_credentials.py", cred_args)
+        return 0
     
     # For sandbox and SSH commands, check Modal authentication first
     if args.command in ["sandbox", "ssh"]:
-        # Check if Modal is authenticated
-        if not check_modal_auth():
-            print("\n⚠️ Please authenticate with Modal before proceeding.")
-            return 1
-        
-        # Try to load credentials and set Modal token if available
+        # Check if using proxy service
+        if hasattr(args, 'use_proxy') and args.use_proxy:
+            # Check if proxy is configured
+            if not check_proxy_config():
+                print("\n⚠️ Modal proxy service is not configured.")
+                print("Please run './gitarsenal.py proxy configure' first.")
+                return 1
+        else:
+            # Check if Modal is authenticated
+            if not check_modal_auth():
+                print("\n⚠️ Please authenticate with Modal before proceeding.")
+                return 1
+            
+            # Try to load credentials and set Modal token if available
+            try:
+                from credentials_manager import CredentialsManager
+                credentials_manager = CredentialsManager()
+                credentials = credentials_manager.load_credentials()
+                
+                if "modal_token" in credentials:
+                    # Set the Modal token in the environment
+                    os.environ["MODAL_TOKEN_ID"] = credentials["modal_token"]
+                    print("✅ Using Modal token from credentials")
+                    
+                    # Try to authenticate with the token
+                    try:
+                        token_result = subprocess.run(
+                            ["modal", "token", "set", credentials["modal_token"]], 
+                            capture_output=True, text=True
+                        )
+                        if token_result.returncode == 0:
+                            print("✅ Successfully authenticated with Modal")
+                    except Exception as e:
+                        print(f"⚠️ Error setting Modal token: {e}")
+            except ImportError:
+                print("⚠️ Could not load credentials manager")
+            except Exception as e:
+                print(f"⚠️ Error loading credentials: {e}")
+    
+    # Handle credentials commands
+    if args.command == "credentials":
         try:
             from credentials_manager import CredentialsManager
             credentials_manager = CredentialsManager()
-            credentials = credentials_manager.load_credentials()
             
-            if "modal_token" in credentials:
-                # Set the Modal token in the environment
-                os.environ["MODAL_TOKEN_ID"] = credentials["modal_token"]
-                print("✅ Using Modal token from credentials")
+            if args.cred_command == "setup":
+                # Set up all credentials
+                print("🔧 Setting up all credentials")
+                
+                # Modal token
+                modal_token = credentials_manager.get_modal_token()
+                if modal_token:
+                    print("✅ Modal token set")
+                else:
+                    print("⚠️ Modal token not set")
+                
+                # OpenAI API key
+                openai_api_key = credentials_manager.get_openai_api_key()
+                if openai_api_key:
+                    print("✅ OpenAI API key set")
+                else:
+                    print("⚠️ OpenAI API key not set")
+                
+                # Hugging Face token
+                huggingface_token = credentials_manager.get_huggingface_token()
+                if huggingface_token:
+                    print("✅ Hugging Face token set")
+                else:
+                    print("⚠️ Hugging Face token not set")
+                
+                # Weights & Biases API key
+                wandb_api_key = credentials_manager.get_wandb_api_key()
+                if wandb_api_key:
+                    print("✅ Weights & Biases API key set")
+                else:
+                    print("⚠️ Weights & Biases API key not set")
+                
+            elif args.cred_command == "set":
+                # Set a specific credential
+                if args.key == "modal_token":
+                    modal_token = credentials_manager.get_modal_token()
+                    if modal_token:
+                        print("✅ Modal token set")
+                    else:
+                        print("⚠️ Modal token not set")
+                elif args.key == "openai_api_key":
+                    openai_api_key = credentials_manager.get_openai_api_key()
+                    if openai_api_key:
+                        print("✅ OpenAI API key set")
+                    else:
+                        print("⚠️ OpenAI API key not set")
+                elif args.key == "huggingface_token":
+                    huggingface_token = credentials_manager.get_huggingface_token()
+                    if huggingface_token:
+                        print("✅ Hugging Face token set")
+                    else:
+                        print("⚠️ Hugging Face token not set")
+                elif args.key == "wandb_api_key":
+                    wandb_api_key = credentials_manager.get_wandb_api_key()
+                    if wandb_api_key:
+                        print("✅ Weights & Biases API key set")
+                    else:
+                        print("⚠️ Weights & Biases API key not set")
+            
+            elif args.cred_command == "get":
+                # Get a specific credential
+                credentials = credentials_manager.load_credentials()
+                if args.key in credentials:
+                    # Mask the credential for security
+                    value = credentials[args.key]
+                    masked_value = value[:4] + "*" * (len(value) - 8) + value[-4:] if len(value) > 8 else "****"
+                    print(f"{args.key}: {masked_value}")
+                else:
+                    print(f"⚠️ {args.key} not found")
+            
+            elif args.cred_command == "clear":
+                # Clear credentials
+                if args.key == "all":
+                    credentials_manager.clear_all_credentials()
+                    print("✅ All credentials cleared")
+                else:
+                    if credentials_manager.clear_credential(args.key):
+                        print(f"✅ {args.key} cleared")
+                    else:
+                        print(f"⚠️ {args.key} not found")
+            
+            elif args.cred_command == "list":
+                # List all credentials
+                credentials = credentials_manager.load_credentials()
+                if credentials:
+                    print("📋 Saved credentials:")
+                    for key in credentials:
+                        print(f"  - {key}")
+                else:
+                    print("⚠️ No credentials found")
+            
+            else:
+                print("⚠️ Unknown credentials command")
+                return 1
                 
-                # Try to authenticate with the token
-                try:
-                    token_result = subprocess.run(
-                        ["modal", "token", "set", credentials["modal_token"]], 
-                        capture_output=True, text=True
-                    )
-                    if token_result.returncode == 0:
-                        print("✅ Successfully authenticated with Modal")
-                except Exception as e:
-                    print(f"⚠️ Error setting Modal token: {e}")
         except ImportError:
-            print("⚠️ Could not load credentials manager")
+            print("❌ Could not import credentials_manager module")
+            return 1
         except Exception as e:
-            print(f"⚠️ Error loading credentials: {e}")
+            print(f"❌ Error: {e}")
+            return 1
     
     # Handle sandbox command
-    if args.command == "sandbox":
-        sandbox_args = []
+    elif args.command == "sandbox":
+        try:
+            # Check if using proxy service
+            if args.use_proxy:
+                # Import proxy client
+                try:
+                    from gitarsenal_proxy_client import GitArsenalProxyClient
+                    client = GitArsenalProxyClient()
+                    
+                    # Create sandbox through proxy
+                    result = client.create_sandbox(
+                        gpu_type=args.gpu,
+                        repo_url=args.repo_url,
+                        repo_name=args.repo_name,
+                        setup_commands=args.setup_commands,
+                        volume_name=args.volume_name,
+                        wait=True  # Always wait for sandbox to be ready
+                    )
+                    
+                    if not result:
+                        print("❌ Failed to create sandbox through proxy service")
+                        return 1
+                    
+                    print("✅ Sandbox created successfully through proxy service")
+                    
+                except ImportError:
+                    print("❌ Could not import gitarsenal_proxy_client module")
+                    print("Please make sure gitarsenal_proxy_client.py is in the same directory")
+                    return 1
+            else:
+                # Import sandbox creation function
+                from test_modalSandboxScript import create_modal_sandbox
+                
+                # Create sandbox directly
+                result = create_modal_sandbox(
+                    args.gpu,
+                    repo_url=args.repo_url,
+                    repo_name=args.repo_name,
+                    setup_commands=args.setup_commands,
+                    volume_name=args.volume_name
+                )
+                
+                if not result:
+                    print("❌ Failed to create sandbox")
+                    return 1
+                
+                print("✅ Sandbox created successfully")
         
-        if args.gpu:
-            sandbox_args.extend(["--gpu", args.gpu])
-        if args.repo_url:
-            sandbox_args.extend(["--repo-url", args.repo_url])
-        if args.repo_name:
-            sandbox_args.extend(["--repo-name", args.repo_name])
-        if args.setup_commands:
-            sandbox_args.extend(["--setup-commands"] + args.setup_commands)
-        if args.volume_name:
-            sandbox_args.extend(["--volume-name", args.volume_name])
-            
-        return run_script("test_modalSandboxScript.py", sandbox_args)
-    
-    # Handle SSH container command
+        except ImportError as e:
+            print(f"❌ Could not import required module: {e}")
+            return 1
+        except Exception as e:
+            print(f"❌ Error: {e}")
+            return 1
+    
+    # Handle SSH command
     elif args.command == "ssh":
-        ssh_args = []
+        try:
+            # Check if using proxy service
+            if args.use_proxy:
+                # Import proxy client
+                try:
+                    from gitarsenal_proxy_client import GitArsenalProxyClient
+                    client = GitArsenalProxyClient()
+                    
+                    # Ensure proxy is configured
+                    config = client.load_config()
+                    if not config.get("proxy_url") or not config.get("api_key"):
+                        print("\n⚠️ Modal proxy service is not fully configured.")
+                        print("Running configuration wizard...")
+                        client.configure(interactive=True)
+                        # Reload config after configuration
+                        config = client.load_config()
+                        if not config.get("proxy_url") or not config.get("api_key"):
+                            print("❌ Proxy configuration failed or was cancelled.")
+                            return 1
+                        print("✅ Proxy configuration completed successfully.")
+                    
+                    # Create SSH container through proxy
+                    print(f"🚀 Creating SSH container through proxy service ({config.get('proxy_url')})")
+                    result = client.create_ssh_container(
+                        gpu_type=args.gpu,
+                        repo_url=args.repo_url,
+                        repo_name=args.repo_name,
+                        setup_commands=args.setup_commands,
+                        volume_name=args.volume_name,
+                        timeout=args.timeout,
+                        wait=True  # Always wait for container to be ready
+                    )
+                    
+                    if not result:
+                        print("❌ Failed to create SSH container through proxy service")
+                        return 1
+                    
+                    print("✅ SSH container created successfully through proxy service")
+                    
+                except ImportError:
+                    print("❌ Could not import gitarsenal_proxy_client module")
+                    print("Please make sure gitarsenal_proxy_client.py is in the same directory")
+                    return 1
+            else:
+                # Import SSH container creation function
+                from test_modalSandboxScript import create_modal_ssh_container
+                
+                # Create SSH container directly
+                result = create_modal_ssh_container(
+                    args.gpu,
+                    repo_url=args.repo_url,
+                    repo_name=args.repo_name,
+                    setup_commands=args.setup_commands,
+                    volume_name=args.volume_name,
+                    timeout_minutes=args.timeout
+                )
+                
+                if not result:
+                    print("❌ Failed to create SSH container")
+                    return 1
+                
+                print("✅ SSH container created successfully")
         
-        if args.gpu:
-            ssh_args.extend(["--gpu", args.gpu])
-        if args.repo_url:
-            ssh_args.extend(["--repo-url", args.repo_url])
-        if args.repo_name:
-            ssh_args.extend(["--repo-name", args.repo_name])
-        if args.setup_commands:
-            ssh_args.extend(["--setup-commands"] + args.setup_commands)
-        if args.volume_name:
-            ssh_args.extend(["--volume-name", args.volume_name])
-        if args.timeout:
-            ssh_args.extend(["--timeout", str(args.timeout)])
-            
-        # Use test_modalSandboxScript.py with SSH mode
-        ssh_args.extend(["--ssh"])
-        return run_script("test_modalSandboxScript.py", ssh_args)
-    
-    return 0
-
-def run_script(script_name, args):
-    """Run a Python script with the given arguments"""
-    # Get the directory of the current script
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    script_path = os.path.join(script_dir, script_name)
+        except ImportError as e:
+            print(f"❌ Could not import required module: {e}")
+            return 1
+        except Exception as e:
+            print(f"❌ Error: {e}")
+            return 1
     
-    # Build the command
-    cmd = [sys.executable, script_path] + args
+    # Handle proxy commands
+    elif args.command == "proxy":
+        if not args.proxy_command:
+            proxy_parser.print_help()
+            return 0
+        
+        try:
+            # Import proxy client
+            from gitarsenal_proxy_client import GitArsenalProxyClient
+            client = GitArsenalProxyClient()
+            
+            if args.proxy_command == "configure":
+                # Configure proxy service
+                client.configure(interactive=True)
+            
+            elif args.proxy_command == "status":
+                # Check proxy service status
+                response = client.health_check()
+                if response["success"]:
+                    print(f"✅ Proxy service is running: {response['data']['message']}")
+                else:
+                    print(f"❌ Proxy service health check failed: {response['error']}")
+                    return 1
+            
+            elif args.proxy_command == "sandbox":
+                # Create sandbox through proxy
+                result = client.create_sandbox(
+                    gpu_type=args.gpu,
+                    repo_url=args.repo_url,
+                    repo_name=args.repo_name,
+                    setup_commands=args.setup_commands,
+                    volume_name=args.volume_name,
+                    wait=args.wait
+                )
+                
+                if not result:
+                    print("❌ Failed to create sandbox through proxy service")
+                    return 1
+                
+                print("✅ Sandbox created successfully through proxy service")
+            
+            elif args.proxy_command == "ssh":
+                # Create SSH container through proxy
+                print("🚀 Creating SSH container through proxy...")
+                
+                # Verify proxy configuration is complete
+                config = client.load_config()
+                if not config.get("proxy_url") or not config.get("api_key"):
+                    print("⚠️ Proxy configuration incomplete. Running configuration wizard...")
+                    client.configure(interactive=True)
+                    # Reload config after configuration
+                    config = client.load_config()
+                    if not config.get("proxy_url") or not config.get("api_key"):
+                        print("❌ Proxy configuration failed or was cancelled.")
+                        return 1
+                    print("✅ Proxy configuration completed successfully.")
+                
+                # Check proxy service health
+                health = client.health_check()
+                if not health["success"]:
+                    print(f"❌ Could not connect to proxy service at {config.get('proxy_url')}")
+                    print(f"   Error: {health.get('error', 'Unknown error')}")
+                    print("   Please check if the proxy service is running.")
+                    return 1
+                
+                print(f"✅ Connected to proxy service at {config.get('proxy_url')}")
+                
+                # Create SSH container
+                result = client.create_ssh_container(
+                    gpu_type=args.gpu,
+                    repo_url=args.repo_url,
+                    repo_name=args.repo_name,
+                    setup_commands=args.setup_commands,
+                    volume_name=args.volume_name,
+                    timeout=args.timeout,
+                    wait=args.wait
+                )
+                
+                if not result:
+                    print("❌ Failed to create SSH container through proxy service")
+                    return 1
+                
+                # Print connection information if available
+                if isinstance(result, dict):
+                    if "container_id" in result:
+                        print(f"📋 Container ID: {result['container_id']}")
+                    if "ssh_password" in result:
+                        print(f"🔐 SSH Password: {result['ssh_password']}")
+                
+                print("✅ SSH container creation process initiated through proxy service")
+                
+                if not args.wait:
+                    print("\n⚠️ Container creation is asynchronous. Use --wait flag to wait for it to be ready.")
+                    print("   You can check the status later using the container ID above.")
+                
+            else:
+                print(f"❌ Unknown proxy command: {args.proxy_command}")
+                proxy_parser.print_help()
+                return 1
+        
+        except ImportError:
+            print("❌ Could not import gitarsenal_proxy_client module")
+            print("Please make sure gitarsenal_proxy_client.py is in the same directory")
+            return 1
+        except Exception as e:
+            print(f"❌ Error: {e}")
+            return 1
     
-    try:
-        # Run the command
-        result = subprocess.run(cmd)
-        return result.returncode
-    except KeyboardInterrupt:
-        print("\n⚠️ Command interrupted by user")
-        return 130
-    except Exception as e:
-        print(f"❌ Error running {script_name}: {e}")
+    else:
+        print(f"❌ Unknown command: {args.command}")
+        parser.print_help()
         return 1
+    
+    return 0
 
 if __name__ == "__main__":
     sys.exit(main())