From 0da50d50af478d36db1ccaa4db3691dca7edc222 Mon Sep 17 00:00:00 2001 From: Nick Petrovic <4001122+nickpetrovic@users.noreply.github.com> Date: Thu, 11 Jan 2024 11:53:54 -0500 Subject: [PATCH] New app config (#41) --- .dockerignore | 1 - .flake8 | 5 - .secrets.template | 54 -- Makefile | 3 - bin/__init__.py | 0 bin/populate_secrets.py | 56 -- docker/Dockerfile.worker | 1 - go.mod | 38 +- go.sum | 93 +- hack/okteto.yaml | 2 - internal/abstractions/image/build.go | 13 +- internal/abstractions/image/image.go | 68 +- internal/common/config.default.yaml | 85 ++ internal/common/config.go | 127 +++ internal/common/redis.go | 181 +--- internal/common/redis_test.go | 38 +- internal/common/registry.go | 27 +- internal/common/secrets.go | 192 ---- internal/common/secrets_test.go | 35 - internal/common/stats.go | 9 +- internal/gateway/gateway.go | 30 +- internal/repository/backend_postgres.go | 29 +- internal/repository/base.go | 6 +- internal/repository/beam_postgres.go | 24 +- internal/repository/metrics_stream_kinesis.go | 23 +- internal/repository/testutils.go | 3 +- internal/repository/worker_pool_redis.go | 43 +- internal/scheduler/backlog_test.go | 2 +- internal/scheduler/pool.go | 41 - internal/scheduler/pool_client.go | 68 -- internal/scheduler/pool_k8s.go | 310 ++---- internal/scheduler/pool_manager.go | 99 +- internal/scheduler/pool_remote.go | 109 --- internal/scheduler/pool_sizing.go | 47 +- internal/scheduler/pool_sizing_test.go | 76 +- internal/scheduler/scheduler.go | 81 +- internal/scheduler/scheduler.proto | 49 - internal/scheduler/scheduler_test.go | 84 +- internal/scheduler/service.go | 102 +- internal/storage/juicefs.go | 19 +- internal/storage/storage.go | 22 +- internal/types/config.go | 155 +++ internal/types/pool.go | 107 --- internal/types/scheduler.go | 1 - internal/worker/creds.go | 16 +- internal/worker/cuda_test.go | 6 +- internal/worker/image.go | 58 +- internal/worker/network.go | 24 +- internal/worker/util.go | 2 +- internal/worker/worker.go | 20 +- manifests/crds/workerpool.yaml | 33 - manifests/k3d/beam.yaml | 17 - proto/scheduler.pb.go | 883 +----------------- proto/scheduler_grpc.pb.go | 112 +-- 54 files changed, 866 insertions(+), 2863 deletions(-) delete mode 100644 .flake8 delete mode 100644 .secrets.template delete mode 100644 bin/__init__.py delete mode 100644 bin/populate_secrets.py create mode 100644 internal/common/config.default.yaml create mode 100644 internal/common/config.go delete mode 100644 internal/common/secrets.go delete mode 100644 internal/common/secrets_test.go delete mode 100644 internal/scheduler/pool_client.go delete mode 100644 internal/scheduler/pool_remote.go create mode 100644 internal/types/config.go delete mode 100644 internal/types/pool.go delete mode 100644 manifests/crds/workerpool.yaml diff --git a/.dockerignore b/.dockerignore index dabed888a..bece379a2 100644 --- a/.dockerignore +++ b/.dockerignore @@ -5,6 +5,5 @@ !go.mod !go.sum !sdk -!.secrets **/__pycache__ *.pyc diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 40f2f39ee..000000000 --- a/.flake8 +++ /dev/null @@ -1,5 +0,0 @@ -[flake8] -ignore = E226,E302,E41,W503,E731 -max-line-length = 100 -exclude = tests/* -max-complexity = 10 \ No newline at end of file diff --git a/.secrets.template b/.secrets.template deleted file mode 100644 index e1f484dda..000000000 --- a/.secrets.template +++ /dev/null @@ -1,54 +0,0 @@ -# code: language=properties -AWS_REGION="us-east-1" -BEAM_AGENT_TOKEN="7ace49a5-d8a6-40fd-a80a-d6752eef5027" -BEAM_AGENT_URL="scheduler.:2003" -BEAM_AGENT_WORKER_NAMESPACE="" -BEAM_CACHE_URL="cache.:2049" -BEAM_DATA_VOLUME_NAME="worker-data" -BEAM_GATEWAY_HOST="beam." -BEAM_GATEWAY_PORT=1993 -BEAM_IMAGESERVICE_BASE_IMAGE_CACHE_SIZE="1" -BEAM_IMAGESERVICE_IMAGE_CREDENTIAL_PROVIDER="docker" -BEAM_IMAGESERVICE_IMAGE_REGISTRY_STORE="local" -BEAM_IMAGESERVICE_IMAGE_TAGS_TO_CACHE="py38,py39" -BEAM_IMAGESERVICE_IMAGE_TLS_VERIFY="false" -BEAM_IMAGESERVICE_SERVICE_TOKEN="token-1234" -BEAM_JUICEFS_AWS_ACCESS_KEY="test" -BEAM_JUICEFS_AWS_SECRET_KEY="test" -BEAM_JUICEFS_REDIS="redis://juicefs-redis-master.:6379/0" -BEAM_JUICEFS_S3_BUCKET="http://localstack.beam:4566/juicefs" -BEAM_NAMESPACE="" -BEAM_RUNNER_BASE_IMAGE_NAME="beam-runner" -BEAM_RUNNER_BASE_IMAGE_REGISTRY="registry.beam.lan:5000" -BEAM_RUNNER_BASE_IMAGE_TAG="latest" -BEAM_WORKER_BACKEND_BASE_URL="http://backend.:8080" -BEAM_WORKER_BEAMAPI_BASE_URL="http://beam-api.:8888" -BEAM_WORKER_IMAGE_NAME="beam-worker" -BEAM_WORKER_IMAGE_REGISTRY="registry.beam.lan:5000" -BEAM_WORKER_IMAGE_TAG="latest" -BEAM_WORKER_IMGFSCACHER_HOST="imgfs-cacher-daemon." -BEAM_WORKER_RESOURCES_ENFORCED="true" -BEAT_REDIS_DIAL_TIMEOUT="3s" -BEAT_REDIS_HOSTS="redis-master.:6379" -BEAT_REDIS_MAX_RETRIES="10" -BEAT_REDIS_MODE="standalone" -DB_HOST="postgres-postgresql." -DB_NAME="main" -DB_PASS="password" -DB_PORT="5432" -DB_USER="root" -DOCKERHUB_PASSWORD="" -DOCKERHUB_USERNAME="" -DOCKER_PASSWORD="" -DOCKER_USERNAME="" -LOCALSTACK_AWS_ACCESS_KEY="test" -LOCALSTACK_AWS_SECRET_KEY="test" -LOCALSTACK_AWS_SESSION="test" -LOCALSTACK_ENDPOINT="http://localstack.:4566" -NAMESPACE="" -REDIS_DIAL_TIMEOUT="3s" -REDIS_HOSTS="redis-master.:6379" -REDIS_MAX_RETRIES="10" -REDIS_MODE="standalone" -WORKER_NAMESPACE="" -WORKER_S2S_TOKEN="RYHSBsqd8hQYFFE1rjeq586OtuN0GUekLmIef81HHddGd0NL7XAhjmfYOZQakn9E" diff --git a/Makefile b/Makefile index e4c24ab2e..bc718f79f 100644 --- a/Makefile +++ b/Makefile @@ -37,8 +37,5 @@ start: stop: cd hack && okteto down --file okteto.yml -secrets: - unset GITHUB_TOKEN && OKTETO_NAMESPACE=beam python3 bin/populate_secrets.py - protocol: cd proto && ./gen.sh diff --git a/bin/__init__.py b/bin/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/bin/populate_secrets.py b/bin/populate_secrets.py deleted file mode 100644 index 05c8acc53..000000000 --- a/bin/populate_secrets.py +++ /dev/null @@ -1,56 +0,0 @@ -#!/usr/bin/python3 - -import subprocess -import os -import json -import sys - -TEMPLATE_FILE_LOCATION = ".secrets.template" -OUTPUT_FILE_LOCATION = ".secrets" - - -def _generate_secrets_file(): - okteto_docker_username = None - okteto_docker_password = None - - # retrieve okteto namespace from environment - okteto_namespace = os.environ.get("OKTETO_NAMESPACE") - if okteto_namespace is None: - print( - "ERROR: unable to load okteto namespace, do you have the OKTETO_NAMESPACE value set in your shell?" - ) - sys.exit(1) - - # load template secrets - template_contents = None - with open(TEMPLATE_FILE_LOCATION, "r") as f: - template_contents = f.read() - - # load okteto context for docker registry credentials - try: - okteto_context = subprocess.check_output( - args=["okteto", "context", "show"], env=os.environ - ) - okteto_context = json.loads(okteto_context.decode("utf-8")) - okteto_docker_username = okteto_context["username"] - okteto_docker_password = okteto_context["token"] - except BaseException: - okteto_docker_username = okteto_namespace - okteto_docker_password = os.environ.get("OKTETO_DOCKER_PASSWORD", "") - - # populate template values - populated_template = ( - template_contents.replace("", okteto_namespace, -1) - .replace("", okteto_docker_username, -1) - .replace("", okteto_docker_password, -1) - ) - - print("\n***** Generated secrets file *****\n") - print(populated_template) - with open(OUTPUT_FILE_LOCATION, "w") as f: - f.write(populated_template) - print("\n**********************************\n") - - -if __name__ == "__main__": - _generate_secrets_file() diff --git a/docker/Dockerfile.worker b/docker/Dockerfile.worker index 9cbfc07d8..9b4592bbb 100644 --- a/docker/Dockerfile.worker +++ b/docker/Dockerfile.worker @@ -66,7 +66,6 @@ RUN GOPRIVATE=github.com/beam-cloud/* \ # ======================== FROM nvidia/cuda:12.3.1-base-ubuntu20.04 AS release FROM release AS dev -ONBUILD COPY .secrets /workspace/.secrets FROM ${BASE_STAGE} AS final diff --git a/go.mod b/go.mod index 30e937616..ca46cc56e 100644 --- a/go.mod +++ b/go.mod @@ -13,16 +13,18 @@ require ( github.com/aws/aws-sdk-go-v2/service/ecr v1.24.4 github.com/aws/aws-sdk-go-v2/service/kinesis v1.20.0 github.com/aws/aws-sdk-go-v2/service/s3 v1.38.3 - github.com/aws/karpenter v0.18.1 github.com/beam-cloud/clip v0.0.0-20230822004255-4ecd939864f9 github.com/beam-cloud/go-runc v0.0.0-20231222221338-b89899f33170 github.com/bsm/redislock v0.9.4 - github.com/gin-gonic/gin v1.9.1 github.com/gofrs/uuid v4.0.0+incompatible github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 github.com/google/uuid v1.5.0 github.com/jmoiron/sqlx v1.3.5 - github.com/joho/godotenv v1.5.1 + github.com/knadh/koanf/parsers/json v0.1.0 + github.com/knadh/koanf/parsers/yaml v0.1.0 + github.com/knadh/koanf/providers/file v0.1.0 + github.com/knadh/koanf/providers/rawbytes v0.1.0 + github.com/knadh/koanf/v2 v2.0.1 github.com/labstack/echo/v4 v4.11.4 github.com/lib/pq v1.10.9 github.com/mholt/archiver/v3 v3.5.1 @@ -34,8 +36,7 @@ require ( github.com/pkg/errors v0.9.1 github.com/pressly/goose/v3 v3.17.0 github.com/prometheus/procfs v0.12.0 - github.com/redis/go-redis/v9 v9.1.0 - github.com/samber/lo v1.39.0 + github.com/redis/go-redis/v9 v9.4.0 github.com/sirupsen/logrus v1.9.3 github.com/stretchr/testify v1.8.4 github.com/tj/assert v0.0.3 @@ -70,29 +71,21 @@ require ( github.com/aws/aws-sdk-go-v2/service/sts v1.21.3 // indirect github.com/aws/smithy-go v1.19.0 // indirect github.com/briandowns/spinner v1.23.0 // indirect - github.com/bytedance/sonic v1.10.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect - github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect - github.com/chenzhuoyu/iasm v0.9.0 // indirect github.com/containerd/console v1.0.3 // indirect github.com/cpuguy83/go-md2man/v2 v2.0.3 // indirect github.com/cyphar/filepath-securejoin v0.2.4 // indirect - github.com/davecgh/go-spew v1.1.1 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/docker/go-units v0.5.0 // indirect github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5 // indirect github.com/emicklei/go-restful/v3 v3.11.0 // indirect github.com/fatih/color v1.16.0 // indirect - github.com/gabriel-vasile/mimetype v1.4.2 // indirect - github.com/gin-contrib/sse v0.1.0 // indirect + github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/go-logr/logr v1.4.1 // indirect github.com/go-openapi/jsonpointer v0.20.2 // indirect github.com/go-openapi/jsonreference v0.20.4 // indirect github.com/go-openapi/swag v0.22.6 // indirect - github.com/go-playground/locales v0.14.1 // indirect - github.com/go-playground/universal-translator v0.18.1 // indirect - github.com/go-playground/validator/v10 v10.15.4 // indirect - github.com/goccy/go-json v0.10.2 // indirect github.com/gogo/protobuf v1.3.2 // indirect github.com/golang-jwt/jwt v3.2.2+incompatible // indirect github.com/golang/protobuf v1.5.3 // indirect @@ -116,15 +109,17 @@ require ( github.com/json-iterator/go v1.1.12 // indirect github.com/karrick/godirwalk v1.17.0 // indirect github.com/klauspost/compress v1.17.4 // indirect - github.com/klauspost/cpuid/v2 v2.2.5 // indirect github.com/klauspost/pgzip v1.2.6 // indirect + github.com/knadh/koanf/maps v0.1.1 // indirect github.com/kylelemons/godebug v1.1.0 // indirect github.com/labstack/gommon v0.4.2 // indirect - github.com/leodido/go-urn v1.2.4 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-sqlite3 v1.14.19 // indirect + github.com/mitchellh/copystructure v1.2.0 // indirect + github.com/mitchellh/mapstructure v1.5.0 // indirect + github.com/mitchellh/reflectwalk v1.0.2 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect @@ -132,15 +127,12 @@ require ( github.com/opencontainers/go-digest v1.0.0 // indirect github.com/opencontainers/image-spec v1.1.0-rc5 // indirect github.com/opencontainers/runc v1.1.10 // indirect - github.com/pelletier/go-toml/v2 v2.1.0 // indirect github.com/pierrec/lz4/v4 v4.1.19 // indirect - github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/rootless-containers/proto v0.1.0 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/sethvargo/go-retry v0.2.4 // indirect github.com/tidwall/btree v1.6.0 // indirect - github.com/twitchyliquid64/golang-asm v0.15.1 // indirect - github.com/ugorji/go/codec v1.2.11 // indirect github.com/ulikunitz/xz v0.5.11 // indirect github.com/urfave/cli v1.22.12 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect @@ -149,13 +141,11 @@ require ( github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect github.com/yuin/gopher-lua v1.1.0 // indirect go.uber.org/multierr v1.11.0 // indirect - golang.org/x/arch v0.5.0 // indirect golang.org/x/crypto v0.17.0 // indirect - golang.org/x/exp v0.0.0-20231219180239-dc181d75b848 // indirect golang.org/x/net v0.19.0 // indirect golang.org/x/oauth2 v0.15.0 // indirect golang.org/x/sync v0.5.0 // indirect - golang.org/x/sys v0.15.0 // indirect + golang.org/x/sys v0.16.0 // indirect golang.org/x/term v0.15.0 // indirect golang.org/x/text v0.14.0 // indirect golang.org/x/time v0.5.0 // indirect diff --git a/go.sum b/go.sum index 948eb8d11..e0e6499a5 100644 --- a/go.sum +++ b/go.sum @@ -80,8 +80,6 @@ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.15.3 h1:90qW9puxI7LgmiYKSPhx6wz4 github.com/aws/aws-sdk-go-v2/service/ssooidc v1.15.3/go.mod h1:kKpyLjToIS7E3z0672lBhxIPD+uoQ9V0MYRYCVGIkO0= github.com/aws/aws-sdk-go-v2/service/sts v1.21.3 h1:s3wBkMxfA/u2EJJl6KRsPcWv858lDHkhinqXyN6fkZI= github.com/aws/aws-sdk-go-v2/service/sts v1.21.3/go.mod h1:b+y9zL57mwCRy6ftp9Nc7CONGHX3sZ50ZCLTrI5xpCc= -github.com/aws/karpenter v0.18.1 h1:fy2dHEjRfWWnLR+GxHC6fIfcyfIBhiCc/iH9/+/d5R4= -github.com/aws/karpenter v0.18.1/go.mod h1:N9ka+O8MD9qgsWp6OVgmSxUHh3nbiJGDEUlJW3WmKyk= github.com/aws/smithy-go v1.14.1/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA= github.com/aws/smithy-go v1.15.0/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA= github.com/aws/smithy-go v1.19.0 h1:KWFKQV80DpP3vJrrA9sVAHQ5gc2z8i4EzrLhLlWXcBM= @@ -93,27 +91,17 @@ github.com/beam-cloud/go-runc v0.0.0-20231222221338-b89899f33170 h1:KYVz18kobBGU github.com/beam-cloud/go-runc v0.0.0-20231222221338-b89899f33170/go.mod h1:aw0zhDi28Hemve0raHcfU9suxZwkCpyNANOEwKZSSXo= github.com/briandowns/spinner v1.23.0 h1:alDF2guRWqa/FOZZYWjlMIx2L6H0wyewPxo/CH4Pt2A= github.com/briandowns/spinner v1.23.0/go.mod h1:rPG4gmXeN3wQV/TsAY4w8lPdIM6RX3yqeBQJSrbXjuE= -github.com/bsm/ginkgo/v2 v2.9.5 h1:rtVBYPs3+TC5iLUVOis1B9tjLTup7Cj5IfzosKtvTJ0= -github.com/bsm/ginkgo/v2 v2.9.5/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= -github.com/bsm/gomega v1.26.0 h1:LhQm+AFcgV2M0WyKroMASzAzCAJVpAxQXv4SaI9a69Y= -github.com/bsm/gomega v1.26.0/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/bsm/redislock v0.9.4 h1:X/Wse1DPpiQgHbVYRE9zv6m070UcKoOGekgvpNhiSvw= github.com/bsm/redislock v0.9.4/go.mod h1:Epf7AJLiSFwLCiZcfi6pWFO/8eAYrYpQXFxEDPoDeAk= -github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM= -github.com/bytedance/sonic v1.10.0-rc/go.mod h1:ElCzW+ufi8qKqNW0FY314xriJhyJhuoJ3gFZdAHF7NM= -github.com/bytedance/sonic v1.10.1 h1:7a1wuFXL1cMy7a3f7/VFcEtriuXQnUBhtoVfOZiaysc= -github.com/bytedance/sonic v1.10.1/go.mod h1:iZcSUejdk5aukTND/Eu/ivjQuEL0Cu9/rf50Hi0u/g4= github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM= github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY= -github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk= -github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d h1:77cEq6EriyTZ0g/qfRdp61a3Uu/AWrgIq2s0ClJV1g0= -github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d/go.mod h1:8EPpVsBuRksnlj1mLy4AWzRNQYxauNi62uWcE3to6eA= -github.com/chenzhuoyu/iasm v0.9.0 h1:9fhXjVzq5hUy2gkhhgHl95zG2cEAhw9OSGs8toWWAwo= -github.com/chenzhuoyu/iasm v0.9.0/go.mod h1:Xjy2NpN3h7aUqeqM+woSuuvxmIe6+DDsiNLIrkAmYog= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= @@ -137,8 +125,9 @@ github.com/cyphar/filepath-securejoin v0.2.2/go.mod h1:FpkQEhXnPnOthhzymB7CGsFk2 github.com/cyphar/filepath-securejoin v0.2.4 h1:Ugdm7cg7i6ZK6x3xDF1oEu1nfkyfH53EtKeQYTC3kyg= github.com/cyphar/filepath-securejoin v0.2.4/go.mod h1:aPGpWjXOXUn2NCNjFvBE6aRxGGx79pTxQpKOJNYHHl4= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/docker/cli v24.0.7+incompatible h1:wa/nIwYFW7BVTGa7SWPVyyXU9lgORqUb1xfI36MSkFg= @@ -168,12 +157,8 @@ github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM= github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= -github.com/gabriel-vasile/mimetype v1.4.2 h1:w5qFW6JKBz9Y393Y4q372O9A7cUSequkh1Q7OhCmWKU= -github.com/gabriel-vasile/mimetype v1.4.2/go.mod h1:zApsH/mKG4w07erKIaJPFiX0Tsq9BFQgN3qGY5GnNgA= -github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= -github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= -github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg= -github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU= +github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA= +github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM= github.com/go-faster/city v1.0.1 h1:4WAxSZ3V2Ws4QRDrscLEDcibJY8uf41H6AhXDrNDcGw= github.com/go-faster/city v1.0.1/go.mod h1:jKcUJId49qdW3L1qKHH/3wPeUstCVpVSXTM6vO3VcTw= github.com/go-faster/errors v0.6.1 h1:nNIPOBkprlKzkThvS/0YaX8Zs9KewLCOSFQS5BU06FI= @@ -190,22 +175,12 @@ github.com/go-openapi/jsonreference v0.20.4 h1:bKlDxQxQJgwpUSgOENiMPzCTBVuc7vTdX github.com/go-openapi/jsonreference v0.20.4/go.mod h1:5pZJyJP2MnYCpoeoMAql78cCHauHj0V9Lhc506VOpw4= github.com/go-openapi/swag v0.22.6 h1:dnqg1XfHXL9aBxSbktBqFR5CxVyVI+7fYWhAf1JOeTw= github.com/go-openapi/swag v0.22.6/go.mod h1:Gl91UqO+btAM0plGGxHqJcQZ1ZTy6jbmridBTsDy8A0= -github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= -github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= -github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= -github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= -github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= -github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= -github.com/go-playground/validator/v10 v10.15.4 h1:zMXza4EpOdooxPel5xDqXEdXG5r+WggpvnAKMsalBjs= -github.com/go-playground/validator/v10 v10.15.4/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU= github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI= github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= -github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= -github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/gofrs/uuid v4.0.0+incompatible h1:1SD/1F5pU8p29ybwgQSwpQk+mwdRrXCYuPhW6m+TnJw= github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= @@ -327,8 +302,6 @@ github.com/jmoiron/sqlx v1.3.5 h1:vFFPA71p1o5gAeqtEAwLU4dnX2napprKtHr7PYIcN3g= github.com/jmoiron/sqlx v1.3.5/go.mod h1:nRVWtLre0KfCLJvgxzCsLVMogSvQ1zNJtpYr2Ccp0mQ= github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901 h1:rp+c0RAYOWj8l6qbCUTSiRLG/iKnW3K3/QfPPuSsBt4= github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901/go.mod h1:Z86h9688Y0wesXCyonoVr47MasHilkuLMqGhRZ4Hpak= -github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= -github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/jonboulle/clockwork v0.4.0 h1:p4Cf1aMWXnXAUh8lVfewRBx1zaTSYKrKMF2g3ST4RZ4= github.com/jonboulle/clockwork v0.4.0/go.mod h1:xgRqUGwRcjKCO1vbZUEtSLrqKoPSsUpK7fnezOII0kc= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= @@ -348,14 +321,22 @@ github.com/klauspost/compress v1.11.4/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYs github.com/klauspost/compress v1.17.4 h1:Ej5ixsIri7BrIjBkRZLTo6ghwrEtHFk7ijlczPW4fZ4= github.com/klauspost/compress v1.17.4/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM= github.com/klauspost/cpuid v1.2.0/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek= -github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= -github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg= -github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= github.com/klauspost/pgzip v1.2.4/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= github.com/klauspost/pgzip v1.2.5/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= github.com/klauspost/pgzip v1.2.6 h1:8RXeL5crjEUFnR2/Sn6GJNWtSQ3Dk8pq4CL3jvdDyjU= github.com/klauspost/pgzip v1.2.6/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= -github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/knadh/koanf/maps v0.1.1 h1:G5TjmUh2D7G2YWf5SQQqSiHRJEjaicvU0KpypqB3NIs= +github.com/knadh/koanf/maps v0.1.1/go.mod h1:npD/QZY3V6ghQDdcQzl1W4ICNVTkohC8E73eI2xW4yI= +github.com/knadh/koanf/parsers/json v0.1.0 h1:dzSZl5pf5bBcW0Acnu20Djleto19T0CfHcvZ14NJ6fU= +github.com/knadh/koanf/parsers/json v0.1.0/go.mod h1:ll2/MlXcZ2BfXD6YJcjVFzhG9P0TdJ207aIBKQhV2hY= +github.com/knadh/koanf/parsers/yaml v0.1.0 h1:ZZ8/iGfRLvKSaMEECEBPM1HQslrZADk8fP1XFUxVI5w= +github.com/knadh/koanf/parsers/yaml v0.1.0/go.mod h1:cvbUDC7AL23pImuQP0oRw/hPuccrNBS2bps8asS0CwY= +github.com/knadh/koanf/providers/file v0.1.0 h1:fs6U7nrV58d3CFAFh8VTde8TM262ObYf3ODrc//Lp+c= +github.com/knadh/koanf/providers/file v0.1.0/go.mod h1:rjJ/nHQl64iYCtAW2QQnF0eSmDEX/YZ/eNFj5yR6BvA= +github.com/knadh/koanf/providers/rawbytes v0.1.0 h1:dpzgu2KO6uf6oCb4aP05KDmKmAmI51k5pe8RYKQ0qME= +github.com/knadh/koanf/providers/rawbytes v0.1.0/go.mod h1:mMTB1/IcJ/yE++A2iEZbY1MLygX7vttU+C+S/YmPu9c= +github.com/knadh/koanf/v2 v2.0.1 h1:1dYGITt1I23x8cfx8ZnldtezdyaZtfAuRtIFOiRzK7g= +github.com/knadh/koanf/v2 v2.0.1/go.mod h1:ZeiIlIDXTE7w1lMT6UVcNiRAS2/rCeLn/GdLNvY1Dus= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= @@ -376,8 +357,6 @@ github.com/labstack/echo/v4 v4.11.4 h1:vDZmA+qNeh1pd/cCkEicDMrjtrnMGQ1QFI9gWN1zG github.com/labstack/echo/v4 v4.11.4/go.mod h1:noh7EvLwqDsmh/X/HWKPUl1AjzJrhyptRyEbQJfxen8= github.com/labstack/gommon v0.4.2 h1:F8qTUNXgG1+6WQmqoUWnz8WiEU60mXVVw0P4ht1WRA0= github.com/labstack/gommon v0.4.2/go.mod h1:QlUFxVM+SNXhDL/Z7YhocGIBYOiwB0mXm1+1bAPHPyU= -github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q= -github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4= github.com/lib/pq v1.0.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= github.com/lib/pq v1.1.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= @@ -406,10 +385,14 @@ github.com/mattn/go-sqlite3 v1.14.19/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= github.com/mholt/archiver/v3 v3.5.1 h1:rDjOBX9JSF5BvoJGvjqK479aL70qh9DIpZCl+k7Clwo= github.com/mholt/archiver/v3 v3.5.1/go.mod h1:e3dqJ7H78uzsRSEACH1joayhuSyhnonssnDhppzS1L4= +github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw= +github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s= github.com/mitchellh/hashstructure/v2 v2.0.2 h1:vGKWl0YJqUNxE8d+h8f6NJLcCJrgbhC4NcD46KavDd4= github.com/mitchellh/hashstructure/v2 v2.0.2/go.mod h1:MG3aRVU/N29oo/V/IhBX8GR/zz4kQkprJgF2EVszyDE= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= +github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ= +github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw= github.com/moby/sys/mountinfo v0.6.2/go.mod h1:IJb6JQeOklcdMU9F5xQ8ZALD+CUr5VlGpwtX+VE0rpI= github.com/moby/sys/mountinfo v0.7.1 h1:/tTvQaSJRr2FshkhXiIpux6fQ2Zvc4j7tAhMTStAG2g= github.com/moby/sys/mountinfo v0.7.1/go.mod h1:IJb6JQeOklcdMU9F5xQ8ZALD+CUr5VlGpwtX+VE0rpI= @@ -454,23 +437,22 @@ github.com/ory/dockertest/v3 v3.10.0 h1:4K3z2VMe8Woe++invjaTB7VRyQXQy5UY+loujO4a github.com/ory/dockertest/v3 v3.10.0/go.mod h1:nr57ZbRWMqfsdGdFNLHz5jjNdDb7VVFnzAeW1n5N1Lg= github.com/paulmach/orb v0.10.0 h1:guVYVqzxHE/CQ1KpfGO077TR0ATHSNjp4s6XGLn3W9s= github.com/paulmach/orb v0.10.0/go.mod h1:5mULz1xQfs3bmQm63QEJA6lNGujuRafwA5S/EnuLaLU= -github.com/pelletier/go-toml/v2 v2.1.0 h1:FnwAJ4oYMvbT/34k9zzHuZNrhlz48GB3/s6at6/MHO4= -github.com/pelletier/go-toml/v2 v2.1.0/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= github.com/pierrec/lz4/v4 v4.1.2/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pierrec/lz4/v4 v4.1.19 h1:tYLzDnjDXh9qIxSTKHwXwOYmm9d887Y7Y1ZkyXYHAN4= github.com/pierrec/lz4/v4 v4.1.19/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= -github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pressly/goose/v3 v3.17.0 h1:fT4CL3LRm4kfyLuPWzDFAoxjR5ZHjeJ6uQhibQtBaIs= github.com/pressly/goose/v3 v3.17.0/go.mod h1:22aw7NpnCPlS86oqkO/+3+o9FuCaJg4ZVWRUO3oGzHQ= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo= github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo= -github.com/redis/go-redis/v9 v9.1.0 h1:137FnGdk+EQdCbye1FW+qOEcY5S+SpY9T0NiuqvtfMY= -github.com/redis/go-redis/v9 v9.1.0/go.mod h1:urWj3He21Dj5k4TK1y59xH8Uj6ATueP8AH1cY3lZl4c= +github.com/redis/go-redis/v9 v9.4.0 h1:Yzoz33UZw9I/mFhx4MNrB6Fk+XHO1VukNcCa1+lwyKk= +github.com/redis/go-redis/v9 v9.4.0/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/fastuuid v1.1.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ= @@ -485,8 +467,6 @@ github.com/rs/zerolog v1.15.0/go.mod h1:xYTKnLHcpfU2225ny5qZjxnj9NvkumZYjJHlAThC github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= -github.com/samber/lo v1.39.0 h1:4gTz1wUhNYLhFSKl6O+8peW0v2F4BCY034GRpU9WnuA= -github.com/samber/lo v1.39.0/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA= github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys= github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= @@ -523,7 +503,6 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= -github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/tidwall/btree v1.6.0 h1:LDZfKfQIBHGHWSwckhXI0RPSXzlo+KYdjK7FWSqOzzg= @@ -534,10 +513,6 @@ github.com/tj/assert v0.0.3/go.mod h1:Ne6X72Q+TB1AteidzQncjw9PabbMp4PBMZ1k+vd1Pv github.com/tj/go-elastic v0.0.0-20171221160941-36157cbbebc2/go.mod h1:WjeM0Oo1eNAjXGDx2yma7uG2XoyRZTq1uv3M/o7imD0= github.com/tj/go-kinesis v0.0.0-20171128231115-08b17f58cb1b/go.mod h1:/yhzCV0xPfx6jb1bBgRFjl5lytqVqZXEaeqWP8lTEao= github.com/tj/go-spin v1.1.0/go.mod h1:Mg1mzmePZm4dva8Qz60H2lHwmJ2loum4VIrLgVnKwh4= -github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= -github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= -github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU= -github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= github.com/ulikunitz/xz v0.5.8/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/ulikunitz/xz v0.5.9/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/ulikunitz/xz v0.5.11 h1:kpFauv27b6ynzBNT/Xy+1k+fK4WswhN/6PN5WhFAGw8= @@ -588,9 +563,6 @@ go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9E go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= go.uber.org/zap v1.13.0/go.mod h1:zwrFLgMcdUuIBviXEYEH1YKNaOBnKXsx2IPda5bBwHM= -golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= -golang.org/x/arch v0.5.0 h1:jpGode6huXQxcskEIpOCvrU+tzo81b6+oFLUYXWtH/Y= -golang.org/x/arch v0.5.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190411191339-88737f569e3a/go.mod h1:WFFai1msRO1wXaEeE5yQxYXgSfI8pQAWXbQop6sCtWE= @@ -608,8 +580,6 @@ golang.org/x/crypto v0.0.0-20220722155217-630584e8d5aa/go.mod h1:IxCIyHEi3zRg3s0 golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k= golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= -golang.org/x/exp v0.0.0-20231219180239-dc181d75b848 h1:+iq7lrkxmFNBM7xx+Rae2W6uyPfhPeDWD+n+JgppptE= -golang.org/x/exp v0.0.0-20231219180239-dc181d75b848/go.mod h1:iRJReGqOEeBhDZGkGbynYwcHlctCvnjTYIamk7uXpHI= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= @@ -671,10 +641,9 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc= -golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= +golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= @@ -807,8 +776,6 @@ modernc.org/strutil v1.2.0 h1:agBi9dp1I+eOnxXeiZawM8F4LawKv4NzGWSaLfyeNZA= modernc.org/strutil v1.2.0/go.mod h1:/mdcBmfOibveCTBxUl5B5l6W+TTH1FXPLHZE6bTosX0= modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= -nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= -rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd h1:EDPBXCAspyGV4jQlpZSudPeMmr1bNJefnuqLsRAsHZo= sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd/go.mod h1:B8JuhiUyNFVKdsE8h686QcCxMaH6HrOAZj4vswFpcB0= sigs.k8s.io/structured-merge-diff/v4 v4.4.1 h1:150L+0vs/8DA78h1u02ooW1/fFq/Lwr+sGiqlzvrtq4= diff --git a/hack/okteto.yaml b/hack/okteto.yaml index 1abb081fb..803546589 100644 --- a/hack/okteto.yaml +++ b/hack/okteto.yaml @@ -6,12 +6,10 @@ dev: sync: - ../:/workspace environment: - DEFAULT_NAMESPACE: ${OKTETO_NAMESPACE} BUILD_COMMAND: | go build -o /workspace/bin/beam /workspace/cmd/beam/main.go go build -o /workspace/bin/worker /workspace/cmd/worker/main.go BUILD_BINARY_PATH: /workspace/bin/beam - CONFIG_PATH: /workspace/.secrets forward: - 1993:1993 - 1994:1994 diff --git a/internal/abstractions/image/build.go b/internal/abstractions/image/build.go index 356cf76c7..416429141 100644 --- a/internal/abstractions/image/build.go +++ b/internal/abstractions/image/build.go @@ -7,7 +7,6 @@ import ( "encoding/hex" "fmt" "log" - "os" "regexp" "strings" "time" @@ -47,13 +46,7 @@ type BuildOpts struct { ForceRebuild bool } -func NewBuilder(scheduler *scheduler.Scheduler, containerRepo repository.ContainerRepository) (*Builder, error) { - storeName := common.Secrets().GetWithDefault("BEAM_IMAGESERVICE_IMAGE_REGISTRY_STORE", "s3") - registry, err := common.NewImageRegistry(storeName) - if err != nil { - return nil, err - } - +func NewBuilder(registry *common.ImageRegistry, scheduler *scheduler.Scheduler, containerRepo repository.ContainerRepository) (*Builder, error) { return &Builder{ scheduler: scheduler, registry: registry, @@ -115,10 +108,6 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co } } - if opts.BaseImageRegistry == "" { - opts.BaseImageRegistry = os.Getenv("BEAM_RUNNER_BASE_IMAGE_REGISTRY") - } - baseImageId, err := b.GetImageId(&BuildOpts{ BaseImageRegistry: opts.BaseImageRegistry, BaseImageName: opts.BaseImageName, diff --git a/internal/abstractions/image/image.go b/internal/abstractions/image/image.go index f14df324a..edfd19685 100644 --- a/internal/abstractions/image/image.go +++ b/internal/abstractions/image/image.go @@ -2,12 +2,12 @@ package image import ( "context" - "fmt" "log" "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/scheduler" + "github.com/beam-cloud/beam/internal/types" pb "github.com/beam-cloud/beam/proto" "github.com/pkg/errors" ) @@ -22,16 +22,28 @@ type RuncImageService struct { pb.UnimplementedImageServiceServer builder *Builder scheduler *scheduler.Scheduler + config types.ImageServiceConfig } -func NewRuncImageService(ctx context.Context, scheduler *scheduler.Scheduler, containerRepo repository.ContainerRepository) (ImageService, error) { - builder, err := NewBuilder(scheduler, containerRepo) +func NewRuncImageService( + ctx context.Context, + config types.ImageServiceConfig, + scheduler *scheduler.Scheduler, + containerRepo repository.ContainerRepository, +) (ImageService, error) { + registry, err := common.NewImageRegistry(config) + if err != nil { + return nil, err + } + + builder, err := NewBuilder(registry, scheduler, containerRepo) if err != nil { return nil, err } return &RuncImageService{ builder: builder, + config: config, }, nil } @@ -39,12 +51,13 @@ func (is *RuncImageService) VerifyImageBuild(ctx context.Context, in *pb.VerifyI var valid bool = true imageId, err := is.builder.GetImageId(&BuildOpts{ - BaseImageName: common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_NAME"), - BaseImageTag: is.getBaseImageTag(in.PythonVersion), - PythonVersion: in.PythonVersion, - PythonPackages: in.PythonPackages, - Commands: in.Commands, - ExistingImageUri: in.ExistingImageUri, + BaseImageTag: is.config.Runner.Tags[in.PythonVersion], + BaseImageName: is.config.Runner.BaseImageName, + BaseImageRegistry: is.config.Runner.BaseImageRegistry, + PythonVersion: in.PythonVersion, + PythonPackages: in.PythonPackages, + Commands: in.Commands, + ExistingImageUri: in.ExistingImageUri, }) if err != nil { valid = false @@ -61,12 +74,13 @@ func (is *RuncImageService) BuildImage(in *pb.BuildImageRequest, stream pb.Image log.Printf("incoming image build request: %+v", in) buildOptions := &BuildOpts{ - BaseImageName: common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_NAME"), - BaseImageTag: is.getBaseImageTag(in.PythonVersion), - PythonVersion: in.PythonVersion, - PythonPackages: in.PythonPackages, - Commands: in.Commands, - ExistingImageUri: in.ExistingImageUri, + BaseImageTag: is.config.Runner.Tags[in.PythonVersion], + BaseImageName: is.config.Runner.BaseImageName, + BaseImageRegistry: is.config.Runner.BaseImageRegistry, + PythonVersion: in.PythonVersion, + PythonPackages: in.PythonPackages, + Commands: in.Commands, + ExistingImageUri: in.ExistingImageUri, } ctx := stream.Context() @@ -96,27 +110,3 @@ func (is *RuncImageService) BuildImage(in *pb.BuildImageRequest, stream pb.Image log.Println("build completed successfully") return nil } - -// Get latest image tag from secrets -func (is *RuncImageService) getBaseImageTag(pythonVersion string) string { - var baseImageTag string - - switch pythonVersion { - case "python3.8": - baseImageTag = fmt.Sprintf("py38-%s", common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_TAG")) - - case "python3.9": - baseImageTag = fmt.Sprintf("py39-%s", common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_TAG")) - - case "python3.10": - baseImageTag = fmt.Sprintf("py310-%s", common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_TAG")) - - case "python3.11": - baseImageTag = fmt.Sprintf("py311-%s", common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_TAG")) - - case "python3.12": - baseImageTag = fmt.Sprintf("py312-%s", common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_TAG")) - } - - return baseImageTag -} diff --git a/internal/common/config.default.yaml b/internal/common/config.default.yaml new file mode 100644 index 000000000..8cd44b9a7 --- /dev/null +++ b/internal/common/config.default.yaml @@ -0,0 +1,85 @@ +debugMode: false +database: + postgres: + host: postgres-postgresql.beam + port: 5432 + username: root + password: password + name: main + timezone: UTC + redis: + mode: single + addrs: + - redis-master.beam:6379 + password: + enableTLS: false + dialTimeout: 3s +storage: + mode: juicefs + fsName: beam-fs + fsPath: /data + objectPath: /data/objects + juicefs: + redisURI: redis://juicefs-redis-master.beam:6379/0 + awsS3Bucket: http://localstack.beam:4566/juicefs + awsAccessKeyID: test + awsSecretAccessKey: test +gateway: + host: beam.beam + port: 1993 +imageService: + cacheURL: + registryStore: local + registryCredentialProvider: docker + registries: + docker: + username: beamcloud + password: + s3: + bucket: beam-images + region: us-east-1 + accessKeyID: test + secretAccessKey: test + runner: + baseImageTag: latest + baseImageName: beam-runner + baseImageRegistry: registry.beam.lan:5000 + tags: + python3.8: py38-latest + python3.9: py39-latest + python3.10: py310-latest + python3.11: py311-latest + python3.12: py312-latest +worker: + pools: + beam-cpu: + jobSpec: + nodeSelector: {} + poolSizing: + defaultWorkerCpu: 1000m + defaultWorkerGpuType: "" + defaultWorkerMemory: 1Gi + minFreeCpu: 1000m + minFreeGpu: 0 + minFreeMemory: 1Gi + # global pool attributes + hostNetwork: false + imageTag: latest + imageName: beam-worker + imageRegistry: registry.beam.lan:5000 + imagePullSecrets: [] + namespace: beam + serviceAccountName: default + # non-standard k8s job spec + resourcesEnforced: false + dataVolumeName: worker-data + defaultWorkerCPURequest: 2000 + defaultWorkerMemoryRequest: 1024 +metrics: + kinesis: + streamName: beam-realtime-metrics + region: us-east-1 + accessKeyID: test + secretAccessKey: test + sessionKey: test + stackEndpoint: http://localstack.beam:4566 diff --git a/internal/common/config.go b/internal/common/config.go new file mode 100644 index 000000000..9a8cfb496 --- /dev/null +++ b/internal/common/config.go @@ -0,0 +1,127 @@ +package common + +import ( + _ "embed" + "errors" + "log" + "os" + "path/filepath" + + "github.com/knadh/koanf/parsers/json" + "github.com/knadh/koanf/parsers/yaml" + "github.com/knadh/koanf/providers/file" + "github.com/knadh/koanf/providers/rawbytes" + "github.com/knadh/koanf/v2" +) + +//go:embed config.default.yaml +var defaultConfig []byte + +// ConfigLoaderFunc is a function type used to load configuration into a Koanf +// instance. It takes a Koanf pointer 'k' as a parameter and returns an error +// if the loading process encounters any issues. +type ConfigLoaderFunc func(k *koanf.Koanf) error + +// ConfigManager is a generic configuration manager that allows handling and +// manipulation of configuration data for various types. It includes a Koanf +// instance ('kf') for managing configuration settings. +type ConfigManager[T any] struct { + kf *koanf.Koanf +} + +// NewConfigManager creates a new instance of the ConfigManager[T] type for +// managing configuration of type 'T'. It initializes the ConfigManager with +// the specified 'T' type, loads a default configuration, and optionally loads +// a user configuration if the 'CONFIG_PATH' environment variable is provided. +// If debug mode is enabled, it prints the current configuration. +func NewConfigManager[T any]() (*ConfigManager[T], error) { + // Initialize a ConfigManager[T] with the specified 'T' type. + cm := &ConfigManager[T]{ + kf: koanf.New("."), + } + + // Load default configuration + err := cm.LoadConfig(YAMLConfigFormat, rawbytes.Provider(defaultConfig)) + if err != nil { + return nil, err + } + + // Load user configuration if provided + cp := os.Getenv("CONFIG_PATH") + ce := filepath.Ext(cp) + if cp != "" && ce != "" { + if err := cm.LoadConfig(ConfigFormat(ce), file.Provider(cp)); err != nil { + return nil, err + } + } + + // If debug mode is enabled, print the current configuration. + if cm.kf.Bool("debugMode") { + log.Println("Debug mode enabled. Current configuration:") + log.Println(cm.Print()) + } + + return cm, nil +} + +// Print returns a string representation of the current configuration state. +func (cm *ConfigManager[T]) Print() string { + return cm.kf.Sprint() +} + +// GetConfig retrieves the current configuration of type 'T' from the ConfigManager. +// It unmarshals the configuration data and returns it. If any errors occur during +// unmarshaling, it logs a fatal error and exits the application. +func (cm *ConfigManager[T]) GetConfig() T { + var c T + err := cm.kf.UnmarshalWithConf("", &c, koanf.UnmarshalConf{Tag: "key", FlatPaths: false}) + if err != nil { + log.Fatal("failed to unmarshal config") + } + + return c +} + +// LoadConfig loads configuration data from a given provider in the specified format +// into the ConfigManager. It obtains a parser for the format, and then loads the +// configuration data. If any errors occur during the loading process, they are +// returned as an error. +func (cm *ConfigManager[T]) LoadConfig(format ConfigFormat, provider koanf.Provider) error { + parser, err := GetConfigParser(format) + if err != nil { + return err + } + + return cm.kf.Load(provider, parser) +} + +var ( + JSONConfigFormat ConfigFormat = "json" + YAMLConfigFormat ConfigFormat = "yaml" + YMLConfigFormat ConfigFormat = "yml" + + parserMap map[ConfigFormat]ParserFunc = map[ConfigFormat]ParserFunc{ + JSONConfigFormat: jsonParserFunc, + YAMLConfigFormat: yamlParserFunc, + YMLConfigFormat: yamlParserFunc, + } +) + +type ConfigFormat string + +type ParserFunc func() (koanf.Parser, error) + +func GetConfigParser(format ConfigFormat) (koanf.Parser, error) { + if parserFunc, ok := parserMap[format]; ok { + return parserFunc() + } + return nil, errors.New("parser not found for format" + string(format)) +} + +func jsonParserFunc() (koanf.Parser, error) { + return json.Parser(), nil +} + +func yamlParserFunc() (koanf.Parser, error) { + return yaml.Parser(), nil +} diff --git a/internal/common/redis.go b/internal/common/redis.go index 0bc3db8bd..78d63ba97 100644 --- a/internal/common/redis.go +++ b/internal/common/redis.go @@ -5,7 +5,6 @@ import ( "crypto/tls" "errors" "fmt" - "log" "reflect" "regexp" "strconv" @@ -13,130 +12,64 @@ import ( "sync" "time" + "github.com/beam-cloud/beam/internal/types" "github.com/bsm/redislock" "github.com/redis/go-redis/v9" ) var ( - ErrChannelClosed = errors.New("redis: channel closed") - ErrConnectionIssue = errors.New("redis: connection issue") - ErrUnknownRedisClient = errors.New("redis: unknown client type") - ErrUnknownRedisOptions = errors.New("redis: unknown options type") - ErrLockNotReleased = errors.New("redislock: lock not released") - RedisModeCluster = "cluster" + ErrChannelClosed = errors.New("redis: channel closed") + ErrConnectionIssue = errors.New("redis: connection issue") + ErrUnknownRedisMode = errors.New("redis: unknown mode") + ErrLockNotReleased = errors.New("redislock: lock not released") ) -type redisCmdable interface { - ClientGetName(ctx context.Context) *redis.StringCmd - Close() error - Del(ctx context.Context, keys ...string) *redis.IntCmd - Do(ctx context.Context, args ...interface{}) *redis.Cmd - Eval(ctx context.Context, script string, keys []string, args ...interface{}) *redis.Cmd - EvalRO(ctx context.Context, script string, keys []string, args ...interface{}) *redis.Cmd - EvalSha(ctx context.Context, sha1 string, keys []string, args ...interface{}) *redis.Cmd - EvalShaRO(ctx context.Context, sha1 string, keys []string, args ...interface{}) *redis.Cmd - Exists(ctx context.Context, keys ...string) *redis.IntCmd - Expire(ctx context.Context, key string, expiration time.Duration) *redis.BoolCmd - ExpireAt(ctx context.Context, key string, tm time.Time) *redis.BoolCmd - Get(ctx context.Context, key string) *redis.StringCmd - HGetAll(ctx context.Context, key string) *redis.MapStringStringCmd - HMSet(ctx context.Context, key string, values ...interface{}) *redis.BoolCmd - HSet(ctx context.Context, key string, values ...interface{}) *redis.IntCmd - Keys(ctx context.Context, pattern string) *redis.StringSliceCmd - LLen(ctx context.Context, key string) *redis.IntCmd - LPop(ctx context.Context, key string) *redis.StringCmd - Ping(ctx context.Context) *redis.StatusCmd - LRange(ctx context.Context, key string, start, stop int64) *redis.StringSliceCmd - PSubscribe(ctx context.Context, channels ...string) *redis.PubSub - Publish(ctx context.Context, channel string, message interface{}) *redis.IntCmd - RPush(ctx context.Context, key string, values ...interface{}) *redis.IntCmd - Scan(ctx context.Context, cursor uint64, match string, count int64) *redis.ScanCmd - ScriptExists(ctx context.Context, hashes ...string) *redis.BoolSliceCmd - ScriptLoad(ctx context.Context, script string) *redis.StringCmd - Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd - SetEx(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd - Subscribe(ctx context.Context, channels ...string) *redis.PubSub - ZAdd(ctx context.Context, key string, members ...redis.Z) *redis.IntCmd - ZCard(ctx context.Context, key string) *redis.IntCmd - ZPopMin(ctx context.Context, key string, count ...int64) *redis.ZSliceCmd - Incr(ctx context.Context, key string) *redis.IntCmd - Decr(ctx context.Context, key string) *redis.IntCmd -} - type RedisClient struct { - // Embedded Redis client. All functions defined on interface will be available - // to this struct. - redisCmdable + redis.UniversalClient } -func WithAddress(a string) func(*RedisOptions) { - return func(ro *RedisOptions) { - ro.Addr = a - } -} +func WithClientName(name string) func(*redis.UniversalOptions) { + // Remove empty spaces and new lines + name = strings.ReplaceAll(name, " ", "") + name = strings.ReplaceAll(name, "\n", "") -func WithClientName(n string) func(*RedisOptions) { - return func(ro *RedisOptions) { - // Remove empty spaces and new lines - n = strings.ReplaceAll(n, " ", "") - n = strings.ReplaceAll(n, "\n", "") + // Remove special characters using a regular expression + reg := regexp.MustCompile("[^a-zA-Z0-9]+") + name = reg.ReplaceAllString(name, "") - // Remove special characters using a regular expression - reg := regexp.MustCompile("[^a-zA-Z0-9]+") - n = reg.ReplaceAllString(n, "") - - ro.ClientName = n + return func(uo *redis.UniversalOptions) { + uo.ClientName = name } } -func NewRedisClient(opts ...func(*RedisOptions)) (*RedisClient, error) { - // Get options from secrets and env vars - options, err := GetRedisConfig() - if err != nil { - return nil, err - } +func NewRedisClient(config types.RedisConfig, options ...func(*redis.UniversalOptions)) (*RedisClient, error) { + opts := &redis.UniversalOptions{} + CopyStruct(&config, opts) - // Apply overrides - for _, opt := range opts { - opt(options) + for _, opt := range options { + opt(opts) } - // Optionally enable TLS - if options.TLSEnabled { - options.TLSConfig = &tls.Config{} + if config.EnableTLS { + opts.TLSConfig = &tls.Config{} } - // Initialize a client - var client *RedisClient - - if options.Mode == RedisModeCluster { - log.Println("Running Redis in cluster mode.") - - ro := &redis.ClusterOptions{} - CopyStruct(options, ro) - ro.Addrs = options.Addrs() - - client = &RedisClient{redis.NewClusterClient(ro)} - } else { - log.Println("Running Redis in standalone mode.") - - ro := &redis.Options{} - CopyStruct(options, ro) - - client = &RedisClient{redis.NewClient(ro)} + var client redis.UniversalClient + switch config.Mode { + case types.RedisModeSingle: + client = redis.NewClient(opts.Simple()) + case types.RedisModeCluster: + client = redis.NewClusterClient(opts.Cluster()) + default: + return nil, ErrUnknownRedisMode } - // Connection test - err = client.Ping(context.TODO()).Err() + err := client.Ping(context.TODO()).Err() if err != nil { return nil, fmt.Errorf("%s: %s", ErrConnectionIssue, err) } - return client, nil -} - -func (r *RedisClient) Close() error { - return r.redisCmdable.Close() + return &RedisClient{UniversalClient: client}, nil } func (r *RedisClient) ToSlice(v interface{}) []interface{} { @@ -180,7 +113,7 @@ func (r *RedisClient) Scan(ctx context.Context, pattern string) ([]string, error defer close(outCh) defer close(errCh) - switch client := r.redisCmdable.(type) { + switch client := r.UniversalClient.(type) { case *redis.Client: scanAndCollect(client) @@ -208,16 +141,16 @@ func (r *RedisClient) Scan(ctx context.Context, pattern string) ([]string, error } func (r *RedisClient) LRange(ctx context.Context, key string, start, stop int64) ([]string, error) { - return r.redisCmdable.LRange(ctx, key, start, stop).Result() + return r.UniversalClient.LRange(ctx, key, start, stop).Result() } func (r *RedisClient) Publish(ctx context.Context, channel string, message interface{}) *redis.IntCmd { - client, ok := r.redisCmdable.(*redis.ClusterClient) + client, ok := r.UniversalClient.(*redis.ClusterClient) if ok { return client.SPublish(ctx, channel, message) } - return r.redisCmdable.Publish(ctx, channel, message) + return r.UniversalClient.Publish(ctx, channel, message) } func (r *RedisClient) PSubscribe(ctx context.Context, channels ...string) (<-chan *redis.Message, <-chan error) { @@ -228,7 +161,7 @@ func (r *RedisClient) PSubscribe(ctx context.Context, channels ...string) (<-cha defer close(outCh) defer close(errCh) - switch client := r.redisCmdable.(type) { + switch client := r.UniversalClient.(type) { case *redis.Client: r.handleChannelSubs(ctx, client.PSubscribe, outCh, errCh, channels...) @@ -256,7 +189,7 @@ func (r *RedisClient) Subscribe(ctx context.Context, channels ...string) (<-chan defer close(outCh) defer close(errCh) - switch client := r.redisCmdable.(type) { + switch client := r.UniversalClient.(type) { case *redis.Client: r.handleChannelSubs(ctx, client.Subscribe, outCh, errCh, channels...) @@ -297,44 +230,6 @@ func (r *RedisClient) handleChannelSubs( } } -// https://pkg.go.dev/github.com/redis/go-redis/v9#Options -type RedisOptions struct { - Mode string `env:"REDIS_MODE"` - Addr string `env:"REDIS_HOSTS"` - ClientName string `env:"REDIS_CLIENT_NAME"` - ContextTimeoutEnabled bool `env:"REDIS_CONTEXT_TIMEOUT_ENABLED"` - MinIdleConns int `env:"REDIS_MIN_IDLE_CONNS"` - MaxIdleConns int `env:"REDIS_MAX_IDLE_CONNS"` - ConnMaxIdleTime time.Duration `env:"REDIS_CONN_MAX_IDLE_TIME"` - ConnMaxLifetime time.Duration `env:"REDIS_CONN_MAX_LIFETIME"` - DialTimeout time.Duration `env:"REDIS_DIAL_TIMEOUT"` - ReadTimeout time.Duration `env:"REDIS_READ_TIMEOUT"` - WriteTimeout time.Duration `env:"REDIS_WRITE_TIMEOUT"` - MaxRedirects int `env:"REDIS_MAX_REDIRECTS"` - MaxRetries int `env:"REDIS_MAX_RETRIES"` - PoolSize int `env:"REDIS_POOL_SIZE"` - Username string `env:"REDIS_USERNAME"` - Password string `env:"REDIS_PASSWORD"` - TLSEnabled bool `env:"REDIS_TLS_ENABLED"` - TLSConfig *tls.Config -} - -func (ro *RedisOptions) Addrs() []string { - return strings.Split(ro.Addr, ",") -} - -func GetRedisConfig() (*RedisOptions, error) { - secrets := Secrets() - options := &RedisOptions{} - - err := secrets.EnvUnmarshal(options) - if err != nil { - return nil, err - } - - return options, nil -} - type RedisLockOptions struct { TtlS int Retries int diff --git a/internal/common/redis_test.go b/internal/common/redis_test.go index 49da81b7d..7fa8c91cb 100644 --- a/internal/common/redis_test.go +++ b/internal/common/redis_test.go @@ -2,14 +2,13 @@ package common import ( "context" - "crypto/tls" "fmt" "math" - "strings" "testing" "time" "github.com/alicebob/miniredis/v2" + "github.com/beam-cloud/beam/internal/types" "github.com/google/uuid" "github.com/stretchr/testify/assert" ) @@ -20,7 +19,7 @@ func NewRedisClientForTest() (*RedisClient, error) { return nil, err } - return NewRedisClient(WithAddress(s.Addr())) + return NewRedisClient(types.RedisConfig{Addrs: []string{s.Addr()}, Mode: types.RedisModeSingle}) } func TestRedisLock(t *testing.T) { @@ -87,38 +86,9 @@ func TestRedisLockWithTTLAndRetry(t *testing.T) { assert.NoError(t, err) } -func TestWithClientName(t *testing.T) { - opts1 := &RedisOptions{} - WithClientName("beam")(opts1) - assert.Equal(t, "beam", opts1.ClientName) - - opts2 := &RedisOptions{} - WithClientName("My ^ App $")(opts2) - assert.Equal(t, "MyApp", opts2.ClientName) - - opts3 := &RedisOptions{} - WithClientName(" g /.\\ []o] o ! @ #d---n$% ^&a*()m-_=e+ ")(opts3) - assert.Equal(t, "goodname", opts3.ClientName) -} - -func TestWithAddress(t *testing.T) { - addr := "redis.beam.cloud:3452" - - opts1 := &RedisOptions{} - WithAddress(addr)(opts1) - assert.Equal(t, addr, opts1.Addr) - assert.Equal(t, []string{addr}, opts1.Addrs()) - - addr = "redis1.beam.cloud:3452,redis2.beam.cloud:3452" - opts2 := &RedisOptions{} - WithAddress(addr)(opts2) - assert.Equal(t, addr, opts2.Addr) - assert.Equal(t, strings.Split(addr, ","), opts2.Addrs()) -} - func TestCopyStruct(t *testing.T) { - options1 := &RedisOptions{ClientName: "hello", PoolSize: 10, ConnMaxLifetime: time.Second, TLSConfig: &tls.Config{MinVersion: tls.VersionTLS13}} - options2 := &RedisOptions{} + options1 := &types.RedisConfig{ClientName: "hello", PoolSize: 10, ConnMaxLifetime: time.Second} + options2 := &types.RedisConfig{} CopyStruct(options1, options2) diff --git a/internal/common/registry.go b/internal/common/registry.go index 252ea4ea9..4f7dc6a61 100644 --- a/internal/common/registry.go +++ b/internal/common/registry.go @@ -13,7 +13,8 @@ import ( "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/feature/s3/manager" "github.com/aws/aws-sdk-go-v2/service/s3" - "github.com/aws/aws-sdk-go-v2/service/s3/types" + awstypes "github.com/aws/aws-sdk-go-v2/service/s3/types" + "github.com/beam-cloud/beam/internal/types" ) const ( @@ -26,18 +27,19 @@ const ( type ImageRegistry struct { store ObjectStore + config types.ImageServiceConfig ImageFileExtension string } -func NewImageRegistry(storeName string) (*ImageRegistry, error) { +func NewImageRegistry(config types.ImageServiceConfig) (*ImageRegistry, error) { var err error var store ObjectStore var imageFileExtension string = "clip" - switch storeName { + switch config.RegistryStore { case s3ImageRegistryStoreName: imageFileExtension = remoteImageFileExtension - store, err = NewS3Store() + store, err = NewS3Store(config.Registries.S3) if err != nil { return nil, err } @@ -51,6 +53,7 @@ func NewImageRegistry(storeName string) (*ImageRegistry, error) { return &ImageRegistry{ store: store, + config: config, ImageFileExtension: imageFileExtension, }, nil } @@ -78,11 +81,8 @@ type ObjectStore interface { Size(ctx context.Context, key string) (int64, error) } -func NewS3Store() (*S3Store, error) { - accessKey := os.Getenv("AWS_ACCESS_KEY_ID") - secretKey := os.Getenv("AWS_SECRET_ACCESS_KEY") - - cfg, err := getAWSConfig(accessKey, secretKey, Secrets().GetWithDefault("AWS_REGION", DefaultAWSRegion)) +func NewS3Store(config types.S3ImageRegistryConfig) (*S3Store, error) { + cfg, err := getAWSConfig(config.AccessKeyID, config.SecretAccessKey, config.Region) if err != nil { return nil, err } @@ -108,6 +108,7 @@ func getAWSConfig(accessKey string, secretKey string, region string) (aws.Config type S3Store struct { client *s3.Client + config types.S3ImageRegistryConfig } func (s *S3Store) Put(ctx context.Context, localPath string, key string) error { @@ -118,10 +119,10 @@ func (s *S3Store) Put(ctx context.Context, localPath string, key string) error { } defer file.Close() - var sse types.ServerSideEncryption = "aws:kms" + var sse awstypes.ServerSideEncryption = "aws:kms" uploader := manager.NewUploader(s.client) _, err = uploader.Upload(ctx, &s3.PutObjectInput{ - Bucket: aws.String(Secrets().Get("BEAM_IMAGESERVICE_IMAGE_REGISTRY_S3_BUCKET")), + Bucket: aws.String(s.config.Bucket), Key: aws.String(key), Body: file, ServerSideEncryption: sse, @@ -143,7 +144,7 @@ func (s *S3Store) Get(ctx context.Context, key string, localPath string) error { downloader := manager.NewDownloader(s.client) _, err = downloader.Download(ctx, f, &s3.GetObjectInput{ - Bucket: aws.String(Secrets().Get("BEAM_IMAGESERVICE_IMAGE_REGISTRY_S3_BUCKET")), + Bucket: aws.String(s.config.Bucket), Key: aws.String(key), }, func(d *manager.Downloader) { d.PartSize = 100 * 1024 * 1024 // 100MiB per part @@ -178,7 +179,7 @@ func (s *S3Store) Size(ctx context.Context, key string) (int64, error) { // headObject returns the metadata of an object func (s *S3Store) headObject(ctx context.Context, key string) (*s3.HeadObjectOutput, error) { return s.client.HeadObject(ctx, &s3.HeadObjectInput{ - Bucket: aws.String(Secrets().Get("BEAM_IMAGESERVICE_IMAGE_REGISTRY_S3_BUCKET")), + Bucket: aws.String(s.config.Bucket), Key: aws.String(key), }) } diff --git a/internal/common/secrets.go b/internal/common/secrets.go deleted file mode 100644 index 8370c93b9..000000000 --- a/internal/common/secrets.go +++ /dev/null @@ -1,192 +0,0 @@ -package common - -import ( - "errors" - "fmt" - "log" - "os" - "reflect" - "strconv" - "sync" - "time" - - "github.com/aws/karpenter/pkg/utils/env" - "github.com/joho/godotenv" -) - -const ( - EnvLocal string = "local" - EnvStaging string = "staging" - EnvProduction string = "production" -) - -// Secrets manager -var secretsFileName string = env.WithDefaultString("CONFIG_PATH", ".secrets") -var secretsStore *SecretStore -var onceSecrets sync.Once - -type SecretStore struct { - loaded bool -} - -func Secrets() *SecretStore { - onceSecrets.Do(func() { - secretsStore = &SecretStore{} - err := secretsStore.load(secretsFileName) - if err != nil { - log.Fatalf("Failed to load secrets: %v\n", err) - } - }) - - return secretsStore -} - -func (s *SecretStore) load(filename string) error { - if _, err := os.Stat(filename); errors.Is(err, os.ErrNotExist) { - return fmt.Errorf("secrets file not found: %v", err) - } - - godotenv.Load(filename) - s.loaded = true - return nil -} - -// Retrieve a secret as a raw string -func (s *SecretStore) Get(key string) string { - if !s.loaded { - s.load(secretsFileName) - } - - return os.Getenv(key) -} - -// Retrieve a secret as a raw string, or return a default value -func (s *SecretStore) GetWithDefault(key string, defaultValue string) string { - val := s.Get(key) - if val == "" { - return defaultValue - } - return val -} - -// Retrieve a secret and convert to a signed integer -func (s *SecretStore) GetInt(key string) int { - val := s.Get(key) - intVal, err := strconv.Atoi(val) - - if err != nil { - log.Printf("Invalid value: %s", val) - return -1 - } - return intVal -} - -// Retrieve a secret and convert to a signed integer -func (s *SecretStore) GetIntWithDefault(key string, defaultValue int) int { - val := s.Get(key) - - intVal, err := strconv.Atoi(val) - if err != nil { - return defaultValue - } - - return intVal -} - -// Load and map env vars to a struct using struct tags. -func (s *SecretStore) EnvUnmarshal(st interface{}) error { - val := reflect.ValueOf(st).Elem() - typ := val.Type() - - for i := 0; i < val.NumField(); i++ { - field := val.Field(i) - if !field.CanSet() { - continue - } - - // Get env var name or continue - envTag := typ.Field(i).Tag.Get("env") - if envTag == "" || envTag == "-" { - continue - } - - // Get env var, use the default, or continue - envVal := os.Getenv(envTag) - if envVal == "" { - defTag := typ.Field(i).Tag.Get("default") - if defTag == "" { - continue - } - envVal = defTag - } - - switch field.Kind() { - - case reflect.String: - field.SetString(envVal) - - case reflect.Int32: - val, err := strconv.ParseInt(envVal, 10, 32) - if err != nil { - return fmt.Errorf("could not convert <%s> to int32: %v", envTag, err) - } - field.SetInt(val) - - case reflect.Int, reflect.Int64: - // Duration is recognized as an int64 - if field.Type() == reflect.ValueOf(time.Duration(0)).Type() { - val, err := time.ParseDuration(envVal) - if err != nil { - return fmt.Errorf("could not convert <%s> to duration: %v", envTag, err) - } - field.Set(reflect.ValueOf(val)) - } else { - val, err := strconv.ParseInt(envVal, 10, 64) - if err != nil { - return fmt.Errorf("could not convert <%s> to int64: %v", envTag, err) - } - field.SetInt(val) - } - - case reflect.Uint32: - val, err := strconv.ParseUint(envVal, 10, 32) - if err != nil { - return fmt.Errorf("could not convert <%s> to uint32: %v", envTag, err) - } - field.SetUint(val) - - case reflect.Uint64: - val, err := strconv.ParseUint(envVal, 10, 64) - if err != nil { - return fmt.Errorf("could not convert <%s> to uint64: %v", envTag, err) - } - field.SetUint(val) - - case reflect.Float32: - val, err := strconv.ParseFloat(envVal, 32) - if err != nil { - return fmt.Errorf("could not convert <%s> to float32: %v", envTag, err) - } - field.SetFloat(val) - - case reflect.Float64: - val, err := strconv.ParseFloat(envVal, 64) - if err != nil { - return fmt.Errorf("could not convert <%s> to float64: %v", envTag, err) - } - field.SetFloat(val) - - case reflect.Bool: - val, err := strconv.ParseBool(envVal) - if err != nil { - return fmt.Errorf("could not convert <%s> to bool: %v", envTag, err) - } - field.SetBool(val) - - default: - log.Printf("parsing tag <%s> of type <%s> is not supported", envTag, field.Kind()) - } - - } - return nil -} diff --git a/internal/common/secrets_test.go b/internal/common/secrets_test.go deleted file mode 100644 index c748e353b..000000000 --- a/internal/common/secrets_test.go +++ /dev/null @@ -1,35 +0,0 @@ -package common - -import ( - "fmt" - "os" - "testing" - "time" - - "github.com/tj/assert" -) - -func TestSecretsEnvUnmarshal(t *testing.T) { - type TestConfig struct { - Host string `env:"TEST_HOSTNAME"` - Port int `env:"TEST_PORT" default:"3306"` - Int32 int32 `env:"TEST_INT32"` - BoolTrue bool `env:"TEST_BOOL_TRUE"` - BoolFalse bool `env:"TEST_BOOL_FALSE"` - Duration time.Duration `env:"TEST_DURATION"` - } - - os.Setenv("TEST_HOSTNAME", "beam.cloud") - os.Setenv("TEST_INT32", "77") - os.Setenv("TEST_BOOL_TRUE", "t") - os.Setenv("TEST_BOOL_FALSE", "0") - os.Setenv("TEST_DURATION", fmt.Sprintf("%v", 3600*time.Second)) - - expect := &TestConfig{Port: 3306, Host: "beam.cloud", Int32: 77, BoolTrue: true, BoolFalse: false, Duration: 3600 * time.Second} - actual := &TestConfig{} - - secrets := Secrets() - err := secrets.EnvUnmarshal(actual) - assert.NoError(t, err) - assert.Equal(t, expect, actual) -} diff --git a/internal/common/stats.go b/internal/common/stats.go index 7dc0c9a6c..e5504388f 100644 --- a/internal/common/stats.go +++ b/internal/common/stats.go @@ -3,7 +3,6 @@ package common import ( "fmt" "io" - "log" "net" "time" ) @@ -18,9 +17,9 @@ var sender *StatsdSender func (s *StatsdSender) begin() { s.queue = make(chan string, 100) if s.uri == "" { - sec := Secrets() - host := sec.GetWithDefault("STATSD_HOST", fmt.Sprintf("beam-statsd.%s.svc.cluster.local", sec.Get("BEAM_NAMESPACE"))) - port := sec.GetIntWithDefault("STATSD_PORT", 8125) + // TODO: Read from config or env + host := "statsd.beam" + port := 8125 s.uri = fmt.Sprintf("%s:%v", host, port) } @@ -41,8 +40,6 @@ func (s *StatsdSender) gauge(metric string, value int) { } func (s *StatsdSender) process() { - log.Println("Connecting to statsd at uri", s.uri) - for msg := range s.queue { if conn, err := net.Dial("udp", s.uri); err == nil { io.WriteString(conn, msg) diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 2accfac91..cb99f810b 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -20,10 +20,12 @@ import ( apiv1 "github.com/beam-cloud/beam/internal/api/v1" "github.com/beam-cloud/beam/internal/auth" - common "github.com/beam-cloud/beam/internal/common" + "github.com/beam-cloud/beam/internal/common" + "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/scheduler" "github.com/beam-cloud/beam/internal/storage" + "github.com/beam-cloud/beam/internal/types" pb "github.com/beam-cloud/beam/proto" "github.com/labstack/echo/v4" "github.com/labstack/echo/v4/middleware" @@ -32,6 +34,7 @@ import ( type Gateway struct { pb.UnimplementedSchedulerServer + config types.AppConfig httpServer *echo.Echo grpcServer *grpc.Server redisClient *common.RedisClient @@ -47,17 +50,23 @@ type Gateway struct { } func NewGateway() (*Gateway, error) { - redisClient, err := common.NewRedisClient(common.WithClientName("BeamGateway")) + configManager, err := common.NewConfigManager[types.AppConfig]() + if err != nil { + return nil, err + } + config := configManager.GetConfig() + + redisClient, err := common.NewRedisClient(config.Database.Redis, common.WithClientName("BeamGateway")) if err != nil { return nil, err } - Scheduler, err := scheduler.NewScheduler() + scheduler, err := scheduler.NewScheduler(config, redisClient) if err != nil { return nil, err } - Storage, err := storage.NewStorage() + storage, err := storage.NewStorage(config.Storage) if err != nil { return nil, err } @@ -67,16 +76,16 @@ func NewGateway() (*Gateway, error) { redisClient: redisClient, ctx: ctx, cancelFunc: cancel, - Storage: Storage, - Scheduler: Scheduler, + Storage: storage, + Scheduler: scheduler, } - beamRepo, err := repository.NewBeamPostgresRepository() + beamRepo, err := repository.NewBeamPostgresRepository(config.Database.Postgres) if err != nil { return nil, err } - backendRepo, err := repository.NewBackendPostgresRepository() + backendRepo, err := repository.NewBackendPostgresRepository(config.Database.Postgres) if err != nil { return nil, err } @@ -84,6 +93,7 @@ func NewGateway() (*Gateway, error) { containerRepo := repository.NewContainerRedisRepository(redisClient) metricsRepo := repository.NewMetricsStatsdRepository() + gateway.config = config gateway.ContainerRepo = containerRepo gateway.BackendRepo = backendRepo gateway.BeamRepo = beamRepo @@ -139,7 +149,7 @@ func (g *Gateway) registerServices() error { pb.RegisterSimpleQueueServiceServer(g.grpcServer, rq) // Register image service - is, err := image.NewRuncImageService(g.ctx, g.Scheduler, g.ContainerRepo) + is, err := image.NewRuncImageService(g.ctx, g.config.ImageService, g.Scheduler, g.ContainerRepo) if err != nil { return err } @@ -167,7 +177,7 @@ func (g *Gateway) registerServices() error { pb.RegisterVolumeServiceServer(g.grpcServer, vs) // Register scheduler - s, err := scheduler.NewSchedulerService() + s, err := scheduler.NewSchedulerService(g.config, g.redisClient) if err != nil { return err } diff --git a/internal/repository/backend_postgres.go b/internal/repository/backend_postgres.go index 9bfb74ec8..675d1935d 100644 --- a/internal/repository/backend_postgres.go +++ b/internal/repository/backend_postgres.go @@ -10,7 +10,6 @@ import ( "log" "strings" - "github.com/beam-cloud/beam/internal/common" _ "github.com/beam-cloud/beam/internal/repository/backend_postgres_migrations" "github.com/beam-cloud/beam/internal/types" "github.com/google/uuid" @@ -23,14 +22,22 @@ type PostgresBackendRepository struct { client *sqlx.DB } -func NewBackendPostgresRepository() (*PostgresBackendRepository, error) { - host := common.Secrets().Get("DB_HOST") - port := common.Secrets().GetInt("DB_PORT") - user := common.Secrets().Get("DB_USER") - password := common.Secrets().Get("DB_PASS") - dbName := common.Secrets().Get("DB_NAME") - - dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=disable TimeZone=UTC", host, user, password, dbName, port) +func NewBackendPostgresRepository(config types.PostgresConfig) (*PostgresBackendRepository, error) { + sslMode := "disable" + if config.EnableTLS { + sslMode = "require" + } + + dsn := fmt.Sprintf( + "host=%s user=%s password=%s dbname=%s port=%d sslmode=%s TimeZone=%s", + config.Host, + config.Username, + config.Password, + config.Name, + config.Port, + sslMode, + config.TimeZone, + ) db, err := sqlx.Connect("postgres", dsn) if err != nil { return nil, err @@ -269,7 +276,7 @@ func (r *PostgresBackendRepository) GetTask(ctx context.Context, externalId stri func (r *PostgresBackendRepository) GetTaskWithRelated(ctx context.Context, externalId string) (*types.TaskWithRelated, error) { var taskWithRelated types.TaskWithRelated query := ` - SELECT w.external_id AS "workspace.external_id", w.name AS "workspace.name", + SELECT w.external_id AS "workspace.external_id", w.name AS "workspace.name", s.external_id AS "stub.external_id", s.name AS "stub.name", s.config AS "stub.config", t.* FROM task t JOIN workspace w ON t.workspace_id = w.id @@ -371,7 +378,7 @@ func (r *PostgresBackendRepository) GetStubByExternalId(ctx context.Context, ext var stub types.StubWithRelated query := ` - SELECT + SELECT s.id, s.external_id, s.name, s.type, s.config, s.config_version, s.object_id, s.workspace_id, s.created_at, s.updated_at, w.id AS "workspace.id", w.external_id AS "workspace.external_id", w.name AS "workspace.name", w.created_at AS "workspace.created_at", w.updated_at AS "workspace.updated_at", o.id AS "object.id", o.external_id AS "object.external_id", o.hash AS "object.hash", o.size AS "object.size", o.workspace_id AS "object.workspace_id", o.created_at AS "object.created_at" diff --git a/internal/repository/base.go b/internal/repository/base.go index 2bff1a00c..07603acce 100644 --- a/internal/repository/base.go +++ b/internal/repository/base.go @@ -88,9 +88,9 @@ type BeamRepository interface { } type WorkerPoolRepository interface { - GetPool(name string) (*types.WorkerPoolResource, error) - GetPools() ([]types.WorkerPoolResource, error) - SetPool(pool *types.WorkerPoolResource) error + GetPool(name string) (*types.WorkerPoolConfig, error) + GetPools() ([]types.WorkerPoolConfig, error) + SetPool(name string, pool *types.WorkerPoolConfig) error RemovePool(name string) error } diff --git a/internal/repository/beam_postgres.go b/internal/repository/beam_postgres.go index 5412c2480..764fdbaa9 100644 --- a/internal/repository/beam_postgres.go +++ b/internal/repository/beam_postgres.go @@ -9,7 +9,6 @@ import ( "strings" "time" - "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/types" "github.com/google/uuid" "gorm.io/driver/postgres" @@ -23,7 +22,7 @@ type PostgresBeamRepository struct { client *gorm.DB } -func NewBeamPostgresRepository() (BeamRepository, error) { +func NewBeamPostgresRepository(config types.PostgresConfig) (BeamRepository, error) { // If client is not nil, it was already initialized. Just return a new PostgresBeamRepository with the existing client. if client != nil { return &PostgresBeamRepository{ @@ -31,13 +30,22 @@ func NewBeamPostgresRepository() (BeamRepository, error) { }, nil } - host := common.Secrets().Get("DB_HOST") - port := common.Secrets().GetInt("DB_PORT") - user := common.Secrets().Get("DB_USER") - password := common.Secrets().Get("DB_PASS") - dbName := common.Secrets().Get("DB_NAME") + sslMode := "prefer" + if config.EnableTLS { + sslMode = "require" + } + + dsn := fmt.Sprintf( + "host=%s user=%s password=%s dbname=%s port=%d sslmode=%s TimeZone=%s", + config.Host, + config.Username, + config.Password, + config.Name, + config.Port, + sslMode, + config.TimeZone, + ) - dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%d sslmode=disable", host, user, password, dbName, port) var err error client, err = gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) if err != nil { diff --git a/internal/repository/metrics_stream_kinesis.go b/internal/repository/metrics_stream_kinesis.go index 9501bda0a..f8d238e86 100644 --- a/internal/repository/metrics_stream_kinesis.go +++ b/internal/repository/metrics_stream_kinesis.go @@ -3,36 +3,31 @@ package repository import ( "context" "encoding/json" - "os" "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/types" "github.com/okteto/okteto/pkg/log" ) -const ( - RealtimeMetricsStreamName = "beam-realtime-metrics" -) - type MetricsStream struct { client types.EventClient } -func NewMetricsStreamRepository(ctx context.Context) *MetricsStream { +func NewMetricsStreamRepository(ctx context.Context, config types.MetricsConfig) *MetricsStream { opts := []common.KinesisClientOption{} - // Localstack is used only for local (okteto) development - localstackEndpoint := common.Secrets().GetWithDefault("LOCALSTACK_ENDPOINT", "") - if localstackEndpoint != "" { - opts = append(opts, common.WithKinesisEndpoint(localstackEndpoint)) + if config.Kinesis.Endpoint != "" { + opts = append(opts, common.WithKinesisEndpoint(config.Kinesis.Endpoint)) + } + if config.Kinesis.AccessKeyID != "" && config.Kinesis.SecretAccessKey != "" { opts = append(opts, common.WithCredentials( - common.Secrets().Get("LOCALSTACK_AWS_ACCESS_KEY"), - common.Secrets().Get("LOCALSTACK_AWS_SECRET_KEY"), - common.Secrets().Get("LOCALSTACK_AWS_SESSION"), + config.Kinesis.AccessKeyID, + config.Kinesis.SecretAccessKey, + config.Kinesis.SessionKey, )) } - client, err := common.NewKinesisClient(ctx, RealtimeMetricsStreamName, os.Getenv("KINESIS_STREAM_REGION"), opts...) // TODO: Override + client, err := common.NewKinesisClient(ctx, config.Kinesis.StreamName, config.Kinesis.Region, opts...) if err != nil { log.Infof("error creating kinesis client: %s", err) } diff --git a/internal/repository/testutils.go b/internal/repository/testutils.go index b20b8789a..e7eefa408 100644 --- a/internal/repository/testutils.go +++ b/internal/repository/testutils.go @@ -3,6 +3,7 @@ package repository import ( "github.com/alicebob/miniredis/v2" "github.com/beam-cloud/beam/internal/common" + "github.com/beam-cloud/beam/internal/types" ) func NewRedisClientForTest() (*common.RedisClient, error) { @@ -11,7 +12,7 @@ func NewRedisClientForTest() (*common.RedisClient, error) { return nil, err } - rdb, err := common.NewRedisClient(common.WithAddress(s.Addr())) + rdb, err := common.NewRedisClient(types.RedisConfig{Addrs: []string{s.Addr()}, Mode: types.RedisModeSingle}) if err != nil { return nil, err } diff --git a/internal/repository/worker_pool_redis.go b/internal/repository/worker_pool_redis.go index f85749a19..34698bc1a 100644 --- a/internal/repository/worker_pool_redis.go +++ b/internal/repository/worker_pool_redis.go @@ -23,13 +23,13 @@ func NewWorkerPoolRedisRepository(rdb *common.RedisClient) WorkerPoolRepository // Get a pool without locking. // Should only be called by public functions of this struct. -func (r *WorkerPoolRedisRepository) getPool(name string) (*types.WorkerPoolResource, error) { +func (r *WorkerPoolRedisRepository) getPool(name string) (*types.WorkerPoolConfig, error) { bytes, err := r.rdb.Get(context.TODO(), common.RedisKeys.WorkerPoolState(name)).Bytes() if err != nil { return nil, err } - p := &types.WorkerPoolResource{} + p := &types.WorkerPoolConfig{} if err := json.Unmarshal(bytes, p); err != nil { return nil, err } @@ -37,7 +37,7 @@ func (r *WorkerPoolRedisRepository) getPool(name string) (*types.WorkerPoolResou return p, nil } -func (r *WorkerPoolRedisRepository) GetPool(name string) (*types.WorkerPoolResource, error) { +func (r *WorkerPoolRedisRepository) GetPool(name string) (*types.WorkerPoolConfig, error) { lockKey := common.RedisKeys.WorkerPoolLock(name) if err := r.lock.Acquire(context.TODO(), lockKey, r.lockOptions); err != nil { return nil, err @@ -47,13 +47,13 @@ func (r *WorkerPoolRedisRepository) GetPool(name string) (*types.WorkerPoolResou return r.getPool(name) } -func (r *WorkerPoolRedisRepository) GetPools() ([]types.WorkerPoolResource, error) { +func (r *WorkerPoolRedisRepository) GetPools() ([]types.WorkerPoolConfig, error) { keys, err := r.rdb.Scan(context.TODO(), common.RedisKeys.WorkerPoolState("*")) if err != nil { return nil, err } - pools := []types.WorkerPoolResource{} + pools := []types.WorkerPoolConfig{} for _, key := range keys { name := strings.Split(key, ":")[2] @@ -68,8 +68,8 @@ func (r *WorkerPoolRedisRepository) GetPools() ([]types.WorkerPoolResource, erro return pools, nil } -func (r *WorkerPoolRedisRepository) SetPool(pool *types.WorkerPoolResource) error { - lockKey := common.RedisKeys.WorkerPoolLock(pool.Name) +func (r *WorkerPoolRedisRepository) SetPool(name string, pool *types.WorkerPoolConfig) error { + lockKey := common.RedisKeys.WorkerPoolLock(name) if err := r.lock.Acquire(context.TODO(), lockKey, r.lockOptions); err != nil { return err } @@ -80,7 +80,7 @@ func (r *WorkerPoolRedisRepository) SetPool(pool *types.WorkerPoolResource) erro return err } - return r.rdb.Set(context.TODO(), common.RedisKeys.WorkerPoolState(pool.Name), bytes, 0).Err() + return r.rdb.Set(context.TODO(), common.RedisKeys.WorkerPoolState(name), bytes, 0).Err() } func (r *WorkerPoolRedisRepository) RemovePool(name string) error { @@ -92,30 +92,3 @@ func (r *WorkerPoolRedisRepository) RemovePool(name string) error { return r.rdb.Del(context.TODO(), common.RedisKeys.WorkerPoolState(name)).Err() } - -// Used when access to a storage backend like Redis is not needed e.g. by the Agent or for testing purposes -type WorkerPoolNoOpRepository struct{} - -func NewWorkerPoolNoOpRepository() WorkerPoolRepository { - return &WorkerPoolNoOpRepository{} -} - -func (r *WorkerPoolNoOpRepository) getPool(name string) (*types.WorkerPoolResource, error) { - return &types.WorkerPoolResource{}, nil -} - -func (r *WorkerPoolNoOpRepository) GetPool(name string) (*types.WorkerPoolResource, error) { - return &types.WorkerPoolResource{}, nil -} - -func (r *WorkerPoolNoOpRepository) GetPools() ([]types.WorkerPoolResource, error) { - return []types.WorkerPoolResource{}, nil -} - -func (r *WorkerPoolNoOpRepository) SetPool(pool *types.WorkerPoolResource) error { - return nil -} - -func (r *WorkerPoolNoOpRepository) RemovePool(name string) error { - return nil -} diff --git a/internal/scheduler/backlog_test.go b/internal/scheduler/backlog_test.go index 23830fba6..36ffa584c 100644 --- a/internal/scheduler/backlog_test.go +++ b/internal/scheduler/backlog_test.go @@ -19,7 +19,7 @@ func TestRequestBacklogOrdering(t *testing.T) { assert.NotNil(t, s) assert.NoError(t, err) - redisClient, err := common.NewRedisClient(common.WithAddress(s.Addr())) + redisClient, err := common.NewRedisClient(types.RedisConfig{Addrs: []string{s.Addr()}, Mode: types.RedisModeSingle}) assert.NotNil(t, redisClient) assert.NoError(t, err) diff --git a/internal/scheduler/pool.go b/internal/scheduler/pool.go index 569915928..be8aa95c4 100644 --- a/internal/scheduler/pool.go +++ b/internal/scheduler/pool.go @@ -1,10 +1,6 @@ package scheduler import ( - "fmt" - "strconv" - - "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/types" ) @@ -38,9 +34,6 @@ type WorkerPoolConfig struct { DataVolumeName string DefaultWorkerCpuRequest int64 DefaultWorkerMemoryRequest int64 - DefaultMaxGpuCpuRequest int64 - DefaultMaxGpuMemoryRequest int64 - AgentToken string } type WorkerPoolCapacity struct { @@ -48,37 +41,3 @@ type WorkerPoolCapacity struct { FreeMemory int64 FreeGpu uint } - -func NewWorkerPoolConfig() (*WorkerPoolConfig, error) { - cfg := &WorkerPoolConfig{} - - var err error - - cfg.DataVolumeName = common.Secrets().GetWithDefault("BEAM_DATA_VOLUME_NAME", "beam-data") - - defaultWorkerCpuRequestStr := common.Secrets().GetWithDefault("DEFAULT_WORKER_CPU_REQUEST", "2000") - cfg.DefaultWorkerCpuRequest, err = strconv.ParseInt(defaultWorkerCpuRequestStr, 10, 64) - if err != nil { - return nil, fmt.Errorf("error parsing DEFAULT_WORKER_CPU_REQUEST: %v", err) - } - - defaultWorkerMemoryRequestStr := common.Secrets().GetWithDefault("DEFAULT_WORKER_MEMORY_REQUEST", "1024") - cfg.DefaultWorkerMemoryRequest, err = strconv.ParseInt(defaultWorkerMemoryRequestStr, 10, 64) - if err != nil { - return nil, fmt.Errorf("error parsing DEFAULT_WORKER_MEMORY_REQUEST: %v", err) - } - - defaultMaxGpuCpuRequestStr := common.Secrets().GetWithDefault("DEFAULT_WORKER_MAX_GPU_CPU_REQUEST", "8000") - cfg.DefaultMaxGpuCpuRequest, err = strconv.ParseInt(defaultMaxGpuCpuRequestStr, 10, 64) - if err != nil { - return nil, fmt.Errorf("error parsing DEFAULT_WORKER_MAX_GPU_CPU_REQUEST: %v", err) - } - - defaultMaxGpuMemoryRequestStr := common.Secrets().GetWithDefault("DEFAULT_WORKER_MAX_GPU_MEMORY_REQUEST", "16384") - cfg.DefaultMaxGpuMemoryRequest, err = strconv.ParseInt(defaultMaxGpuMemoryRequestStr, 10, 64) - if err != nil { - return nil, fmt.Errorf("error parsing DEFAULT_WORKER_MAX_GPU_MEMORY_REQUEST: %v", err) - } - - return cfg, nil -} diff --git a/internal/scheduler/pool_client.go b/internal/scheduler/pool_client.go deleted file mode 100644 index 85ff7a3e0..000000000 --- a/internal/scheduler/pool_client.go +++ /dev/null @@ -1,68 +0,0 @@ -package scheduler - -import ( - "context" - - "github.com/beam-cloud/beam/internal/types" - "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/runtime/schema" - "k8s.io/apimachinery/pkg/runtime/serializer" - "k8s.io/client-go/kubernetes/scheme" - "k8s.io/client-go/rest" -) - -func init() { - // Register the WorkerPool custom resource with the Kubernetes scheme - schemeBuilder := runtime.NewSchemeBuilder(func(s *runtime.Scheme) error { - // Registers WorkerPoolList as-is - s.AddKnownTypes( - schema.GroupVersion{ - Group: types.WorkerPoolSchemaGroupName, - Version: types.WorkerPoolSchemaGroupVersion, - }, - &types.WorkerPoolList{}, - ) - - // Registers WorkerPoolResource as WorkerPool - s.AddKnownTypeWithName( - schema.GroupVersionKind{ - Kind: "WorkerPool", - Group: types.WorkerPoolSchemaGroupName, - Version: types.WorkerPoolSchemaGroupVersion, - }, - &types.WorkerPoolResource{}, - ) - - return nil - }) - - schemeBuilder.AddToScheme(scheme.Scheme) -} - -type WorkerPoolClient struct { - restClient rest.Interface -} - -func NewWorkerPoolClient(cfg *rest.Config) (*WorkerPoolClient, error) { - cfg.GroupVersion = &schema.GroupVersion{Group: types.WorkerPoolSchemaGroupName, Version: types.WorkerPoolSchemaGroupVersion} - cfg.APIPath = "/apis" - cfg.NegotiatedSerializer = serializer.NewCodecFactory(scheme.Scheme) - - restClient, err := rest.RESTClientFor(cfg) - if err != nil { - return nil, err - } - - return &WorkerPoolClient{restClient: restClient}, nil -} - -func (c *WorkerPoolClient) GetWorkerPool(namespace, name string) (*types.WorkerPoolResource, error) { - result := &types.WorkerPoolResource{} - err := c.restClient.Get(). - Namespace(namespace). - Resource(types.WorkerPoolResourcePlural). - Name(name). - Do(context.Background()). - Into(result) - return result, err -} diff --git a/internal/scheduler/pool_k8s.go b/internal/scheduler/pool_k8s.go index 8b0574037..484c960a2 100644 --- a/internal/scheduler/pool_k8s.go +++ b/internal/scheduler/pool_k8s.go @@ -4,14 +4,12 @@ import ( "context" "crypto/sha256" "encoding/hex" - "errors" "fmt" "log" "strconv" "strings" "time" - "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/types" "github.com/google/uuid" @@ -24,85 +22,42 @@ import ( "k8s.io/utils/ptr" ) -func KubernetesWorkerPoolControllerFactory() WorkerPoolControllerFactory { - return func(resource *types.WorkerPoolResource, config WorkerPoolControllerConfig) (WorkerPoolController, error) { - if config, ok := config.(*KubernetesWorkerPoolControllerConfig); ok { - return NewKubernetesWorkerPoolController(resource.Name, config) - } - - return nil, errors.New("kubernetes worker pool controller factory received invalid config") - } -} - -type KubernetesWorkerPoolControllerConfig struct { - IsRemote bool - ImagePullSecrets []string - EnvVars []corev1.EnvVar - HostNetwork bool - Namespace string - WorkerPoolConfig *WorkerPoolConfig - WorkerRepo repository.WorkerRepository -} - -func NewKubernetesWorkerPoolControllerConfig(workerRepo repository.WorkerRepository) (*KubernetesWorkerPoolControllerConfig, error) { - namespace := common.Secrets().GetWithDefault("WORKER_NAMESPACE", "beam") - - workerPoolConfig, err := NewWorkerPoolConfig() - if err != nil { - return nil, err - } - - return &KubernetesWorkerPoolControllerConfig{ - Namespace: namespace, - WorkerPoolConfig: workerPoolConfig, - WorkerRepo: workerRepo, - }, nil -} - type KubernetesWorkerPoolController struct { - name string - controllerConfig *KubernetesWorkerPoolControllerConfig - provisioner string - nodeSelector map[string]string - kubeClient *kubernetes.Clientset - workerPoolClient *WorkerPoolClient - workerPoolConfig *WorkerPoolConfig - workerRepo repository.WorkerRepository + name string + config types.AppConfig + provisioner string + nodeSelector map[string]string + kubeClient *kubernetes.Clientset + workerRepo repository.WorkerRepository } -func NewKubernetesWorkerPoolController(workerPoolName string, config *KubernetesWorkerPoolControllerConfig) (WorkerPoolController, error) { - cfg, err := rest.InClusterConfig() +func NewKubernetesWorkerPoolController(config types.AppConfig, workerPoolName string, workerRepo repository.WorkerRepository) (WorkerPoolController, error) { + kubeConfig, err := rest.InClusterConfig() if err != nil { return nil, err } - workerPoolClient, err := NewWorkerPoolClient(cfg) - if err != nil { - return nil, err - } - - kubeClient, err := kubernetes.NewForConfig(cfg) + kubeClient, err := kubernetes.NewForConfig(kubeConfig) if err != nil { return nil, err } wpc := &KubernetesWorkerPoolController{ - name: workerPoolName, - controllerConfig: config, - workerPoolClient: workerPoolClient, - kubeClient: kubeClient, + name: workerPoolName, + config: config, + kubeClient: kubeClient, } - workerPool, err := wpc.getWorkerPool() + workerPool, err := wpc.GetWorkerPoolConfig() if err != nil { - return nil, err + return nil, fmt.Errorf("worker pool %s not found", workerPoolName) } wpc.name = workerPoolName - wpc.provisioner = workerPool.Spec.JobSpec.NodeSelector[defaultProvisionerLabel] - wpc.nodeSelector = workerPool.Spec.JobSpec.NodeSelector - wpc.workerPoolConfig = config.WorkerPoolConfig - wpc.workerRepo = config.WorkerRepo + wpc.config = config + wpc.provisioner = workerPool.JobSpec.NodeSelector[defaultProvisionerLabel] + wpc.nodeSelector = workerPool.JobSpec.NodeSelector + wpc.workerRepo = workerRepo // Start monitoring worker pool size err = wpc.monitorPoolSize(workerPool) @@ -115,23 +70,29 @@ func NewKubernetesWorkerPoolController(workerPoolName string, config *Kubernetes return wpc, nil } +func (wpc *KubernetesWorkerPoolController) GetWorkerPoolConfig() (*types.WorkerPoolConfig, error) { + w, ok := wpc.config.Worker.Pools[wpc.name] + if !ok { + return nil, fmt.Errorf("worker pool %v not found", wpc.name) + } + return &w, nil +} + func (wpc *KubernetesWorkerPoolController) Name() string { return wpc.name } func (wpc *KubernetesWorkerPoolController) poolId() string { - data := fmt.Sprintf("%s-%s", wpc.workerPoolConfig.AgentToken, wpc.name) - hasher := sha256.New() - hasher.Write([]byte(data)) + hasher.Write([]byte(wpc.name)) hash := hasher.Sum(nil) poolId := hex.EncodeToString(hash[:8]) return poolId } -func (wpc *KubernetesWorkerPoolController) monitorPoolSize(workerPool *types.WorkerPoolResource) error { - config, err := ParsePoolSizingConfig(workerPool.Spec.PoolSizing.Raw) +func (wpc *KubernetesWorkerPoolController) monitorPoolSize(workerPool *types.WorkerPoolConfig) error { + config, err := ParsePoolSizingConfig(workerPool.JobSpec.PoolSizing) if err != nil { return err } @@ -179,28 +140,18 @@ func (wpc *KubernetesWorkerPoolController) AddWorkerWithId(workerId string, cpu } func (wpc *KubernetesWorkerPoolController) addWorkerWithId(workerId string, cpu int64, memory int64, gpuType string) (*types.Worker, error) { - workerPool, err := wpc.getWorkerPool() - if err != nil { - return nil, err - } - // Create a new worker job - job, worker := wpc.createWorkerJob(workerId, cpu, memory, gpuType, workerPool) - - // Set the WorkerPool resouce as the owner of the new job - wpc.setJobOwner(job, workerPool) + job, worker := wpc.createWorkerJob(workerId, cpu, memory, gpuType) // Create the job in the cluster - err = wpc.createJobInCluster(job) - if err != nil { + if err := wpc.createJobInCluster(job); err != nil { return nil, err } worker.PoolId = wpc.poolId() // Add the worker state - err = wpc.workerRepo.AddWorker(worker) - if err != nil { + if err := wpc.workerRepo.AddWorker(worker); err != nil { log.Printf("Unable to create worker: %+v\n", err) return nil, err } @@ -208,17 +159,8 @@ func (wpc *KubernetesWorkerPoolController) addWorkerWithId(workerId string, cpu return worker, nil } -// Gets worker pool (Custom Resource) from Kubernetes API -func (wpc *KubernetesWorkerPoolController) getWorkerPool() (*types.WorkerPoolResource, error) { - workerPool, err := wpc.workerPoolClient.GetWorkerPool(wpc.controllerConfig.Namespace, wpc.name) - if err != nil { - return nil, err - } - return workerPool, nil -} - -func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu int64, memory int64, gpuType string, workerPool *types.WorkerPoolResource) (*batchv1.Job, *types.Worker) { - jobName := fmt.Sprintf("%s-%s-%s", BeamWorkerJobPrefix, workerPool.Name, workerId) +func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu int64, memory int64, gpuType string) (*batchv1.Job, *types.Worker) { + jobName := fmt.Sprintf("%s-%s-%s", BeamWorkerJobPrefix, wpc.name, workerId) labels := map[string]string{ BeamWorkerLabelKey: BeamWorkerLabelValue, } @@ -228,20 +170,21 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu workerGpu := gpuType resourceRequests := corev1.ResourceList{} - if cpu > 0 && cpu > wpc.workerPoolConfig.DefaultWorkerCpuRequest { + + if cpu > 0 && cpu > wpc.config.Worker.DefaultWorkerCPURequest { cpuString := fmt.Sprintf("%dm", cpu) // convert cpu to millicores string resourceRequests[corev1.ResourceCPU] = resource.MustParse(cpuString) } else { - resourceRequests[corev1.ResourceCPU] = resource.MustParse(fmt.Sprintf("%dm", wpc.workerPoolConfig.DefaultWorkerCpuRequest)) - workerCpu = wpc.workerPoolConfig.DefaultWorkerCpuRequest + resourceRequests[corev1.ResourceCPU] = resource.MustParse(fmt.Sprintf("%dm", wpc.config.Worker.DefaultWorkerCPURequest)) + workerCpu = wpc.config.Worker.DefaultWorkerCPURequest } - if memory > 0 && memory > wpc.workerPoolConfig.DefaultWorkerMemoryRequest { + if memory > 0 && memory > wpc.config.Worker.DefaultWorkerMemoryRequest { memoryString := fmt.Sprintf("%dMi", memory) // convert memory to Mi string resourceRequests[corev1.ResourceMemory] = resource.MustParse(memoryString) } else { - resourceRequests[corev1.ResourceMemory] = resource.MustParse(fmt.Sprintf("%dMi", wpc.workerPoolConfig.DefaultWorkerMemoryRequest)) - workerMemory = wpc.workerPoolConfig.DefaultWorkerMemoryRequest + resourceRequests[corev1.ResourceMemory] = resource.MustParse(fmt.Sprintf("%dMi", wpc.config.Worker.DefaultWorkerMemoryRequest)) + workerMemory = wpc.config.Worker.DefaultWorkerMemoryRequest } if gpuType != "" { @@ -249,13 +192,13 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu } workerImage := fmt.Sprintf("%s/%s:%s", - common.Secrets().Get("BEAM_WORKER_IMAGE_REGISTRY"), - common.Secrets().Get("BEAM_WORKER_IMAGE_NAME"), - common.Secrets().Get("BEAM_WORKER_IMAGE_TAG"), + wpc.config.Worker.ImageRegistry, + wpc.config.Worker.ImageName, + wpc.config.Worker.ImageTag, ) resources := corev1.ResourceRequirements{} - if common.Secrets().GetWithDefault("BEAM_WORKER_RESOURCES_ENFORCED", "true") != "false" { + if wpc.config.Worker.ResourcesEnforced { resources.Requests = resourceRequests resources.Limits = resourceRequests } @@ -278,7 +221,7 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu // Add user-defined image pull secrets imagePullSecrets := []corev1.LocalObjectReference{} - for _, s := range wpc.controllerConfig.ImagePullSecrets { + for _, s := range wpc.config.Worker.ImagePullSecrets { imagePullSecrets = append(imagePullSecrets, corev1.LocalObjectReference{Name: s}) } @@ -287,13 +230,12 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu Labels: labels, }, Spec: corev1.PodSpec{ - // TODO: change ServiceAccountName to be pulled from the CR instead of from a secret - ServiceAccountName: common.Secrets().GetWithDefault("BEAM_WORKER_SERVICE_ACCOUNT_NAME", "default"), + ServiceAccountName: wpc.config.Worker.ServiceAccountName, AutomountServiceAccountToken: ptr.To(true), - HostNetwork: wpc.controllerConfig.HostNetwork, + HostNetwork: wpc.config.Worker.HostNetwork, ImagePullSecrets: imagePullSecrets, RestartPolicy: corev1.RestartPolicyOnFailure, - NodeSelector: wpc.getWorkerNodeSelector(), + NodeSelector: map[string]string{}, Containers: containers, Volumes: wpc.getWorkerVolumes(workerMemory), EnableServiceLinks: ptr.To(false), @@ -304,7 +246,7 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu job := &batchv1.Job{ ObjectMeta: metav1.ObjectMeta{ Name: jobName, - Namespace: wpc.controllerConfig.Namespace, + Namespace: wpc.config.Worker.Namespace, Labels: labels, }, Spec: batchv1.JobSpec{ @@ -322,38 +264,12 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu } } -// Set the WorkerPool resource as the owner of the worker job -func (wpc *KubernetesWorkerPoolController) setJobOwner(job *batchv1.Job, workerPool *types.WorkerPoolResource) { - controllerFlag := true - controllerRef := metav1.OwnerReference{ - APIVersion: workerPool.APIVersion, - Kind: workerPool.Kind, - Name: workerPool.Name, - UID: workerPool.UID, - Controller: &controllerFlag, - BlockOwnerDeletion: &controllerFlag, - } - job.SetOwnerReferences(append(job.GetOwnerReferences(), controllerRef)) -} - func (wpc *KubernetesWorkerPoolController) createJobInCluster(job *batchv1.Job) error { - _, err := wpc.kubeClient.BatchV1().Jobs(wpc.controllerConfig.Namespace).Create(context.Background(), job, metav1.CreateOptions{}) + _, err := wpc.kubeClient.BatchV1().Jobs(wpc.config.Worker.Namespace).Create(context.Background(), job, metav1.CreateOptions{}) return err } -func (wpc *KubernetesWorkerPoolController) getWorkerNodeSelector() map[string]string { - stage := common.Secrets().GetWithDefault("STAGE", common.EnvLocal) - - if stage == common.EnvLocal { - return map[string]string{} - } - - return wpc.nodeSelector -} - func (wpc *KubernetesWorkerPoolController) getWorkerVolumes(workerMemory int64) []corev1.Volume { - stage := common.Secrets().GetWithDefault("STAGE", common.EnvLocal) - hostPathType := corev1.HostPathDirectoryOrCreate sharedMemoryLimit := resource.MustParse(fmt.Sprintf("%dMi", workerMemory/2)) @@ -386,48 +302,21 @@ func (wpc *KubernetesWorkerPoolController) getWorkerVolumes(workerMemory int64) }, }, { - Name: wpc.workerPoolConfig.DataVolumeName, + Name: wpc.config.Worker.DataVolumeName, VolumeSource: corev1.VolumeSource{ PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ - ClaimName: wpc.workerPoolConfig.DataVolumeName, + ClaimName: wpc.config.Worker.DataVolumeName, }, }, }, } - // Local volumes - if stage == common.EnvLocal { - return append(volumes, - corev1.Volume{ - Name: imagesVolumeName, - VolumeSource: corev1.VolumeSource{ - PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ - ClaimName: "images", - }, - }, - }, - ) - - } - - // Staging/Production volumes return append(volumes, - corev1.Volume{ - Name: configVolumeName, - VolumeSource: corev1.VolumeSource{ - Secret: &corev1.SecretVolumeSource{ - SecretName: configSecretName, - DefaultMode: func(mode int32) *int32 { return &mode }(420), - Optional: func(optional bool) *bool { return &optional }(false), - }, - }, - }, corev1.Volume{ Name: imagesVolumeName, VolumeSource: corev1.VolumeSource{ - HostPath: &corev1.HostPathVolumeSource{ - Path: "/images", - Type: &hostPathType, + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ + ClaimName: "images", }, }, }, @@ -435,9 +324,7 @@ func (wpc *KubernetesWorkerPoolController) getWorkerVolumes(workerMemory int64) } func (wpc *KubernetesWorkerPoolController) getWorkerVolumeMounts() []corev1.VolumeMount { - stage := common.Secrets().GetWithDefault("STAGE", common.EnvLocal) - - volumeMounts := []corev1.VolumeMount{ + return []corev1.VolumeMount{ { Name: tmpVolumeName, MountPath: "/tmp", @@ -458,52 +345,32 @@ func (wpc *KubernetesWorkerPoolController) getWorkerVolumeMounts() []corev1.Volu Name: "dshm", }, { - Name: wpc.workerPoolConfig.DataVolumeName, + Name: wpc.config.Worker.DataVolumeName, MountPath: "/snapshots", SubPath: "snapshots/", ReadOnly: false, }, { - Name: wpc.workerPoolConfig.DataVolumeName, + Name: wpc.config.Worker.DataVolumeName, MountPath: "/workspaces", SubPath: "workspaces/", ReadOnly: false, }, { - Name: wpc.workerPoolConfig.DataVolumeName, + Name: wpc.config.Worker.DataVolumeName, MountPath: "/volumes", SubPath: "volumes/", ReadOnly: false, }, } - - if stage == common.EnvLocal { - return volumeMounts - } - - // Staging/Production volumes - return append(volumeMounts, - corev1.VolumeMount{ - Name: configVolumeName, - MountPath: "/etc/config", - ReadOnly: true, - }, - ) } func (wpc *KubernetesWorkerPoolController) getWorkerEnvironment(workerId string, cpu int64, memory int64, gpuType string) []corev1.EnvVar { - stage := common.Secrets().GetWithDefault("STAGE", common.EnvLocal) - - // Base environment (common to both local/staging/production) - env := []corev1.EnvVar{ + return []corev1.EnvVar{ { Name: "WORKER_ID", Value: workerId, }, - { - Name: "BEAM_NAMESPACE", - Value: common.Secrets().Get("BEAM_NAMESPACE"), - }, { Name: "CPU_LIMIT", Value: strconv.FormatInt(cpu, 10), @@ -516,14 +383,6 @@ func (wpc *KubernetesWorkerPoolController) getWorkerEnvironment(workerId string, Name: "GPU_TYPE", Value: gpuType, }, - { - Name: "WORKER_S2S_TOKEN", - Value: common.Secrets().Get("WORKER_S2S_TOKEN"), - }, - { - Name: "BEAM_RUNNER_BASE_IMAGE_REGISTRY", - Value: common.Secrets().Get("BEAM_RUNNER_BASE_IMAGE_REGISTRY"), - }, { Name: "POD_IP", ValueFrom: &corev1.EnvVarSource{ @@ -534,68 +393,31 @@ func (wpc *KubernetesWorkerPoolController) getWorkerEnvironment(workerId string, }, { Name: "POD_NAMESPACE", - Value: common.Secrets().Get("WORKER_NAMESPACE"), + Value: wpc.config.Worker.Namespace, }, { Name: "CLUSTER_DOMAIN", Value: defaultClusterDomain, }, - { - Name: "BEAM_CACHE_URL", - Value: common.Secrets().Get("BEAM_CACHE_URL"), - }, { Name: "BEAM_GATEWAY_HOST", - Value: common.Secrets().Get("BEAM_GATEWAY_HOST"), + Value: wpc.config.GatewayService.Host, }, { Name: "BEAM_GATEWAY_PORT", - Value: common.Secrets().Get("BEAM_GATEWAY_PORT"), - }, - { - Name: "STATSD_HOST", - Value: common.Secrets().Get("STATSD_HOST"), - }, - { - Name: "STATSD_PORT", - Value: common.Secrets().Get("STATSD_PORT"), + Value: fmt.Sprint(wpc.config.GatewayService.Port), }, } - - // Add env var to let Worker know it should run in remote mode - if wpc.controllerConfig.IsRemote { - env = append(env, corev1.EnvVar{Name: "WORKER_IS_REMOTE", Value: "true"}) - } - - // Add user-defined env vars - env = append(env, wpc.controllerConfig.EnvVars...) - - // Local environmental variables - if stage == common.EnvLocal { - return env - } - - // Staging/Production - return append( - env, - corev1.EnvVar{ - Name: "CONFIG_PATH", - Value: "/etc/config/settings.env", - }, - corev1.EnvVar{ - Name: "KINESIS_STREAM_REGION", - Value: common.Secrets().Get("KINESIS_STREAM_REGION"), - }, - ) - } +var AddWorkerTimeout = 10 * time.Minute + // deleteStalePendingWorkerJobs ensures that jobs are deleted if they don't // start a pod after a certain amount of time. func (wpc *KubernetesWorkerPoolController) deleteStalePendingWorkerJobs() { ctx := context.Background() maxAge := AddWorkerTimeout - namespace := wpc.controllerConfig.Namespace + namespace := wpc.config.Worker.Namespace ticker := time.NewTicker(time.Minute) defer ticker.Stop() diff --git a/internal/scheduler/pool_manager.go b/internal/scheduler/pool_manager.go index dba6e1a2b..6b11cadae 100644 --- a/internal/scheduler/pool_manager.go +++ b/internal/scheduler/pool_manager.go @@ -1,42 +1,33 @@ package scheduler import ( - "context" - "fmt" - "log" "sync" repo "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/types" - "k8s.io/apimachinery/pkg/runtime/schema" - "k8s.io/apimachinery/pkg/runtime/serializer" - "k8s.io/client-go/kubernetes/scheme" - "k8s.io/client-go/rest" ) -// An interface used by the WorkerPoolControllerFactory type -type WorkerPoolControllerConfig interface{} - -// A function type that is used to initialize a controller through indirection. -// It receives a WorkerPoolResource CR, a WorkerPoolControllerConfig and returns a WorkerPoolController. -type WorkerPoolControllerFactory func(resource *types.WorkerPoolResource, config WorkerPoolControllerConfig) (WorkerPoolController, error) - -// Top level struct for WorkerPoolManager. -// It pairs together a WorkerPoolResource CR and a WorkerPoolController. +// WorkerPool represents a pool of workers with a specific name, configuration, +// and a controller responsible for managing its worker instances. type WorkerPool struct { - Resource *types.WorkerPoolResource + Name string + Config *types.WorkerPoolConfig Controller WorkerPoolController } -// Manages WorkerPools in a distributed system. -// Contains a RWMutex to allows concurrent getting/reading of WorkerPools, but lock when we write them. +// WorkerPoolManager is responsible for managing multiple worker pools. It +// maintains a collection of worker pools, and provides methods to interact with +// and manage these pools. It uses a sync.RWMutex for concurrent access control +// to ensure thread safety. type WorkerPoolManager struct { mu sync.RWMutex pools map[string]*WorkerPool repo repo.WorkerPoolRepository } -// Gets a new WorkerPoolManager. +// NewWorkerPoolManager creates a new instance of WorkerPoolManager with the +// specified worker pool repository ('repo'). It initializes an empty map of +// worker pools and associates it with the repository. func NewWorkerPoolManager(repo repo.WorkerPoolRepository) *WorkerPoolManager { return &WorkerPoolManager{ pools: map[string]*WorkerPool{}, @@ -61,20 +52,21 @@ func (m *WorkerPoolManager) GetPool(name string) (*WorkerPool, bool) { // Set/add WorkerPool. // This will overwrite any existing WorkerPools with the same name defined in WorkerPoolResource. -func (m *WorkerPoolManager) SetPool(resource types.WorkerPoolResource, controller WorkerPoolController) error { +func (m *WorkerPoolManager) SetPool(name string, config *types.WorkerPoolConfig, controller WorkerPoolController) error { m.mu.Lock() defer m.mu.Unlock() pool := &WorkerPool{ - Resource: &resource, + Name: name, + Config: config, Controller: controller, } - if err := m.repo.SetPool(pool.Resource); err != nil { + if err := m.repo.SetPool(pool.Name, pool.Config); err != nil { return err } - m.pools[resource.Name] = pool + m.pools[name] = pool return nil } @@ -88,62 +80,3 @@ func (m *WorkerPoolManager) RemovePool(name string) error { delete(m.pools, name) return m.repo.RemovePool(name) } - -// Loads WorkerPool of a certain factory type. -// This can be called any number of times depending on how many factory controller types you want to load. -func (m *WorkerPoolManager) LoadPools( - factory WorkerPoolControllerFactory, - factoryConfig WorkerPoolControllerConfig, - resources []types.WorkerPoolResource, -) error { - for _, resource := range resources { - controller, err := factory(&resource, factoryConfig) - if err != nil { - return fmt.Errorf("failed to create a controller: %v", err) - } - - err = m.SetPool(resource, controller) - if err != nil { - return fmt.Errorf("failed to set worker pool: %v", err) - } - - if pool, ok := m.pools[resource.Name]; ok { - log.Printf("Loaded WorkerPool: %s\n", pool.Resource.Name) - } else { - return fmt.Errorf("failed to load worker pool: %s", resource.Name) - } - } - - return nil -} - -// Get worker pool resources from Kubernetes -func GetWorkerPoolResources(namespace string) ([]types.WorkerPoolResource, error) { - cfg, err := rest.InClusterConfig() - if err != nil { - return nil, err - } - - // Create a clientset for the worker pool CRD - cfg.GroupVersion = &schema.GroupVersion{Group: types.WorkerPoolSchemaGroupName, Version: types.WorkerPoolSchemaGroupVersion} - cfg.APIPath = "/apis" - cfg.NegotiatedSerializer = serializer.NewCodecFactory(scheme.Scheme) - kubeClient, err := rest.RESTClientFor(cfg) - if err != nil { - return nil, err - } - - // Get all worker pools - result := kubeClient.Get().Namespace(namespace).Resource(types.WorkerPoolResourcePlural).Do(context.Background()) - if result.Error() != nil { - return nil, fmt.Errorf("error getting worker pool resources from kubernetes: %v", result.Error()) - } - - workerPools := types.WorkerPoolList{} - err = result.Into(&workerPools) - if err != nil { - return nil, fmt.Errorf("error deserializing worker pool list: %v", err) - } - - return workerPools.Items, nil -} diff --git a/internal/scheduler/pool_remote.go b/internal/scheduler/pool_remote.go deleted file mode 100644 index 80d757345..000000000 --- a/internal/scheduler/pool_remote.go +++ /dev/null @@ -1,109 +0,0 @@ -package scheduler - -import ( - "context" - "fmt" - "log" - "time" - - "github.com/beam-cloud/beam/internal/repository" - "github.com/beam-cloud/beam/internal/types" - pb "github.com/beam-cloud/beam/proto" - "github.com/google/uuid" -) - -// Determines how long to wait for a new Worker to start up. -// This only affects remote workers. -const AddWorkerTimeout = 10 * time.Minute - -type RemoteWorkerPoolControllerConfig struct { - agent *types.Agent - workerEventStream pb.Scheduler_SubscribeWorkerEventsServer - workerRepo repository.WorkerRepository -} - -type RemoteWorkerPoolController struct { - name string - config *RemoteWorkerPoolControllerConfig -} - -func NewRemoteWorkerPoolController(workerPoolName string, config *RemoteWorkerPoolControllerConfig) WorkerPoolController { - return &RemoteWorkerPoolController{ - name: workerPoolName, - config: config, - } -} - -func (c *RemoteWorkerPoolController) Name() string { - return c.name -} - -func (c *RemoteWorkerPoolController) AddWorker(cpu int64, memory int64, gpuType string) (*types.Worker, error) { - workerId := c.generateWorkerId() - return c.addWorkerWithId(workerId, cpu, memory, gpuType) -} - -func (c *RemoteWorkerPoolController) AddWorkerWithId(workerId string, cpu int64, memory int64, gpuType string) (*types.Worker, error) { - return c.addWorkerWithId(workerId, cpu, memory, gpuType) -} - -func (c *RemoteWorkerPoolController) addWorkerWithId(workerId string, cpu int64, memory int64, gpuType string) (*types.Worker, error) { - worker := &types.Worker{ - Id: workerId, - Status: types.WorkerStatusPending, - Cpu: cpu, - Memory: memory, - Gpu: gpuType, - } - - if err := c.config.workerRepo.AddWorker(worker); err != nil { - return nil, fmt.Errorf("failed to add worker to repo: %v", err) - } - - go func() { - ctx, cancel := context.WithTimeout(context.Background(), AddWorkerTimeout) - defer cancel() - - workerEvent := &pb.SubscribeWorkerEventResponse{ - Worker: &pb.Worker{ - Id: worker.Id, - Cpu: worker.Cpu, - Memory: worker.Memory, - GpuType: worker.Gpu, - }, - WorkerPool: &pb.WorkerPool{Name: c.name}, - } - - if err := c.config.workerEventStream.Send(workerEvent); err != nil { - log.Printf("Failed sending worker event to agent: agent <%s>, worker <%s>, worker pool <%s>: %v\n", c.config.agent.ExternalID, worker.Id, c.name, err) - } - - ticker := time.NewTicker(3 * time.Second) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - log.Printf("Timed out adding worker, removing worker state: agent <%s>, worker <%s>, worker pool <%s>\n", c.config.agent.ExternalID, worker.Id, c.name) - c.config.workerRepo.RemoveWorker(worker) - return - case <-ticker.C: - w, err := c.config.workerRepo.GetWorkerById(worker.Id) - if err == nil && w.Status == types.WorkerStatusAvailable { - log.Printf("Worker added by agent: agent <%s>, worker <%s>, worker pool <%s>\n", c.config.agent.ExternalID, worker.Id, c.name) - return - } - } - } - }() - - return worker, nil -} - -func (c *RemoteWorkerPoolController) FreeCapacity() (*WorkerPoolCapacity, error) { - return &WorkerPoolCapacity{}, nil -} - -func (c *RemoteWorkerPoolController) generateWorkerId() string { - return uuid.New().String()[:8] -} diff --git a/internal/scheduler/pool_sizing.go b/internal/scheduler/pool_sizing.go index 097de6877..f86b3d35c 100644 --- a/internal/scheduler/pool_sizing.go +++ b/internal/scheduler/pool_sizing.go @@ -1,7 +1,6 @@ package scheduler import ( - "encoding/json" "log" "time" @@ -55,51 +54,33 @@ func (s *WorkerPoolSizer) addWorkerIfNeeded(freeCapacity *WorkerPoolCapacity) (* return newWorker, nil } -// Parses a JSON string represented as []byte. +// ParsePoolSizingConfig converts a common.WorkerPoolJobSpecPoolSizingConfig to a types.WorkerPoolSizingConfig. // When a value is not parsable or is invalid, we ignore the error and set a default. -func ParsePoolSizingConfig(raw []byte) (*types.WorkerPoolSizingConfig, error) { - var configMap map[string]interface{} - err := json.Unmarshal(raw, &configMap) - if err != nil { - return nil, err - } - +func ParsePoolSizingConfig(config types.WorkerPoolJobSpecPoolSizingConfig) (*types.WorkerPoolSizingConfig, error) { c := types.NewWorkerPoolSizingConfig() - if val, ok := configMap["minFreeCpu"]; ok { - if minFreeCpu, err := ParseCPU(val); err == nil { - c.MinFreeCpu = minFreeCpu - } + if minFreeCpu, err := ParseCPU(config.MinFreeCPU); err == nil { + c.MinFreeCpu = minFreeCpu } - if val, ok := configMap["minFreeMemory"]; ok { - if minFreeMemory, err := ParseMemory(val); err == nil { - c.MinFreeMemory = minFreeMemory - } + if minFreeMemory, err := ParseMemory(config.MinFreeMemory); err == nil { + c.MinFreeMemory = minFreeMemory } - if val, ok := configMap["minFreeGpu"]; ok { - if minFreeGpu, err := ParseGPU(val); err == nil { - c.MinFreeGpu = minFreeGpu - } + if minFreeGpu, err := ParseGPU(config.MinFreeGPU); err == nil { + c.MinFreeGpu = minFreeGpu } - if val, ok := configMap["defaultWorkerCpu"]; ok { - if defaultWorkerCpu, err := ParseCPU(val); err == nil { - c.DefaultWorkerCpu = defaultWorkerCpu - } + if defaultWorkerCpu, err := ParseCPU(config.DefaultWorkerCPU); err == nil { + c.DefaultWorkerCpu = defaultWorkerCpu } - if val, ok := configMap["defaultWorkerMemory"]; ok { - if defaultWorkerMemory, err := ParseMemory(val); err == nil { - c.DefaultWorkerMemory = defaultWorkerMemory - } + if defaultWorkerMemory, err := ParseMemory(config.DefaultWorkerMemory); err == nil { + c.DefaultWorkerMemory = defaultWorkerMemory } - if val, ok := configMap["defaultWorkerGpuType"]; ok { - if defaultWorkerGpuType, err := ParseGPUType(val); err == nil { - c.DefaultWorkerGpuType = defaultWorkerGpuType.String() - } + if defaultWorkerGpuType, err := ParseGPUType(config.DefaultWorkerGPUType); err == nil { + c.DefaultWorkerGpuType = defaultWorkerGpuType.String() } return c, nil diff --git a/internal/scheduler/pool_sizing_test.go b/internal/scheduler/pool_sizing_test.go index d9a9b42e2..679a2d51f 100644 --- a/internal/scheduler/pool_sizing_test.go +++ b/internal/scheduler/pool_sizing_test.go @@ -15,7 +15,7 @@ func TestAddWorkerIfNeeded(t *testing.T) { assert.NotNil(t, s) assert.Nil(t, err) - redisClient, err := common.NewRedisClient(common.WithAddress(s.Addr())) + redisClient, err := common.NewRedisClient(types.RedisConfig{Addrs: []string{s.Addr()}, Mode: types.RedisModeSingle}) assert.NotNil(t, redisClient) assert.Nil(t, err) @@ -79,11 +79,10 @@ func TestAddWorkerIfNeeded(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - controller, _ := NewWorkerPoolControllerForTest("TestPool", &WorkerPoolControllerConfigForTest{ - namespace: "namespace", - workerPoolConfig: &WorkerPoolConfig{}, - workerRepo: workerRepo, - }) + controller := &WorkerPoolControllerForTest{ + name: "TestPool", + workerRepo: workerRepo, + } sizer := &WorkerPoolSizer{ controller: controller, config: tt.config, @@ -103,14 +102,14 @@ func TestAddWorkerIfNeeded(t *testing.T) { func TestParsePoolSizingConfig(t *testing.T) { tests := []struct { - name string - sizingConfigJSON string - sizingConfigExpected *types.WorkerPoolSizingConfig + name string + sizingConfigHave *types.WorkerPoolJobSpecPoolSizingConfig + sizingConfigWant *types.WorkerPoolSizingConfig }{ { name: "should set defaults when no values provided", - sizingConfigJSON: "{}", - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{}, + sizingConfigWant: &types.WorkerPoolSizingConfig{ MinFreeCpu: 0, MinFreeMemory: 0, MinFreeGpu: 0, @@ -121,17 +120,15 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should ignore bad values and use default values", - sizingConfigJSON: ` - { - "minFreeCpu": "bad value", - "minFreeMemory": "bad value", - "minFreeGpu": "bad value", - "defaultWorkerCpu": -10, - "defaultWorkerMemory": -10, - "defaultWorkerGpuType": "bad value" - } - `, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{ + MinFreeCPU: "bad value", + MinFreeMemory: "bad value", + MinFreeGPU: "bad value", + DefaultWorkerCPU: "bad value", + DefaultWorkerMemory: "bad value", + DefaultWorkerGPUType: "bad value", + }, + sizingConfigWant: &types.WorkerPoolSizingConfig{ MinFreeCpu: 0, MinFreeMemory: 0, MinFreeGpu: 0, @@ -142,14 +139,12 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should parse minFreeCpu minFreeMemory and minFreeGpu", - sizingConfigJSON: ` - { - "minFreeCpu": "132000m", - "minFreeMemory": "100Gi", - "minFreeGpu": 0 - } - `, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{ + MinFreeCPU: "132000m", + MinFreeMemory: "100Gi", + MinFreeGPU: "0", + }, + sizingConfigWant: &types.WorkerPoolSizingConfig{ MinFreeCpu: 132000, MinFreeMemory: 102400, MinFreeGpu: 0, @@ -159,8 +154,8 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should parse defaultWorkerGpuType as T4", - sizingConfigJSON: `{"defaultWorkerGpuType": "T4"}`, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{DefaultWorkerGPUType: "T4"}, + sizingConfigWant: &types.WorkerPoolSizingConfig{ DefaultWorkerCpu: 1000, DefaultWorkerMemory: 1024, DefaultWorkerGpuType: "T4", @@ -168,8 +163,8 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should parse defaultWorkerGpuType as A100-40", - sizingConfigJSON: `{"defaultWorkerGpuType": "A100-40"}`, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{DefaultWorkerGPUType: "A100-40"}, + sizingConfigWant: &types.WorkerPoolSizingConfig{ DefaultWorkerCpu: 1000, DefaultWorkerMemory: 1024, DefaultWorkerGpuType: "A100-40", @@ -177,8 +172,8 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should return empty DefaultWorkerGpuType when using 3060 (str) as defaultWorkerGpuType", - sizingConfigJSON: `{"defaultWorkerGpuType": "3060"}`, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{DefaultWorkerGPUType: "3060"}, + sizingConfigWant: &types.WorkerPoolSizingConfig{ DefaultWorkerCpu: 1000, DefaultWorkerMemory: 1024, DefaultWorkerGpuType: "", @@ -186,8 +181,8 @@ func TestParsePoolSizingConfig(t *testing.T) { }, { name: "should set empty DefaultWorkerGpuType when using 4090 (int) as defaultWorkerGpuType", - sizingConfigJSON: `{"defaultWorkerGpuType": 4090}`, - sizingConfigExpected: &types.WorkerPoolSizingConfig{ + sizingConfigHave: &types.WorkerPoolJobSpecPoolSizingConfig{DefaultWorkerGPUType: "4090"}, + sizingConfigWant: &types.WorkerPoolSizingConfig{ DefaultWorkerCpu: 1000, DefaultWorkerMemory: 1024, DefaultWorkerGpuType: "", @@ -197,11 +192,10 @@ func TestParsePoolSizingConfig(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - raw := []byte(tt.sizingConfigJSON) - config, err := ParsePoolSizingConfig(raw) + config, err := ParsePoolSizingConfig(*tt.sizingConfigHave) assert.NoError(t, err) - assert.Equal(t, tt.sizingConfigExpected, config) + assert.Equal(t, tt.sizingConfigWant, config) }) } } diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index c4d7e81db..a6c0e6adb 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -11,7 +11,6 @@ import ( "github.com/beam-cloud/beam/internal/common" repo "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/types" - "github.com/samber/lo" ) const ( @@ -26,36 +25,19 @@ type Scheduler struct { beamRepo repo.BeamRepository metricsRepo repo.MetricsStatsdRepository eventBus *common.EventBus - redisClient *common.RedisClient } -func NewScheduler() (*Scheduler, error) { - redisClient, err := common.NewRedisClient(common.WithClientName("BeamScheduler")) - if err != nil { - return nil, err - } - +func NewScheduler(config types.AppConfig, redisClient *common.RedisClient) (*Scheduler, error) { eventBus := common.NewEventBus(redisClient) workerRepo := repo.NewWorkerRedisRepository(redisClient) workerPoolRepo := repo.NewWorkerPoolRedisRepository(redisClient) requestBacklog := NewRequestBacklog(redisClient) containerRepo := repo.NewContainerRedisRepository(redisClient) - controllerFactory := KubernetesWorkerPoolControllerFactory() - controllerFactoryConfig, err := NewKubernetesWorkerPoolControllerConfig(workerRepo) - if err != nil { - return nil, err - } - workerPoolManager := NewWorkerPoolManager(workerPoolRepo) - workerPoolResources, err := GetWorkerPoolResources(controllerFactoryConfig.Namespace) - if err != nil { - return nil, err - } - - err = workerPoolManager.LoadPools(controllerFactory, controllerFactoryConfig, workerPoolResources) - if err != nil { - return nil, err + for name, pool := range config.Worker.Pools { + controller, _ := NewKubernetesWorkerPoolController(config, name, workerRepo) + workerPoolManager.SetPool(name, &pool, controller) } return &Scheduler{ @@ -65,7 +47,6 @@ func NewScheduler() (*Scheduler, error) { requestBacklog: requestBacklog, containerRepo: containerRepo, metricsRepo: repo.NewMetricsStatsdRepository(), - redisClient: redisClient, }, nil } @@ -121,23 +102,7 @@ func (s *Scheduler) Stop(containerId string) error { return nil } -// Get a controller. -// When an agent is provided by the user, we attempt to find a worker pool controller associated with that -// agent. If we don't find a controller, we default to returning a Beam hosted controller. func (s *Scheduler) getController(request *types.ContainerRequest) (WorkerPoolController, error) { - if request.Agent != "" { - controller, err := s.getRemoteController(request) - if err != nil { - log.Printf("unable to find remote controllers for user-specified agent <%v>: %v", request.Agent, err) - } else { - return controller, nil - } - } - - return s.getHostedController(request) -} - -func (s *Scheduler) getHostedController(request *types.ContainerRequest) (WorkerPoolController, error) { poolName := "beam-cpu" if request.Gpu != "" { @@ -159,37 +124,6 @@ func (s *Scheduler) getHostedController(request *types.ContainerRequest) (Worker return workerPool.Controller, nil } -func (s *Scheduler) getRemoteController(request *types.ContainerRequest) (WorkerPoolController, error) { - agentName := request.Agent - - agent, err := s.beamRepo.GetAgent(agentName, "") - if err != nil { - return nil, fmt.Errorf("failed to get agent <%s> from database: %v", agentName, err) - } - - if !agent.IsOnline { - return nil, errors.New("unable to use worker pools because agent is not online") - } - - pools, err := agent.GetPools() - if err != nil { - return nil, fmt.Errorf("failed to parse pools on agent <%s>: %v", agentName, err) - } - - if len(pools) == 0 { - return nil, fmt.Errorf("unable to find worker pools because agent <%s> does not have any pools", agentName) - } - - for poolName := range pools { - workerPool, ok := s.workerPoolManager.GetPool(poolName) - if ok { - return workerPool.Controller, nil - } - } - - return nil, fmt.Errorf("unable to find controllers for agent <%s>", agentName) -} - func (s *Scheduler) processRequests() { for { if s.requestBacklog.Len() == 0 { @@ -252,13 +186,6 @@ func (s *Scheduler) selectWorker(request *types.ContainerRequest) (*types.Worker return nil, err } - // When agent is present, filter workers with agents - if request.Agent != "" { - workers = lo.Filter(workers, func(w *types.Worker, _ int) bool { - return w.Agent == request.Agent - }) - } - // Sort workers: available first, then pending sort.Slice(workers, func(i, j int) bool { return workers[i].Status < workers[j].Status diff --git a/internal/scheduler/scheduler.proto b/internal/scheduler/scheduler.proto index 05e85eafa..08b84fe6c 100644 --- a/internal/scheduler/scheduler.proto +++ b/internal/scheduler/scheduler.proto @@ -8,10 +8,6 @@ service Scheduler { rpc GetVersion(VersionRequest) returns (VersionResponse) {} rpc RunContainer(RunContainerRequest) returns (RunContainerResponse) {} rpc StopContainer(StopContainerRequest) returns (StopContainerResponse) {} - // rpc StreamLogs(stream StreamLogRequest) returns (StreamLogResponse); - rpc SubscribeWorkerEvents(SubscribeWorkerEventRequest) - returns (stream SubscribeWorkerEventResponse); - rpc RegisterWorker(RegisterWorkerRequest) returns (RegisterWorkerResponse); } message VersionRequest {} @@ -39,48 +35,3 @@ message StopContainerResponse { bool success = 3; string error = 4; } - -message AgentInfo { - string token = 1; - string version = 2; - string cloud_provider = 3; -} - -message WorkerPool { string name = 1; } - -message Worker { - string id = 1; - int64 cpu = 2; - int64 memory = 3; - string gpu_type = 4; - string status = 5; - string pool_id = 6; -} - -message SubscribeWorkerEventRequest { - AgentInfo agent_info = 1; - repeated WorkerPool worker_pools = 2; -} - -message SubscribeWorkerEventResponse { - Worker worker = 1; - WorkerPool workerPool = 2; -} - -message StatusRequest {} - -message StatusResponse { string status = 1; } - -message StreamLogRequest { - string id = 1; - string log = 2; -} - -message StreamLogResponse { bool ok = 1; } - -message RegisterWorkerRequest { - AgentInfo agent_info = 1; - Worker worker = 2; -} - -message RegisterWorkerResponse {} diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 60859be23..1128ed331 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -2,7 +2,6 @@ package scheduler import ( "context" - "errors" "log" "testing" "time" @@ -12,48 +11,43 @@ import ( repo "github.com/beam-cloud/beam/internal/repository" "github.com/beam-cloud/beam/internal/types" "github.com/google/uuid" + "github.com/knadh/koanf/providers/rawbytes" "github.com/tj/assert" ) -func NewschedulerForTest() (*Scheduler, error) { +func NewSchedulerForTest() (*Scheduler, error) { s, err := miniredis.Run() if err != nil { return nil, err } - rdb, err := common.NewRedisClient(common.WithAddress(s.Addr())) + rdb, err := common.NewRedisClient(types.RedisConfig{Addrs: []string{s.Addr()}, Mode: types.RedisModeSingle}) if err != nil { return nil, err } eventBus := common.NewEventBus(rdb) workerRepo := repo.NewWorkerRedisRepositoryForTest(rdb) - workerPoolRepo := repo.NewWorkerPoolRedisRepository(rdb) containerRepo := repo.NewContainerRedisRepositoryForTest(rdb) requestBacklog := NewRequestBacklogForTest(rdb) - workerPoolConfig := &WorkerPoolConfig{ - DataVolumeName: "beam-data-volume-name", - DefaultWorkerCpuRequest: 1000, - DefaultWorkerMemoryRequest: 1000, - DefaultMaxGpuCpuRequest: 1000, - DefaultMaxGpuMemoryRequest: 1000, - } - - factory := WorkerPoolControllerFactoryForTest() - factoryConfig := &WorkerPoolControllerConfigForTest{ - namespace: "beam-namespace", - workerPoolConfig: workerPoolConfig, - workerRepo: workerRepo, + configManager, err := common.NewConfigManager[types.AppConfig]() + if err != nil { + return nil, err } - workerPoolManager := NewWorkerPoolManager(workerPoolRepo) - workerPoolResources := []types.WorkerPoolResource{ - types.NewWorkerPoolResource("beam-cpu"), - types.NewWorkerPoolResource("beam-a10g"), - types.NewWorkerPoolResource("beam-t4"), + poolJson := []byte(`{"worker":{"pools":{"beam-cpu":{},"beam-a10g":{},"beam-t4":{}}}}}`) + configManager.LoadConfig(common.YAMLConfigFormat, rawbytes.Provider(poolJson)) + config := configManager.GetConfig() + + workerPoolManager := NewWorkerPoolManager(repo.NewWorkerPoolRedisRepository(rdb)) + for name, pool := range config.Worker.Pools { + workerPoolManager.SetPool(name, &pool, &WorkerPoolControllerForTest{ + name: name, + config: config, + workerRepo: workerRepo, + }) } - workerPoolManager.LoadPools(factory, factoryConfig, workerPoolResources) return &Scheduler{ eventBus: eventBus, @@ -65,32 +59,10 @@ func NewschedulerForTest() (*Scheduler, error) { }, nil } -func WorkerPoolControllerFactoryForTest() WorkerPoolControllerFactory { - return func(resource *types.WorkerPoolResource, config WorkerPoolControllerConfig) (WorkerPoolController, error) { - if c, ok := config.(*WorkerPoolControllerConfigForTest); ok { - return NewWorkerPoolControllerForTest(resource.Name, c) - } - - return nil, errors.New("test worker pool controller factory received invalid config") - } -} - -type WorkerPoolControllerConfigForTest struct { - namespace string - workerPoolConfig *WorkerPoolConfig - workerRepo repo.WorkerRepository -} - type WorkerPoolControllerForTest struct { - name string - config *WorkerPoolControllerConfigForTest -} - -func NewWorkerPoolControllerForTest(workerPoolName string, config *WorkerPoolControllerConfigForTest) (WorkerPoolController, error) { - return &WorkerPoolControllerForTest{ - name: workerPoolName, - config: config, - }, nil + name string + config types.AppConfig + workerRepo repo.WorkerRepository } func (wpc *WorkerPoolControllerForTest) generateWorkerId() string { @@ -113,7 +85,7 @@ func (wpc *WorkerPoolControllerForTest) AddWorker(cpu int64, memory int64, gpuTy } // Add the worker state - err := wpc.config.workerRepo.AddWorker(worker) + err := wpc.workerRepo.AddWorker(worker) if err != nil { log.Printf("Unable to create worker: %+v\n", err) return nil, err @@ -131,13 +103,13 @@ func (wpc *WorkerPoolControllerForTest) FreeCapacity() (*WorkerPoolCapacity, err } func TestNewschedulerForTest(t *testing.T) { - wb, err := NewschedulerForTest() + wb, err := NewSchedulerForTest() assert.Nil(t, err) assert.NotNil(t, wb) } func TestRunContainer(t *testing.T) { - wb, err := NewschedulerForTest() + wb, err := NewSchedulerForTest() assert.Nil(t, err) assert.NotNil(t, wb) @@ -161,7 +133,7 @@ func TestRunContainer(t *testing.T) { } func TestProcessRequests(t *testing.T) { - wb, err := NewschedulerForTest() + wb, err := NewSchedulerForTest() assert.Nil(t, err) assert.NotNil(t, wb) @@ -217,7 +189,7 @@ func TestProcessRequests(t *testing.T) { func TestGetController(t *testing.T) { t.Run("returns correct controller", func(t *testing.T) { - wb, _ := NewschedulerForTest() + wb, _ := NewSchedulerForTest() cpuRequest := &types.ContainerRequest{Gpu: ""} cpuController, err := wb.getController(cpuRequest) @@ -239,7 +211,7 @@ func TestGetController(t *testing.T) { }) t.Run("returns error if no suitable controller found", func(t *testing.T) { - wb, _ := NewschedulerForTest() + wb, _ := NewSchedulerForTest() unknownRequest := &types.ContainerRequest{Gpu: "UNKNOWN_GPU"} _, err := wb.getController(unknownRequest) @@ -250,7 +222,7 @@ func TestGetController(t *testing.T) { } func TestSelectGPUWorker(t *testing.T) { - wb, err := NewschedulerForTest() + wb, err := NewSchedulerForTest() assert.Nil(t, err) assert.NotNil(t, wb) @@ -298,7 +270,7 @@ func TestSelectGPUWorker(t *testing.T) { } func TestSelectCPUWorker(t *testing.T) { - wb, err := NewschedulerForTest() + wb, err := NewSchedulerForTest() assert.Nil(t, err) assert.NotNil(t, wb) diff --git a/internal/scheduler/service.go b/internal/scheduler/service.go index c587b4be9..b6e9e2e09 100644 --- a/internal/scheduler/service.go +++ b/internal/scheduler/service.go @@ -2,13 +2,10 @@ package scheduler import ( "context" - "errors" - "log" + "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/types" pb "github.com/beam-cloud/beam/proto" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" ) type SchedulerService struct { @@ -16,16 +13,16 @@ type SchedulerService struct { Scheduler *Scheduler } -func NewSchedulerService() (*SchedulerService, error) { - Scheduler, err := NewScheduler() +func NewSchedulerService(config types.AppConfig, redisClient *common.RedisClient) (*SchedulerService, error) { + scheduler, err := NewScheduler(config, redisClient) if err != nil { return nil, err } - go Scheduler.processRequests() // Start processing ContainerRequests + go scheduler.processRequests() // Start processing ContainerRequests return &SchedulerService{ - Scheduler: Scheduler, + Scheduler: scheduler, }, nil } @@ -92,92 +89,3 @@ func (wbs *SchedulerService) StopContainer(ctx context.Context, in *pb.StopConta Error: "", }, nil } - -// Sets up worker pools reported by the agent/client. -func (wbs *SchedulerService) SubscribeWorkerEvents( - req *pb.SubscribeWorkerEventRequest, - stream pb.Scheduler_SubscribeWorkerEventsServer, -) error { - if req.AgentInfo == nil { - return status.Error(codes.InvalidArgument, "invalid agent token") - } - - // Get and update agent - agent, err := wbs.Scheduler.beamRepo.GetAgentByToken(req.AgentInfo.Token) - if err != nil { - return status.Error(codes.Internal, "invalid agent token") - } - - agent.IsOnline = true - agent.CloudProvider = req.AgentInfo.CloudProvider - agent.Version = req.AgentInfo.Version - workerPools := make(map[string]string, len(req.WorkerPools)) - for _, pool := range req.WorkerPools { - workerPools[pool.Name] = "" - } - agent.SetPools(workerPools) - - agent, err = wbs.Scheduler.beamRepo.UpdateAgent(agent) - if err != nil { - log.Printf("Unable to update agent in database: %v\n", err) - } - - // Register worker pools - for _, pool := range req.WorkerPools { - // Find an existing pool. If found, don't register/overwrite it and prevent it from being removed - // when the client disconects later on. - if _, ok := wbs.Scheduler.workerPoolManager.GetPool(pool.Name); ok { - delete(workerPools, pool.Name) - continue - } - - // No pool registered, lets add one - controller := NewRemoteWorkerPoolController(pool.Name, &RemoteWorkerPoolControllerConfig{ - agent: agent, - workerEventStream: stream, - workerRepo: wbs.Scheduler.workerRepo, - }) - - // Create a WorkerPool and register it with the WorkerPoolManager. - // WorkerPoolManager only needs the pool name to find a controller. So we'll create a - // WorkerPoolResource with just the name and register it with the WorkerPoolManager. - resource := types.NewWorkerPoolResource(pool.Name) - wbs.Scheduler.workerPoolManager.SetPool(resource, controller) - } - - log.Printf("Agent <%v> has connected with pools %+v\n", agent.ExternalID, workerPools) - <-stream.Context().Done() - log.Printf("Agent <%v> has disconnected\n", agent.ExternalID) - - // Clean up - for poolName := range workerPools { - wbs.Scheduler.workerPoolManager.RemovePool(poolName) - } - - agent.IsOnline = false - if _, err = wbs.Scheduler.beamRepo.UpdateAgent(agent); err != nil { - log.Printf("Unable to update agent info on client disconnect: %v\n", err) - } - - return status.Error(codes.Canceled, "client disconnected") -} - -func (wbs *SchedulerService) RegisterWorker(ctx context.Context, in *pb.RegisterWorkerRequest) (*pb.RegisterWorkerResponse, error) { - if in.AgentInfo == nil { - return nil, errors.New("invalid agent token") - } - - _, err := wbs.Scheduler.beamRepo.GetAgentByToken(in.AgentInfo.Token) - if err != nil { - return nil, errors.New("invalid agent token") - } - - return &pb.RegisterWorkerResponse{}, wbs.Scheduler.workerRepo.AddWorker(&types.Worker{ - Id: in.Worker.Id, - Cpu: in.Worker.Cpu, - Memory: in.Worker.Memory, - Gpu: in.Worker.GpuType, - PoolId: in.Worker.PoolId, - Status: types.WorkerStatus(in.Worker.Status), - }) -} diff --git a/internal/storage/juicefs.go b/internal/storage/juicefs.go index 41e93f763..68eadb622 100644 --- a/internal/storage/juicefs.go +++ b/internal/storage/juicefs.go @@ -5,22 +5,25 @@ import ( "log" "os/exec" - "github.com/beam-cloud/beam/internal/common" + "github.com/beam-cloud/beam/internal/types" ) type JuiceFsStorage struct { mountCmd *exec.Cmd + config types.JuiceFSConfig } -func NewJuiceFsStorage() (Storage, error) { - return &JuiceFsStorage{}, nil +func NewJuiceFsStorage(config types.JuiceFSConfig) (Storage, error) { + return &JuiceFsStorage{ + config: config, + }, nil } func (s *JuiceFsStorage) Mount(localPath string) error { s.mountCmd = exec.Command( "juicefs", "mount", - common.Secrets().Get("BEAM_JUICEFS_REDIS"), + s.config.RedisURI, localPath, ) @@ -42,10 +45,10 @@ func (s *JuiceFsStorage) Format(fsName string) error { "juicefs", "format", "--storage", "s3", - "--bucket", common.Secrets().Get("BEAM_JUICEFS_S3_BUCKET"), - "--access-key", common.Secrets().Get("BEAM_JUICEFS_AWS_ACCESS_KEY"), - "--secret-key", common.Secrets().Get("BEAM_JUICEFS_AWS_SECRET_KEY"), - common.Secrets().Get("BEAM_JUICEFS_REDIS"), + "--bucket", s.config.AWSS3Bucket, + "--access-key", s.config.AWSAccessKeyID, + "--secret-key", s.config.AWSSecretAccessKey, + s.config.RedisURI, fsName, ) diff --git a/internal/storage/storage.go b/internal/storage/storage.go index 7dcb43015..7eab37090 100644 --- a/internal/storage/storage.go +++ b/internal/storage/storage.go @@ -4,18 +4,11 @@ import ( "errors" "log" - "github.com/beam-cloud/beam/internal/common" "github.com/beam-cloud/beam/internal/types" ) const ( - DefaultFilesystemName string = types.DefaultFilesystemName - DefaultFilesystemPath string = types.DefaultFilesystemPath - DefaultObjectPath string = types.DefaultObjectPath -) - -const ( - StorageModeJuiceFs string = "JUICEFS" + StorageModeJuiceFS string = "juicefs" ) type Storage interface { @@ -24,25 +17,24 @@ type Storage interface { Unmount(localPath string) error } -func NewStorage() (Storage, error) { - storageMode := common.Secrets().GetWithDefault("BEAM_STORAGE_MODE", StorageModeJuiceFs) +func NewStorage(config types.StorageConfig) (Storage, error) { - switch storageMode { - case StorageModeJuiceFs: - s, err := NewJuiceFsStorage() + switch config.Mode { + case StorageModeJuiceFS: + s, err := NewJuiceFsStorage(config.JuiceFS) if err != nil { return nil, err } // Format filesystem // NOTE: this is a no-op if already formatted - err = s.Format(DefaultFilesystemName) + err = s.Format(config.FilesystemName) if err != nil { log.Fatalf("Unable to format filesystem: %+v\n", err) } // Mount filesystem - err = s.Mount(DefaultFilesystemPath) + err = s.Mount(config.FilesystemPath) if err != nil { log.Fatalf("Unable to mount filesystem: %+v\n", err) } diff --git a/internal/types/config.go b/internal/types/config.go new file mode 100644 index 000000000..082d41bac --- /dev/null +++ b/internal/types/config.go @@ -0,0 +1,155 @@ +package types + +import "time" + +type AppConfig struct { + DebugMode bool `key:"debugMode"` + Database DatabaseConfig `key:"database"` + GatewayService GatewayServiceConfig `key:"gateway"` + ImageService ImageServiceConfig `key:"imageservice"` + Metrics MetricsConfig `key:"metrics"` + Storage StorageConfig `key:"storage"` + Worker WorkerConfig `key:"worker"` +} + +type DatabaseConfig struct { + Redis RedisConfig `key:"redis"` + Postgres PostgresConfig `key:"postgres"` +} + +type RedisMode string + +var ( + RedisModeSingle RedisMode = "single" + RedisModeCluster RedisMode = "cluster" +) + +type RedisConfig struct { + Addrs []string `key:"addrs"` + Mode RedisMode `key:"mode"` + ClientName string `key:"clientName"` + EnableTLS bool `key:"enableTLS"` + MinIdleConns int `key:"minIdleConns"` + MaxIdleConns int `key:"maxIdleConns"` + ConnMaxIdleTime time.Duration `key:"connMaxIdleTime"` + ConnMaxLifetime time.Duration `key:"connMaxLifetime"` + DialTimeout time.Duration `key:"dialTimeout"` + ReadTimeout time.Duration `key:"readTimeout"` + WriteTimeout time.Duration `key:"writeTimeout"` + MaxRedirects int `key:"maxRedirects"` + MaxRetries int `key:"maxRetries"` + PoolSize int `key:"poolSize"` + Username string `key:"username"` + Password string `key:"password"` + RouteByLatency bool `key:"routeByLatency"` +} + +type PostgresConfig struct { + Host string `key:"host"` + Port int `key:"port"` + Name string `key:"name"` + Username string `key:"username"` + Password string `key:"password"` + TimeZone string `key:"timezone"` + EnableTLS bool `key:"enableTLS"` +} + +type GatewayServiceConfig struct { + Host string `key:"host"` + Port int `key:"port"` +} + +type ImageServiceConfig struct { + CacheURL string `key:"cacheURL"` + RegistryStore string `key:"registryStore"` + RegistryCredentialProviderName string `key:"registryCredentialProvider"` + Registries ImageRegistriesConfig `key:"registries"` + EnableTLS bool `key:"enableTLS"` + Runner RunnerConfig `key:"runner"` +} + +type ImageRegistriesConfig struct { + Docker DockerImageRegistryConfig `key:"docker"` + S3 S3ImageRegistryConfig `key:"s3"` +} + +type DockerImageRegistryConfig struct { + Username string `key:"username"` + Password string `key:"password"` +} + +type S3ImageRegistryConfig struct { + AccessKeyID string `key:"accessKeyID"` + SecretAccessKey string `key:"secretAccessKey"` + Bucket string `key:"bucket"` + Region string `key:"region"` +} + +type RunnerConfig struct { + BaseImageName string `key:"baseImageName"` + BaseImageRegistry string `key:"baseImageRegistry"` + BaseImageTag string `key:"baseImageTag"` + Tags map[string]string `key:"tags"` +} + +type StorageConfig struct { + Mode string `key:"mode"` + FilesystemName string `key:"fsName"` + FilesystemPath string `key:"fsPath"` + ObjectPath string `key:"objectPath"` + JuiceFS JuiceFSConfig `key:"juicefs"` +} + +type JuiceFSConfig struct { + RedisURI string `key:"redisURI"` + AWSS3Bucket string `key:"awsS3Bucket"` + AWSAccessKeyID string `key:"awsAccessKeyID"` + AWSSecretAccessKey string `key:"awsSecretAccessKey"` +} + +type WorkerConfig struct { + Pools map[string]WorkerPoolConfig `key:"pools"` + HostNetwork bool `key:"hostNetwork"` + ImageTag string `key:"imageTag"` + ImageName string `key:"imageName"` + ImageRegistry string `key:"imageRegistry"` + ImagePullSecrets []string `key:"imagePullSecrets"` + Namespace string `key:"namespace"` + ServiceAccountName string `key:"serviceAccountName"` + + ResourcesEnforced bool `key:"resourcesEnforced"` + DataVolumeName string `key:"dataVolumeName"` + DefaultWorkerCPURequest int64 `key:"defaultWorkerCPURequest"` + DefaultWorkerMemoryRequest int64 `key:"defaultWorkerMemoryRequest"` +} + +type WorkerPoolConfig struct { + JobSpec WorkerPoolJobSpecConfig `key:"jobSpec"` +} + +type WorkerPoolJobSpecConfig struct { + NodeSelector map[string]string `key:"nodeSelector"` + PoolSizing WorkerPoolJobSpecPoolSizingConfig `key:"poolSizing"` +} + +type WorkerPoolJobSpecPoolSizingConfig struct { + DefaultWorkerCPU string `key:"defaultWorkerCPU"` + DefaultWorkerMemory string `key:"defaultWorkerMemory"` + DefaultWorkerGPUType string `key:"defaultWorkerGPUType"` + MinFreeCPU string `key:"minFreeCPU"` + MinFreeMemory string `key:"minFreeMemory"` + MinFreeGPU string `key:"minFreeGPU"` +} + +type MetricsConfig struct { + Kinesis KinesisConfig `key:"kinesis"` +} + +type KinesisConfig struct { + StreamName string `key:"streamName"` + Region string `key:"region"` + AccessKeyID string `key:"accessKeyID"` + SecretAccessKey string `key:"secretAccessKey"` + SessionKey string `key:"sessionKey"` + Endpoint string `key:"endpoint"` +} diff --git a/internal/types/pool.go b/internal/types/pool.go deleted file mode 100644 index 68de2aedb..000000000 --- a/internal/types/pool.go +++ /dev/null @@ -1,107 +0,0 @@ -package types - -import ( - "fmt" - - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/runtime" -) - -const ( - WorkerPoolSchemaGroupName = "beam.cloud" - WorkerPoolSchemaGroupVersion = "v1" - WorkerPoolResourcePlural = "workerpools" - WorkerPoolResourceSingular = "workerpool" -) - -// WorkerPoolResource is the Go representation of the WorkerPool custom resource -type WorkerPoolResource struct { - metav1.TypeMeta `json:",inline"` - metav1.ObjectMeta `json:"metadata,omitempty"` - - Spec WorkerPoolSpec `json:"spec,omitempty"` -} - -// Gets a new WorkerPoolResource struct. -// A helper function that defines the Name, APIVersion, and Kind of the resource. -func NewWorkerPoolResource(name string) WorkerPoolResource { - return WorkerPoolResource{ - ObjectMeta: metav1.ObjectMeta{ - Name: name, - }, - TypeMeta: metav1.TypeMeta{ - APIVersion: fmt.Sprintf("%v/%v", WorkerPoolSchemaGroupName, WorkerPoolSchemaGroupVersion), - Kind: WorkerPoolResourceSingular, - }, - } -} - -type WorkerPoolSpec struct { - JobSpec JobSpec `json:"jobSpec,omitempty"` - PoolSizing runtime.RawExtension `json:"poolSizing,omitempty"` -} - -type JobSpec struct { - NodeSelector map[string]string `json:"nodeSelector,omitempty"` -} - -type WorkerPoolStatus struct { - CurrentWorkers int `json:"currentWorkers"` -} - -// WorkerPoolList is a list of WorkerPool custom resources -type WorkerPoolList struct { - metav1.TypeMeta `json:",inline"` - metav1.ListMeta `json:"metadata,omitempty"` - Items []WorkerPoolResource `json:"items"` -} - -func (in *WorkerPoolResource) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -func (in *WorkerPoolResource) DeepCopy() *WorkerPoolResource { - if in == nil { - return nil - } - out := new(WorkerPoolResource) - out.TypeMeta = in.TypeMeta - out.ObjectMeta = *in.ObjectMeta.DeepCopy() - - out.Spec.JobSpec.NodeSelector = make(map[string]string, len(in.Spec.JobSpec.NodeSelector)) - for key, val := range in.Spec.JobSpec.NodeSelector { - out.Spec.JobSpec.NodeSelector[key] = val - } - - out.Spec.PoolSizing = runtime.RawExtension{ - Raw: append([]byte(nil), in.Spec.PoolSizing.Raw...), - } - - return out -} - -func (in *WorkerPoolList) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -func (in *WorkerPoolList) DeepCopy() *WorkerPoolList { - if in == nil { - return nil - } - out := new(WorkerPoolList) - out.TypeMeta = in.TypeMeta - out.ListMeta = *in.ListMeta.DeepCopy() - if in.Items != nil { - out.Items = make([]WorkerPoolResource, len(in.Items)) - for i := range in.Items { - out.Items[i] = *in.Items[i].DeepCopy() - } - } - return out -} diff --git a/internal/types/scheduler.go b/internal/types/scheduler.go index ed4ac70f6..9f2a3ee59 100644 --- a/internal/types/scheduler.go +++ b/internal/types/scheduler.go @@ -65,7 +65,6 @@ type ContainerRequest struct { ScheduleTimeout float64 `json:"schedule_timeout"` Mounts []Mount `json:"mounts"` OnScheduleChan chan bool `json:"-"` - Agent string `json:"agent"` } const ContainerExitCodeTtlS int = 300 diff --git a/internal/worker/creds.go b/internal/worker/creds.go index 48f91e2b9..4b683c2b1 100644 --- a/internal/worker/creds.go +++ b/internal/worker/creds.go @@ -8,7 +8,6 @@ import ( awsconfig "github.com/aws/aws-sdk-go-v2/config" ecr "github.com/aws/aws-sdk-go-v2/service/ecr" - "github.com/beam-cloud/beam/internal/common" ) type CredentialProvider interface { @@ -19,6 +18,8 @@ type CredentialProvider interface { // AWS auth provider type AWSCredentialProvider struct { + CredentialProvider + Region string } func (p *AWSCredentialProvider) GetUsername() string { @@ -26,7 +27,7 @@ func (p *AWSCredentialProvider) GetUsername() string { } func (p *AWSCredentialProvider) GetAuthorizationToken() (string, error) { - cfg, err := awsconfig.LoadDefaultConfig(context.TODO(), awsconfig.WithRegion(common.Secrets().Get("AWS_REGION"))) + cfg, err := awsconfig.LoadDefaultConfig(context.TODO(), awsconfig.WithRegion(p.Region)) if err != nil { log.Fatalf("unable to load SDK config, %v", err) return "", nil @@ -49,27 +50,30 @@ func (p *AWSCredentialProvider) GetAuthorizationToken() (string, error) { func (p *AWSCredentialProvider) GetAuthString() (string, error) { token, err := p.GetAuthorizationToken() if err != nil { - return "", err + return "", fmt.Errorf("failed to get aws provider auth string: %v", err) } return token, nil } // Docker provider type DockerCredentialProvider struct { + CredentialProvider + Username string + Password string } func (p *DockerCredentialProvider) GetUsername() string { - return common.Secrets().Get("DOCKER_USERNAME") + return p.Username } func (p *DockerCredentialProvider) GetAuthorizationToken() (string, error) { - return common.Secrets().Get("DOCKER_PASSWORD"), nil + return p.Password, nil } func (p *DockerCredentialProvider) GetAuthString() (string, error) { token, err := p.GetAuthorizationToken() if err != nil { - return "", err + return "", fmt.Errorf("failed to get docker provider auth string: %v", err) } return fmt.Sprintf("%s:%s", p.GetUsername(), token), nil } diff --git a/internal/worker/cuda_test.go b/internal/worker/cuda_test.go index e2abf72ed..bb120e516 100644 --- a/internal/worker/cuda_test.go +++ b/internal/worker/cuda_test.go @@ -27,8 +27,8 @@ func TestInjectCudaEnvVarsNoCudaInImage(t *testing.T) { "NVIDIA_VISIBLE_DEVICES=", "CUDA_VERSION=", "GPU_TYPE=", - "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/cuda-12.2/bin:$PATH", - "LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/lib/worker/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/cuda-12.2/targets/x86_64-linux/lib:$LD_LIBRARY_PATH", + "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/cuda-12.3/bin:$PATH", + "LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/lib/worker/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/cuda-12.3/targets/x86_64-linux/lib:$LD_LIBRARY_PATH", } resultEnv, _ := manager.InjectCudaEnvVars(initialEnv, &ContainerOptions{ @@ -48,7 +48,7 @@ func TestInjectCudaEnvVarsExistingCudaInImage(t *testing.T) { // Set some environment variables to simulate NVIDIA settings os.Setenv("NVIDIA_DRIVER_CAPABILITIES", "all") os.Setenv("NVIDIA_REQUIRE_CUDA", "cuda>=9.0") - os.Setenv("CUDA_VERSION", "12.2") + os.Setenv("CUDA_VERSION", "12.3") expectedEnv := []string{ "INITIAL=1", diff --git a/internal/worker/image.go b/internal/worker/image.go index 29c7c1d1d..1337118ce 100644 --- a/internal/worker/image.go +++ b/internal/worker/image.go @@ -14,6 +14,7 @@ import ( "github.com/beam-cloud/beam/internal/abstractions/image" common "github.com/beam-cloud/beam/internal/common" + types "github.com/beam-cloud/beam/internal/types" "github.com/beam-cloud/clip/pkg/clip" clipCommon "github.com/beam-cloud/clip/pkg/common" "github.com/moby/sys/mountinfo" @@ -27,11 +28,10 @@ import ( ) const ( - imagePullCommand string = "skopeo" - awsCredentialProviderName string = "aws" - imageCachePath string = "/dev/shm/images" - imageAvailableFilename string = "IMAGE_AVAILABLE" - imageMountLockFilename string = "IMAGE_MOUNT_LOCK" + imagePullCommand string = "skopeo" + imageCachePath string = "/dev/shm/images" + imageAvailableFilename string = "IMAGE_AVAILABLE" + imageMountLockFilename string = "IMAGE_MOUNT_LOCK" ) var requiredContainerDirectories []string = []string{"/workspace", "/volumes", "/snapshot"} @@ -45,31 +45,34 @@ type ImageClient struct { CommandTimeout int Debug bool Creds string - VerifyTLS bool + config types.ImageServiceConfig } -func NewImageClient() (*ImageClient, error) { +func NewImageClient(config types.ImageServiceConfig) (*ImageClient, error) { var provider CredentialProvider // Configure image registry credentials - providerName := common.Secrets().GetWithDefault("BEAM_IMAGESERVICE_IMAGE_CREDENTIAL_PROVIDER", "aws") - if providerName == awsCredentialProviderName { - provider = &AWSCredentialProvider{} - } else { - provider = &DockerCredentialProvider{} + switch config.RegistryCredentialProviderName { + case "aws": + provider = &AWSCredentialProvider{ + Region: config.Registries.S3.Region, + } + case "docker": + provider = &DockerCredentialProvider{ + Username: config.Registries.Docker.Username, + Password: config.Registries.Docker.Password, + } + default: + return nil, fmt.Errorf("invalid credential provider name: %s", config.RegistryCredentialProviderName) } - storeName := common.Secrets().GetWithDefault("BEAM_IMAGESERVICE_IMAGE_REGISTRY_STORE", "s3") - registry, err := common.NewImageRegistry(storeName) + registry, err := common.NewImageRegistry(config) if err != nil { return nil, err } - verifyTLS := common.Secrets().GetWithDefault("BEAM_IMAGESERVICE_IMAGE_TLS_VERIFY", "true") != "false" - - cacheUrl, cacheUrlSet := os.LookupEnv("BEAM_CACHE_URL") var cacheClient *CacheClient = nil - if cacheUrlSet && cacheUrl != "" { - cacheClient, err = NewCacheClient(cacheUrl, "") + if config.CacheURL != "" { + cacheClient, err = NewCacheClient(config.CacheURL, "") if err != nil { return nil, err } @@ -84,6 +87,7 @@ func NewImageClient() (*ImageClient, error) { } return &ImageClient{ + config: config, registry: registry, cacheClient: cacheClient, ImagePath: baseImagePath, @@ -91,7 +95,6 @@ func NewImageClient() (*ImageClient, error) { CommandTimeout: -1, Debug: false, Creds: creds, - VerifyTLS: verifyTLS, }, nil } @@ -148,10 +151,6 @@ func (i *ImageClient) PullAndArchiveImage(ctx context.Context, sourceImage strin dest := fmt.Sprintf("oci:%s:%s", baseImage.ImageName, baseImage.ImageTag) args := []string{"copy", fmt.Sprintf("docker://%s", sourceImage), dest} - if !i.VerifyTLS { - args = append(args, []string{"--src-tls-verify=false", "--dest-tls-verify=false"}...) - } - args = append(args, i.args(creds)...) cmd := exec.CommandContext(ctx, i.PullCommand, args...) cmd.Env = os.Environ() @@ -230,6 +229,10 @@ func (i *ImageClient) args(creds *string) (out []string) { out = append(out, "--command-timeout", fmt.Sprintf("%d", i.CommandTimeout)) } + if !i.config.EnableTLS { + out = append(out, []string{"--src-tls-verify=false", "--dest-tls-verify=false"}...) + } + if i.Debug { out = append(out, "--debug") } @@ -286,15 +289,14 @@ func (i *ImageClient) Archive(ctx context.Context, bundlePath string, imageId st }() var err error = nil - archiveStore := common.Secrets().GetWithDefault("BEAM_IMAGESERVICE_IMAGE_REGISTRY_STORE", "s3") - switch archiveStore { + switch i.config.RegistryStore { case "s3": err = clip.CreateAndUploadArchive(clip.CreateOptions{ InputPath: bundlePath, OutputPath: archivePath, }, &clipCommon.S3StorageInfo{ - Bucket: common.Secrets().Get("BEAM_IMAGESERVICE_IMAGE_REGISTRY_S3_BUCKET"), - Region: common.Secrets().Get("BEAM_IMAGESERVICE_IMAGE_REGISTRY_S3_REGION"), + Bucket: i.config.Registries.S3.Bucket, + Region: i.config.Registries.S3.Region, Key: fmt.Sprintf("%s.clip", imageId), }) case "local": diff --git a/internal/worker/network.go b/internal/worker/network.go index 5250470a9..4227a58f2 100644 --- a/internal/worker/network.go +++ b/internal/worker/network.go @@ -5,7 +5,6 @@ import ( "fmt" "net" "os" - "strings" ) func GetRandomFreePort() (int, error) { @@ -23,25 +22,6 @@ func GetRandomFreePort() (int, error) { return l.Addr().(*net.TCPAddr).Port, nil } -func GetPodFQDN() (string, error) { - podIP := os.Getenv("POD_IP") - podNamespace := os.Getenv("POD_NAMESPACE") - clusterDomain := os.Getenv("CLUSTER_DOMAIN") - - if podIP == "" || podNamespace == "" || clusterDomain == "" { - return "", fmt.Errorf("environment variables POD_IP, POD_NAMESPACE, or CLUSTER_DOMAIN are not set") - } - - dnsName := fmt.Sprintf( - "%s.%s.pod.%s", - strings.ReplaceAll(podIP, ".", "-"), - podNamespace, - clusterDomain, - ) - - return dnsName, nil -} - // GetPodIP gets the IP from the POD_IP env var. // Returns an error if it fails to retrieve an IP. func GetPodIP() (string, error) { @@ -62,7 +42,7 @@ func getIPFromInterface(ifname string) (string, error) { } if len(addrs) < 1 { - return "", fmt.Errorf("no IP addresses found on <%s> interface", ifname) + return "", fmt.Errorf("no ip addresses found on <%s> interface", ifname) } ip, _, err := net.ParseCIDR(addrs[0].String()) @@ -77,7 +57,7 @@ func getIPFromInterface(ifname string) (string, error) { func getIPFromEnv(varName string) (string, error) { addr := os.Getenv(varName) if addr == "" { - return "", fmt.Errorf("no IP found in environment variable") + return "", fmt.Errorf("no ip found in environment variable") } ip := net.ParseIP(addr) diff --git a/internal/worker/util.go b/internal/worker/util.go index 045556e3e..15947c70c 100644 --- a/internal/worker/util.go +++ b/internal/worker/util.go @@ -60,7 +60,7 @@ func (fl *FileLock) Release() error { return err } - err = fl.file.Close() + fl.file.Close() fl.file = nil err = os.Remove(fl.path) diff --git a/internal/worker/worker.go b/internal/worker/worker.go index d819955c3..98b529a9f 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -53,6 +53,7 @@ type Worker struct { storage storage.Storage ctx context.Context cancel func() + config types.AppConfig } type ContainerInstance struct { @@ -104,12 +105,18 @@ func NewWorker() (*Worker, error) { return nil, err } - redisClient, err := common.NewRedisClient(common.WithClientName("BeamWorker")) + configManager, err := common.NewConfigManager[types.AppConfig]() if err != nil { return nil, err } + config := configManager.GetConfig() - imageClient, err := NewImageClient() + redisClient, err := common.NewRedisClient(config.Database.Redis, common.WithClientName("BeamWorker")) + if err != nil { + return nil, err + } + + imageClient, err := NewImageClient(config.ImageService) if err != nil { return nil, err } @@ -124,7 +131,7 @@ func NewWorker() (*Worker, error) { return nil, err } - storage, err := storage.NewStorage() + storage, err := storage.NewStorage(config.Storage) if err != nil { return nil, err } @@ -135,12 +142,13 @@ func NewWorker() (*Worker, error) { workerRepo := repo.NewWorkerRedisRepository(redisClient) statsdRepo := repo.NewMetricsStatsdRepository() - workerMetrics := NewWorkerMetrics(ctx, podHostName, statsdRepo, workerRepo, repo.NewMetricsStreamRepository(ctx)) + workerMetrics := NewWorkerMetrics(ctx, podHostName, statsdRepo, workerRepo, repo.NewMetricsStreamRepository(ctx, config.Metrics)) workerMetrics.InitNvml() return &Worker{ ctx: ctx, cancel: cancel, + config: config, userImagePath: imagePath, cpuLimit: cpuLimit, memoryLimit: memoryLimit, @@ -542,8 +550,8 @@ func (s *Worker) getContainerEnvironment(request *types.ContainerRequest, option fmt.Sprintf("BIND_PORT=%d", options.BindPort), fmt.Sprintf("CONTAINER_HOSTNAME=%s", fmt.Sprintf("%s:%d", s.podIPAddr, options.BindPort)), fmt.Sprintf("CONTAINER_ID=%s", request.ContainerId), - fmt.Sprintf("BEAM_GATEWAY_HOST=%s", os.Getenv("BEAM_GATEWAY_HOST")), - fmt.Sprintf("BEAM_GATEWAY_PORT=%s", os.Getenv("BEAM_GATEWAY_PORT")), + fmt.Sprintf("BEAM_GATEWAY_HOST=%s", s.config.GatewayService.Host), + fmt.Sprintf("BEAM_GATEWAY_PORT=%d", s.config.GatewayService.Port), "PYTHONUNBUFFERED=1", } env = append(env, request.Env...) diff --git a/manifests/crds/workerpool.yaml b/manifests/crds/workerpool.yaml deleted file mode 100644 index 7a9ca9b9b..000000000 --- a/manifests/crds/workerpool.yaml +++ /dev/null @@ -1,33 +0,0 @@ -apiVersion: apiextensions.k8s.io/v1 -kind: CustomResourceDefinition -metadata: - name: workerpools.beam.cloud -spec: - group: beam.cloud - versions: - - name: v1 - served: true - storage: true - schema: - openAPIV3Schema: - x-kubernetes-preserve-unknown-fields: true - type: object - properties: - spec: - type: object - x-kubernetes-preserve-unknown-fields: true - properties: - jobSpec: - type: object - properties: - nodeSelector: - type: object - additionalProperties: - type: string - names: - kind: WorkerPool - plural: workerpools - singular: workerpool - shortNames: - - wp - scope: Namespaced diff --git a/manifests/k3d/beam.yaml b/manifests/k3d/beam.yaml index 0da364074..af2449cfa 100644 --- a/manifests/k3d/beam.yaml +++ b/manifests/k3d/beam.yaml @@ -136,23 +136,6 @@ spec: - name: beam port: 1993 --- -apiVersion: beam.cloud/v1 -kind: WorkerPool -metadata: - name: beam-cpu - namespace: beam -spec: - # jobSpec: - # nodeSelector: - # node-label-key: node-label-value - poolSizing: - defaultWorkerCpu: 1000m - defaultWorkerGpuType: "" - defaultWorkerMemory: 1Gi - minFreeCpu: 1000m - minFreeGpu: 0 - minFreeMemory: 1Gi ---- apiVersion: helm.cattle.io/v1 kind: HelmChart metadata: diff --git a/proto/scheduler.pb.go b/proto/scheduler.pb.go index 609cedd09..9fcac8bae 100644 --- a/proto/scheduler.pb.go +++ b/proto/scheduler.pb.go @@ -365,593 +365,6 @@ func (x *StopContainerResponse) GetError() string { return "" } -type AgentInfo struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Token string `protobuf:"bytes,1,opt,name=token,proto3" json:"token,omitempty"` - Version string `protobuf:"bytes,2,opt,name=version,proto3" json:"version,omitempty"` - CloudProvider string `protobuf:"bytes,3,opt,name=cloud_provider,json=cloudProvider,proto3" json:"cloud_provider,omitempty"` -} - -func (x *AgentInfo) Reset() { - *x = AgentInfo{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[6] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *AgentInfo) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*AgentInfo) ProtoMessage() {} - -func (x *AgentInfo) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[6] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use AgentInfo.ProtoReflect.Descriptor instead. -func (*AgentInfo) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{6} -} - -func (x *AgentInfo) GetToken() string { - if x != nil { - return x.Token - } - return "" -} - -func (x *AgentInfo) GetVersion() string { - if x != nil { - return x.Version - } - return "" -} - -func (x *AgentInfo) GetCloudProvider() string { - if x != nil { - return x.CloudProvider - } - return "" -} - -type WorkerPool struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` -} - -func (x *WorkerPool) Reset() { - *x = WorkerPool{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[7] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *WorkerPool) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*WorkerPool) ProtoMessage() {} - -func (x *WorkerPool) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[7] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use WorkerPool.ProtoReflect.Descriptor instead. -func (*WorkerPool) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{7} -} - -func (x *WorkerPool) GetName() string { - if x != nil { - return x.Name - } - return "" -} - -type Worker struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Cpu int64 `protobuf:"varint,2,opt,name=cpu,proto3" json:"cpu,omitempty"` - Memory int64 `protobuf:"varint,3,opt,name=memory,proto3" json:"memory,omitempty"` - GpuType string `protobuf:"bytes,4,opt,name=gpu_type,json=gpuType,proto3" json:"gpu_type,omitempty"` - Status string `protobuf:"bytes,5,opt,name=status,proto3" json:"status,omitempty"` - PoolId string `protobuf:"bytes,6,opt,name=pool_id,json=poolId,proto3" json:"pool_id,omitempty"` -} - -func (x *Worker) Reset() { - *x = Worker{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[8] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *Worker) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*Worker) ProtoMessage() {} - -func (x *Worker) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[8] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use Worker.ProtoReflect.Descriptor instead. -func (*Worker) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{8} -} - -func (x *Worker) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *Worker) GetCpu() int64 { - if x != nil { - return x.Cpu - } - return 0 -} - -func (x *Worker) GetMemory() int64 { - if x != nil { - return x.Memory - } - return 0 -} - -func (x *Worker) GetGpuType() string { - if x != nil { - return x.GpuType - } - return "" -} - -func (x *Worker) GetStatus() string { - if x != nil { - return x.Status - } - return "" -} - -func (x *Worker) GetPoolId() string { - if x != nil { - return x.PoolId - } - return "" -} - -type SubscribeWorkerEventRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - AgentInfo *AgentInfo `protobuf:"bytes,1,opt,name=agent_info,json=agentInfo,proto3" json:"agent_info,omitempty"` - WorkerPools []*WorkerPool `protobuf:"bytes,2,rep,name=worker_pools,json=workerPools,proto3" json:"worker_pools,omitempty"` -} - -func (x *SubscribeWorkerEventRequest) Reset() { - *x = SubscribeWorkerEventRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[9] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *SubscribeWorkerEventRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*SubscribeWorkerEventRequest) ProtoMessage() {} - -func (x *SubscribeWorkerEventRequest) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[9] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use SubscribeWorkerEventRequest.ProtoReflect.Descriptor instead. -func (*SubscribeWorkerEventRequest) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{9} -} - -func (x *SubscribeWorkerEventRequest) GetAgentInfo() *AgentInfo { - if x != nil { - return x.AgentInfo - } - return nil -} - -func (x *SubscribeWorkerEventRequest) GetWorkerPools() []*WorkerPool { - if x != nil { - return x.WorkerPools - } - return nil -} - -type SubscribeWorkerEventResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Worker *Worker `protobuf:"bytes,1,opt,name=worker,proto3" json:"worker,omitempty"` - WorkerPool *WorkerPool `protobuf:"bytes,2,opt,name=workerPool,proto3" json:"workerPool,omitempty"` -} - -func (x *SubscribeWorkerEventResponse) Reset() { - *x = SubscribeWorkerEventResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[10] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *SubscribeWorkerEventResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*SubscribeWorkerEventResponse) ProtoMessage() {} - -func (x *SubscribeWorkerEventResponse) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[10] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use SubscribeWorkerEventResponse.ProtoReflect.Descriptor instead. -func (*SubscribeWorkerEventResponse) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{10} -} - -func (x *SubscribeWorkerEventResponse) GetWorker() *Worker { - if x != nil { - return x.Worker - } - return nil -} - -func (x *SubscribeWorkerEventResponse) GetWorkerPool() *WorkerPool { - if x != nil { - return x.WorkerPool - } - return nil -} - -type StatusRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields -} - -func (x *StatusRequest) Reset() { - *x = StatusRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[11] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *StatusRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*StatusRequest) ProtoMessage() {} - -func (x *StatusRequest) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[11] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use StatusRequest.ProtoReflect.Descriptor instead. -func (*StatusRequest) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{11} -} - -type StatusResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Status string `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` -} - -func (x *StatusResponse) Reset() { - *x = StatusResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[12] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *StatusResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*StatusResponse) ProtoMessage() {} - -func (x *StatusResponse) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[12] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use StatusResponse.ProtoReflect.Descriptor instead. -func (*StatusResponse) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{12} -} - -func (x *StatusResponse) GetStatus() string { - if x != nil { - return x.Status - } - return "" -} - -type StreamLogRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Log string `protobuf:"bytes,2,opt,name=log,proto3" json:"log,omitempty"` -} - -func (x *StreamLogRequest) Reset() { - *x = StreamLogRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[13] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *StreamLogRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*StreamLogRequest) ProtoMessage() {} - -func (x *StreamLogRequest) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[13] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use StreamLogRequest.ProtoReflect.Descriptor instead. -func (*StreamLogRequest) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{13} -} - -func (x *StreamLogRequest) GetId() string { - if x != nil { - return x.Id - } - return "" -} - -func (x *StreamLogRequest) GetLog() string { - if x != nil { - return x.Log - } - return "" -} - -type StreamLogResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` -} - -func (x *StreamLogResponse) Reset() { - *x = StreamLogResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[14] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *StreamLogResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*StreamLogResponse) ProtoMessage() {} - -func (x *StreamLogResponse) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[14] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use StreamLogResponse.ProtoReflect.Descriptor instead. -func (*StreamLogResponse) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{14} -} - -func (x *StreamLogResponse) GetOk() bool { - if x != nil { - return x.Ok - } - return false -} - -type RegisterWorkerRequest struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields - - AgentInfo *AgentInfo `protobuf:"bytes,1,opt,name=agent_info,json=agentInfo,proto3" json:"agent_info,omitempty"` - Worker *Worker `protobuf:"bytes,2,opt,name=worker,proto3" json:"worker,omitempty"` -} - -func (x *RegisterWorkerRequest) Reset() { - *x = RegisterWorkerRequest{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[15] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *RegisterWorkerRequest) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*RegisterWorkerRequest) ProtoMessage() {} - -func (x *RegisterWorkerRequest) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[15] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use RegisterWorkerRequest.ProtoReflect.Descriptor instead. -func (*RegisterWorkerRequest) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{15} -} - -func (x *RegisterWorkerRequest) GetAgentInfo() *AgentInfo { - if x != nil { - return x.AgentInfo - } - return nil -} - -func (x *RegisterWorkerRequest) GetWorker() *Worker { - if x != nil { - return x.Worker - } - return nil -} - -type RegisterWorkerResponse struct { - state protoimpl.MessageState - sizeCache protoimpl.SizeCache - unknownFields protoimpl.UnknownFields -} - -func (x *RegisterWorkerResponse) Reset() { - *x = RegisterWorkerResponse{} - if protoimpl.UnsafeEnabled { - mi := &file_scheduler_proto_msgTypes[16] - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - ms.StoreMessageInfo(mi) - } -} - -func (x *RegisterWorkerResponse) String() string { - return protoimpl.X.MessageStringOf(x) -} - -func (*RegisterWorkerResponse) ProtoMessage() {} - -func (x *RegisterWorkerResponse) ProtoReflect() protoreflect.Message { - mi := &file_scheduler_proto_msgTypes[16] - if protoimpl.UnsafeEnabled && x != nil { - ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) - if ms.LoadMessageInfo() == nil { - ms.StoreMessageInfo(mi) - } - return ms - } - return mi.MessageOf(x) -} - -// Deprecated: Use RegisterWorkerResponse.ProtoReflect.Descriptor instead. -func (*RegisterWorkerResponse) Descriptor() ([]byte, []int) { - return file_scheduler_proto_rawDescGZIP(), []int{16} -} - var File_scheduler_proto protoreflect.FileDescriptor var file_scheduler_proto_rawDesc = []byte{ @@ -988,91 +401,26 @@ var file_scheduler_proto_rawDesc = []byte{ 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x18, 0x0a, 0x07, 0x73, 0x75, 0x63, 0x63, 0x65, 0x73, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x73, 0x75, 0x63, 0x63, 0x65, 0x73, 0x73, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, - 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0x62, 0x0a, 0x09, 0x41, 0x67, 0x65, 0x6e, 0x74, - 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x14, 0x0a, 0x05, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x05, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x18, 0x0a, 0x07, 0x76, 0x65, - 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x76, 0x65, 0x72, - 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x5f, 0x70, 0x72, - 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6c, - 0x6f, 0x75, 0x64, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x22, 0x20, 0x0a, 0x0a, 0x57, - 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x50, 0x6f, 0x6f, 0x6c, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, - 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x8e, 0x01, - 0x0a, 0x06, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x10, 0x0a, 0x03, 0x63, 0x70, 0x75, 0x18, - 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x03, 0x63, 0x70, 0x75, 0x12, 0x16, 0x0a, 0x06, 0x6d, 0x65, - 0x6d, 0x6f, 0x72, 0x79, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x06, 0x6d, 0x65, 0x6d, 0x6f, - 0x72, 0x79, 0x12, 0x19, 0x0a, 0x08, 0x67, 0x70, 0x75, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x18, 0x04, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x67, 0x70, 0x75, 0x54, 0x79, 0x70, 0x65, 0x12, 0x16, 0x0a, - 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73, - 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x17, 0x0a, 0x07, 0x70, 0x6f, 0x6f, 0x6c, 0x5f, 0x69, 0x64, - 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x70, 0x6f, 0x6f, 0x6c, 0x49, 0x64, 0x22, 0x8c, - 0x01, 0x0a, 0x1b, 0x53, 0x75, 0x62, 0x73, 0x63, 0x72, 0x69, 0x62, 0x65, 0x57, 0x6f, 0x72, 0x6b, - 0x65, 0x72, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x33, - 0x0a, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x6e, 0x66, 0x6f, 0x18, 0x01, 0x20, 0x01, - 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x41, - 0x67, 0x65, 0x6e, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x49, - 0x6e, 0x66, 0x6f, 0x12, 0x38, 0x0a, 0x0c, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x5f, 0x70, 0x6f, - 0x6f, 0x6c, 0x73, 0x18, 0x02, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x15, 0x2e, 0x73, 0x63, 0x68, 0x65, - 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x50, 0x6f, 0x6f, 0x6c, - 0x52, 0x0b, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x50, 0x6f, 0x6f, 0x6c, 0x73, 0x22, 0x80, 0x01, - 0x0a, 0x1c, 0x53, 0x75, 0x62, 0x73, 0x63, 0x72, 0x69, 0x62, 0x65, 0x57, 0x6f, 0x72, 0x6b, 0x65, - 0x72, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x29, - 0x0a, 0x06, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, - 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, - 0x72, 0x52, 0x06, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x12, 0x35, 0x0a, 0x0a, 0x77, 0x6f, 0x72, - 0x6b, 0x65, 0x72, 0x50, 0x6f, 0x6f, 0x6c, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x15, 0x2e, - 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, - 0x50, 0x6f, 0x6f, 0x6c, 0x52, 0x0a, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x50, 0x6f, 0x6f, 0x6c, - 0x22, 0x0f, 0x0a, 0x0d, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x22, 0x28, 0x0a, 0x0e, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x34, 0x0a, 0x10, 0x53, - 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4c, 0x6f, 0x67, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, - 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, - 0x10, 0x0a, 0x03, 0x6c, 0x6f, 0x67, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6c, 0x6f, - 0x67, 0x22, 0x23, 0x0a, 0x11, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4c, 0x6f, 0x67, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, 0x20, 0x01, - 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x22, 0x77, 0x0a, 0x15, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, - 0x65, 0x72, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, - 0x33, 0x0a, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x69, 0x6e, 0x66, 0x6f, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, - 0x41, 0x67, 0x65, 0x6e, 0x74, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, - 0x49, 0x6e, 0x66, 0x6f, 0x12, 0x29, 0x0a, 0x06, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x18, 0x02, - 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, - 0x2e, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x52, 0x06, 0x77, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x22, - 0x18, 0x0a, 0x16, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x57, 0x6f, 0x72, 0x6b, 0x65, - 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xbe, 0x03, 0x0a, 0x09, 0x53, 0x63, - 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x12, 0x45, 0x0a, 0x0a, 0x47, 0x65, 0x74, 0x56, 0x65, - 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, - 0x72, 0x2e, 0x56, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x1a, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x56, 0x65, 0x72, - 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x51, - 0x0a, 0x0c, 0x52, 0x75, 0x6e, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x12, 0x1e, - 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x43, 0x6f, - 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, - 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x43, 0x6f, - 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, - 0x00, 0x12, 0x54, 0x0a, 0x0d, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, - 0x65, 0x72, 0x12, 0x1f, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x53, - 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, - 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, - 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x73, - 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x6a, 0x0a, 0x15, 0x53, 0x75, 0x62, 0x73, 0x63, - 0x72, 0x69, 0x62, 0x65, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, - 0x12, 0x26, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x53, 0x75, 0x62, - 0x73, 0x63, 0x72, 0x69, 0x62, 0x65, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x45, 0x76, 0x65, 0x6e, - 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x27, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, - 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x53, 0x75, 0x62, 0x73, 0x63, 0x72, 0x69, 0x62, 0x65, 0x57, 0x6f, - 0x72, 0x6b, 0x65, 0x72, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x30, 0x01, 0x12, 0x55, 0x0a, 0x0e, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x57, - 0x6f, 0x72, 0x6b, 0x65, 0x72, 0x12, 0x20, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, - 0x72, 0x2e, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x57, 0x6f, 0x72, 0x6b, 0x65, 0x72, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, - 0x6c, 0x65, 0x72, 0x2e, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x57, 0x6f, 0x72, 0x6b, - 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x22, 0x5a, 0x20, 0x67, 0x69, - 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x62, 0x65, 0x61, 0x6d, 0x2d, 0x63, 0x6c, - 0x6f, 0x75, 0x64, 0x2f, 0x62, 0x65, 0x61, 0x6d, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x32, 0xfb, 0x01, 0x0a, 0x09, 0x53, 0x63, 0x68, 0x65, + 0x64, 0x75, 0x6c, 0x65, 0x72, 0x12, 0x45, 0x0a, 0x0a, 0x47, 0x65, 0x74, 0x56, 0x65, 0x72, 0x73, + 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, + 0x56, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, + 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x56, 0x65, 0x72, 0x73, 0x69, + 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x51, 0x0a, 0x0c, + 0x52, 0x75, 0x6e, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x12, 0x1e, 0x2e, 0x73, + 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x43, 0x6f, 0x6e, 0x74, + 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x73, + 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x43, 0x6f, 0x6e, 0x74, + 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, + 0x54, 0x0a, 0x0d, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, + 0x12, 0x1f, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x53, 0x74, 0x6f, + 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x20, 0x2e, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x2e, 0x53, 0x74, + 0x6f, 0x70, 0x43, 0x6f, 0x6e, 0x74, 0x61, 0x69, 0x6e, 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x22, 0x5a, 0x20, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, + 0x63, 0x6f, 0x6d, 0x2f, 0x62, 0x65, 0x61, 0x6d, 0x2d, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2f, 0x62, + 0x65, 0x61, 0x6d, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x33, } var ( @@ -1087,48 +435,27 @@ func file_scheduler_proto_rawDescGZIP() []byte { return file_scheduler_proto_rawDescData } -var file_scheduler_proto_msgTypes = make([]protoimpl.MessageInfo, 17) +var file_scheduler_proto_msgTypes = make([]protoimpl.MessageInfo, 6) var file_scheduler_proto_goTypes = []interface{}{ - (*VersionRequest)(nil), // 0: scheduler.VersionRequest - (*VersionResponse)(nil), // 1: scheduler.VersionResponse - (*RunContainerRequest)(nil), // 2: scheduler.RunContainerRequest - (*RunContainerResponse)(nil), // 3: scheduler.RunContainerResponse - (*StopContainerRequest)(nil), // 4: scheduler.StopContainerRequest - (*StopContainerResponse)(nil), // 5: scheduler.StopContainerResponse - (*AgentInfo)(nil), // 6: scheduler.AgentInfo - (*WorkerPool)(nil), // 7: scheduler.WorkerPool - (*Worker)(nil), // 8: scheduler.Worker - (*SubscribeWorkerEventRequest)(nil), // 9: scheduler.SubscribeWorkerEventRequest - (*SubscribeWorkerEventResponse)(nil), // 10: scheduler.SubscribeWorkerEventResponse - (*StatusRequest)(nil), // 11: scheduler.StatusRequest - (*StatusResponse)(nil), // 12: scheduler.StatusResponse - (*StreamLogRequest)(nil), // 13: scheduler.StreamLogRequest - (*StreamLogResponse)(nil), // 14: scheduler.StreamLogResponse - (*RegisterWorkerRequest)(nil), // 15: scheduler.RegisterWorkerRequest - (*RegisterWorkerResponse)(nil), // 16: scheduler.RegisterWorkerResponse + (*VersionRequest)(nil), // 0: scheduler.VersionRequest + (*VersionResponse)(nil), // 1: scheduler.VersionResponse + (*RunContainerRequest)(nil), // 2: scheduler.RunContainerRequest + (*RunContainerResponse)(nil), // 3: scheduler.RunContainerResponse + (*StopContainerRequest)(nil), // 4: scheduler.StopContainerRequest + (*StopContainerResponse)(nil), // 5: scheduler.StopContainerResponse } var file_scheduler_proto_depIdxs = []int32{ - 6, // 0: scheduler.SubscribeWorkerEventRequest.agent_info:type_name -> scheduler.AgentInfo - 7, // 1: scheduler.SubscribeWorkerEventRequest.worker_pools:type_name -> scheduler.WorkerPool - 8, // 2: scheduler.SubscribeWorkerEventResponse.worker:type_name -> scheduler.Worker - 7, // 3: scheduler.SubscribeWorkerEventResponse.workerPool:type_name -> scheduler.WorkerPool - 6, // 4: scheduler.RegisterWorkerRequest.agent_info:type_name -> scheduler.AgentInfo - 8, // 5: scheduler.RegisterWorkerRequest.worker:type_name -> scheduler.Worker - 0, // 6: scheduler.Scheduler.GetVersion:input_type -> scheduler.VersionRequest - 2, // 7: scheduler.Scheduler.RunContainer:input_type -> scheduler.RunContainerRequest - 4, // 8: scheduler.Scheduler.StopContainer:input_type -> scheduler.StopContainerRequest - 9, // 9: scheduler.Scheduler.SubscribeWorkerEvents:input_type -> scheduler.SubscribeWorkerEventRequest - 15, // 10: scheduler.Scheduler.RegisterWorker:input_type -> scheduler.RegisterWorkerRequest - 1, // 11: scheduler.Scheduler.GetVersion:output_type -> scheduler.VersionResponse - 3, // 12: scheduler.Scheduler.RunContainer:output_type -> scheduler.RunContainerResponse - 5, // 13: scheduler.Scheduler.StopContainer:output_type -> scheduler.StopContainerResponse - 10, // 14: scheduler.Scheduler.SubscribeWorkerEvents:output_type -> scheduler.SubscribeWorkerEventResponse - 16, // 15: scheduler.Scheduler.RegisterWorker:output_type -> scheduler.RegisterWorkerResponse - 11, // [11:16] is the sub-list for method output_type - 6, // [6:11] is the sub-list for method input_type - 6, // [6:6] is the sub-list for extension type_name - 6, // [6:6] is the sub-list for extension extendee - 0, // [0:6] is the sub-list for field type_name + 0, // 0: scheduler.Scheduler.GetVersion:input_type -> scheduler.VersionRequest + 2, // 1: scheduler.Scheduler.RunContainer:input_type -> scheduler.RunContainerRequest + 4, // 2: scheduler.Scheduler.StopContainer:input_type -> scheduler.StopContainerRequest + 1, // 3: scheduler.Scheduler.GetVersion:output_type -> scheduler.VersionResponse + 3, // 4: scheduler.Scheduler.RunContainer:output_type -> scheduler.RunContainerResponse + 5, // 5: scheduler.Scheduler.StopContainer:output_type -> scheduler.StopContainerResponse + 3, // [3:6] is the sub-list for method output_type + 0, // [0:3] is the sub-list for method input_type + 0, // [0:0] is the sub-list for extension type_name + 0, // [0:0] is the sub-list for extension extendee + 0, // [0:0] is the sub-list for field type_name } func init() { file_scheduler_proto_init() } @@ -1209,138 +536,6 @@ func file_scheduler_proto_init() { return nil } } - file_scheduler_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*AgentInfo); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*WorkerPool); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[8].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*Worker); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[9].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*SubscribeWorkerEventRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[10].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*SubscribeWorkerEventResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[11].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*StatusRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[12].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*StatusResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[13].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*StreamLogRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[14].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*StreamLogResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[15].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*RegisterWorkerRequest); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } - file_scheduler_proto_msgTypes[16].Exporter = func(v interface{}, i int) interface{} { - switch v := v.(*RegisterWorkerResponse); i { - case 0: - return &v.state - case 1: - return &v.sizeCache - case 2: - return &v.unknownFields - default: - return nil - } - } } type x struct{} out := protoimpl.TypeBuilder{ @@ -1348,7 +543,7 @@ func file_scheduler_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: file_scheduler_proto_rawDesc, NumEnums: 0, - NumMessages: 17, + NumMessages: 6, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/scheduler_grpc.pb.go b/proto/scheduler_grpc.pb.go index 807c43d8b..b78e58a37 100644 --- a/proto/scheduler_grpc.pb.go +++ b/proto/scheduler_grpc.pb.go @@ -19,11 +19,9 @@ import ( const _ = grpc.SupportPackageIsVersion7 const ( - Scheduler_GetVersion_FullMethodName = "/scheduler.Scheduler/GetVersion" - Scheduler_RunContainer_FullMethodName = "/scheduler.Scheduler/RunContainer" - Scheduler_StopContainer_FullMethodName = "/scheduler.Scheduler/StopContainer" - Scheduler_SubscribeWorkerEvents_FullMethodName = "/scheduler.Scheduler/SubscribeWorkerEvents" - Scheduler_RegisterWorker_FullMethodName = "/scheduler.Scheduler/RegisterWorker" + Scheduler_GetVersion_FullMethodName = "/scheduler.Scheduler/GetVersion" + Scheduler_RunContainer_FullMethodName = "/scheduler.Scheduler/RunContainer" + Scheduler_StopContainer_FullMethodName = "/scheduler.Scheduler/StopContainer" ) // SchedulerClient is the client API for Scheduler service. @@ -33,9 +31,6 @@ type SchedulerClient interface { GetVersion(ctx context.Context, in *VersionRequest, opts ...grpc.CallOption) (*VersionResponse, error) RunContainer(ctx context.Context, in *RunContainerRequest, opts ...grpc.CallOption) (*RunContainerResponse, error) StopContainer(ctx context.Context, in *StopContainerRequest, opts ...grpc.CallOption) (*StopContainerResponse, error) - // rpc StreamLogs(stream StreamLogRequest) returns (StreamLogResponse); - SubscribeWorkerEvents(ctx context.Context, in *SubscribeWorkerEventRequest, opts ...grpc.CallOption) (Scheduler_SubscribeWorkerEventsClient, error) - RegisterWorker(ctx context.Context, in *RegisterWorkerRequest, opts ...grpc.CallOption) (*RegisterWorkerResponse, error) } type schedulerClient struct { @@ -73,47 +68,6 @@ func (c *schedulerClient) StopContainer(ctx context.Context, in *StopContainerRe return out, nil } -func (c *schedulerClient) SubscribeWorkerEvents(ctx context.Context, in *SubscribeWorkerEventRequest, opts ...grpc.CallOption) (Scheduler_SubscribeWorkerEventsClient, error) { - stream, err := c.cc.NewStream(ctx, &Scheduler_ServiceDesc.Streams[0], Scheduler_SubscribeWorkerEvents_FullMethodName, opts...) - if err != nil { - return nil, err - } - x := &schedulerSubscribeWorkerEventsClient{stream} - if err := x.ClientStream.SendMsg(in); err != nil { - return nil, err - } - if err := x.ClientStream.CloseSend(); err != nil { - return nil, err - } - return x, nil -} - -type Scheduler_SubscribeWorkerEventsClient interface { - Recv() (*SubscribeWorkerEventResponse, error) - grpc.ClientStream -} - -type schedulerSubscribeWorkerEventsClient struct { - grpc.ClientStream -} - -func (x *schedulerSubscribeWorkerEventsClient) Recv() (*SubscribeWorkerEventResponse, error) { - m := new(SubscribeWorkerEventResponse) - if err := x.ClientStream.RecvMsg(m); err != nil { - return nil, err - } - return m, nil -} - -func (c *schedulerClient) RegisterWorker(ctx context.Context, in *RegisterWorkerRequest, opts ...grpc.CallOption) (*RegisterWorkerResponse, error) { - out := new(RegisterWorkerResponse) - err := c.cc.Invoke(ctx, Scheduler_RegisterWorker_FullMethodName, in, out, opts...) - if err != nil { - return nil, err - } - return out, nil -} - // SchedulerServer is the server API for Scheduler service. // All implementations must embed UnimplementedSchedulerServer // for forward compatibility @@ -121,9 +75,6 @@ type SchedulerServer interface { GetVersion(context.Context, *VersionRequest) (*VersionResponse, error) RunContainer(context.Context, *RunContainerRequest) (*RunContainerResponse, error) StopContainer(context.Context, *StopContainerRequest) (*StopContainerResponse, error) - // rpc StreamLogs(stream StreamLogRequest) returns (StreamLogResponse); - SubscribeWorkerEvents(*SubscribeWorkerEventRequest, Scheduler_SubscribeWorkerEventsServer) error - RegisterWorker(context.Context, *RegisterWorkerRequest) (*RegisterWorkerResponse, error) mustEmbedUnimplementedSchedulerServer() } @@ -140,12 +91,6 @@ func (UnimplementedSchedulerServer) RunContainer(context.Context, *RunContainerR func (UnimplementedSchedulerServer) StopContainer(context.Context, *StopContainerRequest) (*StopContainerResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method StopContainer not implemented") } -func (UnimplementedSchedulerServer) SubscribeWorkerEvents(*SubscribeWorkerEventRequest, Scheduler_SubscribeWorkerEventsServer) error { - return status.Errorf(codes.Unimplemented, "method SubscribeWorkerEvents not implemented") -} -func (UnimplementedSchedulerServer) RegisterWorker(context.Context, *RegisterWorkerRequest) (*RegisterWorkerResponse, error) { - return nil, status.Errorf(codes.Unimplemented, "method RegisterWorker not implemented") -} func (UnimplementedSchedulerServer) mustEmbedUnimplementedSchedulerServer() {} // UnsafeSchedulerServer may be embedded to opt out of forward compatibility for this service. @@ -213,45 +158,6 @@ func _Scheduler_StopContainer_Handler(srv interface{}, ctx context.Context, dec return interceptor(ctx, in, info, handler) } -func _Scheduler_SubscribeWorkerEvents_Handler(srv interface{}, stream grpc.ServerStream) error { - m := new(SubscribeWorkerEventRequest) - if err := stream.RecvMsg(m); err != nil { - return err - } - return srv.(SchedulerServer).SubscribeWorkerEvents(m, &schedulerSubscribeWorkerEventsServer{stream}) -} - -type Scheduler_SubscribeWorkerEventsServer interface { - Send(*SubscribeWorkerEventResponse) error - grpc.ServerStream -} - -type schedulerSubscribeWorkerEventsServer struct { - grpc.ServerStream -} - -func (x *schedulerSubscribeWorkerEventsServer) Send(m *SubscribeWorkerEventResponse) error { - return x.ServerStream.SendMsg(m) -} - -func _Scheduler_RegisterWorker_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { - in := new(RegisterWorkerRequest) - if err := dec(in); err != nil { - return nil, err - } - if interceptor == nil { - return srv.(SchedulerServer).RegisterWorker(ctx, in) - } - info := &grpc.UnaryServerInfo{ - Server: srv, - FullMethod: Scheduler_RegisterWorker_FullMethodName, - } - handler := func(ctx context.Context, req interface{}) (interface{}, error) { - return srv.(SchedulerServer).RegisterWorker(ctx, req.(*RegisterWorkerRequest)) - } - return interceptor(ctx, in, info, handler) -} - // Scheduler_ServiceDesc is the grpc.ServiceDesc for Scheduler service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -271,17 +177,7 @@ var Scheduler_ServiceDesc = grpc.ServiceDesc{ MethodName: "StopContainer", Handler: _Scheduler_StopContainer_Handler, }, - { - MethodName: "RegisterWorker", - Handler: _Scheduler_RegisterWorker_Handler, - }, - }, - Streams: []grpc.StreamDesc{ - { - StreamName: "SubscribeWorkerEvents", - Handler: _Scheduler_SubscribeWorkerEvents_Handler, - ServerStreams: true, - }, }, + Streams: []grpc.StreamDesc{}, Metadata: "scheduler.proto", }