Refactor enums and env vars

polyaxon · Sep 19, 2023 · ade1a7d · ade1a7d
1 parent 1c80ab3
commit ade1a7d
Show file tree

Hide file tree

Showing 36 changed files with 202 additions and 204 deletions.
diff --git a/haupt/haupt/apis/methods/project_resources.py b/haupt/haupt/apis/methods/project_resources.py
@@ -18,7 +18,7 @@
 from haupt.db.managers.bookmarks import bookmark_obj
 from haupt.db.managers.statuses import bulk_new_run_status
 from polyaxon.lifecycle import LifeCycle, LiveState, V1StatusCondition, V1Statuses
-from polyaxon.schemas import V1RunPending
+from polyaxon.polyflow import V1RunPending
 
 
 def create_runs_tags(view, request, *args, **kwargs):

diff --git a/haupt/haupt/apis/serializers/runs.py b/haupt/haupt/apis/serializers/runs.py
@@ -22,8 +22,7 @@
 from polyaxon.exceptions import PolyaxonException
 from polyaxon.lifecycle import ManagedBy
 from polyaxon.polyaxonfile import OperationSpecification
-from polyaxon.polyflow import V1RunEdgeKind
-from polyaxon.schemas import V1RunPending
+from polyaxon.polyflow import V1RunEdgeKind, V1RunPending
 
 
 class BaseRunEdgeSerializer(serializers.ModelSerializer):

diff --git a/haupt/haupt/cli/runners/base.py b/haupt/haupt/cli/runners/base.py
@@ -9,7 +9,7 @@
 
 from haupt.cli.runners.cron import start_cron
 from haupt.cli.runners.manage import migrate
-from polyaxon.env_vars.keys import EV_KEYS_PROXY_LOCAL_PORT
+from polyaxon.env_vars.keys import ENV_KEYS_PROXY_LOCAL_PORT
 
 _logger = logging.getLogger("haupt.cli")
 
@@ -31,7 +31,7 @@ def start_app(
     migrate(migrate_tables=migrate_tables, migrate_db=migrate_db)
     host = host or "0.0.0.0"
     port = int(port or 8000)
-    os.environ[EV_KEYS_PROXY_LOCAL_PORT] = str(port)
+    os.environ[ENV_KEYS_PROXY_LOCAL_PORT] = str(port)
     log_level = log_level or "warning"
     if per_core:
         workers = get_core_workers(per_core=workers or 2, max_workers=max_workers)

diff --git a/haupt/haupt/cli/runners/cron.py b/haupt/haupt/cli/runners/cron.py
@@ -5,8 +5,8 @@
 
 from polyaxon.client import PolyaxonClient
 from polyaxon.constants.globals import DEFAULT, NO_AUTH
-from polyaxon.schemas.cli.client_config import ClientConfig
-from polyaxon.sdk.exceptions import ApiException
+from polyaxon.exceptions import ApiException
+from polyaxon.schemas.client import ClientConfig
 from urllib3.exceptions import HTTPError
 
 _logger = logging.getLogger("haupt.cli.cron")

diff --git a/haupt/haupt/cli/runners/manage.py b/haupt/haupt/cli/runners/manage.py
@@ -7,7 +7,7 @@
 
 from haupt import settings
 from polyaxon.contexts import paths as ctx_paths
-from polyaxon.env_vars.keys import EV_KEYS_SERVICE, EV_KEYS_UI_IN_SANDBOX
+from polyaxon.env_vars.keys import ENV_KEYS_SERVICE, ENV_KEYS_UI_IN_SANDBOX
 from polyaxon.services.values import PolyaxonServices
 
 _logger = logging.getLogger("haupt.cli.manage")
@@ -44,7 +44,7 @@ def migrate(
 
 
 def run_manage(command: str, args: Optional[List[str]] = None):
-    os.environ[EV_KEYS_SERVICE] = PolyaxonServices.API
-    os.environ[EV_KEYS_UI_IN_SANDBOX] = "true"
+    os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.API
+    os.environ[ENV_KEYS_UI_IN_SANDBOX] = "true"
     settings.set_sandbox_config(path=ctx_paths.CONTEXT_ARTIFACTS_ROOT)
     manage(command, args)
diff --git a/haupt/haupt/cli/runners/sandbox.py b/haupt/haupt/cli/runners/sandbox.py
@@ -4,10 +4,9 @@
 
 from haupt import settings
 from haupt.cli.runners.base import start_app
-from haupt.cli.runners.cron import start_cron
 from polyaxon import settings as plx_settings
 from polyaxon.contexts import paths as ctx_paths
-from polyaxon.env_vars.keys import EV_KEYS_SERVICE
+from polyaxon.env_vars.keys import ENV_KEYS_SERVICE
 from polyaxon.services.values import PolyaxonServices
 
 
@@ -19,7 +18,7 @@ def start(
     per_core: bool = False,
     uds: Optional[str] = None,
 ):
-    os.environ[EV_KEYS_SERVICE] = PolyaxonServices.API
+    os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.API
     settings.set_sandbox_config(path=ctx_paths.CONTEXT_ARTIFACTS_ROOT)
     host = host or settings.SANDBOX_CONFIG.host
     port = port or settings.SANDBOX_CONFIG.port

diff --git a/haupt/haupt/cli/runners/streams.py b/haupt/haupt/cli/runners/streams.py
@@ -3,7 +3,7 @@
 from typing import Optional
 
 from haupt.cli.runners.base import start_app
-from polyaxon.env_vars.keys import EV_KEYS_PROXY_STREAMS_TARGET_PORT
+from polyaxon.env_vars.keys import ENV_KEYS_PROXY_STREAMS_TARGET_PORT
 from polyaxon.services.values import PolyaxonServices
 
 
@@ -15,7 +15,7 @@ def start(
     per_core: bool = False,
     uds: Optional[str] = None,
 ):
-    port = port or os.environ.get(EV_KEYS_PROXY_STREAMS_TARGET_PORT)
+    port = port or os.environ.get(ENV_KEYS_PROXY_STREAMS_TARGET_PORT)
     start_app(
         app="haupt.polyconf.asgi.streams:application",
         app_name=PolyaxonServices.STREAMS,

diff --git a/haupt/haupt/cli/runners/viewer.py b/haupt/haupt/cli/runners/viewer.py
@@ -3,7 +3,7 @@
 from typing import Optional
 
 from haupt.cli.runners.base import start_app
-from polyaxon.env_vars.keys import EV_KEYS_PROXY_STREAMS_TARGET_PORT
+from polyaxon.env_vars.keys import ENV_KEYS_PROXY_STREAMS_TARGET_PORT
 from polyaxon.services.values import PolyaxonServices
 
 
@@ -15,7 +15,7 @@ def start(
     per_core: bool = False,
     uds: Optional[str] = None,
 ):
-    port = port or os.environ.get(EV_KEYS_PROXY_STREAMS_TARGET_PORT)
+    port = port or os.environ.get(ENV_KEYS_PROXY_STREAMS_TARGET_PORT)
     start_app(
         app="haupt.polyconf.asgi.viewer:application",
         app_name=PolyaxonServices.STREAMS,

diff --git a/haupt/haupt/cli/sandbox.py b/haupt/haupt/cli/sandbox.py
@@ -42,9 +42,9 @@ def sandbox(
 ):
     """Start a new sandbox session."""
     from haupt.cli.runners.sandbox import start
-    from polyaxon.env_vars.keys import EV_KEYS_SANDBOX_ROOT
+    from polyaxon.env_vars.keys import ENV_KEYS_SANDBOX_ROOT
 
     if path:
-        os.environ[EV_KEYS_SANDBOX_ROOT] = path
+        os.environ[ENV_KEYS_SANDBOX_ROOT] = path
 
     return start(host=host, port=port, workers=workers, per_core=per_core, uds=uds)
diff --git a/haupt/haupt/cli/viewer.py b/haupt/haupt/cli/viewer.py
@@ -88,9 +88,9 @@ def viewer(
 ):
     """Start a new local viewer session."""
     from haupt.cli.runners.viewer import start
-    from polyaxon.env_vars.keys import EV_KEYS_SANDBOX_ROOT
+    from polyaxon.env_vars.keys import ENV_KEYS_SANDBOX_ROOT
 
     if path:
-        os.environ[EV_KEYS_SANDBOX_ROOT] = path
+        os.environ[ENV_KEYS_SANDBOX_ROOT] = path
 
     return start(host=host, port=port, workers=workers, per_core=per_core, uds=uds)
diff --git a/haupt/haupt/common/apis/index/health.py b/haupt/haupt/common/apis/index/health.py
@@ -8,7 +8,7 @@
 from haupt.common.options.registry.installation import ORGANIZATION_KEY
 from haupt.db.managers.dummy_key import get_dummy_key
 from polyaxon.cli.session import get_compatibility
-from polyaxon.schemas.cli.cli_config import CliConfig
+from polyaxon.schemas.cli import CliConfig
 from polyaxon.services.values import PolyaxonServices
 
 

diff --git a/haupt/haupt/db/abstracts/runs.py b/haupt/haupt/db/abstracts/runs.py
@@ -18,8 +18,7 @@
 from haupt.db.abstracts.uid import UuidModel
 from haupt.db.defs import Models
 from polyaxon.lifecycle import ManagedBy, V1Statuses
-from polyaxon.polyflow import V1CloningKind, V1RunKind
-from polyaxon.schemas import V1RunPending
+from polyaxon.polyflow import V1CloningKind, V1RunKind, V1RunPending
 
 
 class BaseRun(

diff --git a/haupt/haupt/db/managers/agents.py b/haupt/haupt/db/managers/agents.py
@@ -22,7 +22,7 @@
     V1Statuses,
 )
 from polyaxon.polyflow import V1RunKind
-from polyaxon.schemas.cli.agent_config import AgentConfig
+from polyaxon.schemas.agent import AgentConfig
 from polyaxon.utils.fqn_utils import get_run_instance
 
 MAX_DELETE_ITEMS = 200

diff --git a/haupt/haupt/db/managers/live_state.py b/haupt/haupt/db/managers/live_state.py
@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import List
 
 from django.conf import settings
 from django.db.models import Q, QuerySet

diff --git a/haupt/haupt/db/managers/runs.py b/haupt/haupt/db/managers/runs.py
@@ -7,8 +7,7 @@
 from haupt.db.abstracts.runs import BaseRun
 from haupt.db.defs import Models
 from polyaxon.lifecycle import LifeCycle, ManagedBy, V1StatusCondition, V1Statuses
-from polyaxon.polyflow import V1CompiledOperation, V1RunKind
-from polyaxon.schemas import V1RunPending
+from polyaxon.polyflow import V1CompiledOperation, V1RunKind, V1RunPending
 
 
 def add_run_contributors(

diff --git a/haupt/haupt/managers/platform.py b/haupt/haupt/managers/platform.py
@@ -5,7 +5,7 @@
 from haupt.schemas.platform_config import PlatformConfig
 from polyaxon.config.manager import ConfigManager
 from polyaxon.config.spec import ConfigSpec
-from polyaxon.env_vars.keys import EV_KEYS_PLATFORM_CONFIG
+from polyaxon.env_vars.keys import ENV_KEYS_PLATFORM_CONFIG
 
 
 class PlatformManager(ConfigManager):
@@ -29,7 +29,7 @@ def base_directory():
         if config_prefix:
             config_module = "{}.{}".format(config_prefix, config_module)
         glob_path = cls.get_global_config_path()
-        platform_config_path = os.getenv(EV_KEYS_PLATFORM_CONFIG) or {}
+        platform_config_path = os.getenv(ENV_KEYS_PLATFORM_CONFIG) or {}
         config_paths = [
             os.environ,
             {

diff --git a/haupt/haupt/orchestration/executor/handlers.py b/haupt/haupt/orchestration/executor/handlers.py
@@ -12,8 +12,7 @@
     META_UPLOAD_ARTIFACTS,
 )
 from polyaxon.lifecycle import LifeCycle, ManagedBy, V1StatusCondition, V1Statuses
-from polyaxon.polyflow import V1RunEdgeKind, V1RunKind
-from polyaxon.schemas import V1RunPending
+from polyaxon.polyflow import V1RunEdgeKind, V1RunKind, V1RunPending
 
 
 class APIHandler:

diff --git a/haupt/haupt/orchestration/operations/service.py b/haupt/haupt/orchestration/operations/service.py
@@ -28,9 +28,9 @@
     V1Operation,
     V1RunEdgeKind,
     V1RunKind,
+    V1RunPending,
     V1ScheduleKind,
 )
-from polyaxon.schemas import V1RunPending
 from polyaxon.schemas.types import V1ArtifactsType
 
 

diff --git a/haupt/haupt/orchestration/scheduler/manager.py b/haupt/haupt/orchestration/scheduler/manager.py
@@ -72,10 +72,10 @@
     V1Optimization,
     V1RunEdgeKind,
     V1RunKind,
+    V1RunPending,
     V1TriggerPolicy,
     dags,
 )
-from polyaxon.schemas import V1RunPending
 from traceml.artifacts import V1ArtifactKind, V1RunArtifact
 
 _logger = logging.getLogger("polyaxon.scheduler")

diff --git a/haupt/haupt/orchestration/scheduler/resolver.py b/haupt/haupt/orchestration/scheduler/resolver.py
@@ -74,9 +74,9 @@
     V1Param,
     V1Plugins,
     V1RunEdgeKind,
+    V1RunPending,
 )
 from polyaxon.polyflow.params import ops_params
-from polyaxon.schemas import V1RunPending
 from polyaxon.schemas.types import V1ArtifactsType
 from polyaxon.utils.fqn_utils import (
     get_project_instance,

diff --git a/haupt/haupt/polyconf/asgi/sandbox.py b/haupt/haupt/polyconf/asgi/sandbox.py
@@ -13,12 +13,12 @@
 
 from haupt import settings
 from polyaxon.contexts import paths as ctx_paths
-from polyaxon.env_vars.keys import EV_KEYS_SERVICE, EV_KEYS_UI_IN_SANDBOX
+from polyaxon.env_vars.keys import ENV_KEYS_SERVICE, ENV_KEYS_UI_IN_SANDBOX
 from polyaxon.services.values import PolyaxonServices
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "haupt.polyconf.settings")
 os.environ.setdefault("ASGI_APPLICATION", "haupt.polyconf.asgi.sandbox.application")
-os.environ[EV_KEYS_SERVICE] = PolyaxonServices.API
-os.environ[EV_KEYS_UI_IN_SANDBOX] = "true"
+os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.API
+os.environ[ENV_KEYS_UI_IN_SANDBOX] = "true"
 settings.set_sandbox_config(path=ctx_paths.CONTEXT_ARTIFACTS_ROOT)
 application = get_asgi_application()
diff --git a/haupt/haupt/polyconf/asgi/server.py b/haupt/haupt/polyconf/asgi/server.py
@@ -11,10 +11,10 @@
 
 from django.core.asgi import get_asgi_application
 
-from polyaxon.env_vars.keys import EV_KEYS_SERVICE
+from polyaxon.env_vars.keys import ENV_KEYS_SERVICE
 from polyaxon.services.values import PolyaxonServices
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "haupt.polyconf.settings")
 os.environ.setdefault("ASGI_APPLICATION", "haupt.polyconf.asgi.server.application")
-os.environ[EV_KEYS_SERVICE] = PolyaxonServices.API
+os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.API
 application = get_asgi_application()
diff --git a/haupt/haupt/polyconf/asgi/streams.py b/haupt/haupt/polyconf/asgi/streams.py
@@ -11,10 +11,10 @@
 
 from django.core.asgi import get_asgi_application
 
-from polyaxon.env_vars.keys import EV_KEYS_SERVICE
+from polyaxon.env_vars.keys import ENV_KEYS_SERVICE
 from polyaxon.services.values import PolyaxonServices
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "haupt.polyconf.settings")
 os.environ.setdefault("ASGI_APPLICATION", "haupt.polyconf.asgi.streams.application")
-os.environ[EV_KEYS_SERVICE] = PolyaxonServices.STREAMS
+os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.STREAMS
 application = get_asgi_application()
diff --git a/haupt/haupt/polyconf/asgi/viewer.py b/haupt/haupt/polyconf/asgi/viewer.py
@@ -14,16 +14,16 @@
 from haupt import settings
 from polyaxon.contexts import paths as ctx_paths
 from polyaxon.env_vars.keys import (
-    EV_KEYS_SANDBOX_IS_LOCAL,
-    EV_KEYS_SERVICE,
-    EV_KEYS_UI_IN_SANDBOX,
+    ENV_KEYS_SANDBOX_IS_LOCAL,
+    ENV_KEYS_SERVICE,
+    ENV_KEYS_UI_IN_SANDBOX,
 )
 from polyaxon.services.values import PolyaxonServices
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "haupt.polyconf.settings")
 os.environ.setdefault("ASGI_APPLICATION", "haupt.polyconf.asgi.viewer.application")
-os.environ[EV_KEYS_SERVICE] = PolyaxonServices.STREAMS
-os.environ[EV_KEYS_UI_IN_SANDBOX] = "true"
-os.environ[EV_KEYS_SANDBOX_IS_LOCAL] = "true"
+os.environ[ENV_KEYS_SERVICE] = PolyaxonServices.STREAMS
+os.environ[ENV_KEYS_UI_IN_SANDBOX] = "true"
+os.environ[ENV_KEYS_SANDBOX_IS_LOCAL] = "true"
 settings.set_sandbox_config(path=ctx_paths.CONTEXT_OFFLINE_ROOT)
 application = get_asgi_application()