Skip to content

Commit

Permalink
Added support for env variables for project and api token (#11)
Browse files Browse the repository at this point in the history
  • Loading branch information
Raalsky authored Aug 9, 2024
1 parent 4c91b15 commit 67f63cb
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 7 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Added support for Forking ([#9](https://github.com/neptune-ai/neptune-client-scale/pull/9))
- Added support for Experiments ([#9](https://github.com/neptune-ai/neptune-client-scale/pull/9))
- Added support for Run resume ([#9](https://github.com/neptune-ai/neptune-client-scale/pull/9))
- Added support for env variables for project and api token ([#11](https://github.com/neptune-ai/neptune-client-scale/pull/11))
31 changes: 24 additions & 7 deletions src/neptune_scale/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

__all__ = ["Run"]

import os
import threading
from contextlib import AbstractContextManager
from datetime import datetime
Expand Down Expand Up @@ -33,6 +34,10 @@
verify_project_qualified_name,
verify_type,
)
from neptune_scale.envs import (
API_TOKEN_ENV_NAME,
PROJECT_ENV_NAME,
)
from neptune_scale.parameters import (
MAX_FAMILY_LENGTH,
MAX_QUEUE_SIZE,
Expand All @@ -48,10 +53,10 @@ class Run(WithResources, AbstractContextManager):
def __init__(
self,
*,
project: str,
api_token: str,
family: str,
run_id: str,
project: str | None = None,
api_token: str | None = None,
resume: bool = False,
as_experiment: str | None = None,
creation_time: datetime | None = None,
Expand All @@ -64,11 +69,13 @@ def __init__(
Initializes a run that logs the model-building metadata to Neptune.
Args:
project: Name of the project where the metadata is logged, in the form `workspace-name/project-name`.
api_token: Your Neptune API token.
family: Identifies related runs. For example, the same value must apply to all runs within a run hierarchy.
Max length: 128 characters.
run_id: Unique identifier of a run. Must be unique within the project. Max length: 128 characters.
project: Name of the project where the metadata is logged, in the form `workspace-name/project-name`.
If not provided, the value of the `NEPTUNE_PROJECT` environment variable is used.
api_token: Your Neptune API token. If not provided, the value of the `NEPTUNE_API_TOKEN` environment
variable is used.
resume: Whether to resume an existing run.
as_experiment: If creating a run as an experiment, ID of an experiment to be associated with the run.
creation_time: Custom creation time of the run.
Expand All @@ -80,10 +87,11 @@ def __init__(
- Maximum size of the queue.
- Exception that made the queue full.
"""
verify_type("api_token", api_token, str)
verify_type("family", family, str)
verify_type("run_id", run_id, str)
verify_type("resume", resume, bool)
verify_type("project", project, (str, type(None)))
verify_type("api_token", api_token, (str, type(None)))
verify_type("as_experiment", as_experiment, (str, type(None)))
verify_type("creation_time", creation_time, (datetime, type(None)))
verify_type("from_run_id", from_run_id, (str, type(None)))
Expand All @@ -102,7 +110,16 @@ def __init__(
if resume and from_step is not None:
raise ValueError("`resume` and `from_step` cannot be used together.")

project = project or os.environ.get(PROJECT_ENV_NAME)
verify_non_empty("project", project)
assert project is not None # mypy
input_project: str = project

api_token = api_token or os.environ.get(API_TOKEN_ENV_NAME)
verify_non_empty("api_token", api_token)
assert api_token is not None # mypy
input_api_token: str = api_token

verify_non_empty("family", family)
verify_non_empty("run_id", run_id)
if as_experiment is not None:
Expand All @@ -115,15 +132,15 @@ def __init__(
verify_max_length("family", family, MAX_FAMILY_LENGTH)
verify_max_length("run_id", run_id, MAX_RUN_ID_LENGTH)

self._project: str = project
self._project: str = input_project
self._family: str = family
self._run_id: str = run_id

self._lock = threading.RLock()
self._operations_queue: OperationsQueue = OperationsQueue(
lock=self._lock, max_size=max_queue_size, max_size_exceeded_callback=max_queue_size_exceeded_callback
)
self._backend: ApiClient = ApiClient(api_token=api_token)
self._backend: ApiClient = ApiClient(api_token=input_api_token)

if not resume:
self._create_run(
Expand Down
3 changes: 3 additions & 0 deletions src/neptune_scale/envs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
PROJECT_ENV_NAME = "NEPTUNE_PROJECT"

API_TOKEN_ENV_NAME = "NEPTUNE_API_TOKEN"

0 comments on commit 67f63cb

Please sign in to comment.