jazari 0.1.0__tar.gz → 0.1.1__tar.gz

{jazari-0.1.0/src/jazari.egg-info → jazari-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jazari
-Version: 0.1.0
+Version: 0.1.1
 Summary: The orchestration layer for modern Slurm clusters.
 Author-email: Levent Ozbek <levent@jazari.run>
 License: MIT

{jazari-0.1.0 → jazari-0.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jazari"
-version = "0.1.0"
+version = "0.1.1"
 description = "The orchestration layer for modern Slurm clusters."
 readme = "README.md"
 requires-python = ">=3.8"
@@ -45,4 +45,4 @@ dependencies = [
 ]
 
 [project.scripts]
-jazari = "jazari.main:app"
+jazari = "jazari.main:app"

jazari-0.1.1/src/jazari/commands/init.py

@@ -0,0 +1,37 @@
+from rich.console import Console
+from rich.prompt import Prompt, Confirm
+from jazari.config_generator import load_config, make_config_dir, CONFIG_FILE
+
+console = Console()
+
+def init_command():
+    """
+    Initialize Jazari config for this cluster.
+    """
+    current_config = load_config()
+
+    console.print('[bold blue]Welcome to Jazari setup.[/bold blue]')
+    console.print(f"This will create a configuration file at [dim]{CONFIG_FILE}[/dim]\n")
+
+    # 1. Ask for Slurm Account
+    default_account = current_config.get("account", "")
+    account = Prompt.ask("Enter your default Slurm account (e.g., def-user)",default=default_account)
+
+    # 2. Ask for default time limit
+    default_time = current_config.get("time", "01:00:00")
+    time_limit = Prompt.ask("Enter default time limit (D-HH:MM)",default=default_time)
+
+    # 3. Ask for W&B tracking default
+    default_wandb = current_config.get("track_wandb", False)
+    track_wandb = Confirm.ask("Enable W&B tracking by default?",default=default_wandb)
+
+    new_config = {
+        "account":      account,
+        "time":         time_limit,
+        "track_wandb":  track_wandb
+    }
+
+    make_config_dir(new_config)
+
+    console.print(f"\n[bold green]✅ Configuration saved![/bold green]")
+    console.print("You can now run jobs without specifying these flags.")

jazari-0.1.1/src/jazari/commands/kill.py

@@ -0,0 +1,32 @@
+import typer, shutil, sys, subprocess
+from rich.console import Console
+from rich.prompt import Confirm
+from typing import Optional
+from jazari.utils import get_current_user
+
+console = Console()
+
+def kill_command(
+        job_id: Optional[str] = typer.Argument(None, help = 'Job ID to cancel.'),
+        cancel_all: bool = typer.Option(False, '--all', help = 'Cancel ALL active jobs.')
+):
+    if not shutil.which("scancel"):
+        console.print("[bold red]Error:[/bold red] 'scancel' command not found.")
+        console.print("Are you running this on a Slurm login node?")
+        sys.exit(1)
+    
+    user = get_current_user()
+
+    if cancel_all:
+        if not Confirm.ask(f"[bold red]⚠️  DANGER:[/bold red] Cancel ALL jobs for '{user}'?"):
+             sys.exit(0)
+        subprocess.run(["scancel", "-u", user], check=True)
+        console.print(f"[bold green]💥 All jobs cancelled.[/bold green]")
+        return
+
+    if job_id:
+        subprocess.run(["scancel", job_id], check=True)
+        console.print(f"[bold green]💥 Signal sent to cancel job {job_id}.[/bold green]")
+        return
+
+    console.print("[yellow]Usage: jazari kill <JOB_ID> or jazari kill --all[/yellow]")

jazari-0.1.1/src/jazari/commands/logs.py

@@ -0,0 +1,42 @@
+import typer, sys, os, subprocess
+from pathlib import Path
+from typing import Optional
+from rich.console import Console
+
+console = Console()
+
+def logs_command(
+    job_id: Optional[str] = typer.Argument(None, help="Job ID to view."),
+    tail: bool = typer.Option(False, "--tail", "-f", help="Follow output live."),
+    error: bool = typer.Option(False, "--error", "-e", help="View .err log instead of .out")
+):
+    """View or follow job logs."""
+    log_dir = Path("logs")
+    if not log_dir.exists():
+        console.print("[red]Error: 'logs/' directory not found.[/red]")
+        sys.exit(1)
+
+    suffix = ".err" if error else ".out"
+    target_file = None
+
+    if job_id:
+        matching = list(log_dir.glob(f"*{job_id}{suffix}"))
+        if not matching:
+             console.print(f"[red]No log found for Job {job_id}[/red]")
+             sys.exit(1)
+        target_file = matching[0]
+    else:
+        files = list(log_dir.glob(f"*{suffix}"))
+        if not files:
+             console.print(f"[yellow]No {suffix} files found.[/yellow]")
+             return
+        files.sort(key=os.path.getmtime)
+        target_file = files[-1]
+
+    console.print(f"[dim]Viewing: {target_file}[/dim]")
+    cmd = ["tail", "-f", str(target_file)] if tail else ["cat", str(target_file)]
+    
+    try:
+        subprocess.run(cmd, check=True)
+    except KeyboardInterrupt:
+        console.print("\n[dim]Stopped.[/dim]")

jazari-0.1.1/src/jazari/commands/run.py

@@ -0,0 +1,94 @@
+import os, typer, sys, shlex, tempfile, subprocess
+from typing import Optional
+from rich.console import Console
+from rich.syntax import Syntax
+from jazari.slurm_generator import generate_sbatch_script
+from jazari.config_generator import load_config, get_config_default
+from jazari.utils import get_wandb_api_key
+
+console = Console()
+CONFIG = load_config()
+
+def run_command(
+    command:     list[str]      = typer.Argument(..., help="The command to run (e.g., python train.py). Use '--' before it."),
+    nodes:       int            = typer.Option(get_config_default("nodes", CONFIG, 1), '--nodes', '-N', help='Number of nodes.'),
+    gpus:        int            = typer.Option(get_config_default("gpus", CONFIG, 1), '--gpus', '-G', help='GPUs per node.'),
+    cpus:        int            = typer.Option(get_config_default("cpus", CONFIG, 1), '--cpus', '-n', help='CPUs per task/GPU'), # Changed -c to -n per our earlier discussion
+    time:        str            = typer.Option(get_config_default("time", CONFIG, '01:00:00'), '--time', '-t', help='Time limit (D-HH:MM)'),
+    name:        str            = typer.Option('jazari_run', '--name', '-j', help='Job name'), # Changed -n to -j to allow CPU shortcut
+    account:     Optional[str]  = typer.Option(get_config_default("account", CONFIG, None), "--account", "-A", help="Slurm account to charge."),
+    track_wandb: bool           = typer.Option(get_config_default("track_wandb", CONFIG, False), "--track-wandb", help="Auto-configure W&B."),
+    push_to_hub: Optional[str]  = typer.Option(None, "--push-to-hub", help="Hugging Face repo ID to upload model to."),   
+    pull_data:   Optional[str]  = typer.Option(None, '--pull-data', help='Huggingface dataset ID to download.'),
+    dry_run:     bool           = typer.Option(False, '--dry-run', help='Print sbatch script without submitting.'),
+    ):
+    '''
+    Launch a distributed training job.
+    '''
+    if not command:
+        console.print("[bold red]Error:[/bold red] You must provide a command to run.")
+        sys.exit(1)
+
+    log_dir = "logs"
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir, exist_ok=True)
+
+    full_command_str = shlex.join(command)
+
+    # --- Weights & Biases ---
+    wandb_key = None
+    if track_wandb:
+        console.print("[dim]W&B tracking enabled.[/dim]")
+        wandb_key = get_wandb_api_key()
+        if not wandb_key:
+            console.print("[bold yellow]Warning:[/bold yellow] Could not find W&B API key locally.")
+            console.print("Please run [green]wandb login[/green] on this machine first.")
+    # -----------------
+
+    console.print(f'[bold]🐘 Generating Slurm script for: {name}[/bold]')
+
+    sbatch_content = generate_sbatch_script(
+        nodes           = nodes,
+        gpus_per_node   = gpus,
+        cpus_per_task   = cpus,
+        time_limit      = time,
+        job_name        = name,
+        account_name    = account,
+        wandb_api_key   = wandb_key,
+        hf_repo_id      = push_to_hub,
+        hf_dataset_id   = pull_data,
+        user_command    = full_command_str
+    )
+
+    if dry_run:
+        console.print('\n[yellow]--- DRY RUN: Generated #SBATCH Script ---[/yellow]')
+        syntax = Syntax(sbatch_content, 'bash', theme = 'monokai', line_numbers = True)
+        console.print(syntax)
+        console.print('[yellow]-----------------------------------------------[/yellow]')
+    else:
+        console.print("[dim]Submitting to Slurm scheduler...[/dim]")
+
+        with tempfile.NamedTemporaryFile(mode='w+', suffix=".sh", delete=False) as temp_file:
+            temp_script_path = temp_file.name
+            temp_file.write(sbatch_content)
+
+        try:
+            result = subprocess.run(
+                ["sbatch", temp_script_path],
+                check           = True,
+                capture_output  = True,
+                text            = True
+            )
+            job_id = result.stdout.strip().split()[-1]
+            console.print(f"[bold green]✅ Job submitted successfully![/bold green] (ID: [bold]{job_id}[/bold])")
+            console.print(f"[dim]View logs: cat logs/{name}-{job_id}.out[/dim]")
+
+        except subprocess.CalledProcessError as e:
+            console.print("[bold red]❌ Failed to submit job.[/bold red]")
+            console.print(f"Sbatch error: {e.stderr}")
+        except FileNotFoundError:
+             console.print("\n[bold red]❌ Error:[/bold red] 'sbatch' command not found.")
+             console.print("Are you running this on a Slurm login node?")
+        finally:
+            if os.path.exists(temp_script_path):
+                os.unlink(temp_script_path)

jazari-0.1.1/src/jazari/commands/status.py

@@ -0,0 +1,60 @@
+import shutil, sys, subprocess
+from rich.console import Console
+from rich.table import Table
+from rich import box
+from jazari.utils import get_current_user  # We will move helper functions to utils.py later
+
+console = Console()
+
+def status_command():
+    if not shutil.which("squeue"):
+        console.print("[bold red]Error:[/bold red] 'squeue' command not found.")
+        console.print("Are you running this on a Slurm login node?")
+        sys.exit(1)
+
+    user = get_current_user()
+
+    squeue_cmd = ["squeue", "-u", user, "-h", "-o", "%i|%j|%T|%M|%D|%R"]
+
+    try:
+        result       = subprocess.run(squeue_cmd, check=True, capture_output=True, text=True)
+        output_lines = result.stdout.strip().split('\n')
+        output_lines = [line for line in output_lines if line.strip()]
+
+        if not output_lines:
+             console.print(f"\n[dim]No active jobs found for user: [bold]{user}[/bold][/dim]\n")
+             return
+
+        table = Table(
+            title        = f"🐘 Cluster Activity for user: [bold green]{user}[/bold green]",
+            box          = box.ROUNDED,
+            header_style = "bold blue",
+            expand       = True
+        )
+        table.add_column("Job ID", style="cyan", no_wrap=True)
+        table.add_column("Name")
+        table.add_column("State")
+        table.add_column("Time Used", justify="right")
+        table.add_column("Nodes", justify="right")
+        table.add_column("Nodelist / Reason", style="dim")
+
+        for line in output_lines:
+            try:
+                job_id, name, state, time_used, nodes, reason = line.split('|')
+
+                # Color-code the state
+                if state == "RUNNING":
+                    state_formatted = f"[bold green]{state}[/bold green]"
+                elif state == "PENDING":
+                    state_formatted = f"[yellow]{state}[/yellow]"
+                else:
+                    state_formatted = f"[red]{state}[/red]"
+
+                table.add_row(job_id, name, state_formatted, time_used, nodes, reason)
+            except ValueError:
+                continue
+
+        console.print(table)
+
+    except subprocess.CalledProcessError as e:
+         console.print(f"[bold red]Error running squeue:[/bold red] {e.stderr}")

{jazari-0.1.0 → jazari-0.1.1}/src/jazari/config_generator.py

@@ -1,16 +1,12 @@
 import os, yaml
 from pathlib import Path
 from rich.console import Console
-# Config helpers
-# defining the path to the config file: ~/.jazari/config.yaml
+
 CONFIG_DIR = os.path.join(Path.home(), '.jazari')
 CONFIG_FILE = os.path.join(CONFIG_DIR, 'config.yaml')
 console = Console()
 
 def load_config():
-    '''
-    Loading config file if exists, otherwise returns empty dict.
-    '''
     if os.path.exists(CONFIG_FILE):
         try:
             with open(CONFIG_FILE, 'r') as f:
@@ -21,16 +17,11 @@ def load_config():
     return {}
 
 def get_config_default(key, CONFIG, default_value):
-    '''
-    Helper to get a default value from config or fall back to a hardcoded one.
-    '''
     return CONFIG.get(key, default_value)
 
 def make_config_dir(new_config):
-    # Creating directory if doesn't exist
     Path(CONFIG_DIR).mkdir(parents = True, exist_ok=True)
-
-    # Writing config to the yaml file
+    
     with open(CONFIG_FILE, 'w') as f:
         yaml.dump(new_config, f, default_flow_style=False)
     

jazari-0.1.1/src/jazari/main.py

@@ -0,0 +1,24 @@
+import typer
+from jazari.commands.status import status_command
+from jazari.commands.kill import kill_command
+from jazari.commands.logs import logs_command
+from jazari.commands.init import init_command
+from jazari.commands.run import run_command
+
+# Initializing the Typer app
+app = typer.Typer(
+    name            = "jazari",
+    help            = "🐘 The orchestration layer for modern Slurm clusters.",
+    add_completion  = False,
+    no_args_is_help = True
+)
+
+# Registering the commands
+app.command(name="status")(status_command)
+app.command(name="kill")(kill_command)
+app.command(name="logs")(logs_command)
+app.command(name="init")(init_command)
+app.command(name="run")(run_command)
+
+if __name__ == "__main__":
+    app()

jazari-0.1.1/src/jazari/utils.py

@@ -0,0 +1,30 @@
+import os, getpass, netrc
+from typing import Optional
+from rich.console import Console
+
+console = Console()
+
+def get_current_user() -> str:
+    """Returns the current username, robustly on Linux environments."""
+    try:
+        return os.environ['USER']
+    except KeyError:
+        return getpass.getuser()
+
+def get_wandb_api_key() -> Optional[str]:
+    """
+    Attempts to read the W&B API key from the local ~/.netrc file.
+    """
+    try:
+        login_info = netrc.netrc().authenticators("api.wandb.ai")
+        if login_info:
+            username, _, api_key = login_info
+            console.print(f"[dim]Found W&B API key for user: {username}[/dim]")
+            return api_key
+        else:
+            return None
+    except FileNotFoundError:
+        return None
+    except Exception as e:
+        console.print(f"[dim bold red]Warning: Could not read W&B credentials: {e}[/dim]")
+        return None

{jazari-0.1.0 → jazari-0.1.1/src/jazari.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jazari
-Version: 0.1.0
+Version: 0.1.1
 Summary: The orchestration layer for modern Slurm clusters.
 Author-email: Levent Ozbek <levent@jazari.run>
 License: MIT

{jazari-0.1.0 → jazari-0.1.1}/src/jazari.egg-info/SOURCES.txt

@@ -5,9 +5,15 @@ src/jazari/__init__.py
 src/jazari/config_generator.py
 src/jazari/main.py
 src/jazari/slurm_generator.py
+src/jazari/utils.py
 src/jazari.egg-info/PKG-INFO
 src/jazari.egg-info/SOURCES.txt
 src/jazari.egg-info/dependency_links.txt
 src/jazari.egg-info/entry_points.txt
 src/jazari.egg-info/requires.txt
-src/jazari.egg-info/top_level.txt
+src/jazari.egg-info/top_level.txt
+src/jazari/commands/init.py
+src/jazari/commands/kill.py
+src/jazari/commands/logs.py
+src/jazari/commands/run.py
+src/jazari/commands/status.py

jazari-0.1.0/src/jazari/main.py

@@ -1,329 +0,0 @@
-import typer, sys, os, tempfile, subprocess, netrc, shutil, getpass
-
-from typing import Optional
-from pathlib import Path
-
-from rich.console import Console
-from rich.syntax import Syntax
-from rich.prompt import Prompt, Confirm
-from rich.table import Table
-from rich import box
-
-from jazari.slurm_generator import generate_sbatch_script
-from jazari.config_generator import load_config, get_config_default, make_config_dir, CONFIG_FILE
-
-# Initializing Rich console for pretty printing
-console = Console()
-
-# Load config at start.
-CONFIG = load_config()
-
-# --- HELPER FUNCTIONS ---
-
-def get_current_user() -> str:
-    """Returns the current username, robustly on Linux environments."""
-    try:
-        return os.environ['USER']
-    except KeyError:
-        return getpass.getuser()
-
-def get_wandb_api_key() -> Optional[str]:
-    """
-    Attempts to read the W&B API key from the local ~/.netrc file.
-    """
-    try:
-        login_info = netrc.netrc().authenticators("api.wandb.ai")
-        if login_info:
-            username, _, api_key = login_info
-            console.print(f"[dim]Found W&B API key for user: {username}[/dim]")
-            return api_key
-        else:
-            return None
-    except FileNotFoundError:
-        return None
-    except Exception as e:
-        console.print(f"[dim bold red]Warning: Could not read W&B credentials: {e}[/dim]")
-        return None
-
-# --- COMMAND FUNCTIONS (No decorators here) ---
-
-def status_command():
-    """
-    View the status of active Slurm jobs for the current user.
-    """
-    # 1. Check if squeue exists available
-    if not shutil.which("squeue"):
-        console.print("[bold red]Error:[/bold red] 'squeue' command not found.")
-        console.print("Are you running this on a Slurm login node?")
-        sys.exit(1)
-
-    user = get_current_user()
-
-    # 2. Define the squeue command
-    # -u: Filter by user
-    # -h: No header (we'll make our own)
-    # -o: Custom format with '|' delimiter.
-    #     %i=ID, %j=Name, %T=State, %M=Time Used, %D=Nodes, %R=Reason/Nodelist
-    squeue_cmd = ["squeue", "-u", user, "-h", "-o", "%i|%j|%T|%M|%D|%R"]
-
-    try:
-        # 3. Run the command and capture output
-        result = subprocess.run(squeue_cmd, check=True, capture_output=True, text=True)
-        output_lines = result.stdout.strip().split('\n')
-
-        # Remove empty lines if squeue returned nothing
-        output_lines = [line for line in output_lines if line.strip()]
-
-        if not output_lines:
-             console.print(f"\n[dim]No active jobs found for user: [bold]{user}[/bold][/dim]\n")
-             return
-
-        # 4. Build the Rich Table
-        table = Table(
-            title=f"🐘 Cluster Activity for user: [bold green]{user}[/bold green]",
-            box=box.ROUNDED,
-            header_style="bold blue",
-            expand=True
-        )
-        table.add_column("Job ID", style="cyan", no_wrap=True)
-        table.add_column("Name")
-        table.add_column("State")
-        table.add_column("Time Used", justify="right")
-        table.add_column("Nodes", justify="right")
-        table.add_column("Nodelist / Reason", style="dim")
-
-        # 5. Parse output and add rows
-        for line in output_lines:
-            try:
-                job_id, name, state, time_used, nodes, reason = line.split('|')
-
-                # Color-code the state
-                if state == "RUNNING":
-                    state_formatted = f"[bold green]{state}[/bold green]"
-                elif state == "PENDING":
-                    state_formatted = f"[yellow]{state}[/yellow]"
-                else:
-                    state_formatted = f"[red]{state}[/red]"
-
-                table.add_row(job_id, name, state_formatted, time_used, nodes, reason)
-            except ValueError:
-                # Handle rare cases where splitting fails
-                continue
-
-        console.print(table)
-
-    except subprocess.CalledProcessError as e:
-         console.print(f"[bold red]Error running squeue:[/bold red] {e.stderr}")
-
-def kill_command(
-        job_id: Optional[str] = typer.Argument(None, help = 'Job ID to cancel.'),
-        cancel_all: bool = typer.Option(False, '--all', help = 'Cancel ALL active jobs.')
-):
-    # Cancel a specific job or all jobs
-    if not shutil.which("scancel"):
-        console.print("[bold red]Error:[/bold red] 'scancel' command not found.")
-        console.print("Are you running this on a Slurm login node?")
-        sys.exit(1)
-    
-    user = get_current_user()
-
-    if cancel_all:
-        if not Confirm.ask(f"[bold red]⚠️  DANGER:[/bold red] Cancel ALL jobs for '{user}'?"):
-             sys.exit(0)
-        subprocess.run(["scancel", "-u", user], check=True)
-        console.print(f"[bold green]💥 All jobs cancelled.[/bold green]")
-        return
-
-    if job_id:
-        subprocess.run(["scancel", job_id], check=True)
-        console.print(f"[bold green]💥 Signal sent to cancel job {job_id}.[/bold green]")
-        return
-
-    console.print("[yellow]Usage: jazari kill <JOB_ID> or jazari kill --all[/yellow]")
-
-def logs_command(
-    job_id: Optional[str] = typer.Argument(None, help="Job ID to view."),
-    tail: bool = typer.Option(False, "--tail", "-f", help="Follow output live."),
-    error: bool = typer.Option(False, "--error", "-e", help="View .err log instead of .out")
-):
-    """View or follow job logs."""
-    log_dir = Path("logs")
-    if not log_dir.exists():
-        console.print("[red]Error: 'logs/' directory not found.[/red]")
-        sys.exit(1)
-
-    suffix = ".err" if error else ".out"
-    target_file = None
-
-    if job_id:
-        matching = list(log_dir.glob(f"*{job_id}{suffix}"))
-        if not matching:
-             console.print(f"[red]No log found for Job {job_id}[/red]")
-             sys.exit(1)
-        target_file = matching[0]
-    else:
-        # Find latest
-        files = list(log_dir.glob(f"*{suffix}"))
-        if not files:
-             console.print(f"[yellow]No {suffix} files found.[/yellow]")
-             return
-        files.sort(key=os.path.getmtime)
-        target_file = files[-1]
-
-    console.print(f"[dim]Viewing: {target_file}[/dim]")
-    cmd = ["tail", "-f", str(target_file)] if tail else ["cat", str(target_file)]
-    
-    try:
-        subprocess.run(cmd, check=True)
-    except KeyboardInterrupt:
-        console.print("\n[dim]Stopped.[/dim]")
-
-def init_command():
-    """
-    Initialize Jazari config for this cluster.
-    Interactively asks the user for defaults and saves them to ~/.jazari/config.yaml.
-    """
-    console.print('[bold blue]Welcome to Jazari setup.[/bold blue]')
-    console.print(f"This will create a configuration file at [dim]{CONFIG_FILE}[/dim]\n")
-
-    # 1. Ask for Slurm Account
-    # We use the current config value as the default prompt if it exists.
-    default_account = CONFIG.get("account", "")
-    account = Prompt.ask(
-        "Enter your default Slurm account (e.g., def-user)",
-        default=default_account
-    )
-
-    # 2. Ask for default time limit
-    default_time = CONFIG.get("time", "01:00:00")
-    time_limit = Prompt.ask(
-        "Enter default time limit (D-HH:MM)",
-        default=default_time
-    )
-
-    # 3. Ask for W&B tracking default
-    default_wandb = CONFIG.get("track_wandb", False)
-    track_wandb = Confirm.ask(
-        "Enable W&B tracking by default?",
-        default=default_wandb
-    )
-
-    # Prepare the config dictionary
-    new_config = {
-        "account": account,
-        "time": time_limit,
-        "track_wandb": track_wandb
-    }
-
-    make_config_dir(new_config)
-
-    console.print(f"\n[bold green]✅ Configuration saved![/bold green]")
-    console.print("You can now run jobs without specifying these flags.")
-
-def run_command(
-    # The default values for these arguments are now pulled dynamically from the
-    # CONFIG dictionary using the get_config_default helper.
-    command:     list[str]      = typer.Argument(..., help="The command to run (e.g., python train.py --batch 64). Use '--' before it if it has flags."),
-    nodes:       int            = typer.Option(get_config_default("nodes", CONFIG, 1), '--nodes', '-N', help='Number of nodes.'),
-    gpus:        int            = typer.Option(get_config_default("gpus", CONFIG, 1), '--gpus', '-G', help='GPUs per node.'),
-    cpus:        int            = typer.Option(get_config_default("cpus", CONFIG, 1), '--cpus', '-c', help='CPUs per task/GPU'),
-    time:        str            = typer.Option(get_config_default("time", CONFIG, '01:00:00'), '--time', '-t', help='Time limit (D-HH:MM)'),
-    name:        str            = typer.Option('jazari_run', '--name', '-n', help='Job name'),
-    account:     Optional[str]  = typer.Option(get_config_default("account", CONFIG, None), "--account", "-A", help="Slurm account to charge."),
-    track_wandb: bool           = typer.Option(get_config_default("track_wandb", CONFIG, False), "--track-wandb", help="Auto-configure W&B."),
-    push_to_hub: Optional[str]  = typer.Option(None, "--push-to-hub", help="Hugging Face repo ID to upload model to (e.g. 'my-org/my-model)."),   
-    pull_data:   Optional[str]  = typer.Option(None, '--pull-data', help='Huggingface dataset ID to download.'),
-    dry_run:     bool           = typer.Option(False, '--dry-run', help='Print sbatch script without submitting.'),
-    ):
-    '''
-    Launch a distributed training job.
-    '''
-    if not command:
-        console.print("[bold red]Error:[/bold red] You must provide a command to run.")
-        sys.exit(1)
-
-    log_dir = "logs"
-    if not os.path.exists(log_dir):
-        os.makedirs(log_dir, exist_ok=True)
-
-    full_command_str = ' '.join(command)
-
-    # --- Weights & Biases ---
-    wandb_key = None
-    if track_wandb:
-        console.print("[dim]W&B tracking enabled.[/dim]")
-        wandb_key = get_wandb_api_key()
-        if not wandb_key:
-            console.print("[bold yellow]Warning:[/bold yellow] Could not find W&B API key locally.")
-            console.print("Please run [green]wandb login[/green] on this machine first.")
-    # -----------------
-
-    console.print(f'[bold]🐘 Generating Slurm script for: {name}[/bold]')
-
-    sbatch_content = generate_sbatch_script(
-        nodes           = nodes,
-        gpus_per_node   = gpus,
-        cpus_per_task   = cpus,
-        time_limit      = time,
-        job_name        = name,
-        account_name    = account,
-        wandb_api_key   = wandb_key,
-        hf_repo_id      = push_to_hub,
-        hf_dataset_id   = pull_data,
-        user_command    = full_command_str
-    )
-
-    if dry_run:
-        console.print('\n[yellow]--- DRY RUN: Generated #SBATCH Script ---[/yellow]')
-        syntax = Syntax(sbatch_content, 'bash', theme = 'monokai', line_numbers = True)
-        console.print(syntax)
-        console.print('[yellow]-----------------------------------------------[/yellow]')
-    else:
-        console.print("[dim]Submitting to Slurm scheduler...[/dim]")
-
-        # --- Submission ---
-        with tempfile.NamedTemporaryFile(mode='w+', suffix=".sh", delete=False) as temp_file:
-            temp_script_path = temp_file.name
-            temp_file.write(sbatch_content)
-
-        try:
-            result = subprocess.run(
-                ["sbatch", temp_script_path],
-                check=True,
-                capture_output=True,
-                text=True
-            )
-            # Extract the job ID from the output "Submitted batch job 12345"
-            job_id = result.stdout.strip().split()[-1]
-            console.print(f"[bold green]✅ Job submitted successfully![/bold green] (ID: [bold]{job_id}[/bold])")
-            console.print(f"[dim]View logs: cat logs/{name}-{job_id}.out[/dim]")
-
-        except subprocess.CalledProcessError as e:
-            console.print("[bold red]❌ Failed to submit job.[/bold red]")
-            console.print(f"Sbatch error: {e.stderr}")
-        except FileNotFoundError:
-             console.print("\n[bold red]❌ Error:[/bold red] 'sbatch' command not found.")
-             console.print("Are you running this on a Slurm login node?")
-        finally:
-            if os.path.exists(temp_script_path):
-                os.unlink(temp_script_path)
-
-# --- APP DEFINITION AND COMMAND REGISTRATION ---
-
-# Initializing the Typer app
-app = typer.Typer(
-    name="jazari",
-    help="🐘 The orchestration layer for modern Slurm clusters.",
-    add_completion=False,
-    no_args_is_help=True
-)
-
-# Explicitly register the commands
-app.command(name="status")(status_command)
-app.command(name="kill")(kill_command)
-app.command(name="logs")(logs_command)
-app.command(name="init")(init_command)
-app.command(name="run")(run_command)
-
-if __name__ == "__main__":
-    app()