{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "\n", "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n", "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n", "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n", "
\n", "\n", "👋 Welcome to Open Universal Machine Intelligence (Oumi)!\n", "\n", "🚀 Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](https://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n", "\n", "🤝 Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n", "\n", "⭐ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Running Jobs Remotely" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Introduction\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In addition to our local training loop, Oumi provides the [launcher module](https://github.com/oumi-ai/oumi/tree/main/src/oumi/launcher) as a simple interface for kicking off jobs on a wide variety of remote hardware. We support various cloud providers (GCP, Runpod, Lambda) out of the box, with the additional flexibility to support your own custom cluster should the need arise! In this tutorial we will focus on running jobs using GCP, but this tutorial applies to all clouds Oumi supports. You can read more about the launcher API [here](https://github.com/oumi-ai/oumi/blob/main/src/oumi/launcher/launcher.py).\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Prerequisites" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This tutorial assumes:\n", "- You have a valid Google Cloud Platform (GCP) project with billing enabled.\n", "- Your GCP project has the `Compute Engine API` enabled.\n", "- You have the following IAM permissions in your project:\n", " - ```bash\n", " roles/browser\n", " roles/compute.admin\n", " roles/serviceusage.serviceUsageConsumer\n", " roles/storage.admin\n", " ```\n", "\n", "You must also authenticate with GCP locally before starting this tutorial:\n", "\n", "```bash\n", "conda install -c conda-forge google-cloud-sdk -y\n", "gcloud init\n", "# Run this if you don't have a credentials file.\n", "# This will generate ~/.config/gcloud/application_default_credentials.json.\n", "gcloud auth application-default login\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Oumi Installation\n", "\n", "First, let's install Oumi. You can find more detailed instructions [here](https://oumi.ai/docs/en/latest/get_started/installation.html)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%pip install oumi" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Creating a Job" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The Oumi Launcher operates using three key concepts:\n", "\n", "1) `Jobs`: A `job` is a unit of work, such as running training or model evaluation. This can be any script you'd like!\n", "2) `Clusters`: A `cluster` is a set of dedicated hardware upon which `jobs` are run. A `cluster` could be as simple as a cloud VM environment.\n", "3) `Clouds` : A `cloud` is a resource provider that manages `clusters`. These include GCP, AWS, Lambda, Runpod, etc.\n", "\n", "When you submit a job to the launcher it will handle queueing your job in the proper cluster's job queue. If your desired Cloud does not have an appropriate cluster for running your job it will try to create one on the fly!\n", "\n", "Start by creating a simple job:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import oumi.launcher as launcher\n", "\n", "job_name = \"Create_a_display_name_for_your_job\"\n", "cloud_name = \"gcp\"\n", "\n", "job = launcher.JobConfig(\n", " name=job_name,\n", " working_dir=\".\",\n", " setup=\"\",\n", " run=\"\",\n", " resources=launcher.JobResources(\n", " # We're using Google Cloud Platform in this example.\n", " cloud=cloud_name,\n", " ),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Congratulations on creating your first job!\n", "\n", "Right now your job has an empty `run` field meaning it won't execute any code at runtime. Let's fix that by adding a few simple echo statements. It's important to note that all lines of `run` will be executed on your cluster directly in the shell--but more on that later." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "env_vars = {\n", " \"TEST_ENV_VARIABLE\": '\"Hello, World!\"',\n", "}\n", "job.envs = env_vars\n", "\n", "run_script = \"\"\"\n", "echo \"$TEST_ENV_VARIABLE\"\n", "\"\"\"\n", "\n", "job.run = run_script" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's also populate `setup`. Like `run`, `setup` is executed in the shell on the cluster. However, for most clouds `setup` is only executed when a cluster is created for the first time. This is where you should `pip install` any dependencies needed by your job's `run` script.\n", "\n", "For now, let's add a simple echo statement:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "setup_script = \"\"\"\n", "echo \"This is a script to help set up your environment for your job.\"\n", "echo \"On most clouds, this is only run during cluster creation.\"\n", "\"\"\"\n", "\n", "job.setup = setup_script" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Running your Job" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now that you have a job, it's time to run it on a cluster. You can use `launcher.up(...)` to launch your job on a cluster. If you don't have any clusters set up yet, the launcher will make a best-effort at spinning up a cluster that meets the requirements you set in your job's `JobResources`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cluster_name = \"your_cluster_name\"\n", "\n", "# If you specify an existing cluster name the launcher will use that cluster.\n", "# Otherwise the launcher will create a new cluster with the specified name.\n", "cluster, job_status = launcher.up(job, cluster_name)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "> You'll notice that the logs from the previous command reference Sky. Individual clouds / clusters in the Oumi launcher may use different libraries for communication and job orchestration. At the time of writing, the GCP cloud implementation leverages Sky Pilot.\n", "\n", "You can get the latest status of your job by querying the job status on your cluster:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "latest_status = cluster.get_job(job_status.id)\n", "\n", "print(latest_status)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "And list the status of all jobs across all clouds and clusters:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "status_list = launcher.status()\n", "\n", "print(status_list)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Another handy utility is the ability to list all active clusters for a cloud. Your new cluster will appear in this list:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "clusters = launcher.get_cloud(cloud_name).list_clusters()\n", "\n", "cluster_names = [cluster.name() for cluster in clusters]\n", "print(cluster_names)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that `launcher.get_cloud(cloud_name)` returned a `BaseCloud` object. You can learn more about the `Cloud` API [here](https://github.com/oumi-ai/oumi/blob/main/src/oumi/core/types/base_cloud.py).\n", "\n", "\n", "You can learn more about the `Cluster` API [here](https://github.com/oumi-ai/oumi/blob/main/src/oumi/core/types/base_cluster.py#L28)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Canceling a Job" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Running jobs can be quickly canceled by using the `cancel` method:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Only run this cell if you want to cancel your job!\n", "final_status = launcher.cancel(job_status.id, cloud_name, job_status.cluster)\n", "\n", "print(final_status)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Cleaning Up" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "After your job is done, make sure you don't forget to turn down your cluster! Most cloud providers will bill you for the time that your cluster is up, whether or not it is actively running jobs:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Cluster names are only unique within a Cloud.\n", "# Specify both the cloud and the cluster you'd like to turn down\n", "launcher.down(cloud_name, cluster_name)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Running a Remote Training Job" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Running training jobs on a remote cluster is simple. Before getting started, we strongly suggest you take a look at our tuning tutorial to learn the ropes of Oumi training.\n", "\n", "You can apply the same methods for local training to a remote job. The following job is a sample script for training Llama-2b on GCP. A few important notes:\n", "\n", "- If you use `${ENV_VAR}` interpolation in your `setup` or `run` script, they must be delimited. e.g. `${ENV_VAR}` -> `\\${ENV_VAR}`\n", "- The job assumes it was kicked off in a `working_dir=.` that contains the Oumi repository. You will see references to local paths like `./configs/examples/misc/sky_init.sh`, etc." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "job_config = launcher.JobConfig(\n", " name=\"llama-2b\",\n", " working_dir=\"..\",\n", " file_mounts={\n", " \"~/.netrc\": \"~/.netrc\" # WandB credentials\n", " },\n", " envs={\n", " \"ACCELERATE_LOG_LEVEL\": \"info\",\n", " },\n", " resources=launcher.JobResources(\n", " # Run on Google Cloud Platform\n", " cloud=\"gcp\",\n", " # Use 4 A100 GPUs\n", " accelerators=\"A100:4\",\n", " ),\n", " setup=\"\"\"\n", "set -e\n", "pip install uv && uv pip install --system oumi[gpu]\n", "\"\"\",\n", " run=\"\"\"\n", "set -e # Exit if any command failed.\n", "\n", "# Run some checks, and export \"OUMI_*\" env vars\n", "source ./configs/examples/misc/sky_init.sh\n", "\n", "set -x\n", "oumi distributed torchrun \\\n", " -m oumi train \\\n", " -c configs/examples/fineweb_ablation_pretraining/fsdp/train.yaml \\\n", " --training.max_steps 20 \\\n", " --training.save_steps 0 \\\n", " --training.save_final_model false\n", "\n", "echo \"Node \\\\${SKYPILOT_NODE_RANK} is all done!\"\n", "\"\"\",\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can kick off this job just as you did before. Note that it requires a cluster with 4 A100 GPUs. You can uncomment the following command and run it to start this training job on GCP:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Uncomment the following line to run training\n", "# cluster, job_status = launcher.up(job_config, \"llama-2b-cluster\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To view your job logs, run the following:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!sky logs llama-2b-cluster" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To turn down your cluster when you're done, run:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "launcher.down(cloud_name, \"llama-2b-cluster\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Advanced Fields" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `JobConfig` used to define a job contains many fields we didn't cover above. See the following definitions to better understand how to set up resourcing for your jobs:\n", "\n", "#### JobConfig\n", "\n", "| **Field Name** | **Type** | **Description** |\n", "|-----------------|-------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n", "| name | Optional[str] (default=None) | The display name of the job. Used for display purposes for most clouds. |\n", "| user | Optional[str] (default=None) | Only used for the `Polaris` cloud. The user that the job will run as. |\n", "| working_dir | str (required) | The local directory containing scripts required to execute the job. This directory will be copied to the remote node. |\n", "| num_nodes | int (required, default=1) | The number of nodes (compute instances) to use for the job. Used during cluster creation. |\n", "| resources | JobResources (required) | The resources required for each node in the job. |\n", "| envs | Dict[str, str] (required, default={}) | The environment variables to set before running the job. |\n", "| file_mounts | Dict[str, str] (required, default={}) | File mounts to attach to the node. For mounting (copying) local directories. The key is the remote path, and the value is the local path. Cannot share a key with `storage_mounts`|\n", "| storage_mounts | Dict[str, StorageMount] (required, default={}) | Storage systems to attach to the node. The key is the remote path, and the value is the storage system to mount. Cannot share a key with `file_mounts` |\n", "| setup | Optional[str] (default=None) | The setup script to run before the job starts. For most clouds this is executed only on cluster creation. ex) `pip install -r requirements.txt` |\n", "| run | str (required) | The script to run on the remote cluster. |\n", "\n", "#### StorageMount\n", "\n", "| **Field Name** | **Type** | **Description** |\n", "|-----------------|-------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n", "| source | str (required) | The remote path to mount the local path to. e.g. 'gs://bucket/path' for GCS, 's3://bucket/path' for S3, or 'r2://path' for R2. |\n", "| store | str (required) | The remote storage solution (Required). Must be one of 's3', 'gcs' or 'r2'. |\n", "\n", "#### JobResources\n", "\n", "| **Field Name** | **Type** | **Description** |\n", "|-----------------|-------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n", "| cloud | str (required) | The cloud used to run the job. |\n", "| region | Optional[str] (default=None) | The region to use (optional). Supported values vary by environment. |\n", "| zone | Optional[str] (default=None) | The zone to use (optional). Supported values vary by environment. |\n", "| accelerators | Optional[str] (default=None) | Accelerator type (optional). Supported values vary by environment. For GCP you may specify the accelerator name and count, e.g. \"V100:4\". |\n", "| cpus | Optional[str] (default=None) | Number of vCPUs to use per node (optional). Sky-based clouds support strings with modifiers, e.g. \"2+\" to indicate at least 2 vCPUs. |\n", "| memory | Optional[str] (default=None) | Memory to allocate per node in GiB (optional). Sky-based clouds support strings with modifiers, e.g. \"256+\" to indicate at least 256 GB. |\n", "| instance_type | Optional[str] (default=None) | Instance type to use (optional). Supported values vary by environment. The instance type is automatically inferred if `accelerators` is specified. |\n", "| use_spot | bool (required, default=False) | Whether the cluster should use spot instances. If unspecified, defaults to False (on-demand instances). |\n", "| disk_size | Optional[int] (default=None) | Disk size in GiB to allocate for OS (mounted at /). Ignored by Polaris. Optional. |\n", "| disk_tier | Optional[str] (default=None) | Disk tier to use for OS (optional). For sky-based clouds this Could be one of 'low', 'medium', 'high' or 'best'. Defaults to 'medium'. Ignored by Polaris. |" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# 🧭 What's Next?\n", "\n", "Congrats on finishing this notebook! Feel free to check out our other [notebooks](https://github.com/oumi-ai/oumi/tree/main/notebooks) in the [Oumi GitHub](https://github.com/oumi-ai/oumi), and give us a star! You can also join the Oumi community over on [Discord](https://discord.gg/oumi).\n", "\n", "📰 Want to keep up with news from Oumi? Subscribe to our [Substack](https://blog.oumi.ai/) and [Youtube](https://www.youtube.com/@Oumi_AI)!\n", "\n", "⚡ Interested in building custom AI in hours, not months? Apply to get [early access](https://oumi.ai/contact?utm_source=oumi_oss_tutorial_remote_jobs) to the Oumi Platform, or [chat with us](https://oumi.ai/book?utm_source=oumi_oss_tutorial_quantization) to learn more!" ] } ], "metadata": { "kernelspec": { "display_name": "oumi", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.11" } }, "nbformat": 4, "nbformat_minor": 2 }