multigeom.pipeline.distribute#

Supporting functions of 2_subspace_alignment.ipynb of multigeom-tutorials: paper Fig. 4D-E, the subspace alignment of the top-k eigenvectors of G(x) and J(x) across encoders and areas, averaged over the test images.

The work is the estimation of G(x) and J(x) for every test image of every (encoder,
area) configuration, then Stage 5 (alignment between configurations), then the plot. The
G(x) estimation dominates: the images are independent (no gradients, no communication
between GPUs), the stage is resumable and every image is seeded on its own (seed + image
index). So the images can be split across workers (GPUs, local or on Modal) without
changing the results, and no DDP / FSDP is involved.

Everything is parametrised by

    configs      (encoder, area), or a list of them: ("inception", "V1") or all four
    n_images     first n test images (the paper uses 200)
    backend      "local" (this machine) or "modal" (serverless GPUs, https://modal.com)
    device       "gpu" or "cpu" (Modal is always "gpu")
    num_workers  how many workers run in parallel (one GPU or one CPU process each)

bundled in a `Scenario`. The G(x) estimation itself is not re-implemented here: the
workers run `multigeom.pipeline.per_image`, so this module cannot diverge from it.

`modal` is imported only inside the Modal functions, so the rest works without it
installed.

Functions

`alive`(pids)	The pids that are still running (also works for processes started by an earlier kernel).
`check_budget`(s, run)	Print the estimate; raise unless run is True and, on Modal, the cost is known and within the budget.
`clear_results`(configs, n_images[, runs_dir, ...])	Delete the per-image files of the first n_images images of the given configurations, so that the G(x) estimation recomputes them from scratch (it skips images that exist).
`estimate_cost`(n_computations, n_workers, ...)	Total GPU-seconds = images x seconds per image + one startup per worker; wall = that / workers.
`is_complete`(configs, n_images[, runs_dir])	True when every image of every configuration has its per-image file.
`launch_local`(workers, *, use_gpu[, runs_dir])	Start one detached process per worker (device cuda:i, or the CPU), so the work goes on if the notebook kernel dies.
`launch_modal`(workers, *[, gpu_type, ...])	Run each worker's shards in its own Modal container with one GPU, in parallel, and wait.
`main`([argv])	Command-line entry point of the workers started by the launchers.
`make_shards`(configs, n_images, num_workers)	Split the (config, image) work into max(num_workers, 1) balanced workers, each a list of shards (one shard = consecutive images of one config).
`modal_clear`(configs, n_images[, volume_name])	Delete the per-image files of the configs on the Modal volume (the containers skip images that exist there).
`modal_download`(configs, n_images[, ...])	Copy the per-image files of the configs from the Modal volume to <runs_dir>/per_image/ (new ones only).
`modal_images_done`(vol, expected)	How many of the expected images have their per-image file on the Modal volume (committed files only).
`modal_upload`(configs[, volume_name])	Copy the checkpoints and test images to a Modal volume (files already there are skipped).
`normalize_configs`(configs)	("inception", "V1") or [("inception", "V1"), ...] -> a tuple of (encoder, area) tuples.
`per_image_command`(shard, *, device[, ...])	Command line of python -m multigeom.pipeline.per_image for one shard (it skips images already done).
`print_estimate`(s)	Print what a scenario would take (images, workers, GPU-time, wall-clock and, on Modal, the cost).
`progress`(configs, n_images[, runs_dir])	Number of per-image files done, per configuration (out of n_images).
`resolve_workers`(s)	The device and number of workers that will really be used: (device, num_workers).
`run_shards`(shards, *, device[, runs_dir, ...])	Run the shards one after the other with multigeom.pipeline.per_image, the G(x) estimation (used by the local workers and the Modal containers).
`scenario_estimate`(s)	Estimate the time and the cost of running the scenario s.
`stop_workers`([procs, runs_dir, grace_s, verbose])	Stop the local workers (POSIX: Linux, macOS).
`wait_for_workers`([procs, configs, n_images, ...])	Block until the local workers have finished (the notebook cell stays busy meanwhile), printing the progress.
`watch_calls`(calls, count_done, bar, *[, ...])	Wait for the Modal function calls to finish while count_done() (images on the volume) drives the progress bar; this runs in the notebook's own thread, so the bar updates in place.

Classes

`CostEstimate`(n_computations, n_workers, ...)	Estimated GPU time and cost of a scenario.
`Scenario`(configs, n_images[, backend, ...])	What to compute and where.
`Shard`(encoder, area, start, end)	Images [start, end) of one configuration.