chore(sdk): add prototype python wrapper for nexus

core/cmd/wandb-core/main.go

@@ -12,6 +12,7 @@ import (
  "github.com/getsentry/sentry-go"
  "github.com/wandb/wandb/core/pkg/observability"
  "github.com/wandb/wandb/core/pkg/server"
+ "github.com/wandb/wandb/core/internal/processlib"
 )
 
 // this is set by the build script and used by the observability package
@@ -33,6 +34,11 @@ func main() {
 
  flag.Parse()
 
+ if *pid != 0 {
+ // Lets make sure this process is killed by the OS (if supported)
+ processlib.ShutdownOnParentDeath(*pid)
+ }
+
  // set up sentry reporting
  observability.InitSentry(*disableAnalytics, commit)
  defer sentry.Flush(2)
@@ -84,7 +90,7 @@ func main() {
  }
  defer trace.Stop()
  }
- serve, err := server.NewServer(ctx, "127.0.0.1:0", *portFilename)
+ serve, err := server.NewServer(ctx, "127.0.0.1:0", *portFilename, *pid)
  if err != nil {
  slog.Error("failed to start server, exiting", "error", err)
  return

core/internal/launcher/launcher.go

@@ -31,6 +31,7 @@ func readLines(path string) ([]string, error) {
 
 type Launcher struct {
  portFilename string
+ pidParent int
 }
 
 func (l *Launcher) tryport() (int, error) {
@@ -80,7 +81,7 @@ func (l *Launcher) prepTempfile() {
 
 func (l *Launcher) LaunchCommand(command string) (*execbin.ForkExecCmd, error) {
  l.prepTempfile()
- args := []string{"--port-filename", l.portFilename}
+ args := []string{"--port-filename", l.portFilename, "--pid", string(l.pidParent)}
  cmd, err := execbin.ForkExecCommand(command, args)
  if err != nil {
  panic(err)
@@ -91,14 +92,14 @@ func (l *Launcher) LaunchCommand(command string) (*execbin.ForkExecCmd, error) {
 func (l *Launcher) LaunchBinary(filePayload []byte) (*execbin.ForkExecCmd, error) {
  l.prepTempfile()
 
- args := []string{"--port-filename", l.portFilename}
+ args := []string{"--port-filename", l.portFilename, "--pid", strconv.Itoa(l.pidParent)}
  cmd, err := execbin.ForkExec(filePayload, args)
  if err != nil {
  panic(err)
  }
  return cmd, err
 }
 
-func NewLauncher() *Launcher {
- return &Launcher{}
+func NewLauncher(pidParent int) *Launcher {
+ return &Launcher{pidParent: pidParent}
 }

core/internal/processlib/processlib_linux.go

@@ -0,0 +1,22 @@
+package processlib
+
+import (
+ "os"
+ "syscall"
+)
+
+const (
+ PRCTL_SYSCALL = 157
+ PR_SET_PDEATHSIG = 1
+)
+
+func ShutdownOnParentDeath(parentPid int) {
+ _, _, errno := syscall.Syscall(uintptr(PRCTL_SYSCALL), uintptr(PR_SET_PDEATHSIG), uintptr(syscall.SIGKILL), 0)
+ if errno != 0 {
+ os.Exit(127 + int(errno))
+ }
+ // One last check... there is a possibility that the parent died right before the syscall was sent
+ if os.Getppid() != parentPid {
+ os.Exit(1)
+ }
+}

core/internal/processlib/processlib_notlinux.go

@@ -0,0 +1,6 @@
+//go:build !linux
+
+package processlib
+
+func ShutdownOnParentDeath(parentPid int) {
+}

core/lang/core/interface_pb.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+ "C"
+ "unsafe"
+ "google.golang.org/protobuf/proto"
+ "github.com/wandb/wandb/core/pkg/service"
+ "github.com/wandb/wandb/core/pkg/gowandb/opts/runopts"
+)
+
+//export pbSessionSetup
+func pbSessionSetup() {
+ wandbcoreSetup()
+}
+
+//export pbSessionTeardown
+func pbSessionTeardown() {
+ // prob dont want this, we could share nexus across "sessions"
+ wandbcoreTeardown()
+}
+
+//export pbRunStart
+func pbRunStart() int {
+ options := []runopts.RunOption{}
+ wandbcoreSetup()
+ run, err := wandbSession.NewRun(options...)
+ if err != nil {
+ panic(err)
+ }
+ num := wandbRuns.Add(run)
+ return num
+}
+
+//export pbRunLog
+func pbRunLog(num int, cBuffer *C.char, cLength C.int) {
+ data := C.GoBytes(unsafe.Pointer(cBuffer), cLength)
+ // Unmarshal protobuf
+ msg := &service.DataRecord{}
+ if err := proto.Unmarshal(data, msg); err != nil {
+ return
+ }
+ // Process data (here simply prepending a string)
+ run := wandbRuns.Get(num)
+
+ dict := make(map[string]interface{})
+ for k, v := range msg.Item {
+ switch value := v.DataType.(type) {
+ case *service.DataValue_ValueInt:
+ dict[k] = value.ValueInt
+ case *service.DataValue_ValueDouble:
+ dict[k] = value.ValueDouble
+ case *service.DataValue_ValueString:
+ dict[k] = value.ValueString
+ }
+ }
+
+ run.Log(dict)
+}
+
+//export pbRunFinish
+func pbRunFinish(num int) {
+ run := wandbRuns.Get(num)
+ run.Finish()
+ wandbRuns.Remove(num)
+}

core/lang/core/runkeeper.go

@@ -21,6 +21,8 @@ func NewRunKeeper() *RunKeeper {
 }
 
 func (k *RunKeeper) Get(num int) *gowandb.Run {
+ k.mutex.Lock()
+ defer k.mutex.Unlock()
  return k.runs[num]
 }
 

core/lang/core/wandbcore.go

@@ -10,6 +10,7 @@ typedef enum {
 import "C"
 
 import (
+ "os"
  "unsafe"
 
  "github.com/wandb/wandb/core/internal/gowandb/internal_runopts"
@@ -31,8 +32,10 @@ func wandbcoreSetup() {
  return
  }
  var err error
+ currentPid := os.Getpid()
  wandbSession, err = gowandb.NewSession(
  sessionopts.WithCoreBinary(coreBinary),
+ sessionopts.WithPidParent(currentPid),
  )
  if err != nil {
  panic(err)

core/lang/py/examples/local.sh

@@ -0,0 +1 @@
+time WANDB_BASE_URL=http://localhost:8080 $*

core/lang/py/examples/multiple.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python
+import wandb
+
+runs = [wandb.init() for x in range(10)]
+for run in runs:
+ run.log({"a": 1, "b": 2, "c": 4.0, "d": "blah"})
+for run in runs:
+ run.finish()

core/lang/py/examples/offline.sh

@@ -0,0 +1 @@
+time WANDB_MODE=offline $*

core/lang/py/examples/setup-nexus-py.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+set -e
+rm -rf nexus-py-1
+virtualenv nexus-py-1
+source nexus-py-1/bin/activate
+pip install --upgrade pip
+cd ../lib
+./build_proto.sh
+./build_lib.sh
+pip install -e .
+echo "Run:"
+echo "source nexus-py-1/bin/activate"

core/lang/py/examples/simple-newapi.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+import wandb
+
+"""
+Top level:
+new_api() -> API (? or sdk or library or core..or not)
+default_api
+default_session
+default_run
+(promote methods from default_api, default_session - and maybe default_run to top level namespace?)
+
+API:
+ new_session -> Session
+
+Session:
+ login()
+ configure_auth()
+ new_run()
+ get_run() # might have mutable and readonly versions of the run? readonly by default?
+ # alternate, prefix with object type? run_new, run_get... dont love
+ # ? are api runs just like runapi --> can we log to a run from the public api? why not?
+
+Run:
+ log()
+ history() -> how does this work for a run in progress
+
+
+"""
+
+api = wandb.new_api()
+session = api.new_session()
+
+run = session.new_run()
+run_id = run.id
+run.log({"a": 1})
+run.finish()
+
+run = session.get_run(run_id)
+for row in run.history():
+ pass

core/lang/py/examples/simple-newsdk.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+import wandb
+
+wb = wandb.new_session()
+run = wb.new_run()
+run.log({"a": 1, "b": 2})
+run.finish()

core/lang/py/examples/simple.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+import wandb
+
+run = wandb.init()
+run.log({"a": 1, "b": 2, "c": 4.0, "d": "blah"})
+run.finish()

core/lang/py/examples/threadpool.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+
+import argparse
+from concurrent.futures import ThreadPoolExecutor
+import time
+import wandb
+
+wandb.setup()
+
+def do_run():
+ run = wandb.init()
+ run.log({"a": 1, "b": 2, "c": 4.0, "d": "blah"})
+ run.finish()
+
+
+parser = argparse.ArgumentParser(description="benchmark wandb performance")
+parser.add_argument("--num-workers", "-n", type=int, default=20)
+args = parser.parse_args()
+
+jobs = []
+with ThreadPoolExecutor(max_workers=args.num_workers) as executor:
+ for _ in range(args.num_workers):
+ jobs.append(executor.submit(do_run))
+
+for job in jobs:
+ _ = job.result()
+print("done?")

core/lang/py/lib/MANIFEST.in

@@ -0,0 +1,5 @@
+recursive-include wandb *.py
+recursive-include wandb *.sh
+recursive-include wandb *.so
+recursive-exclude * __pycache__
+recursive-exclude * *.py[co]

core/lang/py/lib/build_lib.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+set -e
+BASE=../../
+DEST=py/lib/
+cd $BASE/
+./scripts/base-build.sh
+mkdir -p $DEST/wandb/lib
+cp export/lib/libwandb_core.so $DEST/wandb/lib

core/lang/py/lib/build_proto.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -e
+BASE=../../../..
+DEST=core/lang/py/lib/
+# cp $BASE/wandb/proto/*.proto wandb/proto/
+cd $BASE/
+protoc -I=. --python_out=$DEST wandb/proto/wandb_base.proto
+protoc -I=. --python_out=$DEST wandb/proto/wandb_internal.proto
+protoc -I=. --python_out=$DEST wandb/proto/wandb_telemetry.proto
+protoc -I=. --python_out=$DEST wandb/proto/wandb_settings.proto
+protoc -I=. --python_out=$DEST wandb/proto/wandb_server.proto

core/lang/py/lib/pyproject.toml

@@ -0,0 +1,53 @@
+[build-system]
+requires = ['setuptools>61']
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "wandb"
+dynamic = ["version"]
+description = "A CLI and library for interacting with the Weights & Biases API."
+authors = [{ name = "Weights & Biases", email = "[email protected]" }]
+# readme = "package_readme.md"
+# license = { file = "LICENSE" }
+requires-python = ">=3.7"
+dependencies = [
+ "protobuf>=3.12.0,!=4.21.0,<5; python_version < '3.9' and sys_platform == 'linux'",
+ "protobuf>=3.15.0,!=4.21.0,<5; python_version == '3.9' and sys_platform == 'linux'",
+ "protobuf>=3.19.0,!=4.21.0,<5; python_version > '3.9' and sys_platform == 'linux'",
+ "protobuf>=3.19.0,!=4.21.0,<5; sys_platform != 'linux'",
+ "setuptools",
+]
+classifiers = [
+ "Development Status :: 5 - Production/Stable",
+ "Intended Audience :: Developers",
+ "Intended Audience :: Science/Research",
+ "License :: OSI Approved :: MIT License",
+ "Natural Language :: English",
+ "Programming Language :: Python :: 3",
+ "Programming Language :: Python :: 3.7",
+ "Programming Language :: Python :: 3.8",
+ "Programming Language :: Python :: 3.9",
+ "Programming Language :: Python :: 3.10",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
+ "Programming Language :: Python :: 3 :: Only",
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Topic :: Software Development :: Libraries :: Python Modules",
+ "Topic :: System :: Logging",
+ "Topic :: System :: Monitoring",
+]
+
+
+#[project.scripts]
+#wandb = "wandb.cli.cli:cli"
+#wb = "wandb.cli.cli:cli"
+
+[project.urls]
+"Source" = "https://github.com/wandb/wandb"
+"Bug Reports" = "https://github.com/wandb/wandb/issues"
+"Documentation" = "https://docs.wandb.ai/"
+
+[tool.setuptools]
+zip-safe = false
+packages = ["wandb"]
+package-dir = { "wandb" = "wandb" }

core/lang/py/lib/setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup
+
+setup()

core/lang/py/lib/wandb/.gitignore

@@ -0,0 +1,2 @@
+proto/
+lib/