clean main file
This commit is contained in:
40
src/commands/init.py
Normal file
40
src/commands/init.py
Normal file
@@ -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
|
||||||
70
src/commands/upload.py
Normal file
70
src/commands/upload.py
Normal file
@@ -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)
|
||||||
108
src/main.py
108
src/main.py
@@ -1,115 +1,13 @@
|
|||||||
import secrets
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import typer
|
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, init, train, upload
|
||||||
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
|
|
||||||
|
|
||||||
app = typer.Typer(
|
app = typer.Typer(
|
||||||
help="qc-cli: End-to-end model managment for Qualcomm AI Hub.",
|
help="qc-cli: End-to-end model managment for Qualcomm AI Hub.",
|
||||||
no_args_is_help=True,
|
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(infra.app, name="infra")
|
||||||
app.add_typer(train.app, name="train")
|
app.add_typer(train.app, name="train")
|
||||||
app.add_typer(ai_hub.app, name="ai-hub")
|
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)
|
|
||||||
|
|||||||
Reference in New Issue
Block a user