Restructure Python application and library resources#611
Restructure Python application and library resources#611forstmeier wants to merge 1 commit intomasterfrom
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughMajor refactor migrating services from application/* to applications/* and introducing a shared internal library (libraries/python). Large removal of predictionengine, positionmanager, and legacy datamanager code/tests, with new datasets/models/TFT implementations and scripts under libraries/python and applications/models. CI workflows renamed, infrastructure env/image paths updated, and workspace/tooling consolidated. Changes
Sequence Diagram(s)sequenceDiagram
participant C as Client
participant D as Datamanager API
participant A as Alpaca API
C->>D: GET /calendar?start&end
D->>A: get_calendar(start, end)
A-->>D: Calendar data
D-->>C: 200 JSON or 404/500
sequenceDiagram
participant O as Orchestrator
participant M as Models Scripts
participant L as Internal Lib
O->>M: train_tft_model()
M->>M: read_local_data()
M->>L: build dataset/model (Parameters, TFT)
M->>M: train_model()/validate_model()
M->>M: save_model()
M-->>O: workflow complete
Estimated code review effort🎯 5 (Critical) | ⏱️ ~90+ minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
forstmeier
force-pushed
the
07-28-restructure_python_application_and_library_resources
branch
2 times, most recently
from
8765605 to
8460289
Compare
forstmeier
force-pushed
the
07-28-restructure_python_application_and_library_resources
branch
7 times, most recently
from
e675d1a to
17c9a34
Compare
There was a problem hiding this comment.
Pull Request Overview
This PR restructures the Python application architecture by creating an internal library for shared resources and reorganizing the project structure. It updates configuration files to properly support both applications and libraries under the uv package manager, while consolidating common functionality into a centralized library.
- Creates an internal Python library to consolidate shared resources across applications
- Restructures project layout from
application/toapplications/with proper uv workspace configuration - Adds a new
modelsapplication for machine learning workflows with TFT model training capabilities
Reviewed Changes
Copilot reviewed 62 out of 81 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| pyproject.toml | Updates workspace configuration to support applications and libraries structure |
| libraries/python/ | Creates new internal library with shared models, utilities, and ML components |
| applications/ | Restructures and updates all service applications to use the internal library |
| workflows/ | Removes standalone workflow package in favor of models application |
| infrastructure/ | Updates path references to match new applications directory structure |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 47
🔭 Outside diff range comments (8)
applications/positionmanager/src/positionmanager/portfolio.py (3)
11-16: Min/max tickers ctor parameters are unusedThe initializer captures minimum_portfolio_tickers and maximum_portfolio_tickers, but they aren’t enforced in get_optimized_portfolio. Either implement enforcement or remove the parameters to avoid confusion.
Options:
- Enforce max by zeroing smallest weights until non-zero count ≤ maximum_portfolio_tickers.
- Enforce min by boosting weights of top-N assets or adding a post-optimization selection heuristic.
I can provide a small post-processing helper that trims weights to top-N by absolute weight and re-normalizes.
22-23: Validate prediction_weight ∈ [0, 1]Guard against invalid values to avoid unintentionally flipping the sign or overweighting predictions.
Suggested change:
prediction_weight: float = 0.3, ) -> dict[str, int]: - converted_data = historical_data.to_pandas() + if not 0.0 <= prediction_weight <= 1.0: + raise ValueError("prediction_weight must be between 0 and 1 inclusive") + converted_data = historical_data.to_pandas()
54-55: Avoid truncating portfolio value to an intCasting Money to float and then to int truncates cents, systematically under-allocating capital. DiscreteAllocation accepts floats; preserve precision.
- total_portfolio_value=int(float(portfolio_value)), + total_portfolio_value=float(portfolio_value),pyproject.toml (1)
17-28: Pytest rootdir points to a non-existent path; likely breaks test discoverySetting --rootdir=/tests while tests live under applications/*/tests and libraries/python/tests can cause test discovery issues.
Consider removing rootdir or pointing it to the repo root.
addopts = [ "--verbose", "--tb=short", "--strict-markers", "--color=yes", - "--rootdir=/tests", ]libraries/python/src/internal/dates.py (1)
38-45: Type annotations in validator should use datetime.date, not datetime.datetime
startandendaredatetime.date, but the validator type hints aredatetime.datetime. This is misleading and could confuse IDEs/type-checkers.Apply this diff:
- def check_end_after_start( - cls, - end_value: datetime.datetime, - info: core_schema.ValidationInfo, - ) -> datetime.datetime: + def check_end_after_start( + cls, + end_value: datetime.date, + info: core_schema.ValidationInfo, + ) -> datetime.date:applications/datamanager/src/datamanager/main.py (1)
179-191: Content-Type mismatch: writing Arrow stream but declaring file; fix media_type and Content-Disposition.You’re using pa.ipc.RecordBatchStreamWriter (stream format) but returning media_type="application/vnd.apache.arrow.file". Use the stream media type and standard Content-Disposition formatting.
Apply this diff:
- content_disposition = f"attachment; {filename=}" + content_disposition = f'attachment; filename="{filename}"' ... - return Response( - content=sink.getvalue().to_pybytes(), - media_type="application/vnd.apache.arrow.file", + return Response( + content=sink.getvalue().to_pybytes(), + media_type="application/vnd.apache.arrow.stream", headers={ "Content-Disposition": content_disposition, "X-Row-Count": str(data.num_rows), "X-Start-Date": str(start_date), "X-End-Date": str(end_date), }, )applications/positionmanager/tests/test_positionmanager_main.py (1)
88-106: Handle HTTPException in /positions/open so errors return CloudEventVerified: open_position returns CloudEvent for several exceptions but does NOT catch fastapi.HTTPException. The tests set AlpacaClient.get_cash_balance to raise HTTPException, which currently bubbles up (resulting in a 500) — that's why the test expects a 500. To make error handling consistent, convert HTTPException into a CloudEvent error.
Actionable changes:
applications/positionmanager/src/positionmanager/main.py
- In open_position (around @application.post("/positions/open")), include HTTPException in the except clause that handles get_cash_balance (or add a dedicated except HTTPException to return create_cloud_event_error). Example:
except (HTTPException, requests.RequestException, APIError, ValidationError) as e:
return create_cloud_event_error(...)applications/positionmanager/tests/test_positionmanager_main.py
- Update test_open_position_alpaca_error to assert the standardized CloudEvent error response (and status code consistent with your contract) instead of asserting a raw 500 and HTTPException detail.
Reasoning (short): current behavior is intentional given the code, but inconsistent with the rest of the endpoint's CloudEvent error returns — adding HTTPException handling (or a global handler) will align runtime behavior and tests.
applications/positionmanager/src/positionmanager/main.py (1)
175-183: Guard against missing percentile_50 in predictions to avoid KeyError.Use .get and filter only valid entries.
Apply this diff:
- predictions_percentile_50 = { - key: value["percentile_50"] for key, value in predictions.items() - } + predictions_percentile_50 = { + key: value.get("percentile_50") + for key, value in predictions.items() + if isinstance(value, dict) and "percentile_50" in value + }
♻️ Duplicate comments (4)
libraries/python/src/internal/loss_functions.py (1)
20-26: Bug: add() result is unused; loss accumulation is a no-op without assignmentIn tinygrad most ops return new tensors. errors_total.add(...) does not mutate in place; you need to assign back. This matches a prior review comment.
- errors_total.add( + errors_total = errors_total.add( Tensor.where( error > 0, quantile_tensor.mul(error), (quantile_tensor.sub(1)).mul(error), ).mean() )applications/models/pyproject.toml (1)
7-19: Remove duplicate 'internal==0.1.0' dependencyThe 'internal' dependency is listed twice (Lines 7 and 13). Keep a single entry to avoid resolver confusion.
Apply this diff:
dependencies = [ - "internal==0.1.0", + "internal==0.1.0", @@ - "polygon-api-client>=1.14.6", - "internal==0.1.0", + "polygon-api-client>=1.14.6",Additionally, since you’re using a UV workspace source below, you can depend on the workspace package without pinning a version to avoid conflicts between the lock and local workspace: ```diff -[tool.uv.sources] -internal = { workspace = true } +[tool.uv.sources] +internal = { workspace = true }And update the dependency to reference the workspace:
- "internal==0.1.0", + "internal",If you prefer to keep the exact version pin for releases, do it in non-workspace builds only (e.g., via an env-driven dependency override).
libraries/python/src/internal/lstm_network.py (1)
80-80: Tensor.stack likely fails for a single-timestep sequenceTinygrad’s Tensor.stack generally expects multiple tensors. The previous code handled the single-output case via unsqueeze. Restore that branch to prevent runtime errors when sequence_length == 1.
Apply this diff:
- output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) + if len(outputs) == 1: + output_tensor = outputs[0].unsqueeze(dim=1) + else: + output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1)applications/datamanager/src/datamanager/main.py (1)
265-281: Broad exception mapped to 404; return 500 on listing error and clarify log message.An exception when listing S3 objects is not “not found”; it’s an internal error. Also, error message should reflect the actual operation.
Apply this diff:
- try: - blobs = list(s3_client.list_objects(prefix=prefix)) - except Exception as e: # noqa: BLE001 - logger.error(f"Error listing S3 objects: {e}") - return Response( - status_code=status.HTTP_404_NOT_FOUND, - content=json.dumps({"error": "No equity bars found"}), - media_type="application/json", - ) + try: + blobs = list(s3_client.list_objects(prefix=prefix)) + except Exception as e: # noqa: BLE001 + logger.error(f"Error listing equity bars: {e}") + return Response( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content=json.dumps({"error": "Failed to list equity bars"}), + media_type="application/json", + )
🧹 Nitpick comments (47)
applications/positionmanager/src/positionmanager/portfolio.py (3)
26-31: Indexing date column: ensure proper dtypeSetting the index to "date" is correct; consider ensuring a datetime index for robustness with time-based operations.
Optional:
- if "date" in converted_data.columns: - converted_data = converted_data.set_index("date") + if "date" in converted_data.columns: + converted_data["date"] = pd.to_datetime(converted_data["date"], errors="coerce") + converted_data = converted_data.set_index("date").dropna(axis=0, how="any")
46-47: Make risk-free rate configurableHardcoding 2% limits reuse across environments. Consider an argument or class attribute.
- efficient_frontier.max_sharpe(risk_free_rate=0.02) # 2% risk-free rate + efficient_frontier.max_sharpe(risk_free_rate=getattr(self, "risk_free_rate", 0.02))
57-59: Consider lp_portfolio for better discrete allocationgreedy_portfolio is fast but suboptimal. If runtime allows (and solver deps are present), lp_portfolio can yield better allocations and help adhere to constraints (e.g., bounds).
Example:
optimized_portfolio_ticker_share_counts, _ = discrete_allocation.lp_portfolio()Keep greedy as a fallback if solver unavailable.
applications/positionmanager/pyproject.toml (1)
5-5: Exact Python pin may be too strict for deployabilityPinning requires-python to an exact micro version (==3.12.10) can block environments that provide 3.12.11+ security releases. Consider a compatible range:
-requires-python = "==3.12.10" +requires-python = ">=3.12,<3.13"If your infra tooling enforces 3.12.10 specifically (e.g., via container tags), feel free to keep it exact.
libraries/python/pyproject.toml (1)
1-4: Optional: add a build backend for non-uv environmentsIf anyone installs/builds this package outside uv (e.g., plain pip, IDEs, CI builders lacking uv), consider adding a PEP 517 build backend. Example with hatchling:
+[build-system] +requires = ["hatchling>=1.25.0"] +build-backend = "hatchling.build"This won’t affect uv but improves portability.
applications/predictionengine/pyproject.toml (1)
5-5: Avoid exact patch pin for Python; use a compatible rangeExact patch pin in
requires-pythonis brittle and may block installs in equivalent 3.12.x environments (e.g., base images differing by patch). A typical pattern is “>=3.12,<3.13”.-requires-python = "==3.12.10" +requires-python = ">=3.12,<3.13"infrastructure/images.py (1)
15-17: Also validate Dockerfile presence to fail fast with actionable errorDirectory existence is necessary but not sufficient. Checking for a Dockerfile avoids later opaque build errors.
if not service_directory.exists(): message = f"Service directory not found: {service_directory}" raise FileNotFoundError(message) + + dockerfile = service_directory / "Dockerfile" + if not dockerfile.exists(): + raise FileNotFoundError(f"Dockerfile not found: {dockerfile}").github/workflows/teardown_infrastructure.yaml (1)
5-5: Clarify time zone in the comment for cronGitHub cron uses UTC. The comment currently says “EST” which is misleading during daylight savings.
- - cron: '0 23 * * 1,2,3,4,5' # teardown at 6:00 PM EST + - cron: '0 23 * * 1,2,3,4,5' # 23:00 UTC (6pm ET standard time, 7pm ET daylight time).github/workflows/launch_infrastructure.yaml (2)
2-8: Clarify cron timezone in comment to avoid confusionGitHub Actions cron uses UTC. The comment “8:00 AM EST” will be incorrect during daylight saving time (it’ll run 9:00 AM ET). Suggest updating the comment to explicitly note UTC.
- - cron: '0 13 * * 1,2,3,4,5' # launch at 8:00 AM EST + - cron: '0 13 * * 1,2,3,4,5' # 13:00 UTC (8 AM ET standard time / 9 AM ET daylight time)
7-15: Avoid overlapping scheduled runs; add job-level concurrency and minimal permissionsPrevent concurrent “launch” jobs from stepping on each other (e.g., if a previous run overruns) and minimize token scope.
jobs: launch_infrastructure: + permissions: + contents: read name: Launch infrastructure on weekday schedule runs-on: ubuntu-latest environment: pulumi + concurrency: + group: launch_infrastructure-${{ github.ref }} + cancel-in-progress: truepyproject.toml (1)
5-5: Exact Python version pin reduces portabilityrequires-python ==3.12.10 is extremely strict. Unless you require that exact patch version, consider a compatible range (e.g., >=3.12,<3.13) to ease local/CI/environment differences.
-requires-python = "==3.12.10" +requires-python = ">=3.12,<3.13"libraries/python/src/internal/loss_functions.py (1)
16-28: Optional: Vectorize across quantiles to reduce Python overheadYou can compute the loss for all quantiles in one pass by broadcasting targets and quantiles; this is simpler and faster.
- errors_total = Tensor.zeros((1,)) - for index, quantile in enumerate(quantiles): - error = targets.sub(predictions[:, :, index]) - quantile_tensor = Tensor([quantile]) - errors_total = errors_total.add( - Tensor.where( - error > 0, - quantile_tensor.mul(error), - (quantile_tensor.sub(1)).mul(error), - ).mean() - ) - - return errors_total.div(len(quantiles)) # shape: (1,) + # Broadcast targets to (B, O, K) and quantiles to (1, 1, K) + errors = targets[:, :, None].sub(predictions) + q = Tensor(quantiles)[None, None, :] + losses = Tensor.where(errors > 0, q.mul(errors), (q.sub(1)).mul(errors)) + # Mean over batch and output dims, then average over quantiles + return losses.mean(axis=(0, 1)).mean(keepdims=True) # shape: (1,)libraries/python/tests/test_loss_functions.py (1)
32-39: Strengthen assertion for perfect predictionFor zero error, quantile loss should be exactly 0 within numerical tolerance. Tighten the check to assert near-zero.
- assert loss.numpy() >= 0.0 + assert abs(float(loss.numpy())) < 1e-6applications/positionmanager/src/positionmanager/clients.py (3)
120-123: Prefer using DateRange.to_object() for consistency and encapsulationYou already expose a helper to serialize DateRange; use it to avoid duplication and future drift.
- params = { - "start_date": date_range.start.isoformat(), - "end_date": date_range.end.isoformat(), - } + params = date_range.to_object()
126-139: HTTP ergonomics: use raise_for_status and explicit Accept header (optional)
- You can lean on response.raise_for_status() for non-204 cases to simplify error handling.
- If the service returns Arrow stream, advertise it explicitly to guard against content-type drift.
- try: - response = requests.get(endpoint, params=params, timeout=30) + headers = {"Accept": "application/vnd.apache.arrow.stream"} + try: + response = requests.get(endpoint, params=params, headers=headers, timeout=(5, 30)) except requests.RequestException as err: message = f"Data manager service call error: {err}" raise RuntimeError(message) from err if response.status_code == requests.codes["no_content"]: return pl.DataFrame() - if response.status_code != requests.codes["ok"]: - message = f"Data service error: {response.text}, status code: {response.status_code}" # noqa: E501 - raise requests.HTTPError( - message, - response=response, - ) + response.raise_for_status()
144-154: Minor: avoid redundant DataFrame wrapping and favor simpler conversionpl.from_arrow already returns a Polars DataFrame; no need to re-wrap.
- data = pl.DataFrame(pl.from_arrow(table)) + data = pl.from_arrow(table) data = data.with_columns( pl.col("datetime").cast(pl.Datetime).dt.date().alias("date") )libraries/python/src/internal/cloud_event.py (1)
12-21: Optional: accept timezone as a parameter for consistency in non-ET contextsIf these helpers are used outside ET, hardcoding America/New_York can be surprising. Consider adding an optional tz parameter defaulting to ET.
I can draft a small refactor to thread an optional tz: ZoneInfo | None parameter if you want to make this configurable.
Also applies to: 29-38
applications/models/pyproject.toml (2)
5-5: Consider broadening requires-python to allow patch updatesPinning to an exact patch version (==3.12.10) can cause churn and friction in CI images. Unless there’s a known incompatibility, prefer a compatible range like >=3.12,<3.13.
Apply this diff:
-requires-python = "==3.12.10" +requires-python = ">=3.12,<3.13"
8-9: Avoid explicitly depending on botocore alongside boto3boto3 manages its botocore pin. Declaring both can create resolver conflicts. Remove botocore unless you truly use it directly.
Apply this diff:
- "botocore>=1.38.23",applications/models/src/models/get_alpaca_equity_bars.py (1)
55-72: Defensive parsing: optional fields may be missing in barsvw can be missing; using direct indexing can raise KeyError. Prefer .get with default None.
Apply this diff:
- "volume_weighted_average_price": bar["vw"], + "volume_weighted_average_price": bar.get("vw"),libraries/python/src/internal/lstm_network.py (1)
24-30: Avoid shadowing the input_size parameter inside the layer loopShadowing makes the code harder to read. Use a local name for the per-layer input size.
Apply this diff:
- for index in range(layer_count): - input_size = input_size if index == 0 else self.hidden_size + for index in range(layer_count): + layer_input_size = input_size if index == 0 else self.hidden_size self.layers.append( LSTMCell( - input_size=input_size, + input_size=layer_input_size, hidden_size=self.hidden_size, ) )libraries/python/tests/test_lstm_network.py (1)
12-17: Verify the minimum layer_count enforcement in LSTM class.The test correctly initializes LSTM with
layer_count=3, which aligns with the internal implementation that enforces a minimum of 3 layers. The test should also verify that attempting to initialize with fewer layers raises a ValueError.Would you like me to generate an additional test case to verify that the LSTM class properly enforces the minimum layer count requirement?
libraries/python/src/internal/dataset.py (2)
10-26: Add type hints for Scaler attributes.The Scaler class is missing type hints for its attributes
meansandstandard_deviations. Consider adding them for better type safety.class Scaler: def __init__(self) -> None: - pass + self.means: pl.DataFrame | None = None + self.standard_deviations: pl.DataFrame | None = None
116-117: Avoid using boolean literals with noqa comments.The code uses
# noqa: FBT003to suppress linting warnings about boolean literals. Consider using variables for better readability.+ holiday_true = True + holiday_false = False data = ( data.with_columns( pl.col("datetime").dt.weekday().alias("temporary_weekday") ) .with_columns( pl.when( pl.col("is_holiday").is_null() & (pl.col("temporary_weekday") <= friday_number) ) - .then(True) # noqa: FBT003 + .then(holiday_true) .when( pl.col("is_holiday").is_null() & (pl.col("temporary_weekday") > friday_number) ) - .then(False) # noqa: FBT003 + .then(holiday_false)Also applies to: 121-122
libraries/python/tests/test_dataset.py (1)
81-134: Add validation for missing "targets" key in predict mode.The test correctly validates batch structure for prediction mode. Consider also testing that the "targets" key is NOT present in predict mode batches, as it should only be included for train/validate modes.
for batch in batches: assert "encoder_categorical_features" in batch assert "encoder_continuous_features" in batch assert "decoder_categorical_features" in batch assert "static_categorical_features" in batch + # Verify targets are not included in predict mode + assert "targets" not in batchDockerfile.tests (1)
30-30: Potentially unnecessary early copy of sourcesUnless uv needs the full library source to resolve workspace/local path deps during sync, keep full
libraries/copy afteruv syncto avoid cache invalidation on code-only changes.libraries/python/tests/test_mhsa_network.py (3)
9-12: Nit: reuse variables instead of hardcoded literalsPrefer passing
heads_countandembedding_sizevariables to avoid divergence if they change.- attention = MultiHeadSelfAttentionNetwork(heads_count=8, embedding_size=64) + attention = MultiHeadSelfAttentionNetwork( + heads_count=heads_count, embedding_size=embedding_size + )
31-36: Add a negative test for invalid configuration (embedding not divisible by heads)The network should reject incompatible configurations (e.g., embedding_size not divisible by heads_count). Add a test to assert a ValueError (or the specific exception you raise).
Example test to add:
def test_multi_head_attention_invalid_heads_divisibility() -> None: # embedding_size=30 is not divisible by heads_count=8 with pytest.raises(ValueError): _ = MultiHeadSelfAttentionNetwork(heads_count=8, embedding_size=30)
17-22: Optional: seed RNG for determinism (if future assertions go beyond shapes)Currently you only assert shapes, so nondeterminism is fine. If you later assert on values, seed both numpy and tinygrad RNGs to prevent flakes.
libraries/python/src/internal/dates.py (1)
8-31: Consider Pydantic v2 serializers instead of json_encodersIn Pydantic v2,
field_serializeris the preferred way to control JSON serialization.json_encodersmay continue to work, but moving to serializers is more idiomatic and explicit.Example:
class Date(BaseModel): @@ - model_config = {"json_encoders": {datetime.date: lambda d: d.strftime("%Y/%m/%d")}} + @pydantic.field_serializer("date", when_used="json") + def serialize_date(self, value: datetime.date) -> str: + return value.strftime("%Y/%m/%d")Note: add
from pydantic import field_serializerimport or reference viapydantic.field_serializer.libraries/python/src/internal/variable_selection_network.py (1)
1-22: Confirm parameter registration and consider normalizationShort summary: tft_model assigns the two VSNs to self and calls get_parameters(self) — please verify that tinygrad.nn.state.get_parameters actually traverses object attributes and includes parameters from nested nn.Linear instances. I couldn't import tinygrad here to inspect get_parameters, so this is inconclusive from the sandbox.
Points to check:
- libraries/python/src/internal/tft_model.py
- self.encoder_variable_selection_network = VariableSelectionNetwork(...)
- self.decoder_variable_selection_network = VariableSelectionNetwork(...)
- self.parameters = get_parameters(self)
=> Ensure get_parameters(self) includes parameters from those VSN attributes.- libraries/python/src/internal/variable_selection_network.py
- VariableSelectionNetwork.forward currently returns inputs.sigmoid() — if you want a TFT-style attention/gating across features, consider normalizing across the feature dimension (softmax) instead of per-feature sigmoid.
Suggested diff (optional — only if tinygrad.Tensor.softmax supports dim):
def forward(self, inputs: Tensor) -> Tensor: inputs = self.input_layer(inputs) inputs = inputs.relu() inputs = self.output_layer(inputs) - return inputs.sigmoid() + # Normalize across the feature dimension for simplex-style weights + return inputs.softmax(dim=-1) # verify tinygrad's softmax signatureQuick local verification you can run:
- Inspect get_parameters source in your environment:
python - <<'PY'
import inspect, tinygrad.nn.state as s
print(inspect.getsource(s.get_parameters))
PY- Or check whether get_parameters(self) actually contains parameters from self.encoder_variable_selection_network and self.decoder_variable_selection_network at runtime.
If get_parameters does not collect nested attributes, explicitly aggregate VSN parameters (e.g., extend the parameter list with get_parameters(self.encoder_variable_selection_network) and get_parameters(self.decoder_variable_selection_network)).
applications/datamanager/tests/test_datamanager_main.py (2)
21-33: Replacements from SummaryDate to Date are correct; add a JSON serialization test to preserve behaviorThe basics are covered. Add a test to ensure the JSON encoder for
Datestill outputsYYYY/MM/DDas previously asserted forSummaryDate.Example:
def test_date_json_encoder() -> None: d = Date(date=date(2023, 5, 15)) json_data = d.model_dump(mode="json") assert json_data["date"] == "2023/05/15"
156-175: S3Client patching OK; consider adding a test for empty key listing vs missing path semanticsYou handle the "not found" case where
list_objectsreturns empty. Consider also testing whendata_bucket_nameor the path prefix is misconfigured to ensure the endpoint fails fast with a clear error.libraries/python/tests/test_equity_bar.py (2)
133-146: Align test name with behavior: pass a string to exercise the timestamp parserThe test name says “string iso format” but the test passes a date object. Pass a string to cover the validator’s string parsing path.
Apply this diff:
- equity_bar = EquityBar( - ticker="AAPL", - timestamp=date.fromisoformat("2023-06-15"), + equity_bar = EquityBar( + ticker="AAPL", + timestamp="2023-06-15", open_price=150.0, high_price=155.0, low_price=149.0, close_price=153.0, volume=1000000, volume_weighted_average_price=152.5, )
67-81: Message substring assertions look good; consider adding a malformed-date testAsserting substrings keeps this resilient to Pydantic’s error wrappers. Consider adding a test that an invalid date string raises with the expected message to fully cover
validate_timestamp.Example to add elsewhere in this file:
def test_equity_bar_invalid_timestamp_string() -> None: with pytest.raises(ValidationError) as exc_info: EquityBar( ticker="AAPL", timestamp="2023/15/01", # invalid format for the model open_price=150.0, high_price=155.0, low_price=149.0, close_price=153.0, volume=1000000, volume_weighted_average_price=152.5, ) assert "Invalid date format" in str(exc_info.value)libraries/python/tests/test_dates.py (1)
29-33: Use a string input to cover the string-parsing path (test name says “string”)This test currently passes a date object. Pass a string to match the name and to exercise the format parsing logic.
Apply this diff:
-def test_date_string_dash_format() -> None: - date_instance = Date(date=datetime.date.fromisoformat("2023-03-15")) +def test_date_string_dash_format() -> None: + date_instance = Date(date="2023-03-15") # type: ignorelibraries/python/src/internal/mhsa_network.py (1)
61-64: Optional: expose a training flag to control dropoutConsider a
self.trainingflag (and/or a forward arg) to toggle dropout deterministically in tests/inference without changing the rate.I can sketch a minimal training/eval toggle if you want it in this PR.
libraries/python/tests/test_variable_selection_network.py (2)
6-7: Seed RNG for determinism across runsSeeding avoids sporadic failures in environments that analyze exact numeric thresholds or aggregate stats.
Apply this diff:
-rng = Generator(PCG64()) +rng = Generator(PCG64(0))
18-23: Weight-shape assertions may be brittle against tinygrad changestinygrad.nn.Linear’s internal weight layout can vary between (out, in) and (in, out) across versions. These checks may break on upgrades.
Consider replacing weight-shape assertions with behavioral checks (e.g., feed (B, input_dimension) and assert output shape is (B, input_dimension), which you already do below), or gate these asserts behind a detected layout:
# example approach (no diff applied): wshape = vsn.input_layer.weight.shape assert input_dimension in wshape and hidden_size in wshapeWould you like me to adjust the tests to be layout-agnostic?
applications/datamanager/src/datamanager/main.py (2)
255-261: Avoid overriding CloudEvent helper’s timestamp; remove duplicate “date”.create_cloud_event_success already injects a timestamp in data. Passing your own date key overrides it. If overriding is not intentional, drop the custom date field for consistency.
Apply this diff:
- data={ - "date": request_fetch_date, - "count": count, - }, + data={"count": count},
294-333: Prefer specific exception handling for calendar endpoint.Catching bare Exception hides actionable errors (API/HTTP/validation) and makes observability harder. Handle known exceptions and return 500 with clearer messages.
Apply this diff:
- except Exception as e: # noqa: BLE001 + except (requests.RequestException, Exception) as e: # replace Exception with specific Alpaca exceptions if available logger.error(f"Error fetching calendar data: {e}") logger.error(traceback.format_exc()) return Response( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=json.dumps({"error": f"Failed to fetch calendar data: {e!s}"}), media_type="application/json", )applications/positionmanager/tests/test_positionmanager_main.py (1)
21-24: Patch target for requests.get is unnecessary here.DataClient.get_data is already mocked, so requests.get inside the client won’t be called. This patch can be removed to reduce noise.
Apply this diff:
- @patch("applications.positionmanager.src.positionmanager.main.requests.get") def test_open_position_success( self, - mock_requests_get: MagicMock, MockDataClient: MagicMock, # noqa: N803 MockAlpacaClient: MagicMock, # noqa: N803 ) -> None:libraries/python/src/internal/equity_bar.py (1)
40-49: Accepting only YYYY-MM-DD strings is fine; mention it in the field description or reuse internal Date parser.Current validator is strict ISO date. If you want to accept “YYYY/MM/DD” like internal.Date, reuse that parser for consistency; otherwise, clarify the format in the docstring.
Apply this diff to clarify the description:
- timestamp: date = Field(..., description="The date of the bar in YYYY-MM-DD format") + timestamp: date = Field(..., description="The date of the bar in YYYY-MM-DD format (ISO 8601)")applications/models/src/models/train_tft_model.py (1)
143-165: Hard-coded training CSV path; consider parameterizing for portability.Embedding a repo-relative path reduces reusability in Flyte. Expose filepath as a workflow input or env var.
Apply this diff to accept a filepath:
-@workflow -def train_tft_model() -> None: - dataset = read_local_data( - filepath="applications/models/src/models/training_data.csv" - ) # type: ignore[assignment] +@workflow +def train_tft_model(filepath: str = "applications/models/src/models/training_data.csv") -> None: + dataset = read_local_data(filepath=filepath) # type: ignore[assignment]applications/positionmanager/src/positionmanager/main.py (1)
245-256: Avoid duplicating timestamp field; create_cloud_event_success already adds one.Remove redundant "date" key from data to keep payloads consistent across apps.
Apply this diff:
return create_cloud_event_success( application_name="positionmanager", event_metadata=["positions", "open"], data={ - "date": date_range.end.isoformat(), "initial_cash_balance": float(cash_balance), "final_cash_balance": float(final_cash_balance), "optimized_portfolio": optimized_portfolio, "executed_trades": executed_trades, "time_period": date_range.to_object(), }, )libraries/python/src/internal/tft_model.py (2)
85-90: Nit: parameter naming (attention_head_size)This parameter represents a count, not a size. Consider renaming to
attention_head_countinParametersand here for clarity and consistency withMultiHeadSelfAttentionNetwork(heads_count=...). This is a naming-only suggestion; not a functional issue.
247-260: State serialization LGTMThe save/load APIs are consistent with tinygrad’s state utils. Optional: check return value of
load_state_dictfor diagnostics.If desired, capture and log the return for visibility:
- states = safe_load(path_and_file) - load_state_dict(self, states) + states = safe_load(path_and_file) + _ = load_state_dict(self, states)
| self.hidden_size = parameters.hidden_size | ||
| self.batch_size = parameters.input_length | ||
| self.output_length = parameters.output_length | ||
| self.output_size = parameters.output_size | ||
| self.quantiles = parameters.quantiles | ||
|
|
There was a problem hiding this comment.
self.batch_size is incorrectly set to a sequence length (input_length)
Batch size is dynamic per batch and should be inferred from inputs in forward. Setting it to parameters.input_length causes downstream shape bugs (decoder zeros shape, static context view, context expansion, quantiles reshape).
Remove the incorrect attribute and compute batch_size and sequence lengths inside forward.
- self.batch_size = parameters.input_lengthCommittable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 39 to 44, remove the
incorrect assignment self.batch_size = parameters.input_length and any other
places that rely on this static attribute; instead compute batch_size and
sequence lengths inside forward by reading them from the input tensors (e.g.
batch_size = x.size(0), input_length = x.size(1) or equivalent), update decoder
zero tensors, static context view expansions, context expansion logic, and
quantiles reshape calls to use these local values, and delete references to
self.batch_size throughout the class so shapes are derived at runtime per-batch.
| self.output_layer = Linear( | ||
| in_features=parameters.hidden_size, | ||
| out_features=parameters.output_size, | ||
| ) |
There was a problem hiding this comment.
Output layer must emit per-quantile outputs
You later reshape to include a quantile dimension but the layer currently produces only output_size units. This will make the reshape invalid or silently wrong. Make the output size output_size * len(quantiles).
- self.output_layer = Linear(
- in_features=parameters.hidden_size,
- out_features=parameters.output_size,
- )
+ self.output_layer = Linear(
+ in_features=parameters.hidden_size,
+ out_features=parameters.output_size * len(parameters.quantiles),
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self.output_layer = Linear( | |
| in_features=parameters.hidden_size, | |
| out_features=parameters.output_size, | |
| ) | |
| self.output_layer = Linear( | |
| in_features=parameters.hidden_size, | |
| out_features=parameters.output_size * len(parameters.quantiles), | |
| ) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 96 to 99, the output
Linear layer currently uses out_features=parameters.output_size but downstream
code reshapes outputs to include a quantile dimension; change out_features to
output_size multiplied by the number of quantiles (e.g. parameters.output_size *
len(parameters.quantiles) or parameters.output_size * len(self.quantiles)
depending on where quantiles are stored) so the produced tensor can be reshaped
correctly; ensure the quantiles list is accessible at this point and update any
related comments or docs to reflect the per-quantile output shape.
| static_input = static_categorical_features.to( | ||
| device=self.static_context_linear.weight.device | ||
| ) | ||
|
|
||
| static_context = self.static_context_linear(static_input) | ||
|
|
||
| static_context = static_context.view((self.batch_size, self.hidden_size)) | ||
|
|
There was a problem hiding this comment.
Static context shaping bug; avoid hard-coded view with self.batch_size
static_categorical_features likely has shape (batch, 1, static_dim) or (batch, static_dim). Use squeeze(1) if needed and avoid using self.batch_size in view.
- static_input = static_categorical_features.to(
- device=self.static_context_linear.weight.device
- )
-
- static_context = self.static_context_linear(static_input)
-
- static_context = static_context.view((self.batch_size, self.hidden_size))
+ static_input = static_categorical_features.to(
+ device=self.static_context_linear.weight.device
+ )
+ if len(static_input.shape) == 3 and static_input.shape[1] == 1:
+ static_input = static_input.squeeze(1)
+ static_context = self.static_context_linear(static_input) # (batch_size, hidden_size)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| static_input = static_categorical_features.to( | |
| device=self.static_context_linear.weight.device | |
| ) | |
| static_context = self.static_context_linear(static_input) | |
| static_context = static_context.view((self.batch_size, self.hidden_size)) | |
| static_input = static_categorical_features.to( | |
| device=self.static_context_linear.weight.device | |
| ) | |
| if len(static_input.shape) == 3 and static_input.shape[1] == 1: | |
| static_input = static_input.squeeze(1) | |
| static_context = self.static_context_linear(static_input) # (batch_size, hidden_size) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 129 to 136, the
static_context shaping uses a hard-coded self.batch_size which can fail when
static_categorical_features has shape (batch, 1, static_dim) or (batch,
static_dim); instead remove the fixed self.batch_size, squeeze the needless
dimension if present (e.g. .squeeze(1) when second dim == 1) and reshape/view
based on the tensor's actual first dimension or use .view(-1, self.hidden_size)
(or .reshape(-1, self.hidden_size)) to produce (batch, hidden_size); ensure you
perform the squeeze before moving to device/linear if ordering matters and
preserve device consistency.
| def test_dataset_batches() -> None: | ||
| data = pl.DataFrame( | ||
| { | ||
| "timestamp": [ | ||
| 1672531200000, # 2023-01-01 | ||
| 1672617600000, # 2023-01-02 | ||
| 1672704000000, # 2023-01-03 | ||
| ], | ||
| "open_price": [100.0, 101.0, 102.0], | ||
| "high_price": [110.0, 111.0, 112.0], | ||
| "low_price": [90.0, 91.0, 92.0], | ||
| "close_price": [105.0, 106.0, 107.0], | ||
| "volume": [1000.0, 1100.0, 1200.0], | ||
| "volume_weighted_average_price": [105.0, 106.0, 107.0], | ||
| "ticker": ["AAPL", "AAPL", "AAPL"], | ||
| "sector": ["Technology", "Technology", "Technology"], | ||
| "industry": [ | ||
| "Consumer Electronics", | ||
| "Consumer Electronics", | ||
| "Consumer Electronics", | ||
| ], | ||
| "is_holiday": [True, False, False], | ||
| } | ||
| ) | ||
|
|
||
| dataset = TemporalFusionTransformerDataset(data=data) | ||
|
|
||
| expected_input_length = 2 | ||
| expected_output_length = 1 | ||
|
|
||
| batches = dataset.get_batches( | ||
| data_type="predict", | ||
| input_length=expected_input_length, | ||
| output_length=expected_output_length, | ||
| ) | ||
|
|
||
| assert isinstance(batches, list) | ||
| assert len(batches) == 1 | ||
|
|
||
| for batch in batches: | ||
| assert "encoder_categorical_features" in batch | ||
| assert "encoder_continuous_features" in batch | ||
| assert "decoder_categorical_features" in batch | ||
| assert "static_categorical_features" in batch | ||
|
|
||
| encoder_categorical_features = batch["encoder_categorical_features"] | ||
| encoder_continuous_features = batch["encoder_continuous_features"] | ||
| decoder_categorical_features = batch["decoder_categorical_features"] | ||
| static_categorical_features = batch["static_categorical_features"] | ||
|
|
||
| assert encoder_categorical_features.shape[0] == expected_input_length | ||
| assert encoder_continuous_features.shape[0] == expected_input_length | ||
| assert decoder_categorical_features.shape[0] == expected_output_length | ||
| assert static_categorical_features.shape[0] == 1 |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider adding tests for train and validate data types.
The current test only covers data_type="predict". Consider adding test cases for "train" and "validate" modes to ensure complete coverage of the get_batches method.
Would you like me to generate additional test cases for the "train" and "validate" data types to ensure comprehensive coverage of the batch generation functionality?
🤖 Prompt for AI Agents
In libraries/python/tests/test_dataset.py around lines 81 to 134, the test only
exercises get_batches for data_type="predict"; add two additional subtests
calling dataset.get_batches with data_type="train" and data_type="validate"
(using the same expected_input_length and expected_output_length) and perform
the same assertions as the existing predict test: verify the result is a list,
non-empty (expecting one batch for this small fixture), and for each batch
assert the presence of "encoder_categorical_features",
"encoder_continuous_features", "decoder_categorical_features", and
"static_categorical_features" and that their row counts equal
expected_input_length, expected_input_length, expected_output_length, and 1
respectively.
| hidden_size: int = 64 | ||
| output_size: int = 1 # closing price | ||
| lstm_layer_count: int = 3 | ||
| attention_head_size: int = 4 | ||
| dropout_rate: float = 0.1 | ||
| quantiles: list[float] = [0.1, 0.5, 0.9] | ||
| decoder_categorical_dimension: int |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid mutable default for quantiles in a Pydantic model
Using a list literal as a default creates a shared mutable default across instances. Prefer Field(default_factory=...).
Apply this diff:
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
@@
- quantiles: list[float] = [0.1, 0.5, 0.9]
+ quantiles: list[float] = Field(default_factory=lambda: [0.1, 0.5, 0.9])📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| hidden_size: int = 64 | |
| output_size: int = 1 # closing price | |
| lstm_layer_count: int = 3 | |
| attention_head_size: int = 4 | |
| dropout_rate: float = 0.1 | |
| quantiles: list[float] = [0.1, 0.5, 0.9] | |
| decoder_categorical_dimension: int | |
| from pydantic import BaseModel, Field | |
| hidden_size: int = 64 | |
| output_size: int = 1 # closing price | |
| lstm_layer_count: int = 3 | |
| attention_head_size: int = 4 | |
| dropout_rate: float = 0.1 | |
| quantiles: list[float] = Field(default_factory=lambda: [0.1, 0.5, 0.9]) | |
| decoder_categorical_dimension: int |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 20 to 26, the
Pydantic model uses a list literal for the default of `quantiles`, which is a
mutable shared default; change it to use Field with a default_factory to return
a new list per instance (e.g. import Field from pydantic and replace the literal
with Field(default_factory=lambda: [0.1, 0.5, 0.9])) so each model instance gets
its own list.
| def forward(self, inputs: dict[str, Tensor]) -> dict[str, Tensor]: | ||
| # NOTE: rename all variables | ||
| # NOTE: potentially remove unused variables | ||
| encoder_categorical_features = inputs["encoder_categorical_features"] | ||
| encoder_continuous_features = inputs["encoder_continuous_features"] | ||
|
|
||
| encoder_input = encoder_categorical_features.cat( | ||
| encoder_continuous_features, | ||
| dim=2, | ||
| ) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Derive dynamic batch/sequence sizes at runtime
Compute batch_size and encoder sequence length from the provided inputs; they are used repeatedly in this method.
encoder_categorical_features = inputs["encoder_categorical_features"]
encoder_continuous_features = inputs["encoder_continuous_features"]
+ # Dynamic sizes
+ batch_size = encoder_categorical_features.shape[0]
+ encoder_seq_len = encoder_categorical_features.shape[1]
+
encoder_input = encoder_categorical_features.cat(
encoder_continuous_features,
dim=2,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def forward(self, inputs: dict[str, Tensor]) -> dict[str, Tensor]: | |
| # NOTE: rename all variables | |
| # NOTE: potentially remove unused variables | |
| encoder_categorical_features = inputs["encoder_categorical_features"] | |
| encoder_continuous_features = inputs["encoder_continuous_features"] | |
| encoder_input = encoder_categorical_features.cat( | |
| encoder_continuous_features, | |
| dim=2, | |
| ) | |
| def forward(self, inputs: dict[str, Tensor]) -> dict[str, Tensor]: | |
| # NOTE: rename all variables | |
| # NOTE: potentially remove unused variables | |
| encoder_categorical_features = inputs["encoder_categorical_features"] | |
| encoder_continuous_features = inputs["encoder_continuous_features"] | |
| # Dynamic sizes | |
| batch_size = encoder_categorical_features.shape[0] | |
| encoder_seq_len = encoder_categorical_features.shape[1] | |
| encoder_input = encoder_categorical_features.cat( | |
| encoder_continuous_features, | |
| dim=2, | |
| ) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 103 to 113, the
method currently doesn't derive dynamic dimensions and repeatedly assumes sizes;
compute batch_size = inputs["encoder_categorical_features"].size(0) and
encoder_seq_len = inputs["encoder_categorical_features"].size(1) (or equivalent
based on tensor shape) at runtime and use these variables wherever batch or
encoder sequence length is needed in the method, replacing hard-coded or
repeatedly accessed size calls to improve clarity and avoid redundant indexing.
| decoder_categorical_features = inputs["decoder_categorical_features"] | ||
| decoder_continuous_features = Tensor.zeros( | ||
| self.batch_size, decoder_categorical_features.shape[1], 0 | ||
| ) # not currently used | ||
|
|
||
| decoder_input = decoder_categorical_features.cat( | ||
| decoder_continuous_features, | ||
| dim=2, | ||
| ) | ||
|
|
There was a problem hiding this comment.
Fix decoder zeros shape; first dimension must be batch size (not input length)
Currently uses self.batch_size which is wrong. Also capture decoder sequence length.
decoder_categorical_features = inputs["decoder_categorical_features"]
- decoder_continuous_features = Tensor.zeros(
- self.batch_size, decoder_categorical_features.shape[1], 0
- ) # not currently used
+ decoder_seq_len = decoder_categorical_features.shape[1]
+ decoder_continuous_features = Tensor.zeros(
+ batch_size, decoder_seq_len, 0
+ ) # not currently usedCommittable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 114 to 123, the zero
tensor uses self.batch_size incorrectly; derive batch_size and decoder_seq_len
from decoder_categorical_features (batch_size =
decoder_categorical_features.shape[0], decoder_seq_len =
decoder_categorical_features.shape[1]) and create the zeros tensor with shape
(batch_size, decoder_seq_len, 0) (or use a zeros_like approach for clarity)
before concatenating along dim=2 so the first two dimensions match
decoder_categorical_features.
| encoder_static_context = static_context.unsqueeze(1).expand( | ||
| -1, self.batch_size, -1 | ||
| ) | ||
|
|
||
| decoder_static_context = static_context.unsqueeze(1).expand( | ||
| -1, self.output_length, -1 | ||
| ) |
There was a problem hiding this comment.
Use dynamic sequence lengths for static context expansion
self.batch_size is not a sequence length. Expand over encoder/decoder sequence lengths.
- encoder_static_context = static_context.unsqueeze(1).expand(
- -1, self.batch_size, -1
- )
+ encoder_static_context = static_context.unsqueeze(1).expand(
+ -1, encoder_seq_len, -1
+ )
- decoder_static_context = static_context.unsqueeze(1).expand(
- -1, self.output_length, -1
- )
+ decoder_static_context = static_context.unsqueeze(1).expand(
+ -1, decoder_seq_len, -1
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| encoder_static_context = static_context.unsqueeze(1).expand( | |
| -1, self.batch_size, -1 | |
| ) | |
| decoder_static_context = static_context.unsqueeze(1).expand( | |
| -1, self.output_length, -1 | |
| ) | |
| encoder_static_context = static_context.unsqueeze(1).expand( | |
| -1, encoder_seq_len, -1 | |
| ) | |
| decoder_static_context = static_context.unsqueeze(1).expand( | |
| -1, decoder_seq_len, -1 | |
| ) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 145 to 151, the
static context is being expanded using self.batch_size for the encoder sequence
dimension which is incorrect; change the encoder expansion to use the encoder
sequence length (e.g., self.input_length or the actual encoder tensor's sequence
length) so it becomes static_context.unsqueeze(1).expand(-1, encoder_seq_len,
-1). Leave the decoder expansion using self.output_length or similarly derive it
dynamically from the decoder/input tensors to ensure the expanded shapes match
encoder/decoder time dimensions.
| decoder_attended = attended_output[:, -self.output_length :, :] | ||
|
|
||
| output = self.pre_output_layer(decoder_attended).relu() | ||
|
|
||
| predictions = self.output_layer(output) | ||
|
|
||
| quantiles = predictions.reshape( | ||
| self.batch_size, | ||
| self.output_length, | ||
| self.output_size, | ||
| len(self.quantiles), | ||
| ) | ||
|
|
||
| return { | ||
| "predictions": predictions, # shape: (batch_size, output_length, output_size) # noqa: E501 | ||
| "quantiles": quantiles, # shape: (batch_size, output_length, output_size, len(quantiles)) # noqa: E501 | ||
| } |
There was a problem hiding this comment.
Quantile tensor construction and point predictions
- Reshape uses
self.batch_sizeand assumes output layer already expanded quantile dimension; both are incorrect as-is. - Provide a correct reshape with dynamic
batch_size, and produce a point prediction (e.g., median) for convenience.
- output = self.pre_output_layer(decoder_attended).relu()
-
- predictions = self.output_layer(output)
-
- quantiles = predictions.reshape(
- self.batch_size,
- self.output_length,
- self.output_size,
- len(self.quantiles),
- )
-
- return {
- "predictions": predictions, # shape: (batch_size, output_length, output_size)
- "quantiles": quantiles, # shape: (batch_size, output_length, output_size, len(quantiles))
- }
+ output = self.pre_output_layer(decoder_attended).relu()
+
+ raw = self.output_layer(output) # (batch_size, output_length, output_size * Q)
+ quantiles = raw.reshape(
+ batch_size,
+ self.output_length,
+ self.output_size,
+ len(self.quantiles),
+ )
+
+ # Use median quantile for point predictions if available
+ try:
+ median_index = self.quantiles.index(0.5)
+ except ValueError:
+ median_index = len(self.quantiles) // 2
+ predictions = quantiles[:, :, :, median_index]
+
+ return {
+ "predictions": predictions, # (batch_size, output_length, output_size)
+ "quantiles": quantiles, # (batch_size, output_length, output_size, Q)
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| decoder_attended = attended_output[:, -self.output_length :, :] | |
| output = self.pre_output_layer(decoder_attended).relu() | |
| predictions = self.output_layer(output) | |
| quantiles = predictions.reshape( | |
| self.batch_size, | |
| self.output_length, | |
| self.output_size, | |
| len(self.quantiles), | |
| ) | |
| return { | |
| "predictions": predictions, # shape: (batch_size, output_length, output_size) # noqa: E501 | |
| "quantiles": quantiles, # shape: (batch_size, output_length, output_size, len(quantiles)) # noqa: E501 | |
| } | |
| decoder_attended = attended_output[:, -self.output_length :, :] | |
| output = self.pre_output_layer(decoder_attended).relu() | |
| raw = self.output_layer(output) # (batch_size, output_length, output_size * Q) | |
| quantiles = raw.reshape( | |
| batch_size, | |
| self.output_length, | |
| self.output_size, | |
| len(self.quantiles), | |
| ) | |
| # Use median quantile for point predictions if available | |
| try: | |
| median_index = self.quantiles.index(0.5) | |
| except ValueError: | |
| median_index = len(self.quantiles) // 2 | |
| predictions = quantiles[:, :, :, median_index] | |
| return { | |
| "predictions": predictions, # (batch_size, output_length, output_size) | |
| "quantiles": quantiles, # (batch_size, output_length, output_size, Q) | |
| } |
🤖 Prompt for AI Agents
libraries/python/src/internal/tft_model.py around lines 173 to 189: the current
reshape hardcodes self.batch_size and assumes the output_layer already expanded
the quantile dimension; instead reshape dynamically by inferring the batch
dimension (use -1 for batch when calling reshape/view) to produce quantiles with
shape (batch, output_length, output_size, len(self.quantiles)), and add a
convenience point prediction returned (e.g., the median) by selecting the center
quantile along the last dimension; ensure you use self.output_length and
self.output_size for the other dimensions and len(self.quantiles) for the
quantile axis.
| loss = quantile_loss( | ||
| outputs["predictions"], | ||
| inputs["targets"], | ||
| self.quantiles, | ||
| ) | ||
|
|
||
| optimizer.zero_grad() | ||
| _ = loss.backward() | ||
| optimizer.step() | ||
|
|
There was a problem hiding this comment.
Incorrect tensor shapes passed to quantile_loss
quantile_loss expects (batch, output_size, Q) predictions and (batch, output_size) targets. Your predictions currently lack a quantile dimension and include a time dimension. Use the quantiles tensor from forward and flatten the time dimension.
- loss = quantile_loss(
- outputs["predictions"],
- inputs["targets"],
- self.quantiles,
- )
+ Q = len(self.quantiles)
+ preds = outputs["quantiles"].reshape(-1, self.output_size, Q)
+ targets = inputs["targets"].reshape(-1, self.output_size)
+ loss = quantile_loss(
+ preds,
+ targets,
+ self.quantiles,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| loss = quantile_loss( | |
| outputs["predictions"], | |
| inputs["targets"], | |
| self.quantiles, | |
| ) | |
| optimizer.zero_grad() | |
| _ = loss.backward() | |
| optimizer.step() | |
| Q = len(self.quantiles) | |
| preds = outputs["quantiles"].reshape(-1, self.output_size, Q) | |
| targets = inputs["targets"].reshape(-1, self.output_size) | |
| loss = quantile_loss( | |
| preds, | |
| targets, | |
| self.quantiles, | |
| ) | |
| optimizer.zero_grad() | |
| _ = loss.backward() | |
| optimizer.step() |
| loss = quantile_loss( | ||
| outputs["predictions"], | ||
| inputs["targets"], | ||
| self.quantiles, | ||
| ) | ||
|
|
||
| total_loss += loss.numpy().item() | ||
| batch_count += 1 | ||
|
|
||
| return total_loss / batch_count |
There was a problem hiding this comment.
Validation loss uses wrong tensor shapes
Same issue as training: use quantiles and flatten time dimension before computing loss.
- loss = quantile_loss(
- outputs["predictions"],
- inputs["targets"],
- self.quantiles,
- )
-
- total_loss += loss.numpy().item()
+ Q = len(self.quantiles)
+ preds = outputs["quantiles"].reshape(-1, self.output_size, Q)
+ targets = inputs["targets"].reshape(-1, self.output_size)
+ loss = quantile_loss(
+ preds,
+ targets,
+ self.quantiles,
+ )
+ total_loss += loss.numpy().item()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| loss = quantile_loss( | |
| outputs["predictions"], | |
| inputs["targets"], | |
| self.quantiles, | |
| ) | |
| total_loss += loss.numpy().item() | |
| batch_count += 1 | |
| return total_loss / batch_count | |
| Q = len(self.quantiles) | |
| preds = outputs["quantiles"].reshape(-1, self.output_size, Q) | |
| targets = inputs["targets"].reshape(-1, self.output_size) | |
| loss = quantile_loss( | |
| preds, | |
| targets, | |
| self.quantiles, | |
| ) | |
| total_loss += loss.numpy().item() | |
| batch_count += 1 | |
| return total_loss / batch_count |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 236 to 245, the
validation loss is computed with mismatched tensor shapes; reshape the tensors
to flatten the time dimension and align the quantile dimension before calling
quantile_loss — specifically permute/reshape outputs["predictions"] to
(batch*time, num_quantiles) and flatten inputs["targets"] to (batch*time,) (or
expand to (batch*time, num_quantiles) if quantile_loss expects that), ensure
quantiles are broadcastable to the prediction shape, then call quantile_loss,
convert the resulting tensor to a scalar and update total_loss and batch_count
as before.
There was a problem hiding this comment.
Actionable comments posted: 39
🔭 Outside diff range comments (7)
infrastructure/environment_variables.py (1)
20-32: Preserve secret semantics and add a temporary alias for backward compatibility.Wrap the final mapping in Output.secret to ensure the entire map is treated as a secret and not logged. Also add ALPACA_API_KEY as an alias to reduce rollout risk; you can remove it after all consumers migrate.
- return pulumi.Output.all( - [ - ("ALPACA_API_KEY_ID", alpaca_api_key_id), - ("ALPACA_API_SECRET", alpaca_api_secret), - ("DATA_BUCKET_NAME", data_bucket_name), - ("POLYGON_API_KEY", polygon_api_key), - ("DUCKDB_ACCESS_KEY", duckdb_access_key), - ("DUCKDB_SECRET", duckdb_secret), - ("AWS_REGION", aws_region), - ("DUCKDB_USER_ACCESS_KEY_ID", duckdb_user_access_key.id), - ("DUCKDB_USER_ACCESS_KEY_SECRET", duckdb_user_access_key.secret), - ] - ).apply(lambda secrets: dict(secrets)) + return pulumi.Output.secret( + pulumi.Output.all( + [ + ("ALPACA_API_KEY_ID", alpaca_api_key_id), + # Temporary alias to ease migration; remove after downstream updates. + ("ALPACA_API_KEY", alpaca_api_key_id), + ("ALPACA_API_SECRET", alpaca_api_secret), + ("DATA_BUCKET_NAME", data_bucket_name), + ("POLYGON_API_KEY", polygon_api_key), + ("DUCKDB_ACCESS_KEY", duckdb_access_key), + ("DUCKDB_SECRET", duckdb_secret), + ("AWS_REGION", aws_region), + ("DUCKDB_USER_ACCESS_KEY_ID", duckdb_user_access_key.id), + ("DUCKDB_USER_ACCESS_KEY_SECRET", duckdb_user_access_key.secret), + ] + ).apply(lambda secrets: dict(secrets)) + )applications/datamanager/pyproject.toml (1)
6-19: Action: Add alpaca-py to applications/datamanager/pyproject.tomldatamanager imports alpaca.trading (applications/datamanager/src/datamanager/clients.py) but applications/datamanager/pyproject.toml doesn't declare alpaca-py — this will raise ImportError at runtime. I verified the repo already declares alpaca-py in applications/positionmanager/pyproject.toml ("alpaca-py>=0.15.0") and uv.lock pins alpaca-py to 0.42.0; pick the repo-wide constraint you prefer.
Files to update:
- applications/datamanager/pyproject.toml — add dependency
- applications/datamanager/src/datamanager/clients.py — (imports alpaca.trading)
Suggested diff (align version with positionmanager or uv.lock as desired):
dependencies = [ "fastapi>=0.115.12", "uvicorn>=0.34.2", "duckdb>=1.2.2", "polars>=1.29.0", "pyarrow>=20.0.0", "loguru>=0.7.3", "requests>=2.31.0", "prometheus-fastapi-instrumentator>=7.1.0", "cloudevents>=1.12.0", "polygon-api-client>=1.14.6", "boto3>=1.38.23", + "alpaca-py>=0.15.0", "internal==0.1.0", ]libraries/python/src/internal/dates.py (1)
38-49: Fix validator type annotations and avoid redundant @classmethod.The end validator is declared on a date field but typed as datetime.datetime. Also, field_validator returns a classmethod in Pydantic v2; the explicit @classmethod is redundant and unconventional.
Apply:
- @field_validator("end") - @classmethod - def check_end_after_start( - cls, - end_value: datetime.datetime, - info: core_schema.ValidationInfo, - ) -> datetime.datetime: + @field_validator("end") + def check_end_after_start( + cls, + end_value: datetime.date, + info: core_schema.ValidationInfo, + ) -> datetime.date: start_value = info.data.get("start") if start_value and end_value <= start_value: msg = "End date must be after start date." raise ValueError(msg) return end_valuepyproject.toml (1)
22-28: Pytest rootdir points to an absolute path that likely doesn't exist"--rootdir=/tests" will set the repository root for pytest to a non-existent absolute directory on most CI/dev setups, breaking discovery/imports. You already set testpaths; remove rootdir to let pytest infer the correct project root.
addopts = [ "--verbose", "--tb=short", "--strict-markers", "--color=yes", - "--rootdir=/tests", ]libraries/python/src/internal/lstm_network.py (1)
68-72: Dropout is applied to the wrong tensor and toggles training mode unexpectedly
- Dropout should be applied to the layer’s output (layer_hidden_state), not directly to hidden_state[index] after it’s stored.
- Calling .train() on a Tensor here is surprising and leaves training state altered for downstream ops.
Refactor to apply dropout to layer_hidden_state without mutating training mode.
- if self.dropout_rate > 0.0 and index < self.layer_count - 1: - hidden_state[index].train() - hidden_state[index] = hidden_state[index].dropout(self.dropout_rate) + if self.dropout_rate > 0.0 and index < self.layer_count - 1: + layer_hidden_state = layer_hidden_state.dropout(self.dropout_rate) + + # propagate the (possibly dropped-out) hidden state forward and into states + hidden_state[index] = layer_hidden_state + cell_state[index] = layer_cell_stateNote: This moves state assignment after optional dropout to keep state and forward input consistent.
Also applies to: 56-66
applications/datamanager/src/datamanager/main.py (1)
179-185: Fix Content-Disposition header and Arrow media type
- f-string debug syntax {filename=} yields Content-Disposition: attachment; filename='…' which is invalid.
- You’re writing Arrow IPC stream (RecordBatchStreamWriter), so media type should be application/vnd.apache.arrow.stream.
Apply this patch:
- filename = f"equity_bars_{start_date}_{end_date}.arrow" - content_disposition = f"attachment; {filename=}" + filename = f"equity_bars_{start_date}_{end_date}.arrow" + content_disposition = f'attachment; filename="{filename}"' ... - return Response( + return Response( content=sink.getvalue().to_pybytes(), - media_type="application/vnd.apache.arrow.file", + media_type="application/vnd.apache.arrow.stream", headers={ "Content-Disposition": content_disposition, "X-Row-Count": str(data.num_rows), "X-Start-Date": str(start_date), "X-End-Date": str(end_date), }, )applications/positionmanager/pyproject.toml (1)
5-5: Python pin (==3.12.10) conflicts with test image (Python 3.13) — fix requiredShort: Multiple pyproject.toml files (and uv.lock) pin requires-python = "==3.12.10", while the test images use Python 3.13 — this will cause lock/uv resolution failures in CI/tests.
Files to update / check:
- pyproject.toml (root) — requires-python = "==3.12.10"
- libraries/python/pyproject.toml — requires-python = "==3.12.10"
- applications/positionmanager/pyproject.toml — requires-python = "==3.12.10"
- applications/predictionengine/pyproject.toml — requires-python = "==3.12.10"
- applications/datamanager/pyproject.toml — requires-python = "==3.12.10"
- applications/models/pyproject.toml — requires-python = "==3.12.10"
- infrastructure/pyproject.toml — requires-python = "==3.12.10"
- uv.lock — requires-python = "==3.12.10" (must be regenerated after changes)
- Dockerfiles using Python 3.13:
- Dockerfile.tests (root): FROM python:3.13
- applications/datamanager/Dockerfile.test: FROM python:3.13-slim
Proposed change (preferred):
requires-python = ">=3.12,<3.14"After making the pin change in all pyproject.toml files, regenerate the lockfile (e.g., run the project’s uv sync command used in Dockerfile.tests) and re-run tests.
Alternative: align test images to Python 3.12.10 (change Dockerfile.tests and any test images) instead of loosening the pyproject pins.
♻️ Duplicate comments (4)
applications/models/pyproject.toml (1)
7-19: Remove duplicate internal dependency; consider relying on the workspace source.The dependency "internal==0.1.0" is listed twice, which can cause resolver noise. Also, since you’ve configured [tool.uv.sources] to use the workspace, you can avoid pinning or relax to a range to prevent version skew between the workspace and this app.
Apply this minimal fix:
dependencies = [ - "internal==0.1.0", + "internal==0.1.0", "boto3>=1.38.23", "botocore>=1.38.23", "requests>=2.31.0", "pyarrow>=20.0.0", "polygon-api-client>=1.14.6", - "internal==0.1.0", "flytekit>=1.16.1", "polars>=1.29.0", "loguru>=0.7.3", "pydantic>=2.8.2", "wandb>=0.21.1", ]Optional follow-ups:
- Prefer unpinned internal to align with the workspace source:
- Replace "internal==0.1.0" with "internal" or "internal>=0.1.0".
- Consider relaxing requires-python to a minor range (e.g., ">=3.12,<3.13") unless you have a strict patch-level constraint.
libraries/python/src/internal/loss_functions.py (1)
16-26: Bug: accumulation doesn’t update errors_total (non-inplace add).Tensor.add likely returns a new Tensor; the current code discards the result, yielding an incorrect (always-zero) loss.
Apply this fix:
- errors_total = Tensor.zeros((1,)) + errors_total = Tensor.zeros((1,)) for index, quantile in enumerate(quantiles): error = targets.sub(predictions[:, :, index]) quantile_tensor = Tensor([quantile]) - errors_total.add( + errors_total = errors_total.add( Tensor.where( error > 0, quantile_tensor.mul(error), (quantile_tensor.sub(1)).mul(error), ).mean() )libraries/python/src/internal/lstm_network.py (1)
76-81: Stacking outputs will fail/behave oddly for a single timestep; handle len==1 safelyUsing Tensor.stack(outputs[0], *outputs[1:], dim=1) assumes ≥1 outputs and relies on tinygrad’s varargs semantics; single-element behavior can be inconsistent and the current pattern is brittle. Stack from the whole list and special-case len==1 to unsqueeze.
- output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) + if len(outputs) == 1: + output_tensor = outputs[0].unsqueeze(dim=1) + else: + output_tensor = Tensor.stack(*outputs, dim=1)applications/datamanager/src/datamanager/main.py (1)
266-281: Don’t return 404 on S3 listing exceptions; distinguish errors from “not found”An exception while listing objects is an operational error (500), not a “no data” condition (404). Only return 404 when the list is empty without exceptions.
- try: - blobs = list(s3_client.list_objects(prefix=prefix)) - except Exception as e: # noqa: BLE001 - logger.error(f"Error listing S3 objects: {e}") - return Response( - status_code=status.HTTP_404_NOT_FOUND, - content=json.dumps({"error": "No equity bars found"}), - media_type="application/json", - ) + try: + blobs = list(s3_client.list_objects(prefix=prefix)) + except Exception as e: # noqa: BLE001 + logger.error(f"Error listing S3 objects: {e}") + return Response( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content=json.dumps({"error": f"Failed to list equity bars: {e!s}"}), + media_type="application/json", + )Optionally catch boto3’s ClientError explicitly if available.
🧹 Nitpick comments (48)
.gitignore (1)
6-8: Scope CSV and wandb ignores to avoid masking legitimate data files.Top-level ignores for *.csv and wandb/ can inadvertently hide intentionally tracked assets elsewhere in the repo. Recommend scoping to the models app (or a dedicated data/output directory) to prevent accidental omissions.
Apply this diff to narrow the patterns:
-*.egg-info -wandb/ -*.csv +*.egg-info +# Scope experiment artifacts to models app only (adjust as needed) +applications/models/wandb/ +applications/models/**/*.csvOptional:
- If you want to allow CSVs in repo-wide docs/examples, consider adding exceptions with negate patterns (!docs/**/*.csv).
Also applies to: 10-10, 19-21
applications/datamanager/pyproject.toml (1)
18-18: Avoid hard pinning the internal workspace dependency to reduce churn.You’re already sourcing internal from the uv workspace; pinning to an exact version adds friction when you bump the library.
Apply one of the following:
- Prefer letting the workspace source define the version:
- "internal==0.1.0", + "internal",
- Or allow forward-compatible bumps:
- "internal==0.1.0", + "internal>=0.1.0",libraries/python/src/internal/dates.py (1)
8-31: Date model looks good — add backward-compatibility alias (SummaryDate → Date)Renaming SummaryDate → Date is a breaking change. I searched the repo for "SummaryDate" (no matches found), but absence of matches here doesn't guarantee no external consumers — add a soft-landing alias if downstream code may reference the old name.
- File to update:
- libraries/python/src/internal/dates.py — add alias at the bottom of the module
Suggested diff:
+# Backward-compatibility alias (remove after migration) +SummaryDate = Datelibraries/python/pyproject.toml (1)
1-12: Right direction for a shared internal library; consider optionalizing heavy deps.This centralizes shared code under uv. To reduce install footprint for services not using ML, consider moving tinygrad (and possibly numpy) under an optional dependency group, e.g., [project.optional-dependencies]. Keep pydantic, cloudevents, polars as core if widely used.
Example:
[project] name = "internal" version = "0.1.0" description = "Shared Python resources" requires-python = "==3.12.10" -dependencies = [ - "pydantic>=2.8.2", - "cloudevents>=1.12.0", - "tinygrad>=0.10.3", - "numpy>=2.2.6", - "polars>=1.29.0", -] +dependencies = [ + "pydantic>=2.8.2", + "cloudevents>=1.12.0", + "polars>=1.29.0", +] + +[project.optional-dependencies] +ml = [ + "numpy>=2.2.6", + "tinygrad>=0.10.3", +]If you make this change, ensure packages that require ML features add internal[ml] to their dependencies.
.github/workflows/teardown_infrastructure.yaml (1)
1-9: Clarify the schedule comment for DST and simplify the if guard (optional).
- The comment says “6:00 PM EST,” but 23:00 UTC is 7:00 PM during EDT. Prefer “ET” or specify both.
- The if condition is redundant for a workflow that only has a single schedule; it’s harmless but can be removed.
Apply:
name: Teardown infrastructure on: schedule: - - cron: '0 23 * * 1,2,3,4,5' # teardown at 6:00 PM EST + - cron: '0 23 * * 1,2,3,4,5' # 23:00 UTC (6pm ET standard / 7pm EDT) jobs: teardown_infrastructure: - name: Teardown infrastructure on weekday schedule - if: github.event.schedule == '0 23 * * 1,2,3,4,5' + name: Teardown infrastructure on weekday schedule.github/workflows/launch_infrastructure.yaml (1)
2-8: Cron comment references EST; consider DST implications.The cron at Line 5 runs at 13:00 UTC, which is 8:00 AM EST but 9:00 AM EDT. If you intend 8:00 AM local ET year-round, update the comment and/or cron strategy to avoid confusion.
libraries/python/tests/test_lstm_network.py (1)
39-47: Strengthen state validation in tests.Optionally assert that each state list matches layer_count to catch regressions (e.g., hidden/cell state list length).
You can extend assertions like this:
- assert isinstance(hidden_state, tuple) - assert len(hidden_state) == expected_hidden_state + assert isinstance(hidden_state, tuple) + assert len(hidden_state) == expected_hidden_state # (hidden_state, cell_state) + hs, cs = hidden_state + assert isinstance(hs, list) and isinstance(cs, list) + assert len(hs) == lstm.layer_count + assert len(cs) == lstm.layer_countAlso applies to: 72-77, 87-91, 103-108
applications/models/src/models/get_alpaca_calendar.py (1)
33-37: Optional: allow output path override (CLI arg or env).Hard-coding "calendar.csv" to CWD can be inconvenient in CI or when organizing datasets.
Example: accept an OUTPUT_PATH env var with a default.
- calendar_data.write_csv("calendar.csv") + output_path = os.getenv("OUTPUT_PATH", "calendar.csv") + calendar_data.write_csv(output_path) - logger.info("Calendar data has been written to calendar.csv") + logger.info(f"Calendar data has been written to {output_path}")libraries/python/src/internal/cloud_event.py (1)
7-11: Type generality and docstrings for maintainability.Consider using Sequence[str] for event_metadata and add short docstrings clarifying attribute/type composition and timezone choice.
Also applies to: 24-28
libraries/python/src/internal/loss_functions.py (1)
16-28: Optional: vectorize across quantiles to reduce Python loop overhead.You can stack quantiles into a Tensor and compute the piecewise loss in one pass.
Sketch:
# q: (Q,), error: (B, O, Q) q = Tensor(quantiles) error = targets.unsqueeze(-1).sub(predictions) loss_q = Tensor.where(error > 0, q * error, (q - 1) * error).mean(axis=(0, 1)) # (Q,) return loss_q.mean().reshape((1,))libraries/python/src/internal/variable_selection_network.py (2)
1-2: Unify tinygrad imports for consistency with other internal modulesOther internal modules (e.g., mhsa_network.py) import Linear directly from tinygrad.nn. Consider aligning import style here to avoid mixing patterns across the codebase.
-from tinygrad import nn -from tinygrad.tensor import Tensor +from tinygrad.nn import Linear +from tinygrad.tensor import Tensor
17-21: Add a runtime shape check and minimal docstringShape mismatches can propagate and be hard to debug later. A fast check keeps failures local. Adding a short docstring clarifies expected input/output shapes.
- def forward(self, inputs: Tensor) -> Tensor: + def forward(self, inputs: Tensor) -> Tensor: + """ + Args: + inputs: Tensor of shape (..., input_dimension) + Returns: + Tensor of shape (..., input_dimension) with values in [0, 1] + """ + # Quick runtime validation + if inputs.shape[-1] != self.input_dimension: + raise ValueError( + f"Expected inputs.shape[-1] == {self.input_dimension}, got {inputs.shape[-1]}" + ) inputs = self.input_layer(inputs) inputs = inputs.relu() inputs = self.output_layer(inputs) return inputs.sigmoid()libraries/python/tests/test_mhsa_network.py (3)
8-15: Add a negative-case test for non-divisible head dimensionsThis guards the constructor validation.
def test_multi_head_attention_initialization() -> None: heads_count = 8 embedding_size = 64 attention = MultiHeadSelfAttentionNetwork(heads_count=8, embedding_size=64) assert attention.heads_count == heads_count assert attention.embedding_size == embedding_size + + +def test_multi_head_attention_invalid_dims_raises() -> None: + # 30 is not divisible by 8 + try: + _ = MultiHeadSelfAttentionNetwork(heads_count=8, embedding_size=30) + assert False, "Expected ValueError for invalid heads/embedding ratio" + except ValueError: + pass
17-29: Assert attention weights normalization across key dimensionValidates the softmax property: sums should be ~1 across the last axis.
def test_multi_head_attention_forward() -> None: attention = MultiHeadSelfAttentionNetwork(heads_count=4, embedding_size=32) input_tensor = Tensor(rng.standard_normal((2, 10, 32))) output, attention_weights = attention.forward(input_tensor) batch_size = 2 heads_count = 4 assert output.shape == (batch_size, 10, 32) assert attention_weights.shape[0] == batch_size assert attention_weights.shape[1] == heads_count + # Softmax normalization: sum over keys dimension (-1) ≈ 1 + sums = attention_weights.sum(axis=-1).numpy() + assert (abs(sums - 1.0) < 1e-5).all()
64-73: Exercise dropout path without changing shapesA simple sanity check that enabling dropout preserves shapes and runs without error.
def test_multi_head_attention_batch_processing() -> None: attention = MultiHeadSelfAttentionNetwork(heads_count=2, embedding_size=32) for batch_size in [1, 2, 4, 8]: input_tensor = Tensor(rng.standard_normal((batch_size, 5, 32))) output, attention_weights = attention.forward(input_tensor) assert output.shape == (batch_size, 5, 32) assert attention_weights.shape[0] == batch_size + + +def test_multi_head_attention_dropout_shapes() -> None: + attention = MultiHeadSelfAttentionNetwork(heads_count=2, embedding_size=32, dropout_rate=0.5) + input_tensor = Tensor(rng.standard_normal((2, 6, 32))) + output, attention_weights = attention.forward(input_tensor) + assert output.shape == (2, 6, 32) + assert attention_weights.shape[:2] == (2, 2)applications/models/src/models/get_alpaca_equity_bars.py (2)
21-27: Use requests params rather than manual query constructionrequests will handle encoding and edge cases for you; also use headers per-call to ensure validation is performed.
- response = requests.get(url, headers=headers, timeout=30) + response = requests.get(url, headers=_alpaca_headers(), timeout=30)
80-85: Prefer UTC internally; only localize if needed for loggingUsing NY time to compute a 6-year window is okay, but it can lead to the 'Z' bug above and DST edge conditions. Consider computing in UTC and converting only when necessary for human-oriented reporting.
- end = datetime.now(tz=ZoneInfo("America/New_York")) + end = datetime.now(tz=ZoneInfo("UTC")) start = end - timedelta(days=365 * 6)libraries/python/src/internal/lstm_network.py (1)
32-38: Add basic input validation for forward()Forward assumes a 3D input (batch, time, features). Add a quick shape check and explicit int cast for sequence_length for clarity.
def forward( self, - inputs: Tensor, + inputs: Tensor, state: tuple[list[Tensor], list[Tensor]] | None = None, ) -> tuple[Tensor, tuple[list[Tensor], list[Tensor]]]: - batch_size, sequence_length, _ = inputs.shape + if len(inputs.shape) != 3: + raise ValueError(f"Expected inputs with shape (batch, time, features), got {inputs.shape}") + batch_size, sequence_length, _ = inputs.shape + sequence_length = int(sequence_length)Also applies to: 39-50, 53-55
.mise.toml (2)
44-50: Make port mapping configurable; avoid hard-coding 8080Parameterize host/container ports to avoid collisions and allow non-8080 services.
[tasks."application:service:run:production"] description = "Run the service application" run = """ docker run \ --env-file .env \ - --publish 8080:8080 \ + --publish {{arg(name="host_port", default="8080")}}:{{arg(name="container_port", default="8080")}} \ pocketsizefund/{{arg(name="application_name")}}:latest \ """
55-57: Path updates to applications/ look good; add a unit-tests task using uv for faster local runsThe cd path changes are correct. Given our team guideline (remembered from prior feedback) to use mise tasks for tests/format/lint, consider adding a fast unit test task that runs directly with uv/pytest (complementary to the docker-based integration harness).
Suggested addition (new task):
[tasks."python:test:unit"] description = "Run Python unit tests locally with uv/pytest" run = """ uv run pytest -q --disable-warnings --maxfail=1 """Also consider a coverage variant if desired.
Also applies to: 62-64, 69-71, 76-78
applications/datamanager/src/datamanager/main.py (3)
193-205: Keep error telemetry but avoid double-logging full tracebacksYou already log the exception message and a traceback. Consider reducing to one structured log with e and traceback once to keep logs concise, or keep as-is if you rely on message + trace separation.
208-223: CloudEvent success payload includes a date twice; clarify intentcreate_cloud_event_success injects a "date" (now) and you pass "date" (request_fetch_date). The merge order makes your provided date win, which is fine, but the implicit override is non-obvious. Either:
- remove "date" from the data payload here (let the helper set now), or
- add a different key like "requested_date" to avoid ambiguity.
- data={ - "date": request_fetch_date, - "count": count, - }, + data={ + "requested_date": request_fetch_date, + "count": count, + },If you do want the event data's "date" to be the requested date, consider documenting that in the helper or reversing the merge order in create_cloud_event_success.
Also applies to: 255-262
294-333: New /calendar endpoint: consider basic validation and rate limitingEndpoint looks good. Consider:
- Validating start_date <= end_date and max range to avoid heavy upstream calls.
- Returning an empty list (200) instead of 404 is sometimes preferable for “no results” endpoints; verify with API consumers.
Would you like me to add unit tests for this endpoint (success, empty, provider error) mirroring the equity bars suite?
applications/datamanager/tests/test_datamanager_main.py (3)
92-114: Good coverage for /equity-bars database error and empty-data pathsMocks and assertions match endpoint behavior. Consider asserting media_type and headers on success paths in a separate test for stronger contract coverage.
Also applies to: 115-138
156-173: DELETE /equity-bars tests: success and not-found cases are solidThe mocks and expectations reflect the refactored Date model. If you adopt the 500-on-exception change, add a test to assert 500 for an S3 listing error.
Also applies to: 174-193
196-207: POST /equity-bars/fetch tests look good; consider adding /calendar testsAdd tests for:
- GET /calendar success: mock AlpacaClient.get_calendar to return data; assert 200 and JSON structure.
- GET /calendar empty: return []; assert 404 (or 200 if you change semantics).
- GET /calendar provider error: raise Exception; assert 500.
I can draft these test cases if you’d like.
Also applies to: 232-261
libraries/python/src/internal/summaries.py (1)
4-6: BarsSummary model is fine; consider tightening validation and export surfaceCurrent implementation matches the tests (string date, int count). If you want to harden the model without changing semantics:
- Forbid unexpected fields.
- Optionally validate date strings follow YYYY-MM-DD (keeps current usage intact if consumers already use this format).
-from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, field_validator +import re @@ class BarsSummary(BaseModel): - date: str - count: int + model_config = ConfigDict(extra="forbid") + + date: str + count: int + + @field_validator("date") + def _validate_date_format(cls, v: str) -> str: # noqa: N805 + # Accept strictly YYYY-MM-DD; relax if broader formats are needed + if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", v): + raise ValueError("date must match YYYY-MM-DD") + return vIf you prefer not to enforce the date format yet, at least add
extra="forbid"to prevent accidental fields from being accepted.applications/positionmanager/tests/test_positionmanager_main.py (2)
46-59: Mocks for Arrow/Polars/Optimizer look correct; consider autospec for interface safetyThe patch targets and order look right. To reduce brittle tests and catch interface drift, use
autospec=True(where possible) and assert calls on these mocks to ensure your code exercised the expected paths.
71-87: Strengthen assertions to ensure side effects (order placement) occurAfter the POST succeeds, assert that the optimizer and Alpaca client were actually used with expected arguments.
@@ response = client.post("/positions/open", json=cloud_event_data) assert response.status_code == status.HTTP_200_OK response_data = response.json() assert "source" in response_data assert "type" in response_data + # Strengthen: verify business interactions occurred + MockPortfolioOptimizer.assert_called_once() + mock_optimizer_instance.get_optimized_portfolio.assert_called_once() + # Ensure we attempted to place orders using the computed allocation + assert mock_alpaca_instance.place_notional_order.call_count >= 1libraries/python/tests/test_dataset.py (2)
49-79: Consider asserting expected dimension counts, not just presenceYou can make this test more robust by checking counts match the current feature schema (helps detect accidental regressions in feature sets).
@@ dimensions = dataset.get_dimensions() - assert "encoder_categorical_features" in dimensions - assert "encoder_continuous_features" in dimensions - assert "decoder_categorical_features" in dimensions - assert "decoder_continuous_features" in dimensions - assert "static_categorical_features" in dimensions - assert "static_continuous_features" in dimensions + assert dimensions["encoder_categorical_features"] > 0 + assert dimensions["encoder_continuous_features"] > 0 + assert dimensions["decoder_categorical_features"] > 0 + assert dimensions["decoder_continuous_features"] == 0 + assert dimensions["static_categorical_features"] > 0 + assert dimensions["static_continuous_features"] == 0
81-135: LGTM: Prediction batch shape checks hit key invariantsGreat coverage of encoder/decoder/static partitions and lengths.
As a follow-up, consider adding negative tests to lock error semantics:
- invalid data_type raises ValueError,
- insufficient days for requested input/output lengths raises ValueError,
- missing/extra columns raises ValueError.
Example additions (outside the selected lines):
def test_get_batches_invalid_type_raises(dataframe_factory): dataset = TemporalFusionTransformerDataset(data=dataframe_factory(days=3)) with pytest.raises(ValueError): dataset.get_batches(data_type="oops") def test_get_batches_insufficient_days_raises(): data = pl.DataFrame({ "timestamp": [1672531200000], # single day "open_price": [100.0], "high_price": [110.0], "low_price": [90.0], "close_price": [105.0], "volume": [1000.0], "volume_weighted_average_price": [105.0], "ticker": ["AAPL"], "sector": ["Technology"], "industry": ["Consumer Electronics"], "is_holiday": [False], }) dataset = TemporalFusionTransformerDataset(data=data) with pytest.raises(ValueError): dataset.get_batches(data_type="predict", input_length=2, output_length=1)applications/positionmanager/src/positionmanager/clients.py (1)
120-123: Use DateRange.to_object() to construct request params (and avoid duplication).You already expose a stable payload shape via
DateRange.to_object(). Using it here reduces duplication and centralizes the date formatting.- params = { - "start_date": date_range.start.isoformat(), - "end_date": date_range.end.isoformat(), - } + params = date_range.to_object()libraries/python/tests/test_dates.py (2)
10-19: Patch datetime with wraps to avoid over-mocking the module.Patching the entire
internal.dates.datetimemodule can inadvertently mock attributes used elsewhere. Usingwraps=datetimepreserves real behavior except for what you override, making the test more robust.-def test_date_default() -> None: - with patch("internal.dates.datetime") as mock_datetime: +def test_date_default() -> None: + with patch("internal.dates.datetime", wraps=datetime) as mock_datetime: mock_datetime.datetime.now.return_value = datetime.datetime( 2023, 5, 15, 10, 30, 0, tzinfo=ZoneInfo("America/New_York") ) - mock_datetime.date = datetime.date
29-33: Align test name with behavior: pass a string when testing “string dash format”.The current test name suggests a string input, but a date object is provided. Either rename the test or pass the string.
-def test_date_string_dash_format() -> None: - date_instance = Date(date=datetime.date.fromisoformat("2023-03-15")) +def test_date_string_dash_format() -> None: + date_instance = Date(date="2023-03-15")libraries/python/tests/test_loss_functions.py (2)
17-19: Tighten the scalar-shape assertion.The current shape assertion is too permissive and could mask regressions. Assert that the numpy view is a scalar or a single-element vector.
- assert isinstance(loss, Tensor) - assert len(loss.shape) == 0 or loss.shape in [(), (0,), (1,)] + assert isinstance(loss, Tensor) + assert loss.numpy().shape in [(), (1,)]
28-30: Make shape assertion consistent and strict.Same rationale as above.
- assert isinstance(loss, Tensor) - assert len(loss.shape) == 0 or loss.shape in [(), (0,), (1,)] + assert isinstance(loss, Tensor) + assert loss.numpy().shape in [(), (1,)]libraries/python/tests/test_variable_selection_network.py (1)
56-58: Add a finiteness check to catch NaNs/Infs.This small assertion hardens the test against numerical issues without being brittle.
output_numpy = output.numpy() assert np.all(output_numpy >= 0.0) assert np.all(output_numpy <= 1.0) + assert np.isfinite(output_numpy).all()libraries/python/tests/test_equity_bar.py (2)
133-146: Fix mismatch between test name and input under test (pass a string, not a date).The test is named “…_timestamp_string_iso_format” but passes a date object created via date.fromisoformat. Pass the raw string to exercise the validator’s string parsing path.
Apply this diff:
- equity_bar = EquityBar( - ticker="AAPL", - timestamp=date.fromisoformat("2023-06-15"), - open_price=150.0, - high_price=155.0, - low_price=149.0, - close_price=153.0, - volume=1000000, - volume_weighted_average_price=152.5, - ) + equity_bar = EquityBar( + ticker="AAPL", + timestamp="2023-06-15", + open_price=150.0, + high_price=155.0, + low_price=149.0, + close_price=153.0, + volume=1000000, + volume_weighted_average_price=152.5, + )
170-185: Consider adding tests for volume and VWAP non-negativity.Volume and VWAP should never be negative in practice. The model currently doesn’t validate these. Recommend adding tests to guard these invariants; see model-side suggestion in internal/equity_bar.py.
Would you like me to add parametrized tests for negative volume/VWAP as follow-ups?
applications/datamanager/tests/test_datamanager_models.py (2)
20-23: Be explicit with zero-padded dates to avoid locale/implementation quirks.While Python’s strptime usually accepts “2023-5-15” for "%Y-%m-%d", using “2023-05-15” ensures strict adherence to the documented format and avoids edge-case failures.
Apply this diff:
- date_with_string = Date(date="2023-5-15") # type: ignore + date_with_string = Date(date="2023-05-15") # type: ignore
85-90: Confirm domain semantics for negative bar counts.The test asserts negative counts are acceptable. If BarsSummary.count is intended to be non-negative (typical for aggregate counts), consider adding validation to forbid negatives. If negatives are used for sentinel/error values, keep as-is and document.
libraries/python/src/internal/equity_bar.py (1)
40-49: Optional: accept datetime inputs in timestamp validator.If upstream may pass datetime (not just date or ISO strings), extend the validator to accept datetime and call .date().
Proposed snippet (outside selected range):
from datetime import date, datetime @field_validator("timestamp", mode="before") @classmethod def validate_timestamp(cls, value: date | datetime | str) -> date: if isinstance(value, date) and not isinstance(value, datetime): return value if isinstance(value, datetime): return value.date() try: return date.fromisoformat(value) except ValueError as e: raise ValueError("Invalid date format: expected YYYY-MM-DD") from eapplications/models/src/models/train_tft_model.py (3)
32-43: Use total_seconds() for accurate read timing..seconds drops minutes/hours and ignores subseconds. Prefer total_seconds().
Apply this diff:
- runtime_seconds = (datetime.now(tz=timezone) - start_time).seconds + runtime_seconds = int((datetime.now(tz=timezone) - start_time).total_seconds())
111-126: Use total_seconds() for accurate validation timing.Apply this diff:
- runtime_seconds = (datetime.now(tz=timezone) - start_time).seconds + runtime_seconds = int((datetime.now(tz=timezone) - start_time).total_seconds())
130-139: Use total_seconds() for accurate save timing.Apply this diff:
- runtime_seconds = (datetime.now(tz=timezone) - start_time).seconds + runtime_seconds = int((datetime.now(tz=timezone) - start_time).total_seconds())libraries/python/src/internal/dataset.py (2)
10-26: Add type hints forScalerattributesThe
Scalerclass usesself.meansandself.standard_deviationswithout type declarations, which could lead to typing issues and makes the code less self-documenting.class Scaler: def __init__(self) -> None: - pass + self.means: pl.DataFrame | None = None + self.standard_deviations: pl.DataFrame | None = None def fit(self, data: pl.DataFrame) -> None:
108-126: Consider extracting the magic number for FridayThe Friday number (4) is hardcoded and lacks context. Consider using a named constant for clarity.
- friday_number = 4 + FRIDAY_WEEKDAY_INDEX = 4 # Monday=0, ..., Friday=4 # set is_holiday value for missing weekdays data = ( data.with_columns( pl.col("datetime").dt.weekday().alias("temporary_weekday") ) .with_columns( pl.when( pl.col("is_holiday").is_null() - & (pl.col("temporary_weekday") <= friday_number) + & (pl.col("temporary_weekday") <= FRIDAY_WEEKDAY_INDEX) ) .then(True) # noqa: FBT003 .when( pl.col("is_holiday").is_null() - & (pl.col("temporary_weekday") > friday_number) + & (pl.col("temporary_weekday") > FRIDAY_WEEKDAY_INDEX) )applications/positionmanager/src/positionmanager/main.py (1)
168-173: Consider extracting the type check to avoid duplicationThe event data validation logic appears multiple times. Consider creating a helper function to validate event data format.
Create a helper function at the module level:
def validate_event_data(event_data: Any) -> dict: """Validate that event data is a dictionary.""" if not isinstance(event_data, dict): raise ValueError("Invalid event data format, expected a dictionary.") return event_dataThen use it in the handler:
try: - event_data = event.data - if not isinstance(event_data, dict): - return create_cloud_event_error( - application_name="positionmanager", - event_metadata=["positions", "open"], - error_message="Invalid event data format, expected a dictionary.", - ) + event_data = validate_event_data(event.data)
| class AlpacaClient: | ||
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: | ||
| self.trading_client = TradingClient( | ||
| api_key=alpaca_api_key_id, | ||
| secret_key=alpaca_api_secret, | ||
| ) | ||
|
|
There was a problem hiding this comment.
Initialize TradingClient with explicit paper mode and raw_data=False to match attribute access.
By default, alpaca-py may return raw dicts (raw_data=True), which would break attribute access (calendar_entry.date/open/close). Also, explicitly enabling paper mode avoids accidental live trading.
Apply:
class AlpacaClient:
def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None:
- self.trading_client = TradingClient(
- api_key=alpaca_api_key_id,
- secret_key=alpaca_api_secret,
- )
+ self.trading_client = TradingClient(
+ api_key=alpaca_api_key_id,
+ secret_key=alpaca_api_secret,
+ paper=True,
+ raw_data=False,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| class AlpacaClient: | |
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: | |
| self.trading_client = TradingClient( | |
| api_key=alpaca_api_key_id, | |
| secret_key=alpaca_api_secret, | |
| ) | |
| class AlpacaClient: | |
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: | |
| self.trading_client = TradingClient( | |
| api_key=alpaca_api_key_id, | |
| secret_key=alpaca_api_secret, | |
| paper=True, | |
| raw_data=False, | |
| ) |
🤖 Prompt for AI Agents
In applications/datamanager/src/datamanager/clients.py around lines 27 to 33,
the TradingClient is constructed without explicit paper mode and raw_data
settings which can cause returned objects to be raw dicts (breaking attribute
access) and risk accidental live trading; update the TradingClient
initialization to set paper=True and raw_data=False so the client uses paper
trading and returns model objects with attribute access.
| "list[Calendar]", self.trading_client.get_calendar(request) | ||
| ) | ||
|
|
||
| calendar_list = [] | ||
| for calendar_entry in calendar_data: | ||
| calendar_dict = { | ||
| "date": calendar_entry.date, | ||
| "open": calendar_entry.open, | ||
| "close": calendar_entry.close, | ||
| } | ||
| calendar_list.append(calendar_dict) | ||
|
|
||
| return calendar_list | ||
|
|
There was a problem hiding this comment.
Ensure get_calendar returns strings as annotated and is robust to raw_data mode.
The return type is list[dict[str, str]], but date/open/close may be date/time objects (or dict values if raw_data=True). Normalize to strings.
Apply:
- calendar_list = []
- for calendar_entry in calendar_data:
- calendar_dict = {
- "date": calendar_entry.date,
- "open": calendar_entry.open,
- "close": calendar_entry.close,
- }
- calendar_list.append(calendar_dict)
+ calendar_list: list[dict[str, str]] = []
+ for entry in calendar_data:
+ # Support both model objects and raw dicts
+ if isinstance(entry, dict):
+ date_val = entry.get("date")
+ open_val = entry.get("open")
+ close_val = entry.get("close")
+ else:
+ date_val = getattr(entry, "date", None)
+ open_val = getattr(entry, "open", None)
+ close_val = getattr(entry, "close", None)
+
+ calendar_list.append(
+ {
+ "date": str(date_val) if date_val is not None else "",
+ "open": str(open_val) if open_val is not None else "",
+ "close": str(close_val) if close_val is not None else "",
+ }
+ )Given there’s another AlpacaClient in applications/positionmanager with overlapping concerns, consider centralizing an Alpaca client wrapper in libraries/python/src/internal to avoid duplication and keep config (paper/raw_data) consistent across services.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| calendar_data: list[Calendar] = cast( | |
| "list[Calendar]", self.trading_client.get_calendar(request) | |
| ) | |
| calendar_list = [] | |
| for calendar_entry in calendar_data: | |
| calendar_dict = { | |
| "date": calendar_entry.date, | |
| "open": calendar_entry.open, | |
| "close": calendar_entry.close, | |
| } | |
| calendar_list.append(calendar_dict) | |
| return calendar_list | |
| calendar_data: list[Calendar] = cast( | |
| "list[Calendar]", self.trading_client.get_calendar(request) | |
| ) | |
| calendar_list: list[dict[str, str]] = [] | |
| for entry in calendar_data: | |
| # Support both model objects and raw dicts | |
| if isinstance(entry, dict): | |
| date_val = entry.get("date") | |
| open_val = entry.get("open") | |
| close_val = entry.get("close") | |
| else: | |
| date_val = getattr(entry, "date", None) | |
| open_val = getattr(entry, "open", None) | |
| close_val = getattr(entry, "close", None) | |
| calendar_list.append( | |
| { | |
| "date": str(date_val) if date_val is not None else "", | |
| "open": str(open_val) if open_val is not None else "", | |
| "close": str(close_val) if close_val is not None else "", | |
| } | |
| ) | |
| return calendar_list |
| end = datetime.now(tz=ZoneInfo("America/New_York")) | ||
| start = end - timedelta(days=365 * 6) | ||
|
|
||
| base_url = "https://api.alpaca.markets/v2/calendar" | ||
| params = { | ||
| "start": f"{start.isoformat()}Z", | ||
| "end": f"{end.isoformat()}Z", | ||
| } | ||
| url = f"{base_url}?{urllib.parse.urlencode(params)}" |
There was a problem hiding this comment.
Incorrect use of 'Z' with America/New_York timestamps; send date-only parameters.
You're appending 'Z' (UTC designator) to New York local datetimes, producing invalid timestamps (e.g., ...-04:00Z). The Alpaca Calendar API expects YYYY-MM-DD dates. Use date-only parameters to avoid timezone issues.
Apply this diff:
- end = datetime.now(tz=ZoneInfo("America/New_York"))
- start = end - timedelta(days=365 * 6)
+ end = datetime.now(tz=ZoneInfo("America/New_York"))
+ start = end - timedelta(days=365 * 6)
@@
- params = {
- "start": f"{start.isoformat()}Z",
- "end": f"{end.isoformat()}Z",
- }
+ params = {
+ "start": start.date().isoformat(),
+ "end": end.date().isoformat(),
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| end = datetime.now(tz=ZoneInfo("America/New_York")) | |
| start = end - timedelta(days=365 * 6) | |
| base_url = "https://api.alpaca.markets/v2/calendar" | |
| params = { | |
| "start": f"{start.isoformat()}Z", | |
| "end": f"{end.isoformat()}Z", | |
| } | |
| url = f"{base_url}?{urllib.parse.urlencode(params)}" | |
| end = datetime.now(tz=ZoneInfo("America/New_York")) | |
| start = end - timedelta(days=365 * 6) | |
| base_url = "https://api.alpaca.markets/v2/calendar" | |
| params = { | |
| "start": start.date().isoformat(), | |
| "end": end.date().isoformat(), | |
| } | |
| url = f"{base_url}?{urllib.parse.urlencode(params)}" |
🤖 Prompt for AI Agents
In applications/models/src/models/get_alpaca_calendar.py around lines 11 to 19,
the code currently appends a 'Z' to America/New_York datetimes which yields
invalid mixed timezone strings; change the parameters to date-only strings in
YYYY-MM-DD format (use start.date().isoformat() and end.date().isoformat()),
remove the trailing 'Z', and build the query params with those date strings so
the Alpaca Calendar API receives start and end as date-only values.
| headers = { | ||
| "accept": "application/json", | ||
| "APCA-API-KEY-ID": os.getenv("APCA_API_KEY_ID"), | ||
| "APCA-API-SECRET-KEY": os.getenv("APCA_API_SECRET_KEY"), | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Validate required Alpaca API credentials before making the request.
If env vars are missing, the request will 401 without a clear error. Fail fast with a descriptive message. Also consider a User-Agent for observability.
- headers = {
- "accept": "application/json",
- "APCA-API-KEY-ID": os.getenv("APCA_API_KEY_ID"),
- "APCA-API-SECRET-KEY": os.getenv("APCA_API_SECRET_KEY"),
- }
+ api_key_id = os.getenv("APCA_API_KEY_ID") or os.getenv("ALPACA_API_KEY_ID")
+ api_secret_key = os.getenv("APCA_API_SECRET_KEY") or os.getenv("ALPACA_API_SECRET_KEY")
+ if not api_key_id or not api_secret_key:
+ raise RuntimeError(
+ "Missing Alpaca API credentials. Set APCA_API_KEY_ID/APCA_API_SECRET_KEY "
+ "or ALPACA_API_KEY_ID/ALPACA_API_SECRET_KEY."
+ )
+
+ headers = {
+ "accept": "application/json",
+ "APCA-API-KEY-ID": api_key_id,
+ "APCA-API-SECRET-KEY": api_secret_key,
+ "User-Agent": "pocketsizefund-models/1.0 (+https://github.com/pocketsizefund/pocketsizefund)",
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| headers = { | |
| "accept": "application/json", | |
| "APCA-API-KEY-ID": os.getenv("APCA_API_KEY_ID"), | |
| "APCA-API-SECRET-KEY": os.getenv("APCA_API_SECRET_KEY"), | |
| } | |
| api_key_id = os.getenv("APCA_API_KEY_ID") or os.getenv("ALPACA_API_KEY_ID") | |
| api_secret_key = os.getenv("APCA_API_SECRET_KEY") or os.getenv("ALPACA_API_SECRET_KEY") | |
| if not api_key_id or not api_secret_key: | |
| raise RuntimeError( | |
| "Missing Alpaca API credentials. Set APCA_API_KEY_ID/APCA_API_SECRET_KEY " | |
| "or ALPACA_API_KEY_ID/ALPACA_API_SECRET_KEY." | |
| ) | |
| headers = { | |
| "accept": "application/json", | |
| "APCA-API-KEY-ID": api_key_id, | |
| "APCA-API-SECRET-KEY": api_secret_key, | |
| "User-Agent": "pocketsizefund-models/1.0 (+https://github.com/pocketsizefund/pocketsizefund)", | |
| } |
🤖 Prompt for AI Agents
In applications/models/src/models/get_alpaca_calendar.py around lines 21 to 25,
the code reads Alpaca credentials directly from env vars into headers without
validating them; add an explicit check that APCA_API_KEY_ID and
APCA_API_SECRET_KEY are present (fail fast by raising a clear exception or
logging an error and exiting) before constructing the request headers, and
include a User-Agent header (e.g., app name/version) for observability so
requests fail with a descriptive message instead of a generic 401.
💡 Verification agent
🧩 Analysis chain
Confirm environment variable naming across the repo (APCA_ vs ALPACA_).
This script reads APCA_; other components may use ALPACA_ variants after recent refactors. Ensure consistency to avoid prod misconfig.
Run this to audit env-var usage:
🏁 Script executed:
#!/bin/bash
# Find all Alpaca env-var names
rg -n --pretty $'APCA_API_KEY_ID|APCA_API_SECRET_KEY|ALPACA_API_KEY_ID|ALPACA_API_SECRET_KEY|alpaca_api_key_id|alpaca_api_secret_key'Length of output: 1120
Use ALPACA_ env vars consistently — update get_alpaca_calendar.py*
The repo standard is ALPACA_API_KEY_ID / ALPACA_API_SECRET_KEY; only get_alpaca_calendar.py is using APCA_* env var names and should be updated. HTTP header names ("APCA-API-KEY-ID" / "APCA-API-SECRET-KEY") are correct and should remain unchanged.
Files to change / verify:
- applications/models/src/models/get_alpaca_calendar.py — lines ~23-24: replace APCA_* env var names with ALPACA_*.
- Evidence of ALPACA_* usage elsewhere:
- applications/models/src/models/get_alpaca_equity_bars.py:13 (os.getenv("ALPACA_API_KEY_ID"))
- applications/datamanager/src/datamanager/main.py:60 (os.getenv("ALPACA_API_KEY_ID", ""))
- applications/positionmanager/src/positionmanager/main.py:66,124,262 (os.getenv("ALPACA_API_KEY_ID", ""))
- infrastructure/environment_variables.py:8,22 (require_secret / mapping for "ALPACA_API_KEY_ID")
Suggested change (in get_alpaca_calendar.py):
Replace:
"APCA-API-KEY-ID": os.getenv("APCA_API_KEY_ID"),
"APCA-API-SECRET-KEY": os.getenv("APCA_API_SECRET_KEY"),
With:
"APCA-API-KEY-ID": os.getenv("ALPACA_API_KEY_ID"),
"APCA-API-SECRET-KEY": os.getenv("ALPACA_API_SECRET_KEY"),
Also verify CI/infra secrets and any deployment configs use ALPACA_* (or update them accordingly) to avoid runtime breakage.
🤖 Prompt for AI Agents
In applications/models/src/models/get_alpaca_calendar.py around lines 21 to 25
the code is reading environment variables using the APCA_* names which are
inconsistent with the repo standard; change the env var keys used in os.getenv
to ALPACA_API_KEY_ID and ALPACA_API_SECRET_KEY (leave the HTTP header names
"APCA-API-KEY-ID" and "APCA-API-SECRET-KEY" unchanged) so the headers are
populated from os.getenv("ALPACA_API_KEY_ID") and
os.getenv("ALPACA_API_SECRET_KEY"); after making the change, verify CI/infra
secrets and deployment configs use ALPACA_* names to avoid runtime breakage.
| response = requests.get(url, headers=headers, timeout=10) | ||
|
|
||
| response.raise_for_status() | ||
|
|
||
| body = response.json() | ||
|
|
||
| calendar_data = pl.DataFrame(body) | ||
|
|
||
| calendar_data.write_csv("calendar.csv") | ||
|
|
||
| logger.info("Calendar data has been written to calendar.csv") |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Wrap network call and JSON parsing with error handling and logging.
Surface actionable errors and preserve stack traces. Keep success path unchanged.
- response = requests.get(url, headers=headers, timeout=10)
-
- response.raise_for_status()
-
- body = response.json()
-
- calendar_data = pl.DataFrame(body)
-
- calendar_data.write_csv("calendar.csv")
-
- logger.info("Calendar data has been written to calendar.csv")
+ try:
+ response = requests.get(url, headers=headers, timeout=10)
+ response.raise_for_status()
+ body = response.json()
+ except requests.RequestException as exc:
+ logger.exception(f"Failed to fetch Alpaca calendar: {exc}")
+ raise
+
+ calendar_data = pl.DataFrame(body)
+ calendar_data.write_csv("calendar.csv")
+ logger.info("Calendar data has been written to calendar.csv")Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In applications/models/src/models/get_alpaca_calendar.py around lines 27 to 37,
the network request and JSON parsing are unprotected; wrap the requests.get
call, response.raise_for_status, and response.json() in a try/except that
catches request and JSON errors (requests.exceptions.RequestException,
ValueError/JSONDecodeError, plus a general Exception fallback), log the full
exception with stack trace (using logger.exception or logger.error with
exc_info=True) and surface an actionable error (re-raise or return an error
result) while leaving the existing success path (creating the DataFrame, writing
CSV, and the info log) unchanged.
| def train( | ||
| self, | ||
| inputs_list: list[dict[str, Tensor]], | ||
| epoch_count: int, | ||
| learning_rate: float = 1e-3, | ||
| ) -> dict[str, list[float]]: | ||
| optimizer = Adam(params=self.parameters, lr=learning_rate) | ||
|
|
||
| losses: list[float] = [] | ||
|
|
||
| for _ in range(epoch_count): | ||
| epoch_loss: float = 0.0 | ||
|
|
||
| for inputs in inputs_list: | ||
| outputs = self.forward(inputs) | ||
|
|
||
| loss = quantile_loss( | ||
| outputs["predictions"], | ||
| inputs["targets"], | ||
| self.quantiles, | ||
| ) | ||
|
|
||
| optimizer.zero_grad() | ||
| _ = loss.backward() | ||
| optimizer.step() | ||
|
|
||
| epoch_loss += loss.numpy().item() | ||
|
|
||
| average_epoch_loss: float = epoch_loss / len(inputs_list) | ||
| losses.append(average_epoch_loss) | ||
|
|
||
| return { | ||
| "losses": losses, | ||
| } |
There was a problem hiding this comment.
Use quantile outputs for loss and fold time into batch to match loss expectations.
quantile_loss expects predictions with a quantile dimension and targets without it. Currently you pass predictions (no quantile dim). Flatten the time axis and use outputs["quantiles"].
Apply this diff:
- losses: list[float] = []
+ losses: list[float] = []
@@
- for inputs in inputs_list:
- outputs = self.forward(inputs)
-
- loss = quantile_loss(
- outputs["predictions"],
- inputs["targets"],
- self.quantiles,
- )
+ for inputs in inputs_list:
+ outputs = self.forward(inputs)
+ # reshape to (B*T_out, output_size, Q) and (B*T_out, output_size)
+ q = outputs["quantiles"]
+ t = inputs["targets"]
+ bt = q.shape[0] * q.shape[1]
+ q2 = q.reshape(bt, self.output_size, len(self.quantiles))
+ t2 = t.reshape(bt, self.output_size)
+ loss = quantile_loss(q2, t2, self.quantiles)
@@
- epoch_loss += loss.numpy().item()
+ epoch_loss += float(loss.numpy().item())📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def train( | |
| self, | |
| inputs_list: list[dict[str, Tensor]], | |
| epoch_count: int, | |
| learning_rate: float = 1e-3, | |
| ) -> dict[str, list[float]]: | |
| optimizer = Adam(params=self.parameters, lr=learning_rate) | |
| losses: list[float] = [] | |
| for _ in range(epoch_count): | |
| epoch_loss: float = 0.0 | |
| for inputs in inputs_list: | |
| outputs = self.forward(inputs) | |
| loss = quantile_loss( | |
| outputs["predictions"], | |
| inputs["targets"], | |
| self.quantiles, | |
| ) | |
| optimizer.zero_grad() | |
| _ = loss.backward() | |
| optimizer.step() | |
| epoch_loss += loss.numpy().item() | |
| average_epoch_loss: float = epoch_loss / len(inputs_list) | |
| losses.append(average_epoch_loss) | |
| return { | |
| "losses": losses, | |
| } | |
| def train( | |
| self, | |
| inputs_list: list[dict[str, Tensor]], | |
| epoch_count: int, | |
| learning_rate: float = 1e-3, | |
| ) -> dict[str, list[float]]: | |
| optimizer = Adam(params=self.parameters, lr=learning_rate) | |
| losses: list[float] = [] | |
| for _ in range(epoch_count): | |
| epoch_loss: float = 0.0 | |
| for inputs in inputs_list: | |
| outputs = self.forward(inputs) | |
| # reshape to (B*T_out, output_size, Q) and (B*T_out, output_size) | |
| q = outputs["quantiles"] | |
| t = inputs["targets"] | |
| bt = q.shape[0] * q.shape[1] | |
| q2 = q.reshape(bt, self.output_size, len(self.quantiles)) | |
| t2 = t.reshape(bt, self.output_size) | |
| loss = quantile_loss(q2, t2, self.quantiles) | |
| optimizer.zero_grad() | |
| _ = loss.backward() | |
| optimizer.step() | |
| epoch_loss += float(loss.numpy().item()) | |
| average_epoch_loss: float = epoch_loss / len(inputs_list) | |
| losses.append(average_epoch_loss) | |
| return { | |
| "losses": losses, | |
| } |
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 191 to 224, the
training loop passes outputs["predictions"] (no quantile dimension) to
quantile_loss and does not fold the time axis into the batch, causing a shape
mismatch with quantile_loss which expects predictions with a quantile dim and
targets without it; update the code to use outputs["quantiles"] instead of
outputs["predictions"] and reshape both predictions and targets by merging batch
and time dimensions before calling quantile_loss (e.g., reshape quantile
predictions to (batch*time, num_quantiles) and targets to (batch*time,) or
(batch*time, 1) as required), then compute loss, backpropagate, and continue as
before.
| def validate( | ||
| self, | ||
| inputs_list: list[dict[str, Tensor]], | ||
| ) -> float: | ||
| total_loss = 0.0 | ||
| batch_count = 0 | ||
|
|
||
| for inputs in inputs_list: | ||
| outputs = self.forward(inputs) | ||
|
|
||
| loss = quantile_loss( | ||
| outputs["predictions"], | ||
| inputs["targets"], | ||
| self.quantiles, | ||
| ) | ||
|
|
||
| total_loss += loss.numpy().item() | ||
| batch_count += 1 | ||
|
|
||
| return total_loss / batch_count | ||
|
|
There was a problem hiding this comment.
Validation path mirrors training bug; use quantiles and fold time.
Same shape issue as train(). Fix accordingly.
Apply this diff:
- for inputs in inputs_list:
- outputs = self.forward(inputs)
-
- loss = quantile_loss(
- outputs["predictions"],
- inputs["targets"],
- self.quantiles,
- )
-
- total_loss += loss.numpy().item()
+ for inputs in inputs_list:
+ outputs = self.forward(inputs)
+ q = outputs["quantiles"]
+ t = inputs["targets"]
+ bt = q.shape[0] * q.shape[1]
+ q2 = q.reshape(bt, self.output_size, len(self.quantiles))
+ t2 = t.reshape(bt, self.output_size)
+ loss = quantile_loss(q2, t2, self.quantiles)
+ total_loss += float(loss.numpy().item())Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In libraries/python/src/internal/tft_model.py around lines 226 to 246, the
validate() method currently mirrors the training bug by passing mismatched
shapes to quantile_loss; fix it by using the model's quantiles and folding the
time dimension the same way train() does: for each batch slice or inputs,
select/slice the forecast time fold from outputs["predictions"] so its time axis
matches inputs["targets"], reshape or expand inputs["targets"] to include the
quantile axis (or transpose predictions to [batch, time, quantile]) and then
call quantile_loss with self.quantiles so shapes align before accumulating loss
and averaging over batches.
| class VariableSelectionNetwork: | ||
| def __init__(self, input_dimension: int, hidden_size: int) -> None: | ||
| self.input_layer = nn.Linear( | ||
| in_features=input_dimension, | ||
| out_features=hidden_size, | ||
| ) | ||
|
|
||
| self.output_layer = nn.Linear( | ||
| in_features=hidden_size, | ||
| out_features=input_dimension, | ||
| ) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Store input_dimension and validate shapes early
Capturing the declared input dimension and validating inputs at runtime helps catch silent shape mismatches, especially when integrating into TFT where concatenation dims can drift.
class VariableSelectionNetwork:
def __init__(self, input_dimension: int, hidden_size: int) -> None:
- self.input_layer = nn.Linear(
+ self.input_dimension = input_dimension
+ self.hidden_size = hidden_size
+
+ self.input_layer = Linear(
in_features=input_dimension,
out_features=hidden_size,
)
- self.output_layer = nn.Linear(
+ self.output_layer = Linear(
in_features=hidden_size,
out_features=input_dimension,
)Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In libraries/python/src/internal/variable_selection_network.py around lines 5 to
16, the class constructor does not persist the declared input_dimension nor
validate incoming tensor shapes; update __init__ to store self.input_dimension =
input_dimension (and optionally self.hidden_size = hidden_size) and then in the
forward method validate that the input tensor's last dimension equals
self.input_dimension, raising a clear ValueError if it does not so shape
mismatches are caught early; ensure any other internal layers that depend on
input_dimension use the stored attribute for consistency.
| @@ -1,19 +1,16 @@ | |||
| import numpy as np | |||
| import pytest | |||
| from internal.loss_functions import quantile_loss | |||
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider adding validation in quantile_loss for predictions vs quantiles size (and fix accumulation).
The internal implementation would benefit from:
- Validating
predictions.shape[-1] == len(quantiles)with a clear error. - Accumulating loss correctly instead of ignoring
add’s return value.
Proposed change in libraries/python/src/internal/loss_functions.py (outside this diff):
def quantile_loss(
predictions: Tensor, # shape: (batch_size, output_size, len(quantiles))
targets: Tensor, # shape: (batch_size, output_size)
quantiles: list[float] | None = None,
) -> Tensor:
if quantiles is None:
quantiles = [0.1, 0.5, 0.9]
if not all(0 <= q <= 1 for q in quantiles):
message = "All quantiles must be between 0 and 1"
raise ValueError(message)
+ if predictions.shape[-1] != len(quantiles):
+ message = (
+ f"Predictions last dimension ({predictions.shape[-1]}) must equal "
+ f"number of quantiles ({len(quantiles)})"
+ )
+ raise ValueError(message)
+
- errors_total = Tensor.zeros((1,))
+ errors_total = Tensor.zeros((1,))
for index, quantile in enumerate(quantiles):
error = targets.sub(predictions[:, :, index])
quantile_tensor = Tensor([quantile])
- errors_total.add(
+ errors_total = errors_total.add(
Tensor.where(
error > 0,
quantile_tensor.mul(error),
(quantile_tensor.sub(1)).mul(error),
).mean()
)
return errors_total.div(len(quantiles)) # shape: (1,)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from internal.loss_functions import quantile_loss | |
| def quantile_loss( | |
| predictions: Tensor, # shape: (batch_size, output_size, len(quantiles)) | |
| targets: Tensor, # shape: (batch_size, output_size) | |
| quantiles: list[float] | None = None, | |
| ) -> Tensor: | |
| if quantiles is None: | |
| quantiles = [0.1, 0.5, 0.9] | |
| if not all(0 <= q <= 1 for q in quantiles): | |
| message = "All quantiles must be between 0 and 1" | |
| raise ValueError(message) | |
| if predictions.shape[-1] != len(quantiles): | |
| message = ( | |
| f"Predictions last dimension ({predictions.shape[-1]}) must equal " | |
| f"number of quantiles ({len(quantiles)})" | |
| ) | |
| raise ValueError(message) | |
| errors_total = Tensor.zeros((1,)) | |
| for index, quantile in enumerate(quantiles): | |
| error = targets.sub(predictions[:, :, index]) | |
| quantile_tensor = Tensor([quantile]) | |
| errors_total = errors_total.add( | |
| Tensor.where( | |
| error > 0, | |
| quantile_tensor.mul(error), | |
| (quantile_tensor.sub(1)).mul(error), | |
| ).mean() | |
| ) | |
| return errors_total.div(len(quantiles)) # shape: (1,) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/loss_functions.py around the quantile_loss
function (update the implementation referenced by tests in
libraries/python/tests/test_loss_functions.py line 3), validate that
predictions.shape[-1] == len(quantiles) and raise a ValueError with a clear
message if not, and fix the accumulation by using the return value from add
(e.g., loss = loss + add(... ) or loss += add(...)) so each quantile
contribution is actually accumulated into the total loss.
| @@ -3,39 +3,19 @@ name = "pocketsizefund" | |||
| version = "20250602.4" | |||
| description = "Open source quantitative hedge fund 🍊" | |||
| requires-python = "==3.12.10" | |||
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid pinning requires-python to a patch release
Pinning to "==3.12.10" is brittle and can make local/dev CI setups unnecessarily hard. Prefer a compatible range (e.g., >=3.12,<3.13) unless you truly require a specific patch.
-requires-python = "==3.12.10"
+requires-python = ">=3.12,<3.13"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| requires-python = "==3.12.10" | |
| requires-python = ">=3.12,<3.13" |
🤖 Prompt for AI Agents
In pyproject.toml around line 5, the requires-python field is pinned to a
specific patch ("==3.12.10"), which is brittle; change it to a compatible range
such as ">=3.12,<3.13" (or another appropriate minor-range) to allow any 3.12.x
interpreter while preventing upgrades to 3.13, updating the requires-python
value accordingly in the file.
There was a problem hiding this comment.
Actionable comments posted: 4
🔭 Outside diff range comments (1)
libraries/python/src/internal/lstm_network.py (1)
68-73: Dropout applied incorrectly and toggling training on a Tensor is invalid
- Calling train() on a Tensor is not a valid tinygrad pattern (training mode is global via Tensor.training or managed at module level).
- Dropout should affect the input to the next layer, not mutate the retained hidden_state itself. Currently, layer_input is set to the non-dropped tensor, so dropout has no effect downstream.
- if self.dropout_rate > 0.0 and index < self.layer_count - 1: - hidden_state[index].train() - hidden_state[index] = hidden_state[index].dropout(self.dropout_rate) - - layer_input = layer_hidden_state + if self.dropout_rate > 0.0 and index < self.layer_count - 1 and Tensor.training: + # Apply dropout to the inter-layer signal only (do not corrupt the recurrent state). + layer_input = layer_hidden_state.dropout(self.dropout_rate) + else: + layer_input = layer_hidden_state
♻️ Duplicate comments (5)
libraries/python/src/internal/loss_functions.py (1)
9-15: Add essential input/shape validation to prevent silent broadcasting and index errors.Ensure non-empty quantiles and validate shapes to match (B, O, Q) vs (B, O). This prevents out-of-bounds indexing at Line 18 and subtle miscomputations due to broadcasting.
Apply:
if quantiles is None: quantiles = [0.1, 0.5, 0.9] + if len(quantiles) == 0: + raise ValueError("quantiles must be non-empty") + if not all(0 <= q <= 1 for q in quantiles): message = "All quantiles must be between 0 and 1" raise ValueError(message) + # Validate shapes: predictions (B, O, Q), targets (B, O) + if len(predictions.shape) != 3: + raise ValueError("predictions must be 3D: (batch_size, output_size, num_quantiles)") + if len(targets.shape) != 2: + raise ValueError("targets must be 2D: (batch_size, output_size)") + if predictions.shape[0] != targets.shape[0] or predictions.shape[1] != targets.shape[1]: + raise ValueError(f"Shape mismatch: predictions {predictions.shape[:2]} vs targets {targets.shape}") + if predictions.shape[2] != len(quantiles): + raise ValueError( + f"Quantile count mismatch: predictions last dim={predictions.shape[2]} vs len(quantiles)={len(quantiles)}" + )libraries/python/src/internal/lstm_network.py (1)
13-16: Don’t hard-enforce layer_count >= 3; allow 1–2 (warn instead)
This breaks common LSTM configs and back-compat. Keep default=3 if desired, but only error for <1 and warn for 1–2.- minimum_layer_count = 3 - if layer_count < minimum_layer_count: - message = f"Layer count must be at least {minimum_layer_count}" - raise ValueError(message) + if layer_count < 1: + raise ValueError("layer_count must be at least 1") + if layer_count < 3: + import warnings + warnings.warn( + "Using fewer than 3 LSTM layers is supported but may reduce model capacity; " + "defaults and some models expect 3 layers.", + UserWarning, + )applications/datamanager/src/datamanager/clients.py (2)
47-55: Normalize calendar fields: handle None and prefer ISO format; remain robust if raw_data is dictsReturning
str(None)produces "None". Prefer""for missing values, and.isoformat()for date/time objects. Also tolerate raw dict entries ifraw_data=Trueslips through.- calendar_list: list[dict[str, str]] = [] - for calendar_entry in calendar_data: - calendar_dict = { - "date": str(calendar_entry.date), - "open": str(calendar_entry.open), - "close": str(calendar_entry.close), - } - calendar_list.append(calendar_dict) + calendar_list: list[dict[str, str]] = [] + for entry in calendar_data: + # Support both model objects and raw dicts + if isinstance(entry, dict): + date_val = entry.get("date") + open_val = entry.get("open") + close_val = entry.get("close") + else: + date_val = getattr(entry, "date", None) + open_val = getattr(entry, "open", None) + close_val = getattr(entry, "close", None) + + def fmt(val: object) -> str: + return ( + val.isoformat() # type: ignore[attr-defined] + if hasattr(val, "isoformat") + else (str(val) if val is not None else "") + ) + + calendar_list.append( + { + "date": fmt(date_val), + "open": fmt(open_val), + "close": fmt(close_val), + } + )
33-36: Set paper=True and raw_data=False on TradingClient to avoid live trading and ensure model objectsTo prevent accidental live trading and guarantee attribute access (vs raw dicts), initialize
TradingClientwith explicitpaper=Trueandraw_data=False.- self.trading_client: TradingClient = TradingClient( - api_key=alpaca_api_key_id, - secret_key=alpaca_api_secret, - ) + self.trading_client: TradingClient = TradingClient( + api_key=alpaca_api_key_id, + secret_key=alpaca_api_secret, + paper=True, + raw_data=False, + )applications/datamanager/src/datamanager/main.py (1)
273-284: Fix broken exception block (syntax error) and return consistent 500 responseThere’s a stray logger call inside a Response constructor and two competing returns. CI failures cite unclosed parenthesis and invalid argument order. Replace the block with a single, coherent 500 response.
try: blobs = list(s3_client.list_objects(prefix=prefix)) - except Exception as e: # noqa: BLE001 - logger.error(f"Error listing S3 objects: {e}") - return Response( - status_code=status.HTTP_404_NOT_FOUND, - logger.error(f"Error listing equity bars: {e}") - return Response( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - content=json.dumps({"error": f"Failed to list equity bars: {e!s}"}), - media_type="application/json", - ) + except Exception as e: # noqa: BLE001 + logger.error(f"Error listing equity bars: {e}") + return Response( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content=json.dumps({"error": f"Failed to list equity bars: {e!s}"}), + media_type="application/json", + )
🧹 Nitpick comments (4)
libraries/python/src/internal/loss_functions.py (3)
1-8: Generalize type annotation to accept tuples and other sequences.Many callers use tuples; widen the annotation to Sequence[float] to avoid type friction.
Apply:
-from tinygrad.tensor import Tensor +from typing import Sequence +from tinygrad.tensor import Tensor @@ - quantiles: list[float] | None = None, + quantiles: Sequence[float] | None = None,
16-28: Optional: vectorize across quantiles to reduce Python overhead and Tensor allocations.Removes the Python loop and constructs a broadcasted quantile tensor. Semantics preserve the current “mean across batch/output and quantiles.” If you revert to sum semantics, multiply the final mean by len(quantiles).
Apply:
- errors_total = Tensor.zeros((1,)) - for index, quantile in enumerate(quantiles): - error = targets.sub(predictions[:, :, index]) - quantile_tensor = Tensor([quantile]) - errors_total = errors_total.add( - Tensor.where( - error > 0, - quantile_tensor.mul(error), - (quantile_tensor.sub(1)).mul(error), - ).mean() - ) - - return errors_total.div(len(quantiles)) # shape: (1,) + # Vectorized over quantiles: predictions (B, O, Q), targets (B, O, 1), quantiles (1, 1, Q) + quantiles_tensor = Tensor(quantiles).reshape(1, 1, -1) + error = targets.reshape(targets.shape[0], targets.shape[1], 1).sub(predictions) + loss_per_elem = Tensor.where( + error > 0, + quantiles_tensor.mul(error), + quantiles_tensor.sub(1).mul(error), + ) + return loss_per_elem.mean() # mean over batch, output, and quantiles
4-8: Nit: add a short docstring clarifying expected shapes and the loss formula.Improves readability for new contributors.
Apply:
def quantile_loss( predictions: Tensor, # shape: (batch_size, output_size, len(quantiles)) targets: Tensor, # shape: (batch_size, output_size) quantiles: list[float] | None = None, ) -> Tensor: + """ + Compute the multi-quantile pinball loss. + Expected shapes: + - predictions: (B, O, Q) one prediction per quantile + - targets: (B, O) + - quantiles: sequence of floats in [0, 1] + Returns a (1,) tensor containing the mean loss across batch/output and quantiles. + """libraries/python/src/internal/lstm_network.py (1)
39-47: Initialize state with matching device/dtype as inputs
Creating zeros without matching inputs’ device/dtype can cause subtle bugs (mixed device/dtype). Prefer initializing with inputs’ properties (e.g., using device=inputs.device, dtype=inputs.dtype or an inputs-backed helper, if available).Can you confirm tinygrad’s Tensor.zeros supports device/dtype kwargs in your pinned version? If unsure, I can help propose a safe helper to mirror inputs’ properties.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (7)
.claude/settings.local.json(1 hunks).github/workflows/teardown_infrastructure.yaml(1 hunks)applications/datamanager/src/datamanager/clients.py(2 hunks)applications/datamanager/src/datamanager/main.py(5 hunks)applications/models/pyproject.toml(1 hunks)libraries/python/src/internal/loss_functions.py(1 hunks)libraries/python/src/internal/lstm_network.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- applications/models/pyproject.toml
- .claude/settings.local.json
- .github/workflows/teardown_infrastructure.yaml
🧰 Additional context used
🧬 Code Graph Analysis (4)
libraries/python/src/internal/loss_functions.py (3)
application/predictionengine/src/predictionengine/loss_function.py (1)
quantile_loss(8-29)application/predictionengine/tests/test_loss_function.py (7)
test_quantile_loss_multiple_samples(24-32)test_quantile_loss_basic(13-21)test_quantile_loss_shapes(56-63)test_quantile_loss_different_quantiles(45-53)test_quantile_loss_perfect_prediction(35-42)test_quantile_loss_shape_mismatch(66-72)test_quantile_loss_invalid_quantiles(75-81)application/predictionengine/src/predictionengine/miniature_temporal_fusion_transformer.py (1)
validate(148-161)
libraries/python/src/internal/lstm_network.py (2)
application/predictionengine/src/predictionengine/long_short_term_memory.py (3)
LongShortTermMemory(5-69)forward(22-69)__init__(6-20)application/predictionengine/tests/test_long_short_term_memory.py (4)
test_lstm_forward(31-43)test_lstm_multiple_layers(58-67)test_lstm_initialization(14-28)test_lstm_single_timestep(70-78)
applications/datamanager/src/datamanager/main.py (3)
libraries/python/src/internal/cloud_event.py (2)
create_cloud_event_error(24-38)create_cloud_event_success(7-21)libraries/python/src/internal/dates.py (1)
Date(8-31)applications/datamanager/src/datamanager/clients.py (7)
AlpacaClient(27-27)AlpacaClient(28-56)PolygonClient(14-24)S3Client(59-83)get_all_equity_bars(18-24)list_objects(65-73)get_calendar(37-56)
applications/datamanager/src/datamanager/clients.py (2)
applications/datamanager/src/datamanager/main.py (1)
get_calendar(298-335)application/positionmanager/src/positionmanager/clients.py (2)
AlpacaClient(16-103)__init__(17-33)
🪛 Ruff (0.12.2)
applications/datamanager/src/datamanager/main.py
279-279: SyntaxError: Positional argument cannot follow keyword argument
279-280: SyntaxError: Expected ')', found newline
applications/datamanager/src/datamanager/clients.py
28-28: SyntaxError: Expected an indented block after class definition
🪛 GitHub Actions: Test and coverage check
applications/datamanager/src/datamanager/main.py
[error] 277-277: SyntaxError: '(' was never closed. Likely a missing closing parenthesis in a return statement.
🪛 GitHub Actions: Code quality check
applications/datamanager/src/datamanager/main.py
[error] 279-279: Ruff format failed. Command: 'ruff format'. Parse error in applications/datamanager/src/datamanager/main.py:279:9 - Positional argument cannot follow keyword argument.
applications/datamanager/src/datamanager/clients.py
[error] 28-28: Ruff format failed. Command: 'ruff format'. Parse error in applications/datamanager/src/datamanager/clients.py:28:1 - Expected an indented block after class definition.
🔇 Additional comments (10)
libraries/python/src/internal/loss_functions.py (3)
9-11: Confirm change to default quantiles (from 0.25/0.5/0.75 to 0.1/0.5/0.9).This is a behavioral change from the previous implementation (predictionengine) and will affect training dynamics for any caller not explicitly passing quantiles.
If this wasn’t intentional, revert to the prior defaults:
- if quantiles is None: - quantiles = [0.1, 0.5, 0.9] + if quantiles is None: + quantiles = [0.25, 0.5, 0.75]
28-28: Confirm intentional change: returning the mean across quantiles (new) vs sum (previous).The previous implementation summed per-quantile means; this one averages across quantiles. That changes loss scale by a factor of len(quantiles) and can impact training unless learning rates/expectations were adjusted.
If you want backward compatibility with the old behavior, apply:
- return errors_total.div(len(quantiles)) # shape: (1,) + return errors_total # shape: (1,)
20-26: Good: accumulation uses the returned Tensor from .add().Tinygrad ops return new tensors; reassigning to errors_total is the correct pattern.
libraries/python/src/internal/lstm_network.py (3)
25-30: Constructor now uses explicit kwargs for LSTMCell — good change
This improves readability and guards against positional arg ordering changes.
76-78: Good guard against empty timestep outputs
Raising early prevents a later stack/indexing failure and keeps the error message clear.
32-36: Forward signature/state shape changed — call sites verified, no action requiredChecked the repo: the LSTM.forward signature returns/accepts (hidden_state:list[Tensor], cell_state:list[Tensor]) and the only call sites are already using the tuple-of-two convention.
Files inspected:
- libraries/python/src/internal/lstm_network.py — forward signature returns (hidden_state, cell_state) as lists.
- libraries/python/src/internal/tft_model.py — uses:
encoder_output, (h_n, c_n) = self.lstm_encoder.forward(...)and passes(h_n, c_n)into the decoder.- libraries/python/tests/test_lstm_network.py — tests call
lstm.forward(...)and assert the returned state is a tuple of length 2.No other usages of LSTM.forward or
state=were found. Since callers and tests already match the new per-layer list shape, no changes are required.applications/datamanager/src/datamanager/clients.py (1)
29-32: Good: fail fast on missing Alpaca credentialsEarly validation with a clear
ValueErroris appropriate; this prevents confusing runtime failures.applications/datamanager/src/datamanager/main.py (3)
59-63: Good: Alpaca client wired at startup with env-driven credentialsThis aligns with fail-fast behavior given the constructor’s validation. Ensure deployment envs set both variables to avoid startup failure.
216-220: Sane default for fetch_dateDefaulting to “today” in America/New_York when
fetch_dateis None is sensible and keeps behavior explicit.
297-336: Endpoint looks good; status handling and logging are clearThe
/calendarendpoint behavior is straightforward: 404 on empty, 200 with JSON on success, 500 on exception. Once the client’s calendar serialization is hardened (see clients.py), this route should be robust.
| class AlpacaClient: | ||
| class AlpacaClient: | ||
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: |
There was a problem hiding this comment.
Fix duplicated class declaration causing SyntaxError (CI is failing)
There are two consecutive class AlpacaClient: lines. Ruff/CI error confirms a parse error at Line 28. Remove the duplicate declaration.
Apply:
-class AlpacaClient:
class AlpacaClient:
def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| class AlpacaClient: | |
| class AlpacaClient: | |
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: | |
| class AlpacaClient: | |
| def __init__(self, alpaca_api_key_id: str, alpaca_api_secret: str) -> None: |
🧰 Tools
🪛 Ruff (0.12.2)
28-28: SyntaxError: Expected an indented block after class definition
🪛 GitHub Actions: Code quality check
[error] 28-28: Ruff format failed. Command: 'ruff format'. Parse error in applications/datamanager/src/datamanager/clients.py:28:1 - Expected an indented block after class definition.
🤖 Prompt for AI Agents
In applications/datamanager/src/datamanager/clients.py around lines 27 to 29,
there are two consecutive identical class declarations "class AlpacaClient:"
causing a SyntaxError; remove the duplicate declaration so only a single "class
AlpacaClient:" remains before the __init__ definition and ensure indentation of
the __init__ method follows that single class declaration.
| return create_cloud_event_success( | ||
| application_name="datamanager", | ||
| event_metadata=["equity", "bars", "created"], | ||
| data={ | ||
| "date": request_summary_date, | ||
| "date": request_fetch_date, | ||
| "count": count, | ||
| }, | ||
| ) |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Avoid overriding CloudEvent’s timestamp: rename “date” in data payload
create_cloud_event_success already injects a "date" timestamp into the data payload. Providing "date" here will override it. Use a distinct key (e.g., "request_date").
- return create_cloud_event_success(
- application_name="datamanager",
- event_metadata=["equity", "bars", "created"],
- data={
- "date": request_fetch_date,
- "count": count,
- },
- )
+ return create_cloud_event_success(
+ application_name="datamanager",
+ event_metadata=["equity", "bars", "created"],
+ data={
+ "request_date": request_fetch_date,
+ "count": count,
+ },
+ )Would you like me to sweep the codebase for other CloudEvent payloads that might be overriding this reserved key?
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Searching for occurrences of 'create_cloud_event_success' ==="
rg -n --hidden --no-ignore -S "create_cloud_event_success" || true
echo
echo "=== Searching for definition 'def create_cloud_event_success' ==="
rg -n --hidden --no-ignore -S "def create_cloud_event_success" || true
echo
echo "=== Showing context (±8 lines) around matches ==="
rg -n -C8 --hidden --no-ignore -S "create_cloud_event_success" || true
# If files found, list the files and show occurrences of 'date' inside those files to check whether the helper inserts a 'date' key
files=$(rg -l --hidden --no-ignore -S "create_cloud_event_success" || true)
if [ -n "$files" ]; then
echo
echo "=== Inspecting 'date' mentions in files that reference create_cloud_event_success ==="
for f in $files; do
echo "---- $f ----"
rg -n --hidden --no-ignore -S "date" "$f" || true
echo
done
else
echo
echo "No files reference 'create_cloud_event_success' (or rg didn't find matches)."
fiLength of output: 14304
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Showing libraries/python/src/internal/cloud_event.py ==="
nl -ba libraries/python/src/internal/cloud_event.py | sed -n '1,240p'Length of output: 313
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== libraries/python/src/internal/cloud_event.py ==="
sed -n '1,240p' libraries/python/src/internal/cloud_event.py || trueLength of output: 1191
Avoid overriding CloudEvent's timestamp — rename payload "date" to "request_date"
create_cloud_event_success injects "date" into event.data and merges caller data after it ({"date": ..., **data}), so any "date" key you pass will override the generated timestamp. Rename the payload key to avoid accidental overrides.
Files to review:
- libraries/python/src/internal/cloud_event.py — create_cloud_event_success sets "date" before merging caller data.
- applications/datamanager/src/datamanager/main.py (lines ~255–262) — currently passes "date": request_fetch_date.
- applications/positionmanager/src/positionmanager/main.py (lines ~245–252 and ~279–286) — passes "date" in event payloads.
Suggested change (datamanager):
- return create_cloud_event_success(
- application_name="datamanager",
- event_metadata=["equity", "bars", "created"],
- data={
- "date": request_fetch_date,
- "count": count,
- },
- )
+ return create_cloud_event_success(
+ application_name="datamanager",
+ event_metadata=["equity", "bars", "created"],
+ data={
+ "request_date": request_fetch_date,
+ "count": count,
+ },
+ )Also review the positionmanager occurrences and decide whether those should intentionally override the CloudEvent timestamp or be renamed (e.g., event_date/request_date).
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return create_cloud_event_success( | |
| application_name="datamanager", | |
| event_metadata=["equity", "bars", "created"], | |
| data={ | |
| "date": request_summary_date, | |
| "date": request_fetch_date, | |
| "count": count, | |
| }, | |
| ) | |
| return create_cloud_event_success( | |
| application_name="datamanager", | |
| event_metadata=["equity", "bars", "created"], | |
| data={ | |
| "request_date": request_fetch_date, | |
| "count": count, | |
| }, | |
| ) |
🤖 Prompt for AI Agents
In applications/datamanager/src/datamanager/main.py around lines 255 to 262, the
payload is passing "date": request_fetch_date which will override the
CloudEvent's generated timestamp; rename the payload key to "request_date" (or
similar, e.g., event_date) and update the data dict to use that new key so
create_cloud_event_success can safely inject its own "date"; additionally scan
the two occurrences in applications/positionmanager/src/positionmanager/main.py
(around lines 245–252 and 279–286) and either rename those payload keys or
confirm they intentionally override the event timestamp, then update
callers/tests accordingly.
| if len(outputs) == 1: | ||
| output_tensor = outputs[0].unsqueeze(1) | ||
| output_tensor = outputs[0].unsqueeze(dim=1) | ||
| else: | ||
| output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Prefer positional arg for unsqueeze to avoid API mismatch
tinygrad’s unsqueeze often mirrors PyTorch’s signature but keyword support isn’t guaranteed; using a positional arg is safer across versions.
- if len(outputs) == 1:
- output_tensor = outputs[0].unsqueeze(dim=1)
+ if len(outputs) == 1:
+ output_tensor = outputs[0].unsqueeze(1)Run this to see if keyword-based unsqueeze is used elsewhere (and decide whether to standardize):
🏁 Script executed:
#!/bin/bash
rg -n -S "unsqueeze\\(" -A 0 -B 2Length of output: 1054
Prefer positional arg for unsqueeze to avoid API mismatch
tinygrad may not accept the keyword form (dim=…); standardize to positional calls. My search found one keyword usage; other occurrences already use positional.
- libraries/python/src/internal/lstm_network.py:81 — change outputs[0].unsqueeze(dim=1) to outputs[0].unsqueeze(1)
- libraries/python/src/internal/tft_model.py:145 — static_context.unsqueeze(1) (already positional)
- libraries/python/src/internal/tft_model.py:163 — static_context.unsqueeze(1) (already positional)
Suggested diff:
- if len(outputs) == 1:
- output_tensor = outputs[0].unsqueeze(dim=1)
+ if len(outputs) == 1:
+ output_tensor = outputs[0].unsqueeze(1)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if len(outputs) == 1: | |
| output_tensor = outputs[0].unsqueeze(1) | |
| output_tensor = outputs[0].unsqueeze(dim=1) | |
| else: | |
| output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) | |
| if len(outputs) == 1: | |
| output_tensor = outputs[0].unsqueeze(1) | |
| else: | |
| output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) |
🤖 Prompt for AI Agents
In libraries/python/src/internal/lstm_network.py around lines 80 to 83, the call
uses a keyword argument unsqueeze(dim=1) which may be incompatible with
tinygrad's API; change outputs[0].unsqueeze(dim=1) to the positional form
outputs[0].unsqueeze(1) so it matches other usages and avoids the API mismatch.
cd2b10c to
17c9a34
Compare
forstmeier
force-pushed
the
07-28-restructure_python_application_and_library_resources
branch
from
17c9a34 to
fd0bad4
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🔭 Outside diff range comments (2)
libraries/python/src/internal/lstm_network.py (1)
68-73: Fix dropout: currently mutates recurrent state and isn’t applied to inter-layer input
- Dropout is applied to hidden_state (corrupts recurrent state for the next timestep) and is not used as the input to the next layer (layer_input remains pre-dropout). This breaks the intended “between-layers” dropout semantics.
- Calling train() on a Tensor is not a valid tinygrad pattern; rely on global training mode if needed.
Apply dropout to layer_input and leave hidden/cell states untouched:
- if self.dropout_rate > 0.0 and index < self.layer_count - 1: - hidden_state[index].train() - hidden_state[index] = hidden_state[index].dropout(self.dropout_rate) - - layer_input = layer_hidden_state + layer_input = layer_hidden_state + if self.dropout_rate > 0.0 and index < self.layer_count - 1: + # Apply dropout between layers, not on the recurrent state. + layer_input = layer_input.dropout(self.dropout_rate)infrastructure/environment_variables.py (1)
8-32: Rename infra secret: ALPACA_API_SECRET → ALPACA_API_SECRET_KEYInfra currently provides ALPACA_API_SECRET but application code reads ALPACA_API_SECRET_KEY — this mismatch will cause runtime failures. Update infra or callers so names match.
Files to update:
- infrastructure/environment_variables.py — change the secret name and the mapping.
- applications/models/src/models/get_alpaca_equity_bars.py — reads os.getenv("ALPACA_API_SECRET_KEY")
- applications/models/src/models/get_alpaca_calendar.py — reads os.getenv("ALPACA_API_SECRET_KEY")
Suggested minimal change (infra):
--- a/infrastructure/environment_variables.py +++ b/infrastructure/environment_variables.py @@ - alpaca_api_secret = configuration.require_secret("ALPACA_API_SECRET") + alpaca_api_secret = configuration.require_secret("ALPACA_API_SECRET_KEY") @@ - ("ALPACA_API_SECRET", alpaca_api_secret), + ("ALPACA_API_SECRET_KEY", alpaca_api_secret),Also ensure any deployment/configuration (Pulumi config, secret manager entries) is updated to use ALPACA_API_SECRET_KEY.
♻️ Duplicate comments (17)
Dockerfile.tests (1)
1-1: Align base image Python version with requires-python across workspace (likely 3.12.10)
Multiple pyprojects previously declaredrequires-python = "==3.12.10". Usingpython:3.13can cause uv resolution to fail under strict constraints. If the workspace still pins 3.12.10, switch the base image; if you upgraded the constraint to 3.13, ignore this.Proposed change:
-FROM python:3.13 +FROM python:3.12.10Optionally slim:
-FROM python:3.13 +FROM python:3.12.10-slimVerify the constraints:
#!/bin/bash set -euo pipefail echo "requires-python across workspace:" rg -n --type toml --glob '**/pyproject.toml' '^\s*requires-python\s*=.*' -n -C1libraries/python/tests/test_loss_functions.py (1)
53-61: Seed the RNG for reproducible tests; add negative tests for shape/quantile-count mismatches
- Reproducibility: As written, tests are nondeterministic across runs. Seed the generator once so failures are reproducible.
- Negative tests: You removed shape-mismatch tests, and the implementation currently doesn’t validate count/shape mismatches. Reintroducing these tests will catch silent broadcasting or indexing errors and guard future refactors.
Seed suggestion (outside this hunk, near Line 7):
rng = Generator(PCG64(42))Add these tests to the bottom of the file (pairs with adding validation in the implementation; see next comment):
def test_quantile_loss_quantile_count_mismatch_raises() -> None: predictions = Tensor([[[1.0, 1.2, 1.3]]]) # 3 quantiles in predictions targets = Tensor([[2.0]]) quantiles = [0.25, 0.5] # only 2 quantiles provided with pytest.raises(ValueError, match="Quantile count mismatch"): quantile_loss(predictions, targets, quantiles) def test_quantile_loss_shape_mismatch_raises() -> None: predictions = Tensor([[[1.0, 1.2, 1.3]]]) # (1,1,3) targets = Tensor([[2.0], [3.0]]) # (2,1) mismatched batch quantiles = [0.25, 0.5, 0.75] with pytest.raises(ValueError, match="Shape mismatch"): quantile_loss(predictions, targets, quantiles)If you want, I can open a follow-up PR or push a commit wiring these in alongside the implementation validation.
libraries/python/src/internal/lstm_network.py (3)
80-82: Use positional arg for unsqueeze to avoid tinygrad API mismatchPrefer positional form; some tinygrad versions don’t accept dim= keyword here.
- output_tensor = outputs[0].unsqueeze(dim=1) + output_tensor = outputs[0].unsqueeze(1)
13-16: Don’t hard-enforce layer_count >= 3 — allow >= 1 (optionally warn for < 3)Enforcing a minimum of 3 is a breaking change and blocks common 1–2 layer LSTMs. Keep default=3 but allow shallower configs; optionally warn.
- minimum_layer_count = 3 - if layer_count < minimum_layer_count: - message = f"Layer count must be at least {minimum_layer_count}" - raise ValueError(message) + if layer_count < 1: + raise ValueError("layer_count must be at least 1") + if layer_count < 3: + import warnings + warnings.warn( + "Using fewer than 3 LSTM layers is supported but may reduce capacity; " + "defaults and some models expect 3.", + UserWarning, + )If you intentionally require 3+, please confirm all call sites/tests never pass 1–2. To verify:
#!/bin/bash set -euo pipefail rg -nP '\bLSTM\s*\(' -C3 rg -nP 'layer_count\s*=\s*[12]\b' -n -C2
49-50: Validate external state length to match layer_countGuard against mismatched state lists to avoid indexing errors and silent shape bugs.
else: hidden_state, cell_state = state + expected_len = self.layer_count + if len(hidden_state) != expected_len or len(cell_state) != expected_len: + raise ValueError( + f"State lengths must equal layer_count ({expected_len}); " + f"got hidden={len(hidden_state)}, cell={len(cell_state)}" + )applications/models/src/models/train_tft_model.py (4)
1-11: Fix import sorting to satisfy pipeline requirements.The pipeline is failing due to Ruff I001 import ordering violation. Imports need to be grouped and sorted: standard library, third-party, then local imports.
Apply this diff to fix the import order:
-from datetime import datetime -from zoneinfo import ZoneInfo - -import polars as pl -from flytekit import task, workflow -from internal.dataset import TemporalFusionTransformerDataset -from internal.tft_model import Parameters, TemporalFusionTransformer -from loguru import logger - -import wandb +from datetime import datetime +from zoneinfo import ZoneInfo + +import polars as pl +from flytekit import task, workflow +from loguru import logger +import wandb + +from internal.dataset import TemporalFusionTransformerDataset +from internal.tft_model import Parameters, TemporalFusionTransformer
20-27: Move wandb initialization inside train_model to avoid side effects.Module-level wandb initialization creates side effects during import that can break Flyte task registration and distributed execution.
Apply this diff to remove module-level wandb initialization:
-if wandb.run is not None: - wandb.finish() # close active run if it exists - -run = wandb.init( - project="Pocket Size Fund", - config=configuration, - name=f"tft-model-run-{datetime.now(tz=ZoneInfo('America/New_York')).strftime('%Y-%m-%d_%H-%M-%S')}", -) +# wandb run will be initialized within train_model to avoid import-time side effectsAnd update the train_model function to initialize wandb locally (see next comment).
84-99: Fix multiple issues in train_model function.Several issues need to be addressed:
- Handle the return type from
model.train()which returns{"losses": list[float]}according to the TFT model interface- Move wandb initialization into this function to avoid global state conflicts
- Use
total_seconds()for accurate timing calculationApply this diff:
losses = model.train( inputs_list=batches, epoch_count=epoch_count, learning_rate=learning_rate, ) - for loss in losses: - run.log({"loss": loss}) + # Initialize wandb run within the task + if wandb.run is not None: + wandb.finish() + run = wandb.init( + project="Pocket Size Fund", + config=configuration, + name=f"tft-model-run-{datetime.now(tz=timezone).strftime('%Y-%m-%d_%H-%M-%S')}", + ) + + # Handle dict return type from model.train + loss_list = losses["losses"] if isinstance(losses, dict) else losses + for loss_value in loss_list: + run.log({"loss": loss_value}) run.finish() - runtime_seconds = (datetime.now(tz=timezone) - start_time).seconds + runtime_seconds = int((datetime.now(tz=timezone) - start_time).total_seconds())
111-116: Fix validation batch construction to use correct input length.The validation batches should use
model.parameters.input_lengthor the model'sinput_lengthattribute, notmodel.batch_size(which doesn't exist in the TFT model interface).Based on the TFT model interface from the relevant code snippets, apply this diff:
batches = data.get_batches( data_type="validation", validation_split=validation_split, - input_length=model.input_length, - output_length=model.output_length, + input_length=model.parameters.input_length, + output_length=model.parameters.output_length, )libraries/python/src/internal/mhsa_network.py (2)
13-21: Guard against invalid heads_count/embedding_size and bound dropout_rateAvoid ZeroDivisionError when heads_count == 0 and validate dropout_rate range. Also prefer casting dropout_rate to float once.
) -> None: - if embedding_size % heads_count != 0: + if heads_count <= 0: + raise ValueError("heads_count must be a positive integer") + if embedding_size <= 0: + raise ValueError("embedding_size must be a positive integer") + if embedding_size % heads_count != 0: message = "Embedding dimension must be divisible by heads count" raise ValueError(message) self.heads_count = heads_count self.embedding_size = embedding_size self.heads_dimension = embedding_size // heads_count - self.dropout_rate = dropout_rate + self.dropout_rate = float(dropout_rate) + if not (0.0 <= self.dropout_rate < 1.0): + raise ValueError("dropout_rate must be in the range [0.0, 1.0).")
30-33: Fast-fail on input embedding mismatchRaise a clear error if the last dim of inputs doesn’t match embedding_size.
def forward(self, inputs: Tensor) -> tuple[Tensor, Tensor]: - batch_size, sequence_length, _ = inputs.shape + if inputs.shape[-1] != self.embedding_size: + raise ValueError( + f"Last dimension of inputs must equal embedding_size " + f"({inputs.shape[-1]} != {self.embedding_size})" + ) + batch_size, sequence_length, _ = inputs.shapelibraries/python/tests/test_dataset.py (1)
81-134: Test coverage is incomplete for batch generation.The current test only covers
data_type="predict". Theget_batchesmethod likely supports "train" and "validate" modes as well, which should be tested to ensure complete coverage.Would you like me to generate additional test cases for the "train" and "validate" data types to ensure comprehensive coverage of the batch generation functionality?
infrastructure/environment_variables.py (2)
22-24: Export APCA_ vars too to align with Alpaca’s standard and mixed callers*There’s mixed usage of ALPACA_* and APCA_* names across the codebase. Exporting APCA_* here avoids breaking services while you migrate callers.
Apply this diff:
@@ - ("ALPACA_API_KEY_ID", alpaca_api_key_id), - ("ALPACA_API_SECRET", alpaca_api_secret), + ("ALPACA_API_KEY_ID", alpaca_api_key_id), + ("ALPACA_API_SECRET", alpaca_api_secret), + ("ALPACA_API_SECRET_KEY", alpaca_api_secret), + ("APCA_API_KEY_ID", alpaca_api_key_id), + ("APCA_API_SECRET_KEY", alpaca_api_secret),Optionally, plan a follow-up PR to unify all consumers on APCA_* and then remove ALPACA_*.
8-9: Mismatch with consumers: ALPACA_API_SECRET vs ALPACA_API_SECRET_KEYDownstream code (e.g., models scripts) reads ALPACA_API_SECRET_KEY, but this file only provisions ALPACA_API_SECRET. This will yield missing credentials at runtime.
Apply this safe, backwards-compatible change to export both variants:
@@ - ("ALPACA_API_KEY_ID", alpaca_api_key_id), - ("ALPACA_API_SECRET", alpaca_api_secret), + ("ALPACA_API_KEY_ID", alpaca_api_key_id), + ("ALPACA_API_SECRET", alpaca_api_secret), + ("ALPACA_API_SECRET_KEY", alpaca_api_secret),applications/models/src/models/get_alpaca_equity_bars.py (2)
22-32: Credential env-var mismatch and lack of validation; add fallbacks and fail fastThis script reads ALPACA_API_SECRET_KEY, but infra exports ALPACA_API_SECRET. Also, if env vars are unset, clients are created with None. Validate and support both names.
Apply this diff:
@@ - alpaca_trading_client = TradingClient( - api_key=os.getenv("ALPACA_API_KEY_ID"), - secret_key=os.getenv("ALPACA_API_SECRET_KEY"), - paper=os.getenv("ALPACA_PAPER", "true").lower() == "true", - ) + api_key_id = os.getenv("ALPACA_API_KEY_ID") or os.getenv("APCA_API_KEY_ID") + api_secret = ( + os.getenv("ALPACA_API_SECRET_KEY") + or os.getenv("ALPACA_API_SECRET") + or os.getenv("APCA_API_SECRET_KEY") + ) + if not api_key_id or not api_secret: + raise RuntimeError( + "Missing Alpaca API credentials. Set ALPACA_API_KEY_ID and " + "ALPACA_API_SECRET (or *_SECRET_KEY/APCA_* equivalents)." + ) + + alpaca_trading_client = TradingClient( + api_key=api_key_id, + secret_key=api_secret, + paper=os.getenv("ALPACA_PAPER", "true").lower() == "true", + ) @@ - alpaca_data_client = StockHistoricalDataClient( - api_key=os.getenv("ALPACA_API_KEY_ID"), - secret_key=os.getenv("ALPACA_API_SECRET_KEY"), - sandbox=os.getenv("ALPACA_PAPER", "true").lower() == "true", - ) + alpaca_data_client = StockHistoricalDataClient( + api_key=api_key_id, + secret_key=api_secret, + )Note: See next comment about verifying whether StockHistoricalDataClient supports a sandbox parameter.
93-100: Fix timestamp parsing: datetime.fromisoformat() can’t parse 'Z'Alpaca timestamps commonly end with 'Z'. datetime.fromisoformat('...Z') raises ValueError. Normalize 'Z' to '+00:00' before parsing.
Apply this diff:
- equity_bars_data = equity_bars_data.with_columns( - ( - pl.col("timestamp").map_elements( - lambda x: int(datetime.fromisoformat(x).timestamp() * 1000) - ) - ).alias("timestamp") - ) + equity_bars_data = equity_bars_data.with_columns( + ( + pl.col("timestamp").map_elements( + lambda x: int(datetime.fromisoformat(x.replace("Z", "+00:00")).timestamp() * 1000) + ) + ).alias("timestamp") + )applications/models/src/models/get_alpaca_calendar.py (1)
13-17: Credential env-var mismatch and missing validationThis script reads ALPACA_API_SECRET_KEY but infra exports ALPACA_API_SECRET. Add fallbacks and fail fast when credentials are missing to avoid confusing 401s.
Apply this diff:
@@ - alpaca_client = TradingClient( - api_key=os.getenv("ALPACA_API_KEY_ID"), - secret_key=os.getenv("ALPACA_API_SECRET_KEY"), - paper=os.getenv("ALPACA_PAPER", "true").lower() == "true", - ) + api_key_id = os.getenv("ALPACA_API_KEY_ID") or os.getenv("APCA_API_KEY_ID") + api_secret = ( + os.getenv("ALPACA_API_SECRET_KEY") + or os.getenv("ALPACA_API_SECRET") + or os.getenv("APCA_API_SECRET_KEY") + ) + if not api_key_id or not api_secret: + raise RuntimeError( + "Missing Alpaca API credentials. Set ALPACA_API_KEY_ID and " + "ALPACA_API_SECRET (or *_SECRET_KEY/APCA_* equivalents)." + ) + alpaca_client = TradingClient( + api_key=api_key_id, + secret_key=api_secret, + paper=os.getenv("ALPACA_PAPER", "true").lower() == "true", + )
🧹 Nitpick comments (14)
Dockerfile.tests (1)
30-30: Pin to lockfile during CI installsTo ensure reproducible installs in CI, consider enforcing the lockfile.
Apply this diff:
-RUN uv sync --all-packages --dev +RUN uv sync --all-packages --dev --frozenlibraries/python/tests/test_loss_functions.py (4)
11-16: Strengthen the basic test by asserting non-negativity (pinball loss property)Quantile loss should be non-negative when quantiles ∈ [0,1]. Add an assertion to make this explicit.
loss = quantile_loss(predictions, targets, quantiles) - assert isinstance(loss, Tensor) + assert isinstance(loss, Tensor) + assert loss.numpy() >= 0.0 # pinball loss is non-negative assert len(loss.shape) == 0 or loss.shape in [(), (0,), (1,)]
33-40: Assert perfect prediction is ~0 (use tolerance) instead of only non-negativeThis is a good place to be strict: perfect prediction should yield zero loss. Using a tolerance keeps the test robust to tinygrad/numpy scalar vs array nuances.
loss = quantile_loss(predictions, targets, quantiles) - assert loss.numpy() >= 0.0 + # Use tolerance to avoid flakiness across scalar/array representations + assert np.allclose(loss.numpy(), 0.0, atol=1e-7)
64-69: Broaden invalid-quantiles coverage with parametrizeAlso test the <0 boundary to complement the >1 case and reduce duplication.
-def test_quantile_loss_invalid_quantiles() -> None: - predictions = Tensor([[[1.0, 1.2, 1.3]]]) - targets = Tensor([[2.0]]) - quantiles = [0.25, 1.5, 0.75] # Invalid quantile > 1 - with pytest.raises(ValueError, match="All quantiles must be between 0 and 1"): - quantile_loss(predictions, targets, quantiles) +@pytest.mark.parametrize("bad_q", [-0.1, 1.5]) +def test_quantile_loss_invalid_quantiles(bad_q: float) -> None: + predictions = Tensor([[[1.0, 1.2, 1.3]]]) + targets = Tensor([[2.0]]) + quantiles = [0.25, bad_q, 0.75] + with pytest.raises(ValueError, match="All quantiles must be between 0 and 1"): + quantile_loss(predictions, targets, quantiles)
3-3: Pair tests with implementation validation: enforce quantile count and shape constraintsThe current implementation in libraries/python/src/internal/loss_functions.py (per the snippet) lacks:
- A check that predictions.shape[-1] equals len(quantiles)
- A check that predictions.shape[:2] matches targets.shape
Without these, the new tests for mismatches would fail to catch silent broadcasting/indexing bugs.
Proposed update to libraries/python/src/internal/loss_functions.py (outside this file):
def quantile_loss( predictions: Tensor, # shape: (batch_size, output_size, len(quantiles)) targets: Tensor, # shape: (batch_size, output_size) quantiles: list[float] | None = None, ) -> Tensor: if quantiles is None: quantiles = [0.1, 0.5, 0.9] if not all(0 <= q <= 1 for q in quantiles): raise ValueError("All quantiles must be between 0 and 1") # NEW: shape validations if predictions.shape[-1] != len(quantiles): raise ValueError( f"Quantile count mismatch: predictions last dimension ({predictions.shape[-1]}) " f"vs number of quantiles ({len(quantiles)})" ) if predictions.shape[:2] != targets.shape: raise ValueError( f"Shape mismatch: predictions batch/output {predictions.shape[:2]} " f"vs targets {targets.shape}" ) errors_total = Tensor.zeros((1,)) for index, quantile in enumerate(quantiles): error = targets.sub(predictions[:, :, index]) quantile_tensor = Tensor([quantile]) errors_total = errors_total.add( Tensor.where( error > 0, quantile_tensor.mul(error), (quantile_tensor.sub(1)).mul(error), ).mean() ) return errors_total.div(len(quantiles)) # shape: (1,)Would you like me to push commits adding the validations and the two negative tests?
libraries/python/src/internal/lstm_network.py (3)
23-30: Avoid shadowing the input_size parameter for readabilityUse a local variable for per-layer input size.
- for index in range(layer_count): - input_size = input_size if index == 0 else self.hidden_size - self.layers.append( - LSTMCell( - input_size=input_size, - hidden_size=self.hidden_size, - ) - ) + for index in range(layer_count): + layer_input_size = input_size if index == 0 else self.hidden_size + self.layers.append( + LSTMCell( + input_size=layer_input_size, + hidden_size=self.hidden_size, + ) + )
18-20: Validate dropout_rate rangeEnsure the configured dropout probability is valid.
self.hidden_size = hidden_size self.layer_count = layer_count self.dropout_rate = dropout_rate + if not (0.0 <= self.dropout_rate < 1.0): + raise ValueError("dropout_rate must be in [0.0, 1.0).")
82-83: Minor: simplify stacking callCleaner and avoids special-casing the first element.
- output_tensor = Tensor.stack(outputs[0], *outputs[1:], dim=1) + output_tensor = Tensor.stack(*outputs, dim=1)applications/models/src/models/train_tft_model.py (1)
38-38: Use total_seconds() for accurate timing calculations.Using
.secondsonly returns the seconds component and ignores sub-second precision and minutes/hours. Use.total_seconds()for accurate timing.Apply this diff to fix timing calculations:
- runtime_seconds = (datetime.now(tz=timezone) - start_time).seconds + runtime_seconds = int((datetime.now(tz=timezone) - start_time).total_seconds())This applies to lines 38, 95, 122, and 136.
Also applies to: 95-95, 122-122, 136-136
libraries/python/src/internal/mhsa_network.py (2)
28-29: Use a Python float for scale to avoid dtype/device frictionScale is a constant; storing it as a Tensor is unnecessary and can cause minor dtype/device mismatches. A plain float works with broadcasting.
- self.scale = Tensor(self.heads_dimension**0.5, dtype=dtypes.float32) + self.scale = float(self.heads_dimension ** 0.5)No other code changes needed; division by a Python float will broadcast as expected.
Also applies to: 44-46
6-12: Add a concise class docstring (shapes, returns, dropout semantics)Helps downstream users and aligns with internal library standards.
-class MultiHeadSelfAttentionNetwork: +class MultiHeadSelfAttentionNetwork: + """ + Multi-head self-attention block. + + Args: + heads_count: Number of attention heads. + embedding_size: Input/output embedding dimension per token. + dropout_rate: Dropout probability applied to attention weights and final projection. + + Forward: + inputs: Tensor of shape (batch, seq_len, embedding_size) + Returns: (output, attention_weights) + - output shape: (batch, seq_len, embedding_size) + - attention_weights shape: (batch, heads_count, seq_len, seq_len) + """infrastructure/environment_variables.py (1)
20-32: Preserve secrecy of the combined env dict explicitlyEven though secrets generally propagate through Pulumi Outputs, explicitly marking the final dict as secret avoids accidental exposure.
Apply this diff:
- return pulumi.Output.all( + return pulumi.Output.secret(pulumi.Output.all( [ ("ALPACA_API_KEY_ID", alpaca_api_key_id), ("ALPACA_API_SECRET", alpaca_api_secret), ("DATA_BUCKET_NAME", data_bucket_name), ("POLYGON_API_KEY", polygon_api_key), ("DUCKDB_ACCESS_KEY", duckdb_access_key), ("DUCKDB_SECRET", duckdb_secret), ("AWS_REGION", aws_region), ("DUCKDB_USER_ACCESS_KEY_ID", duckdb_user_access_key.id), ("DUCKDB_USER_ACCESS_KEY_SECRET", duckdb_user_access_key.secret), ] - ).apply(lambda secrets: dict(secrets)) + ).apply(lambda secrets: dict(secrets)))applications/models/src/models/get_alpaca_equity_bars.py (1)
56-71: Prefer enum constants over constructing from stringsTimeFrameUnit("Day") and Adjustment("All") rely on string coercion. Use constants to avoid subtle breakage across versions.
Apply this diff:
- timeframe=TimeFrame( - amount=1, - unit=TimeFrameUnit("Day"), - ), - adjustment=Adjustment("All"), + timeframe=TimeFrame(amount=1, unit=TimeFrameUnit.Day), + adjustment=Adjustment.ALL,applications/models/src/models/get_alpaca_calendar.py (1)
22-30: Wrap network call with error handling to surface actionable failuresWrap get_calendar in try/except to log and preserve stack traces. Keeps success path unchanged.
Apply this diff:
- calendars: list[Calendar] = cast( - "list[Calendar]", - alpaca_client.get_calendar( - GetCalendarRequest( - start=start, - end=end, - ) - ), - ) + try: + calendars: list[Calendar] = cast( + "list[Calendar]", + alpaca_client.get_calendar( + GetCalendarRequest( + start=start, + end=end, + ) + ), + ) + except Exception as exc: + logger.exception(f"Failed to fetch Alpaca calendar: {exc}") + raise
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (78)
.claude/settings.local.json(1 hunks).github/workflows/launch_infrastructure.yaml(1 hunks).github/workflows/teardown_infrastructure.yaml(1 hunks).gitignore(2 hunks).mise.toml(2 hunks)Dockerfile.tests(1 hunks)README.md(0 hunks)application/datamanager/.dockerignore(0 hunks)application/datamanager/Dockerfile(0 hunks)application/datamanager/Dockerfile.test(0 hunks)application/datamanager/compose.yaml(0 hunks)application/datamanager/features/environment.py(0 hunks)application/datamanager/features/equity_bars.feature(0 hunks)application/datamanager/features/health.feature(0 hunks)application/datamanager/features/steps/equity_bars_steps.py(0 hunks)application/datamanager/features/steps/health_steps.py(0 hunks)application/datamanager/mise.toml(0 hunks)application/datamanager/pyproject.toml(0 hunks)application/datamanager/src/datamanager/clients.py(0 hunks)application/datamanager/src/datamanager/main.py(0 hunks)application/datamanager/tests/test_datamanager_main.py(0 hunks)application/datamanager/tests/test_datamanager_models.py(0 hunks)application/positionmanager/Dockerfile(0 hunks)application/positionmanager/pyproject.toml(0 hunks)application/positionmanager/src/positionmanager/__init__.py(0 hunks)application/positionmanager/src/positionmanager/clients.py(0 hunks)application/positionmanager/src/positionmanager/main.py(0 hunks)application/positionmanager/src/positionmanager/portfolio.py(0 hunks)application/positionmanager/tests/test_positionmanager_main.py(0 hunks)application/predictionengine/Dockerfile(0 hunks)application/predictionengine/compose.yaml(0 hunks)application/predictionengine/pyproject.toml(0 hunks)application/predictionengine/src/predictionengine/dataset.py(0 hunks)application/predictionengine/src/predictionengine/gated_residual_network.py(0 hunks)application/predictionengine/src/predictionengine/loss_function.py(0 hunks)application/predictionengine/src/predictionengine/main.py(0 hunks)application/predictionengine/src/predictionengine/miniature_temporal_fusion_transformer.py(0 hunks)application/predictionengine/src/predictionengine/multi_head_self_attention.py(0 hunks)application/predictionengine/src/predictionengine/post_processor.py(0 hunks)application/predictionengine/src/predictionengine/ticker_embedding.py(0 hunks)application/predictionengine/tests/test_dataset.py(0 hunks)application/predictionengine/tests/test_gated_residual_network.py(0 hunks)application/predictionengine/tests/test_post_processor.py(0 hunks)application/predictionengine/tests/test_predictionengine_main.py(0 hunks)application/predictionengine/tests/test_ticker_embedding.py(0 hunks)applications/datamanager/pyproject.toml(1 hunks)applications/models/pyproject.toml(1 hunks)applications/models/src/models/get_alpaca_calendar.py(1 hunks)applications/models/src/models/get_alpaca_equity_bars.py(1 hunks)applications/models/src/models/train_tft_model.py(1 hunks)applications/portfoliomanager/pyproject.toml(1 hunks)cli/datamanager.py(0 hunks)cli/pyproject.toml(0 hunks)infrastructure/environment_variables.py(2 hunks)infrastructure/images.py(1 hunks)libraries/python/pyproject.toml(1 hunks)libraries/python/src/internal/cloud_event.py(1 hunks)libraries/python/src/internal/dataset.py(1 hunks)libraries/python/src/internal/dates.py(2 hunks)libraries/python/src/internal/equity_bar.py(1 hunks)libraries/python/src/internal/loss_functions.py(1 hunks)libraries/python/src/internal/lstm_network.py(2 hunks)libraries/python/src/internal/mhsa_network.py(1 hunks)libraries/python/src/internal/money.py(0 hunks)libraries/python/src/internal/summaries.py(1 hunks)libraries/python/src/internal/tft_model.py(1 hunks)libraries/python/src/internal/variable_selection_network.py(1 hunks)libraries/python/tests/test_dataset.py(1 hunks)libraries/python/tests/test_dates.py(1 hunks)libraries/python/tests/test_equity_bar.py(1 hunks)libraries/python/tests/test_loss_functions.py(4 hunks)libraries/python/tests/test_lstm_network.py(5 hunks)libraries/python/tests/test_mhsa_network.py(5 hunks)libraries/python/tests/test_variable_selection_network.py(1 hunks)pyproject.toml(1 hunks)workflows/fetch_data.py(0 hunks)workflows/pyproject.toml(0 hunks)workflows/train_predictionengine.py(0 hunks)
💤 Files with no reviewable changes (45)
- cli/pyproject.toml
- application/datamanager/.dockerignore
- application/predictionengine/tests/test_post_processor.py
- application/predictionengine/src/predictionengine/miniature_temporal_fusion_transformer.py
- application/datamanager/features/environment.py
- application/predictionengine/pyproject.toml
- application/datamanager/compose.yaml
- application/positionmanager/pyproject.toml
- application/datamanager/features/steps/health_steps.py
- application/positionmanager/src/positionmanager/init.py
- application/datamanager/Dockerfile.test
- application/predictionengine/tests/test_predictionengine_main.py
- application/predictionengine/src/predictionengine/ticker_embedding.py
- application/datamanager/Dockerfile
- README.md
- application/datamanager/features/health.feature
- application/positionmanager/Dockerfile
- application/datamanager/tests/test_datamanager_models.py
- workflows/pyproject.toml
- application/predictionengine/compose.yaml
- application/datamanager/src/datamanager/clients.py
- application/predictionengine/src/predictionengine/gated_residual_network.py
- application/datamanager/mise.toml
- application/datamanager/features/equity_bars.feature
- application/positionmanager/tests/test_positionmanager_main.py
- application/positionmanager/src/positionmanager/portfolio.py
- application/datamanager/pyproject.toml
- workflows/train_predictionengine.py
- application/predictionengine/src/predictionengine/loss_function.py
- application/predictionengine/Dockerfile
- application/predictionengine/src/predictionengine/dataset.py
- application/datamanager/src/datamanager/main.py
- application/datamanager/tests/test_datamanager_main.py
- application/positionmanager/src/positionmanager/clients.py
- libraries/python/src/internal/money.py
- application/predictionengine/src/predictionengine/post_processor.py
- application/predictionengine/tests/test_gated_residual_network.py
- application/positionmanager/src/positionmanager/main.py
- application/predictionengine/src/predictionengine/main.py
- cli/datamanager.py
- application/predictionengine/tests/test_ticker_embedding.py
- application/datamanager/features/steps/equity_bars_steps.py
- workflows/fetch_data.py
- application/predictionengine/tests/test_dataset.py
- application/predictionengine/src/predictionengine/multi_head_self_attention.py
✅ Files skipped from review due to trivial changes (2)
- applications/portfoliomanager/pyproject.toml
- applications/datamanager/pyproject.toml
🚧 Files skipped from review as they are similar to previous changes (22)
- libraries/python/src/internal/summaries.py
- libraries/python/pyproject.toml
- libraries/python/src/internal/cloud_event.py
- infrastructure/images.py
- .github/workflows/launch_infrastructure.yaml
- libraries/python/src/internal/tft_model.py
- pyproject.toml
- .mise.toml
- libraries/python/src/internal/variable_selection_network.py
- .claude/settings.local.json
- libraries/python/tests/test_dates.py
- libraries/python/tests/test_mhsa_network.py
- libraries/python/src/internal/dataset.py
- .gitignore
- libraries/python/src/internal/dates.py
- libraries/python/tests/test_lstm_network.py
- libraries/python/tests/test_variable_selection_network.py
- .github/workflows/teardown_infrastructure.yaml
- applications/models/pyproject.toml
- libraries/python/src/internal/equity_bar.py
- libraries/python/tests/test_equity_bar.py
- libraries/python/src/internal/loss_functions.py
🧰 Additional context used
🧬 Code Graph Analysis (6)
libraries/python/tests/test_dataset.py (1)
libraries/python/src/internal/dataset.py (2)
TemporalFusionTransformerDataset(28-326)get_dimensions(240-248)
libraries/python/tests/test_loss_functions.py (2)
libraries/python/src/internal/loss_functions.py (1)
quantile_loss(4-28)application/predictionengine/src/predictionengine/loss_function.py (1)
quantile_loss(8-29)
applications/models/src/models/get_alpaca_equity_bars.py (4)
application/datamanager/src/datamanager/main.py (2)
get_equity_bars(149-202)fetch_equity_bars(206-269)workflows/train_predictionengine.py (1)
fetch_data(22-67)application/datamanager/tests/test_datamanager_main.py (1)
test_fetch_equity_bars_success(199-226)cli/datamanager.py (1)
get_equity_bars(54-59)
applications/models/src/models/train_tft_model.py (2)
libraries/python/src/internal/dataset.py (3)
TemporalFusionTransformerDataset(28-326)get_dimensions(240-248)get_batches(250-326)libraries/python/src/internal/tft_model.py (5)
Parameters(19-33)TemporalFusionTransformer(37-271)train(192-225)validate(227-246)save(248-253)
libraries/python/src/internal/lstm_network.py (2)
application/predictionengine/src/predictionengine/long_short_term_memory.py (3)
LongShortTermMemory(5-69)forward(22-69)__init__(6-20)application/predictionengine/tests/test_long_short_term_memory.py (4)
test_lstm_forward(31-43)test_lstm_consistency(81-92)test_lstm_multiple_layers(58-67)test_lstm_single_timestep(70-78)
libraries/python/src/internal/mhsa_network.py (2)
application/predictionengine/src/predictionengine/multi_head_self_attention.py (3)
MultiHeadSelfAttention(8-66)__init__(9-30)forward(32-66)application/predictionengine/tests/test_multi_head_self_attention.py (6)
test_multi_head_attention_different_heads(34-45)test_multi_head_attention_forward(20-31)test_multi_head_attention_batch_processing(67-75)test_multi_head_attention_single_sequence(48-54)test_multi_head_attention_initialization(11-17)test_multi_head_attention_longer_sequences(57-64)
🪛 GitHub Check: Run Python quality checks
applications/models/src/models/train_tft_model.py
[failure] 1-10: Ruff (I001)
applications/models/src/models/train_tft_model.py:1:1: I001 Import block is un-sorted or un-formatted
🪛 GitHub Actions: Code quality check
applications/models/src/models/train_tft_model.py
[error] 1-1: Command 'ruff check --output-format=github .' failed. Ruff I001: Import block is un-sorted or un-formatted.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Run Python tests
🔇 Additional comments (17)
Dockerfile.tests (3)
20-28: Good caching-friendly COPY of workspace pyproject.toml filesCopying only the pyproject.toml files for each workspace member before
uv syncis the right move to preserve Docker layer caching. Paths also reflect the new applications/* and libraries/python layout.
32-34: Nice: post-sync copy of full sources and inclusion of libraries testsCopying full sources after
uv synckeeps the dependency layer cacheable, and addingCOPY libraries/ensures internal library tests are available in the image.
20-24: All referenced pyproject.toml files are present — no action neededVerified with the provided script output; all three files exist:
- applications/datamanager/pyproject.toml — OK
- applications/portfoliomanager/pyproject.toml — OK
- applications/models/pyproject.toml — OK
libraries/python/tests/test_loss_functions.py (3)
3-3: Import path migration looks correct; ensure test environment resolves internal moduleThe import reflects the new internal library layout. Confirm that the test runner installs/paths include libraries/python/src so
internal.loss_functionsresolves under uv/pyproject configuration.
22-27: Multi-sample, multi-quantile case reads wellCovers the batched path with Q=3 and matches the new interface. No issues spotted.
43-49: Broader quantile coverage LGTMNice coverage for Q=5; this exercises the loop over quantiles and keeps the non-negativity check.
libraries/python/src/internal/mhsa_network.py (3)
37-43: LGTM on head reshaping and transpositionThe view/transpose pipeline correctly produces (batch, heads, seq_len, head_dim) for Q/K/V; aligns with expected attention shapes.
44-49: Scaled dot-product attention is correctDividing by sqrt(d_k) and softmax along the last axis matches spec.
64-64: Overall: solid, readable implementationClear structure, minimal allocations, and returns both output and attention weights. With the input/param validations added, this should be robust.
libraries/python/tests/test_dataset.py (5)
1-2: LGTM!Clean imports and appropriate testing module structure.
5-47: LGTM!The test validates the basic instantiation and attribute presence after data loading. The test data includes all required columns with realistic financial data across multiple tickers and time periods.
49-79: LGTM!Comprehensive validation of the dimensions dictionary structure, ensuring all required feature group keys are present in the returned dimensions.
111-115: Signature verified — get_batches accepts data_type, input_length, and output_lengthThe test call matches the implementation.
- Location: libraries/python/src/internal/dataset.py —
def get_batches(self, data_type: str = "train", validation_split: float = 0.8, input_length: int = 35, output_length: int = 7) -> list[dict[str, Tensor]](defined at/around line 250).
126-134: Validate .shape assertions — OK (no change required)get_batches constructs tinygrad.Tensor instances from numpy arrays, and tinygrad.Tensor (v0.10.3) exposes a .shape property — the assertions in test_dataset.py are therefore meaningful.
- libraries/python/src/internal/dataset.py — get_batches: creates Tensor(...) for encoder_categorical_features, encoder_continuous_features, decoder_categorical_features, static_categorical_features, and targets.
- libraries/python/tests/test_dataset.py — test_dataset_batches: asserts .shape[0] on those tensors.
applications/models/src/models/get_alpaca_equity_bars.py (3)
34-43: Check GetAssetsRequest.attributes type; prefer enums/list over stringPassing attributes="has_options" may be brittle if the SDK expects a list or enum values (e.g., AssetAttributes.HAS_OPTIONS).
Please confirm the expected type in the installed alpaca-py version. If it expects enums, adjust to:
from alpaca.trading.requests import GetAssetsRequest from alpaca.trading.enums import AssetAttributes GetAssetsRequest( status=AssetStatus.ACTIVE, asset_class=AssetClass.US_EQUITY, attributes=[AssetAttributes.HAS_OPTIONS], )
112-118: LGTM: robust aggregation with saved-files guardGood job tracking written files and guarding against the empty case before concatenation. Prevents file-not-found errors and clarifies logs.
28-32: No change required — StockHistoricalDataClient supports thesandboxparameterVerified: alpaca-py’s StockHistoricalDataClient constructor includes a sandbox: bool parameter; the instantiation in applications/models/src/models/get_alpaca_equity_bars.py (the existing sandbox=os.getenv("ALPACA_PAPER", "true").lower() == "true") is valid.
| end = datetime.now(tz=ZoneInfo("America/New_York")) | ||
| start = end - timedelta(days=365 * 6) | ||
|
|
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Confirm GetCalendarRequest expects date vs datetime
The Trading calendar API typically works with date-only parameters. Passing timezone-aware datetimes should work via SDK coercion, but confirm your alpaca-py version’s types to avoid subtle TZ issues.
Run:
If it expects date, consider start.date() / end.date().
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find alpaca-py in pyproject to infer version:"
rg -n -C2 -S 'alpaca-py' --type toml
echo
echo "Search for Calendar/GetCalendarRequest usage hints:"
rg -n -C2 -S 'GetCalendarRequest\('Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- Show target file ----"
sed -n '1,240p' applications/models/src/models/get_alpaca_calendar.py || true
echo
echo "---- Search for GetCalendarRequest / GetCalendar usage ----"
rg -n --hidden -S 'GetCalendarRequest|GetCalendar' || true
echo
echo "---- Search for alpaca package usages/imports (python files, pyproject, requirements, lockfiles) ----"
rg -n --hidden -S "from .*alpaca|import .*alpaca|alpaca-py|alpaca_trade_api|alpaca" --glob '!**/node_modules/**' --glob '!**/.git/**' || true
echo
echo "---- Show pyproject.toml and poetry.lock if present ----"
for f in pyproject.toml poetry.lock Pipfile requirements.txt requirements-dev.txt; do
if [ -f "$f" ]; then
echo "---- $f ----"
sed -n '1,240p' "$f"
fi
done
echo
echo "---- Search installed package typing hints / stubs in repo (patterns like GetCalendarRequest) ----"
rg -n --hidden -S 'GetCalendarRequest' || trueLength of output: 9356
🌐 Web query:
alpaca-py "GetCalendarRequest" start end date datetime "alpaca.trading.requests"
💡 Result:
Short answer — use alpaca.trading.requests.GetCalendarRequest with Python date objects (datetime.date). You can pass datetime.datetime.date() if you have datetimes. Example usage:
Notes from the docs: GetCalendarRequest signature is
class alpaca.trading.requests.GetCalendarRequest(*, start: Optional[date] = None, end: Optional[date] = None). The calendar API can be filtered by these start/end dates. (alpaca.markets)
Examples
- Using datetime.date directly
from datetime import date
from alpaca.trading.client import TradingClient
from alpaca.trading.requests import GetCalendarRequest
client = TradingClient(API_KEY, API_SECRET, paper=True)
req = GetCalendarRequest(start=date(2025, 8, 1), end=date(2025, 8, 15))
calendar = client.get_calendar(req)
for day in calendar:
print(day) # Calendar objects (has date, open, close, etc.)- Converting from datetime.datetime
from datetime import datetime
# assume dt_start/dt_end are datetime objects
req = GetCalendarRequest(start=dt_start.date(), end=dt_end.date())
calendar = client.get_calendar(req)- Converting from ISO string
from datetime import datetime
req = GetCalendarRequest(
start=datetime.strptime("2025-08-01", "%Y-%m-%d").date(),
end=datetime.strptime("2025-08-15", "%Y-%m-%d").date(),
)
calendar = client.get_calendar(req)Other tips
- start and end are date objects (no time component). (alpaca.markets)
- If you omit both start and end you get the full calendar (docs note coverage from 1970 through 2029). (alpaca.markets)
If you want, tell me what form your dates are in (strings, datetimes, pandas timestamps) and I’ll give the exact conversion code you need.
Citations:
- 1: https://alpaca.markets/sdks/python/api_reference/trading/requests.html?utm_source=chatgpt.com
- 2: https://alpaca.markets/sdks/python/api_reference/trading/requests.html?utm_source=chatgpt.com
- 3: https://alpaca.markets/sdks/python/api_reference/trading/calendar.html?utm_source=chatgpt.com
Convert start/end to date objects for GetCalendarRequest
GetCalendarRequest.start/end expect datetime.date (no time/tz). Passing timezone-aware datetimes risks subtle TZ/truncation issues — convert to .date().
- File: applications/models/src/models/get_alpaca_calendar.py — where end/start are computed and passed to GetCalendarRequest.
Minimal fix (change only the request args):
GetCalendarRequest(
start=start.date(),
end=end.date(),
)
Or clearer refactor:
end_dt = datetime.now(tz=ZoneInfo("America/New_York"))
start_dt = end_dt - timedelta(days=365 * 6)
alpaca_client.get_calendar(
GetCalendarRequest(start=start_dt.date(), end=end_dt.date())
)
🤖 Prompt for AI Agents
In applications/models/src/models/get_alpaca_calendar.py around lines 19 to 21,
the code computes timezone-aware datetimes for start/end and passes them to
GetCalendarRequest which expects datetime.date; update the call to pass dates
instead of datetimes by converting to .date() (minimal change: pass start.date()
and end.date() into GetCalendarRequest), or rename the variables to
end_dt/start_dt and call .date() when building the request to avoid
timezone/truncation issues.
| if self.dropout_rate > 0: | ||
| attention_weights = attention_weights.dropout(self.dropout_rate) | ||
|
|
||
| attention_output = attention_weights.matmul(value_weights) | ||
|
|
||
| attention_output = attention_output.transpose(1, 2).reshape( | ||
| batch_size, sequence_length, self.embedding_size | ||
| ) | ||
|
|
||
| output = self.fully_connected_out(attention_output) | ||
|
|
||
| if self.dropout_rate > 0: | ||
| output = output.dropout(self.dropout_rate) | ||
|
|
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Dropout behavior: ensure it only applies during training
Confirm that tinygrad’s Tensor.dropout respects the global training flag. If not, gate dropout explicitly to avoid applying it during evaluation.
To verify, please check tinygrad’s dropout semantics and whether tests toggle training mode:
🏁 Script executed:
#!/bin/bash
# 1) Inspect usage of Tensor.dropout and training flag across the repo
rg -n -C2 'Tensor\.training|\.dropout\(' --type=py
# 2) Optionally, inspect tinygrad docs or source if vendored (if present in repo)
fd tinygrad | head -n 50Length of output: 1396
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== mhsa_network.py (head) ==="
sed -n '1,240p' libraries/python/src/internal/mhsa_network.py || true
echo
echo "=== lstm_network.py (head) ==="
sed -n '1,240p' libraries/python/src/internal/lstm_network.py || true
echo
echo "=== Search for Tensor class or dropout implementation ==="
rg -n --hidden -S 'class\s+Tensor\b' --type=py || true
rg -n --hidden -S 'def\s+dropout\(' --type=py || true
echo
echo "=== Search for imports mentioning tinygrad ==="
rg -n --hidden -S '\btinygrad\b' --type=py || true
echo
echo "=== Search for .train() / .eval() / .training usage ==="
rg -n --hidden -S '\.train\(' --type=py || true
rg -n --hidden -S '\.eval\(' --type=py || true
rg -n --hidden -S '\.training\b' --type=py || true
echo
echo "=== Check repo dependency files for tinygrad (requirements/pyproject) ==="
rg -n --hidden -S 'tinygrad' -g '!**/node_modules/**' || trueLength of output: 10032
🌐 Web query:
tinygrad Tensor.dropout implementation
💡 Result:
Below is a compact, practical implementation of dropout you can drop into tinygrad's Tensor class. It uses "inverted dropout" (scale by 1/(1-p) at train time) so you don't need to change anything at eval time. The mask is a non-trainable Tensor so autograd just multiplies incoming gradients by the same mask during backward.
Notes:
- This is a safe, non-destructive default (no inplace changes) — an optional inplace branch is included if you need it.
- You can pass broadcast_dims to create a mask that broadcasts across some dimensions (useful for channel/spatial dropout).
- If your tinygrad build uses different attribute names (e.g. .data vs .numpy()), adapt the small accesses accordingly.
Implementation (add to Tensor class):
import numpy as np
from tinygrad.tensor import Tensor # adjust import path as needed
def dropout(self, p: float = 0.5, training: bool = True, inplace: bool = False,
broadcast_dims: tuple = None, seed: int = None):
"""
Inverted dropout:
- p: dropout probability (0 <= p < 1)
- training: if False, returns self unchanged
- inplace: if True, tries to modify self in-place (may break autograd if used incorrectly)
- broadcast_dims: tuple of dimension indices to broadcast mask across (e.g. (0,2,3) for channel-wise mask on NCHW)
- seed: optional RNG seed for reproducibility
"""
if not training or p == 0.0:
return self
if p < 0.0 or p >= 1.0:
raise ValueError("p must be in [0, 1).")
if seed is not None:
rng = np.random.RandomState(seed)
rand = rng.rand
else:
rand = np.random.rand
# determine mask shape: either full shape or shape with ones for broadcast dims
if broadcast_dims is None:
mask_shape = tuple(self.shape)
else:
mask_shape = tuple(1 if i in broadcast_dims else s for i, s in enumerate(self.shape))
# create mask, scale by 1/(1-p) for inverted dropout
# use underlying dtype from tensor (adjust if your API differs)
dtype = getattr(self, "data", None)
if dtype is not None:
npdtype = self.data.dtype
else:
# fallback if .data not present
npdtype = self.numpy().dtype
mask_np = (rand(*mask_shape) > p).astype(npdtype) / (1.0 - p)
mask_t = Tensor(mask_np, requires_grad=False)
if inplace:
# attempt in-place multiplication (modify internal storage)
# Make sure your tinygrad supports .data modification safely
self.data *= mask_t.data
return self
else:
return self * mask_t
Example usage:
x = Tensor(np.ones((4, 3, 8, 8)), requires_grad=True) # sample NCHW
y = x.dropout(p=0.2, training=True) # elementwise dropout
# channel-wise dropout (same mask for every spatial location and batch):
y_ch = x.dropout(p=0.2, training=True, broadcast_dims=(0,2,3))
Why this works with autograd:
- mask_t is created with requires_grad=False so it is treated as a constant; the forward multiplies x by mask_t.
- during backward, gradients are multiplied by the same mask (and scaled), which implements dropout's gradient behavior.
Small test (sanity):
import numpy as np
x = Tensor(np.ones((10000,)), requires_grad=True)
y = x.dropout(p=0.5, training=True, seed=42)
print(y.numpy().mean()) # should be close to 1.0 because of inverted scaling
If you want the exact tinygrad integration (e.g. to use tinygrad RNG ops instead of numpy or to match your library internals) paste your current Tensor class here and I can adapt the patch precisely.
Dropout must be gated to run only in training mode (fix required)
MHSA currently calls Tensor.dropout unconditionally; ensure dropout runs only during training (either pass a training flag into forward and to Tensor.dropout, or conditionally call dropout). LSTM calls .train() on the hidden state before dropout, so follow a consistent, explicit pattern.
Files to change:
- libraries/python/src/internal/mhsa_network.py — forward (attention_weights and output dropout calls, ~lines 49–63)
- Update callers/tests (e.g. libraries/python/tests/test_mhsa_network.py) to pass the training flag or use model.train()/eval() as appropriate.
Suggested change (minimal example — prefer explicit training arg):
Replace forward signature and the unconditional drops:
from
def forward(self, inputs: Tensor) -> tuple[Tensor, Tensor]:
...
attention_weights = attention_scores.softmax(axis=-1)
if self.dropout_rate > 0:
attention_weights = attention_weights.dropout(self.dropout_rate)
...
if self.dropout_rate > 0:
output = output.dropout(self.dropout_rate)
return output, attention_weights
with
def forward(self, inputs: Tensor, training: bool = True) -> tuple[Tensor, Tensor]:
...
attention_weights = attention_scores.softmax(axis=-1)
if self.dropout_rate > 0 and training:
attention_weights = attention_weights.dropout(self.dropout_rate)
...
if self.dropout_rate > 0 and training:
output = output.dropout(self.dropout_rate)
return output, attention_weights
Or, if tinygrad's Tensor.dropout supports a training kwarg, use:
attention_weights = attention_weights.dropout(self.dropout_rate, training=training)
output = output.dropout(self.dropout_rate, training=training)
Remember to update all call sites/tests to pass the training flag (or call the appropriate .train()/.eval() API) so inference doesn't apply dropout.
🤖 Prompt for AI Agents
In libraries/python/src/internal/mhsa_network.py around lines 50 to 63, dropout
is applied unconditionally to attention_weights and output; change the forward
signature to accept a training: bool = True (or use Tensor.dropout's training
kwarg if available) and guard both dropout calls so they only run when
self.dropout_rate > 0 and training is True (or pass training into dropout).
Update all callers and tests (e.g., libraries/python/tests/test_mhsa_network.py)
to pass the training flag or switch to model.train()/eval() so inference does
not apply dropout.
|
Closing because it's gotten too convoluted with review feedback - opening a fresh one to do reviews there. |
forstmeier
deleted the
07-28-restructure_python_application_and_library_resources
branch
Overview
Changes
pyproject.tomlconfigurations for "applications" and "libraries"create_cloud_event_*)models"application" resources for initial model training workflow scriptapplicationtoapplications(for my own sanity)Comments
This is a large block of updates but necessary to improve the
uvconfigurations. Additionally, there's a newmodelspackage that has stuff for TFT in it (still not fully training correctly but I'll debug that and have it working soon). I also need to do a manual test/refactor of the services inapplicationsindependent of infrastructure being up so be forewarned. The Flyte stuff was complaining about the values being passed between the tasks for some reason too.Summary by CodeRabbit
New Features
Refactor
Chores
Removed
Tests