File tree Expand file tree Collapse file tree 1 file changed +51
-0
lines changed
pathwaysutils/experimental/shared_pathways_service Expand file tree Collapse file tree 1 file changed +51
-0
lines changed Original file line number Diff line number Diff line change 1+ # Shared Pathways Service
2+
3+ Shared pathways service is a multi-tenant Pathways cluster with dedicated TPU
4+ resources. This eliminates the need for complex cloud setup, allowing you to
5+ get started from a familiar local environment (like a laptop or cloud VM) with
6+ minimal overhead: Just wrap your Python entrypoint in a
7+ ` with isc_pathways.connect(): ` block!.
8+
9+ ## Requirements
10+
11+ Make sure that your cluster is running the Resource Manager and Worker pods.
12+ If not, you can use [ pw-service-example.yaml] ( yamls/pw-service-example.yaml ) .
13+ Make sure to modify the following values to deploy these pods:
14+
15+ - A unique Jobset name for the cluster's Pathways pods
16+ - GCS bucket path
17+ - TPU type and topology
18+ - Number of slices
19+
20+ These fields are highlighted in the YAML file with trailing comments for easier
21+ understanding.
22+
23+ ## Instructions
24+
25+ 1 . Clone ` pathwaysutils ` .
26+
27+ ` git clone https://github.com/AI-Hypercomputer/pathways-utils.git `
28+
29+ 2 . Import ` isc_pathways.py ` and move your workload under
30+ ` with isc_pathways.connect() ` statement. Refer to
31+ [ run_connect_example.py] ( run_connect_example.py ) for reference. Example code:
32+
33+ ```
34+ from pathwaysutils.experimental.shared_pathways_service import isc_pathways
35+
36+ with isc_pathways.connect(
37+ "my-cluster",
38+ "my-project",
39+ "region",
40+ "gs://user-bucket",
41+ "pathways-cluster-pathways-head-0-0.pathways-cluster:29001",
42+ {"tpuv6e:2x2": 2},
43+ ) as tm:
44+ import jax.numpy as jnp
45+ import pathwaysutils
46+ import pprint
47+
48+ pathwaysutils.initialize()
49+ orig_matrix = jnp.zeros(5)
50+ ...
51+ ```
You can’t perform that action at this time.
0 commit comments