diff --git a/src/commands/init.py b/src/commands/init.py new file mode 100644 index 0000000..0a01e2a --- /dev/null +++ b/src/commands/init.py @@ -0,0 +1,40 @@ +import secrets +from pathlib import Path + +import typer +import yaml + +from src.commands.utils import CONSOLE +from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config + +app = typer.Typer() + + +@app.command() +def init( + output: str = typer.Option("config.yaml", help="Destination path for the config file"), + force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"), +) -> None: + """Write a starter config.yaml to the current directory.""" + dest = Path(output) + if dest.exists() and not force: + CONSOLE.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.") + raise typer.Exit(1) + + config = _new_isolated_config() + dest.parent.mkdir(parents=True, exist_ok=True) + config_data = config.model_dump(mode="json") + config_data["sagemaker"].pop("role_name", None) + with open(dest, "w") as f: + yaml.safe_dump(config_data, f, sort_keys=False) + + CONSOLE.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]") + CONSOLE.print("Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands.") + + +def _new_isolated_config() -> Config: + suffix = secrets.token_hex(6) + namespace = f"{GENERATED_STACK_PREFIX}{suffix}" + config = Config(infra=InfraConfig(stack_name=namespace)) + config.s3 = S3Config(bucket=f"{namespace}-data") + return config diff --git a/src/commands/upload.py b/src/commands/upload.py new file mode 100644 index 0000000..d48223e --- /dev/null +++ b/src/commands/upload.py @@ -0,0 +1,70 @@ +from pathlib import Path + +import typer +from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn + +from src.aws import s3 as s3_ops +from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg + +app = typer.Typer() + + +@app.command() +def upload( + path: Path = typer.Argument(..., help="Local file or directory to upload"), + s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"), + config: str = CONFIG_OPT, +) -> None: + """Upload a local file or directory to S3.""" + cfg = load_cfg(config) + + if path.is_file(): + key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}" + try: + with CONSOLE.status(f"Uploading {path.name}..."): + uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key) + except Exception as e: + CONSOLE.print(f"[red]Upload failed: {e}[/red]") + raise typer.Exit(1) + + CONSOLE.print(f"[green]✓[/green] {path.name} -> {uri}") + return + + if path.is_dir(): + if s3_key is not None: + CONSOLE.print("[red]--s3-key can only be used when uploading a single file.[/red]") + raise typer.Exit(1) + + files = [file for file in path.rglob("*") if file.is_file()] + if not files: + CONSOLE.print("[yellow]No files found in directory.[/yellow]") + raise typer.Exit(0) + + prefix = cfg.s3.data_prefix + CONSOLE.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") + try: + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + console=CONSOLE, + ) as progress: + task = progress.add_task("Uploading...", total=len(files)) + count = s3_ops.upload_dir( + cfg.aws.region, + cfg.aws.profile, + cfg.s3.bucket, + str(path), + prefix, + on_progress=lambda: progress.advance(task), + ) + except Exception as e: + CONSOLE.print(f"[red]Upload failed: {e}[/red]") + raise typer.Exit(1) + + CONSOLE.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") + return + + CONSOLE.print(f"[red]Path not found: {path}[/red]") + raise typer.Exit(1) diff --git a/src/main.py b/src/main.py index 23b172c..61df2ee 100644 --- a/src/main.py +++ b/src/main.py @@ -1,115 +1,13 @@ -import secrets -from pathlib import Path - import typer -import yaml -from rich.console import Console -from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn -from src.aws import s3 as s3_ops -from src.commands import ai_hub, infra, train -from src.commands.utils import CONFIG_OPT, load_cfg -from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config +from src.commands import ai_hub, infra, init, train, upload app = typer.Typer( help="qc-cli: End-to-end model managment for Qualcomm AI Hub.", no_args_is_help=True, ) +app.add_typer(init.app) +app.add_typer(upload.app) app.add_typer(infra.app, name="infra") app.add_typer(train.app, name="train") app.add_typer(ai_hub.app, name="ai-hub") - -console = Console() - - -@app.command() -def init( - output: str = typer.Option("config.yaml", help="Destination path for the config file"), - force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"), -) -> None: - """Write a starter config.yaml to the current directory.""" - dest = Path(output) - if dest.exists() and not force: - console.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.") - raise typer.Exit(1) - - config = _new_isolated_config() - dest.parent.mkdir(parents=True, exist_ok=True) - config_data = config.model_dump(mode="json") - config_data["sagemaker"].pop("role_name", None) - with open(dest, "w") as f: - yaml.safe_dump(config_data, f, sort_keys=False) - - console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]") - console.print( - "Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands." - ) - - -def _new_isolated_config() -> Config: - suffix = secrets.token_hex(6) - namespace = f"{GENERATED_STACK_PREFIX}{suffix}" - config = Config(infra=InfraConfig(stack_name=namespace)) - config.s3 = S3Config(bucket=f"{namespace}-data") - return config - - -@app.command() -def upload( - path: Path = typer.Argument(..., help="Local file or directory to upload"), - s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"), - config: str = CONFIG_OPT, -) -> None: - """Upload a local file or directory to S3.""" - cfg = load_cfg(config) - - if path.is_file(): - key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}" - try: - with console.status(f"Uploading {path.name}..."): - uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key) - except Exception as e: - console.print(f"[red]Upload failed: {e}[/red]") - raise typer.Exit(1) - - console.print(f"[green]✓[/green] {path.name} -> {uri}") - return - - if path.is_dir(): - if s3_key is not None: - console.print("[red]--s3-key can only be used when uploading a single file.[/red]") - raise typer.Exit(1) - - files = [file for file in path.rglob("*") if file.is_file()] - if not files: - console.print("[yellow]No files found in directory.[/yellow]") - raise typer.Exit(0) - - prefix = cfg.s3.data_prefix - console.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") - try: - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TaskProgressColumn(), - console=console, - ) as progress: - task = progress.add_task("Uploading...", total=len(files)) - count = s3_ops.upload_dir( - cfg.aws.region, - cfg.aws.profile, - cfg.s3.bucket, - str(path), - prefix, - on_progress=lambda: progress.advance(task), - ) - except Exception as e: - console.print(f"[red]Upload failed: {e}[/red]") - raise typer.Exit(1) - - console.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") - return - - console.print(f"[red]Path not found: {path}[/red]") - raise typer.Exit(1)