scripts/launcher.py

launcher

Command-line entry point: python scripts/launcher.py <mode> config.yaml.

Stage: launch. Builds LauncherConfig (model + data) from the YAML, sets up the Lightning Trainer, W&B logging and callbacks, then trains, evaluates, caches backbone outputs or exports the model (.pt2).

LauncherParams = namedtuple('LauncherParams', ['launcher_config', 'launcher_logger']) module-attribute

Modes = Literal['train', 'evaluate', 'cache', 'compile'] module-attribute

args = _parse_args() module-attribute

launcher(mode, config, experiment_dir=None)

Run one mode (train/evaluate/cache/compile) end to end from a YAML config.

Builds config, callbacks and Trainer (GPU fp16 or CPU fp32, DDP), then dispatches. Note: training_examples_per_eval must be a multiple of the global batch size.

Source code in scripts/launcher.py

def launcher(
    mode: Modes,
    config: Mapping | str | Path,
    experiment_dir: Optional[str | Path] = None,
):
    """Run one mode (train/evaluate/cache/compile) end to end from a YAML config.

    Builds config, callbacks and Trainer (GPU fp16 or CPU fp32, DDP), then dispatches.
    Note: training_examples_per_eval must be a multiple of the global batch size.
    """
    # Reject unknown modes before any setup.
    if mode not in typing.get_args(Modes):
        raise ValueError(
            f"Unrecognized mode {mode}. Valid modes are: {typing.get_args(Modes)}."
        )
    # Build config + logger and callbacks, then log run metadata to W&B.
    params = _setup_launcher(config, experiment_dir)
    callbacks = _setup_callbacks(params.launcher_config)
    _log_to_wandb(*params)

    # Device count and eval cadence come from the config's training params.
    num_gpus = params.launcher_config.num_gpus
    training_params = params.launcher_config.training_params
    training_examples_per_eval = training_params.training_examples_per_eval

    # Validation cadence: every epoch, or every N examples (converted to batches).
    if training_examples_per_eval is None:
        val_args = dict(check_val_every_n_epoch=1, val_check_interval=None)
    elif training_examples_per_eval % training_params._global_batch_size != 0:
        raise ValueError(
            f"Invalid training_examples_per_eval: {training_examples_per_eval}"
            f" is not divisible by global_batch_size = "
            f"{training_params._global_batch_size}"
        )
    else:
        val_args = dict(
            check_val_every_n_epoch=None,
            val_check_interval=training_examples_per_eval
            // training_params._global_batch_size,
        )

    # Trainer: GPU/CPU, multi-node DDP, grad accumulation, cadence from above.
    trainer = Trainer(
        accelerator="gpu" if num_gpus > 0 else "cpu",
        num_nodes=int(os.environ.get("NUM_NODES", 1)),
        # HARDCODED: DDP with find_unused_parameters=True; fp16 on GPU, fp32 on CPU.
        strategy=DDPStrategy(
            find_unused_parameters=True
        ),  # LoRA seems to need this set to True
        precision=16 if num_gpus > 0 else 32,
        logger=params.launcher_logger.wandb_logger
        if params.launcher_logger is not None
        else None,
        accumulate_grad_batches=training_params.total_accum_grad_batches,
        max_epochs=params.launcher_config.training_params.max_epochs,
        log_every_n_steps=5,
        devices=num_gpus if num_gpus > 0 else "auto",
        callbacks=callbacks,
        **val_args,
    )

    # Dispatch to the requested mode; compile does not use the Trainer.
    if mode == "train":
        run_training(trainer, params.launcher_config)
    elif mode == "evaluate":
        run_evaluation(trainer, params.launcher_config)
    elif mode == "cache":
        run_caching(trainer, params.launcher_config, params.launcher_logger)
    elif mode == "compile":
        run_compilation(params.launcher_config)

    # Log the metadata tables of the splits actually used.
    if mode != "compile":
        _log_datasets_to_wandb(*params)

run_training(trainer, launcher_config)

Fit on train/val, reload the best checkpoint, then test it on the test split.

Source code in scripts/launcher.py

def run_training(trainer: Trainer, launcher_config: LauncherConfig):
    """Fit on train/val, reload the best checkpoint, then test it on the test split."""
    # One loader per split; only train is in training mode.
    train_loader = launcher_config.get_data_loader("train", training_mode=True)
    val_loader = launcher_config.get_data_loader("val", training_mode=False)
    test_loader = launcher_config.get_data_loader("test", training_mode=False)
    # Fit on train, validating on val; optionally resume a saved training state.
    trainer.fit(
        launcher_config.model,
        train_loader,
        val_loader,
        ckpt_path=launcher_config.restore_training_state_ckpt_path,
    )

    #
    # Load the best checkpoint's weights, then evaluate on the test split.
    launcher_config.model.load_state_dict(
        torch.load(trainer.checkpoint_callback.best_model_path)["state_dict"]
    )
    trainer.test(launcher_config.model, test_loader)

run_evaluation(trainer, launcher_config)

Test a trained checkpoint (model_params.ckpt_path) on the test split.

Source code in scripts/launcher.py

def run_evaluation(trainer: Trainer, launcher_config: LauncherConfig):
    """Test a trained checkpoint (model_params.ckpt_path) on the test split."""
    if launcher_config.model_params.ckpt_path == "":
        raise ValueError("Need to set a checkpoint path (`ckpt_path`) for evaluation.")

    # Run the test split with the checkpointed model.
    test_loader = launcher_config.get_data_loader("test", training_mode=False)
    trainer.test(
        launcher_config.model,
        test_loader,
        ckpt_path=launcher_config.restore_training_state_ckpt_path,
    )

run_caching(trainer, launcher_config, launcher_logger)

Run the backbone over train/val/test and store its outputs in W&B (cache mode).

Requires GPUs and no backbone_cache in data_params; a dummy predict() starts DDP.

Source code in scripts/launcher.py

def run_caching(
    trainer: Trainer, launcher_config: LauncherConfig, launcher_logger: LauncherLogger
):
    """Run the backbone over train/val/test and store its outputs in W&B (cache mode).

    Requires GPUs and no backbone_cache in data_params; a dummy predict() starts DDP.
    """
    # prohibit running caching if cache is set
    assert "backbone_cache" not in launcher_config.config["data_params"]

    # Collect dataset and loader kwargs (eval mode) for each split.
    datasets, data_loader_args = [], []
    for split in ["train", "val", "test"]:
        datasets.append(launcher_config.get_dataset(split))
        data_loader_args.append(launcher_config.get_data_loader_kwargs(split, False))

    # Warning: before running caching procedure, the distributed environment
    # should be properly initialized. It happens under the hood when we call
    # trainer.fit(), .test() or .predict(), so one of these methods should be called
    # prior to caching. For this reason we call .predict() with 1 sample per gpus.
    dummy_loader = torch.utils.data.DataLoader(datasets[-1], **data_loader_args[-1])
    samples = [next(iter(dummy_loader))] * torch.cuda.device_count()
    sample_loader = torch.utils.data.DataLoader(
        samples, batch_size=1, collate_fn=lambda x: x[0]
    )
    trainer.predict(
        launcher_config.model,
        sample_loader,
        ckpt_path=launcher_config.restore_training_state_ckpt_path,
    )

    assert torch.distributed.is_initialized()

    # replace sampler with the distributed one
    for dataset, arg in zip(datasets, data_loader_args):
        arg["sampler"] = UuidSampler(dataset.datasets["metadata"], distributed=True)

    # Move the model to this rank's GPU in eval mode.
    launcher_config.model.cuda(torch.distributed.get_rank()).eval()

    # Save backbone weights and cache its features for all splits to W&B.
    state_dict = launcher_config.model.backbone.state_dict(prefix="backbone.")
    cache_backbone(
        launcher_config.model.compute_features_to_cache,
        launcher_logger.wandb_logger if launcher_logger else None,
        datasets,
        data_loader_args,
        state_dict,
    )

run_compilation(launcher_config)

Export the model with torch.export on one test window, check outputs, log .pt2 to W&B.

Note: uses only the audio feature; needs an active wandb.run.

Source code in scripts/launcher.py

def run_compilation(launcher_config: LauncherConfig):
    """Export the model with torch.export on one test window, check outputs, log .pt2 to W&B.

    Note: uses only the audio feature; needs an active wandb.run.
    """

    # we need exactly one sample for running a compiler
    dummy_dataset = launcher_config.get_dataset("test")
    dummy_data_loader = launcher_config.get_data_loader_kwargs("test", False)
    dummy_loader = torch.utils.data.DataLoader(dummy_dataset, **dummy_data_loader)
    batch = next(iter(dummy_loader))
    sample_input = ({"audio": batch["features"]["audio"][:1, ...]}, [1])

    #
    # Export in eval mode, using the one-window sample as the example input.
    launcher_config.model.eval()
    model_compiled = torch.export.export(launcher_config.model, sample_input)

    # sanity check
    output_original, _ = launcher_config.model(*sample_input)
    output_compiled, _ = model_compiled.module()(*sample_input)
    for key in output_compiled:
        if not torch.allclose(output_compiled[key], output_original[key]):
            raise ValueError("Computation results don't match")

    #
    # Save the .pt2 in a temp dir, upload as a W&B model artifact, then clean up.
    temp_dir = tempfile.mkdtemp()
    path = os.path.join(temp_dir, "model.pt2")

    torch.export.save(model_compiled, path)

    artifact = wandb.Artifact(name=f"model-{wandb.run.id}", type="model")
    artifact.add_file(path)
    wandb.log_artifact(artifact)

    shutil.rmtree(temp_dir, ignore_errors=True)

setup_multi_node_environment()

Wait for all Kubernetes job pods, then set MASTER_ADDR/PORT, WORLD_SIZE, NODE_RANK.

Source code in scripts/launcher.py

def setup_multi_node_environment():
    """Wait for all Kubernetes job pods, then set MASTER_ADDR/PORT, WORLD_SIZE, NODE_RANK."""

    # this is a bit hacky.
    # kubernetes package is only required inside docker container worker,
    # not necessarily in the kipy
    from kubernetes import client, config

    #
    # Use the pod's in-cluster credentials to query the Kubernetes API.
    config.load_incluster_config()
    core_v1 = client.CoreV1Api()

    #
    # Job name and expected pod count come from the environment.
    job_name = os.environ["JOB_NAME"]
    num_nodes = int(os.environ.get("NUM_NODES", 1))

    #
    # Poll until every pod has an IP; abort after the timeout.
    start_time = time.time()
    while True:

        # get pods associated with this job
        pods = core_v1.list_namespaced_pod(
            # HARDCODED: Kubernetes namespace "default".
            namespace="default", label_selector=f"job-name={job_name}"
        )

        # Pods that already have an IP count as started.
        active_nodes = sum(pod.status.pod_ip is not None for pod in pods.items)
        if active_nodes == num_nodes:
            break

        print(f"{active_nodes}/{num_nodes} nodes are ready. Waiting ...", flush=True)

        # we give it at most 20 mins to warmup.
        # In most cases it should be less than few mins.
        if time.time() - start_time > 20 * 60:
            print(
                f"Max wait time has been exceeded. Target num nodes {num_nodes} has not been reached. Shutting down."
            )
            exit(1)

        time.sleep(10)

    # setup env variables for multi-node training
    # node with JOB_COMPLETION_INDEX == 0 is a master node
    os.environ["MASTER_ADDR"] = [
        pod.status.pod_ip
        for pod in pods.items
        if pod.metadata.annotations["batch.kubernetes.io/job-completion-index"] == "0"
    ][0]
    # HARDCODED: DDP master port 12355.
    os.environ["MASTER_PORT"] = "12355"  # some arbitrary recommended port
    os.environ["WORLD_SIZE"] = str(int(num_nodes * torch.cuda.device_count()))
    os.environ["NODE_RANK"] = os.environ["JOB_COMPLETION_INDEX"]