Implement python training code

Co-authored-by: Evaan Yaser Ahmed <ahmedeva@grinnell.edu>

Author: aaronmondal

.bazelignore

@@ -0,0 +1 @@
+.direnv

.gitignore

@@ -1,6 +1,5 @@
 .direnv/
 .pre-commit-config.yaml
-MODULE.bazel.lock
 bazel-*
 compile_commands.json
 lre.bazelrc

finetuning_on_cpu/.bazelrc

@@ -0,0 +1,32 @@
+# Pick up system specific configs.
+common --enable_platform_specific_config
+
+# Don't inherit PATH and LD_LIBRARY_PATH.
+build --incompatible_strict_action_env
+
+# Forbid network access unless explicitly enabled.
+build --sandbox_default_allow_network=false
+
+# Use correct runfile locations.
+build --nolegacy_external_runfiles
+
+# Enable sandboxing for exclusive tests like GPU performance tests.
+test --incompatible_exclusive_test_sandboxed
+
+# Simulate rules_python's initialization with a script.
+common --@rules_python//python/config_settings:bootstrap_impl=script
+
+# Remote optimizations.
+build --remote_download_minimal
+
+# Smaller profiling. Careful. Disabling this might explode remote cache usage.
+build --slim_profile
+build --experimental_profile_include_target_label
+build --noexperimental_profile_include_primary_output
+
+# Use the container image specified in `platforms/BUILD.bazel`.
+build:remote --extra_execution_platforms=@//platforms:python-remote
+build:remote --extra_toolchains=@zig_sdk//toolchain:linux_amd64_musl
+
+# Add user specific flags.
+try-import %workspace%/user.bazelrc

finetuning_on_cpu/.bazelversion

@@ -0,0 +1 @@
+8.2.1

finetuning_on_cpu/.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+bazel-*
+.env
+user.bazelrc

finetuning_on_cpu/BUILD.bazel

@@ -0,0 +1,19 @@
+load("@rules_python//python/uv:lock.bzl", "lock")
+
+exports_files(
+    ["python_provider.sh"],
+    visibility = ["//visibility:public"],
+)
+
+[
+    lock(
+        name = "requirements_{}".format(os),
+        srcs = ["pyproject.toml"],
+        out = "requirements_{}.lock".format(os),
+        build_constraints = ["@platforms//os:{}".format(os)],
+    )
+    for os in [
+        "linux",
+        "macos",
+    ]
+]

finetuning_on_cpu/MODULE.bazel

@@ -0,0 +1,42 @@
+bazel_dep(name = "hermetic_cc_toolchain", version = "4.0.0")
+bazel_dep(name = "platforms", version = "0.0.11")
+bazel_dep(name = "rules_cc", version = "0.1.1")
+bazel_dep(name = "rules_python", version = "1.4.1")
+bazel_dep(name = "rules_shell", version = "0.4.1")
+bazel_dep(name = "bazel_skylib", version = "1.7.1")
+
+# Minimal C++ toolchain required internally by rules_python.
+cc_toolchains = use_extension("@hermetic_cc_toolchain//toolchain:ext.bzl", "toolchains")
+use_repo(cc_toolchains, "zig_sdk")
+
+# Python configuration
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(python_version = "3.13")
+
+# Python packages
+pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
+pip.parse(
+    # CPU-only torch is not available in the regular pypi repositories.
+    experimental_index_url = "https://pypi.org/simple",
+    experimental_index_url_overrides = {
+        "torch": "https://download.pytorch.org/whl/cpu",
+    },
+    hub_name = "pypi",
+    python_version = "3.13",
+    # This file may be regenerated via `bazel run //:requirements_macos.update`.
+    requirements_darwin = "//:requirements_macos.lock",
+    # This file may be regenerated via `bazel run //:requirements_linux.update`.
+    requirements_linux = "//:requirements_linux.lock",
+)
+use_repo(pip, "pypi")
+
+# UV toolchain config for regenerating the lockfile
+uv = use_extension(
+    "@rules_python//python/uv:uv.bzl",
+    "uv",
+    dev_dependency = True,
+)
+uv.configure(
+    # Latest version that supports the old URL schema used in rules_python.
+    version = "0.6.5",
+)

finetuning_on_cpu/MODULE.bazel.lock

@@ -0,0 +1,67 @@
+# Train a machine learning model with Bazel and NativeLink
+
+An example project for training a machine learning model with Bazel.
+
+## ❄️ Setup
+
+1. Either have Bazelisk installed or enter the nix flake from this repository.
+2. Create a `user.bazelrc` in this directory and paste the values from the
+   `Quickstart` section at <https://app.nativelink.com> into it:
+
+    ```bash
+    # user.bazelrc
+    build --remote_cache=grpcs://TODO
+    build --remote_header=x-nativelink-api-key=TODO
+    build --bes_backend=grpcs://TODO
+    build --bes_header=x-nativelink-api-key=TODO
+    build --bes_results_url=https://app.nativelink.com/TODO
+    build --remote_timeout=600
+    build --remote_executor=grpcs://TODO
+    ```
+
+## 🚄 Training
+
+Run a local invocation:
+
+```bash
+bazel run training
+```
+
+Now run a remote invocation. The first time you do this it'll be slow (~15min)
+as the runner needs to spin up and fetch the various toolchains:
+
+```bash
+bazel test training --config=remote --test_output=all
+```
+
+Run the above command a second time. Since the runner is now warm it'll
+immediately start the test (~5min).
+
+## ✨ Keeping dependencies up-to-date
+
+The dependencies for this project are tracked in `pyproject.toml` and locked in
+`requirements_{linux±|macos}.lock` via a Bazel-wrapped `uv` toolchain. To
+regenerate the lockfiles:
+
+```bash
+# On linux
+bazel run //:requirements_linux.update
+
+# On macos
+bazel run //:requirements_macos.update
+```
+
+On darwin you might have to disable remote execution temporarily for this.
+
+## 🧹 Fixing formatting and linting
+
+Use the flake if you develop on this code. Don't use local ruff installations as
+versions other than the one used in the flake might format code differently.
+
+```bash
+# Check lints
+ruff check
+
+# Format
+ruff format
+```

finetuning_on_cpu/README.md

@@ -0,0 +1,12 @@
+platform(
+    name = "python-remote",
+    constraint_values = [
+        "@platforms//os:linux",
+        "@platforms//cpu:x86_64",
+    ],
+    exec_properties = {
+        "OSFamily": "Linux",
+        "ISA": "x86-64",
+        "container-image": "ubuntu:24.04",
+    },
+)