Merge pull request #1819 from transformerlab/add/lab-task-init-cli

deep1401 · web-flow · commit 7a6f2d51e83a · 2026-04-21T12:48:42.000-06:00
Add lab task init to CLI
diff --git a/.agents/skills/transformerlab-cli/SKILL.md b/.agents/skills/transformerlab-cli/SKILL.md
@@ -102,6 +102,20 @@ lab job download JOB_ID --file "*.csv" -o ./results
 
 ## Creating Tasks
 
+### Scaffold a new task with `lab task init`
+
+**When the user asks to create, initialize, or start a new task, always use `lab task init`** rather than writing `task.yaml` / `main.py` by hand. It scaffolds both files with sensible defaults in the current directory so the user has a working starting point.
+
+```bash
+mkdir my-task && cd my-task
+lab task init            # writes task.yaml + main.py with defaults (skips existing files)
+lab task init --interactive   # prompts for name, CPUs, memory, setup, and run command
+```
+
+- Default mode is non-interactive. It creates `task.yaml` (with `name`, `resources: {cpus: 2, memory: 4}`, and `run: python main.py`) and a starter `main.py`. Existing files are skipped, not overwritten.
+- `--interactive` writes only `task.yaml` (no `main.py`) and prompts for each field. In this mode `task.yaml` will prompt before overwrite.
+- After init, edit `main.py`, customize `task.yaml`, then run `lab task add .` to upload it.
+
 ### task.yaml Structure
 
 Full docs: https://lab.cloud/for-teams/running-a-task/task-yaml-structure
@@ -291,6 +305,7 @@ This applies to launching jobs, fetching logs, checking cluster status, and ever
 | `lab version` | Show CLI version | No |
 | `lab task list` | List tasks in current experiment | Yes |
 | `lab task info <id>` | Get task details | Yes |
+| `lab task init` | Scaffold `task.yaml` + `main.py` in the current directory (`--interactive` to prompt) | No |
 | `lab task add [dir]` | Add task from directory or `--from-git` URL (`--no-interactive`, `--dry-run`) | Yes |
 | `lab task delete <id>` | Delete a task (`--no-interactive` to skip confirmation) | Yes |
 | `lab task queue <id>` | Queue task on compute provider | Yes |
diff --git a/.agents/skills/transformerlab-cli/references/commands.md b/.agents/skills/transformerlab-cli/references/commands.md
@@ -98,6 +98,17 @@ Get details for a specific task.
 {"id": 1, "name": "my-task", "type": "TRAINING", "config": {...}, ...}
 ```
 
+### `task init`
+
+Scaffold a new task in the current directory. **Prefer this over writing `task.yaml` by hand.** Does not require an experiment to be set.
+
+| Option | Description |
+|---|---|
+| (default, no flag) | Non-interactive. Creates `task.yaml` with defaults (`name` = folder name, `cpus: 2`, `memory: 4`, `run: python main.py`) and a starter `main.py`. Existing files are skipped, not overwritten. |
+| `--interactive` | Prompts for task name, CPUs, memory, accelerators, setup, and run command. Writes only `task.yaml` (no `main.py`). Prompts before overwriting an existing `task.yaml`. |
+
+**JSON output (default):** `{"created": ["task.yaml", "main.py"], "skipped": [], "path": "/abs/dir"}`
+
 ### `task add [directory]`
 
 Add a new task from a local directory containing `task.yaml`, or from a Git repository.
diff --git a/cli/pyproject.toml b/cli/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab-cli"
-version = "0.0.15"
+version = "0.0.16"
 description = "Transformer Lab CLI"
 requires-python = ">=3.10"
 authors = [{ name = "Transformer Lab", email = "hello@transformerlab.ai" }]
@@ -43,7 +43,7 @@ lab = "transformerlab_cli.main:app"
 where = ["src"]
 
 [tool.setuptools.package-data]
-"*" = ["*.tcss", "*.css"]
+"*" = ["*.tcss", "*.css", "templates/task_init/*"]
 
 [dependency-groups]
 dev = ["pytest", "textual>=6.11.0", "typer[testing]"]
diff --git a/cli/src/transformerlab_cli/commands/task.py b/cli/src/transformerlab_cli/commands/task.py
@@ -2,6 +2,7 @@
 import json
 import os
 import zipfile
+from pathlib import Path
 
 import httpx
 import typer
@@ -234,6 +235,159 @@ def command_task_list():
     list_tasks(output_format=cli_state.output_format, experiment_id=current_experiment)
 
 
+TASK_INIT_TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates" / "task_init"
+
+
+def _render_task_yaml_template(task_name: str) -> str:
+    template = (TASK_INIT_TEMPLATES_DIR / "task.yaml").read_text(encoding="utf-8")
+    return template.replace("{{TASK_NAME}}", task_name)
+
+
+def _main_py_template() -> str:
+    return (TASK_INIT_TEMPLATES_DIR / "main.py").read_text(encoding="utf-8")
+
+
+def _write_task_yaml(path: str, data: dict) -> None:
+    yaml_text = yaml.safe_dump(data, sort_keys=False, default_flow_style=False)
+    with open(path, "w", encoding="utf-8", newline="\n") as f:
+        f.write(yaml_text)
+
+
+def _print_next_steps(include_main_py: bool) -> None:
+    console.print("\nNext steps:")
+    if include_main_py:
+        console.print("- Edit [bold]main.py[/bold] with your task code")
+    console.print("- Customize [bold]task.yaml[/bold] (resources, setup, parameters)")
+    console.print("- Run: [bold]lab task add .[/bold]")
+    console.print("- Docs: https://lab.cloud/for-teams/running-a-task/task-yaml-structure")
+
+
+def _task_init_default(task_yaml_path: str, main_py_path: str, folder_name: str) -> None:
+    if os.path.exists(task_yaml_path):
+        if cli_state.output_format == "json":
+            print(json.dumps({"error": "task.yaml already exists"}))
+        else:
+            console.print(
+                f"[error]Error:[/error] [bold]{task_yaml_path}[/bold] already exists. "
+                "Refusing to overwrite. Remove it first or run `lab task init` in an empty directory."
+            )
+        raise typer.Exit(1)
+
+    with open(task_yaml_path, "w", encoding="utf-8", newline="\n") as f:
+        f.write(_render_task_yaml_template(folder_name))
+
+    main_py_existed = os.path.exists(main_py_path)
+    if not main_py_existed:
+        with open(main_py_path, "w", encoding="utf-8", newline="\n") as f:
+            f.write(_main_py_template())
+
+    if cli_state.output_format == "json":
+        created = ["task.yaml"] if main_py_existed else ["task.yaml", "main.py"]
+        skipped = ["main.py"] if main_py_existed else []
+        print(json.dumps({"created": created, "skipped": skipped, "path": os.path.dirname(task_yaml_path)}))
+        return
+
+    console.print("[success]✓[/success] Created [bold]task.yaml[/bold]")
+    if main_py_existed:
+        console.print("[warning]•[/warning] Skipped [bold]main.py[/bold] (already exists)")
+    else:
+        console.print("[success]✓[/success] Created [bold]main.py[/bold]")
+
+    console.print(f"\nLocation: [bold]{os.path.dirname(task_yaml_path)}[/bold]")
+    _print_next_steps(include_main_py=not main_py_existed)
+
+
+def _task_init_interactive(task_yaml_path: str, folder_name: str) -> None:
+    if os.path.exists(task_yaml_path):
+        if cli_state.output_format == "json":
+            print(json.dumps({"error": "task.yaml already exists"}))
+            raise typer.Exit(1)
+        should_overwrite = typer.confirm("task.yaml already exists. Overwrite?", default=False)
+        if not should_overwrite:
+            console.print("[warning]Cancelled.[/warning]")
+            raise typer.Exit(0)
+
+    task_name = typer.prompt("Task name", default=folder_name).strip() or folder_name
+
+    cpus = typer.prompt("CPUs", default="2").strip()
+    memory = typer.prompt("Memory (GB)", default="4").strip()
+    accelerators = typer.prompt("Accelerators (optional)", default="", show_default=False).strip()
+
+    setup = ""
+    run = ""
+
+    if cli_state.output_format != "json" and os.isatty(0) and os.isatty(1):
+        edited = typer.edit(
+            "\n".join(
+                [
+                    "# Define the commands for your task below.",
+                    "# This YAML snippet will be parsed and merged into task.yaml.",
+                    "",
+                    "setup: |",
+                    "  # Optional: install deps, download data, etc.",
+                    "  ",
+                    "run: |",
+                    "  # Required: the main command to execute",
+                    "  ",
+                    "",
+                ]
+            )
+        )
+        if edited:
+            try:
+                edited_obj = yaml.safe_load(edited)
+                if isinstance(edited_obj, dict):
+                    setup_val = edited_obj.get("setup")
+                    run_val = edited_obj.get("run")
+                    if isinstance(setup_val, str):
+                        setup = setup_val.rstrip()
+                    if isinstance(run_val, str):
+                        run = run_val.rstrip()
+            except yaml.YAMLError:
+                pass
+
+    if not setup.strip():
+        setup = typer.prompt("Setup command (optional)", default="", show_default=False).rstrip()
+
+    while not run.strip():
+        run = typer.prompt("Run command", default="", show_default=False).rstrip()
+
+    task_yaml: dict = {
+        "name": task_name,
+        "resources": {"cpus": cpus, "memory": memory},
+        "run": run,
+    }
+    if accelerators:
+        task_yaml["resources"]["accelerators"] = accelerators
+    if setup.strip():
+        task_yaml["setup"] = setup
+
+    _write_task_yaml(task_yaml_path, task_yaml)
+
+    if cli_state.output_format == "json":
+        print(json.dumps({"path": task_yaml_path}))
+        return
+
+    console.print(f"[success]✓[/success] Wrote [bold]{task_yaml_path}[/bold]")
+    _print_next_steps(include_main_py=False)
+
+
+@app.command("init")
+def command_task_init(
+    interactive: bool = typer.Option(False, "--interactive", help="Prompt for task settings instead of using defaults"),
+):
+    """Initialize a task.yaml and main.py in the current directory."""
+    cwd = os.getcwd()
+    task_yaml_path = os.path.join(cwd, "task.yaml")
+    main_py_path = os.path.join(cwd, "main.py")
+    folder_name = os.path.basename(cwd).strip() or "my-task"
+
+    if interactive:
+        _task_init_interactive(task_yaml_path, folder_name)
+    else:
+        _task_init_default(task_yaml_path, main_py_path, folder_name)
+
+
 @app.command("add")
 def command_task_add(
     task_directory: str = typer.Argument(None, help="Path to the task directory containing task.yaml"),
diff --git a/cli/src/transformerlab_cli/templates/task_init/main.py b/cli/src/transformerlab_cli/templates/task_init/main.py
@@ -0,0 +1,76 @@
+"""Sample Transformer Lab task.
+
+This file demonstrates the core Lab SDK calls you will use in a real task.
+Replace the sleep loop with your actual training / inference / eval code.
+
+Common SDK calls:
+  - lab.init()                    attach this process to the job
+  - lab.log(message)              stream a line to the task log
+  - lab.update_progress(n)        set job progress (0..100)
+  - lab.get_config()              read `parameters:` from task.yaml
+  - lab.save_artifact(path, name) attach a file to the job (summaries,
+                                  plots, CSVs, etc.) for later download
+  - lab.save_checkpoint(dir, name) persist a training checkpoint
+  - lab.save_model(dir, name=...)  register a final trained model
+  - lab.finish(message)            mark the job SUCCESS
+  - lab.error(message)             mark the job FAILED
+
+Docs: https://lab.cloud/for-teams/running-a-task/lab-sdk
+"""
+
+import json
+import time
+from pathlib import Path
+
+from lab import lab
+
+
+def main():
+    # Always call init() first. Without it, log / progress / artifact
+    # calls will not be associated with the job.
+    lab.init()
+
+    try:
+        # Read parameters defined under `parameters:` in task.yaml.
+        # Returns a dict; use .get() with a default so the task still
+        # runs when a parameter is missing.
+        config = lab.get_config() or {}
+        total_steps = int(config.get("total_steps", 5))
+
+        lab.log(f"Starting task with config: {config}")
+        lab.update_progress(5)
+
+        # ----- Fake training loop -----
+        # Replace this block with your real work. In a real task you
+        # would typically also call:
+        #   lab.save_checkpoint(checkpoint_dir, f"checkpoint-{step}")
+        # every N steps so the job can resume on retry.
+        for step in range(1, total_steps + 1):
+            time.sleep(1)  # pretend we are doing something expensive
+            progress = 5 + int(step / total_steps * 90)  # 5..95
+            lab.log(f"step {step}/{total_steps}")
+            lab.update_progress(progress)
+
+        # ----- Save an artifact -----
+        # Artifacts are any files you want attached to the job for later
+        # download (metrics, plots, CSVs, etc.). The second argument is
+        # the name the artifact appears under in the UI / CLI.
+        summary_path = Path("training_summary.json")
+        summary_path.write_text(json.dumps({"steps": total_steps, "ok": True}, indent=2))
+        lab.save_artifact(str(summary_path), "training_summary.json")
+
+        # If you trained a model, you would also call:
+        #   lab.save_model(model_dir, name="my_trained_model")
+
+        lab.update_progress(100)
+        lab.finish("Task completed successfully")
+
+    except Exception as e:
+        # Any unhandled exception should be reported via lab.error so
+        # the job is marked FAILED instead of being stuck in RUNNING.
+        lab.error(f"Task failed: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()
diff --git a/cli/src/transformerlab_cli/templates/task_init/task.yaml b/cli/src/transformerlab_cli/templates/task_init/task.yaml
@@ -0,0 +1,28 @@
+# Transformer Lab task configuration
+# Docs: https://lab.cloud/for-teams/running-a-task/task-yaml-structure
+
+name: {{TASK_NAME}}
+
+# Compute requested for this task. Memory and disk_space are in GB.
+resources:
+  cpus: 2
+  memory: 4
+  # accelerators: "H100:1"   # e.g. 1x H100 GPU; format is "TYPE:COUNT"
+  # disk_space: 100
+  # num_nodes: 1             # >1 for distributed training
+
+# Commands that run on the remote machine BEFORE `run`. Use this for
+# installing deps, downloading data, etc.
+setup: |
+  pip install transformerlab
+
+# Main entry point. This is the command that executes your task.
+run: python main.py
+
+# Parameters are available inside main.py via `lab.get_config()`.
+# parameters:
+#   learning_rate: 0.001
+#   batch_size: 32
+
+# envs:
+#   HF_TOKEN: "${HF_TOKEN}"
diff --git a/cli/tests/commands/test_task_init.py b/cli/tests/commands/test_task_init.py
