Add LeptonExecutor support

[roclark]

The LeptonExecutor adds support for launching jobs on NVIDIA DGX Cloud Lepton using the Lepton Python SDK. The Lepton SDK needs to be installed with pip install nemo-run[lepton] or pip install leptonai and authenticate with the cluster using the command available on the DGX Cloud Lepton UI in the Settings > Tokens page.

The LeptonExecutor can be defined using the following example:

def your_lepton_executor(nodes: int, gpus_per_node: int, container_image: str):
    # Ensure these are set correctly for your DGX Cloud environment
    # You might fetch these from environment variables or a config file
    resource_shape = "gpu.8xh100-80gb" # Replace with your desired resource shape representing the number of GPUs in a pod
    node_group = "my-node-group"  # The node group to run the job in
    nemo_run_dir = "/nemo-workspace/nemo-run"  # The NeMo-Run directory where experiments are saved
    # Define the remote storage directory that will be mounted in the job pods
    # Ensure the path specified here contains your NEMORUN_HOME
    storage_path = "/nemo-workspace" # The remote storage directory to mount in jobs
    mount_path = "/nemo-workspace" # The path where the remote storage directory will be mounted inside the container

    executor = run.LeptonExecutor(
        resource_shape=resource_shape,
        node_group=node_group,
        container_image=container_image,
        nodes=nodes,
        nemo_run_dir=nemo_run_dir,
        gpus_per_node=gpus_per_node,
        mounts=[{"path": storage_path, "mount_path": mount_path}],
        # Optional: Add custom environment variables or PyTorch specs if needed
        env_vars=common_envs(),
        # packager=run.GitArchivePackager() # Choose appropriate packager
    )
    return executor

# Example usage:
executor = your_lepton_executor(nodes=4, gpus_per_node=8, container_image="your-nemo-image")

Launching jobs with the LeptonExecutor will queue a Batch Job on the DGX Cloud Lepton cluster in the specified node group.

[hemildesai]

Mostly LGTM, can you make sure the format, lint and test checks pass?

[hemildesai]

Also, can you add some basic tests or create an issue to add tests for the lepton executor later?

[roclark]

Thanks for the review, @hemildesai! I will clear up the failing tests here and add some new tests with the PR early next week. I changed up some more functionality today based on feedback I got offline on the storage API as well.

[roclark]

@hemildesai, added some tests and cleared up everything else that was failing in the PR. Let me know if there's anything else I need to change, thanks! :)

[hemildesai]

LGTM, just one small comment.

[roclark]

Alright, should be resolved now.

[hemildesai]

Looks great, thanks 🎉